additional test coverage

pola-rs · May 12, 2024 · 30f1d19 · 30f1d19
1 parent a1d34d1
commit 30f1d19
Show file tree

Hide file tree

Showing 2 changed files with 50 additions and 13 deletions.
diff --git a/crates/polars-sql/src/context.rs b/crates/polars-sql/src/context.rs
@@ -8,8 +8,8 @@ use polars_plan::prelude::*;
 use sqlparser::ast::{
     Distinct, ExcludeSelectItem, Expr as SQLExpr, FunctionArg, GroupByExpr, JoinOperator,
     ObjectName, ObjectType, Offset, OrderByExpr, Query, Select, SelectItem, SetExpr, SetOperator,
-    SetQuantifier, Statement, TableAlias, TableFactor, TableWithJoins, Value as SQLValue,
-    WildcardAdditionalOptions,
+    SetQuantifier, Statement, TableAlias, TableFactor, TableWithJoins, UnaryOperator,
+    Value as SQLValue, WildcardAdditionalOptions,
 };
 use sqlparser::dialect::GenericDialect;
 use sqlparser::parser::{Parser, ParserOptions};
@@ -412,23 +412,31 @@ impl SQLContext {
         match &select_stmt.group_by {
             // Standard "GROUP BY x, y, z" syntax
             GroupByExpr::Expressions(group_by_exprs) => {
-                group_by_keys = group_by_exprs.iter()
+                group_by_keys = group_by_exprs
+                    .iter()
                     .map(|e| match e {
+                        SQLExpr::UnaryOp {
+                            op: UnaryOperator::Minus,
+                            expr,
+                        } if matches!(**expr, SQLExpr::Value(SQLValue::Number(_, _))) => {
+                            if let SQLExpr::Value(SQLValue::Number(ref idx, _)) = **expr {
+                                Err(polars_err!(
+                                ComputeError:
+                                "group_by error: expected a positive integer or valid expression; got -{}",
+                                idx
+                                ))
+                            } else {
+                                unreachable!()
+                            }
+                        },
                         SQLExpr::Value(SQLValue::Number(idx, _)) => {
-                            let idx = match idx.parse::<usize>() {
-                                Ok(0) | Err(_) => Err(polars_err!(
-                                    ComputeError:
-                                    "group_by error: a positive number or an expression expected, got {}",
-                                    idx
-                                )),
-                                Ok(idx) => Ok(idx),
-                            }?;
                             // note: sql queries are 1-indexed
+                            let idx = idx.parse::<usize>().unwrap();
                             Ok(projections[idx - 1].clone())
                         },
-                        SQLExpr::Value(_) => Err(polars_err!(
+                        SQLExpr::Value(v) => Err(polars_err!(
                             ComputeError:
-                            "group_by error: a positive number or an expression expected",
+                            "group_by error: expected a positive integer or valid expression; got {}", v,
                         )),
                         _ => parse_sql_expr(e, self, schema.as_deref()),
                     })

diff --git a/py-polars/tests/unit/sql/test_group_by.py b/py-polars/tests/unit/sql/test_group_by.py
@@ -6,6 +6,7 @@
 import pytest
 
 import polars as pl
+from polars.exceptions import ComputeError
 from polars.testing import assert_frame_equal
 
 
@@ -196,3 +197,31 @@ def test_group_by_ordinal_position() -> None:
             SELECT c, total_b FROM grp ORDER BY c"""
         )
         assert_frame_equal(res2, expected)
+
+
+def test_group_by_errors() -> None:
+    df = pl.DataFrame(
+        {
+            "a": ["xx", "yy", "xx"],
+            "b": [10, 20, 30],
+            "c": [99, 99, 66],
+        }
+    )
+
+    with pytest.raises(
+        ComputeError,
+        match=r"expected a positive integer or valid expression; got -99",
+    ):
+        df.sql("SELECT a, SUM(b) FROM self GROUP BY -99, a")
+
+    with pytest.raises(
+        ComputeError,
+        match=r"expected a positive integer or valid expression; got '!!!'",
+    ):
+        df.sql("SELECT a, SUM(b) FROM self GROUP BY a, '!!!'")
+
+    with pytest.raises(
+        ComputeError,
+        match=r"'a' should participate in the GROUP BY clause or an aggregate function",
+    ):
+        df.sql("SELECT a, SUM(b) FROM self GROUP BY b")