Merge branch 'main' into upstream-case-field-metadata

Author: osipovartem

.github/workflows/docs.yaml

@@ -67,10 +67,12 @@ jobs:
         run: |
           set -x
           cd asf-site/
+          # don't delete old versions of the docs (see versions/ in asf-site)
           rsync \
             -a \
             --delete \
             --exclude '/.git/' \
+            --exclude '/versions/' \
             ../docs/build/html/ \
             ./
           cp ../.asf.yaml .

benchmarks/sql_benchmarks/nlj/benchmarks/q18.benchmark

@@ -0,0 +1,12 @@
+name Q18
+group nlj
+
+expect_plan NestedLoopJoinExec
+expect_plan join_type=RightSemi
+
+run
+-- Q18: RIGHT SEMI 4K x 4K | no matches; the maximum sum is 4096 * 4096 - 1
+SELECT r.value
+FROM range(4096) AS l
+RIGHT SEMI JOIN range(4096) AS r
+ON l.value * 4096 + r.value >= 4096 * 4096;

benchmarks/sql_benchmarks/nlj/benchmarks/q19.benchmark

@@ -0,0 +1,12 @@
+name Q19
+group nlj
+
+expect_plan NestedLoopJoinExec
+expect_plan join_type=RightAnti
+
+run
+-- Q19: RIGHT ANTI 4K x 4K | no matches; the maximum sum is 4096 * 4096 - 1
+SELECT r.value
+FROM range(4096) AS l
+RIGHT ANTI JOIN range(4096) AS r
+ON l.value * 4096 + r.value >= 4096 * 4096;

benchmarks/sql_benchmarks/nlj/benchmarks/q20.benchmark

@@ -0,0 +1,28 @@
+name Q20
+group nlj
+
+expect_plan NestedLoopJoinExec
+expect_plan join_type=RightSemi
+
+run
+-- Q20: RIGHT SEMI 2K x 2K | staged matches with 260-byte filter strings
+-- With batch_size=8192, candidates shrink 2048 -> 1024 -> 512 -> ... -> 16.
+-- Under that setting, the multi-left-row path starts at 512 candidates.
+SELECT r.value
+FROM (
+    SELECT repeat('x', 256) || lpad(CAST(
+        CASE WHEN value < 128 THEN 1024
+             WHEN value < 256 THEN 1536
+             WHEN value < 384 THEN 1792
+             WHEN value < 512 THEN 1920
+             WHEN value < 640 THEN 1984
+             WHEN value < 768 THEN 2016
+             ELSE 2032 END AS VARCHAR), 4, '0') AS threshold
+    FROM range(2048)
+) AS l
+RIGHT SEMI JOIN (
+    SELECT value,
+           repeat('x', 256) || lpad(CAST(value AS VARCHAR), 4, '0') AS key
+    FROM range(2048)
+) AS r
+ON l.threshold > r.key;

benchmarks/sql_benchmarks/nlj/benchmarks/q21.benchmark

@@ -0,0 +1,28 @@
+name Q21
+group nlj
+
+expect_plan NestedLoopJoinExec
+expect_plan join_type=RightAnti
+
+run
+-- Q21: RIGHT ANTI 2K x 2K | staged matches with 260-byte filter strings
+-- With batch_size=8192, candidates shrink 2048 -> 1024 -> 512 -> ... -> 16.
+-- Under that setting, the multi-left-row path starts at 512 candidates.
+SELECT r.value
+FROM (
+    SELECT repeat('x', 256) || lpad(CAST(
+        CASE WHEN value < 128 THEN 1024
+             WHEN value < 256 THEN 1536
+             WHEN value < 384 THEN 1792
+             WHEN value < 512 THEN 1920
+             WHEN value < 640 THEN 1984
+             WHEN value < 768 THEN 2016
+             ELSE 2032 END AS VARCHAR), 4, '0') AS threshold
+    FROM range(2048)
+) AS l
+RIGHT ANTI JOIN (
+    SELECT value,
+           repeat('x', 256) || lpad(CAST(value AS VARCHAR), 4, '0') AS key
+    FROM range(2048)
+) AS r
+ON l.threshold > r.key;

benchmarks/sql_benchmarks/nlj/benchmarks/q22.benchmark

@@ -0,0 +1,29 @@
+name Q22
+group nlj
+
+expect_plan NestedLoopJoinExec
+expect_plan join_type=RightSemi
+
+# Materialize the threshold so the filter keeps both join inputs.
+# Keep the left input smaller so the default optimizer preserves RightSemi.
+init
+CREATE TABLE nlj_repeated_left AS
+SELECT 3277 AS threshold FROM range(4096);
+
+init
+CREATE TABLE nlj_repeated_right AS
+SELECT value AS id FROM range(8192);
+
+run
+-- Repeated 40% matches: only the first probe adds matches.
+-- With one partition and batch_size=8192, 4915 survivors exceed the 4096-row
+-- compaction threshold, so every left row probes the full right batch.
+SELECT r.id
+FROM nlj_repeated_left l
+RIGHT SEMI JOIN nlj_repeated_right r ON l.threshold > r.id;
+
+cleanup
+DROP TABLE nlj_repeated_left;
+
+cleanup
+DROP TABLE nlj_repeated_right;

benchmarks/sql_benchmarks/nlj/benchmarks/q23.benchmark

@@ -0,0 +1,32 @@
+name Q23
+group nlj
+
+expect_plan NestedLoopJoinExec
+expect_plan join_type=RightAnti
+
+init
+CREATE TABLE nlj_final_left AS
+SELECT arrow_cast('m', 'Utf8') AS threshold;
+
+# Materialize 32 MiB of Utf8 strings outside the measured query. Compaction
+# would copy these strings; different first bytes keep comparison cheap.
+init
+CREATE TABLE nlj_final_right AS
+SELECT value AS id,
+       arrow_cast((CASE WHEN value < 4096 THEN 'a' ELSE 'z' END)
+                  || repeat('x', 4095), 'Utf8') AS payload
+FROM range(8192);
+
+run
+-- With one partition and batch_size=8192, half the batch survives the only
+-- left row. Compacting after that final probe would needlessly copy 16 MiB.
+-- Return IDs only to exclude string copying from output materialization.
+SELECT r.id
+FROM nlj_final_left l
+RIGHT ANTI JOIN nlj_final_right r ON l.threshold > r.payload;
+
+cleanup
+DROP TABLE nlj_final_left;
+
+cleanup
+DROP TABLE nlj_final_right;

datafusion/functions-nested/src/set_ops.rs

@@ -234,15 +234,15 @@ impl ScalarUDFImpl for ArrayIntersect {
 
 #[user_doc(
     doc_section(label = "Array Functions"),
-    description = "Returns distinct values from the array after removing duplicates.",
+    description = "Returns distinct values from the array after removing duplicates, in the order they first appear.",
     syntax_example = "array_distinct(array)",
     sql_example = r#"```sql
 > select array_distinct([1, 3, 2, 3, 1, 2, 4]);
-+---------------------------------+
-| array_distinct(List([1,2,3,4])) |
-+---------------------------------+
-| [1, 2, 3, 4]                    |
-+---------------------------------+
++---------------------------------------+
+| array_distinct(List([1,3,2,3,1,2,4])) |
++---------------------------------------+
+| [1, 3, 2, 4]                          |
++---------------------------------------+
 ```"#,
     argument(
         name = "array",
@@ -297,7 +297,7 @@ impl ScalarUDFImpl for ArrayDistinct {
 }
 
 /// array_distinct SQL function
-/// example: from list [1, 3, 2, 3, 1, 2, 4] to [1, 2, 3, 4]
+/// example: from list [1, 3, 2, 3, 1, 2, 4] to [1, 3, 2, 4]
 fn array_distinct_inner(args: &[ArrayRef]) -> Result<ArrayRef> {
     let [array] = take_function_args("array_distinct", args)?;
     match array.data_type() {

datafusion/optimizer/src/push_down_filter.rs

@@ -1135,8 +1135,8 @@ impl OptimizerRule for PushDownFilter {
                 // multiple window functions, each with potentially different partition keys.
                 // Therefore, we need to ensure that any potential partition key returned is used in
                 // ALL window functions. Otherwise, filters cannot be pushed by through that column.
-                fn extract_partition_keys(func: &WindowFunction) -> HashSet<Column> {
-                    expr_columns(&func.params.partition_by)
+                fn extract_partition_keys(func: &WindowFunction) -> HashSet<&Expr> {
+                    func.params.partition_by.iter().collect()
                 }
 
                 let potential_partition_keys = window
@@ -1172,21 +1172,23 @@ impl OptimizerRule for PushDownFilter {
                 let mut keep_predicates = vec![];
                 let mut push_predicates = vec![];
                 for expr in predicates {
-                    let cols = expr.column_refs();
-                    if cols.iter().all(|c| potential_partition_keys.contains(c)) {
+                    // A volatile predicate has to stay above the window: pushing it
+                    // changes which rows the window function sees.
+                    if !expr.is_volatile()
+                        && reads_only_partition_keys(&expr, &potential_partition_keys)?
+                    {
                         push_predicates.push(expr);
                     } else {
                         keep_predicates.push(expr);
                     }
                 }
 
                 // Unlike with aggregations, there are no cases where we have to replace, e.g.,
-                // `a+b` with Column(a)+Column(b). This is because partition expressions are not
-                // available as standalone columns to the user. For example, while an aggregation on
-                // `a+b` becomes Column(a + b), in a window partition it becomes
-                // `func() PARTITION BY [a + b] ...`. Thus, filters on expressions always remain in
-                // place, so we can use `push_predicates` directly. This is consistent with other
-                // optimizers, such as the one used by Postgres.
+                // `a+b` with Column(a+b). This is because partition expressions are not available
+                // as standalone columns to the user: while an aggregation on `a+b` becomes
+                // Column(a + b), in a window partition it stays `func() PARTITION BY [a + b] ...`.
+                // That is why the predicate is matched against the key expressions themselves and
+                // can be pushed unchanged.
 
                 // If we have a filter to push, we push it down to the input of the aggregate
                 let result = if let Some(predicate) = conjunction(push_predicates) {
@@ -1585,6 +1587,54 @@ fn with_filters(predicates: Vec<Expr>, plan: LogicalPlan) -> LogicalPlan {
     }
 }
 
+/// Can `expr` be evaluated below a window with these `PARTITION BY` keys?
+///
+/// A predicate that reads only the partition keys is constant within each
+/// partition, so filtering before the window drops whole partitions and leaves
+/// the surviving rows' window values unchanged. "Reads only the keys" is checked
+/// structurally: every column reference must sit inside a subtree that is equal
+/// to one of the keys. Given `PARTITION BY a, b + c`:
+///
+/// * `a < 5`, `b + c = 4` and `(b + c) + 1 > 10` can be pushed down
+/// * `d < 5` and `b < 5` cannot (`b` on its own is not a key), and neither can
+///   `c + b = 4` (the match is structural, `c + b` is not `b + c`)
+///
+/// A predicate containing a subquery is never pushed: what the subquery reads is
+/// not visible from the expression tree.
+fn reads_only_partition_keys(
+    expr: &Expr,
+    partition_keys: &HashSet<&Expr>,
+) -> Result<bool> {
+    let mut reads_something_else = false;
+    expr.apply(|node| {
+        Ok(if partition_keys.contains(&node) {
+            // the whole key was matched, so whatever it reads is accounted for
+            TreeNodeRecursion::Jump
+        } else if reads_beyond_this_node(node) {
+            reads_something_else = true;
+            TreeNodeRecursion::Stop
+        } else {
+            TreeNodeRecursion::Continue
+        })
+    })?;
+    Ok(!reads_something_else)
+}
+
+/// Does this node read data that walking its children cannot account for? A
+/// column reads itself; a subquery reads whatever its plan reads, which
+/// `Expr::apply` does not visit.
+fn reads_beyond_this_node(node: &Expr) -> bool {
+    matches!(
+        node,
+        Expr::Column(_)
+            | Expr::OuterReferenceColumn(..)
+            | Expr::ScalarSubquery(_)
+            | Expr::Exists(_)
+            | Expr::InSubquery(_)
+            | Expr::SetComparison(_)
+    )
+}
+
 fn expr_columns(exprs: &[Expr]) -> HashSet<Column> {
     exprs
         .iter()
@@ -1956,10 +2006,12 @@ mod tests {
         )
     }
 
-    /// verifies that filters on partition expressions are not pushed, as the single expression
-    /// column is not available to the user, unlike with aggregations
+    /// verifies that a filter on an expression partition key is pushed; the
+    /// remaining shapes (mixed keys, operand order, subqueries, volatile
+    /// predicates, several windows) are covered in
+    /// `push_down_filter_regression.slt`
     #[test]
-    fn filter_expression_keep_window() -> Result<()> {
+    fn filter_expression_move_window() -> Result<()> {
         let table_scan = test_table_scan()?;
 
         let window = Expr::from(WindowFunction::new(
@@ -1975,49 +2027,14 @@ mod tests {
 
         let plan = LogicalPlanBuilder::from(table_scan)
             .window(vec![window])?
-            // unlike with aggregations, single partition column "test.a + test.b" is not available
-            // to the plan, so we use multiple columns when filtering
-            .filter(add(col("a"), col("b")).gt(lit(10i64)))?
-            .build()?;
-
-        assert_optimized_plan_equal!(
-            plan,
-            @r"
-        Filter: test.a + test.b > Int64(10)
-          WindowAggr: windowExpr=[[rank() PARTITION BY [test.a + test.b] ORDER BY [test.c ASC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]]
-            TableScan: test
-        "
-        )
-    }
-
-    /// verifies that filters are not pushed on order by columns (that are not used in partitioning)
-    #[test]
-    fn filter_order_keep_window() -> Result<()> {
-        let table_scan = test_table_scan()?;
-
-        let window = Expr::from(WindowFunction::new(
-            WindowFunctionDefinition::WindowUDF(
-                datafusion_functions_window::rank::rank_udwf(),
-            ),
-            vec![],
-        ))
-        .partition_by(vec![col("a")])
-        .order_by(vec![col("c").sort(true, true)])
-        .build()
-        .unwrap();
-
-        let plan = LogicalPlanBuilder::from(table_scan)
-            .window(vec![window])?
-            .filter(col("c").gt(lit(10i64)))?
+            .filter(add(col("a"), col("b")).gt(lit(10i64)))? // a + b > 10
             .build()?;
-        assert_plan_not_transformed!(plan.clone());
 
         assert_optimized_plan_equal!(
             plan,
             @r"
-        Filter: test.c > Int64(10)
-          WindowAggr: windowExpr=[[rank() PARTITION BY [test.a] ORDER BY [test.c ASC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]]
-            TableScan: test
+        WindowAggr: windowExpr=[[rank() PARTITION BY [test.a + test.b] ORDER BY [test.c ASC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]]
+          TableScan: test, full_filters=[test.a + test.b > Int64(10)]
         "
         )
     }