fix: Ensure height is maintained in SQL SELECT 1 FROM (#20241)

pola-rs · Dec 10, 2024 · 296b4d4 · 296b4d4
1 parent a21ac2e
commit 296b4d4
Show file tree

Hide file tree

Showing 5 changed files with 37 additions and 3 deletions.
diff --git a/crates/polars-plan/src/plans/aexpr/scalar.rs b/crates/polars-plan/src/plans/aexpr/scalar.rs
@@ -8,10 +8,14 @@ pub fn is_scalar_ae(node: Node, expr_arena: &Arena<AExpr>) -> bool {
         AExpr::Literal(lv) => lv.is_scalar(),
         AExpr::Function { options, input, .. }
         | AExpr::AnonymousFunction { options, input, .. } => {
-            if options.is_elementwise() || !options.flags.contains(FunctionFlags::CHANGES_LENGTH) {
+            if options.flags.contains(FunctionFlags::RETURNS_SCALAR) {
+                true
+            } else if options.is_elementwise()
+                || !options.flags.contains(FunctionFlags::CHANGES_LENGTH)
+            {
                 input.iter().all(|e| e.is_scalar(expr_arena))
             } else {
-                options.flags.contains(FunctionFlags::RETURNS_SCALAR)
+                false
             }
         },
         AExpr::BinaryExpr { left, right, .. } => {

diff --git a/crates/polars-sql/src/context.rs b/crates/polars-sql/src/context.rs
@@ -721,7 +721,9 @@ impl SQLContext {
 
             // Final/selected cols, accounting for 'SELECT *' modifiers
             let mut retained_cols = Vec::with_capacity(projections.len());
+            let mut retained_names = Vec::with_capacity(projections.len());
             let have_order_by = query.order_by.is_some();
+            let mut all_literal = true;
 
             // Note: if there is an 'order by' then we project everything (original cols
             // and new projections) and *then* select the final cols; the retained cols
@@ -735,11 +737,13 @@ impl SQLContext {
                 if select_modifiers.matches_ilike(&name)
                     && !select_modifiers.exclude.contains(&name)
                 {
+                    all_literal &= expr_to_leaf_column_names_iter(p).next().is_none();
                     retained_cols.push(if have_order_by {
                         col(name.as_str())
                     } else {
                         p.clone()
                     });
+                    retained_names.push(col(name));
                 }
             }
 
@@ -755,7 +759,12 @@ impl SQLContext {
             }
 
             lf = self.process_order_by(lf, &query.order_by, Some(&retained_cols))?;
-            lf = lf.select(retained_cols);
+
+            if all_literal && !have_order_by {
+                lf = lf.with_columns(retained_cols).select(retained_names);
+            } else {
+                lf = lf.select(retained_cols);
+            }
 
             if !select_modifiers.rename.is_empty() {
                 lf = lf.rename(

diff --git a/crates/polars-sql/tests/simple_exprs.rs b/crates/polars-sql/tests/simple_exprs.rs
@@ -202,6 +202,7 @@ fn test_literal_exprs() {
     let df_pl = df
         .lazy()
         .select(&[
+            Expr::Nth(0),
             lit(1i64).alias("int_lit"),
             lit(1.0).alias("float_lit"),
             lit("foo").alias("string_lit"),
@@ -210,6 +211,10 @@ fn test_literal_exprs() {
             lit(Duration::parse("1q2w1d50s")).alias("duration_lit"),
         ])
         .collect()
+        .unwrap()
+        .lazy()
+        .drop([Expr::Nth(0)])
+        .collect()
         .unwrap();
     assert!(df_sql.equals_missing(&df_pl));
 }

diff --git a/py-polars/tests/unit/sql/test_literals.py b/py-polars/tests/unit/sql/test_literals.py
@@ -168,3 +168,10 @@ def test_select_from_table_with_reserved_names() -> None:
         eager=True,
     )
     assert out.rows() == [(5, 2)]
+
+
+def test_select_literal_20058() -> None:
+    assert_frame_equal(
+        pl.LazyFrame({"a": [1, 2, 3]}).sql("SELECT 1 FROM self").collect(),
+        pl.select(literal=pl.Series([1, 1, 1], dtype=pl.Int32)),
+    )
diff --git a/py-polars/tests/unit/test_scalar.py b/py-polars/tests/unit/test_scalar.py
@@ -1,6 +1,7 @@
 import pytest
 
 import polars as pl
+from polars.testing import assert_frame_equal
 
 
 @pytest.mark.may_fail_auto_streaming
@@ -70,3 +71,11 @@ def test_scalar_len_20046() -> None:
     )
 
     assert q.select(pl.len()).collect().item() == 3
+
+
+def test_scalar_identification_function_expr_in_binary() -> None:
+    x = pl.Series("x", [1, 2, 3])
+    assert_frame_equal(
+        pl.select(x).with_columns(o=pl.col("x").null_count() > 0),
+        pl.select(x, o=False),
+    )