feat: Support default alias for non-column refs for sql select

Author: cancai

src/daft-sql/src/expr_name.rs

@@ -0,0 +1,139 @@
+use sqlparser::ast::{
+    BinaryOperator, Expr as SQLExpr, FunctionArg, FunctionArgExpr, FunctionArguments, Ident,
+    ObjectName, ObjectNamePart, UnaryOperator, Value, ValueWithSpan,
+};
+
+fn ident_to_string(ident: &Ident) -> String {
+    ident.value.clone()
+}
+
+fn object_name_to_string(name: &ObjectName) -> String {
+    name.0
+        .iter()
+        .map(|part| match part {
+            ObjectNamePart::Identifier(ident) => ident_to_string(ident),
+            ObjectNamePart::Function(func) => func.name.value.clone(),
+        })
+        .collect::<Vec<_>>()
+        .join(".")
+}
+
+fn value_to_string(v: &Value) -> String {
+    match v {
+        Value::Number(n, _) => n.clone(),
+        Value::Boolean(b) => b.to_string(),
+        Value::Null => "null".to_string(),
+        Value::SingleQuotedString(s) | Value::DoubleQuotedString(s) => format!("\"{s}\""),
+        other => format!("{other}"),
+    }
+}
+
+fn value_with_span_to_string(v: &ValueWithSpan) -> String {
+    value_to_string(&v.value)
+}
+
+fn binary_op_to_string(op: &BinaryOperator) -> String {
+    match op {
+        BinaryOperator::Plus => "+".to_string(),
+        BinaryOperator::Minus => "-".to_string(),
+        BinaryOperator::Multiply => "*".to_string(),
+        BinaryOperator::Divide => "/".to_string(),
+        BinaryOperator::Modulo => "%".to_string(),
+        BinaryOperator::StringConcat => "||".to_string(),
+        other => format!("{other}"),
+    }
+}
+
+fn unary_op_to_string(op: &UnaryOperator) -> String {
+    match op {
+        UnaryOperator::Plus => "+".to_string(),
+        UnaryOperator::Minus => "-".to_string(),
+        UnaryOperator::Not => "not".to_string(),
+        other => format!("{other}"),
+    }
+}
+
+fn expr_to_string(e: &SQLExpr) -> String {
+    match e {
+        SQLExpr::Identifier(ident) => ident_to_string(ident),
+        SQLExpr::CompoundIdentifier(idents) => idents
+            .iter()
+            .map(ident_to_string)
+            .collect::<Vec<_>>()
+            .join("."),
+        SQLExpr::Value(v) => value_with_span_to_string(v),
+
+        SQLExpr::UnaryOp { op, expr } => {
+            let op = unary_op_to_string(op);
+            format!("({op} {})", expr_to_string(expr))
+        }
+        SQLExpr::BinaryOp { left, op, right } => {
+            let op = binary_op_to_string(op);
+            format!("({} {op} {})", expr_to_string(left), expr_to_string(right))
+        }
+
+        SQLExpr::Nested(inner) => expr_to_string(inner),
+        SQLExpr::Cast {
+            expr, data_type, ..
+        } => {
+            format!("cast({} as {data_type})", expr_to_string(expr))
+        }
+
+        SQLExpr::Function(func) => {
+            let fn_name = object_name_to_string(&func.name).to_lowercase();
+
+            // Special-case COUNT(*) to preserve the star.
+            if fn_name == "count"
+                && matches!(
+                    &func.args,
+                    FunctionArguments::List(args)
+                        if args.args.len() == 1
+                            && matches!(&args.args[0], FunctionArg::Unnamed(FunctionArgExpr::Wildcard))
+                )
+            {
+                return "count(*)".to_string();
+            }
+
+            let args: Vec<String> = match &func.args {
+                FunctionArguments::None => vec![],
+                FunctionArguments::Subquery(_) => vec!["<subquery>".to_string()],
+                FunctionArguments::List(args) => args
+                    .args
+                    .iter()
+                    .map(|arg| match arg {
+                        FunctionArg::Named { arg, .. } | FunctionArg::ExprNamed { arg, .. } => {
+                            match arg {
+                                FunctionArgExpr::Expr(e) => expr_to_string(e),
+                                FunctionArgExpr::QualifiedWildcard(obj) => {
+                                    format!("{}.*", object_name_to_string(obj))
+                                }
+                                FunctionArgExpr::Wildcard => "*".to_string(),
+                            }
+                        }
+                        FunctionArg::Unnamed(arg) => match arg {
+                            FunctionArgExpr::Expr(e) => expr_to_string(e),
+                            FunctionArgExpr::QualifiedWildcard(obj) => {
+                                format!("{}.*", object_name_to_string(obj))
+                            }
+                            FunctionArgExpr::Wildcard => "*".to_string(),
+                        },
+                    })
+                    .collect(),
+            };
+
+            format!("{fn_name}({})", args.join(", "))
+        }
+
+        // If we haven't normalized a specific expression yet, fall back to sqlparser's Display
+        // to avoid changing the set of supported queries.
+        other => format!("{other}"),
+    }
+}
+
+/// Returns a stable, user-facing name for an unnamed SQL projection expression.
+///
+/// This is intended to be shared across SQL (AST -> name) and, in the future,
+/// DataFrame (DSL Expr -> name) so that default column names follow consistent rules.
+pub fn normalized_sql_expr_name(expr: &SQLExpr) -> String {
+    expr_to_string(expr)
+}

src/daft-sql/src/lib.rs

@@ -2,6 +2,7 @@ pub mod error;
 pub mod functions;
 
 mod exec;
+mod expr_name;
 mod modules;
 mod planner;
 mod schema;

src/daft-sql/src/planner.rs

@@ -38,8 +38,8 @@ use sqlparser::{
 };
 
 use crate::{
-    column_not_found_err, error::*, invalid_operation_err, schema::sql_dtype_to_dtype,
-    statement::Statement, table_not_found_err, unsupported_sql_err,
+    column_not_found_err, error::*, expr_name::normalized_sql_expr_name, invalid_operation_err,
+    schema::sql_dtype_to_dtype, statement::Statement, table_not_found_err, unsupported_sql_err,
 };
 
 /// Bindings are used to lookup in-scope tables, views, and columns (targets T).
@@ -1192,7 +1192,11 @@ impl SQLPlanner<'_> {
                 self.bound_columns.insert(alias.clone(), expr.clone());
                 Ok(vec![expr.alias(alias)])
             }
-            SelectItem::UnnamedExpr(expr) => self.plan_expr(expr).map(|e| vec![e]),
+            SelectItem::UnnamedExpr(expr) => {
+                let name = normalized_sql_expr_name(expr);
+                let expr = self.plan_expr(expr)?;
+                Ok(vec![expr.alias(name)])
+            }
 
             SelectItem::Wildcard(wildcard_opts) => {
                 check_wildcard_options(wildcard_opts)?;

tests/sql/test_default_column_alias.py

@@ -0,0 +1,69 @@
+from __future__ import annotations
+
+import daft
+
+
+def test_sql_default_column_name_agg() -> None:
+    df = daft.from_pydict({"a": [1, 2, 3]})
+
+    res = daft.sql("select sum(a) from df", df=df).to_pydict()
+
+    assert res == {"sum(a)": [6]}
+
+
+def test_sql_default_column_name_count() -> None:
+    df = daft.from_pydict({"a": [1, 2, 3]})
+
+    res = daft.sql("select count(*) from df", df=df).to_pydict()
+
+    assert res == {"count(*)": [3]}
+
+    res = daft.sql("select count(1) from df", df=df).to_pydict()
+
+    assert res == {"count(1)": [3]}
+
+    res = daft.sql("select count(a) from df", df=df).to_pydict()
+
+    assert res == {"count(a)": [3]}
+
+
+def test_sql_default_column_name_non_agg_functions() -> None:
+    df = daft.from_pydict({"a": ["AbC", None]})
+
+    res = daft.sql("select upper(a) from df", df=df).to_pydict()
+    assert res == {"upper(a)": ["ABC", None]}
+
+    res = daft.sql("select coalesce(a, 'x') from df", df=df).to_pydict()
+    assert res == {'coalesce(a, "x")': ["AbC", "x"]}
+
+    # String concatenation operator `||` is not supported; use concat().
+    res = daft.sql("select concat(upper(a), 'x') from df", df=df).to_pydict()
+    assert res == {'concat(upper(a), "x")': ["ABCx", None]}
+
+
+def test_sql_default_column_name_binary_op_is_parenthesized() -> None:
+    df = daft.from_pydict({"a": [1, 2, 3]})
+
+    res = daft.sql("select a + 1 from df", df=df).to_pydict()
+
+    assert res == {"(a + 1)": [2, 3, 4]}
+
+
+def test_sql_default_column_name_string_literal_is_quoted() -> None:
+    df = daft.from_pydict({"a": [1]})
+
+    res = daft.sql("select 'a' from df", df=df).to_pydict()
+
+    assert res == {'"a"': ["a"]}
+
+
+def test_sql_explicit_alias_overrides_default_name() -> None:
+    df = daft.from_pydict({"a": [1, 2, 3]})
+
+    res = daft.sql("select sum(a) as s from df", df=df).to_pydict()
+
+    assert res == {"s": [6]}
+
+    res = daft.sql("select sum(a) as S from df", df=df).to_pydict()
+
+    assert res == {"S": [6]}