diff --git a/src/read/optimizers.rs b/src/read/optimizers.rs index 8afaeb01..d94f3460 100644 --- a/src/read/optimizers.rs +++ b/src/read/optimizers.rs @@ -1998,14 +1998,16 @@ mod ordered_union_for_topk_tests { // SELECT aliases; a relation without them is left alone rather than guessed at. use datafusion::sql::sqlparser::ast::{ - Expr as SqlExpr, Function, FunctionArg, FunctionArgExpr, FunctionArgumentList, FunctionArguments, Ident, ObjectName, ObjectNamePart, Query, Select, - SelectItem, SetExpr, Statement, TableFactor, Value as SqlValue, ValueWithSpan, VisitMut, VisitorMut, + Expr as SqlExpr, Function, FunctionArg, FunctionArgExpr, FunctionArgumentList, FunctionArguments, GroupByExpr, Ident, ObjectName, ObjectNamePart, OrderBy, + OrderByKind, Query, Select, SelectItem, SetExpr, Statement, TableFactor, Value as SqlValue, ValueWithSpan, VisitMut, VisitorMut, }; /// Cheap guard so the parse/unparse round trip only happens for statements that /// could possibly need it. pub fn might_need_rewrite(sql: &str) -> bool { - sql.to_ascii_lowercase().contains("row_to_json") + static ORDINAL: std::sync::LazyLock = + std::sync::LazyLock::new(|| regex::Regex::new(r"(?i)\border\s+by\s+(?:[^;]*?,\s*)?\d").expect("valid regex")); + sql.to_ascii_lowercase().contains("row_to_json") || ORDINAL.is_match(sql) } /// Returns the rewritten statement only when something actually changed, so an @@ -2025,11 +2027,38 @@ impl VisitorMut for RowToJsonRecord { fn post_visit_query(&mut self, query: &mut Query) -> std::ops::ControlFlow { self.rewrite_set_expr(query.body.as_mut()); + self.rewrite_order_ordinals(query); std::ops::ControlFlow::Continue(()) } } impl RowToJsonRecord { + /// DataFusion rejects `GROUP BY 1 ORDER BY 2` when select item 2 is a scalar over an + /// aggregate (`round(avg(x))`) yet accepts the same item spelled out, so an ordinal naming + /// an unaliased aggregate output is replaced by its expression. Group keys already plan + /// and stay untouched, as do a union's ordinals, which name its output. + fn rewrite_order_ordinals(&mut self, query: &mut Query) { + let (SetExpr::Select(select), Some(OrderBy { kind: OrderByKind::Expressions(exprs), .. })) = (query.body.as_ref(), query.order_by.as_mut()) else { + return; + }; + let GroupByExpr::Expressions(keys, _) = &select.group_by else { return }; + if keys.is_empty() { + return; + } + let ordinal = |expr: &SqlExpr| match expr { + SqlExpr::Value(ValueWithSpan { value: SqlValue::Number(n, _), .. }) => n.parse::().ok(), + _ => None, + }; + for order in exprs { + let Some(n) = ordinal(&order.expr) else { continue }; + let Some(SelectItem::UnnamedExpr(expr)) = n.checked_sub(1).and_then(|i| select.projection.get(i)) else { continue }; + if !matches!(expr, SqlExpr::Identifier(_) | SqlExpr::CompoundIdentifier(_)) && !keys.iter().any(|key| ordinal(key) == Some(n) || key == expr) { + order.expr = expr.clone(); + self.rewrote = true; + } + } + } + /// A UNION's branches are `SetExpr`s, not `Query`s, so they are never /// reached by matching on `query.body` alone. fn rewrite_set_expr(&mut self, body: &mut SetExpr) { @@ -2165,6 +2194,25 @@ mod row_to_json_record_tests { fn declines_to_rewrite(sql: &str) -> Option { rewrite_sql(sql) } + + #[test_case::test_case("SELECT k, round(avg(v)) FROM t GROUP BY 1 ORDER BY 2 DESC" => Some(vec!["b".into(), "a".into()]) ; "scalar over aggregate")] + #[test_case::test_case("SELECT k, round(avg(v)) FROM t GROUP BY 1 ORDER BY 2, 1" => Some(vec!["a".into(), "b".into()]) ; "mixed ordinals")] + #[test_case::test_case("SELECT k, avg(v) FROM t GROUP BY 1 ORDER BY 2 DESC" => Some(vec!["b".into(), "a".into()]) ; "a bare aggregate keeps its order")] + #[test_case::test_case("SELECT k, avg(v) FROM t GROUP BY 1 ORDER BY 1 DESC" => None ; "a column ordinal is left alone")] + #[test_case::test_case("SELECT upper(k), count(*) FROM t GROUP BY 1 ORDER BY 1" => None ; "a computed group key is left alone")] + #[test_case::test_case("SELECT upper(k), count(*) FROM t GROUP BY upper(k) ORDER BY 1" => None ; "a group key spelled out is left alone")] + #[test_case::test_case("SELECT k FROM t UNION ALL SELECT k FROM t ORDER BY 1" => None ; "a union's ordinal is left alone")] + #[tokio::test] + async fn order_by_ordinal_over_a_scalar_aggregate_plans(sql: &str) -> Option> { + assert!(might_need_rewrite(sql)); + let rewritten = rewrite_sql(sql)?; + let ctx = datafusion::prelude::SessionContext::new(); + ctx.sql("CREATE TABLE t (k VARCHAR, v DOUBLE) AS VALUES ('a', 1.0), ('b', 2.0)").await.expect("table"); + let batches = ctx.sql(&rewritten).await.expect("the rewrite plans").collect().await.expect("runs"); + let keys = arrow::compute::concat_batches(&batches[0].schema(), &batches).expect("concat"); + let keys = arrow::compute::cast(keys.column(0), &arrow::datatypes::DataType::Utf8).expect("cast"); + Some(arrow::array::AsArray::as_string::(&keys).iter().map(|k| k.unwrap_or_default().to_string()).collect()) + } } // Rewrites `EXISTS(q)` in a projection to an executable scalar expression: diff --git a/tests/suite/pg_client_compat.rs b/tests/suite/pg_client_compat.rs index 7022b33f..2be751c3 100644 --- a/tests/suite/pg_client_compat.rs +++ b/tests/suite/pg_client_compat.rs @@ -260,6 +260,28 @@ async fn pgadmin_dashboard_rewrites_every_union_branch() -> Result<()> { Ok(()) } +/// DataFusion rejects `ORDER BY ` naming a scalar over an aggregate; the AST rewrite +/// spells the ordinal out. Pinned on both protocols, which parse through different entry points. +#[tokio::test(flavor = "multi_thread")] +async fn order_by_ordinal_over_a_scalar_aggregate_plans_on_both_protocols() -> Result<()> { + let server = TestServer::start().await?; + let client = server.client().await?; + let sql = "WITH t(k, v) AS (VALUES ('a', 1.0), ('b', 2.0), ('b', 4.0)) SELECT k, round(avg(v))::TEXT FROM t GROUP BY 1 ORDER BY 2 DESC"; + let extended = client.query(sql, &[]).await?.iter().map(|row| row.get::<_, String>(0)).collect::>(); + assert_eq!(extended, ["b", "a"], "extended protocol"); + let simple = client + .simple_query(sql) + .await? + .iter() + .filter_map(|message| match message { + tokio_postgres::SimpleQueryMessage::Row(row) => row.get(0).map(str::to_owned), + _ => None, + }) + .collect::>(); + assert_eq!(simple, ["b", "a"], "simple protocol"); + Ok(()) +} + #[tokio::test(flavor = "multi_thread")] async fn aggregate_ordering_preserves_input_values_and_output_aliases() -> Result<()> { let server = TestServer::start().await?;