From ac78aa073f019953ddacd2ea1fc055dca4a04a15 Mon Sep 17 00:00:00 2001 From: Oleks V Date: Tue, 1 Sep 2026 18:34:58 +0000 Subject: [PATCH 01/37] feat: extend SLT tests for hash join vs PWMJ, hash join vs SMJ (#24805) ## Which issue does this PR close? - Closes #24641 . ## Rationale for this change `PiecewiseMergeJoinExec` and `SortMergeJoinExec` must return the same results as the mature join implementations they can be swapped for, across every batch boundary. Today that equivalence is checked ad hoc. Using the `# configMatrix:` directive from #24493, one `.slt` file can assert it directly: run the same queries once per join implementation and once per batch size, and require identical output. ## What changes are included in this PR? ## What is the testing strategy for this PR? ## Are there any user-facing changes? --- .../test_files/mark_join_matrix.slt | 202 ++++++ .../piecewise_merge_join_batches.slt | 103 +++ .../piecewise_merge_join_matrix.slt | 661 ++++++++++++++++++ .../test_files/sort_merge_join_batches.slt | 65 ++ .../test_files/sort_merge_join_matrix.slt | 315 +++++++++ 5 files changed, 1346 insertions(+) create mode 100644 datafusion/sqllogictest/test_files/mark_join_matrix.slt create mode 100644 datafusion/sqllogictest/test_files/piecewise_merge_join_batches.slt create mode 100644 datafusion/sqllogictest/test_files/piecewise_merge_join_matrix.slt create mode 100644 datafusion/sqllogictest/test_files/sort_merge_join_batches.slt create mode 100644 datafusion/sqllogictest/test_files/sort_merge_join_matrix.slt diff --git a/datafusion/sqllogictest/test_files/mark_join_matrix.slt b/datafusion/sqllogictest/test_files/mark_join_matrix.slt new file mode 100644 index 0000000000000..d394d527d5d13 --- /dev/null +++ b/datafusion/sqllogictest/test_files/mark_join_matrix.slt @@ -0,0 +1,202 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at + +# http://www.apache.org/licenses/LICENSE-2.0 + +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +# Mark-join correctness across a config matrix. A mark join adds a boolean `mark` +# per left row and comes from EXISTS/IN/NOT EXISTS/NOT IN inside a disjunction +# (`WHERE OR EXISTS(..)`; see mark_join in decorrelate_predicate_subquery.rs). +# For an equijoin mark, prefer_hash_join toggles the operator: false -> +# SortMergeJoinExec LeftMark, true -> HashJoinExec RightMark (inputs swapped), so +# sweeping {true,false} x batch_size {1,2,100,8192} cross-checks both directions +# and must agree row-for-row. Range mark joins run on NestedLoopJoin regardless +# (Part 2, batch_size only). SMJ needs target_partitions>1 and repartition_joins, +# set below (not swept). Matrix rules: no EXPLAIN, no in-file SET of a swept knob, +# rowsort every multi-row query. + +# configMatrix: datafusion.optimizer.prefer_hash_join=true,false +# configMatrix: datafusion.execution.batch_size=1,2,100,8192 + +statement ok +set datafusion.execution.target_partitions = 4; + +statement ok +set datafusion.optimizer.repartition_joins = true; + +# ------------------------------------------------------------------ +# Fixtures: duplicate and NULL keys on the subquery side, a NULL key and an +# unmatched key on the outer side. +# ------------------------------------------------------------------ +statement ok +CREATE TABLE mk_l(k INT, v INT) AS VALUES + (1, 10), (2, 20), (3, 30), (4, 40), (NULL, 50); + +statement ok +CREATE TABLE mk_r(k INT) AS VALUES (1), (2), (2), (NULL); + +# mk_r without the NULL, for NOT IN cases that would otherwise be swallowed by +# three-valued logic. +statement ok +CREATE TABLE mk_r_nn(k INT) AS VALUES (1), (2), (2); + +statement ok +CREATE TABLE mk_empty(k INT); + +# ================================================================== +# Part 1: Equijoin mark joins (LeftMark on SMJ vs RightMark on HashJoin) +# ================================================================== + +# EXISTS in a disjunction with a sometimes-true predicate. mark(EXISTS k in +# {1,2}) is true for k=1,2; the predicate v>35 is true for k=4 and k=NULL. +query II rowsort +SELECT l.k, l.v FROM mk_l l +WHERE l.v > 35 OR EXISTS (SELECT 1 FROM mk_r r WHERE l.k = r.k); +---- +1 10 +2 20 +4 40 +NULL 50 + +# NOT EXISTS in a disjunction: mark is negated, true for k not in {1,2}. +query II rowsort +SELECT l.k, l.v FROM mk_l l +WHERE l.v > 35 OR NOT EXISTS (SELECT 1 FROM mk_r r WHERE l.k = r.k); +---- +3 30 +4 40 +NULL 50 + +# Predicate never true (no negative k), so the result isolates the mark: the +# EXISTS rows k in {1,2}. Exercises the mark column with the OR contributing +# nothing. +query II rowsort +SELECT l.k, l.v FROM mk_l l +WHERE l.k < 0 OR EXISTS (SELECT 1 FROM mk_r r WHERE l.k = r.k); +---- +1 10 +2 20 + +# Same, negated: isolates NOT EXISTS. The NULL-keyed left row never matches, so +# it is kept. +query II rowsort +SELECT l.k, l.v FROM mk_l l +WHERE l.k < 0 OR NOT EXISTS (SELECT 1 FROM mk_r r WHERE l.k = r.k); +---- +3 30 +4 40 +NULL 50 + +# IN in a disjunction: same matches as EXISTS here; the NULL in the subquery adds +# no true values. +query II rowsort +SELECT l.k, l.v FROM mk_l l +WHERE l.v > 35 OR l.k IN (SELECT r.k FROM mk_r r); +---- +1 10 +2 20 +4 40 +NULL 50 + +# NOT IN over a NULL-free subquery: mark(NOT IN) is true for k not in {1,2}. +query II rowsort +SELECT l.k, l.v FROM mk_l l +WHERE l.v > 35 OR l.k NOT IN (SELECT r.k FROM mk_r_nn r); +---- +3 30 +4 40 +NULL 50 + +# https://github.com/apache/datafusion/issues/24854 +# query II rowsort +# SELECT l.k, l.v FROM mk_l l +# WHERE l.v > 35 OR l.k NOT IN (SELECT r.k FROM mk_r r); +# ---- +# 4 40 +# NULL 50 + +# Empty subquery: EXISTS is always false, so the result is just the predicate. +query II rowsort +SELECT l.k, l.v FROM mk_l l +WHERE l.v > 35 OR EXISTS (SELECT 1 FROM mk_empty r WHERE l.k = r.k); +---- +4 40 +NULL 50 + +# Empty subquery, negated: NOT EXISTS is always true, so every row survives. +query II rowsort +SELECT l.k, l.v FROM mk_l l +WHERE l.k < 0 OR NOT EXISTS (SELECT 1 FROM mk_empty r WHERE l.k = r.k); +---- +1 10 +2 20 +3 30 +4 40 +NULL 50 + +# Multiple equi keys in the mark correlation. +statement ok +CREATE TABLE mk2_l(k1 INT, k2 INT, v INT) AS VALUES + (1, 1, 10), (1, 2, 20), (2, 2, 30), (3, 3, 40); + +statement ok +CREATE TABLE mk2_r(k1 INT, k2 INT) AS VALUES (1, 1), (2, 2), (1, 9); + +query III rowsort +SELECT l.k1, l.k2, l.v FROM mk2_l l +WHERE l.v > 35 OR EXISTS (SELECT 1 FROM mk2_r r WHERE l.k1 = r.k1 AND l.k2 = r.k2); +---- +1 1 10 +2 2 30 +3 3 40 + +# Cross-table filter in the correlation (regression #21197): unmatched mark rows +# produce null right indices that must not corrupt non-nullable left columns. +# EXISTS holds for the k=2 rows (partner (2,99) differs in d); k=1's partner +# equals its d, so false. +statement ok +CREATE TABLE mkf_l(k INT, d INT) AS VALUES (1, 10), (2, 20), (2, 25), (3, 30); + +statement ok +CREATE TABLE mkf_r(k INT, d INT) AS VALUES (1, 10), (2, 20), (2, 99); + +query II rowsort +SELECT l.k, l.d FROM mkf_l l +WHERE l.d < 0 OR EXISTS (SELECT 1 FROM mkf_r r WHERE l.k = r.k AND r.d <> l.d); +---- +2 20 +2 25 + +# ================================================================== +# Part 2: Range mark joins (RightMark on NestedLoopJoin) +# ================================================================== +# No equi key, so these run on NestedLoopJoin in both combinations; only +# batch_size varies. mk_r non-null keys = {1,2,2}. + +# EXISTS l.k > r.k is true for k>1, i.e. k in {2,3,4}. Predicate never true. +query II rowsort +SELECT l.k, l.v FROM mk_l l +WHERE l.v > 100 OR EXISTS (SELECT 1 FROM mk_r r WHERE l.k > r.k); +---- +2 20 +3 30 +4 40 + +# NOT EXISTS of the same range: true for k=1 (no smaller r) and the NULL key. +query II rowsort +SELECT l.k, l.v FROM mk_l l +WHERE l.k < 0 OR NOT EXISTS (SELECT 1 FROM mk_r r WHERE l.k > r.k); +---- +1 10 +NULL 50 diff --git a/datafusion/sqllogictest/test_files/piecewise_merge_join_batches.slt b/datafusion/sqllogictest/test_files/piecewise_merge_join_batches.slt new file mode 100644 index 0000000000000..db27c9fc8bb60 --- /dev/null +++ b/datafusion/sqllogictest/test_files/piecewise_merge_join_batches.slt @@ -0,0 +1,103 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at + +# http://www.apache.org/licenses/LICENSE-2.0 + +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +# Plan proof for piecewise_merge_join_matrix.slt. Not a matrix file: it fixes +# batch_size=2 and enable_piecewise_merge_join=true so it can EXPLAIN what the +# matrix relies on but cannot check itself -- the streamed side is a range +# LazyMemoryExec with batch_size=2 (so range(3,8), 5 rows, arrives as 3 batches) +# feeding a PiecewiseMergeJoin, not a fallback. A VALUES source would not show +# batch_size in the plan. + +statement ok +set datafusion.optimizer.enable_piecewise_merge_join = true; + +statement ok +set datafusion.execution.batch_size = 2; + +statement ok +CREATE TABLE pb_l(id INT, v INT) AS + SELECT CAST(value AS INT) AS id, CAST(value AS INT) AS v FROM range(1, 11); + +# Existence (LeftSemi): range streamed side (batch_size=2) -> LeftSemi PWMJ. +query TT +EXPLAIN SELECT count(*) FROM pb_l l WHERE EXISTS (SELECT 1 FROM range(3, 8) r WHERE l.v > r.value); +---- +logical_plan +01)Projection: count(Int64(1)) AS count(*) +02)--Aggregate: groupBy=[[]], aggr=[[count(Int64(1))]] +03)----Projection: +04)------LeftSemi Join: Filter: CAST(l.v AS Int64) > __correlated_sq_1.value +05)--------SubqueryAlias: l +06)----------TableScan: pb_l projection=[v] +07)--------SubqueryAlias: __correlated_sq_1 +08)----------SubqueryAlias: r +09)------------TableScan: range() projection=[value] +physical_plan +01)ProjectionExec: expr=[count(Int64(1))@0 as count(*)] +02)--AggregateExec: mode=Final, gby=[], aggr=[count(Int64(1))] +03)----CoalescePartitionsExec +04)------AggregateExec: mode=Partial, gby=[], aggr=[count(Int64(1))] +05)--------ProjectionExec: expr=[] +06)----------PiecewiseMergeJoin: operator=Gt, join_type=LeftSemi, on=(CAST(v AS Int64) > value) +07)------------SortPreservingMergeExec: [CAST(v@0 AS Int64) ASC] +08)--------------SortExec: expr=[CAST(v@0 AS Int64) ASC], preserve_partitioning=[true] +09)----------------DataSourceExec: partitions=4, partition_sizes=[2, 1, 1, 1] +10)------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 +11)--------------LazyMemoryExec: partitions=1, batch_generators=[range: start=3, end=8, batch_size=2] + +# Same result the matrix asserts (v > min(3) -> {4..10} = 7). +query I +SELECT count(*) FROM pb_l l WHERE EXISTS (SELECT 1 FROM range(3, 8) r WHERE l.v > r.value); +---- +7 + +# Classic Inner: same range streamed side -> Inner PWMJ. +query TT +EXPLAIN SELECT count(*) FROM pb_l l JOIN range(3, 8) r ON l.v < r.value; +---- +logical_plan +01)Projection: count(Int64(1)) AS count(*) +02)--Aggregate: groupBy=[[]], aggr=[[count(Int64(1))]] +03)----Projection: +04)------Inner Join: Filter: CAST(l.v AS Int64) < r.value +05)--------SubqueryAlias: l +06)----------TableScan: pb_l projection=[v] +07)--------SubqueryAlias: r +08)----------TableScan: range() projection=[value] +physical_plan +01)ProjectionExec: expr=[count(Int64(1))@0 as count(*)] +02)--AggregateExec: mode=Final, gby=[], aggr=[count(Int64(1))] +03)----CoalescePartitionsExec +04)------AggregateExec: mode=Partial, gby=[], aggr=[count(Int64(1))] +05)--------ProjectionExec: expr=[] +06)----------PiecewiseMergeJoin: operator=Lt, join_type=Inner, on=(CAST(v AS Int64) < value) +07)------------SortPreservingMergeExec: [CAST(v@0 AS Int64) DESC] +08)--------------SortExec: expr=[CAST(v@0 AS Int64) DESC], preserve_partitioning=[true] +09)----------------DataSourceExec: partitions=4, partition_sizes=[2, 1, 1, 1] +10)------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 +11)--------------LazyMemoryExec: partitions=1, batch_generators=[range: start=3, end=8, batch_size=2] + +query I +SELECT count(*) FROM pb_l l JOIN range(3, 8) r ON l.v < r.value; +---- +20 + +statement ok +RESET datafusion.execution.batch_size; + +statement ok +RESET datafusion.optimizer.enable_piecewise_merge_join; diff --git a/datafusion/sqllogictest/test_files/piecewise_merge_join_matrix.slt b/datafusion/sqllogictest/test_files/piecewise_merge_join_matrix.slt new file mode 100644 index 0000000000000..13aeeddfc2afa --- /dev/null +++ b/datafusion/sqllogictest/test_files/piecewise_merge_join_matrix.slt @@ -0,0 +1,661 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at + +# http://www.apache.org/licenses/LICENSE-2.0 + +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +# PiecewiseMergeJoin correctness across a config matrix. A range join with no +# equi key runs on PiecewiseMergeJoin when enable_piecewise_merge_join=true and +# on NestedLoopJoin when false, so sweeping {true,false} x batch_size +# {1,2,100,8192} requires the two to agree row-for-row. PWMJ covers +# Inner/Left/Right/Full (Part 1) and LeftSemi/LeftAnti (Part 2); explicit RIGHT +# SEMI/ANTI JOIN and Mark (OR EXISTS) stay on NLJ. Matrix rules: no EXPLAIN, no +# in-file SET of a swept knob, rowsort every multi-row query. + +# configMatrix: datafusion.optimizer.enable_piecewise_merge_join=true,false +# configMatrix: datafusion.execution.batch_size=1,2,100,8192 + +# ================================================================== +# Part 1: Classic range joins (Inner / Left / Right / Full) +# ================================================================== +# cj_l.lv = {1,3,5,NULL}, cj_r.rv = {3,5,NULL}. A NULL key matches nothing, so it +# appears only as an unmatched row on an outer side. +statement ok +CREATE TABLE cj_l(lid INT, lv INT); + +statement ok +INSERT INTO cj_l VALUES (1, 1), (2, 3), (3, 5), (4, NULL); + +statement ok +CREATE TABLE cj_r(rid INT, rv INT); + +statement ok +INSERT INTO cj_r VALUES (1, 3), (2, 5), (3, NULL); + +# INNER, all four operators. `<=`/`>=` differ from `<`/`>` through the equal +# values (lv=3=rv, lv=5=rv). +query IIII rowsort +SELECT l.lid, l.lv, r.rid, r.rv FROM cj_l l JOIN cj_r r ON l.lv < r.rv; +---- +1 1 1 3 +1 1 2 5 +2 3 2 5 + +query IIII rowsort +SELECT l.lid, l.lv, r.rid, r.rv FROM cj_l l JOIN cj_r r ON l.lv <= r.rv; +---- +1 1 1 3 +1 1 2 5 +2 3 1 3 +2 3 2 5 +3 5 2 5 + +query IIII rowsort +SELECT l.lid, l.lv, r.rid, r.rv FROM cj_l l JOIN cj_r r ON l.lv > r.rv; +---- +3 5 1 3 + +query IIII rowsort +SELECT l.lid, l.lv, r.rid, r.rv FROM cj_l l JOIN cj_r r ON l.lv >= r.rv; +---- +2 3 1 3 +3 5 1 3 +3 5 2 5 + +# LEFT: the INNER `<` rows plus every unmatched left row (lv=5 and lv=NULL) with +# a NULL right side. +query IIII rowsort +SELECT l.lid, l.lv, r.rid, r.rv FROM cj_l l LEFT JOIN cj_r r ON l.lv < r.rv; +---- +1 1 1 3 +1 1 2 5 +2 3 2 5 +3 5 NULL NULL +4 NULL NULL NULL + +# RIGHT: INNER rows plus unmatched right rows (NULL left). rid=3 has a NULL key +# and must still be emitted (regression: streamed-side NULLs were dropped). +query IIII rowsort +SELECT l.lid, l.lv, r.rid, r.rv FROM cj_l l RIGHT JOIN cj_r r ON l.lv < r.rv; +---- +1 1 1 3 +1 1 2 5 +2 3 2 5 +NULL NULL 3 NULL + +# RIGHT with a different operator flips the buffered-side sort direction; the +# NULL-keyed right row rid=3 is still emitted unmatched. +query IIII rowsort +SELECT l.lid, l.lv, r.rid, r.rv FROM cj_l l RIGHT JOIN cj_r r ON l.lv >= r.rv; +---- +2 3 1 3 +3 5 1 3 +3 5 2 5 +NULL NULL 3 NULL + +# FULL: the INNER rows plus unmatched rows from both sides, including a NULL key +# on each side. +query IIII rowsort +SELECT l.lid, l.lv, r.rid, r.rv FROM cj_l l FULL JOIN cj_r r ON l.lv < r.rv; +---- +1 1 1 3 +1 1 2 5 +2 3 2 5 +3 5 NULL NULL +4 NULL NULL NULL +NULL NULL 3 NULL + +# Empty right side: LEFT keeps every left row with NULLs, INNER is empty. +statement ok +CREATE TABLE cj_empty(rid INT, rv INT); + +query IIII rowsort +SELECT l.lid, l.lv, r.rid, r.rv FROM cj_l l JOIN cj_empty r ON l.lv < r.rv; +---- + +query IIII rowsort +SELECT l.lid, l.lv, r.rid, r.rv FROM cj_l l LEFT JOIN cj_empty r ON l.lv < r.rv; +---- +1 1 NULL NULL +2 3 NULL NULL +3 5 NULL NULL +4 NULL NULL NULL + +# Empty left side: RIGHT keeps every right row with NULLs, including rid=3's NULL +# key. +query IIII rowsort +SELECT l.rid, l.rv, r.rid, r.rv FROM cj_empty l RIGHT JOIN cj_r r ON l.rv < r.rv; +---- +NULL NULL 1 3 +NULL NULL 2 5 +NULL NULL 3 NULL + +# Duplicate keys on both sides: each qualifying left row must pair with the whole +# matching right suffix. lv=1 (twice) is below rv=2 (twice) -> 4 rows; lv=3 has +# no larger right value. +statement ok +CREATE TABLE cj_dup_l(lid INT, lv INT); + +statement ok +INSERT INTO cj_dup_l VALUES (1, 1), (2, 1), (3, 3); + +statement ok +CREATE TABLE cj_dup_r(rid INT, rv INT); + +statement ok +INSERT INTO cj_dup_r VALUES (1, 2), (2, 2); + +query IIII rowsort +SELECT l.lid, l.lv, r.rid, r.rv FROM cj_dup_l l JOIN cj_dup_r r ON l.lv < r.rv; +---- +1 1 1 2 +1 1 2 2 +2 1 1 2 +2 1 2 2 + +# ================================================================== +# Part 2: Existence joins (LeftSemi / LeftAnti; Right/Mark fall back to NLJ) +# ================================================================== +# EXISTS -> LeftSemi, NOT EXISTS -> LeftAnti (the only existence joins PWMJ +# implements). Explicit RIGHT SEMI/ANTI JOIN and Mark (OR EXISTS) stay on NLJ; +# covered at the end. + +# ------------------------------------------------------------------ +# Fixtures +# ------------------------------------------------------------------ +statement ok +CREATE TABLE ej_l(id INT, v INT); + +statement ok +INSERT INTO ej_l VALUES (1, 5), (2, 4), (3, 2), (4, 1); + +statement ok +CREATE TABLE ej_r(v INT); + +statement ok +INSERT INTO ej_r VALUES (2), (3), (4); + +# ------------------------------------------------------------------ +# Basic: every operator, both LeftSemi (EXISTS) and LeftAnti (NOT EXISTS) +# ------------------------------------------------------------------ +# ej_l.v = {5,4,2,1}, ej_r.v = {2,3,4}. + +# `<` : 2<3 and 1<2 match; 5 and 4 exceed every right value. +query I rowsort +SELECT l.id FROM ej_l l WHERE EXISTS (SELECT 1 FROM ej_r r WHERE l.v < r.v); +---- +3 +4 + +query I rowsort +SELECT l.id FROM ej_l l WHERE NOT EXISTS (SELECT 1 FROM ej_r r WHERE l.v < r.v); +---- +1 +2 + +# `<=` : additionally admits v=4 through the equal right value 4. +query I rowsort +SELECT l.id FROM ej_l l WHERE EXISTS (SELECT 1 FROM ej_r r WHERE l.v <= r.v); +---- +2 +3 +4 + +query I rowsort +SELECT l.id FROM ej_l l WHERE NOT EXISTS (SELECT 1 FROM ej_r r WHERE l.v <= r.v); +---- +1 + +# `>` : 5 and 4 exceed some right value; 2 and 1 do not. +query I rowsort +SELECT l.id FROM ej_l l WHERE EXISTS (SELECT 1 FROM ej_r r WHERE l.v > r.v); +---- +1 +2 + +query I rowsort +SELECT l.id FROM ej_l l WHERE NOT EXISTS (SELECT 1 FROM ej_r r WHERE l.v > r.v); +---- +3 +4 + +# `>=` : additionally admits v=2 through the equal right value 2. +query I rowsort +SELECT l.id FROM ej_l l WHERE EXISTS (SELECT 1 FROM ej_r r WHERE l.v >= r.v); +---- +1 +2 +3 + +query I rowsort +SELECT l.id FROM ej_l l WHERE NOT EXISTS (SELECT 1 FROM ej_r r WHERE l.v >= r.v); +---- +4 + +# ------------------------------------------------------------------ +# Correlation written inner-column-first, and with an expression +# ------------------------------------------------------------------ +# `r.v < l.v` is the same join as `l.v > r.v`; the planner flips the operator so +# the marked (left) side stays buffered. Same rows as `l.v > r.v` above. +query I rowsort +SELECT l.id FROM ej_l l WHERE EXISTS (SELECT 1 FROM ej_r r WHERE r.v < l.v); +---- +1 +2 + +# An expression on the streamed side: `l.v < r.v + 1` is `l.v <= r.v` over ints. +query I rowsort +SELECT l.id FROM ej_l l WHERE EXISTS (SELECT 1 FROM ej_r r WHERE l.v < r.v + 1); +---- +2 +3 +4 + +# ------------------------------------------------------------------ +# NULL semantics on the buffered (left) key +# ------------------------------------------------------------------ +# A NULL key satisfies no comparison, so it is excluded from EXISTS and kept by +# NOT EXISTS. ej_ln.v = {NULL,3,NULL,1}, right non-null key = {2}. +statement ok +CREATE TABLE ej_ln(id INT, v INT); + +statement ok +INSERT INTO ej_ln VALUES (1, NULL), (2, 3), (3, NULL), (4, 1); + +statement ok +CREATE TABLE ej_rn(v INT); + +statement ok +INSERT INTO ej_rn VALUES (2), (NULL); + +# `<` : only v=1 is below 2. Both NULL-keyed left rows stay in NOT EXISTS. +query I rowsort +SELECT l.id FROM ej_ln l WHERE EXISTS (SELECT 1 FROM ej_rn r WHERE l.v < r.v); +---- +4 + +query I rowsort +SELECT l.id FROM ej_ln l WHERE NOT EXISTS (SELECT 1 FROM ej_rn r WHERE l.v < r.v); +---- +1 +2 +3 + +# `>` : only v=3 exceeds 2. +query I rowsort +SELECT l.id FROM ej_ln l WHERE EXISTS (SELECT 1 FROM ej_rn r WHERE l.v > r.v); +---- +2 + +query I rowsort +SELECT l.id FROM ej_ln l WHERE NOT EXISTS (SELECT 1 FROM ej_rn r WHERE l.v > r.v); +---- +1 +3 +4 + +# ------------------------------------------------------------------ +# Empty inputs +# ------------------------------------------------------------------ +statement ok +CREATE TABLE ej_empty(v INT); + +# Empty streamed (right) side: no key can match, so EXISTS is empty and NOT +# EXISTS keeps every left row. +query I rowsort +SELECT l.id FROM ej_l l WHERE EXISTS (SELECT 1 FROM ej_empty r WHERE l.v > r.v); +---- + +query I rowsort +SELECT l.id FROM ej_l l WHERE NOT EXISTS (SELECT 1 FROM ej_empty r WHERE l.v > r.v); +---- +1 +2 +3 +4 + +statement ok +CREATE TABLE ej_lempty(id INT, v INT); + +# Empty buffered (left) side: nothing to emit either way. +query I rowsort +SELECT l.id FROM ej_lempty l WHERE EXISTS (SELECT 1 FROM ej_r r WHERE l.v > r.v); +---- + +query I rowsort +SELECT l.id FROM ej_lempty l WHERE NOT EXISTS (SELECT 1 FROM ej_r r WHERE l.v > r.v); +---- + +# ------------------------------------------------------------------ +# All-NULL sides (distinct early-exit triggers in the existence path) +# ------------------------------------------------------------------ +# An all-NULL buffered side saturates the watermark before the first poll, so no +# streamed batch is read; EXISTS is still empty. +statement ok +CREATE TABLE ej_l_allnull(id INT, v INT); + +statement ok +INSERT INTO ej_l_allnull VALUES (1, NULL), (2, NULL); + +query I rowsort +SELECT l.id FROM ej_l_allnull l WHERE EXISTS (SELECT 1 FROM ej_r r WHERE l.v > r.v); +---- + +query I rowsort +SELECT l.id FROM ej_l_allnull l WHERE NOT EXISTS (SELECT 1 FROM ej_r r WHERE l.v > r.v); +---- +1 +2 + +# An all-NULL streamed side never lowers the watermark, so it behaves like an +# empty streamed side. +statement ok +CREATE TABLE ej_r_allnull(v INT); + +statement ok +INSERT INTO ej_r_allnull VALUES (NULL), (NULL); + +query I rowsort +SELECT l.id FROM ej_l l WHERE EXISTS (SELECT 1 FROM ej_r_allnull r WHERE l.v > r.v); +---- + +query I rowsort +SELECT l.id FROM ej_l l WHERE NOT EXISTS (SELECT 1 FROM ej_r_allnull r WHERE l.v > r.v); +---- +1 +2 +3 +4 + +# ------------------------------------------------------------------ +# Duplicate buffered keys straddling the match boundary +# ------------------------------------------------------------------ +# The first matching buffered row is found by binary search; it must return the +# FIRST index of a run of equal keys or earlier duplicates vanish from EXISTS. +# ej_dup_l.v = {5,5,3,3,1}; the deciding streamed key is 3, so `<=`/`>=` land +# inside the run of 3s. +statement ok +CREATE TABLE ej_dup_l(id INT, v INT); + +statement ok +INSERT INTO ej_dup_l VALUES (1, 5), (2, 5), (3, 3), (4, 3), (5, 1); + +statement ok +CREATE TABLE ej_dup_r(v INT); + +statement ok +INSERT INTO ej_dup_r VALUES (3); + +# `<` : only v=1 is below 3. +query I rowsort +SELECT l.id FROM ej_dup_l l WHERE EXISTS (SELECT 1 FROM ej_dup_r r WHERE l.v < r.v); +---- +5 + +# `<=` : both v=3 rows must appear. +query I rowsort +SELECT l.id FROM ej_dup_l l WHERE EXISTS (SELECT 1 FROM ej_dup_r r WHERE l.v <= r.v); +---- +3 +4 +5 + +# `>` : only the two v=5 rows exceed 3. +query I rowsort +SELECT l.id FROM ej_dup_l l WHERE EXISTS (SELECT 1 FROM ej_dup_r r WHERE l.v > r.v); +---- +1 +2 + +# `>=` : the boundary is inside the run of 3s from the other direction. +query I rowsort +SELECT l.id FROM ej_dup_l l WHERE EXISTS (SELECT 1 FROM ej_dup_r r WHERE l.v >= r.v); +---- +1 +2 +3 +4 + +query I rowsort +SELECT l.id FROM ej_dup_l l WHERE NOT EXISTS (SELECT 1 FROM ej_dup_r r WHERE l.v >= r.v); +---- +5 + +# ------------------------------------------------------------------ +# Subquery filter -> repartitioned streamed side (multi-partition final pass) +# ------------------------------------------------------------------ +# A predicate on the inner table is pushed to the scan and lets the streamed +# side repartition, so the final existence pass is coordinated across streamed +# partitions. `r.v > 0` keeps all of {2,3,4}, so the rows match `l.v > r.v`. +query I rowsort +SELECT l.id FROM ej_l l WHERE EXISTS (SELECT 1 FROM ej_r r WHERE l.v > r.v AND r.v > 0); +---- +1 +2 + +query I rowsort +SELECT l.id FROM ej_l l WHERE NOT EXISTS (SELECT 1 FROM ej_r r WHERE l.v > r.v AND r.v > 0); +---- +3 +4 + +# ------------------------------------------------------------------ +# Type coverage: the existence path shares the join comparator +# ------------------------------------------------------------------ + +# Date32 key. +statement ok +CREATE TABLE ej_dl(id INT, d DATE); + +statement ok +INSERT INTO ej_dl VALUES (1, DATE '2022-04-23'), (2, DATE '2022-04-28'), (3, DATE '2022-04-18'); + +statement ok +CREATE TABLE ej_dr(d DATE); + +statement ok +INSERT INTO ej_dr VALUES (DATE '2022-04-20'), (DATE '2022-04-26'); + +query I rowsort +SELECT l.id FROM ej_dl l WHERE EXISTS (SELECT 1 FROM ej_dr r WHERE l.d > r.d); +---- +1 +2 + +query I rowsort +SELECT l.id FROM ej_dl l WHERE NOT EXISTS (SELECT 1 FROM ej_dr r WHERE l.d > r.d); +---- +3 + +# Float key with negative zero: -0.0 and +0.0 compare equal in SQL, so +# `-0.0 < 0.0` is false and `-0.0 <= 0.0` is true. The comparator normalizes the +# sign of zero before comparing. +statement ok +CREATE TABLE ej_fl(id INT, v DOUBLE); + +statement ok +INSERT INTO ej_fl VALUES (1, -0.0), (2, 2.5); + +statement ok +CREATE TABLE ej_fr(v DOUBLE); + +statement ok +INSERT INTO ej_fr VALUES (0.0); + +query I rowsort +SELECT l.id FROM ej_fl l WHERE EXISTS (SELECT 1 FROM ej_fr r WHERE l.v < r.v); +---- + +query I rowsort +SELECT l.id FROM ej_fl l WHERE EXISTS (SELECT 1 FROM ej_fr r WHERE l.v <= r.v); +---- +1 + +# String key. +statement ok +CREATE TABLE ej_sl(id INT, s VARCHAR); + +statement ok +INSERT INTO ej_sl VALUES (1, 'apple'), (2, 'cherry'), (3, 'mango'); + +statement ok +CREATE TABLE ej_sr(s VARCHAR); + +statement ok +INSERT INTO ej_sr VALUES ('banana'), ('lemon'); + +query I rowsort +SELECT l.id FROM ej_sl l WHERE EXISTS (SELECT 1 FROM ej_sr r WHERE l.s > r.s); +---- +2 +3 + +query I rowsort +SELECT l.id FROM ej_sl l WHERE NOT EXISTS (SELECT 1 FROM ej_sr r WHERE l.s > r.s); +---- +1 + +# Dictionary-encoded key: no typed arrow min/max kernel, so the extreme key per +# streamed batch is chosen through the generic ScalarValue path. +statement ok +CREATE TABLE ej_dict_l AS + SELECT column1 AS id, arrow_cast(column2, 'Dictionary(Int32, Utf8)') AS v + FROM (VALUES (1, 'a'), (2, 'c'), (3, 'e'), (4, NULL)); + +statement ok +CREATE TABLE ej_dict_r AS + SELECT arrow_cast(column1, 'Dictionary(Int32, Utf8)') AS v FROM (VALUES ('c')); + +query I rowsort +SELECT l.id FROM ej_dict_l l WHERE EXISTS (SELECT 1 FROM ej_dict_r r WHERE l.v > r.v); +---- +3 + +query I rowsort +SELECT l.id FROM ej_dict_l l WHERE NOT EXISTS (SELECT 1 FROM ej_dict_r r WHERE l.v > r.v); +---- +1 +2 +4 + +query I rowsort +SELECT l.id FROM ej_dict_l l WHERE EXISTS (SELECT 1 FROM ej_dict_r r WHERE l.v <= r.v); +---- +1 +2 + +# ------------------------------------------------------------------ +# Multi-batch stress: streamed side from range() (fragments by batch_size) +# ------------------------------------------------------------------ +# range() is a LazyMemoryExec that emits ceil(n / batch_size) batches, so +# batch_size genuinely fragments the streamed side (unlike a VALUES table, whose +# split is incidental). Plan proof in piecewise_merge_join_batches.slt. Buffered +# v = 1..10; streamed range(3,8) = {3,4,5,6,7}; counts are matrix-invariant. +statement ok +CREATE TABLE ej_range_l(id INT, v INT) AS + SELECT CAST(value AS INT) AS id, CAST(value AS INT) AS v FROM range(1, 11); + +# Existence, `>`: v > min(3) -> {4..10} = 7. +query I +SELECT count(*) FROM ej_range_l l WHERE EXISTS (SELECT 1 FROM range(3, 8) r WHERE l.v > r.value); +---- +7 + +query I +SELECT count(*) FROM ej_range_l l WHERE NOT EXISTS (SELECT 1 FROM range(3, 8) r WHERE l.v > r.value); +---- +3 + +# `<`: v < max(7) -> {1..6} = 6. +query I +SELECT count(*) FROM ej_range_l l WHERE EXISTS (SELECT 1 FROM range(3, 8) r WHERE l.v < r.value); +---- +6 + +query I +SELECT count(*) FROM ej_range_l l WHERE NOT EXISTS (SELECT 1 FROM range(3, 8) r WHERE l.v < r.value); +---- +4 + +# `>=` with a pushed-down inner filter that keeps every row, so the streamed side +# also repartitions. v >= min(3) -> {3..10} = 8. +query I +SELECT count(*) FROM ej_range_l l WHERE EXISTS (SELECT 1 FROM range(3, 8) r WHERE l.v >= r.value AND r.value > 0); +---- +8 + +query I +SELECT count(*) FROM ej_range_l l WHERE NOT EXISTS (SELECT 1 FROM range(3, 8) r WHERE l.v >= r.value AND r.value > 0); +---- +2 + +# Classic range join over the same range() streamed side. INNER l.v < r.value: +# sum over v of |{r : r > v}| = 5+5+4+3+2+1 = 20. +query I +SELECT count(*) FROM ej_range_l l JOIN range(3, 8) r ON l.v < r.value; +---- +20 + +# LEFT adds the unmatched left rows (v >= 7, which is v in {7,8,9,10}): 20 + 4 = 24. +query I +SELECT count(*) FROM ej_range_l l LEFT JOIN range(3, 8) r ON l.v < r.value; +---- +24 + +# Existence result feeding an aggregate: EXISTS l.v>r.v over the small fixture is +# {1,2}, so the count is 2. +query I +SELECT count(*) FROM ej_l l WHERE EXISTS (SELECT 1 FROM ej_r r WHERE l.v > r.v); +---- +2 + +# ------------------------------------------------------------------ +# Other existence joins: RightSemi / RightAnti / Mark +# ------------------------------------------------------------------ +# PWMJ rejects these, so both matrix combinations run them on NestedLoopJoin; +# only batch_size varies. Included so the file covers every existence join type. +statement ok +CREATE TABLE ej_rs_l(v INT); + +statement ok +INSERT INTO ej_rs_l VALUES (3), (6); + +statement ok +CREATE TABLE ej_rs_r(id INT, v INT); + +statement ok +INSERT INTO ej_rs_r VALUES (1, 1), (2, 4), (3, 7); + +# RIGHT SEMI: keep right rows that have some larger left value. left={3,6}: +# right 1 (l=3,6), right 4 (l=6) qualify; right 7 has none. +query I rowsort +SELECT r.id FROM ej_rs_l l RIGHT SEMI JOIN ej_rs_r r ON l.v > r.v; +---- +1 +2 + +# RIGHT ANTI: the complement -> only right 7. +query I rowsort +SELECT r.id FROM ej_rs_l l RIGHT ANTI JOIN ej_rs_r r ON l.v > r.v; +---- +3 + +# Mark join: EXISTS used inside a disjunction produces a LeftMark. `l.v > r.v` +# over the fixture yields {1,2}; `l.id = 4` adds id 4. +query I rowsort +SELECT l.id FROM ej_l l WHERE l.id = 4 OR EXISTS (SELECT 1 FROM ej_r r WHERE l.v > r.v); +---- +1 +2 +4 diff --git a/datafusion/sqllogictest/test_files/sort_merge_join_batches.slt b/datafusion/sqllogictest/test_files/sort_merge_join_batches.slt new file mode 100644 index 0000000000000..eb26d0a715104 --- /dev/null +++ b/datafusion/sqllogictest/test_files/sort_merge_join_batches.slt @@ -0,0 +1,65 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at + +# http://www.apache.org/licenses/LICENSE-2.0 + +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +# Plan proof for sort_merge_join_matrix.slt's cross-batch coverage. SMJ sorts +# both inputs and SortExec re-batches to batch_size, so equal-key runs span +# batches even from a VALUES source. Not a matrix file: it fixes batch_size=1 to +# assert (via EXPLAIN ANALYZE) that the join then receives one row per batch +# (input_batches=16 = input_rows), i.e. the runs of five 1s and three 2s are +# split across batches. Fixtures match sort_merge_join_matrix.slt. + +statement ok +set datafusion.optimizer.prefer_hash_join = false; + +statement ok +set datafusion.execution.target_partitions = 4; + +statement ok +set datafusion.optimizer.repartition_joins = true; + +statement ok +set datafusion.execution.batch_size = 1; + +statement ok +CREATE TABLE smj_dup_l(k INT) AS VALUES + (1), (1), (1), (1), (1), (2), (2), (2), (9); + +statement ok +CREATE TABLE smj_dup_r(k INT) AS VALUES + (1), (1), (1), (1), (2), (2), (3); + +query TT +EXPLAIN ANALYZE SELECT count(*) FROM smj_dup_l l JOIN smj_dup_r r ON l.k = r.k; +---- +Plan with Metrics +01) +02) +03) +04) +05)SortMergeJoinExec:input_batches=16, input_rows=16 + + +query I +SELECT count(*) FROM smj_dup_l l JOIN smj_dup_r r ON l.k = r.k; +---- +26 + +statement ok +RESET datafusion.execution.batch_size; + +statement ok +RESET datafusion.optimizer.prefer_hash_join; diff --git a/datafusion/sqllogictest/test_files/sort_merge_join_matrix.slt b/datafusion/sqllogictest/test_files/sort_merge_join_matrix.slt new file mode 100644 index 0000000000000..7f1857782cb38 --- /dev/null +++ b/datafusion/sqllogictest/test_files/sort_merge_join_matrix.slt @@ -0,0 +1,315 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at + +# http://www.apache.org/licenses/LICENSE-2.0 + +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +# Sort-merge-join correctness across a config matrix. Every equijoin runs on +# SortMergeJoin (prefer_hash_join=false) and HashJoin (true), so sweeping +# {true,false} x batch_size {1,2,100,8192} requires them to agree row-for-row. +# batch_size=1 forces equal-key runs to span single-row batches (historically +# bug-prone). SMJ needs target_partitions>1 and repartition_joins, set below (not +# swept). Null-aware NOT IN always uses HashJoin, so it is out of scope (see +# null_aware_anti_join.slt). Matrix rules: no EXPLAIN, no in-file SET of a swept +# knob, rowsort every multi-row query. + +# configMatrix: datafusion.optimizer.prefer_hash_join=true,false +# configMatrix: datafusion.execution.batch_size=1,2,100,8192 + +statement ok +set datafusion.execution.target_partitions = 4; + +statement ok +set datafusion.optimizer.repartition_joins = true; + +# ------------------------------------------------------------------ +# Fixtures: duplicate keys, an unmatched key on each side, and a NULL key on +# each side (a NULL key never matches under default equality). +# ------------------------------------------------------------------ +statement ok +CREATE TABLE smj_l(k INT, v VARCHAR) AS VALUES + (1, 'a'), (1, 'b'), (2, 'c'), (3, 'd'), (NULL, 'e'); + +statement ok +CREATE TABLE smj_r(k INT, w VARCHAR) AS VALUES + (1, 'x'), (1, 'y'), (2, 'z'), (4, 'q'), (NULL, 'n'); + +# ------------------------------------------------------------------ +# Basic: Inner / Left / Right / Full on a single equi key +# ------------------------------------------------------------------ +# INNER: k=1 gives 2x2, k=2 gives 1x1; k=3, k=4 and both NULLs are unmatched. +query ITT rowsort +SELECT l.k, l.v, r.w FROM smj_l l JOIN smj_r r ON l.k = r.k; +---- +1 a x +1 a y +1 b x +1 b y +2 c z + +# LEFT: the INNER rows plus unmatched left rows (k=3, NULL) padded with NULL. +query ITT rowsort +SELECT l.k, l.v, r.w FROM smj_l l LEFT JOIN smj_r r ON l.k = r.k; +---- +1 a x +1 a y +1 b x +1 b y +2 c z +3 d NULL +NULL e NULL + +# RIGHT: the INNER rows plus unmatched right rows (k=4, NULL) padded with NULL. +query TIT rowsort +SELECT l.v, r.k, r.w FROM smj_l l RIGHT JOIN smj_r r ON l.k = r.k; +---- +NULL 4 q +NULL NULL n +a 1 x +a 1 y +b 1 x +b 1 y +c 2 z + +# FULL: INNER rows plus every unmatched row from both sides. +query ITT rowsort +SELECT l.k, l.v, r.w FROM smj_l l FULL JOIN smj_r r ON l.k = r.k; +---- +1 a x +1 a y +1 b x +1 b y +2 c z +3 d NULL +NULL NULL n +NULL NULL q +NULL e NULL + +# ------------------------------------------------------------------ +# Existence joins over an equi key: LeftSemi / LeftAnti / RightSemi / RightAnti +# ------------------------------------------------------------------ +# LeftSemi (EXISTS): left rows whose key matches, each kept once (k=1 twice +# because the left side has two such rows, not because the right side does). +query IT rowsort +SELECT l.k, l.v FROM smj_l l WHERE EXISTS (SELECT 1 FROM smj_r r WHERE l.k = r.k); +---- +1 a +1 b +2 c + +# LeftAnti (NOT EXISTS): the complementary left rows, including the NULL key. +query IT rowsort +SELECT l.k, l.v FROM smj_l l WHERE NOT EXISTS (SELECT 1 FROM smj_r r WHERE l.k = r.k); +---- +3 d +NULL e + +# RightSemi: right rows whose key matches some left row. +query IT rowsort +SELECT r.k, r.w FROM smj_l l RIGHT SEMI JOIN smj_r r ON l.k = r.k; +---- +1 x +1 y +2 z + +# RightAnti: the complementary right rows, including the NULL key. +query IT rowsort +SELECT r.k, r.w FROM smj_l l RIGHT ANTI JOIN smj_r r ON l.k = r.k; +---- +4 q +NULL n + +# ------------------------------------------------------------------ +# Null-safe equality: `IS NOT DISTINCT FROM` makes NULL match NULL +# ------------------------------------------------------------------ +# The INNER result gains the (NULL,e)-(NULL,n) pair that plain `=` drops. +query ITT rowsort +SELECT l.k, l.v, r.w FROM smj_l l JOIN smj_r r ON l.k IS NOT DISTINCT FROM r.k; +---- +1 a x +1 a y +1 b x +1 b y +2 c z +NULL e n + +# ------------------------------------------------------------------ +# Multiple equi keys +# ------------------------------------------------------------------ +statement ok +CREATE TABLE smj_mk_l(k1 INT, k2 INT, v INT) AS VALUES + (1, 1, 10), (1, 2, 20), (2, 2, 30), (NULL, 1, 40); + +statement ok +CREATE TABLE smj_mk_r(k1 INT, k2 INT, w INT) AS VALUES + (1, 1, 100), (1, 2, 200), (2, 2, 300), (1, 3, 400); + +query IIII rowsort +SELECT l.k1, l.k2, l.v, r.w +FROM smj_mk_l l JOIN smj_mk_r r ON l.k1 = r.k1 AND l.k2 = r.k2; +---- +1 1 10 100 +1 2 20 200 +2 2 30 300 + +query IIII rowsort +SELECT l.k1, l.k2, l.v, r.w +FROM smj_mk_l l LEFT JOIN smj_mk_r r ON l.k1 = r.k1 AND l.k2 = r.k2; +---- +1 1 10 100 +1 2 20 200 +2 2 30 300 +NULL 1 40 NULL + +# ------------------------------------------------------------------ +# Equi key plus a non-equi join filter in the ON clause +# ------------------------------------------------------------------ +statement ok +CREATE TABLE smj_a(k INT, x INT) AS VALUES (1, 10), (1, 20), (2, 30), (3, 40); + +statement ok +CREATE TABLE smj_b(k INT, y INT) AS VALUES (1, 15), (1, 25), (2, 5), (4, 50); + +# INNER: only k=1 has matches, kept where x < y. +query IIII rowsort +SELECT l.k, l.x, r.k, r.y FROM smj_a l JOIN smj_b r ON l.k = r.k AND l.x < r.y; +---- +1 10 1 15 +1 10 1 25 +1 20 1 25 + +# LEFT: rows whose key never matches, or matches but fails the filter, keep the +# left row with NULLs (k=2's x=30 fails 30<5; k=3 has no right key). +query III rowsort +SELECT l.k, l.x, r.y FROM smj_a l LEFT JOIN smj_b r ON l.k = r.k AND l.x < r.y; +---- +1 10 15 +1 10 25 +1 20 25 +2 30 NULL +3 40 NULL + +# FULL with the same filter: unmatched rows from both sides survive. +query IIII rowsort +SELECT l.k, l.x, r.k, r.y FROM smj_a l FULL JOIN smj_b r ON l.k = r.k AND l.x < r.y; +---- +1 10 1 15 +1 10 1 25 +1 20 1 25 +2 30 NULL NULL +3 40 NULL NULL +NULL NULL 2 5 +NULL NULL 4 50 + +# ------------------------------------------------------------------ +# Non-integer (string) key +# ------------------------------------------------------------------ +statement ok +CREATE TABLE smj_sk_l(k VARCHAR, v INT) AS VALUES ('a', 1), ('b', 2), ('a', 3); + +statement ok +CREATE TABLE smj_sk_r(k VARCHAR, w INT) AS VALUES ('a', 10), ('c', 20); + +query TII rowsort +SELECT l.k, l.v, r.w FROM smj_sk_l l JOIN smj_sk_r r ON l.k = r.k; +---- +a 1 10 +a 3 10 + +query TII rowsort +SELECT l.k, l.v, r.w FROM smj_sk_l l FULL JOIN smj_sk_r r ON l.k = r.k; +---- +NULL NULL 20 +a 1 10 +a 3 10 +b 2 NULL + +# ------------------------------------------------------------------ +# Empty inputs +# ------------------------------------------------------------------ +statement ok +CREATE TABLE smj_empty(k INT, w VARCHAR); + +# INNER against an empty side is empty; LEFT keeps every left row with NULLs. +query ITT rowsort +SELECT l.k, l.v, r.w FROM smj_l l JOIN smj_empty r ON l.k = r.k; +---- + +query ITT rowsort +SELECT l.k, l.v, r.w FROM smj_l l LEFT JOIN smj_empty r ON l.k = r.k; +---- +1 a NULL +1 b NULL +2 c NULL +3 d NULL +NULL e NULL + +# Empty left side with a RIGHT join keeps every right row with NULLs. +query IIT rowsort +SELECT l.k, r.k, r.w FROM smj_empty l RIGHT JOIN smj_r r ON l.k = r.k; +---- +NULL 1 x +NULL 1 y +NULL 2 z +NULL 4 q +NULL NULL n + +# ------------------------------------------------------------------ +# Duplicate keys across batch boundaries (verified by count) +# ------------------------------------------------------------------ +# SortExec re-batches to batch_size, so equal-key runs span batches (proof in +# sort_merge_join_batches.slt). Left: five k=1, three k=2, one unmatched k=9; +# right: four k=1, two k=2, one unmatched k=3. INNER = 5*4 + 3*2 = 26. +statement ok +CREATE TABLE smj_dup_l(k INT) AS VALUES + (1), (1), (1), (1), (1), (2), (2), (2), (9); + +statement ok +CREATE TABLE smj_dup_r(k INT) AS VALUES + (1), (1), (1), (1), (2), (2), (3); + +query I +SELECT count(*) FROM smj_dup_l l JOIN smj_dup_r r ON l.k = r.k; +---- +26 + +# LEFT = INNER + unmatched left k=9 = 27. +query I +SELECT count(*) FROM smj_dup_l l LEFT JOIN smj_dup_r r ON l.k = r.k; +---- +27 + +# RIGHT = INNER + unmatched right k=3 = 27. +query I +SELECT count(*) FROM smj_dup_l l RIGHT JOIN smj_dup_r r ON l.k = r.k; +---- +27 + +# FULL = INNER + both unmatched = 28. +query I +SELECT count(*) FROM smj_dup_l l FULL JOIN smj_dup_r r ON l.k = r.k; +---- +28 + +# LeftSemi = left rows with a match = 5 + 3 = 8. +query I +SELECT count(*) FROM smj_dup_l l WHERE EXISTS (SELECT 1 FROM smj_dup_r r WHERE l.k = r.k); +---- +8 + +# LeftAnti = unmatched left rows = 1 (k=9). +query I +SELECT count(*) FROM smj_dup_l l WHERE NOT EXISTS (SELECT 1 FROM smj_dup_r r WHERE l.k = r.k); +---- +1 From 124291e113b8a0991a3abad585703c4a55b948cf Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Tue, 1 Sep 2026 19:31:44 +0000 Subject: [PATCH 02/37] Align metadata propagation through Physical and Logical casts (#23169) ## Which issue does this PR close? - Closes #22079 - Closes #24724 ## Rationale for this change The logical `Expr::Cast` and `Expr::TryCast` have a `FieldRef` target that was added in https://github.com/apache/datafusion/pull/18136 so that logical casts can express a cast to an extension type. In combination with a SQL type planner ( https://github.com/apache/datafusion/pull/20676 ) and an optimizer rule, this enabled casting to/from extension types with custom semantics to actually occur. The ability to do this was reverted by https://github.com/apache/datafusion/pull/20836 (which removed the original test) and I am not sure that ability ever made it into a release. When investigating this issue, it became clear the logical and physical cast behaviour had diverged with respect to the target field. ## What changes are included in this PR? This PR strips specific metadata keys (extension name and extension metadata) when propagating metadata from the source of a cast to the target (because doing so may result in an invalid destination field that consumers could reject), and propagates all metadata from the (logical) cast target field (e.g., so that a cast to an extension type represented by the cast target field will have a `to_field()` that communicates the extension type). For the physical cast, this behaviour is replicated exactly (I hope). Note that actually casting to an extension type can be implemented with an optimizer rule, planner, or by the mechanism I have in the works in https://github.com/apache/datafusion/pull/21071 . ## Are these changes tested? Yes ## Are there any user-facing changes? It was in practice not common to create a `Expr::Cast` with field metadata internally and thus I don't think users will see metadata changes from the inclusion of metadata from the target field. I would be surprised if stripping the extension name/metadata from the source was disruptive (it was more likely to have caused errors). Superceeds an earlier but similar attempt ( https://github.com/apache/datafusion/pull/22162 ). --------- Co-authored-by: Andrew Lamb Co-authored-by: Tim Saucer --- Cargo.lock | 1 + .../custom_data_source/custom_file_casts.rs | 5 +- datafusion/expr/src/expr_schema.rs | 167 +++++++- datafusion/functions/src/core/arrow_cast.rs | 26 +- .../functions/src/core/arrow_try_cast.rs | 24 +- .../physical-expr-adapter/src/rewrite.rs | 7 +- datafusion/physical-expr/Cargo.toml | 1 + .../physical-expr/src/expressions/cast.rs | 391 ++++++++++++++---- .../physical-expr/src/expressions/mod.rs | 2 +- .../physical-expr/src/expressions/try_cast.rs | 354 +++++++++++++++- datafusion/physical-expr/src/planner.rs | 159 +++++-- .../proto/src/logical_plan/from_proto.rs | 6 +- .../tests/cases/roundtrip_logical_plan.rs | 26 ++ .../cast_extension_type_metadata.slt | 54 ++- 14 files changed, 1069 insertions(+), 154 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 7a68fcb5e63d7..52797394237c0 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2402,6 +2402,7 @@ name = "datafusion-physical-expr" version = "55.0.0" dependencies = [ "arrow", + "arrow-schema", "criterion", "datafusion-common", "datafusion-expr", diff --git a/datafusion-examples/examples/custom_data_source/custom_file_casts.rs b/datafusion-examples/examples/custom_data_source/custom_file_casts.rs index 71addc6d1bcb0..202c0a71257e9 100644 --- a/datafusion-examples/examples/custom_data_source/custom_file_casts.rs +++ b/datafusion-examples/examples/custom_data_source/custom_file_casts.rs @@ -188,10 +188,11 @@ impl PhysicalExprAdapter for CustomCastsPhysicalExprAdapter { if let Some(cast) = expr.downcast_ref::() { let input_data_type = cast.expr().data_type(&self.physical_file_schema)?; - let output_data_type = cast.target_field().data_type(); + let output_field = cast.target_field(); if !cast.is_bigger_cast(&input_data_type) { return not_impl_err!( - "Unsupported CAST from {input_data_type} to {output_data_type}" + "Unsupported CAST from {input_data_type} to {}", + output_field.data_type() ); } } diff --git a/datafusion/expr/src/expr_schema.rs b/datafusion/expr/src/expr_schema.rs index 36b76f076d26a..75b3a60af1465 100644 --- a/datafusion/expr/src/expr_schema.rs +++ b/datafusion/expr/src/expr_schema.rs @@ -33,6 +33,7 @@ use crate::{ use arrow::compute::can_cast_types; use arrow::datatypes::FieldRef; use arrow::datatypes::{DataType, Field}; +use arrow_schema::extension::{EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_NAME_KEY}; use datafusion_common::datatype::FieldExt; use datafusion_common::{ Column, DataFusionError, ExprSchema, Result, ScalarValue, Spans, TableReference, @@ -71,18 +72,46 @@ pub trait ExprSchemable { -> Result<(DataType, bool)>; } -/// Derives the output field for a cast expression from the source field. +/// Derives the output field for a cast expression from the source and target fields. +/// +/// Metadata handling: +/// - Type-only casts (i.e., target_field == DataType::SomeDataType.into_nullable_field()) +/// propagate non extension-type metadata from the source. This is for backward compatibility +/// (casts have propagated source metadata for many if not all previous versions), recognizing +/// that the return type of `::` should have the +/// return type of `` (e.g., casting arrow.json to utf8). +/// - All other casts preserve target metadata exactly. This ensures in particular that output +/// metadata when casting to an extension type contains the extension information in the +/// output field. Callers that wish to have some mix of source and target metadata can use +/// Alias or construct an output field themselves (whose metadata will be used directly). +/// /// For `TryCast`, `force_nullable` is `true` since a failed cast returns NULL. fn cast_output_field( source_field: &FieldRef, - target_type: &DataType, + target_field: &FieldRef, force_nullable: bool, ) -> Arc { + // Check if this is a "type-only" cast (target_field == DataType::X.into_nullable_field()) + let is_type_only = target_field.name().is_empty() + && target_field.is_nullable() + && target_field.metadata().is_empty(); + + let metadata = if is_type_only { + // Type-only cast: propagate source metadata, stripping extension type keys + let mut meta = source_field.metadata().clone(); + meta.remove(EXTENSION_TYPE_NAME_KEY); + meta.remove(EXTENSION_TYPE_METADATA_KEY); + meta + } else { + // Explicit target field: use target metadata exactly + target_field.metadata().clone() + }; + let mut f = source_field .as_ref() .clone() - .with_data_type(target_type.clone()) - .with_metadata(source_field.metadata().clone()); + .with_data_type(target_field.data_type().clone()) + .with_metadata(metadata); if force_nullable { f = f.with_nullable(true); } @@ -478,7 +507,8 @@ impl ExprSchemable for Expr { /// - **Aliases**: Merge underlying expr metadata with alias-specific metadata, preferring the alias metadata /// - **Binary expressions**: field metadata is empty /// - **Boolean expressions**: field metadata is empty - /// - **Cast expressions**: determined by the input expression's field metadata handling + /// - **Cast expressions**: Type-only casts pass through source metadata (stripping extension + /// type keys); casts with explicit target fields use target metadata exactly /// - **Scalar functions**: Generate metadata via function's [`return_field_from_args`] method, /// with the default implementation returning empty field metadata /// - **Aggregate functions**: Generate metadata via function's [`return_field`] method, @@ -623,20 +653,16 @@ impl ExprSchemable for Expr { func.return_field_from_args(args) } // _ => Ok((self.get_type(schema)?, self.nullable(schema)?)), - Expr::Cast(Cast { expr, field }) => { - expr.to_field(schema).map(|(_table_ref, src)| { - cast_output_field(&src, field.data_type(), false) - }) - } + Expr::Cast(Cast { expr, field }) => expr + .to_field(schema) + .map(|(_table_ref, src)| cast_output_field(&src, field, false)), Expr::Placeholder(Placeholder { id: _, field: Some(field), }) => Ok(Arc::clone(field).renamed(&schema_name)), - Expr::TryCast(TryCast { expr, field }) => { - expr.to_field(schema).map(|(_table_ref, src)| { - cast_output_field(&src, field.data_type(), true) - }) - } + Expr::TryCast(TryCast { expr, field }) => expr + .to_field(schema) + .map(|(_table_ref, src)| cast_output_field(&src, field, true)), Expr::LambdaVariable(LambdaVariable { field: Some(field), .. }) => Ok(Arc::clone(field).renamed(&schema_name)), @@ -1427,4 +1453,115 @@ mod tests { assert_eq!(meta, expr.metadata(&schema).unwrap()); } + + #[test] + fn test_cast_and_try_cast_extension_type_metadata() { + use crate::expr::{Cast, TryCast}; + use arrow_schema::extension::{ + EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_NAME_KEY, + }; + + // Helper to build either Cast or TryCast expression + fn make_cast_expr( + expr: Expr, + target_field: FieldRef, + use_try_cast: bool, + ) -> Expr { + if use_try_cast { + Expr::TryCast(TryCast { + expr: Box::new(expr), + field: target_field, + }) + } else { + Expr::Cast(Cast { + expr: Box::new(expr), + field: target_field, + }) + } + } + + // Run the same test logic for both Cast and TryCast + for use_try_cast in [false, true] { + let cast_name = if use_try_cast { "TryCast" } else { "Cast" }; + + // Create a schema with a field that has extension type metadata + let mut source_meta = HashMap::new(); + source_meta.insert( + EXTENSION_TYPE_NAME_KEY.to_string(), + "arrow.uuid".to_string(), + ); + source_meta.insert("custom_key".to_string(), "custom_value".to_string()); + + let source_field = Field::new("foo", DataType::FixedSizeBinary(16), false) + .with_metadata(source_meta); + + let schema = MockExprSchema::new() + .with_data_type(DataType::FixedSizeBinary(16)) + .with_metadata(FieldMetadata::from(source_field.metadata().clone())); + + // Test 1: Cast to a type without extension metadata strips extension metadata + // but preserves non-extension metadata + let cast_expr = make_cast_expr( + col("foo"), + Arc::new(Field::new("", DataType::Utf8, true)), + use_try_cast, + ); + + let (_, result_field) = cast_expr.to_field(&schema).unwrap(); + assert!( + result_field + .metadata() + .get(EXTENSION_TYPE_NAME_KEY) + .is_none(), + "{cast_name}: Extension type name should be stripped when target has no extension metadata" + ); + assert_eq!( + result_field.metadata().get("custom_key"), + Some(&"custom_value".to_string()), + "{cast_name}: Non-extension metadata should be preserved" + ); + if use_try_cast { + assert!( + result_field.is_nullable(), + "TryCast result should be nullable" + ); + } + + // Test 2: Cast to a field with explicit metadata uses target metadata exactly + let mut target_meta = HashMap::new(); + target_meta.insert( + EXTENSION_TYPE_NAME_KEY.to_string(), + "arrow.json".to_string(), + ); + target_meta.insert(EXTENSION_TYPE_METADATA_KEY.to_string(), "{}".to_string()); + + let target_field = + Field::new("", DataType::Utf8, true).with_metadata(target_meta); + + let cast_expr = + make_cast_expr(col("foo"), Arc::new(target_field), use_try_cast); + + let (_, result_field) = cast_expr.to_field(&schema).unwrap(); + assert_eq!( + result_field.metadata().get(EXTENSION_TYPE_NAME_KEY), + Some(&"arrow.json".to_string()), + "{cast_name}: Extension type name should come from target field" + ); + assert_eq!( + result_field.metadata().get(EXTENSION_TYPE_METADATA_KEY), + Some(&"{}".to_string()), + "{cast_name}: Extension type metadata should come from target field" + ); + assert!( + result_field.metadata().get("custom_key").is_none(), + "{cast_name}: Source metadata should NOT propagate when target has explicit metadata" + ); + if use_try_cast { + assert!( + result_field.is_nullable(), + "TryCast result should be nullable" + ); + } + } + } } diff --git a/datafusion/functions/src/core/arrow_cast.rs b/datafusion/functions/src/core/arrow_cast.rs index 0b67883c17c87..929a92abcca23 100644 --- a/datafusion/functions/src/core/arrow_cast.rs +++ b/datafusion/functions/src/core/arrow_cast.rs @@ -27,8 +27,8 @@ use datafusion_common::{ use datafusion_expr::simplify::{ExprSimplifyResult, SimplifyContext}; use datafusion_expr::{ - Coercion, ColumnarValue, Documentation, Expr, ReturnFieldArgs, ScalarFunctionArgs, - ScalarUDFImpl, Signature, TypeSignatureClass, Volatility, + Coercion, ColumnarValue, Documentation, Expr, ExprSchemable, ReturnFieldArgs, + ScalarFunctionArgs, ScalarUDFImpl, Signature, TypeSignatureClass, Volatility, }; use datafusion_macros::user_doc; @@ -161,9 +161,27 @@ impl ScalarUDFImpl for ArrowCastFunc { let [source_arg, type_arg] = take_function_args(self.name(), args)?; let target_type = data_type_from_type_arg(self.name(), &type_arg)?; let source_type = info.get_data_type(&source_arg)?; + + // We can skip the cast only if: + // 1. The source and target types are the same + // 2. The source has no extension metadata that needs to be stripped let new_expr = if source_type == target_type { - // the argument's data type is already the correct type - source_arg + // Check if source has extension metadata + let source_field = source_arg.to_field(info.schema())?; + let has_extension_metadata = source_field + .1 + .metadata() + .contains_key("ARROW:extension:name"); + if has_extension_metadata { + // Need to create a cast to strip extension metadata + Expr::Cast(datafusion_expr::Cast { + expr: Box::new(source_arg), + field: target_type.into_nullable_field_ref(), + }) + } else { + // the argument's data type is already the correct type + source_arg + } } else { // Use an actual cast to get the correct type Expr::Cast(datafusion_expr::Cast { diff --git a/datafusion/functions/src/core/arrow_try_cast.rs b/datafusion/functions/src/core/arrow_try_cast.rs index d27b29ba5736d..0914695b60370 100644 --- a/datafusion/functions/src/core/arrow_try_cast.rs +++ b/datafusion/functions/src/core/arrow_try_cast.rs @@ -26,8 +26,8 @@ use datafusion_common::{ use datafusion_expr::simplify::{ExprSimplifyResult, SimplifyContext}; use datafusion_expr::{ - Coercion, ColumnarValue, Documentation, Expr, ReturnFieldArgs, ScalarFunctionArgs, - ScalarUDFImpl, Signature, TypeSignatureClass, Volatility, + Coercion, ColumnarValue, Documentation, Expr, ExprSchemable, ReturnFieldArgs, + ScalarFunctionArgs, ScalarUDFImpl, Signature, TypeSignatureClass, Volatility, }; use datafusion_macros::user_doc; @@ -134,8 +134,26 @@ impl ScalarUDFImpl for ArrowTryCastFunc { let target_type = data_type_from_type_arg(self.name(), &type_arg)?; let source_type = info.get_data_type(&source_arg)?; + + // We can skip the cast only if: + // 1. The source and target types are the same + // 2. The source has no extension metadata that needs to be stripped let new_expr = if source_type == target_type { - source_arg + // Check if source has extension metadata + let source_field = source_arg.to_field(info.schema())?; + let has_extension_metadata = source_field + .1 + .metadata() + .contains_key("ARROW:extension:name"); + if has_extension_metadata { + // Need to create a try_cast to strip extension metadata + Expr::TryCast(datafusion_expr::TryCast { + expr: Box::new(source_arg), + field: target_type.into_nullable_field_ref(), + }) + } else { + source_arg + } } else { Expr::TryCast(datafusion_expr::TryCast { expr: Box::new(source_arg), diff --git a/datafusion/physical-expr-adapter/src/rewrite.rs b/datafusion/physical-expr-adapter/src/rewrite.rs index 7345a587ee6a4..3d31feca96031 100644 --- a/datafusion/physical-expr-adapter/src/rewrite.rs +++ b/datafusion/physical-expr-adapter/src/rewrite.rs @@ -319,6 +319,9 @@ mod tests { assert_eq!(source.name(), "__datafusion_file_row_index"); assert_eq!(source.index(), 2); + // The row index column is at index 2, beyond the user-visible schema. + // When the source column lookup fails, the field name is empty. The + // correct field name would be provided by a parent projection/alias. let input_schema = Schema::new(vec![ Field::new("value", DataType::Int64, true), Field::new("__datafusion_file_row_index", DataType::Int64, false) @@ -328,9 +331,11 @@ mod tests { )])), ]); let return_field = expr.return_field(&input_schema)?; - assert_eq!(return_field.name(), "file_row_index"); + // Field name is empty because column index 2 is beyond the schema + assert_eq!(return_field.name(), ""); assert_eq!(return_field.data_type(), &DataType::Int64); assert!(return_field.is_nullable()); + // Exact target field does not preserve source metadata assert!(return_field.metadata().is_empty()); Ok(()) } diff --git a/datafusion/physical-expr/Cargo.toml b/datafusion/physical-expr/Cargo.toml index 65ef2a3ceb216..0588a777230fb 100644 --- a/datafusion/physical-expr/Cargo.toml +++ b/datafusion/physical-expr/Cargo.toml @@ -51,6 +51,7 @@ proto = [ [dependencies] arrow = { workspace = true } +arrow-schema = { workspace = true } datafusion-common = { workspace = true } datafusion-expr = { workspace = true } datafusion-expr-common = { workspace = true } diff --git a/datafusion/physical-expr/src/expressions/cast.rs b/datafusion/physical-expr/src/expressions/cast.rs index cb3103d38c52a..565d396dad1b5 100644 --- a/datafusion/physical-expr/src/expressions/cast.rs +++ b/datafusion/physical-expr/src/expressions/cast.rs @@ -15,6 +15,7 @@ // specific language governing permissions and limitations // under the License. +use std::collections::HashMap; use std::fmt; use std::hash::Hash; use std::sync::Arc; @@ -22,8 +23,9 @@ use std::sync::Arc; use crate::physical_expr::PhysicalExpr; use arrow::compute::{CastOptions, can_cast_types}; -use arrow::datatypes::{DataType, DataType::*, FieldRef, Schema}; +use arrow::datatypes::{DataType, DataType::*, Field, FieldRef, Schema}; use arrow::record_batch::RecordBatch; +use arrow_schema::extension::{EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_NAME_KEY}; use datafusion_common::datatype::DataTypeExt; use datafusion_common::format::DEFAULT_FORMAT_OPTIONS; use datafusion_common::nested_struct::{ @@ -59,8 +61,18 @@ fn can_cast_named_struct_types(source: &DataType, target: &DataType) -> bool { pub struct CastExpr { /// The expression to cast pub expr: Arc, - /// Field metadata describing the desired output after casting + /// The target field. + /// + /// For a type-only cast (see [`CastExpr::new`]) this is a field synthesized + /// from the target data type alone and only its data type is meaningful. + /// For a cast built from an explicit field (see + /// [`CastExpr::new_with_target_field`]) its metadata and nullability are + /// applied to the output field as-is. target_field: FieldRef, + /// Whether `target_field` was supplied by the caller (as opposed to being + /// synthesized from a `DataType`), and therefore whether its metadata and + /// nullability describe the output field exactly. + explicit_target: bool, /// Cast options cast_options: CastOptions<'static>, } @@ -68,8 +80,12 @@ pub struct CastExpr { // Manually derive PartialEq and Hash to work around https://github.com/rust-lang/rust/issues/78808 impl PartialEq for CastExpr { fn eq(&self, other: &Self) -> bool { + // Compare the semantically meaningful parts of the target field only: + // the field name never affects the output of this expression. self.expr.eq(&other.expr) - && self.target_field.eq(&other.target_field) + && self.cast_type().eq(other.cast_type()) + && self.target_metadata().eq(&other.target_metadata()) + && self.target_nullable().eq(&other.target_nullable()) && self.cast_options.eq(&other.cast_options) } } @@ -77,7 +93,17 @@ impl PartialEq for CastExpr { impl Hash for CastExpr { fn hash(&self, state: &mut H) { self.expr.hash(state); - self.target_field.hash(state); + self.cast_type().hash(state); + // Hash the metadata by iterating over sorted keys for deterministic ordering + if let Some(metadata) = self.target_metadata() { + let mut entries: Vec<_> = metadata.iter().collect(); + entries.sort_by_key(|(k, _)| *k); + for (k, v) in entries { + k.hash(state); + v.hash(state); + } + } + self.target_nullable().hash(state); self.cast_options.hash(state); } } @@ -85,39 +111,39 @@ impl Hash for CastExpr { impl CastExpr { /// Create a new `CastExpr` using only a `DataType`. /// - /// This constructor is provided for compatibility with existing call sites - /// that only know the target type. It synthesizes a ``Field`` with the - /// given type (**nullable by default**) and no name metadata. Callers that - /// already have a `FieldRef` (for example, coming from schema inference or a - /// resolved column) should prefer [`CastExpr::new_with_target_field`], which - /// preserves the field's name, nullability, and other metadata. In other - /// words: + /// This constructor creates a type-only cast where metadata and nullability + /// are passed through from the source expression (with extension type keys + /// stripped from metadata). This is the most common use case when you only + /// need to change the data type. /// - /// * use `new()` when only a `DataType` is available and you want the legacy - /// semantics of a type-only cast - /// * use `new_with_target_field()` when you need explicit field - /// metadata/name/nullability preserved + /// For explicit control over the output field's metadata and nullability, + /// use [`CastExpr::new_with_target_field`] or the individual builder methods. pub fn new( expr: Arc, cast_type: DataType, cast_options: Option>, ) -> Self { - Self::new_with_target_field( + Self { expr, - cast_type.into_nullable_field_ref(), - cast_options, - ) + target_field: cast_type.into_nullable_field_ref(), + explicit_target: false, + cast_options: cast_options.unwrap_or(DEFAULT_CAST_OPTIONS), + } } /// Create a new `CastExpr` with an explicit target `FieldRef`. /// - /// The provided `target_field` is used verbatim for the expression's - /// return schema, so the field's name, nullability, and other metadata are - /// preserved. This is the preferred constructor when the caller already - /// has field information (for example, during logical-to-physical planning). + /// The provided `target_field` determines the output characteristics: + /// - The field's data type becomes the cast target type + /// - The field's metadata is used exactly as provided + /// - The field's nullability is preserved + /// + /// This is the preferred constructor when the caller has explicit field + /// information that should be used exactly (for example, during schema + /// enforcement or adapter layers). /// - /// See [`CastExpr::new`] for the compatibility constructor that only accepts - /// a `DataType`. + /// See [`CastExpr::new`] for type-only casts where source metadata should + /// pass through. pub fn new_with_target_field( expr: Arc, target_field: FieldRef, @@ -126,6 +152,7 @@ impl CastExpr { Self { expr, target_field, + explicit_target: true, cast_options: cast_options.unwrap_or(DEFAULT_CAST_OPTIONS), } } @@ -140,7 +167,30 @@ impl CastExpr { self.target_field.data_type() } - /// Field metadata describing the output column after casting. + /// Explicit metadata for the output field, or `None` to pass through source metadata. + pub fn target_metadata(&self) -> Option<&HashMap> { + self.explicit_target.then(|| self.target_field.metadata()) + } + + /// Explicit nullability for the output field, or `None` to pass through source nullability. + pub fn target_nullable(&self) -> Option { + self.explicit_target + .then(|| self.target_field.is_nullable()) + } + + /// The target field this cast was constructed with. + /// + /// For a type-only cast this is a field synthesized from the target data + /// type alone; only its data type is meaningful. Note that the returned + /// field may not match what `return_field()` returns when evaluated against + /// a schema, since `return_field()` may incorporate source field information. + /// + /// Prefer [`cast_type()`], [`target_metadata()`], and [`target_nullable()`] + /// for direct access to the individual components. + /// + /// [`cast_type()`]: CastExpr::cast_type + /// [`target_metadata()`]: CastExpr::target_metadata + /// [`target_nullable()`]: CastExpr::target_nullable pub fn target_field(&self) -> &FieldRef { &self.target_field } @@ -150,19 +200,53 @@ impl CastExpr { &self.cast_options } + /// Whether this cast has explicit metadata (vs pass-through from source). + pub fn has_explicit_metadata(&self) -> bool { + self.explicit_target + } + + /// Whether this cast has explicit nullability (vs pass-through from source). + pub fn has_explicit_nullability(&self) -> bool { + self.explicit_target + } + fn resolved_target_field(&self, input_schema: &Schema) -> Result { - if is_default_target_field(&self.target_field) { - self.expr.return_field(input_schema).map(|field| { - Arc::new( - field - .as_ref() - .clone() - .with_data_type(self.cast_type().clone()), + // Try to get the source field for the name. If the target field is + // explicit, we can fall back to an empty name if the source lookup fails + // (e.g., for virtual row-index columns appended at scan time). + let source_result = self.expr.return_field(input_schema); + + if self.explicit_target { + // Metadata and nullability come from the target field verbatim + let name = source_result + .as_ref() + .map(|f| f.name().to_string()) + .unwrap_or_default(); + return Ok(Arc::new( + Field::new( + name, + self.cast_type().clone(), + self.target_field.is_nullable(), ) - }) - } else { - Ok(Arc::clone(&self.target_field)) + .with_metadata(self.target_field.metadata().clone()), + )); } + + // Type-only cast: pass through the source metadata and nullability, + // stripping extension type keys (the cast is to a plain storage type). + source_result.map(|source_field| { + let mut metadata = source_field.metadata().clone(); + metadata.remove(EXTENSION_TYPE_NAME_KEY); + metadata.remove(EXTENSION_TYPE_METADATA_KEY); + + Arc::new( + source_field + .as_ref() + .clone() + .with_data_type(self.cast_type().clone()) + .with_metadata(metadata), + ) + }) } /// Check if casting from the specified source type to the target type is a @@ -191,12 +275,6 @@ impl CastExpr { } } -fn is_default_target_field(target_field: &FieldRef) -> bool { - target_field.name().is_empty() - && target_field.is_nullable() - && target_field.metadata().is_empty() -} - pub(crate) fn is_order_preserving_cast_family( source_type: &DataType, target_type: &DataType, @@ -267,11 +345,12 @@ impl PhysicalExpr for CastExpr { self: Arc, children: Vec>, ) -> Result> { - Ok(Arc::new(CastExpr::new_with_target_field( - Arc::clone(&children[0]), - Arc::clone(&self.target_field), - Some(self.cast_options.clone()), - ))) + Ok(Arc::new(CastExpr { + expr: Arc::clone(&children[0]), + target_field: Arc::clone(&self.target_field), + explicit_target: self.explicit_target, + cast_options: self.cast_options.clone(), + })) } fn evaluate_bounds(&self, children: &[&Interval]) -> Result { @@ -372,32 +451,61 @@ pub fn cast_with_options( cast_type: DataType, cast_options: Option>, ) -> Result> { - cast_with_target_field( - expr, - input_schema, - cast_type.into_nullable_field_ref(), - cast_options, - ) + let expr_type = expr.data_type(input_schema)?; + + // If the types match, no cast is needed for a type-only cast + if expr_type == cast_type { + return Ok(Arc::clone(&expr)); + } + + let can_build_cast = if requires_nested_struct_cast(&expr_type, &cast_type) { + can_cast_named_struct_types(&expr_type, &cast_type) + } else { + can_cast_types(&expr_type, &cast_type) + }; + + if !can_build_cast { + return not_impl_err!("Unsupported CAST from {expr_type} to {cast_type}"); + } + + Ok(Arc::new(CastExpr::new(expr, cast_type, cast_options))) } /// Return a PhysicalExpression representing `expr` casted to `target_field`, /// preserving any explicit field semantics such as name, nullability, and /// metadata. /// -/// If the input expression already has the same data type, this helper still -/// preserves an explicit `target_field` by constructing a field-aware -/// [`CastExpr`]. Only the default synthesized field created by the legacy -/// type-only API is elided back to the original child expression. +/// If the input expression already has the same data type and the target field +/// has no explicit metadata or nullability constraints, the original expression +/// is returned unchanged. pub fn cast_with_target_field( expr: Arc, input_schema: &Schema, - target_field: FieldRef, + target_field: &FieldRef, cast_options: Option>, ) -> Result> { let expr_type = expr.data_type(input_schema)?; let cast_type = target_field.data_type(); - if expr_type == *cast_type && is_default_target_field(&target_field) { - return Ok(Arc::clone(&expr)); + + // Check if this is a "default" target field (type-only cast with no explicit + // metadata or nullability constraints). This is the field created by + // `into_nullable_field_ref()` when only a DataType is known. + let is_type_only = target_field.name().is_empty() + && target_field.is_nullable() + && target_field.metadata().is_empty(); + + // For same-type casts, we can skip creating a CastExpr only if: + // 1. The target is type-only (no explicit metadata) + // 2. The source has no extension metadata that needs to be stripped + // Otherwise we need the CastExpr to strip extension metadata from the source. + if expr_type == *cast_type && is_type_only { + let source_field = expr.return_field(input_schema)?; + let has_extension_metadata = source_field + .metadata() + .contains_key(EXTENSION_TYPE_NAME_KEY); + if !has_extension_metadata { + return Ok(Arc::clone(&expr)); + } } let can_build_cast = if requires_nested_struct_cast(&expr_type, cast_type) { @@ -415,11 +523,22 @@ pub fn cast_with_target_field( return not_impl_err!("Unsupported CAST from {expr_type} to {cast_type}"); } - Ok(Arc::new(CastExpr::new_with_target_field( - expr, - target_field, - cast_options, - ))) + // For type-only casts, use CastExpr::new which preserves source metadata/nullability. + // For explicit target fields, use new_with_target_field which applies the target's + // extension metadata and nullability. + if is_type_only { + Ok(Arc::new(CastExpr::new( + expr, + cast_type.clone(), + cast_options, + ))) + } else { + Ok(Arc::new(CastExpr::new_with_target_field( + expr, + Arc::clone(target_field), + cast_options, + ))) + } } /// Return a PhysicalExpression representing `expr` casted to @@ -990,26 +1109,32 @@ mod tests { #[test] fn field_aware_cast_preserves_target_field_semantics() -> Result<()> { + // Target field metadata should be preserved exactly (no merging with source). let metadata = HashMap::from([("target_meta".to_string(), "1".to_string())]); for (child_nullable, target_nullable) in [(true, false), (false, true)] { let schema = Schema::new(vec![Field::new("a", Int32, child_nullable)]); + let target_field = Arc::new( + Field::new("cast_target", Int64, target_nullable) + .with_metadata(metadata.clone()), + ); let expr = CastExpr::new_with_target_field( col("a", &schema)?, - Arc::new( - Field::new("cast_target", Int64, target_nullable) - .with_metadata(metadata.clone()), - ), + Arc::clone(&target_field), None, ); let field = expr.return_field(&schema)?; - assert_eq!(field.name(), "cast_target"); + // Field name comes from source + assert_eq!(field.name(), "a"); assert_eq!(field.data_type(), &Int64); + // Nullability comes from target assert_eq!(field.is_nullable(), target_nullable); + // Target metadata should be preserved exactly assert_eq!( - field.metadata().get("target_meta").map(String::as_str), - Some("1") + field.metadata().get("target_meta"), + Some(&"1".to_string()), + "Target metadata should be preserved exactly" ); assert_eq!(expr.nullable(&schema)?, child_nullable || target_nullable); } @@ -1017,6 +1142,38 @@ mod tests { Ok(()) } + #[test] + fn target_field_accessor_returns_the_constructed_field() -> Result<()> { + let schema = Schema::new(vec![Field::new("a", Int32, true)]); + let metadata = HashMap::from([("target_meta".to_string(), "1".to_string())]); + let target_field = + Arc::new(Field::new("cast_target", Int64, false).with_metadata(metadata)); + + let expr = CastExpr::new_with_target_field( + col("a", &schema)?, + Arc::clone(&target_field), + None, + ); + + // The field is returned verbatim, including its name. + assert_eq!(expr.target_field(), &target_field); + assert_eq!(expr.cast_type(), &Int64); + assert_eq!(expr.target_metadata(), Some(target_field.metadata())); + assert_eq!(expr.target_nullable(), Some(false)); + assert!(expr.has_explicit_metadata()); + assert!(expr.has_explicit_nullability()); + + // A type-only cast reports no explicit target. + let type_only = CastExpr::new(col("a", &schema)?, Int64, None); + assert_eq!(type_only.cast_type(), &Int64); + assert_eq!(type_only.target_metadata(), None); + assert_eq!(type_only.target_nullable(), None); + assert!(!type_only.has_explicit_metadata()); + assert!(!type_only.has_explicit_nullability()); + + Ok(()) + } + #[test] fn type_only_cast_preserves_legacy_field_name_and_nullability() -> Result<()> { let schema = Schema::new(vec![Field::new("a", Int32, false)]); @@ -1170,7 +1327,8 @@ mod tests { let literal = Arc::new(crate::expressions::Literal::new(ScalarValue::Struct( Arc::new(scalar_struct), ))); - let expr = CastExpr::new_with_target_field(literal, Arc::new(target_field), None); + let target_field = Arc::new(target_field); + let expr = CastExpr::new_with_target_field(literal, target_field, None); let batch = RecordBatch::new_empty(schema); let result = expr.evaluate(&batch)?; @@ -1216,6 +1374,93 @@ mod tests { Ok(()) } + #[test] + fn type_only_cast_strips_extension_metadata() -> Result<()> { + // When using type-only cast (new()), extension metadata from source should NOT propagate + let source_meta = HashMap::from([ + ( + EXTENSION_TYPE_NAME_KEY.to_string(), + "arrow.uuid".to_string(), + ), + ("custom_key".to_string(), "custom_value".to_string()), + ]); + let schema = Schema::new(vec![ + Field::new("a", FixedSizeBinary(16), false).with_metadata(source_meta), + ]); + + let expr = CastExpr::new(col("a", &schema)?, Utf8, None); + + let field = expr.return_field(&schema)?; + assert!( + field.metadata().get(EXTENSION_TYPE_NAME_KEY).is_none(), + "Type-only cast should strip extension type name from source" + ); + assert_eq!( + field.metadata().get("custom_key"), + Some(&"custom_value".to_string()), + "Type-only cast should preserve non-extension metadata" + ); + + Ok(()) + } + + #[test] + fn field_aware_cast_uses_exact_target_metadata() -> Result<()> { + // When using field-aware cast, target's metadata should be used exactly + let source_meta = HashMap::from([ + ( + EXTENSION_TYPE_NAME_KEY.to_string(), + "source.type".to_string(), + ), + ("source_key".to_string(), "source_value".to_string()), + ]); + let target_meta = HashMap::from([ + ( + EXTENSION_TYPE_NAME_KEY.to_string(), + "target.type".to_string(), + ), + ( + EXTENSION_TYPE_METADATA_KEY.to_string(), + "target_ext_meta".to_string(), + ), + ("target_key".to_string(), "target_value".to_string()), + ]); + let schema = Schema::new(vec![ + Field::new("a", FixedSizeBinary(16), false).with_metadata(source_meta), + ]); + + let target_field = + Arc::new(Field::new("b", Utf8, true).with_metadata(target_meta)); + let expr = CastExpr::new_with_target_field( + col("a", &schema)?, + Arc::clone(&target_field), + None, + ); + + let field = expr.return_field(&schema)?; + assert_eq!( + field.metadata().get(EXTENSION_TYPE_NAME_KEY), + Some(&"target.type".to_string()), + "Field-aware cast should use target's extension type name" + ); + assert_eq!( + field.metadata().get(EXTENSION_TYPE_METADATA_KEY), + Some(&"target_ext_meta".to_string()), + "Field-aware cast should use target's extension type metadata" + ); + assert!( + field.metadata().get("source_key").is_none(), + "Field-aware cast should NOT preserve source metadata" + ); + assert_eq!( + field.metadata().get("target_key"), + Some(&"target_value".to_string()), + "Field-aware cast should preserve target's non-extension metadata" + ); + + Ok(()) + } + #[test] fn test_check_bigger_cast_precision_loss() { use DataType::*; diff --git a/datafusion/physical-expr/src/expressions/mod.rs b/datafusion/physical-expr/src/expressions/mod.rs index 521b8b87e305c..f2f9285de560a 100644 --- a/datafusion/physical-expr/src/expressions/mod.rs +++ b/datafusion/physical-expr/src/expressions/mod.rs @@ -61,7 +61,7 @@ pub use no_op::NoOp; pub use not::{NotExpr, not}; pub(crate) use similar_to_pattern::translate_scalar; pub use similar_to_pattern::{SqlSimilarToPattern, sql_similar_to_regex}; -pub use try_cast::{TryCastExpr, try_cast}; +pub use try_cast::{TryCastExpr, try_cast, try_cast_with_target_field}; pub use unknown_column::UnKnownColumn; pub(crate) use cast::cast_with_target_field; diff --git a/datafusion/physical-expr/src/expressions/try_cast.rs b/datafusion/physical-expr/src/expressions/try_cast.rs index 65b953fd181b7..c054026724fb1 100644 --- a/datafusion/physical-expr/src/expressions/try_cast.rs +++ b/datafusion/physical-expr/src/expressions/try_cast.rs @@ -15,6 +15,7 @@ // specific language governing permissions and limitations // under the License. +use std::collections::HashMap; use std::fmt; use std::hash::Hash; use std::sync::Arc; @@ -22,40 +23,94 @@ use std::sync::Arc; use crate::PhysicalExpr; use arrow::compute; use arrow::compute::CastOptions; -use arrow::datatypes::{DataType, FieldRef, Schema}; +use arrow::datatypes::{DataType, Field, FieldRef, Schema}; use arrow::record_batch::RecordBatch; +use arrow_schema::extension::{EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_NAME_KEY}; use compute::can_cast_types; +use datafusion_common::datatype::DataTypeExt; use datafusion_common::format::DEFAULT_FORMAT_OPTIONS; use datafusion_common::{Result, not_impl_err}; use datafusion_expr::ColumnarValue; /// TRY_CAST expression casts an expression to a specific data type and returns NULL on invalid cast -#[derive(Debug, Eq)] +#[derive(Debug, Clone, Eq)] pub struct TryCastExpr { /// The expression to cast expr: Arc, - /// The data type to cast to - cast_type: DataType, + /// The target field. + /// + /// For a type-only cast (see [`TryCastExpr::new`]) this is a field + /// synthesized from the target data type alone and only its data type is + /// meaningful. For a cast built from an explicit field (see + /// [`TryCastExpr::new_with_target_field`]) its metadata is applied to the + /// output field as-is. + target_field: FieldRef, + /// Whether `target_field` was supplied by the caller (as opposed to being + /// synthesized from a `DataType`), and therefore whether its metadata + /// describes the output field exactly. + explicit_target: bool, } // Manually derive PartialEq and Hash to work around https://github.com/rust-lang/rust/issues/78808 impl PartialEq for TryCastExpr { fn eq(&self, other: &Self) -> bool { - self.expr.eq(&other.expr) && self.cast_type == other.cast_type + // Compare the semantically meaningful parts of the target field only: + // the field name never affects the output of this expression. + self.expr.eq(&other.expr) + && self.cast_type() == other.cast_type() + && self.target_metadata() == other.target_metadata() } } impl Hash for TryCastExpr { fn hash(&self, state: &mut H) { self.expr.hash(state); - self.cast_type.hash(state); + self.cast_type().hash(state); + // Hash the metadata by iterating over sorted keys for deterministic ordering + if let Some(metadata) = self.target_metadata() { + let mut entries: Vec<_> = metadata.iter().collect(); + entries.sort_by_key(|(k, _)| *k); + for (k, v) in entries { + k.hash(state); + v.hash(state); + } + } } } impl TryCastExpr { - /// Create a new CastExpr + /// Create a new `TryCastExpr` using only a `DataType`. + /// + /// This constructor creates a type-only cast where metadata is passed through + /// from the source expression (with extension type keys stripped). + /// TRY_CAST results are always nullable since failed casts return NULL. pub fn new(expr: Arc, cast_type: DataType) -> Self { - Self { expr, cast_type } + Self { + expr, + target_field: cast_type.into_nullable_field_ref(), + explicit_target: false, + } + } + + /// Create a new `TryCastExpr` with an explicit target `FieldRef`. + /// + /// The provided `target_field` determines the output characteristics: + /// - The field's data type becomes the cast target type + /// - The field's metadata is used exactly as provided + /// + /// TRY_CAST results are always nullable since failed casts return NULL. + /// + /// See [`TryCastExpr::new`] for type-only casts where source metadata should + /// pass through. + pub fn new_with_target_field( + expr: Arc, + target_field: FieldRef, + ) -> Self { + Self { + expr, + target_field, + explicit_target: true, + } } /// The expression to cast @@ -65,19 +120,33 @@ impl TryCastExpr { /// The data type to cast to pub fn cast_type(&self) -> &DataType { - &self.cast_type + self.target_field.data_type() + } + + /// Explicit metadata for the output field, or `None` to pass through source metadata. + pub fn target_metadata(&self) -> Option<&HashMap> { + self.explicit_target.then(|| self.target_field.metadata()) + } + + /// The target field this cast was constructed with. + /// + /// For a type-only cast this is a field synthesized from the target data + /// type alone; only its data type is meaningful. TRY_CAST results are + /// always nullable regardless of the target field's nullability. + pub fn target_field(&self) -> &FieldRef { + &self.target_field } } impl fmt::Display for TryCastExpr { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { - write!(f, "TRY_CAST({} AS {})", self.expr, self.cast_type) + write!(f, "TRY_CAST({} AS {})", self.expr, self.cast_type()) } } impl PhysicalExpr for TryCastExpr { fn data_type(&self, _input_schema: &Schema) -> Result { - Ok(self.cast_type.clone()) + Ok(self.cast_type().clone()) } fn nullable(&self, _input_schema: &Schema) -> Result { @@ -90,14 +159,41 @@ impl PhysicalExpr for TryCastExpr { safe: true, format_options: DEFAULT_FORMAT_OPTIONS, }; - value.cast_to(&self.cast_type, Some(&options)) + value.cast_to(self.cast_type(), Some(&options)) } fn return_field(&self, input_schema: &Schema) -> Result { - self.expr - .return_field(input_schema) - .map(|f| f.as_ref().clone().with_data_type(self.cast_type.clone())) - .map(Arc::new) + // If metadata is explicit, we can build the field without source + // (though we still try to get source for the name) + let source_result = self.expr.return_field(input_schema); + + if let Some(metadata) = self.target_metadata() { + // Explicit metadata: use it exactly, TRY_CAST is always nullable + let name = source_result + .as_ref() + .map(|f| f.name().to_string()) + .unwrap_or_default(); + return Ok(Arc::new( + Field::new(name, self.cast_type().clone(), true) + .with_metadata(metadata.clone()), + )); + } + + // Pass-through metadata from source (stripping extension keys) + source_result.map(|source_field| { + let mut metadata = source_field.metadata().clone(); + metadata.remove(EXTENSION_TYPE_NAME_KEY); + metadata.remove(EXTENSION_TYPE_METADATA_KEY); + + Arc::new( + source_field + .as_ref() + .clone() + .with_data_type(self.cast_type().clone()) + .with_nullable(true) // TRY_CAST is always nullable + .with_metadata(metadata), + ) + }) } fn children(&self) -> Vec<&Arc> { @@ -108,16 +204,17 @@ impl PhysicalExpr for TryCastExpr { self: Arc, children: Vec>, ) -> Result> { - Ok(Arc::new(TryCastExpr::new( - Arc::clone(&children[0]), - self.cast_type.clone(), - ))) + Ok(Arc::new(TryCastExpr { + expr: Arc::clone(&children[0]), + target_field: Arc::clone(&self.target_field), + explicit_target: self.explicit_target, + })) } fn fmt_sql(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { write!(f, "TRY_CAST(")?; self.expr.fmt_sql(f)?; - write!(f, " AS {:?})", self.cast_type) + write!(f, " AS {:?})", self.cast_type()) } #[cfg(feature = "proto")] @@ -190,6 +287,60 @@ pub fn try_cast( } } +/// Return a PhysicalExpression representing `expr` casted to `target_field`, +/// preserving any explicit field semantics such as metadata. +/// +/// TRY_CAST results are always nullable since failed casts return NULL. +/// +/// If the input expression already has the same data type, the target field +/// has no explicit metadata constraints, and the source has no extension +/// metadata to strip, the original expression is returned unchanged. +pub fn try_cast_with_target_field( + expr: Arc, + input_schema: &Schema, + target_field: &FieldRef, +) -> Result> { + let expr_type = expr.data_type(input_schema)?; + let cast_type = target_field.data_type(); + + // Check if this is a "default" target field (type-only cast with no explicit + // metadata constraints). This is the field created by `into_nullable_field_ref()` + // when only a DataType is known. + let is_type_only = target_field.name().is_empty() + && target_field.is_nullable() + && target_field.metadata().is_empty(); + + // For same-type casts, we can skip creating a TryCastExpr only if: + // 1. The target is type-only (no explicit metadata) + // 2. The source has no extension metadata that needs to be stripped + // Otherwise we need the TryCastExpr to strip extension metadata from the source. + if expr_type == *cast_type && is_type_only { + let source_field = expr.return_field(input_schema)?; + let has_extension_metadata = source_field + .metadata() + .contains_key(EXTENSION_TYPE_NAME_KEY); + if !has_extension_metadata { + return Ok(Arc::clone(&expr)); + } + } + + if !can_cast_types(&expr_type, cast_type) { + return not_impl_err!("Unsupported TRY_CAST from {expr_type} to {cast_type}"); + } + + // For type-only casts, use TryCastExpr::new which preserves source metadata. + // For explicit target fields, use new_with_target_field which applies the target's + // metadata exactly. + if is_type_only { + Ok(Arc::new(TryCastExpr::new(expr, cast_type.clone()))) + } else { + Ok(Arc::new(TryCastExpr::new_with_target_field( + expr, + Arc::clone(target_field), + ))) + } +} + #[cfg(test)] mod tests { use super::*; @@ -642,6 +793,167 @@ mod tests { Ok(()) } + + #[test] + fn field_aware_try_cast_uses_exact_target_metadata() -> Result<()> { + // When using field-aware cast, target's metadata should be used exactly + let source_meta = HashMap::from([ + ( + EXTENSION_TYPE_NAME_KEY.to_string(), + "source.type".to_string(), + ), + ("source_key".to_string(), "source_value".to_string()), + ]); + let target_meta = HashMap::from([ + ( + EXTENSION_TYPE_NAME_KEY.to_string(), + "target.type".to_string(), + ), + ( + EXTENSION_TYPE_METADATA_KEY.to_string(), + "target_ext_meta".to_string(), + ), + ("target_key".to_string(), "target_value".to_string()), + ]); + let schema = Schema::new(vec![ + Field::new("a", DataType::FixedSizeBinary(16), false) + .with_metadata(source_meta), + ]); + + let target_field = + Arc::new(Field::new("b", DataType::Utf8, true).with_metadata(target_meta)); + let expr = TryCastExpr::new_with_target_field( + col("a", &schema)?, + Arc::clone(&target_field), + ); + + let field = expr.return_field(&schema)?; + assert_eq!( + field.metadata().get(EXTENSION_TYPE_NAME_KEY), + Some(&"target.type".to_string()), + "Field-aware try_cast should use target's extension type name" + ); + assert_eq!( + field.metadata().get(EXTENSION_TYPE_METADATA_KEY), + Some(&"target_ext_meta".to_string()), + "Field-aware try_cast should use target's extension type metadata" + ); + assert!( + field.metadata().get("source_key").is_none(), + "Field-aware try_cast should NOT preserve source metadata" + ); + assert_eq!( + field.metadata().get("target_key"), + Some(&"target_value".to_string()), + "Field-aware try_cast should preserve target's non-extension metadata" + ); + // TRY_CAST is always nullable + assert!(field.is_nullable()); + + Ok(()) + } + + #[test] + fn field_aware_try_cast_preserves_target_field_semantics() -> Result<()> { + // Target field metadata should be preserved exactly (no merging with source). + // TRY_CAST is always nullable regardless of target field's nullability. + let metadata = HashMap::from([("target_meta".to_string(), "1".to_string())]); + + for child_nullable in [true, false] { + let schema = + Schema::new(vec![Field::new("a", DataType::Int32, child_nullable)]); + let target_field = Arc::new( + Field::new("cast_target", DataType::Int64, false) // target says non-nullable + .with_metadata(metadata.clone()), + ); + let expr = TryCastExpr::new_with_target_field( + col("a", &schema)?, + Arc::clone(&target_field), + ); + + let field = expr.return_field(&schema)?; + // Field name comes from source + assert_eq!(field.name(), "a"); + assert_eq!(field.data_type(), &DataType::Int64); + // TRY_CAST is ALWAYS nullable (ignores target field's nullability) + assert!(field.is_nullable(), "TRY_CAST should always be nullable"); + // Target metadata should be preserved exactly + assert_eq!( + field.metadata().get("target_meta"), + Some(&"1".to_string()), + "Target metadata should be preserved exactly" + ); + assert!( + expr.nullable(&schema)?, + "TRY_CAST should always be nullable" + ); + } + + Ok(()) + } + + #[test] + fn type_only_try_cast_strips_extension_keys() -> Result<()> { + // Type-only cast should strip extension keys but preserve other source metadata + let source_meta = HashMap::from([ + ( + EXTENSION_TYPE_NAME_KEY.to_string(), + "source.extension".to_string(), + ), + ( + EXTENSION_TYPE_METADATA_KEY.to_string(), + "ext_meta".to_string(), + ), + ("custom_key".to_string(), "custom_value".to_string()), + ]); + let schema = Schema::new(vec![ + Field::new("a", DataType::Int32, false).with_metadata(source_meta), + ]); + + let expr = TryCastExpr::new(col("a", &schema)?, DataType::Int64); + let field = expr.return_field(&schema)?; + + // Extension keys should be stripped + assert!( + field.metadata().get(EXTENSION_TYPE_NAME_KEY).is_none(), + "Type-only try_cast should strip extension type name" + ); + assert!( + field.metadata().get(EXTENSION_TYPE_METADATA_KEY).is_none(), + "Type-only try_cast should strip extension type metadata" + ); + // Non-extension metadata should pass through + assert_eq!( + field.metadata().get("custom_key"), + Some(&"custom_value".to_string()), + "Type-only try_cast should preserve non-extension metadata" + ); + // Field name preserved, type changed, always nullable + assert_eq!(field.name(), "a"); + assert_eq!(field.data_type(), &DataType::Int64); + assert!(field.is_nullable()); + + Ok(()) + } + + #[test] + fn type_only_try_cast_is_always_nullable() -> Result<()> { + // TRY_CAST is always nullable even when source is non-nullable + let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); + let expr = TryCastExpr::new(col("a", &schema)?, DataType::Int64); + + let field = expr.return_field(&schema)?; + + assert_eq!(field.name(), "a"); + assert_eq!(field.data_type(), &DataType::Int64); + assert!(field.is_nullable(), "TRY_CAST should always be nullable"); + assert!( + expr.nullable(&schema)?, + "TRY_CAST should always be nullable" + ); + + Ok(()) + } } #[cfg(all(test, feature = "proto"))] diff --git a/datafusion/physical-expr/src/planner.rs b/datafusion/physical-expr/src/planner.rs index 9a0bdc33da8e9..8c5347db742bb 100644 --- a/datafusion/physical-expr/src/planner.rs +++ b/datafusion/physical-expr/src/planner.rs @@ -27,7 +27,7 @@ use crate::{ use arrow::datatypes::Schema; use datafusion_common::config::ConfigOptions; use datafusion_common::datatype::FieldExt; -use datafusion_common::metadata::{FieldMetadata, format_type_and_metadata}; +use datafusion_common::metadata::FieldMetadata; use datafusion_common::{ DFSchema, Result, ScalarValue, TableReference, ToDFSchema, exec_err, internal_datafusion_err, not_impl_err, plan_datafusion_err, plan_err, @@ -387,23 +387,11 @@ pub fn create_physical_expr( Expr::Cast(Cast { expr, field }) => expressions::cast_with_target_field( create_physical_expr(expr, input_dfschema, execution_props, planning_ctx)?, input_schema, - Arc::clone(field), + field, None, ), Expr::TryCast(TryCast { expr, field }) => { - if !field.metadata().is_empty() { - let (_, src_field) = expr.to_field(input_dfschema)?; - return plan_err!( - "TryCast from {} to {} is not supported", - format_type_and_metadata( - src_field.data_type(), - Some(src_field.metadata()), - ), - format_type_and_metadata(field.data_type(), Some(field.metadata())) - ); - } - - expressions::try_cast( + expressions::try_cast_with_target_field( create_physical_expr( expr, input_dfschema, @@ -411,7 +399,7 @@ pub fn create_physical_expr( planning_ctx, )?, input_schema, - field.data_type().clone(), + field, ) } Expr::Not(expr) => expressions::not(create_physical_expr( @@ -743,6 +731,7 @@ pub fn logical2physical(expr: &Expr, schema: &Schema) -> Arc { mod tests { use arrow::array::{ArrayRef, BooleanArray, RecordBatch, StringArray}; use arrow::datatypes::{DataType, Field}; + use arrow_schema::extension::{EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_NAME_KEY}; use datafusion_common::HashMap; use datafusion_expr::physical_planning_context::{ ScalarSubqueryResults, SubqueryIndex, @@ -771,6 +760,14 @@ mod tests { .expect("planner should lower logical CAST to CastExpr") } + fn as_planner_try_cast( + physical: &Arc, + ) -> &expressions::TryCastExpr { + physical + .downcast_ref::() + .expect("planner should lower logical TRY_CAST to TryCastExpr") + } + #[test] fn test_create_physical_expr_scalar_input_output() -> Result<()> { let expr = col("letter").eq(lit("A")); @@ -833,9 +830,21 @@ mod tests { #[test] fn test_cast_lowering_preserves_target_field_metadata() -> Result<()> { let schema = test_cast_schema(); + + // Target field with both extension metadata and custom metadata. + // With exact target metadata semantics, all target metadata should propagate. let target_field = Arc::new( - Field::new("cast_target", DataType::Int64, true) - .with_metadata([("target_meta".to_string(), "1".to_string())].into()), + Field::new("cast_target", DataType::Int64, true).with_metadata( + [ + ( + EXTENSION_TYPE_NAME_KEY.to_string(), + "arrow.json".to_string(), + ), + (EXTENSION_TYPE_METADATA_KEY.to_string(), "{}".to_string()), + ("custom_target_meta".to_string(), "custom_value".to_string()), + ] + .into(), + ), ); let cast_expr = Expr::Cast(Cast::new_from_field( Box::new(col("a")), @@ -845,8 +854,35 @@ mod tests { let physical = lower_cast_expr(&cast_expr, &schema)?; let cast = as_planner_cast(&physical); - assert_eq!(cast.target_field(), &target_field); - assert_eq!(physical.return_field(&schema)?, target_field); + // The CastExpr stores the target type and all target metadata + assert_eq!(cast.cast_type(), &DataType::Int64); + let target_metadata = cast.target_metadata().expect("should have metadata"); + assert_eq!( + target_metadata.get(EXTENSION_TYPE_NAME_KEY), + Some(&"arrow.json".to_string()) + ); + assert_eq!( + target_metadata.get(EXTENSION_TYPE_METADATA_KEY), + Some(&"{}".to_string()) + ); + assert_eq!(cast.target_nullable(), Some(true)); + + // return_field should have all target metadata (exact semantics) + let returned = physical.return_field(&schema)?; + assert_eq!( + returned.metadata().get(EXTENSION_TYPE_NAME_KEY), + Some(&"arrow.json".to_string()) + ); + assert_eq!( + returned.metadata().get(EXTENSION_TYPE_METADATA_KEY), + Some(&"{}".to_string()) + ); + // All target metadata should propagate with exact semantics + assert_eq!( + returned.metadata().get("custom_target_meta"), + Some(&"custom_value".to_string()), + "All target metadata should propagate with exact semantics" + ); assert!(physical.nullable(&schema)?); Ok(()) @@ -872,22 +908,85 @@ mod tests { #[test] fn test_cast_lowering_preserves_same_type_field_semantics() -> Result<()> { let schema = test_cast_schema(); + + // Same-type cast with extension metadata on target. + // With exact target metadata semantics, all target metadata should propagate. let target_field = Arc::new( Field::new("same_type_cast", DataType::Int32, true).with_metadata( - [("target_meta".to_string(), "same-type".to_string())].into(), + [ + ( + EXTENSION_TYPE_NAME_KEY.to_string(), + "arrow.opaque".to_string(), + ), + ("custom_meta".to_string(), "custom_value".to_string()), + ] + .into(), ), ); - let cast_expr = Expr::Cast(Cast::new_from_field( - Box::new(col("a")), - Arc::clone(&target_field), - )); - let physical = lower_cast_expr(&cast_expr, &schema)?; - let cast = as_planner_cast(&physical); + for use_try_cast in [false, true] { + // For error labelling + let cast_name = if use_try_cast { "TRY_CAST" } else { "CAST" }; - assert_eq!(cast.target_field(), &target_field); - assert_eq!(physical.return_field(&schema)?, target_field); - assert!(physical.nullable(&schema)?); + let cast_expr = if use_try_cast { + Expr::TryCast(TryCast::new_from_field( + Box::new(col("a")), + Arc::clone(&target_field), + )) + } else { + Expr::Cast(Cast::new_from_field( + Box::new(col("a")), + Arc::clone(&target_field), + )) + }; + + let physical = lower_cast_expr(&cast_expr, &schema)?; + + // Extract common fields - both CastExpr and TryCastExpr have these + let (cast_type, target_metadata, target_nullable) = if use_try_cast { + let cast = as_planner_try_cast(&physical); + (cast.cast_type(), cast.target_metadata(), None) + } else { + let cast = as_planner_cast(&physical); + ( + cast.cast_type(), + cast.target_metadata(), + cast.target_nullable(), + ) + }; + + // Verify the physical expression stores correct metadata (same for both) + assert_eq!(cast_type, &DataType::Int32, "{cast_name}: cast_type"); + let target_metadata = target_metadata.expect("should have metadata"); + assert_eq!( + target_metadata.get(EXTENSION_TYPE_NAME_KEY), + Some(&"arrow.opaque".to_string()), + "{cast_name}: extension type name" + ); + + // Only CastExpr tracks target_nullable (TryCast is always nullable) + if !use_try_cast { + assert_eq!(target_nullable, Some(true), "{cast_name}: target_nullable"); + } + + // return_field should have all target metadata (exact semantics) + let returned = physical.return_field(&schema)?; + assert_eq!( + returned.metadata().get(EXTENSION_TYPE_NAME_KEY), + Some(&"arrow.opaque".to_string()), + "{cast_name}: return_field extension type name" + ); + // All target metadata should propagate with exact semantics + assert_eq!( + returned.metadata().get("custom_meta"), + Some(&"custom_value".to_string()), + "{cast_name}: All target metadata should propagate with exact semantics" + ); + assert!( + physical.nullable(&schema)?, + "{cast_name}: should be nullable" + ); + } Ok(()) } diff --git a/datafusion/proto/src/logical_plan/from_proto.rs b/datafusion/proto/src/logical_plan/from_proto.rs index a1a1ff6f04fe4..1eac6f974beed 100644 --- a/datafusion/proto/src/logical_plan/from_proto.rs +++ b/datafusion/proto/src/logical_plan/from_proto.rs @@ -429,7 +429,8 @@ pub fn parse_expr( let data_type: DataType = cast.arrow_type.as_ref().required("arrow_type")?; let field = data_type .into_nullable_field() - .with_nullable(cast.nullable.unwrap_or(true)); + .with_nullable(cast.nullable.unwrap_or(true)) + .with_metadata(cast.metadata.clone()); Ok(Expr::Cast(Cast::new_from_field(expr, Arc::new(field)))) } ExprType::TryCast(cast) => { @@ -442,7 +443,8 @@ pub fn parse_expr( let data_type: DataType = cast.arrow_type.as_ref().required("arrow_type")?; let field = data_type .into_nullable_field() - .with_nullable(cast.nullable.unwrap_or(true)); + .with_nullable(cast.nullable.unwrap_or(true)) + .with_metadata(cast.metadata.clone()); Ok(Expr::TryCast(TryCast::new_from_field( expr, Arc::new(field), diff --git a/datafusion/proto/tests/cases/roundtrip_logical_plan.rs b/datafusion/proto/tests/cases/roundtrip_logical_plan.rs index 7c916db9c3cd1..750b20323ad2e 100644 --- a/datafusion/proto/tests/cases/roundtrip_logical_plan.rs +++ b/datafusion/proto/tests/cases/roundtrip_logical_plan.rs @@ -2693,6 +2693,19 @@ fn roundtrip_cast() { let ctx = SessionContext::new(); roundtrip_expr_test(test_expr, ctx); + + let field = + Field::new("", DataType::Boolean, false).with_metadata(HashMap::from([( + String::from("key"), + String::from("value"), + )])); + let test_expr = Expr::Cast(Cast::new_from_field( + Box::new(lit(1.0_f32)), + Arc::new(field), + )); + + let ctx = SessionContext::new(); + roundtrip_expr_test(test_expr, ctx); } #[test] @@ -2703,6 +2716,19 @@ fn roundtrip_try_cast() { let ctx = SessionContext::new(); roundtrip_expr_test(test_expr, ctx); + let field = + Field::new("", DataType::Boolean, false).with_metadata(HashMap::from([( + String::from("key"), + String::from("value"), + )])); + let test_expr = Expr::TryCast(TryCast::new_from_field( + Box::new(lit(1.0_f32)), + Arc::new(field), + )); + + let ctx = SessionContext::new(); + roundtrip_expr_test(test_expr, ctx); + let test_expr = Expr::TryCast(TryCast::new(Box::new(lit("not a bool")), DataType::Boolean)); diff --git a/datafusion/sqllogictest/test_files/cast_extension_type_metadata.slt b/datafusion/sqllogictest/test_files/cast_extension_type_metadata.slt index 425d8ac16eaee..01a19454e9a80 100644 --- a/datafusion/sqllogictest/test_files/cast_extension_type_metadata.slt +++ b/datafusion/sqllogictest/test_files/cast_extension_type_metadata.slt @@ -45,5 +45,55 @@ FROM ( ---- 00010203040506070809000102030506 arrow.uuid -statement error DataFusion error: Optimizer rule 'simplify_expressions' failed[\s\S]*TryCast from FixedSizeBinary\(16\) to FixedSizeBinary\(16\)<\{"ARROW:extension:name": "arrow\.uuid"\}> is not supported -SELECT TRY_CAST(arrow_cast(X'00010203040506070809000102030506', 'FixedSizeBinary(16)') AS UUID); +# TRY_CAST to extension type should also preserve extension metadata +query ?T +SELECT + TRY_CAST( + arrow_cast(X'00010203040506070809000102030506', 'FixedSizeBinary(16)') + AS UUID + ), + arrow_metadata( + TRY_CAST( + arrow_cast(X'00010203040506070809000102030506', 'FixedSizeBinary(16)') + AS UUID + ), + 'ARROW:extension:name' + ); +---- +00010203040506070809000102030506 arrow.uuid + +# TRY_CAST to UUID from a subquery +query ?T +SELECT + TRY_CAST(raw AS UUID), + arrow_metadata(TRY_CAST(raw AS UUID), 'ARROW:extension:name') +FROM ( + VALUES ( + arrow_cast(X'00010203040506070809000102030506', 'FixedSizeBinary(16)') + ) +) AS uuids(raw); +---- +00010203040506070809000102030506 arrow.uuid + +# arrow_cast from UUID to same underlying type (FixedSizeBinary(16)) strips +# extension metadata (type-only cast semantics) +query ?T +SELECT + arrow_cast(uuid_val, 'FixedSizeBinary(16)'), + arrow_metadata(arrow_cast(uuid_val, 'FixedSizeBinary(16)'), 'ARROW:extension:name') +FROM ( + SELECT CAST(arrow_cast(X'00010203040506070809000102030506', 'FixedSizeBinary(16)') AS UUID) AS uuid_val +); +---- +00010203040506070809000102030506 NULL + +# arrow_cast to a different type strips extension metadata (type-only cast semantics) +query ?T +SELECT + arrow_cast(uuid_val, 'Binary'), + arrow_metadata(arrow_cast(uuid_val, 'Binary'), 'ARROW:extension:name') +FROM ( + SELECT CAST(arrow_cast(X'00010203040506070809000102030506', 'FixedSizeBinary(16)') AS UUID) AS uuid_val +); +---- +00010203040506070809000102030506 NULL From f1e186dc3b3dc526441266ab0045c8547377e4c3 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Tue, 1 Sep 2026 19:56:14 +0000 Subject: [PATCH 03/37] chore(deps-dev): bump browserslist from 4.28.1 to 4.28.8 in /datafusion/wasmtest/datafusion-wasm-app (#24862) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Bumps [browserslist](https://github.com/browserslist/browserslist) from 4.28.1 to 4.28.8.
Release notes

Sourced from browserslist's releases.

4.28.8

  • Fixed including kaios in baseline queries (by @​Jaybhade).

4.28.7

4.28.6

4.28.5

4.28.4

  • Fixed SyntaxError regression of 4.28.3.

4.28.3

  • Fixed baseline query case-insensitivity (by @​swwind).

4.28.2

Changelog

Sourced from browserslist's changelog.

4.28.8

  • Fixed including kaios in baseline queries (by @​Jaybhade).

4.28.7

4.28.6

4.28.5

4.28.4

  • Fixed SyntaxError regression of 4.28.3.

4.28.3

  • Fixed baseline query case-insensitivity (by @​swwind).

4.28.2

Commits
  • f2f2e6c Release 4.28.8 version
  • d0787c8 Update dependencies
  • fcf8fa9 Merge pull request #939 from Jaybhade/fix/baseline-kaios-without-downstream
  • 57ecd64 fix: support "including kaios" without downstream
  • 093a0f6 Update EM banner
  • b637868 Release 4.28.7 version
  • 313f465 Update dependencies
  • c935c5a Fix regexp performance
  • d7e9e65 Rewrite structure parsing to make it always fast
  • ec4a55e Fix import order
  • Additional commits viewable in compare view
Maintainer changes

This version was pushed to npm by GitHub Actions, a new releaser for browserslist since your current version.


[![Dependabot compatibility score](https://dependabot-badges.githubapp.com/badges/compatibility_score?dependency-name=browserslist&package-manager=npm_and_yarn&previous-version=4.28.1&new-version=4.28.8)](https://docs.github.com/en/github/managing-security-vulnerabilities/about-dependabot-security-updates#about-compatibility-scores) Dependabot will resolve any conflicts with this PR as long as you don't alter it yourself. You can also trigger a rebase manually by commenting `@dependabot rebase`. [//]: # (dependabot-automerge-start) [//]: # (dependabot-automerge-end) ---
Dependabot commands and options
You can trigger Dependabot actions by commenting on this PR: - `@dependabot rebase` will rebase this PR - `@dependabot recreate` will recreate this PR, overwriting any edits that have been made to it - `@dependabot show ignore conditions` will show all of the ignore conditions of the specified dependency - `@dependabot ignore this major version` will close this PR and stop Dependabot creating any more for this major version (unless you reopen the PR or upgrade to it yourself) - `@dependabot ignore this minor version` will close this PR and stop Dependabot creating any more for this minor version (unless you reopen the PR or upgrade to it yourself) - `@dependabot ignore this dependency` will close this PR and stop Dependabot creating any more for this dependency (unless you reopen the PR or upgrade to it yourself) You can disable automated security fix PRs for this repo from the [Security Alerts page](https://github.com/apache/datafusion/network/alerts).
Signed-off-by: dependabot[bot] Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> --- .../datafusion-wasm-app/package-lock.json | 113 ++++++++++-------- 1 file changed, 63 insertions(+), 50 deletions(-) diff --git a/datafusion/wasmtest/datafusion-wasm-app/package-lock.json b/datafusion/wasmtest/datafusion-wasm-app/package-lock.json index a43843bd0fe9f..a85846fd8bb39 100644 --- a/datafusion/wasmtest/datafusion-wasm-app/package-lock.json +++ b/datafusion/wasmtest/datafusion-wasm-app/package-lock.json @@ -755,12 +755,16 @@ "dev": true }, "node_modules/baseline-browser-mapping": { - "version": "2.9.19", - "resolved": "https://registry.npmjs.org/baseline-browser-mapping/-/baseline-browser-mapping-2.9.19.tgz", - "integrity": "sha512-ipDqC8FrAl/76p2SSWKSI+H9tFwm7vYqXQrItCuiVPt26Km0jS+NzSsBWAaBusvSbQcfJG+JitdMm+wZAgTYqg==", + "version": "2.11.20", + "resolved": "https://registry.npmjs.org/baseline-browser-mapping/-/baseline-browser-mapping-2.11.20.tgz", + "integrity": "sha512-H0ulySigv6icDJ1F7SjtdCD6PrhTpdYCmP0CactWy1+ekh0AFd0o1Wn5T8b+hnTmdBx19u9yhL6wvCylXMY7zw==", "dev": true, + "license": "Apache-2.0", "bin": { - "baseline-browser-mapping": "dist/cli.js" + "baseline-browser-mapping": "dist/cli.cjs" + }, + "engines": { + "node": ">=6.0.0" } }, "node_modules/batch": { @@ -830,9 +834,9 @@ } }, "node_modules/browserslist": { - "version": "4.28.1", - "resolved": "https://registry.npmjs.org/browserslist/-/browserslist-4.28.1.tgz", - "integrity": "sha512-ZC5Bd0LgJXgwGqUknZY/vkUQ04r8NXnJZ3yYi4vDmSiZmC/pdSN0NbNRPxZpbtO4uAfDUAFffO8IZoM3Gj8IkA==", + "version": "4.28.8", + "resolved": "https://registry.npmjs.org/browserslist/-/browserslist-4.28.8.tgz", + "integrity": "sha512-V2NpofLblG64mfOtSgDhOJESZEGogzDMBv/q+W6oc4LXWP/q75eOXoOaaOu1EOadB9U4Bwx/e0yzbvwKH8zalA==", "dev": true, "funding": [ { @@ -848,12 +852,13 @@ "url": "https://github.com/sponsors/ai" } ], + "license": "MIT", "dependencies": { - "baseline-browser-mapping": "^2.9.0", - "caniuse-lite": "^1.0.30001759", - "electron-to-chromium": "^1.5.263", - "node-releases": "^2.0.27", - "update-browserslist-db": "^1.2.0" + "baseline-browser-mapping": "^2.11.12", + "caniuse-lite": "^1.0.30001809", + "electron-to-chromium": "^1.5.402", + "node-releases": "^2.0.53", + "update-browserslist-db": "^1.3.0" }, "bin": { "browserslist": "cli.js" @@ -931,9 +936,9 @@ } }, "node_modules/caniuse-lite": { - "version": "1.0.30001768", - "resolved": "https://registry.npmjs.org/caniuse-lite/-/caniuse-lite-1.0.30001768.tgz", - "integrity": "sha512-qY3aDRZC5nWPgHUgIB84WL+nySuo19wk0VJpp/XI9T34lrvkyhRvNVOFJOp2kxClQhiFBu+TaUSudf6oa3vkSA==", + "version": "1.0.30001810", + "resolved": "https://registry.npmjs.org/caniuse-lite/-/caniuse-lite-1.0.30001810.tgz", + "integrity": "sha512-TITQPUkaz+aVk5GL6NhOdwk1aEaNTSDPsGFWrTuhKGtjTF70jL/Oht2W4c6rXUe5fu7Ie19VIahAXHIIiWWNeg==", "dev": true, "funding": [ { @@ -948,7 +953,8 @@ "type": "github", "url": "https://github.com/sponsors/ai" } - ] + ], + "license": "CC-BY-4.0" }, "node_modules/chokidar": { "version": "5.0.0", @@ -1283,10 +1289,11 @@ "dev": true }, "node_modules/electron-to-chromium": { - "version": "1.5.286", - "resolved": "https://registry.npmjs.org/electron-to-chromium/-/electron-to-chromium-1.5.286.tgz", - "integrity": "sha512-9tfDXhJ4RKFNerfjdCcZfufu49vg620741MNs26a9+bhLThdB+plgMeou98CAaHu/WATj2iHOOHTp1hWtABj2A==", - "dev": true + "version": "1.5.419", + "resolved": "https://registry.npmjs.org/electron-to-chromium/-/electron-to-chromium-1.5.419.tgz", + "integrity": "sha512-nHMPn8x4yCxCI0iSnL+LlHL5sUoUfjLXkcRIagZ4GBdrfFLFaiLNvzJWbJqZhFT9IAhw5tUSNlhggWN+otvp/A==", + "dev": true, + "license": "ISC" }, "node_modules/encodeurl": { "version": "2.0.0", @@ -1363,6 +1370,7 @@ "resolved": "https://registry.npmjs.org/escalade/-/escalade-3.2.0.tgz", "integrity": "sha512-WUj2qlxaQtO4g6Pq5c29GTcWGDyd8itL8zTlipgECz3JesAiiOKotd8JU6otB3PACgG6xkJUyVhboMS+bje/jA==", "dev": true, + "license": "MIT", "engines": { "node": ">=6" } @@ -2689,10 +2697,14 @@ "dev": true }, "node_modules/node-releases": { - "version": "2.0.27", - "resolved": "https://registry.npmjs.org/node-releases/-/node-releases-2.0.27.tgz", - "integrity": "sha512-nmh3lCkYZ3grZvqcCH+fjmQ7X+H0OeZgP40OierEaAptX4XofMh5kwNbWh7lBduUzCcV/8kZ+NDLCwm2iorIlA==", - "dev": true + "version": "2.0.54", + "resolved": "https://registry.npmjs.org/node-releases/-/node-releases-2.0.54.tgz", + "integrity": "sha512-YHs7BmmcsdAI5Ozuf8JZo6PT0mv2GIWC9vMfvUC3dp65M8hn7Ux8CPL+2oBI7juNuj9d0ndhTcznq2ODBps9cQ==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">=18" + } }, "node_modules/normalize-path": { "version": "3.0.0", @@ -3687,9 +3699,9 @@ } }, "node_modules/update-browserslist-db": { - "version": "1.2.3", - "resolved": "https://registry.npmjs.org/update-browserslist-db/-/update-browserslist-db-1.2.3.tgz", - "integrity": "sha512-Js0m9cx+qOgDxo0eMiFGEueWztz+d4+M3rGlmKPT+T4IS/jP4ylw3Nwpu6cpTTP8R1MAC1kF4VbdLt3ARf209w==", + "version": "1.3.2", + "resolved": "https://registry.npmjs.org/update-browserslist-db/-/update-browserslist-db-1.3.2.tgz", + "integrity": "sha512-UQ+MSxlhRm1bzjhU+DcuXfjFO1FzNtqhK5+9Yvlp90ItDLk5vT932A0rFu619nf7RVS+Y/VeaUW1jaRDqZ8VJw==", "dev": true, "funding": [ { @@ -3705,6 +3717,7 @@ "url": "https://github.com/sponsors/ai" } ], + "license": "MIT", "dependencies": { "escalade": "^3.2.0", "picocolors": "^1.1.1" @@ -4742,9 +4755,9 @@ } }, "baseline-browser-mapping": { - "version": "2.9.19", - "resolved": "https://registry.npmjs.org/baseline-browser-mapping/-/baseline-browser-mapping-2.9.19.tgz", - "integrity": "sha512-ipDqC8FrAl/76p2SSWKSI+H9tFwm7vYqXQrItCuiVPt26Km0jS+NzSsBWAaBusvSbQcfJG+JitdMm+wZAgTYqg==", + "version": "2.11.20", + "resolved": "https://registry.npmjs.org/baseline-browser-mapping/-/baseline-browser-mapping-2.11.20.tgz", + "integrity": "sha512-H0ulySigv6icDJ1F7SjtdCD6PrhTpdYCmP0CactWy1+ekh0AFd0o1Wn5T8b+hnTmdBx19u9yhL6wvCylXMY7zw==", "dev": true }, "batch": { @@ -4798,16 +4811,16 @@ } }, "browserslist": { - "version": "4.28.1", - "resolved": "https://registry.npmjs.org/browserslist/-/browserslist-4.28.1.tgz", - "integrity": "sha512-ZC5Bd0LgJXgwGqUknZY/vkUQ04r8NXnJZ3yYi4vDmSiZmC/pdSN0NbNRPxZpbtO4uAfDUAFffO8IZoM3Gj8IkA==", + "version": "4.28.8", + "resolved": "https://registry.npmjs.org/browserslist/-/browserslist-4.28.8.tgz", + "integrity": "sha512-V2NpofLblG64mfOtSgDhOJESZEGogzDMBv/q+W6oc4LXWP/q75eOXoOaaOu1EOadB9U4Bwx/e0yzbvwKH8zalA==", "dev": true, "requires": { - "baseline-browser-mapping": "^2.9.0", - "caniuse-lite": "^1.0.30001759", - "electron-to-chromium": "^1.5.263", - "node-releases": "^2.0.27", - "update-browserslist-db": "^1.2.0" + "baseline-browser-mapping": "^2.11.12", + "caniuse-lite": "^1.0.30001809", + "electron-to-chromium": "^1.5.402", + "node-releases": "^2.0.53", + "update-browserslist-db": "^1.3.0" } }, "buffer-from": { @@ -4858,9 +4871,9 @@ } }, "caniuse-lite": { - "version": "1.0.30001768", - "resolved": "https://registry.npmjs.org/caniuse-lite/-/caniuse-lite-1.0.30001768.tgz", - "integrity": "sha512-qY3aDRZC5nWPgHUgIB84WL+nySuo19wk0VJpp/XI9T34lrvkyhRvNVOFJOp2kxClQhiFBu+TaUSudf6oa3vkSA==", + "version": "1.0.30001810", + "resolved": "https://registry.npmjs.org/caniuse-lite/-/caniuse-lite-1.0.30001810.tgz", + "integrity": "sha512-TITQPUkaz+aVk5GL6NhOdwk1aEaNTSDPsGFWrTuhKGtjTF70jL/Oht2W4c6rXUe5fu7Ie19VIahAXHIIiWWNeg==", "dev": true }, "chokidar": { @@ -5091,9 +5104,9 @@ "dev": true }, "electron-to-chromium": { - "version": "1.5.286", - "resolved": "https://registry.npmjs.org/electron-to-chromium/-/electron-to-chromium-1.5.286.tgz", - "integrity": "sha512-9tfDXhJ4RKFNerfjdCcZfufu49vg620741MNs26a9+bhLThdB+plgMeou98CAaHu/WATj2iHOOHTp1hWtABj2A==", + "version": "1.5.419", + "resolved": "https://registry.npmjs.org/electron-to-chromium/-/electron-to-chromium-1.5.419.tgz", + "integrity": "sha512-nHMPn8x4yCxCI0iSnL+LlHL5sUoUfjLXkcRIagZ4GBdrfFLFaiLNvzJWbJqZhFT9IAhw5tUSNlhggWN+otvp/A==", "dev": true }, "encodeurl": { @@ -5965,9 +5978,9 @@ "dev": true }, "node-releases": { - "version": "2.0.27", - "resolved": "https://registry.npmjs.org/node-releases/-/node-releases-2.0.27.tgz", - "integrity": "sha512-nmh3lCkYZ3grZvqcCH+fjmQ7X+H0OeZgP40OierEaAptX4XofMh5kwNbWh7lBduUzCcV/8kZ+NDLCwm2iorIlA==", + "version": "2.0.54", + "resolved": "https://registry.npmjs.org/node-releases/-/node-releases-2.0.54.tgz", + "integrity": "sha512-YHs7BmmcsdAI5Ozuf8JZo6PT0mv2GIWC9vMfvUC3dp65M8hn7Ux8CPL+2oBI7juNuj9d0ndhTcznq2ODBps9cQ==", "dev": true }, "normalize-path": { @@ -6663,9 +6676,9 @@ "dev": true }, "update-browserslist-db": { - "version": "1.2.3", - "resolved": "https://registry.npmjs.org/update-browserslist-db/-/update-browserslist-db-1.2.3.tgz", - "integrity": "sha512-Js0m9cx+qOgDxo0eMiFGEueWztz+d4+M3rGlmKPT+T4IS/jP4ylw3Nwpu6cpTTP8R1MAC1kF4VbdLt3ARf209w==", + "version": "1.3.2", + "resolved": "https://registry.npmjs.org/update-browserslist-db/-/update-browserslist-db-1.3.2.tgz", + "integrity": "sha512-UQ+MSxlhRm1bzjhU+DcuXfjFO1FzNtqhK5+9Yvlp90ItDLk5vT932A0rFu619nf7RVS+Y/VeaUW1jaRDqZ8VJw==", "dev": true, "requires": { "escalade": "^3.2.0", From 37ee169ce505c674325438acfb0731e5b36af465 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Tue, 1 Sep 2026 19:56:43 +0000 Subject: [PATCH 04/37] chore(deps): bump the all-other-cargo-deps group with 7 updates (#24840) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Bumps the all-other-cargo-deps group with 7 updates: | Package | From | To | | --- | --- | --- | | [flate2](https://github.com/rust-lang/flate2-rs) | `1.1.9` | `1.1.10` | | [indexmap](https://github.com/indexmap-rs/indexmap) | `2.14.0` | `2.14.1` | | [log](https://github.com/rust-lang/log) | `0.4.33` | `0.4.34` | | [uuid](https://github.com/uuid-rs/uuid) | `1.24.1` | `1.26.0` | | [crc32fast](https://github.com/srijs/rust-crc32fast) | `1.5.0` | `1.5.1` | | [twox-hash](https://github.com/shepmaster/twox-hash) | `2.1.3` | `2.1.4` | | [syn](https://github.com/dtolnay/syn) | `3.0.3` | `3.0.4` | Updates `flate2` from 1.1.9 to 1.1.10
Release notes

Sourced from flate2's releases.

1.1.10

What's Changed

New Contributors

Full Changelog: https://github.com/rust-lang/flate2-rs/compare/1.1.9...1.1.10

Commits
  • ed93d4f Merge pull request #558 from rust-lang/lib-doc-update
  • fb5228d Merge pull request #559 from bushrat011899/no_std
  • 6ed3ba3 Add executable no_std example
  • faed8a0 Expand CI to test no_std compatibility and correctness
  • 2ba8e7e Add unstable no_std support behind flate2_unstable_nightly_alloc_io
  • 3fe1126 Centralize usage of std for error and io
  • 98e313a Add GzHeader::mtime_as_duration
  • 0642965 Switch to core implicit prelude and only use std where required
  • 454a63c Remove left-over dbg! statement
  • 2a490b7 Add runtime_detection feature
  • Additional commits viewable in compare view

Updates `indexmap` from 2.14.0 to 2.14.1
Changelog

Sourced from indexmap's changelog.

2.14.1 (2026-08-28)

  • Simplify comparisons where Equivalent isn't needed (Q = K).
  • Unify index assertions for bounds checks.
  • Fix (or expect) clippy lints.
Commits
  • fdf7e17 Merge pull request #449 from cuviper/release-2.14.1
  • ada540e Release 2.14.1
  • af93b43 expect clippy::redundant_slicing in tests
  • c95da18 fix clippy::derivable_impls
  • 2196365 fix clippy::useless_vec (and more) in tests
  • 1c2be7b use inherent usize::MAX
  • 64f4a06 fix clippy::int_plus_one
  • 41760c5 fix clippy::map_entry
  • be7ffd0 expect clippy::unnecessary_get_then_check in benches
  • bb35663 expect clippy::reversed_empty_ranges in tests
  • Additional commits viewable in compare view

Updates `log` from 0.4.33 to 0.4.34
Release notes

Sourced from log's releases.

0.4.34

What's Changed

New Contributors

Full Changelog: https://github.com/rust-lang/log/compare/0.4.33...0.4.34

Changelog

Sourced from log's changelog.

[0.4.34] - 2026-08-22

What's Changed

New Contributors

Full Changelog: https://github.com/rust-lang/log/compare/0.4.33...0.4.34

Commits

Updates `uuid` from 1.24.1 to 1.26.0
Release notes

Sourced from uuid's releases.

v1.26.0

What's Changed

Full Changelog: https://github.com/uuid-rs/uuid/compare/1.25.0...v1.26.0

1.25.0

What's Changed

New Contributors

Full Changelog: https://github.com/uuid-rs/uuid/compare/v1.24.1...1.25.0

Commits
  • cdc96a8 Merge pull request #905 from uuid-rs/cargo/v1.26.0
  • 34e4f49 don't test macros under miri
  • d9e7242 update nightly used for miri
  • ec16819 prepare for 1.26.0 release
  • 162cd20 Merge pull request #904 from ChrisJr404/v7-additional-precision-bits
  • 97eceff Add ContextV7::with_additional_precision_bits for microsecond clocks
  • 302e0bf Merge pull request #903 from uuid-rs/cargo/1.25.0
  • b7ccde8 prepare for 1.25.0 release
  • c62dffb Merge pull request #902 from ChrisJr404/serde-bytes-module
  • 8c198b2 Add a serde::bytes module that encodes as a byte string
  • See full diff in compare view

Updates `crc32fast` from 1.5.0 to 1.5.1
Commits
  • a150f65 release 1.5.1
  • f066e8d perf(simd): widen x86 folds, add ARM 3-way, and speed up small inputs (#56)
  • d5c123d consolidate dword load in baseline implementation (#55)
  • 50e2046 downgrade msrv ci run to just cargo build
  • See full diff in compare view

Updates `twox-hash` from 2.1.3 to 2.1.4
Changelog

Sourced from twox-hash's changelog.

2.1.4 - 2026-08-27

Changed

  • Documentation added about the stability of the hashing algorithms.
Commits
  • 6f866bf Release version 2.1.4
  • bcfd930 Update the changelog
  • 188f698 Merge pull request #125 from shepmaster/32-bit-consistency
  • e37e1ef Document the stability of the algorithms and caveats
  • a08cde9 Run the tests on a 32-bit platform (via Miri)
  • 9cba544 Adjust test to compile when usize is 32-bit
  • 6f8020e Merge pull request #126 from shepmaster/maint
  • a001afb Upgrade GHA to actions/checkout@v7
  • 4c480b6 Apply some extra Clippy lints
  • See full diff in compare view

Updates `syn` from 3.0.3 to 3.0.4
Release notes

Sourced from syn's releases.

3.0.4

  • Allow safe fn in impl Parse for ForeignItemFn (#2078)
Commits
  • b5d62a6 Release 3.0.4
  • abf019c Merge pull request #2078 from dtolnay/foreginitemfn
  • d454333 Allow safe fn in impl Parse for ForeignItemFn
  • 8011b1c Update test suite to nightly-2026-08-18
  • 56a8d83 Raise rayon thread size for tests
  • f2c5c50 Ignore assert_is_empty pedantic clippy lint
  • 0eba76d Update test suite to nightly-2026-08-05
  • baaebce Update test suite to nightly-2026-07-25
  • b886a38 Update test suite to nightly-2026-07-24
  • 3c41416 Update test suite to nightly-2026-07-23
  • See full diff in compare view

Dependabot will resolve any conflicts with this PR as long as you don't alter it yourself. You can also trigger a rebase manually by commenting `@dependabot rebase`. [//]: # (dependabot-automerge-start) [//]: # (dependabot-automerge-end) ---
Dependabot commands and options
You can trigger Dependabot actions by commenting on this PR: - `@dependabot rebase` will rebase this PR - `@dependabot recreate` will recreate this PR, overwriting any edits that have been made to it - `@dependabot show ignore conditions` will show all of the ignore conditions of the specified dependency - `@dependabot ignore major version` will close this group update PR and stop Dependabot creating any more for the specific dependency's major version (unless you unignore this specific dependency's major version or upgrade to it yourself) - `@dependabot ignore minor version` will close this group update PR and stop Dependabot creating any more for the specific dependency's minor version (unless you unignore this specific dependency's minor version or upgrade to it yourself) - `@dependabot ignore ` will close this group update PR and stop Dependabot creating any more for the specific dependency (unless you unignore this specific dependency or upgrade to it yourself) - `@dependabot unignore ` will remove all of the ignore conditions of the specified dependency - `@dependabot unignore ` will remove the ignore condition of the specified dependency and ignore conditions
Signed-off-by: dependabot[bot] Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> --- Cargo.lock | 92 +++++++++++++++++++++++++++--------------------------- 1 file changed, 46 insertions(+), 46 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 52797394237c0..5aee4c4e85ea7 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -226,7 +226,7 @@ dependencies = [ "bzip2", "crc", "flate2", - "indexmap 2.14.0", + "indexmap 2.14.1", "liblzma", "rand 0.9.4", "serde", @@ -356,7 +356,7 @@ dependencies = [ "arrow-select", "chrono", "half", - "indexmap 2.14.0", + "indexmap 2.14.1", "itoa", "lexical-core", "memchr", @@ -511,7 +511,7 @@ checksum = "82f6aeea286b8eb4dd3431a1be1b59d290ace00f5bfd8e2a159bc2a05e2c1667" dependencies = [ "proc-macro2", "quote", - "syn 3.0.3", + "syn 3.0.4", ] [[package]] @@ -1335,7 +1335,7 @@ dependencies = [ "heck", "proc-macro2", "quote", - "syn 3.0.3", + "syn 3.0.4", ] [[package]] @@ -1510,9 +1510,9 @@ checksum = "217698eaf96b4a3f0bc4f3662aaa55bdf913cd54d7204591faa790070c6d0853" [[package]] name = "crc32fast" -version = "1.5.0" +version = "1.5.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9481c1c90cbf2ac953f07c8d4a58aa3945c425b7185c9154d67a65e4230da511" +checksum = "8498c871161e1742aaa9d52551b2d6ebdd4c3d45a3be423e3728f33b955be550" dependencies = [ "cfg-if", ] @@ -1746,7 +1746,7 @@ dependencies = [ "futures", "glob", "half", - "indexmap 2.14.0", + "indexmap 2.14.1", "insta", "itertools 0.15.0", "liblzma", @@ -1897,7 +1897,7 @@ dependencies = [ "half", "hashbrown 0.17.1", "hex", - "indexmap 2.14.0", + "indexmap 2.14.1", "insta", "itertools 0.15.0", "libc", @@ -2173,7 +2173,7 @@ dependencies = [ "datafusion-proto-common", "datafusion-proto-models", "env_logger", - "indexmap 2.14.0", + "indexmap 2.14.1", "insta", "itertools 0.15.0", "recursive", @@ -2187,7 +2187,7 @@ version = "55.0.0" dependencies = [ "arrow", "datafusion-common", - "indexmap 2.14.0", + "indexmap 2.14.1", "insta", "itertools 0.15.0", "rstest", @@ -2367,7 +2367,7 @@ version = "55.0.0" dependencies = [ "datafusion-doc", "quote", - "syn 3.0.3", + "syn 3.0.4", ] [[package]] @@ -2388,7 +2388,7 @@ dependencies = [ "datafusion-physical-expr", "datafusion-sql", "env_logger", - "indexmap 2.14.0", + "indexmap 2.14.1", "insta", "itertools 0.15.0", "log", @@ -2413,7 +2413,7 @@ dependencies = [ "datafusion-proto-models", "half", "hashbrown 0.17.1", - "indexmap 2.14.0", + "indexmap 2.14.1", "insta", "itertools 0.15.0", "parking_lot", @@ -2448,7 +2448,7 @@ dependencies = [ "datafusion-expr-common", "datafusion-proto-models", "hashbrown 0.17.1", - "indexmap 2.14.0", + "indexmap 2.14.1", "itertools 0.15.0", "parking_lot", "pin-project", @@ -2506,7 +2506,7 @@ dependencies = [ "futures", "half", "hashbrown 0.17.1", - "indexmap 2.14.0", + "indexmap 2.14.1", "insta", "itertools 0.15.0", "log", @@ -2659,7 +2659,7 @@ dependencies = [ "datafusion-functions-nested", "datafusion-functions-window", "env_logger", - "indexmap 2.14.0", + "indexmap 2.14.1", "insta", "itertools 0.15.0", "log", @@ -3060,9 +3060,9 @@ dependencies = [ [[package]] name = "flate2" -version = "1.1.9" +version = "1.1.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "843fba2746e448b37e26a819579957415c8cef339bf08564fe8b7ddbd959573c" +checksum = "6e634e2e0ebac1ee034020da1ca582e17ffe4e0f5e985823721e168928136dcb" dependencies = [ "crc32fast", "miniz_oxide", @@ -3303,7 +3303,7 @@ dependencies = [ "futures-core", "futures-sink", "http 1.4.0", - "indexmap 2.14.0", + "indexmap 2.14.1", "slab", "tokio", "tokio-util", @@ -3735,9 +3735,9 @@ dependencies = [ [[package]] name = "indexmap" -version = "2.14.0" +version = "2.14.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d466e9454f08e4a911e14806c24e16fba1b4c121d1ea474396f396069cf949d9" +checksum = "07aa2048142242915a31d35844fb311e0e53fcca590c3a0a40dcf1b841fa09eb" dependencies = [ "equivalent", "hashbrown 0.17.1", @@ -4073,9 +4073,9 @@ dependencies = [ [[package]] name = "log" -version = "0.4.33" +version = "0.4.34" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0ceec5bc11778974d1bcb055b18002eba7f4b3518b6a0081b3af5f21666da9ad" +checksum = "f9f8bd3e56ce4dfc153cf470fffbfa98c7620958b312ca5c3a4b8d5181fd13c6" [[package]] name = "lru-slab" @@ -4151,9 +4151,9 @@ dependencies = [ [[package]] name = "miniz_oxide" -version = "0.8.9" +version = "0.9.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1fa76a2c86f704bdb222d66965fb3d63269ce38518b83cb0575fca855ebb6316" +checksum = "b63fbc4a50860e98e7b2aa7804ded1db5cbc3aff9193adaff57a6931bf7c4b4c" dependencies = [ "adler2", "simd-adler32", @@ -4636,7 +4636,7 @@ checksum = "8701b58ea97060d5e5b155d383a69952a60943f0e6dfe30b04c287beb0b27455" dependencies = [ "fixedbitset", "hashbrown 0.15.5", - "indexmap 2.14.0", + "indexmap 2.14.1", "serde", ] @@ -5606,7 +5606,7 @@ checksum = "e7a5d71263a5a7d47b41f6b3f06ba276f10cc18b0931f1799f710578e2309348" dependencies = [ "proc-macro2", "quote", - "syn 3.0.3", + "syn 3.0.4", ] [[package]] @@ -5626,7 +5626,7 @@ version = "1.0.151" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c841b55ecdae098c80dcae9cf767f6f8a0c2cdb3416bbef72181df4d0fe73f14" dependencies = [ - "indexmap 2.14.0", + "indexmap 2.14.1", "itoa", "memchr", "serde", @@ -5689,7 +5689,7 @@ dependencies = [ "chrono", "hex", "indexmap 1.9.3", - "indexmap 2.14.0", + "indexmap 2.14.1", "schemars 0.9.0", "schemars 1.2.1", "serde_core", @@ -5716,7 +5716,7 @@ version = "0.9.34+deprecated" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6a8b1a1a2ebf674015cc02edccce75287f1a0130d394307b36743c2f5d504b47" dependencies = [ - "indexmap 2.14.0", + "indexmap 2.14.1", "itoa", "ryu", "serde", @@ -6041,7 +6041,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e620ff4d5c02fd6f7752931aa74b16a26af66a63022cc1ad412c77edbe0bab47" dependencies = [ "heck", - "indexmap 2.14.0", + "indexmap 2.14.1", "pbjson 0.8.0", "pbjson-build 0.8.0", "pbjson-types", @@ -6080,9 +6080,9 @@ dependencies = [ [[package]] name = "syn" -version = "3.0.3" +version = "3.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "53e9bae58849f64dfa4f5d5ae372c8341f7305f82a3868709269343628b659a3" +checksum = "e6275cddf4610d1775e6d1fe9469b2e77d0f39fd98fb7450901b821e0c53649f" dependencies = [ "proc-macro2", "quote", @@ -6205,7 +6205,7 @@ checksum = "bc04cd3e1236dd4a98afca4569f2deb3f120e5422a4023be2cb683f8486292af" dependencies = [ "proc-macro2", "quote", - "syn 3.0.3", + "syn 3.0.4", ] [[package]] @@ -6388,7 +6388,7 @@ version = "1.1.4+spec-1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3aace63f4bbcdfc2c965b059de67119c89c4017a70d633be6c104910f67056f5" dependencies = [ - "indexmap 2.14.0", + "indexmap 2.14.1", "serde_core", "serde_spanned", "toml_datetime", @@ -6412,7 +6412,7 @@ version = "0.25.11+spec-1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0b59c4d22ed448339746c59b905d24568fcbb3ab65a500494f7b8c3e97739f2b" dependencies = [ - "indexmap 2.14.0", + "indexmap 2.14.1", "toml_datetime", "toml_parser", "winnow", @@ -6481,7 +6481,7 @@ checksum = "ebe5ef63511595f1344e2d5cfa636d973292adc0eec1f0ad45fae9f0851ab1d4" dependencies = [ "futures-core", "futures-util", - "indexmap 2.14.0", + "indexmap 2.14.1", "pin-project-lite", "slab", "sync_wrapper", @@ -6587,9 +6587,9 @@ checksum = "e421abadd41a4225275504ea4d6566923418b7f05506fbc9c0fe86ba7396114b" [[package]] name = "twox-hash" -version = "2.1.3" +version = "2.1.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8464ec13c3691491391d9fce00f6416c9a48e46972f72d7865688be2080192c9" +checksum = "5283634e518fe9e82c7b20520bb4bc209009fd16c82077c802f8111ecbb0117a" dependencies = [ "rand 0.10.1", ] @@ -6782,9 +6782,9 @@ checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821" [[package]] name = "uuid" -version = "1.24.1" +version = "1.26.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2cefc03fd367c0c6d4305de1b312cf00248c4114f4a0418ce6a6af769e3b0bd9" +checksum = "b5772d71c9be8a8a6ac2117d949c5b224c1b72241bb611d9a3012edcf8af7812" dependencies = [ "getrandom 0.4.2", "js-sys", @@ -6990,7 +6990,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "bb0e353e6a2fbdc176932bbaab493762eb1255a7900fe0fea1a2f96c296cc909" dependencies = [ "anyhow", - "indexmap 2.14.0", + "indexmap 2.14.1", "wasm-encoder", "wasmparser", ] @@ -7016,7 +7016,7 @@ checksum = "47b807c72e1bac69382b3a6fb3dbe8ea4c0ed87ff5629b8685ae6b9a611028fe" dependencies = [ "bitflags", "hashbrown 0.15.5", - "indexmap 2.14.0", + "indexmap 2.14.1", "semver", ] @@ -7393,7 +7393,7 @@ checksum = "b7c566e0f4b284dd6561c786d9cb0142da491f46a9fbed79ea69cdad5db17f21" dependencies = [ "anyhow", "heck", - "indexmap 2.14.0", + "indexmap 2.14.1", "prettyplease", "syn 2.0.119", "wasm-metadata", @@ -7424,7 +7424,7 @@ checksum = "9d66ea20e9553b30172b5e831994e35fbde2d165325bec84fc43dbf6f4eb9cb2" dependencies = [ "anyhow", "bitflags", - "indexmap 2.14.0", + "indexmap 2.14.1", "log", "serde", "serde_derive", @@ -7443,7 +7443,7 @@ checksum = "ecc8ac4bc1dc3381b7f59c34f00b67e18f910c2c0f50015669dde7def656a736" dependencies = [ "anyhow", "id-arena", - "indexmap 2.14.0", + "indexmap 2.14.1", "log", "semver", "serde", From 873308958b615166a60491098810773cfa9c8fa5 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Tue, 1 Sep 2026 19:57:06 +0000 Subject: [PATCH 05/37] chore(deps): bump taiki-e/install-action from 2.86.5 to 2.87.1 (#24845) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Bumps [taiki-e/install-action](https://github.com/taiki-e/install-action) from 2.86.5 to 2.87.1.
Release notes

Sourced from taiki-e/install-action's releases.

2.87.1

  • Update uv@latest to 0.12.7.

  • Update typos@latest to 1.49.1.

  • Update syft@latest to 1.51.1.

  • Update prek@latest to 0.5.0.

  • Update d2@latest to 0.8.2.

  • Update cargo-zigbuild@latest to 0.23.3.

  • Update cargo-rdme@latest to 2.2.2.

  • Update biome@latest to 2.5.11.

2.87.0

  • Support kache. (#1980, thanks @​ChrisJr404)

  • Update vacuum@latest to 0.30.1.

  • Update uv@latest to 0.12.6.

  • Update mise@latest to 2026.8.14.

  • Update editorconfig-checker@latest to 3.11.2.

2.86.8

  • Update wasmtime@latest to 48.0.1.

  • Update wasm-tools@latest to 1.258.0.

  • Update oxfmt@latest to 1.80.0.

  • Update mise@latest to 2026.8.12.

  • Update kingfisher@latest to 2.0.0.

  • Update cargo-zigbuild@latest to 0.23.2.

2.86.7

  • Update tombi@latest to 1.4.1.

  • Update rafn@latest to 0.1.5.

  • Update cargo-binstall@latest to 1.22.0.

2.86.6

  • Update dprint@latest to 0.56.1.

... (truncated)

Changelog

Sourced from taiki-e/install-action's changelog.

Changelog

All notable changes to this project will be documented in this file.

This project adheres to Semantic Versioning.

[Unreleased]

[2.87.2] - 2026-08-30

  • Update typos@latest to 1.50.0.

  • Update tombi@latest to 1.5.0.

  • Update shfmt@latest to 3.14.0.

[2.87.1] - 2026-08-29

  • Update uv@latest to 0.12.7.

  • Update typos@latest to 1.49.1.

  • Update syft@latest to 1.51.1.

  • Update prek@latest to 0.5.0.

  • Update d2@latest to 0.8.2.

  • Update cargo-zigbuild@latest to 0.23.3.

  • Update cargo-rdme@latest to 2.2.2.

  • Update biome@latest to 2.5.11.

[2.87.0] - 2026-08-27

  • Support kache. (#1980, thanks @​ChrisJr404)

  • Update vacuum@latest to 0.30.1.

  • Update uv@latest to 0.12.6.

  • Update mise@latest to 2026.8.14.

  • Update editorconfig-checker@latest to 3.11.2.

... (truncated)

Commits

[![Dependabot compatibility score](https://dependabot-badges.githubapp.com/badges/compatibility_score?dependency-name=taiki-e/install-action&package-manager=github_actions&previous-version=2.86.5&new-version=2.87.1)](https://docs.github.com/en/github/managing-security-vulnerabilities/about-dependabot-security-updates#about-compatibility-scores) Dependabot will resolve any conflicts with this PR as long as you don't alter it yourself. You can also trigger a rebase manually by commenting `@dependabot rebase`. [//]: # (dependabot-automerge-start) [//]: # (dependabot-automerge-end) ---
Dependabot commands and options
You can trigger Dependabot actions by commenting on this PR: - `@dependabot rebase` will rebase this PR - `@dependabot recreate` will recreate this PR, overwriting any edits that have been made to it - `@dependabot show ignore conditions` will show all of the ignore conditions of the specified dependency - `@dependabot ignore this major version` will close this PR and stop Dependabot creating any more for this major version (unless you reopen the PR or upgrade to it yourself) - `@dependabot ignore this minor version` will close this PR and stop Dependabot creating any more for this minor version (unless you reopen the PR or upgrade to it yourself) - `@dependabot ignore this dependency` will close this PR and stop Dependabot creating any more for this dependency (unless you reopen the PR or upgrade to it yourself)
Signed-off-by: dependabot[bot] Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> --- .github/workflows/audit.yml | 2 +- .github/workflows/breaking_changes_detector.yml | 2 +- .github/workflows/dependencies.yml | 2 +- .github/workflows/dev.yml | 6 +++--- .github/workflows/docs.yaml | 2 +- .github/workflows/docs_pr.yaml | 2 +- .github/workflows/rust.yml | 8 ++++---- 7 files changed, 12 insertions(+), 12 deletions(-) diff --git a/.github/workflows/audit.yml b/.github/workflows/audit.yml index 494e469a65ac9..e3071c1fc8ce7 100644 --- a/.github/workflows/audit.yml +++ b/.github/workflows/audit.yml @@ -45,7 +45,7 @@ jobs: steps: - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 - name: Install cargo-audit - uses: taiki-e/install-action@ba47c86ac325773530516bb756137ac718732518 # v2.86.5 + uses: taiki-e/install-action@742a3317eac7bd62f91cd888b4eead5e784ba833 # v2.87.1 with: tool: cargo-audit - name: Run audit check diff --git a/.github/workflows/breaking_changes_detector.yml b/.github/workflows/breaking_changes_detector.yml index 9e23fef8a6c8a..a44a50a53f1e8 100644 --- a/.github/workflows/breaking_changes_detector.yml +++ b/.github/workflows/breaking_changes_detector.yml @@ -89,7 +89,7 @@ jobs: - name: Install cargo-semver-checks if: steps.changed_crates.outputs.packages != '' - uses: taiki-e/install-action@ba47c86ac325773530516bb756137ac718732518 # v2.86.5 + uses: taiki-e/install-action@742a3317eac7bd62f91cd888b4eead5e784ba833 # v2.87.1 with: tool: cargo-semver-checks diff --git a/.github/workflows/dependencies.yml b/.github/workflows/dependencies.yml index d5acaa6338db7..726d0783eee39 100644 --- a/.github/workflows/dependencies.yml +++ b/.github/workflows/dependencies.yml @@ -63,7 +63,7 @@ jobs: steps: - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 - name: Install cargo-machete - uses: taiki-e/install-action@ba47c86ac325773530516bb756137ac718732518 # v2.86.5 + uses: taiki-e/install-action@742a3317eac7bd62f91cd888b4eead5e784ba833 # v2.87.1 with: tool: cargo-machete@0.9 - name: Detect unused dependencies diff --git a/.github/workflows/dev.yml b/.github/workflows/dev.yml index fbbb1e2a9d691..83ebb6e78b678 100644 --- a/.github/workflows/dev.yml +++ b/.github/workflows/dev.yml @@ -38,7 +38,7 @@ jobs: steps: - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 - name: Install HawkEye - uses: taiki-e/install-action@ba47c86ac325773530516bb756137ac718732518 # v2.86.5 + uses: taiki-e/install-action@742a3317eac7bd62f91cd888b4eead5e784ba833 # v2.87.1 with: tool: hawkeye@7.0.0 - name: Run license header check @@ -66,7 +66,7 @@ jobs: source ci/scripts/utils/tool_versions.sh echo "LYCHEE_VERSION=${LYCHEE_VERSION}" >> "$GITHUB_ENV" - name: Install lychee - uses: taiki-e/install-action@ba47c86ac325773530516bb756137ac718732518 # v2.86.5 + uses: taiki-e/install-action@742a3317eac7bd62f91cd888b4eead5e784ba833 # v2.87.1 with: tool: lychee@${{ env.LYCHEE_VERSION }} - name: Run markdown link check @@ -91,7 +91,7 @@ jobs: # it may cause checks to fail more often. # We can upgrade it manually once a while. - name: Install typos - uses: taiki-e/install-action@ba47c86ac325773530516bb756137ac718732518 # v2.86.5 + uses: taiki-e/install-action@742a3317eac7bd62f91cd888b4eead5e784ba833 # v2.87.1 with: tool: typos@1.37.0 - name: Run typos check diff --git a/.github/workflows/docs.yaml b/.github/workflows/docs.yaml index c52ed5681c64f..4eda5ac4590a2 100644 --- a/.github/workflows/docs.yaml +++ b/.github/workflows/docs.yaml @@ -53,7 +53,7 @@ jobs: sudo apt-get update sudo apt-get install -y graphviz - name: Install cargo-depgraph - uses: taiki-e/install-action@ba47c86ac325773530516bb756137ac718732518 # v2.86.5 + uses: taiki-e/install-action@742a3317eac7bd62f91cd888b4eead5e784ba833 # v2.87.1 with: tool: cargo-depgraph@1.6 diff --git a/.github/workflows/docs_pr.yaml b/.github/workflows/docs_pr.yaml index 2cbb64274a526..097cfe130b5b0 100644 --- a/.github/workflows/docs_pr.yaml +++ b/.github/workflows/docs_pr.yaml @@ -59,7 +59,7 @@ jobs: sudo apt-get update sudo apt-get install -y graphviz - name: Install cargo-depgraph - uses: taiki-e/install-action@ba47c86ac325773530516bb756137ac718732518 # v2.86.5 + uses: taiki-e/install-action@742a3317eac7bd62f91cd888b4eead5e784ba833 # v2.87.1 with: tool: cargo-depgraph@1.6 - name: Build docs html and check for warnings diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml index 9de0c80e7a1f0..7925fdd817f0c 100644 --- a/.github/workflows/rust.yml +++ b/.github/workflows/rust.yml @@ -346,7 +346,7 @@ jobs: - name: Install llvm-tools-preview run: rustup component add llvm-tools-preview - name: Install cargo-llvm-cov - uses: taiki-e/install-action@ba47c86ac325773530516bb756137ac718732518 # v2.86.5 + uses: taiki-e/install-action@742a3317eac7bd62f91cd888b4eead5e784ba833 # v2.87.1 with: tool: cargo-llvm-cov - name: Rust Dependency Cache @@ -491,7 +491,7 @@ jobs: sudo apt-get update -qq sudo apt-get install -y -qq clang - name: Setup wasm-pack - uses: taiki-e/install-action@ba47c86ac325773530516bb756137ac718732518 # v2.86.5 + uses: taiki-e/install-action@742a3317eac7bd62f91cd888b4eead5e784ba833 # v2.87.1 with: tool: wasm-pack - name: Run tests with headless mode @@ -717,7 +717,7 @@ jobs: with: rust-version: stable - name: Install taplo - uses: taiki-e/install-action@ba47c86ac325773530516bb756137ac718732518 # v2.86.5 + uses: taiki-e/install-action@742a3317eac7bd62f91cd888b4eead5e784ba833 # v2.87.1 with: tool: taplo-cli@0.9 # if you encounter an error, try running 'taplo format' to fix the formatting automatically. @@ -802,7 +802,7 @@ jobs: - name: Setup Rust toolchain uses: ./.github/actions/setup-builder - name: Install cargo-msrv - uses: taiki-e/install-action@ba47c86ac325773530516bb756137ac718732518 # v2.86.5 + uses: taiki-e/install-action@742a3317eac7bd62f91cd888b4eead5e784ba833 # v2.87.1 with: tool: cargo-msrv From ad5820c84ed00d31125842e62a5124303a39c8d2 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Tue, 1 Sep 2026 19:57:36 +0000 Subject: [PATCH 06/37] chore(deps): bump the codeql-actions group across 1 directory with 2 updates (#24844) Bumps the codeql-actions group with 2 updates in the / directory: [github/codeql-action/init](https://github.com/github/codeql-action) and [github/codeql-action/analyze](https://github.com/github/codeql-action). Updates `github/codeql-action/init` from 4.37.8 to 4.37.9
Release notes

Sourced from github/codeql-action/init's releases.

v4.37.9

  • Update default CodeQL bundle version to 2.26.4. #4106
Changelog

Sourced from github/codeql-action/init's changelog.

CodeQL Action Changelog

See the releases page for the relevant changes to the CodeQL CLI and language packs.

[UNRELEASED]

No user facing changes.

4.37.9 - 26 Aug 2026

  • Update default CodeQL bundle version to 2.26.4. #4106

4.37.8 - 21 Aug 2026

No user facing changes.

4.37.7 - 13 Aug 2026

  • Update default CodeQL bundle version to 2.26.3. #4085

4.37.6 - 04 Aug 2026

  • Changed the default filepath for the new remote file address format that was introduced in CodeQL Action 4.37.0 / 3.37.0 to .github/codeql-config.yml to align it with the suggested path that is used elsewhere. #4070

4.37.5 - 03 Aug 2026

  • Fixed a bug where a network error while streaming the download of the CodeQL bundle could terminate the init Action instead of falling back to downloading the bundle before extracting it. #4061

4.37.4 - 29 Jul 2026

  • This version of the CodeQL Action adds support for the tools input for the codeql-action/init step to be specified using a github-codeql-tools repository property. This feature will gradually be rolled out following the release of this version. Once rolled out, this allows for the CodeQL CLI version that is used in GitHub-managed workflows, such as Default Setup, to be set to a custom value. For example, customers who run into issues with rate limits when a new CodeQL CLI version is released can set the value to toolcache to always use the CodeQL CLI version that is available in the runner toolcache. For Advanced Setup workflows, the value provided for tools in the workflow definition always takes precedence unless the value of the repository property starts with !. #4037
  • Update default CodeQL bundle version to 2.26.2. #4051

4.37.3 - 22 Jul 2026

No user facing changes.

4.37.2 - 21 Jul 2026

  • The new address format for the config-file input that was introduced in CodeQL Action 4.37.0 is now enabled by default. In addition to the format described there, the remote= prefix can now be used to explicitly indicate that the input refers to a remote file. All previous input formats continue to be accepted as well. #4023
  • The CodeQL Action can now make use of configured private registries in Default Setup to retrieve CodeQL configuration files from remote repositories that require authentication. This will allow customers to store their CodeQL configuration in a single repository that can then be referenced by Default Setup workflows in other repositories. We expect to roll this and other, related changes out to everyone in July. #4007

4.37.1 - 16 Jul 2026

  • Upcoming breaking change: Add a deprecation warning for customers using CodeQL version 2.20.6 and earlier. These versions of CodeQL were discontinued on 1 July 2026 alongside GitHub Enterprise Server 3.16, and will be unsupported by the next minor release of the CodeQL Action. #3956
  • Update default CodeQL bundle version to 2.26.1. #4019

4.37.0 - 08 Jul 2026

  • Update default CodeQL bundle version to 2.26.0. #3995

... (truncated)

Commits
  • cdf488f Merge pull request #4107 from github/update-v4.37.9-920ba7cd1
  • 7243f38 Update changelog for v4.37.9
  • 920ba7c Merge pull request #4106 from github/update-bundle/codeql-bundle-v2.26.4
  • ecfa6e1 Add changelog note
  • adcdf4a Update default bundle to codeql-bundle-v2.26.4
  • 486fec2 Merge pull request #4099 from github/update-supported-enterprise-server-versions
  • 134624c Merge pull request #4101 from github/dependabot/npm_and_yarn/npm-minor-457d82...
  • ff43db8 Merge pull request #4103 from github/mergeback/v4.37.8-to-main-db488dde
  • 4605e03 Rebuild
  • 099c869 Update changelog and version after v4.37.8
  • Additional commits viewable in compare view

Updates `github/codeql-action/analyze` from 4.37.8 to 4.37.9
Release notes

Sourced from github/codeql-action/analyze's releases.

v4.37.9

  • Update default CodeQL bundle version to 2.26.4. #4106
Changelog

Sourced from github/codeql-action/analyze's changelog.

CodeQL Action Changelog

See the releases page for the relevant changes to the CodeQL CLI and language packs.

[UNRELEASED]

No user facing changes.

4.37.9 - 26 Aug 2026

  • Update default CodeQL bundle version to 2.26.4. #4106

4.37.8 - 21 Aug 2026

No user facing changes.

4.37.7 - 13 Aug 2026

  • Update default CodeQL bundle version to 2.26.3. #4085

4.37.6 - 04 Aug 2026

  • Changed the default filepath for the new remote file address format that was introduced in CodeQL Action 4.37.0 / 3.37.0 to .github/codeql-config.yml to align it with the suggested path that is used elsewhere. #4070

4.37.5 - 03 Aug 2026

  • Fixed a bug where a network error while streaming the download of the CodeQL bundle could terminate the init Action instead of falling back to downloading the bundle before extracting it. #4061

4.37.4 - 29 Jul 2026

  • This version of the CodeQL Action adds support for the tools input for the codeql-action/init step to be specified using a github-codeql-tools repository property. This feature will gradually be rolled out following the release of this version. Once rolled out, this allows for the CodeQL CLI version that is used in GitHub-managed workflows, such as Default Setup, to be set to a custom value. For example, customers who run into issues with rate limits when a new CodeQL CLI version is released can set the value to toolcache to always use the CodeQL CLI version that is available in the runner toolcache. For Advanced Setup workflows, the value provided for tools in the workflow definition always takes precedence unless the value of the repository property starts with !. #4037
  • Update default CodeQL bundle version to 2.26.2. #4051

4.37.3 - 22 Jul 2026

No user facing changes.

4.37.2 - 21 Jul 2026

  • The new address format for the config-file input that was introduced in CodeQL Action 4.37.0 is now enabled by default. In addition to the format described there, the remote= prefix can now be used to explicitly indicate that the input refers to a remote file. All previous input formats continue to be accepted as well. #4023
  • The CodeQL Action can now make use of configured private registries in Default Setup to retrieve CodeQL configuration files from remote repositories that require authentication. This will allow customers to store their CodeQL configuration in a single repository that can then be referenced by Default Setup workflows in other repositories. We expect to roll this and other, related changes out to everyone in July. #4007

4.37.1 - 16 Jul 2026

  • Upcoming breaking change: Add a deprecation warning for customers using CodeQL version 2.20.6 and earlier. These versions of CodeQL were discontinued on 1 July 2026 alongside GitHub Enterprise Server 3.16, and will be unsupported by the next minor release of the CodeQL Action. #3956
  • Update default CodeQL bundle version to 2.26.1. #4019

4.37.0 - 08 Jul 2026

  • Update default CodeQL bundle version to 2.26.0. #3995

... (truncated)

Commits
  • cdf488f Merge pull request #4107 from github/update-v4.37.9-920ba7cd1
  • 7243f38 Update changelog for v4.37.9
  • 920ba7c Merge pull request #4106 from github/update-bundle/codeql-bundle-v2.26.4
  • ecfa6e1 Add changelog note
  • adcdf4a Update default bundle to codeql-bundle-v2.26.4
  • 486fec2 Merge pull request #4099 from github/update-supported-enterprise-server-versions
  • 134624c Merge pull request #4101 from github/dependabot/npm_and_yarn/npm-minor-457d82...
  • ff43db8 Merge pull request #4103 from github/mergeback/v4.37.8-to-main-db488dde
  • 4605e03 Rebuild
  • 099c869 Update changelog and version after v4.37.8
  • Additional commits viewable in compare view

Dependabot will resolve any conflicts with this PR as long as you don't alter it yourself. You can also trigger a rebase manually by commenting `@dependabot rebase`. [//]: # (dependabot-automerge-start) [//]: # (dependabot-automerge-end) ---
Dependabot commands and options
You can trigger Dependabot actions by commenting on this PR: - `@dependabot rebase` will rebase this PR - `@dependabot recreate` will recreate this PR, overwriting any edits that have been made to it - `@dependabot show ignore conditions` will show all of the ignore conditions of the specified dependency - `@dependabot ignore major version` will close this group update PR and stop Dependabot creating any more for the specific dependency's major version (unless you unignore this specific dependency's major version or upgrade to it yourself) - `@dependabot ignore minor version` will close this group update PR and stop Dependabot creating any more for the specific dependency's minor version (unless you unignore this specific dependency's minor version or upgrade to it yourself) - `@dependabot ignore ` will close this group update PR and stop Dependabot creating any more for the specific dependency (unless you unignore this specific dependency or upgrade to it yourself) - `@dependabot unignore ` will remove all of the ignore conditions of the specified dependency - `@dependabot unignore ` will remove the ignore condition of the specified dependency and ignore conditions
Signed-off-by: dependabot[bot] Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> --- .github/workflows/codeql.yml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/codeql.yml b/.github/workflows/codeql.yml index b9f90dbd8eb14..1a7d075d9e886 100644 --- a/.github/workflows/codeql.yml +++ b/.github/workflows/codeql.yml @@ -45,11 +45,11 @@ jobs: persist-credentials: false - name: Initialize CodeQL - uses: github/codeql-action/init@db488ddef3bf6cb639b32c2e9a7c0a7ea8271d28 # v4 + uses: github/codeql-action/init@cdf488f595d80d6e07e03d4674febd5ab45fa938 # v4 with: languages: actions - name: Perform CodeQL Analysis - uses: github/codeql-action/analyze@db488ddef3bf6cb639b32c2e9a7c0a7ea8271d28 # v4 + uses: github/codeql-action/analyze@cdf488f595d80d6e07e03d4674febd5ab45fa938 # v4 with: category: "/language:actions" From cafb80161649868d3aba8b5e728d0c611f112005 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Tue, 1 Sep 2026 20:44:46 +0000 Subject: [PATCH 07/37] chore(deps): bump blake2 from 0.10.6 to 0.11.0 (#24841) Bumps [blake2](https://github.com/RustCrypto/hashes) from 0.10.6 to 0.11.0.
Commits

[![Dependabot compatibility score](https://dependabot-badges.githubapp.com/badges/compatibility_score?dependency-name=blake2&package-manager=cargo&previous-version=0.10.6&new-version=0.11.0)](https://docs.github.com/en/github/managing-security-vulnerabilities/about-dependabot-security-updates#about-compatibility-scores) Dependabot will resolve any conflicts with this PR as long as you don't alter it yourself. You can also trigger a rebase manually by commenting `@dependabot rebase`. [//]: # (dependabot-automerge-start) [//]: # (dependabot-automerge-end) ---
Dependabot commands and options
You can trigger Dependabot actions by commenting on this PR: - `@dependabot rebase` will rebase this PR - `@dependabot recreate` will recreate this PR, overwriting any edits that have been made to it - `@dependabot show ignore conditions` will show all of the ignore conditions of the specified dependency - `@dependabot ignore this major version` will close this PR and stop Dependabot creating any more for this major version (unless you reopen the PR or upgrade to it yourself) - `@dependabot ignore this minor version` will close this PR and stop Dependabot creating any more for this minor version (unless you reopen the PR or upgrade to it yourself) - `@dependabot ignore this dependency` will close this PR and stop Dependabot creating any more for this dependency (unless you reopen the PR or upgrade to it yourself)
Signed-off-by: dependabot[bot] Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> --- Cargo.lock | 7 +++---- datafusion/functions/Cargo.toml | 2 +- 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 5aee4c4e85ea7..3077074e862fa 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1023,11 +1023,11 @@ checksum = "c4512299f36f043ab09a583e57bceb5a5aab7a73db1805848e8fef3c9e8c78b3" [[package]] name = "blake2" -version = "0.10.6" +version = "0.11.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "46502ad458c9a52b69d4d4d32775c788b7a1b85e8bc9d482d92250fc0e3f8efe" +checksum = "5b5d4d889834ee8ecfc0f8426ad30faf7cdcb10f741a8e6d7224d95325479f6f" dependencies = [ - "digest 0.10.7", + "digest 0.11.2", ] [[package]] @@ -2777,7 +2777,6 @@ checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" dependencies = [ "block-buffer 0.10.4", "crypto-common 0.1.7", - "subtle", ] [[package]] diff --git a/datafusion/functions/Cargo.toml b/datafusion/functions/Cargo.toml index 00d1531a3a771..27c644042becb 100644 --- a/datafusion/functions/Cargo.toml +++ b/datafusion/functions/Cargo.toml @@ -68,7 +68,7 @@ name = "datafusion_functions" arrow = { workspace = true } arrow-buffer = { workspace = true } base64 = { version = "0.23", optional = true } -blake2 = { version = "^0.10.2", optional = true } +blake2 = { version = "^0.11.0", optional = true } blake3 = { version = "1.8", optional = true } chrono = { workspace = true } chrono-tz = { version = "0.10.4", optional = true } From fefc227edb894a469f06037b111ec1057e04bcc7 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Tue, 1 Sep 2026 20:45:26 +0000 Subject: [PATCH 08/37] chore(deps-dev): bump the all-npm-deps group in /datafusion/wasmtest/datafusion-wasm-app with 2 updates (#24839) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Bumps the all-npm-deps group in /datafusion/wasmtest/datafusion-wasm-app with 2 updates: [webpack](https://github.com/webpack/webpack) and [webpack-cli](https://github.com/webpack/webpack-cli). Updates `webpack` from 5.109.2 to 5.110.1
Release notes

Sourced from webpack's releases.

v5.110.1

Patch Changes

v5.110.0

Minor Changes

  • Wrap concatenated modules in lazy __webpack_require__.cw accessors and inline require(), keeping a wrapped body's names and side effects intact. (by @​hai-x in #21519)

  • Add performance hints reporting what a build costs: duplicate packages and modules, circular dependencies, broad contexts, large modules and chunks, hotspots, eval, missing PURE annotations, polyfills, redundant dynamic imports, OS-dependent rules, cache effectiveness, how chunks load, what splitting refused, why an optimization was skipped, and rules, defines, externals, aliases and barrel reexports nothing uses. An oversized asset names its largest modules, and an entrypoint carrying the runtime recommends optimization.runtimeChunk. Enable every check not set individually with performance.all, report hints in stats only with performance.hints: "stats", and get them in a stable order that leaves the build hashes unchanged. (by @​alexander-akait in #21841)

  • Add the descriptionRelativePath module rule condition. (by @​alexander-akait in #21705)

  • Add OS-independent glob matching to module rules. (by @​alexander-akait in #21771)

  • Report inner-graph, AMD and bare module bailouts in optimizationBailout. (by @​alexander-akait in #21740)

  • Allow marking externals as side-effect-free with a sideEffects flag. (by @​alexander-akait in #21712)

  • Give externals the original request of a context module element. (by @​alexander-akait in #21780)

  • Add the externalsPresets.nodeModules preset with an allowlist option to externalize installed packages, replacing the webpack-node-externals plugin. (by @​alexander-akait in #21569)

  • Add output.library.umdAmdContainer for an AMD-style loader branch in UMD. (by @​hai-x in #21770)

  • Resolve @custom-media values that are true / false or name another custom media. (by @​alexander-akait in #21624)

  • Add the __webpack_css_server_styles__ module variable to read the CSS collected while rendering without a DOM, and keep that CSS in the order the styles were applied. (by @​alexander-akait in #21576)

  • Patch the HTML <head> in place on hot update instead of forcing a full reload, including when a <script> that never executed is removed. (by @​alexander-akait in #21624)

  • Scope counter names in CSS modules; fix the counter() counter-style and animation timeline keywords. (by @​alexander-akait in #21600)

  • Derive import defer / import source from the target and fix the source phase. (by @​alexander-akait in #21810)

  • Emit analyzable ESM urls for chunks, assets, styles, workers and wasm. (by @​alexander-akait in #21788)

  • Tree shake CommonJS: module.exports object literals, exports destructured from a require() binding, unused method requires, and unused side-effect-free require() calls and reexports. (by @​alexander-akait in #21841)

  • Resolve relative entry baseUri values and bake one side of a hash cycle. (by @​alexander-akait in #21750)

  • Minify CSS further, only where the document is unchanged: shorthands and box longhands, font-weight, <position> and font-stretch keywords, colors (polar, Lab and hsl() converted to hex), numbers, times, zero units, calc() and every math function the spec names folded over constants, selector lists, An+B, keyframe selectors, media-feature ranges, unicode-range, transition layers, display, transforms, gradients, font families, identical repeated declarations, and rules an identical later one makes dead. Abilities are read off the target browsers, vendorPrefixes adds and drops vendor prefixes for them, and rewriteCustomProperties shortens custom property values. Minification never changes whether a declaration parses, and beautifying keeps every rule. (by @​alexander-akait in #21841)

  • Safely minify CSS (with source maps) and HTML assets when optimization.minimize is enabled, unless a minimizer is already configured for them, making only transformations an engine cannot tell apart. Every rewrite is named as an option, so it can be switched off. (by @​alexander-akait in #21841)

... (truncated)

Changelog

Sourced from webpack's changelog.

5.110.1

Patch Changes

5.110.0

Minor Changes

  • Wrap concatenated modules in lazy __webpack_require__.cw accessors and inline require(), keeping a wrapped body's names and side effects intact. (by @​hai-x in #21519)

  • Add performance hints reporting what a build costs: duplicate packages and modules, circular dependencies, broad contexts, large modules and chunks, hotspots, eval, missing PURE annotations, polyfills, redundant dynamic imports, OS-dependent rules, cache effectiveness, how chunks load, what splitting refused, why an optimization was skipped, and rules, defines, externals, aliases and barrel reexports nothing uses. An oversized asset names its largest modules, and an entrypoint carrying the runtime recommends optimization.runtimeChunk. Enable every check not set individually with performance.all, report hints in stats only with performance.hints: "stats", and get them in a stable order that leaves the build hashes unchanged. (by @​alexander-akait in #21841)

  • Add the descriptionRelativePath module rule condition. (by @​alexander-akait in #21705)

  • Add OS-independent glob matching to module rules. (by @​alexander-akait in #21771)

  • Report inner-graph, AMD and bare module bailouts in optimizationBailout. (by @​alexander-akait in #21740)

  • Allow marking externals as side-effect-free with a sideEffects flag. (by @​alexander-akait in #21712)

  • Give externals the original request of a context module element. (by @​alexander-akait in #21780)

  • Add the externalsPresets.nodeModules preset with an allowlist option to externalize installed packages, replacing the webpack-node-externals plugin. (by @​alexander-akait in #21569)

  • Add output.library.umdAmdContainer for an AMD-style loader branch in UMD. (by @​hai-x in #21770)

  • Resolve @custom-media values that are true / false or name another custom media. (by @​alexander-akait in #21624)

  • Add the __webpack_css_server_styles__ module variable to read the CSS collected while rendering without a DOM, and keep that CSS in the order the styles were applied. (by @​alexander-akait in #21576)

  • Patch the HTML <head> in place on hot update instead of forcing a full reload, including when a <script> that never executed is removed. (by @​alexander-akait in #21624)

  • Scope counter names in CSS modules; fix the counter() counter-style and animation timeline keywords. (by @​alexander-akait in #21600)

  • Derive import defer / import source from the target and fix the source phase. (by @​alexander-akait in #21810)

  • Emit analyzable ESM urls for chunks, assets, styles, workers and wasm. (by @​alexander-akait in #21788)

  • Tree shake CommonJS: module.exports object literals, exports destructured from a require() binding, unused method requires, and unused side-effect-free require() calls and reexports. (by @​alexander-akait in #21841)

  • Resolve relative entry baseUri values and bake one side of a hash cycle. (by @​alexander-akait in #21750)

  • Minify CSS further, only where the document is unchanged: shorthands and box longhands, font-weight, <position> and font-stretch keywords, colors (polar, Lab and hsl() converted to hex), numbers, times, zero units, calc() and every math function the spec names folded over constants, selector lists, An+B, keyframe selectors, media-feature ranges, unicode-range, transition layers, display, transforms, gradients, font families, identical repeated declarations, and rules an identical later one makes dead. Abilities are read off the target browsers, vendorPrefixes adds and drops vendor prefixes for them, and rewriteCustomProperties shortens custom property values. Minification never changes whether a declaration parses, and beautifying keeps every rule. (by @​alexander-akait in #21841)

... (truncated)

Commits
  • 0b2952e chore(release): new release (#21846)
  • 2a1fda4 fix: throw a SyntaxError from a module that failed to parse (#21847)
  • 281efa6 test(css): cover a loader-added BOM in the css pipeline (#21848)
  • 36ded2c fix: accept the optimization.minimize true shorthand after normalization (#21...
  • c150225 fix(cache): accept a managedPaths RegExp without a capture group, and speed u...
  • 3a7c0e6 chore(release): new release (#21545)
  • 75891c8 chore(changesets): union same-subject entries (#21841)
  • 6b96de7 feat(config): emit no development source map for library builds (#21839)
  • ff0c832 feat(css,html): reach embedded source from cssMinify and htmlMinify (#21838)
  • 06bedfb chore(deps): bump test/wpt in the dependencies group (#21836)
  • Additional commits viewable in compare view

Updates `webpack-cli` from 7.2.2 to 7.2.3
Release notes

Sourced from webpack-cli's releases.

webpack-cli@7.2.3

Patch Changes

  • fix: resolve the webpack-dev-server type from its default export, so the types work with both v5 and v6 (by @​bjohansebas in #4834)

  • feat: allow toml@5 as a peer dependency for TOML configuration files (by @​alexander-akait in #4837)

Changelog

Sourced from webpack-cli's changelog.

7.2.3

Patch Changes

  • fix: resolve the webpack-dev-server type from its default export, so the types work with both v5 and v6 (by @​bjohansebas in #4834)

  • feat: allow toml@5 as a peer dependency for TOML configuration files (by @​alexander-akait in #4837)

Commits
  • 7d40e4e chore(release): new release (#4839)
  • 1f6593a ci: use the input names changesets/action v2 expects (#4838)
  • 99cfc4f build(deps): update dependencies (#4837)
  • 11be634 feat(create-webpack-app): stop asking about HTML and CSS (#4836)
  • a2418aa feat(create-webpack-app): use webpack's native CSS and HTML support in init t...
  • 3664b9d chore: update webpack-dev-server to v6 and test against v5 and v6 (#4834)
  • ce1a219 ci: node 26 (#4763)
  • 0cfc077 chore(deps): bump changesets/action in the dependencies group (#4831)
  • ef262eb chore(deps): bump changesets/action in the dependencies group (#4830)
  • d90f5ab chore: add gitignore to ignore autogenerated build & fix codecov option (#4828)
  • Additional commits viewable in compare view

Dependabot will resolve any conflicts with this PR as long as you don't alter it yourself. You can also trigger a rebase manually by commenting `@dependabot rebase`. [//]: # (dependabot-automerge-start) [//]: # (dependabot-automerge-end) ---
Dependabot commands and options
You can trigger Dependabot actions by commenting on this PR: - `@dependabot rebase` will rebase this PR - `@dependabot recreate` will recreate this PR, overwriting any edits that have been made to it - `@dependabot show ignore conditions` will show all of the ignore conditions of the specified dependency - `@dependabot ignore major version` will close this group update PR and stop Dependabot creating any more for the specific dependency's major version (unless you unignore this specific dependency's major version or upgrade to it yourself) - `@dependabot ignore minor version` will close this group update PR and stop Dependabot creating any more for the specific dependency's minor version (unless you unignore this specific dependency's minor version or upgrade to it yourself) - `@dependabot ignore ` will close this group update PR and stop Dependabot creating any more for the specific dependency (unless you unignore this specific dependency or upgrade to it yourself) - `@dependabot unignore ` will remove all of the ignore conditions of the specified dependency - `@dependabot unignore ` will remove the ignore condition of the specified dependency and ignore conditions
Signed-off-by: dependabot[bot] Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> --- .../datafusion-wasm-app/package-lock.json | 186 ++++++------------ .../wasmtest/datafusion-wasm-app/package.json | 4 +- 2 files changed, 65 insertions(+), 125 deletions(-) diff --git a/datafusion/wasmtest/datafusion-wasm-app/package-lock.json b/datafusion/wasmtest/datafusion-wasm-app/package-lock.json index a85846fd8bb39..c0e1d2e6abf14 100644 --- a/datafusion/wasmtest/datafusion-wasm-app/package-lock.json +++ b/datafusion/wasmtest/datafusion-wasm-app/package-lock.json @@ -13,8 +13,8 @@ }, "devDependencies": { "copy-webpack-plugin": "14.0.0", - "webpack": "5.109.2", - "webpack-cli": "7.2.2", + "webpack": "5.110.1", + "webpack-cli": "7.2.3", "webpack-dev-server": "6.0.0" } }, @@ -36,6 +36,7 @@ "resolved": "https://registry.npmjs.org/@jridgewell/gen-mapping/-/gen-mapping-0.3.13.tgz", "integrity": "sha512-2kkt/7niJ6MgEPxF0bYdQ6etZaA+fQvDcLKckhy1yIQOzaoKjBBjSj63/aLVjYE3qhRt5dvM+uUyfCg6UKCBbA==", "dev": true, + "license": "MIT", "dependencies": { "@jridgewell/sourcemap-codec": "^1.5.0", "@jridgewell/trace-mapping": "^0.3.24" @@ -46,6 +47,7 @@ "resolved": "https://registry.npmjs.org/@jridgewell/resolve-uri/-/resolve-uri-3.1.2.tgz", "integrity": "sha512-bRISgCIjP20/tbWSPWMEi54QVPRZExkuD9lJL+UIxUKtwVJA8wW1Trb1jMs1RFXo1CBTNZ/5hpC9QvmKWdopKw==", "dev": true, + "license": "MIT", "engines": { "node": ">=6.0.0" } @@ -55,22 +57,25 @@ "resolved": "https://registry.npmjs.org/@jridgewell/source-map/-/source-map-0.3.11.tgz", "integrity": "sha512-ZMp1V8ZFcPG5dIWnQLr3NSI1MiCU7UETdS/A0G8V/XWHvJv3ZsFqutJn1Y5RPmAPX6F3BiE397OqveU/9NCuIA==", "dev": true, + "license": "MIT", "dependencies": { "@jridgewell/gen-mapping": "^0.3.5", "@jridgewell/trace-mapping": "^0.3.25" } }, "node_modules/@jridgewell/sourcemap-codec": { - "version": "1.5.5", - "resolved": "https://registry.npmjs.org/@jridgewell/sourcemap-codec/-/sourcemap-codec-1.5.5.tgz", - "integrity": "sha512-cYQ9310grqxueWbl+WuIUIaiUaDcj7WOq5fVhEljNVgRfOUhY9fy2zTvfoqWsnebh8Sl70VScFbICvJnLKB0Og==", - "dev": true + "version": "1.6.0", + "resolved": "https://registry.npmjs.org/@jridgewell/sourcemap-codec/-/sourcemap-codec-1.6.0.tgz", + "integrity": "sha512-T7jf+5zgsZHwNJ4lvQ7/aezbyk0nNX+zJVWpmHA7VYsEx7a7qr5Rg5IbtJFqkgze5Y2sruq1RUY8Q837Od7iFw==", + "dev": true, + "license": "MIT" }, "node_modules/@jridgewell/trace-mapping": { "version": "0.3.31", "resolved": "https://registry.npmjs.org/@jridgewell/trace-mapping/-/trace-mapping-0.3.31.tgz", "integrity": "sha512-zzNR+SdQSDJzc8joaeP8QQoCQr8NuYx2dIIytl1QeBEZHJ9uW6hebsrYgbz8hJwUQao3TWCMtmfV8Nu1twOLAw==", "dev": true, + "license": "MIT", "dependencies": { "@jridgewell/resolve-uri": "^3.1.0", "@jridgewell/sourcemap-codec": "^1.4.14" @@ -670,6 +675,7 @@ "resolved": "https://registry.npmjs.org/acorn/-/acorn-8.18.0.tgz", "integrity": "sha512-lGq+9yr1/GuAWaVYIHRjvvySG5/4VfKIvC8EWxStPdcDh/Ka7FG3twP6v4d5BkravUilhIAsG4Qj83t02LWUPQ==", "dev": true, + "license": "MIT", "bin": { "acorn": "bin/acorn" }, @@ -871,7 +877,8 @@ "version": "1.1.2", "resolved": "https://registry.npmjs.org/buffer-from/-/buffer-from-1.1.2.tgz", "integrity": "sha512-E+XQCRwSbaaiChtv6k6Dwgc+bx+Bs6vuKJHHl5kox/BaKbhiXzqQOwK4cO22yElGp2OCmjwVhT3HmxgyPGnJfQ==", - "dev": true + "dev": true, + "license": "MIT" }, "node_modules/bundle-name": { "version": "4.1.0", @@ -1001,7 +1008,8 @@ "version": "2.20.3", "resolved": "https://registry.npmjs.org/commander/-/commander-2.20.3.tgz", "integrity": "sha512-GpVkmM8vF2vQUkj2LvZmD35JxeJOLCwJ9cUkugyk2nuhbv3+mJvpLYYt+0+USMxE+oj+ey/lJEnhZw75x/OMcQ==", - "dev": true + "dev": true, + "license": "MIT" }, "node_modules/compressible": { "version": "2.0.18", @@ -1381,49 +1389,6 @@ "integrity": "sha512-NiSupZ4OeuGwr68lGIeym/ksIZMJodUGOSCZ/FSnTxcrekbvqrgdUxlJOMpijaKZVjAJrWrGs/6Jy8OMuyj9ow==", "dev": true }, - "node_modules/eslint-scope": { - "version": "5.1.1", - "resolved": "https://registry.npmjs.org/eslint-scope/-/eslint-scope-5.1.1.tgz", - "integrity": "sha512-2NxwbF/hZ0KpepYN0cNbo+FN6XoK7GaHlQhgx/hIZl6Va0bF45RQOOwhLIy8lQDbuCiadSLCBnH2CFYquit5bw==", - "dev": true, - "dependencies": { - "esrecurse": "^4.3.0", - "estraverse": "^4.1.1" - }, - "engines": { - "node": ">=8.0.0" - } - }, - "node_modules/esrecurse": { - "version": "4.3.0", - "resolved": "https://registry.npmjs.org/esrecurse/-/esrecurse-4.3.0.tgz", - "integrity": "sha512-KmfKL3b6G+RXvP8N1vr3Tq1kL/oCFgn2NYXEtqP8/L3pKapUA4G8cFVaoF3SU323CD4XypR/ffioHmkti6/Tag==", - "dev": true, - "dependencies": { - "estraverse": "^5.2.0" - }, - "engines": { - "node": ">=4.0" - } - }, - "node_modules/esrecurse/node_modules/estraverse": { - "version": "5.3.0", - "resolved": "https://registry.npmjs.org/estraverse/-/estraverse-5.3.0.tgz", - "integrity": "sha512-MMdARuVEQziNTeJD8DgMqmhwR11BRQ/cBP+pLtYdSTnf3MIO8fFeiINEbX36ZdNlfU/7A9f3gUw49B3oQsvwBA==", - "dev": true, - "engines": { - "node": ">=4.0" - } - }, - "node_modules/estraverse": { - "version": "4.3.0", - "resolved": "https://registry.npmjs.org/estraverse/-/estraverse-4.3.0.tgz", - "integrity": "sha512-39nnKffWz8xN1BU/2c79n9nB9HDzo0niYUqx6xyqUnyoAnQyyWpOTdZEeiCch8BBu515t4wp9ZmgVfVhn9EBpw==", - "dev": true, - "engines": { - "node": ">=4.0" - } - }, "node_modules/etag": { "version": "1.8.1", "resolved": "https://registry.npmjs.org/etag/-/etag-1.8.1.tgz", @@ -1686,6 +1651,7 @@ "resolved": "https://registry.npmjs.org/has-flag/-/has-flag-4.0.0.tgz", "integrity": "sha512-EykJT/Q1KjTWctppgIAgfSO0tKVuZUjhgMr17kqTumMl6Afv3EISleU7qZUzoXDFTAHTDC4NOoG/ZxU3EvlMPQ==", "dev": true, + "license": "MIT", "engines": { "node": ">=8" } @@ -1988,6 +1954,7 @@ "resolved": "https://registry.npmjs.org/jest-worker/-/jest-worker-27.5.1.tgz", "integrity": "sha512-7vuh85V5cdDofPyxn58nrPjBktZo0u9x1g8WtjQol+jZDaE+fhN+cIvTj11GndBnMnyfrUOG1sZQxCdjKh+DKg==", "dev": true, + "license": "MIT", "dependencies": { "@types/node": "*", "merge-stream": "^2.0.0", @@ -2565,7 +2532,8 @@ "version": "2.0.0", "resolved": "https://registry.npmjs.org/merge-stream/-/merge-stream-2.0.0.tgz", "integrity": "sha512-abv/qOcuPfk3URPfDzmZU1LKmuw8kT+0nIHvKrKgFrwifol/doWcdA4ZqsWQ8ENrFKkd67Mfpo/LovbIUsbt3w==", - "dev": true + "dev": true, + "license": "MIT" }, "node_modules/micromatch": { "version": "4.0.8", @@ -2602,15 +2570,16 @@ } }, "node_modules/minimizer-webpack-plugin": { - "version": "5.6.1", - "resolved": "https://registry.npmjs.org/minimizer-webpack-plugin/-/minimizer-webpack-plugin-5.6.1.tgz", - "integrity": "sha512-DoeAZz8Q1C1znwsUzej1fdoi4jCf7/+Em27ouLqfK/+3m8G+D7yDhUwrc3CNhjSzGUN1kn7Iv4sWmjflQHenpw==", + "version": "5.8.0", + "resolved": "https://registry.npmjs.org/minimizer-webpack-plugin/-/minimizer-webpack-plugin-5.8.0.tgz", + "integrity": "sha512-2cT9+goJfBhtMz+gJqejSf09ClgmYhPhccRb/fb0ztVbixt0BkO8mRuI26FoPPuUNkRk6iEDryWAIouc5W8eJA==", "dev": true, + "license": "MIT", "dependencies": { - "@jridgewell/trace-mapping": "^0.3.25", + "@jridgewell/trace-mapping": "^0.3.31", "jest-worker": "^27.4.5", - "schema-utils": "^4.3.0", - "terser": "^5.31.1" + "schema-utils": "^4.3.3", + "terser": "^5.51.0" }, "engines": { "node": ">= 10.13.0" @@ -3461,6 +3430,7 @@ "resolved": "https://registry.npmjs.org/source-map/-/source-map-0.6.1.tgz", "integrity": "sha512-UjgapumWlbMhkBgzT7Ykc5YXUT46F0iKu8SGXq0bcwP5dz/h0Plj6enJqjz1Zbq2l5WaqYnrVbwWOWMyF3F47g==", "dev": true, + "license": "BSD-3-Clause", "engines": { "node": ">=0.10.0" } @@ -3470,6 +3440,7 @@ "resolved": "https://registry.npmjs.org/source-map-support/-/source-map-support-0.5.21.tgz", "integrity": "sha512-uBHU3L3czsIyYXKX88fdrGovxdSCoTGDRZ6SYXtSRxLZUzHg5P/66Ht6uoUlHu9EZod+inXhKo3qQgwXUT/y1w==", "dev": true, + "license": "MIT", "dependencies": { "buffer-from": "^1.0.0", "source-map": "^0.6.0" @@ -3489,6 +3460,7 @@ "resolved": "https://registry.npmjs.org/supports-color/-/supports-color-8.1.1.tgz", "integrity": "sha512-MpUEN2OodtUzxvKQl72cUF7RQ5EiHsGvSsVG0ia9c5RbWGL2CI4C7EpPS8UTBIplnlzZiNuV56w+FuNxy3ty2Q==", "dev": true, + "license": "MIT", "dependencies": { "has-flag": "^4.0.0" }, @@ -3525,10 +3497,11 @@ } }, "node_modules/terser": { - "version": "5.50.0", - "resolved": "https://registry.npmjs.org/terser/-/terser-5.50.0.tgz", - "integrity": "sha512-CN9BVxWhgS/hRxtUMjtC2uRWSTcSfQFHMDWma6sKKfIivCD91sM+FOPfvwoaRMqCSrUpe1nv3jDamd9eEQ4y+w==", + "version": "5.51.2", + "resolved": "https://registry.npmjs.org/terser/-/terser-5.51.2.tgz", + "integrity": "sha512-bWnjSNscmuI+GJze6ZupnHP8G/cTcsJF+bXCeQknk2SHQsgbNJnLrqiH9jZ2W4STPVXH2mDKKRX3iwPhc9Cn/Q==", "dev": true, + "license": "BSD-2-Clause", "dependencies": { "@jridgewell/source-map": "^0.3.3", "acorn": "^8.15.0", @@ -3751,10 +3724,11 @@ } }, "node_modules/webpack": { - "version": "5.109.2", - "resolved": "https://registry.npmjs.org/webpack/-/webpack-5.109.2.tgz", - "integrity": "sha512-U9/cvLzxObKNEZ9+TtdqrHM5/9z3lgl2c+c4BzbqGxFQvQvBAq87yql5A8pQ+rrMbS496MZJeF5enVBndIy2hw==", + "version": "5.110.1", + "resolved": "https://registry.npmjs.org/webpack/-/webpack-5.110.1.tgz", + "integrity": "sha512-gInQB+jxXxgnZyvPwuzT5NGQmECDqeu85oxcrjinrYHqPoBex0hCAN2SFTJVyPVrK0Pq9E44VFP+e89fAc10/w==", "dev": true, + "license": "MIT", "dependencies": { "@types/estree": "^1.0.8", "@types/json-schema": "^7.0.15", @@ -3766,11 +3740,10 @@ "chrome-trace-event": "^1.0.2", "enhanced-resolve": "^5.24.4", "es-module-lexer": "^2.1.0", - "eslint-scope": "5.1.1", "events": "^3.2.0", "graceful-fs": "^4.2.11", "mime-db": "^1.54.0", - "minimizer-webpack-plugin": "^5.6.1", + "minimizer-webpack-plugin": "^5.7.0", "neo-async": "^2.6.2", "schema-utils": "^4.3.3", "tapable": "^2.3.0", @@ -3794,10 +3767,11 @@ } }, "node_modules/webpack-cli": { - "version": "7.2.2", - "resolved": "https://registry.npmjs.org/webpack-cli/-/webpack-cli-7.2.2.tgz", - "integrity": "sha512-lD0pALneslq8FfV+rwvm1BMW0AFAJrHHhNupAGN4asYjMvqrtRsenU4iKpiBo09gS4ntMxKGUxl9jhTEzVt0oA==", + "version": "7.2.3", + "resolved": "https://registry.npmjs.org/webpack-cli/-/webpack-cli-7.2.3.tgz", + "integrity": "sha512-vDFU7jrfCctnN7jJQWPl+V26B51GLp11prVZXg50oeonsgeBzSTJEWmXkjHsOjTgMjlPOQM8GWLh33X9RN/0ow==", "dev": true, + "license": "MIT", "dependencies": { "@discoveryjs/json-ext": "^1.1.0", "commander": "^14.0.3", @@ -3821,7 +3795,7 @@ "peerDependencies": { "js-yaml": "^4.0.0 || ^5.0.0", "json5": "^2.2.3", - "toml": "^3.0.0 || ^4.0.0", + "toml": "^3.0.0 || ^4.0.0 || ^5.0.0", "webpack": "^5.101.0", "webpack-bundle-analyzer": "^4.0.0 || ^5.0.0", "webpack-dev-server": "^5.0.0 || ^6.0.0" @@ -4090,9 +4064,9 @@ } }, "@jridgewell/sourcemap-codec": { - "version": "1.5.5", - "resolved": "https://registry.npmjs.org/@jridgewell/sourcemap-codec/-/sourcemap-codec-1.5.5.tgz", - "integrity": "sha512-cYQ9310grqxueWbl+WuIUIaiUaDcj7WOq5fVhEljNVgRfOUhY9fy2zTvfoqWsnebh8Sl70VScFbICvJnLKB0Og==", + "version": "1.6.0", + "resolved": "https://registry.npmjs.org/@jridgewell/sourcemap-codec/-/sourcemap-codec-1.6.0.tgz", + "integrity": "sha512-T7jf+5zgsZHwNJ4lvQ7/aezbyk0nNX+zJVWpmHA7VYsEx7a7qr5Rg5IbtJFqkgze5Y2sruq1RUY8Q837Od7iFw==", "dev": true }, "@jridgewell/trace-mapping": { @@ -5170,39 +5144,6 @@ "integrity": "sha512-NiSupZ4OeuGwr68lGIeym/ksIZMJodUGOSCZ/FSnTxcrekbvqrgdUxlJOMpijaKZVjAJrWrGs/6Jy8OMuyj9ow==", "dev": true }, - "eslint-scope": { - "version": "5.1.1", - "resolved": "https://registry.npmjs.org/eslint-scope/-/eslint-scope-5.1.1.tgz", - "integrity": "sha512-2NxwbF/hZ0KpepYN0cNbo+FN6XoK7GaHlQhgx/hIZl6Va0bF45RQOOwhLIy8lQDbuCiadSLCBnH2CFYquit5bw==", - "dev": true, - "requires": { - "esrecurse": "^4.3.0", - "estraverse": "^4.1.1" - } - }, - "esrecurse": { - "version": "4.3.0", - "resolved": "https://registry.npmjs.org/esrecurse/-/esrecurse-4.3.0.tgz", - "integrity": "sha512-KmfKL3b6G+RXvP8N1vr3Tq1kL/oCFgn2NYXEtqP8/L3pKapUA4G8cFVaoF3SU323CD4XypR/ffioHmkti6/Tag==", - "dev": true, - "requires": { - "estraverse": "^5.2.0" - }, - "dependencies": { - "estraverse": { - "version": "5.3.0", - "resolved": "https://registry.npmjs.org/estraverse/-/estraverse-5.3.0.tgz", - "integrity": "sha512-MMdARuVEQziNTeJD8DgMqmhwR11BRQ/cBP+pLtYdSTnf3MIO8fFeiINEbX36ZdNlfU/7A9f3gUw49B3oQsvwBA==", - "dev": true - } - } - }, - "estraverse": { - "version": "4.3.0", - "resolved": "https://registry.npmjs.org/estraverse/-/estraverse-4.3.0.tgz", - "integrity": "sha512-39nnKffWz8xN1BU/2c79n9nB9HDzo0niYUqx6xyqUnyoAnQyyWpOTdZEeiCch8BBu515t4wp9ZmgVfVhn9EBpw==", - "dev": true - }, "etag": { "version": "1.8.1", "resolved": "https://registry.npmjs.org/etag/-/etag-1.8.1.tgz", @@ -5938,15 +5879,15 @@ } }, "minimizer-webpack-plugin": { - "version": "5.6.1", - "resolved": "https://registry.npmjs.org/minimizer-webpack-plugin/-/minimizer-webpack-plugin-5.6.1.tgz", - "integrity": "sha512-DoeAZz8Q1C1znwsUzej1fdoi4jCf7/+Em27ouLqfK/+3m8G+D7yDhUwrc3CNhjSzGUN1kn7Iv4sWmjflQHenpw==", + "version": "5.8.0", + "resolved": "https://registry.npmjs.org/minimizer-webpack-plugin/-/minimizer-webpack-plugin-5.8.0.tgz", + "integrity": "sha512-2cT9+goJfBhtMz+gJqejSf09ClgmYhPhccRb/fb0ztVbixt0BkO8mRuI26FoPPuUNkRk6iEDryWAIouc5W8eJA==", "dev": true, "requires": { - "@jridgewell/trace-mapping": "^0.3.25", + "@jridgewell/trace-mapping": "^0.3.31", "jest-worker": "^27.4.5", - "schema-utils": "^4.3.0", - "terser": "^5.31.1" + "schema-utils": "^4.3.3", + "terser": "^5.51.0" } }, "ms": { @@ -6563,9 +6504,9 @@ "dev": true }, "terser": { - "version": "5.50.0", - "resolved": "https://registry.npmjs.org/terser/-/terser-5.50.0.tgz", - "integrity": "sha512-CN9BVxWhgS/hRxtUMjtC2uRWSTcSfQFHMDWma6sKKfIivCD91sM+FOPfvwoaRMqCSrUpe1nv3jDamd9eEQ4y+w==", + "version": "5.51.2", + "resolved": "https://registry.npmjs.org/terser/-/terser-5.51.2.tgz", + "integrity": "sha512-bWnjSNscmuI+GJze6ZupnHP8G/cTcsJF+bXCeQknk2SHQsgbNJnLrqiH9jZ2W4STPVXH2mDKKRX3iwPhc9Cn/Q==", "dev": true, "requires": { "@jridgewell/source-map": "^0.3.3", @@ -6701,9 +6642,9 @@ } }, "webpack": { - "version": "5.109.2", - "resolved": "https://registry.npmjs.org/webpack/-/webpack-5.109.2.tgz", - "integrity": "sha512-U9/cvLzxObKNEZ9+TtdqrHM5/9z3lgl2c+c4BzbqGxFQvQvBAq87yql5A8pQ+rrMbS496MZJeF5enVBndIy2hw==", + "version": "5.110.1", + "resolved": "https://registry.npmjs.org/webpack/-/webpack-5.110.1.tgz", + "integrity": "sha512-gInQB+jxXxgnZyvPwuzT5NGQmECDqeu85oxcrjinrYHqPoBex0hCAN2SFTJVyPVrK0Pq9E44VFP+e89fAc10/w==", "dev": true, "requires": { "@types/estree": "^1.0.8", @@ -6716,11 +6657,10 @@ "chrome-trace-event": "^1.0.2", "enhanced-resolve": "^5.24.4", "es-module-lexer": "^2.1.0", - "eslint-scope": "5.1.1", "events": "^3.2.0", "graceful-fs": "^4.2.11", "mime-db": "^1.54.0", - "minimizer-webpack-plugin": "^5.6.1", + "minimizer-webpack-plugin": "^5.7.0", "neo-async": "^2.6.2", "schema-utils": "^4.3.3", "tapable": "^2.3.0", @@ -6737,9 +6677,9 @@ } }, "webpack-cli": { - "version": "7.2.2", - "resolved": "https://registry.npmjs.org/webpack-cli/-/webpack-cli-7.2.2.tgz", - "integrity": "sha512-lD0pALneslq8FfV+rwvm1BMW0AFAJrHHhNupAGN4asYjMvqrtRsenU4iKpiBo09gS4ntMxKGUxl9jhTEzVt0oA==", + "version": "7.2.3", + "resolved": "https://registry.npmjs.org/webpack-cli/-/webpack-cli-7.2.3.tgz", + "integrity": "sha512-vDFU7jrfCctnN7jJQWPl+V26B51GLp11prVZXg50oeonsgeBzSTJEWmXkjHsOjTgMjlPOQM8GWLh33X9RN/0ow==", "dev": true, "requires": { "@discoveryjs/json-ext": "^1.1.0", diff --git a/datafusion/wasmtest/datafusion-wasm-app/package.json b/datafusion/wasmtest/datafusion-wasm-app/package.json index c5723b079b6d0..ab03cbf969b81 100644 --- a/datafusion/wasmtest/datafusion-wasm-app/package.json +++ b/datafusion/wasmtest/datafusion-wasm-app/package.json @@ -27,8 +27,8 @@ "datafusion-wasmtest": "../pkg" }, "devDependencies": { - "webpack": "5.109.2", - "webpack-cli": "7.2.2", + "webpack": "5.110.1", + "webpack-cli": "7.2.3", "webpack-dev-server": "6.0.0", "copy-webpack-plugin": "14.0.0" } From 1608ee70faa06063a38f01e8c57c1a6879708cd6 Mon Sep 17 00:00:00 2001 From: Yongting You <2010youy01@gmail.com> Date: Wed, 2 Sep 2026 01:08:05 +0000 Subject: [PATCH 09/37] refactor(hash-aggr): Let partial-reduce mode aggregation early emit when OOM (#24785) ## Which issue does this PR close? Part of https://github.com/apache/datafusion/issues/22710 ## Rationale for this change See https://github.com/apache/datafusion/pull/24486 for the rationale. That discussion covers the issue and three potential solutions: error, early-emit, and spill. This PR implements the early-emit behavior agreed on there. ## What changes are included in this PR? Key changes: In datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs - top comment explains the high-level ideas - To understand the implementation, start from `poll_next()` and follow along ## Are these changes tested? UTs (note partial-reduce aggregation can't be planned from SQL, so we can't do sqllogictests here) ## Are there any user-facing changes? No --- .../physical-plan/src/aggregates/mod.rs | 185 ++++++++++++-- .../src/aggregates/partial_reduce_stream.rs | 235 +++++++++++++++--- 2 files changed, 359 insertions(+), 61 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index f00f1a160e27a..6da7ee1018dc5 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -187,7 +187,6 @@ use datafusion_common::{ assert_eq_or_internal_err, internal_err, not_impl_err, }; use datafusion_execution::TaskContext; -use datafusion_execution::memory_pool::MemoryLimit; use datafusion_expr::{Accumulator, Aggregate}; use datafusion_physical_expr::aggregate::AggregateFunctionExpr; use datafusion_physical_expr::equivalence::ProjectionMapping; @@ -1308,12 +1307,7 @@ impl AggregateExec { && self.group_by.is_single() } - fn should_use_partial_reduce_hash_stream(&self, context: &TaskContext) -> bool { - // TODO: implement memory-limited path and remove this limitation - if matches!(context.memory_pool().memory_limit(), MemoryLimit::Finite(_)) { - return false; - } - + fn should_use_partial_reduce_hash_stream(&self, _context: &TaskContext) -> bool { self.mode == AggregateMode::PartialReduce && self.limit_options.is_none() && self.input_order_mode == InputOrderMode::Linear @@ -3241,7 +3235,7 @@ mod tests { use arrow::array::{ BooleanArray, DictionaryArray, Float32Array, Float64Array, Int32Array, - Int64Array, StructArray, UInt32Array, UInt64Array, + Int64Array, NullArray, StructArray, UInt32Array, UInt64Array, }; use arrow::compute::{SortOptions, concat_batches}; use arrow::datatypes::Int32Type; @@ -4586,6 +4580,15 @@ mod tests { } fn partial_reduce_test_aggregate() -> Result { + partial_reduce_test_aggregate_with_batches(1) + } + + /// Partial-reduce aggregate over `num_input_batches` identical input batches + /// of partial states, each reducing to groups `1, 2, 3` with sums + /// `50, 20, 30`. + fn partial_reduce_test_aggregate_with_batches( + num_input_batches: usize, + ) -> Result { let schema = Arc::new(Schema::new(vec![ Field::new("a", DataType::UInt32, false), Field::new("b", DataType::Float64, false), @@ -4618,7 +4621,7 @@ mod tests { ], )?; let partial_reduce_input = TestMemoryExec::try_new_exec( - &[vec![partial_state_batch]], + &[vec![partial_state_batch; num_input_batches]], Arc::clone(&partial_schema), None, )?; @@ -4654,26 +4657,162 @@ mod tests { Ok(()) } - /// Spilling behavior is not implemented for partial-reduce stream yet, so fall - /// back to the existing `GroupedHashAggregateStream` + /// Partial-reduce hash aggregation emits its accumulated partial states early + /// under memory pressure instead of failing, and the early-emitted states + /// still merge into the correct result. #[tokio::test] - async fn partial_reduce_aggregate_with_memory_limit_planning() -> Result<()> { - let partial_reduce = partial_reduce_test_aggregate()?; + async fn partial_reduce_aggregate_with_memory_limit_emits_early() -> Result<()> { + let num_input_batches = 3; + let partial_reduce = + partial_reduce_test_aggregate_with_batches(num_input_batches)?; let runtime = RuntimeEnvBuilder::new() .with_memory_limit(1, 1.0) .build_arc()?; - let task_ctx = - Arc::new( - TaskContext::default() - .with_session_config(SessionConfig::new().set_bool( - "datafusion.execution.enable_migration_aggregate", - true, - )) - .with_runtime(runtime), - ); + // A batch size smaller than the number of flushed groups also covers + // splitting one flush across several output batches. + let batch_size = 2; + let task_ctx = Arc::new( + TaskContext::default() + .with_session_config(migrated_hash_session_config(batch_size)) + .with_runtime(runtime), + ); let stream = partial_reduce.execute_typed(0, &task_ctx)?; - assert!(matches!(stream, StreamType::GroupedHash(_))); + assert!(matches!(stream, StreamType::PartialReduceHash(_))); + let stream: SendableRecordBatchStream = stream.into(); + let output = collect(stream).await?; + + // The table is flushed after every input batch, so each of the three + // groups is emitted once per input batch instead of being merged into a + // single row. Each flush is sliced into batches of 2 and 1 rows. + assert_eq!(output.len(), 2 * num_input_batches); + assert_snapshot!(batches_to_string(&output), @r" + +---+-------------+ + | a | SUM(b)[sum] | + +---+-------------+ + | 1 | 50.0 | + | 2 | 20.0 | + | 3 | 30.0 | + | 1 | 50.0 | + | 2 | 20.0 | + | 3 | 30.0 | + | 1 | 50.0 | + | 2 | 20.0 | + | 3 | 30.0 | + +---+-------------+ + "); + + Ok(()) + } + + /// Same shape as [`partial_reduce_test_aggregate_with_batches`], but with multiple + /// group keys. + fn partial_reduce_test_aggregate_rows_multi_group_keys( + num_input_batches: usize, + ) -> Result { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::UInt32, false), + Field::new("n", DataType::Null, true), + Field::new("b", DataType::Float64, false), + ])); + let group_by = PhysicalGroupBy::new_single(vec![ + (col("a", &schema)?, "a".to_string()), + (col("n", &schema)?, "n".to_string()), + ]); + let aggregates: Vec> = vec![Arc::new( + AggregateExprBuilder::new(sum_udaf(), vec![col("b", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("SUM(b)") + .build()?, + )]; + + let empty_input = + TestMemoryExec::try_new_exec(&[vec![]], Arc::clone(&schema), None)?; + let partial = AggregateExec::try_new( + AggregateMode::Partial, + group_by.clone(), + aggregates.clone(), + vec![None], + empty_input, + Arc::clone(&schema), + )?; + let partial_schema = partial.schema(); + let partial_state_batch = RecordBatch::try_new( + Arc::clone(&partial_schema), + vec![ + Arc::new(UInt32Array::from(vec![1, 2, 1, 3])), + Arc::new(NullArray::new(4)), + Arc::new(Float64Array::from(vec![10.0, 20.0, 40.0, 30.0])), + ], + )?; + let partial_reduce_input = TestMemoryExec::try_new_exec( + &[vec![partial_state_batch; num_input_batches]], + Arc::clone(&partial_schema), + None, + )?; + + AggregateExec::try_new( + AggregateMode::PartialReduce, + group_by, + aggregates, + vec![None], + partial_reduce_input, + partial_schema, + ) + } + + #[tokio::test] + async fn partial_reduce_aggregate_with_memory_limit_emits_early_multi_group_keys() + -> Result<()> { + let num_input_batches = 3; + let partial_reduce = + partial_reduce_test_aggregate_rows_multi_group_keys(num_input_batches)?; + + // Pin the representation: this is exactly the condition + // `new_group_values` uses to pick `GroupValuesRows` over + // `GroupValuesColumn`. If a `Null` `GroupColumn` is ever added, this + // assertion fires and the test stops covering the row-encoded path. + let group_schema = partial_reduce + .group_by + .group_schema(&partial_reduce.schema())?; + assert!( + !group_values::multi_group_by::supported_schema(&group_schema), + "expected the Null group column to force the GroupValuesRows fallback" + ); + + let runtime = RuntimeEnvBuilder::new() + .with_memory_limit(1, 1.0) + .build_arc()?; + let batch_size = 2; + let task_ctx = Arc::new( + TaskContext::default() + .with_session_config(migrated_hash_session_config(batch_size)) + .with_runtime(runtime), + ); + + let stream = partial_reduce.execute_typed(0, &task_ctx)?; + assert!(matches!(stream, StreamType::PartialReduceHash(_))); + let stream: SendableRecordBatchStream = stream.into(); + let output = collect(stream).await?; + + // Same flush cadence as the column-backed test: one flush per input + // batch, each sliced into batches of 2 and 1 rows. + assert_eq!(output.len(), 2 * num_input_batches); + assert_snapshot!(batches_to_string(&output), @r" + +---+---+-------------+ + | a | n | SUM(b)[sum] | + +---+---+-------------+ + | 1 | | 50.0 | + | 2 | | 20.0 | + | 3 | | 30.0 | + | 1 | | 50.0 | + | 2 | | 20.0 | + | 3 | | 30.0 | + | 1 | | 50.0 | + | 2 | | 20.0 | + | 3 | | 30.0 | + +---+---+-------------+ + "); Ok(()) } diff --git a/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs b/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs index 9eedf868477a0..d8f1447bc9521 100644 --- a/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs +++ b/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs @@ -28,7 +28,7 @@ use std::task::{Context, Poll}; use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; -use datafusion_common::Result; +use datafusion_common::{DataFusionError, Result}; use datafusion_execution::TaskContext; use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; use futures::stream::{Stream, StreamExt}; @@ -68,6 +68,15 @@ use crate::{InputOrderMode, RecordBatchStream, SendableRecordBatchStream}; /// This stage is useful for tree-reduce plans. It consumes the same schema as /// a final aggregate stage, but emits the same schema as a partial aggregate /// stage. +/// +/// # Memory Management +/// +/// If the memory reservation cannot grow after aggregating an input batch, all +/// accumulated partial states are emitted immediately, and the remaining input +/// is aggregated with an empty table. This repeats until the input ends. +/// +/// See [`crate::aggregates::AggregateMode::PartialReduce`] for why it's allowed +/// to emit the same group multiple times. pub(crate) struct PartialReduceHashAggregateStream { /// Output schema: group columns followed by partial aggregate state columns. schema: SchemaRef, @@ -75,6 +84,9 @@ pub(crate) struct PartialReduceHashAggregateStream { /// Input batches containing partial aggregate state rows. input: SendableRecordBatchStream, + /// Target output batch size from configuration. + batch_size: usize, + /// Execution metrics shared with the aggregate plan node. baseline_metrics: BaselineMetrics, @@ -93,10 +105,22 @@ enum PartialReduceHashAggregateState { ReadingInput { hash_table: AggregateHashTable, }, + /// A fully materialized partial-state batch being emitted incrementally + /// because the table ran out of memory while reading input. + EmittingOnMemoryPressure { + hash_table: AggregateHashTable, + // After each incremental emitting step, `remaining_groups` is updated + // with batch slicing. + remaining_groups: RecordBatch, + }, ProducingOutput { hash_table: AggregateHashTable, }, Done, + /// Sentinel state to use when returning error from any other states, because: + /// - It explicitly releases state-owned resources immediately + /// - More defensive against accidentally resuming execution after error + Error, } type PartialReduceHashAggregatePoll = Poll>>; @@ -111,28 +135,34 @@ type PartialReduceHashAggregateStateTransition = ControlFlow< impl PartialReduceHashAggregateState { fn hash_table(&self) -> &AggregateHashTable { match self { - Self::ReadingInput { hash_table } | Self::ProducingOutput { hash_table } => { - hash_table + Self::ReadingInput { hash_table } + | Self::EmittingOnMemoryPressure { hash_table, .. } + | Self::ProducingOutput { hash_table } => hash_table, + Self::Done | Self::Error => { + unreachable!("Done and Error states do not hold a hash table") } - Self::Done => unreachable!("Done state does not hold a hash table"), } } fn hash_table_mut(&mut self) -> &mut AggregateHashTable { match self { - Self::ReadingInput { hash_table } | Self::ProducingOutput { hash_table } => { - hash_table + Self::ReadingInput { hash_table } + | Self::EmittingOnMemoryPressure { hash_table, .. } + | Self::ProducingOutput { hash_table } => hash_table, + Self::Done | Self::Error => { + unreachable!("Done and Error states do not hold a hash table") } - Self::Done => unreachable!("Done state does not hold a hash table"), } } fn into_hash_table(self) -> AggregateHashTable { match self { - Self::ReadingInput { hash_table } | Self::ProducingOutput { hash_table } => { - hash_table + Self::ReadingInput { hash_table } + | Self::EmittingOnMemoryPressure { hash_table, .. } + | Self::ProducingOutput { hash_table } => hash_table, + Self::Done | Self::Error => { + unreachable!("Done and Error states do not hold a hash table") } - Self::Done => unreachable!("Done state does not hold a hash table"), } } @@ -178,19 +208,25 @@ impl PartialReduceHashAggregateStream { Ok(Self { schema, input, + batch_size, baseline_metrics, reservation, state: Some(PartialReduceHashAggregateState::ReadingInput { hash_table }), }) } - fn start_output( - &mut self, - hash_table: &mut AggregateHashTable, - ) -> Result<()> { + fn close_input(&mut self) { let input_schema = self.input.schema(); self.input = Box::pin(EmptyRecordBatchStream::new(input_schema)); - hash_table.start_output() + } + + fn break_with_err( + error: DataFusionError, + ) -> PartialReduceHashAggregateStateTransition { + ControlFlow::Break(( + Poll::Ready(Some(Err(error))), + PartialReduceHashAggregateState::Error, + )) } /// Handle ReadingInput state - aggregate partial state batches into the hash table. @@ -219,46 +255,126 @@ impl PartialReduceHashAggregateStream { timer.done(); if let Err(e) = result { - return ControlFlow::Break(( - Poll::Ready(Some(Err(e))), - original_state, - )); + return Self::break_with_err(e); } - if let Err(e) = self - .reservation - .try_resize(original_state.hash_table().memory_size()) - { - return ControlFlow::Break(( - Poll::Ready(Some(Err(e))), - original_state, - )); - } - - ControlFlow::Continue(original_state) - } - Poll::Ready(Some(Err(e))) => { - ControlFlow::Break((Poll::Ready(Some(Err(e))), original_state)) + // Update the memory reservation. If OOM, do early emit. + self.resize_or_emit_early(original_state) } + Poll::Ready(Some(Err(e))) => Self::break_with_err(e), // Input ends, move to output state Poll::Ready(None) => { + self.close_input(); let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); let timer = elapsed_compute.timer(); - let result = self.start_output(original_state.hash_table_mut()); + let result = original_state.hash_table_mut().start_output(); timer.done(); match result { Ok(()) => { ControlFlow::Continue(original_state.into_producing_output()) } - Err(e) => { - ControlFlow::Break((Poll::Ready(Some(Err(e))), original_state)) - } + Err(e) => Self::break_with_err(e), } } } } + /// Update the memory reservation. If the reservation succeeds, continue reading + /// input. If OOM, clear the aggregated states in the hash table, and early emit + /// them immediately. + /// + /// Returns the next state; the caller finishes the intended task based on it. + /// + /// The reservation is left at its pre-emission size while the states are being + /// emitted, because the cleared states are still held in memory as + /// `remaining_groups`. The reservation will be reset after exiting the + /// `EmittingOnMemoryPressure` state. + /// + /// # Implementation Note + /// All accumulated states are materialized at once, and then sliced into + /// `batch_size` output batches. Emit them incrementally after blocked state + /// management is ready. + /// + /// Issue: + fn resize_or_emit_early( + &mut self, + mut original_state: PartialReduceHashAggregateState, + ) -> PartialReduceHashAggregateStateTransition { + let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); + let _timer = elapsed_compute.timer(); // Stop on drop + let resize_result = self + .reservation + .try_resize(original_state.hash_table().memory_size()); + + let oom = match resize_result { + Ok(()) => return ControlFlow::Continue(original_state), + Err(e @ DataFusionError::ResourcesExhausted(_)) => e, + Err(e) => return Self::break_with_err(e), + }; + + let state_batch_result = original_state.hash_table_mut().take_state_batch(); + + match state_batch_result { + Ok(Some(remaining_groups)) => ControlFlow::Continue( + PartialReduceHashAggregateState::EmittingOnMemoryPressure { + hash_table: original_state.into_hash_table(), + remaining_groups, + }, + ), + // No accumulated group to emit, so early emission cannot release any + // memory: report the original error. + Ok(None) => Self::break_with_err(oom), + Err(e) => Self::break_with_err(e), + } + } + + /// Handle EmittingOnMemoryPressure state - emit a materialized partial-state + /// batch in `batch_size`(from configuration) slices. After all slices are + /// emitted, update the memory reservation and resume reading input. + /// + /// See comments at `poll_next()` for details. + /// + /// Returns the next operator state with control flow decision. + fn handle_emitting_on_memory_pressure( + &mut self, + original_state: PartialReduceHashAggregateState, + ) -> PartialReduceHashAggregateStateTransition { + let PartialReduceHashAggregateState::EmittingOnMemoryPressure { + hash_table, + remaining_groups: batch, + } = original_state + else { + unreachable!("expected the EmittingOnMemoryPressure state") + }; + + let (output_batch, next_state) = if batch.num_rows() <= self.batch_size { + // Go back to `ReadingInput` + ( + batch, + PartialReduceHashAggregateState::ReadingInput { hash_table }, + ) + } else { + // More batches to output, continue in the current state. + let remaining = + batch.slice(self.batch_size, batch.num_rows() - self.batch_size); + let output = batch.slice(0, self.batch_size); + ( + output, + PartialReduceHashAggregateState::EmittingOnMemoryPressure { + hash_table, + remaining_groups: remaining, + }, + ) + }; + + debug_assert!(output_batch.num_rows() > 0); + ControlFlow::Break(( + Poll::Ready(Some(Ok(output_batch.record_output(&self.baseline_metrics)))), + next_state, + )) + } + /// Handle ProducingOutput state - emit merged partial aggregate state batches. /// /// See comments at `poll_next()` for details. @@ -281,6 +397,8 @@ impl PartialReduceHashAggregateStream { match result { Ok(Some(batch)) => { + // The output is already materialized, so a failed resize cannot + // be acted on: keep the reservation as is and finish the output. let _ = self .reservation .try_resize(original_state.hash_table().memory_size()); @@ -300,7 +418,7 @@ impl PartialReduceHashAggregateStream { let _ = self.reservation.try_resize(0); ControlFlow::Continue(original_state.into_done()) } - Err(e) => ControlFlow::Break((Poll::Ready(Some(Err(e))), original_state)), + Err(e) => Self::break_with_err(e), } } } @@ -325,10 +443,22 @@ impl Stream for PartialReduceHashAggregateStream { /// Aggregate one partial-state input batch, update the inner aggregate /// hash table, and continue with the next input batch. /// + /// -> EmittingOnMemoryPressure + /// The table cannot reserve enough memory. Materialize all accumulated + /// partial states and begin emitting them incrementally. + /// /// -> ProducingOutput /// Input was exhausted. Move to the next state to start outputting /// merged partial aggregate states. /// + /// EmittingOnMemoryPressure + /// -> EmittingOnMemoryPressure + /// One batch-sized slice was yielded; repeat until all materialized + /// partial states are emitted. + /// + /// -> ReadingInput + /// The materialized states were emitted; continue with the empty table. + /// /// ProducingOutput /// -> ProducingOutput /// One merged partial-state output batch was yielded; repeat to @@ -337,6 +467,13 @@ impl Stream for PartialReduceHashAggregateStream { /// -> Done /// All merged partial-state output was emitted. /// + /// Any active state + /// -> Error + /// An error drops state-owned resources before it is returned. + /// + /// Error + /// -> (end) + /// /// Done /// -> (end) /// ``` @@ -354,9 +491,18 @@ impl Stream for PartialReduceHashAggregateStream { state @ PartialReduceHashAggregateState::ReadingInput { .. } => { self.handle_reading_input(cx, state) } + state @ PartialReduceHashAggregateState::EmittingOnMemoryPressure { + .. + } => self.handle_emitting_on_memory_pressure(state), state @ PartialReduceHashAggregateState::ProducingOutput { .. } => { self.handle_producing_output(state) } + state @ PartialReduceHashAggregateState::Error => { + self.close_input(); + self.reservation.free(); + self.state = Some(state); + return Poll::Ready(None); + } state @ PartialReduceHashAggregateState::Done => { let _ = self.reservation.try_resize(0); self.state = Some(state); @@ -368,6 +514,19 @@ impl Stream for PartialReduceHashAggregateStream { ControlFlow::Continue(next_state) => { self.state = Some(next_state); } + ControlFlow::Break((Poll::Ready(Some(Err(e))), next_state)) => { + debug_assert!(matches!( + next_state, + PartialReduceHashAggregateState::Error + )); + + // The handler has already discarded its state-owned resources. + // Release the remaining stream-owned resources before returning. + self.close_input(); + self.reservation.free(); + self.state = Some(PartialReduceHashAggregateState::Error); + return Poll::Ready(Some(Err(e))); + } ControlFlow::Break((poll, next_state)) => { self.state = Some(next_state); return poll; From adb071111f43474aa0010004c3e041d886be7929 Mon Sep 17 00:00:00 2001 From: Stefan Wang <1fannnw@gmail.com> Date: Wed, 2 Sep 2026 05:43:10 +0000 Subject: [PATCH 10/37] fix: dedup commutative AND/OR operands when canonicalize is disabled (#23615) ## Which issue does this PR close? - Closes #14943. Revives #21870, which was closed by the stale bot. The direction was agreed in that thread, and `sql_planner` was benchmarked there ([results](https://github.com/apache/datafusion/pull/21870#issuecomment-4338614104), no measurable change). ## Rationale for this change `SimplifyExpressions` disables canonicalization for `LogicalPlan::Join` (see #8780), so the AND/OR dedup in `expr_contains_inner` cannot recognize duplicates that differ only by commutative operand order (`A = B` vs `B = A`). Simplifying inside a join filter, the delta-rs MERGE case reported in #14943, keeps the duplicate across simplifier cycles because nothing normalizes operand order first. The issue thread pointed at the fix: CSE already dedups `A = B` / `B = A` via `NormalizeEq`, so this routes the simplifier's leaf comparison through the same trait. ## What changes are included in this PR? `expr_contains_inner` compares leaves with `Expr::normalize_eq` instead of `==`. `NormalizeEq` handles `+`, `*`, `&`, `|`, `^`, `=`, and `!=` commutatively and falls back to structural `==` for everything else. Non-commutative rules and the existing `!needle.is_volatile()` guard are unchanged. A regression test covers the `AND`, `OR`, and 3-conjunct nested forms. `delete_xor_in_complex_expr` uses `normalize_eq` as well. `expr_contains` guards the `BitwiseXor` rules, which then hand deletion off to that helper. Leaving it structural let the guard fire on operands the helper could not delete, so the rule rebuilt its input and still reported `Transformed::yes`, spinning the simplifier to its cycle limit without changing the result. XOR is the only `expr_contains` caller shaped that way. The other eight return `*left`/`*right` directly, so a normalized match there is always a real change. `NormalizeEq for Expr` also now compares scalar/aggregate/window functions by full identity rather than `name()` alone. It previously treated two distinct functions that share a display name as equal, so routing the dedup through it collapsed `regex_udf(x) AND regex_udf(x)` (two different UDFs, same name) into one predicate. Registry-parsed SQL reuses one instance per name and is unaffected. This only stops merging genuinely distinct same-named functions, which programmatically built plans (delta-rs MERGE, the #14943 reporter) can produce. ## Are these changes tested? Each fix has a regression that fails without it. `test_simplify_swapped_operands_in_and_or_no_canonicalize` fails on `main` (the duplicate passes through unchanged) and passes with the fix; the `simplify_expr.slt` cases cover the join path end to end. `test_simplify_swapped_operands_in_xor_no_canonicalize` pins the cycle behavior: with `delete_xor_in_complex_expr` left structural it returns the input unchanged after `cycles=3`, and with the fix it folds to `0` in `cycles=2`. The existing `test_parameterized_scalar_udf` guards the function-identity fix. `datafusion-optimizer`, `datafusion-common`, `datafusion-expr`, and the full `sqllogictest` suite pass. `cargo fmt`/`clippy -D warnings` are clean. `sql_planner` shows no measurable change. ## Are there any user-facing changes? No public API changes. AND/OR and XOR chains containing commutative-equivalent duplicates now collapse even when the simplifier's canonicalizer is disabled (currently the `LogicalPlan::Join` path). Canonicalize-on paths produce the same output as before. --------- Signed-off-by: 1fanwang <1fannnw@gmail.com> Signed-off-by: Stefan Wang <1fannnw@gmail.com> Co-authored-by: Andrew Lamb --- datafusion/expr/src/expr.rs | 35 +++++++- .../simplify_expressions/expr_simplifier.rs | 85 +++++++++++++++++++ .../src/simplify_expressions/utils.rs | 20 ++++- .../sqllogictest/test_files/simplify_expr.slt | 54 ++++++++++++ 4 files changed, 187 insertions(+), 7 deletions(-) diff --git a/datafusion/expr/src/expr.rs b/datafusion/expr/src/expr.rs index 3b708b4edeec9..ad491c9782637 100644 --- a/datafusion/expr/src/expr.rs +++ b/datafusion/expr/src/expr.rs @@ -2514,7 +2514,7 @@ impl NormalizeEq for Expr { args: other_args, }), ) => { - self_func.name() == other_func.name() + self_func == other_func && self_args.len() == other_args.len() && self_args .iter() @@ -2545,7 +2545,7 @@ impl NormalizeEq for Expr { }, }), ) => { - self_func.name() == other_func.name() + self_func == other_func && self_distinct == other_distinct && self_null_treatment == other_null_treatment && self_args.len() == other_args.len() @@ -2598,7 +2598,7 @@ impl NormalizeEq for Expr { }, } = other.as_ref(); - self_fun.name() == other_fun.name() + self_fun == other_fun && self_window_frame == other_window_frame && match (self_filter, other_filter) { (Some(a), Some(b)) => a.normalize_eq(b), @@ -2611,10 +2611,12 @@ impl NormalizeEq for Expr { .iter() .zip(other_args.iter()) .all(|(a, b)| a.normalize_eq(b)) + && self_partition_by.len() == other_partition_by.len() && self_partition_by .iter() .zip(other_partition_by.iter()) .all(|(a, b)| a.normalize_eq(b)) + && self_order_by.len() == other_order_by.len() && self_order_by .iter() .zip(other_order_by.iter()) @@ -3831,6 +3833,7 @@ pub fn physical_name(expr: &Expr) -> Result { #[cfg(test)] mod test { use crate::expr_fn::col; + use crate::test::function_stub::max_udaf; use crate::{ ColumnarValue, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Volatility, case, lit, placeholder, qualified_wildcard, wildcard, wildcard_with_options, @@ -4345,6 +4348,32 @@ mod test { use super::*; use crate::logical_plan::{EmptyRelation, LogicalPlan}; + #[test] + fn normalize_eq_window_function_over_clause_lengths() { + let window = |partition_by: Vec, order_by: Vec| { + let mut window = WindowFunction::new(max_udaf(), vec![col("value")]); + window.params.partition_by = partition_by; + window.params.order_by = order_by; + Expr::from(window) + }; + let base = window(vec![col("a")], vec![Sort::new(col("a"), true, true)]); + + let extra_partition = window( + vec![col("a"), col("b")], + vec![Sort::new(col("a"), true, true)], + ); + assert!(!base.normalize_eq(&extra_partition)); + + let extra_order = window( + vec![col("a")], + vec![ + Sort::new(col("a"), true, true), + Sort::new(col("b"), true, true), + ], + ); + assert!(!base.normalize_eq(&extra_order)); + } + #[test] fn test_display_wildcard() { assert_eq!(format!("{}", wildcard()), "*"); diff --git a/datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs b/datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs index dadea4784802a..ec01becb8f7a5 100644 --- a/datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs +++ b/datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs @@ -2514,6 +2514,77 @@ mod tests { } } + #[test] + fn test_simplify_swapped_operands_in_and_or_no_canonicalize() { + // Regression test for https://github.com/apache/datafusion/issues/14943 + // + // `SimplifyExpressions` disables canonicalization for `LogicalPlan::Join` + // (see https://github.com/apache/datafusion/pull/8780), so commutative + // operands like `A = B` and `B = A` cannot be normalized to a single + // form before AND/OR dedup runs. The dedup itself must therefore + // recognize commutative equivalence directly. + + // c3 = 5 AND 5 = c3 --> c3 = 5 + let expr = col("c3_non_null") + .eq(lit(5_i64)) + .and(lit(5_i64).eq(col("c3_non_null"))); + let expected = col("c3_non_null").eq(lit(5_i64)); + assert_eq!(simplify_no_canonicalize(expr), expected); + + // 5 = c3 AND c3 = 5 --> 5 = c3 + let expr = lit(5_i64) + .eq(col("c3_non_null")) + .and(col("c3_non_null").eq(lit(5_i64))); + let expected = lit(5_i64).eq(col("c3_non_null")); + assert_eq!(simplify_no_canonicalize(expr), expected); + + // c3 = 5 OR 5 = c3 --> c3 = 5 + let expr = col("c3_non_null") + .eq(lit(5_i64)) + .or(lit(5_i64).eq(col("c3_non_null"))); + let expected = col("c3_non_null").eq(lit(5_i64)); + assert_eq!(simplify_no_canonicalize(expr), expected); + + // (c3 = 5 AND c4 > 0) AND (5 = c3) --> c3 = 5 AND c4 > 0 + let expr = col("c3_non_null") + .eq(lit(5_i64)) + .and(col("c4_non_null").gt(lit(0_u32))) + .and(lit(5_i64).eq(col("c3_non_null"))); + let expected = col("c3_non_null") + .eq(lit(5_i64)) + .and(col("c4_non_null").gt(lit(0_u32))); + assert_eq!(simplify_no_canonicalize(expr), expected); + } + + #[test] + fn test_simplify_swapped_operands_in_xor_no_canonicalize() { + // `expr_contains` guards the XOR rules, so `delete_xor_in_complex_expr` has to + // use the same equality relation. Comparing structurally there while the guard + // normalizes makes the rule rebuild its input and still report `Transformed::yes`, + // which spins the simplifier until it hits the cycle limit. + // + // Operands are non-nullable so the expected results do not depend on how the XOR + // rules treat a NULL operand, which is a separate question from the cycle. + + // (c4_non_null + 1) ^ (1 + c4_non_null) --> 0 + let expr = bitwise_xor( + col("c4_non_null") + lit(1_u32), + lit(1_u32) + col("c4_non_null"), + ); + let (simplified, cycles) = simplify_no_canonicalize_with_cycle_count(expr); + assert_eq!(simplified, lit(0_u32)); + assert_eq!(cycles, 2); + + // (c4_non_null + 1) ^ ((1 + c4_non_null) ^ c4_non_null) --> c4_non_null + let expr = bitwise_xor( + col("c4_non_null") + lit(1_u32), + bitwise_xor(lit(1_u32) + col("c4_non_null"), col("c4_non_null")), + ); + let (simplified, cycles) = simplify_no_canonicalize_with_cycle_count(expr); + assert_eq!(simplified, col("c4_non_null")); + assert_eq!(cycles, 2); + } + #[test] fn test_simplify_eq_not_self() { // `expr_a`: column `c2` is nullable, so `c2 = c2` simplifies to `c2 IS NOT NULL OR NULL` @@ -3667,6 +3738,20 @@ mod tests { try_simplify(expr).unwrap() } + fn simplify_no_canonicalize(expr: Expr) -> Expr { + simplify_no_canonicalize_with_cycle_count(expr).0 + } + + fn simplify_no_canonicalize_with_cycle_count(expr: Expr) -> (Expr, u32) { + let schema = expr_test_schema(); + let (transformed, count) = + ExprSimplifier::new(SimplifyContext::builder().with_schema(schema).build()) + .with_canonicalize(false) + .simplify_with_cycle_count_transformed(expr) + .unwrap(); + (transformed.data, count) + } + fn try_simplify_with_cycle_count(expr: Expr) -> Result<(Expr, u32)> { let schema = expr_test_schema(); let simplifier = diff --git a/datafusion/optimizer/src/simplify_expressions/utils.rs b/datafusion/optimizer/src/simplify_expressions/utils.rs index 89bb762d59ce2..78d801630c7ba 100644 --- a/datafusion/optimizer/src/simplify_expressions/utils.rs +++ b/datafusion/optimizer/src/simplify_expressions/utils.rs @@ -17,6 +17,7 @@ //! Utility functions for expression simplification +use datafusion_common::cse::NormalizeEq; use datafusion_common::{Result, ScalarValue, internal_err}; use datafusion_expr::{ Case, Expr, Like, Operator, @@ -26,13 +27,18 @@ use datafusion_expr::{ /// returns true if `needle` is found in a chain of search_op /// expressions. Such as: (A AND B) AND C +/// +/// Equality uses [`NormalizeEq`] so commutative operands (`A = B` vs `B = A`) +/// are recognized as the same predicate even when the canonicalizer is +/// disabled, notably for `LogicalPlan::Join` filters (see +/// ). fn expr_contains_inner(expr: &Expr, needle: &Expr, search_op: Operator) -> bool { match expr { Expr::BinaryExpr(BinaryExpr { left, op, right }) if *op == search_op => { expr_contains_inner(left, needle, search_op) || expr_contains_inner(right, needle, search_op) } - _ => expr == needle, + _ => expr.normalize_eq(needle), } } @@ -43,6 +49,12 @@ pub fn expr_contains(expr: &Expr, needle: &Expr, search_op: Operator) -> bool { /// Deletes all 'needles' or remains one 'needle' that are found in a chain of xor /// expressions. Such as: A ^ (A ^ (B ^ A)) +/// +/// Matching uses [`NormalizeEq`] to stay consistent with the [`expr_contains`] +/// guard on the XOR rules. A structural comparison here would let the guard fire +/// on operands this function then fails to delete, so the rule would rebuild its +/// input and still report a transformation, spinning the simplifier until it hits +/// the cycle limit. pub fn delete_xor_in_complex_expr(expr: &Expr, needle: &Expr, is_left: bool) -> Expr { /// Deletes recursively 'needles' in a chain of xor expressions fn recursive_delete_xor_in_expr( @@ -56,10 +68,10 @@ pub fn delete_xor_in_complex_expr(expr: &Expr, needle: &Expr, is_left: bool) -> { let left_expr = recursive_delete_xor_in_expr(left, needle, xor_counter); let right_expr = recursive_delete_xor_in_expr(right, needle, xor_counter); - if left_expr == *needle { + if left_expr.normalize_eq(needle) { *xor_counter += 1; return right_expr; - } else if right_expr == *needle { + } else if right_expr.normalize_eq(needle) { *xor_counter += 1; return left_expr; } @@ -76,7 +88,7 @@ pub fn delete_xor_in_complex_expr(expr: &Expr, needle: &Expr, is_left: bool) -> let mut xor_counter: i32 = 0; let result_expr = recursive_delete_xor_in_expr(expr, needle, &mut xor_counter); - if result_expr == *needle { + if result_expr.normalize_eq(needle) { return needle.clone(); } else if xor_counter % 2 == 0 { if is_left { diff --git a/datafusion/sqllogictest/test_files/simplify_expr.slt b/datafusion/sqllogictest/test_files/simplify_expr.slt index 158096328e960..b24a1a6cff975 100644 --- a/datafusion/sqllogictest/test_files/simplify_expr.slt +++ b/datafusion/sqllogictest/test_files/simplify_expr.slt @@ -165,6 +165,60 @@ physical_plan 01)ProjectionExec: expr=[column1@0 = 1 as opt1, column1@0 = 2 AND column1@0 != 2 as noopt1, column1@0 = 4 as opt2, column1@0 != 5 AND column1@0 = 5 as noopt2] 02)--DataSourceExec: partitions=1, partition_sizes=[1] +# Dedup commutative AND/OR operands in join conditions +# (canonicalization is disabled for joins, so `A = B` vs `B = A` must be +# recognized directly). See https://github.com/apache/datafusion/issues/14943 + +statement ok +create table t1(a int) as values (1), (2); + +statement ok +create table t2(b int) as values (1), (3); + +# `t1.a = t2.b OR t2.b = t1.a` simplifies to `t1.a = t2.b`, allowing a +# HashJoin instead of a NestedLoopJoin +query TT +explain select * from t1 join t2 on t1.a = t2.b or t2.b = t1.a; +---- +logical_plan +01)Inner Join: t1.a = t2.b +02)--TableScan: t1 projection=[a] +03)--TableScan: t2 projection=[b] +physical_plan +01)HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(a@0, b@0)] +02)--DataSourceExec: partitions=1, partition_sizes=[1] +03)--DataSourceExec: partitions=1, partition_sizes=[1] + +query II +select * from t1 join t2 on t1.a = t2.b or t2.b = t1.a; +---- +1 1 + +# swapped commutative operands (`t1.a + t2.b` vs `t2.b + t1.a`) dedup within +# a non-equi join filter +query TT +explain select * from t1 join t2 on t1.a > t2.b and (t1.a + t2.b > 1 or t2.b + t1.a > 1); +---- +logical_plan +01)Inner Join: Filter: t1.a > t2.b AND t1.a + t2.b > Int32(1) +02)--TableScan: t1 projection=[a] +03)--TableScan: t2 projection=[b] +physical_plan +01)NestedLoopJoinExec: join_type=Inner, filter=a@0 > b@1 AND a@0 + b@1 > 1 +02)--DataSourceExec: partitions=1, partition_sizes=[1] +03)--DataSourceExec: partitions=1, partition_sizes=[1] + +query II +select * from t1 join t2 on t1.a > t2.b and (t1.a + t2.b > 1 or t2.b + t1.a > 1); +---- +2 1 + +statement ok +drop table t1; + +statement ok +drop table t2; + # Identity Date cast in a comparison predicate. # `cast(d AS date)` where `d` is already Date32 is an identity cast and should # fold away, so the predicate compares against the bare column `d`. This enables From 3a4c310da72ea419e777bf56bbedaf1fb83e087e Mon Sep 17 00:00:00 2001 From: Oleks V Date: Wed, 2 Sep 2026 07:33:47 +0000 Subject: [PATCH 11/37] minor: reenable failing mark join test (#24863) ## Which issue does this PR close? - Related #24854 . ## Rationale for this change Reenable test for slt join matrix test, after #21585 merged ## What changes are included in this PR? ## What is the testing strategy for this PR? ## Are there any user-facing changes? --- .../sqllogictest/test_files/mark_join_matrix.slt | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/datafusion/sqllogictest/test_files/mark_join_matrix.slt b/datafusion/sqllogictest/test_files/mark_join_matrix.slt index d394d527d5d13..a17227df48e4b 100644 --- a/datafusion/sqllogictest/test_files/mark_join_matrix.slt +++ b/datafusion/sqllogictest/test_files/mark_join_matrix.slt @@ -118,13 +118,12 @@ WHERE l.v > 35 OR l.k NOT IN (SELECT r.k FROM mk_r_nn r); 4 40 NULL 50 -# https://github.com/apache/datafusion/issues/24854 -# query II rowsort -# SELECT l.k, l.v FROM mk_l l -# WHERE l.v > 35 OR l.k NOT IN (SELECT r.k FROM mk_r r); -# ---- -# 4 40 -# NULL 50 +query II rowsort +SELECT l.k, l.v FROM mk_l l +WHERE l.v > 35 OR l.k NOT IN (SELECT r.k FROM mk_r r); +---- +4 40 +NULL 50 # Empty subquery: EXISTS is always false, so the result is just the predicate. query II rowsort From 3b63300849b5b016f9bcff03803b6267f2c74cba Mon Sep 17 00:00:00 2001 From: Thor <8681572+thorfour@users.noreply.github.com> Date: Wed, 2 Sep 2026 12:52:13 +0000 Subject: [PATCH 12/37] fix: support RunEndEncoded arrays in function and window frame type coercion (#24565) ## Which issue does this PR close? - Closes #24564. ## Rationale for this change Adds missing support for coercing REE arrays and extracting range windows from REE arrays. ## What changes are included in this PR? Added case statements for supporting REE arrays ## Are these changes tested? Yes a unit test is included. ## Are there any user-facing changes? No --- .../expr/src/type_coercion/functions.rs | 54 ++++++++++++++++ .../optimizer/src/analyzer/type_coercion.rs | 7 ++- .../test_files/run_end_encoded.slt | 62 +++++++++++++++++++ 3 files changed, 122 insertions(+), 1 deletion(-) diff --git a/datafusion/expr/src/type_coercion/functions.rs b/datafusion/expr/src/type_coercion/functions.rs index 781559ddd0c5c..8e86cb3685e90 100644 --- a/datafusion/expr/src/type_coercion/functions.rs +++ b/datafusion/expr/src/type_coercion/functions.rs @@ -1158,6 +1158,16 @@ fn coerced_from<'a>( { Some(type_into.clone()) } + (_, RunEndEncoded(_, value_type)) + if coerced_from(type_into, value_type.data_type()).is_some() => + { + Some(type_into.clone()) + } + (RunEndEncoded(_, value_type), _) + if coerced_from(value_type.data_type(), type_from).is_some() => + { + Some(type_into.clone()) + } // coerced into type_into (Int8, Null | Int8) => Some(type_into.clone()), (Int16, Null | Int8 | Int16 | UInt8) => Some(type_into.clone()), @@ -1617,6 +1627,50 @@ mod tests { ); } + #[test] + fn test_coerced_from_run_end_encoded() { + let run_end_encoded_of = |value_type: DataType| { + DataType::RunEndEncoded( + Field::new("run_ends", DataType::Int32, false).into(), + Field::new("values", value_type, true).into(), + ) + }; + + let type_into = run_end_encoded_of(DataType::UInt32); + let type_from = DataType::Int64; + assert_eq!(coerced_from(&type_into, &type_from), None); + + let type_from = run_end_encoded_of(DataType::UInt32); + let type_into = DataType::Int64; + assert_eq!( + coerced_from(&type_into, &type_from), + Some(type_into.clone()) + ); + + // Signature candidates for functions like `date_bin` are plain + // Timestamp, but a REE-encoded column (e.g. a segment written with + // REE-dict encoding for that field) should still coerce against + // them via the wrapped value type. + let type_from = + run_end_encoded_of(DataType::Timestamp(TimeUnit::Nanosecond, None)); + let type_into = DataType::Timestamp(TimeUnit::Nanosecond, None); + assert_eq!( + coerced_from(&type_into, &type_from), + Some(type_into.clone()) + ); + + // The reverse direction: a plain type coercing into an REE target + // (e.g. a signature that happens to require RunEndEncoded) should + // succeed whenever the plain type coerces into the wrapped value + // type. + let type_into = run_end_encoded_of(DataType::Int64); + let type_from = DataType::Int32; + assert_eq!( + coerced_from(&type_into, &type_from), + Some(type_into.clone()) + ); + } + #[test] fn test_get_valid_types_array_and_array() -> Result<()> { let function = "array_and_array"; diff --git a/datafusion/optimizer/src/analyzer/type_coercion.rs b/datafusion/optimizer/src/analyzer/type_coercion.rs index 0f82f0b0df764..13b9541b6a71a 100644 --- a/datafusion/optimizer/src/analyzer/type_coercion.rs +++ b/datafusion/optimizer/src/analyzer/type_coercion.rs @@ -1088,6 +1088,8 @@ fn extract_window_frame_target_type(col_type: &DataType) -> Result { Ok(DataType::Interval(IntervalUnit::MonthDayNano)) } else if let DataType::Dictionary(_, value_type) = col_type { extract_window_frame_target_type(value_type) + } else if let DataType::RunEndEncoded(_, value_type) = col_type { + extract_window_frame_target_type(value_type.data_type()) } else { internal_err!("Cannot run range queries on datatype: {col_type}") } @@ -1113,8 +1115,11 @@ fn coerce_window_frame( // `current_value ± offset`, so it is only meaningful for target // types that support arithmetic. Other orderable target types can // still use free range frames, whose bounds require comparison only. + // REE arrays are not supported by arrow's numeric kernesl. + // Tracked at https://github.com/apache/arrow-rs/issues/10891). let supports_offset_arithmetic = - target_type.is_numeric() || is_interval(&target_type); + !matches!(col_type, DataType::RunEndEncoded(_, _)) + && (target_type.is_numeric() || is_interval(&target_type)); if !supports_offset_arithmetic && !window_frame.free_range() { return plan_err!( "RANGE with offset PRECEDING/FOLLOWING is not supported for ORDER BY type {target_type}" diff --git a/datafusion/sqllogictest/test_files/run_end_encoded.slt b/datafusion/sqllogictest/test_files/run_end_encoded.slt index b5909bc7c430d..094f6ede1b294 100644 --- a/datafusion/sqllogictest/test_files/run_end_encoded.slt +++ b/datafusion/sqllogictest/test_files/run_end_encoded.slt @@ -85,3 +85,65 @@ false false true true true true true true + +# date_bin's signature is a plain Timestamp, not RunEndEncoded, so this only +# plans if coerced_from() unwraps the REE value type the same way it does for +# Dictionary (see test_coerced_from_run_end_encoded) +statement ok +CREATE TABLE ree_timestamps AS +SELECT + arrow_cast( + arrow_cast(ts, 'Timestamp(Nanosecond, None)'), + 'RunEndEncoded("run_ends": non-null Int32, "values": Timestamp(Nanosecond, None))' + ) AS ts +FROM (VALUES + ('2023-12-04T00:00:00'), + ('2023-12-04T00:05:00'), + ('2023-12-04T00:35:00') +) AS t(ts); + +query P rowsort +SELECT date_bin('30 minutes', ts) FROM ree_timestamps; +---- +2023-12-04T00:00:00 +2023-12-04T00:00:00 +2023-12-04T00:30:00 + +# RANGE window frame ordered by an REE-encoded column: extract_window_frame_target_type() +# unwraps RunEndEncoded the same way it does for Dictionary, so a frame that +# doesn't need offset arithmetic. +query II +SELECT + temperature, + SUM(temperature) OVER ( + ORDER BY arrow_cast( + temperature, + 'RunEndEncoded("run_ends": non-null Int32, "values": Int64)' + ) + RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW + ) AS windowed_sum +FROM sensor_readings +ORDER BY temperature; +---- +20 20 +22 42 +23 65 +24 89 + +# An offset RANGE bound (`k PRECEDING`/`k FOLLOWING`, unlike CURRENT ROW or +# UNBOUNDED above) is computed by adding/subtracting the offset from the +# current row's value, and arrow's numeric kernels don't support arithmetic +# on a still-RunEndEncoded-wrapped value (tracked at +# https://github.com/apache/arrow-rs/issues/10891). +query error RANGE with offset PRECEDING/FOLLOWING is not supported for ORDER BY type +SELECT + temperature, + SUM(temperature) OVER ( + ORDER BY arrow_cast( + temperature, + 'RunEndEncoded("run_ends": non-null Int32, "values": Int64)' + ) + RANGE BETWEEN 1 PRECEDING AND 1 FOLLOWING + ) AS windowed_sum +FROM sensor_readings +ORDER BY temperature; From 406c0c668cb679f96b0406ef4a6837fd5d3bff64 Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Wed, 2 Sep 2026 13:45:29 +0000 Subject: [PATCH 13/37] chore: convert `FinalHashAggregateStream` to async generators and cleanup (#24874) ## Which issue does this PR close? Related to: - #23974 ## Rationale for this change Cleanup the code and remove state ## What changes are included in this PR? changed `FinalHashAggregateStream` to async generator and remove unneeded code due to that The first commit in this PR is `FinalHashAggregateStream` ## Are these changes tested? existing tests ## Are there any user-facing changes? nope --- .../src/aggregates/hash_stream.rs | 646 +++++------------- .../physical-plan/src/aggregates/mod.rs | 2 +- 2 files changed, 176 insertions(+), 472 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/hash_stream.rs b/datafusion/physical-plan/src/aggregates/hash_stream.rs index 340bf5cfc12d6..edf084ad328bc 100644 --- a/datafusion/physical-plan/src/aggregates/hash_stream.rs +++ b/datafusion/physical-plan/src/aggregates/hash_stream.rs @@ -32,9 +32,12 @@ use std::task::{Context, Poll}; use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; -use datafusion_common::{DataFusionError, Result, internal_datafusion_err, internal_err}; -use datafusion_execution::TaskContext; +use datafusion_common::{ + DataFusionError, Result, assert_ne_or_internal_err, internal_datafusion_err, + internal_err, +}; use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; +use datafusion_execution::{TaskContext, TryEmitter, async_try_stream}; use datafusion_physical_expr::PhysicalSortExpr; use datafusion_physical_expr::expressions::Column; use datafusion_physical_expr_common::sort_expr::LexOrdering; @@ -53,7 +56,7 @@ use crate::metrics::{ use crate::sorts::IncrementalSortIterator; use crate::sorts::streaming_merge::{SortedSpillFile, StreamingMergeBuilder}; use crate::spill::spill_manager::SpillManager; -use crate::stream::EmptyRecordBatchStream; +use crate::stream::{EmptyRecordBatchStream, RecordBatchStreamAdapter}; use crate::{InputOrderMode, RecordBatchStream, SendableRecordBatchStream, metrics}; /// Hash aggregation is implemented in two stages: partial and final. This @@ -249,47 +252,15 @@ pub(crate) struct FinalHashAggregateStream { /// See comments for the same variable in [`PartialHashAggregateStream`]. group_values_soft_limit: Option, - /// Tracks the high-level stream lifecycle. The hash table owns the lower-level + /// The hash table owns the lower-level /// state for emitting output batches. - state: Option, -} - -/// States for final hash aggregation processing. -// The typestate pattern is used in case the inner logic becomes more complex in -// the future. -enum FinalHashAggregateState { - ReadingInput { - hash_table: AggregateHashTable, - /// `None` if spilling is not supported by the configured `DiskManager`. - spill_context: Option>, - }, - Spilling { - hash_table: AggregateHashTable, - spill_context: Box, - }, - ProducingOutput { - hash_table: AggregateHashTable, - }, - PreparingMergeInput { - hash_table: AggregateHashTable, - spill_context: Box, - }, - MergingSpills { - stream: SendableRecordBatchStream, - }, - Done, - /// Sentinel state to use when returning error from any other states, because: - /// - It explicitly releases state-owned resources immediately - /// - More defensive against accidentally resuming execution after error - Error, + /// + /// This will be None when creating the stream + hash_table: Option>, + /// `None` if spilling is not supported by the configured `DiskManager`. + spill_context: Option>, } -type FinalHashAggregatePoll = Poll>>; -type FinalHashAggregateStateTransition = ControlFlow< - (FinalHashAggregatePoll, FinalHashAggregateState), - FinalHashAggregateState, ->; - impl FinalSpillContext { fn new( agg: &AggregateExec, @@ -1057,29 +1028,50 @@ impl FinalHashAggregateStream { baseline_metrics, reservation, group_values_soft_limit: agg.limit_options().map(|config| config.limit()), - state: Some(FinalHashAggregateState::ReadingInput { - hash_table, - spill_context, - }), + hash_table: Some(hash_table), + spill_context, }) } - fn close_input(&mut self) { - let input_schema = self.input.schema(); - self.input = Box::pin(EmptyRecordBatchStream::new(input_schema)); + pub(crate) fn into_stream(self) -> SendableRecordBatchStream { + let schema = Arc::clone(&self.schema); + + Box::pin(RecordBatchStreamAdapter::new(schema, self.create_stream())) } - fn break_with_err(error: DataFusionError) -> FinalHashAggregateStateTransition { - ControlFlow::Break(( - Poll::Ready(Some(Err(error))), - FinalHashAggregateState::Error, - )) + /// Entry point for the final hash aggregate flow + /// + /// See comments in [`FinalHashAggregateStream`] for high-level ideas. + fn create_stream(mut self) -> impl Stream> { + async_try_stream(|emitter| async move { + let mut hash_table = self + .hash_table + .take() + .expect("hash_table should not be None"); + + let mut spill_context = self.spill_context.take(); + + self.consume_input(&mut hash_table, &mut spill_context) + .await?; + self.close_input(); + + match spill_context.filter(|s| s.has_spills()) { + // - If spilled before, perform merging spill runs + Some(spill_context) => { + self.produce_output_from_spills(hash_table, spill_context, emitter) + .await? + } + // Either all the input fit in memory or hit soft group limit with no spilling + None => self.produce_output_from_memory(hash_table, emitter).await?, + } + + Ok(()) + }) } - fn break_with_internal_err( - message: impl std::fmt::Display, - ) -> FinalHashAggregateStateTransition { - Self::break_with_err(internal_datafusion_err!("{message}")) + fn close_input(&mut self) { + let input_schema = self.input.schema(); + self.input = Box::pin(EmptyRecordBatchStream::new(input_schema)); } /// See comments in [`Self::group_values_soft_limit`] for details. @@ -1088,14 +1080,6 @@ impl FinalHashAggregateStream { .is_some_and(|limit| limit <= hash_table.building_group_count()) } - fn start_output( - &mut self, - hash_table: &mut AggregateHashTable, - ) -> Result<()> { - self.close_input(); - hash_table.start_output() - } - /// Reserve memory for the current aggregate table. fn reservation_size_for_table( hash_table: &AggregateHashTable, @@ -1116,454 +1100,174 @@ impl FinalHashAggregateStream { } } - /// Handle ReadingInput state - aggregate partial state batches into the hash table. - /// - /// See comments at `poll_next()` for details. + /// Read input stream, if no memory, then spill and continue reading - aggregate partial state batches into the hash table. /// - /// Returns the next operator state with control flow decision. - fn handle_reading_input( + /// Spilling: The table cannot reserve enough memory. + /// Move all current states into one fully group-key-sorted spill run. + async fn consume_input( &mut self, - cx: &mut Context<'_>, - original_state: FinalHashAggregateState, - ) -> FinalHashAggregateStateTransition { - let FinalHashAggregateState::ReadingInput { - mut hash_table, - spill_context, - } = original_state - else { - return Self::break_with_internal_err( - "Final hash aggregate stream expected ReadingInput state", - ); - }; + hash_table: &mut AggregateHashTable, + spill_context: &mut Option>, + ) -> Result<()> { + let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); - match self.input.poll_next_unpin(cx) { - Poll::Pending => ControlFlow::Break(( - Poll::Pending, - FinalHashAggregateState::ReadingInput { - hash_table, - spill_context, - }, - )), - Poll::Ready(Some(Ok(batch))) => { - let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); - let timer = elapsed_compute.timer(); - let result = hash_table.aggregate_batch(&batch); - timer.done(); + while let Some(batch) = self.input.next().await.transpose()? { + let _timer = elapsed_compute.timer(); + hash_table.aggregate_batch(&batch)?; + + // Soft group limits are usually small and rarely coincide with + // spilling. Once spilling has occurred, skip this optimization to + // make the internal logic simpler. + let spilled = spill_context + .as_ref() + .is_some_and(|context| context.has_spills()); + if self.hit_soft_group_limit(hash_table) && !spilled { + break; + } - if let Err(e) = result { - return Self::break_with_err(e); - } + // Check memory reservation, and potentially spill. + let resize_result = + self.reservation + .try_resize(Self::reservation_size_for_table( + hash_table, + spill_context.as_deref(), + )); + + match resize_result { + Ok(()) => {} + + // The table cannot reserve enough memory. + // Move all current states into one fully group-key-sorted spill run. + Err(e @ DataFusionError::ResourcesExhausted(_)) => { + // OOM and don't support spilling from configuration + let spill_context = spill_context.as_mut().ok_or_else(|| e.context( + "Final hash aggregate cannot spill because temporary files are not enabled in the DiskManager", + ))?; + + // Sanity check: impossible to OOM when there is no group aggregated. + assert_ne_or_internal_err!( + hash_table.building_group_count(), + 0, + "Final hash aggregate ran out of memory with no aggregated groups" + ); - // Soft group limits are usually small and rarely coincide with - // spilling. Once spilling has occurred, skip this optimization to - // make the internal logic simpler. - let spilled = spill_context - .as_ref() - .is_some_and(|context| context.has_spills()); - if self.hit_soft_group_limit(&hash_table) && !spilled { - let timer = elapsed_compute.timer(); - let result = self.start_output(&mut hash_table); - timer.done(); + // Sorts and spills one complete in-memory state run - return match result { - Ok(()) => ControlFlow::Continue( - FinalHashAggregateState::ProducingOutput { hash_table }, - ), - Err(e) => Self::break_with_err(e), - }; - } + // Go to the next state to perform spilling the aggregated + // groups so far. + let result = spill_context.spill_table(hash_table); - // Check memory reservation, and potentially spill. - let timer = elapsed_compute.timer(); - let resize_result = + // Spilling shrinks the aggregate table and releases its accumulated + // memory. Update the reservation accordingly. self.reservation - .try_resize(Self::reservation_size_for_table( - &hash_table, - spill_context.as_deref(), - )); - timer.done(); - match resize_result { - Ok(()) => {} - Err(e @ DataFusionError::ResourcesExhausted(_)) => { - // OOM and don't support spilling from configuration - let Some(spill_context) = spill_context else { - return Self::break_with_err(e.context( - "Final hash aggregate cannot spill because temporary files are not enabled in the DiskManager", - )); - }; - // Sanity check: impossible to OOM when there is no group aggregated. - if hash_table.building_group_count() == 0 { - return Self::break_with_internal_err( - "Final hash aggregate ran out of memory with no aggregated groups", - ); - } - // Go to the next state to perform spilling the aggregated - // groups so far. - return ControlFlow::Continue( - FinalHashAggregateState::Spilling { - hash_table, - spill_context, - }, - ); - } - Err(e) => return Self::break_with_err(e), - } + .try_resize(hash_table.memory_size()) + .map_err(|e| { + e.context( + "Decreasing allocation after spilling should succeed", + ) + })?; - ControlFlow::Continue(FinalHashAggregateState::ReadingInput { - hash_table, - spill_context, - }) - } - Poll::Ready(Some(Err(e))) => Self::break_with_err(e), - // Input done, move to next state: - // - If spilled before, perform merging spill runs - // - If not spilled, start producing outputs - Poll::Ready(None) => { - self.close_input(); - match spill_context { - Some(spill_context) if spill_context.has_spills() => { - ControlFlow::Continue( - FinalHashAggregateState::PreparingMergeInput { - hash_table, - spill_context, - }, - ) - } - _ => { - let elapsed_compute = - self.baseline_metrics.elapsed_compute().clone(); - let timer = elapsed_compute.timer(); - let result = hash_table.start_output(); - timer.done(); - - match result { - Ok(()) => ControlFlow::Continue( - FinalHashAggregateState::ProducingOutput { hash_table }, - ), - Err(e) => Self::break_with_err(e), - } - } + result?; + + // One sorted run was written; resume reading the original input. } + Err(e) => return Err(e), } } + + Ok(()) } - /// Sorts and spills one complete in-memory state run, then resumes input. - /// - /// See comments at `poll_next()` for details. - /// - /// Returns the next operator state with control flow decision. - fn handle_spilling( + /// Produce output from spills + /// 1. Spill in progress in-memory hash table + /// 2. Switch to ordered final stream + /// 3. passthrough stream output + async fn produce_output_from_spills( &mut self, - original_state: FinalHashAggregateState, - ) -> FinalHashAggregateStateTransition { - let FinalHashAggregateState::Spilling { - mut hash_table, - mut spill_context, - } = original_state - else { - return Self::break_with_internal_err( - "Final hash aggregate stream expected Spilling state", - ); - }; - - // Sanity check: it is impossible to OOM when the table is empty. - if hash_table.building_group_count() == 0 { - return Self::break_with_internal_err( - "Final hash aggregation entered Spilling with an empty table", - ); - } - + mut hash_table: AggregateHashTable, + mut spill_context: Box, + mut emitter: TryEmitter, + ) -> Result<()> { let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); let timer = elapsed_compute.timer(); - let mut result = spill_context.spill_table(&mut hash_table); - // Spilling shrinks the aggregate table and releases its accumulated - // memory. Update the reservation accordingly. - if let Err(e) = self.reservation.try_resize(hash_table.memory_size()) { - result = - Err(e.context("Decreasing allocation after spilling should succeed")); - } + // Input was exhausted after spilling. Spill the last in-memory run + spill_context.spill_table(&mut hash_table)?; + + // Construct the ordered input used to merge all spill files. + let mut output_stream = + self.switch_to_ordered_final_stream(hash_table, spill_context)?; timer.done(); - match result { - // Finished spilling the aggregate table, continue aggregating from input. - Ok(()) => ControlFlow::Continue(FinalHashAggregateState::ReadingInput { - hash_table, - spill_context: Some(spill_context), - }), - Err(e) => Self::break_with_err(e), + // Forwards output from the fully ordered stream that consumes the merged + // spill runs. + // + // Not wrapping in a timer and not record output batches since this is now `merge_stream` responsibility + // we just pass through + while let Some(batch) = output_stream.next().await.transpose()? { + emitter.emit(batch).await; } + + Ok(()) } - /// 1. Spills the last in-memory run. - /// 2. Constructs a globally ordered input stream by applying a sort-preserving + /// 1. Constructs a globally ordered input stream by applying a sort-preserving /// merge to all spills. - /// 3. Constructs a replay stream: an ordered final aggregate stream over the + /// 2. Constructs a replay stream: an ordered final aggregate stream over the /// fully ordered input constructed from the spills. /// - /// See comments at `poll_next()` for details. - /// - /// Returns the next operator state with control flow decision. - fn handle_preparing_merge_input( + /// Returns the replay stream + fn switch_to_ordered_final_stream( &mut self, - original_state: FinalHashAggregateState, - ) -> FinalHashAggregateStateTransition { - let FinalHashAggregateState::PreparingMergeInput { - mut hash_table, - mut spill_context, - } = original_state - else { - return Self::break_with_internal_err( - "Final hash aggregate stream expected PreparingMergeInput state", - ); - }; - - let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); - let timer = elapsed_compute.timer(); - let replay = match spill_context.spill_table(&mut hash_table) { - Ok(()) => { - let metrics = OrderedAggregateTableMetrics::from_hash_table(&hash_table); - drop(hash_table); - match self.reservation.try_resize(0) { - Ok(()) => (*spill_context).into_replay_stream( - &self.baseline_metrics, - metrics, - self.reservation.new_empty(), - ), - Err(e) => Err(e), - } - } - Err(e) => Err(e), - }; - timer.done(); - - match replay { - Ok(stream) => { - ControlFlow::Continue(FinalHashAggregateState::MergingSpills { stream }) - } - Err(e) => Self::break_with_err(e), - } - } - - /// Forwards output from the fully ordered stream that consumes the merged - /// spill runs. - /// - /// See comments at `poll_next()` for details. - /// - /// Returns the next operator state with control flow decision. - fn handle_merging_spills( - &mut self, - cx: &mut Context<'_>, - original_state: FinalHashAggregateState, - ) -> FinalHashAggregateStateTransition { - let FinalHashAggregateState::MergingSpills { mut stream } = original_state else { - return Self::break_with_internal_err( - "Final hash aggregate stream expected MergingSpills state", - ); - }; - - match stream.poll_next_unpin(cx) { - Poll::Pending => ControlFlow::Break(( - Poll::Pending, - FinalHashAggregateState::MergingSpills { stream }, - )), - Poll::Ready(Some(Ok(batch))) => ControlFlow::Break(( - Poll::Ready(Some(Ok(batch))), - FinalHashAggregateState::MergingSpills { stream }, - )), - Poll::Ready(Some(Err(e))) => Self::break_with_err(e), - Poll::Ready(None) => ControlFlow::Continue(FinalHashAggregateState::Done), - } + hash_table: AggregateHashTable, + spill_context: Box, + ) -> Result { + let metrics = OrderedAggregateTableMetrics::from_hash_table(&hash_table); + drop(hash_table); + self.reservation.try_resize(0)?; + spill_context.into_replay_stream( + &self.baseline_metrics, + metrics, + self.reservation.new_empty(), + ) } - /// Handle ProducingOutput state - emit final aggregate value batches. - /// - /// See comments at `poll_next()` for details. - /// - /// Returns the next operator state with control flow decision. - fn handle_producing_output( + /// Emit final aggregate value batches: + /// Input was exhausted without spilling, or the soft group limit was reached. + async fn produce_output_from_memory( &mut self, - original_state: FinalHashAggregateState, - ) -> FinalHashAggregateStateTransition { - let FinalHashAggregateState::ProducingOutput { mut hash_table } = original_state - else { - return Self::break_with_internal_err( - "Final hash aggregate stream expected ProducingOutput state", - ); - }; - + mut hash_table: AggregateHashTable, + mut emitter: TryEmitter, + ) -> Result<()> { let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); - let timer = elapsed_compute.timer(); - let result = hash_table.next_output_batch(); - timer.done(); - - match result { - Ok(Some(batch)) => { - let next_state = if hash_table.is_done() { - drop(hash_table); - if let Err(e) = self.reservation.try_resize(0) { - return Self::break_with_err(e); - } - FinalHashAggregateState::Done - } else { - if let Err(e) = self.reservation.try_resize(hash_table.memory_size()) - { - return Self::break_with_err(e); - } - FinalHashAggregateState::ProducingOutput { hash_table } - }; - ControlFlow::Break(( - Poll::Ready(Some(Ok(batch.record_output(&self.baseline_metrics)))), - next_state, - )) - } - Err(e) => Self::break_with_err(e), - Ok(None) => { - drop(hash_table); - let next_state = FinalHashAggregateState::Done; - if let Err(e) = self.reservation.try_resize(0) { - return Self::break_with_err(e); - } - ControlFlow::Continue(next_state) - } - } - } -} - -impl Stream for FinalHashAggregateStream { - type Item = Result; + let mut timer = elapsed_compute.timer(); + hash_table.start_output()?; - /// Entry point for the final hash aggregate state machine. - /// - /// See comments in [`FinalHashAggregateStream`] for high-level ideas. - /// - /// State transition graph: - /// - /// ```text - /// (start) - /// -> ReadingInput - /// The stream starts by polling partial-state input and aggregating - /// those states into the final hash table. - /// - /// ReadingInput - /// -> ReadingInput - /// Aggregate one partial-state input batch. If it fits in memory, - /// continue with the next input batch. - /// -> Spilling - /// The table cannot reserve enough memory. Move all current states into - /// one fully group-key-sorted spill run. - /// -> ProducingOutput - /// Input was exhausted without spilling, or the soft group limit was - /// reached. Start outputting final aggregate values. - /// -> PreparingMergeInput - /// Input was exhausted after spilling. Spill the last in-memory run and - /// construct the ordered input used to merge all spill files. - /// - /// Spilling - /// -> ReadingInput - /// One sorted run was written; resume reading the original input. - /// - /// PreparingMergeInput - /// Spill the final in-memory run and build the input ordered replay stream. - /// -> MergingSpills - /// The final run was spilled and the ordered replay stream was built. - /// - /// MergingSpills - /// Aggregate the merged spill runs and emit final results. - /// -> MergingSpills - /// Forward one result batch from the fully ordered replay stream that - /// consumes the sort-preserving merge. - /// -> Done - /// The merged spill input was fully aggregated. - /// - /// ProducingOutput - /// -> ProducingOutput - /// One final output batch was yielded; repeat to continue producing - /// output incrementally. - /// -> Done - /// All final output was emitted. - /// - /// Any active state - /// -> Error - /// An error drops state-owned resources before it is returned. - /// - /// Error - /// -> (end) - /// - /// Done - /// -> (end) - /// ``` - fn poll_next( - mut self: std::pin::Pin<&mut Self>, - cx: &mut Context<'_>, - ) -> Poll> { loop { - let cur_state = self - .state - .take() - .expect("FinalHashAggregateStream state should not be None"); - - let next_state = match cur_state { - state @ FinalHashAggregateState::ReadingInput { .. } => { - self.handle_reading_input(cx, state) - } - state @ FinalHashAggregateState::Spilling { .. } => { - self.handle_spilling(state) - } - state @ FinalHashAggregateState::PreparingMergeInput { .. } => { - self.handle_preparing_merge_input(state) - } - state @ FinalHashAggregateState::MergingSpills { .. } => { - self.handle_merging_spills(cx, state) - } - state @ FinalHashAggregateState::ProducingOutput { .. } => { - self.handle_producing_output(state) - } - state @ FinalHashAggregateState::Error => { - self.close_input(); - self.reservation.free(); - self.state = Some(state); - return Poll::Ready(None); - } - state @ FinalHashAggregateState::Done => { - let _ = self.reservation.try_resize(0); - self.state = Some(state); - return Poll::Ready(None); - } + let Some(batch) = hash_table.next_output_batch()? else { + // Only reachable when the table held no groups at all: a + // non-empty table always reports its last batch together with + // the `Done` state, which the `try_resize` below already zeroes. + self.reservation.try_resize(0)?; + return Ok(()); }; - match next_state { - ControlFlow::Continue(next_state) => { - self.state = Some(next_state); - } - ControlFlow::Break((Poll::Ready(Some(Err(e))), next_state)) => { - debug_assert!(matches!(next_state, FinalHashAggregateState::Error)); + // The table hands over its groups as they are materialized and + // reports a size of 0 once it reaches `Done`, so this releases the + // reservation before the final batch goes downstream. + self.reservation.try_resize(hash_table.memory_size())?; - // The handler has already discarded its state-owned resources. - // Release the remaining stream-owned resources before returning. - self.close_input(); - self.reservation.free(); - self.state = Some(FinalHashAggregateState::Error); - return Poll::Ready(Some(Err(e))); - } - ControlFlow::Break((poll, next_state)) => { - self.state = Some(next_state); - return poll; - } - } + timer.done(); + emitter + .emit(batch.record_output(&self.baseline_metrics)) + .await; + timer = elapsed_compute.timer(); } } } -impl RecordBatchStream for FinalHashAggregateStream { - fn schema(&self) -> SchemaRef { - Arc::clone(&self.schema) - } -} - #[cfg(test)] mod tests { use std::sync::Arc; diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 6da7ee1018dc5..a6b5819916275 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -715,7 +715,7 @@ impl From for SendableRecordBatchStream { StreamType::AggregateStream(stream) => Box::pin(stream), StreamType::PartialHash(stream) => Box::pin(stream), StreamType::PartialReduceHash(stream) => Box::pin(stream), - StreamType::FinalHash(stream) => Box::pin(stream), + StreamType::FinalHash(stream) => stream.into_stream(), StreamType::SingleHash(stream) => Box::pin(stream), StreamType::OrderedPartialAggregate(stream) => stream.into_stream(), StreamType::OrderedFinalAggregate(stream) => Box::pin(stream), From bce3c27f88c45e0aae52fab08fb5f36c48d809e6 Mon Sep 17 00:00:00 2001 From: Jay Zhan Date: Wed, 2 Sep 2026 14:29:00 +0000 Subject: [PATCH 14/37] perf: Avoid cloning EquivalenceProperties in ordering satisfaction checks (#24800) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Which issue does this PR close? - Closes #. ## Rationale for this change Physical planning asks "is this ordering already satisfied?" constantly — sort removal, `EnforceSorting`, `EnforceDistribution`, and the requirement checks for windows, joins and aggregates all call into `EquivalenceProperties::ordering_satisfy`, `ordering_satisfy_requirement` and `extract_common_sort_prefix`. Each of those calls deep-clones the entire `EquivalenceProperties` — every equivalence class, every equivalent ordering, and the normalized ordering cache — before doing anything else, even when it never modifies the copy. The clone exists for a real reason: as the check walks a multi-key ordering left to right, it registers each satisfied key as a constant so the next key is evaluated within that key's tie group. That mutates state, so it needs its own copy. But two cases pay for it and get nothing back: 1. **A single-key check never mutates anything.** There is no "next key" to set up for, so the whole clone is wasted. This is the most common shape of these calls. 2. **The last key of *any* check registers constants nobody reads.** After the final key is verified, the code still calls `add_satisfied_key_constants`, which rebuilds the ordering cache and re-runs ordering discovery — and then the object is dropped. ## What changes are included in this PR? Two changes in `EquivalenceProperties`, to `ordering_satisfy_requirement` and `common_sort_prefix_length` (the latter backs `ordering_satisfy`, `extract_common_sort_prefix` and `reorder`): - **Clone on first write instead of up front.** The loop borrows `self` and clones only when it actually needs to register a constant. Single-key checks never clone at all. - **Skip the registration after the last key.** Nothing reads it. Plus a new criterion benchmark, `equivalence_properties`, covering these entry points. This only changes *when* the copy is made — the results of these functions are unchanged. ## Metrics Apple M4 Pro, rustc 1.97.0, criterion. All changes significant at p = 0.00. Properties under test: 3 equivalent orderings (`[c0,c1,c2,c3]`, `[c4,c5]`, `[c6]`) and a varying number of equivalence classes. **At 8 equivalence classes:** | benchmark | before | after | change | |---|---:|---:|---| | `ordering_satisfy` — 1 key | 2.72 µs | 0.41 µs | **−84.9%** | | `ordering_satisfy` — 1 key, unsatisfied | 1.47 µs | 0.41 µs | **−72.5%** | | `ordering_satisfy_requirement` — 1 key | 2.70 µs | 0.36 µs | **−86.3%** | | `ordering_satisfy_requirement` — 4 keys | 7.15 µs | 6.07 µs | −13.8% | | `ordering_satisfy` — 4 keys | 6.99 µs | 6.20 µs | −11.6% | | `extract_common_sort_prefix` — 4 keys | 7.22 µs | 6.38 µs | −9.2% | **How it scales** (`ordering_satisfy`, 1 key): | equivalence classes | before | after | change | |---:|---:|---:|---| | 2 | 2.44 µs | 0.43 µs | −82.5% | | 8 | 2.72 µs | 0.41 µs | −84.9% | | 32 | 4.63 µs | 0.41 µs | **−90.8%** | Reading the tables: for an *N*-key check the work goes from `1 clone + N registrations` to `(N > 1 ? 1 : 0) clones + (N − 1) registrations`. - **1-key checks** drop both the clone and the registration. Note the "after" column is flat at ~0.41 µs regardless of how many equivalence classes exist — with the clone gone, the check no longer scales with the size of the equivalence group at all. The "before" column does, which is why the win grows from −82% to −91%. - **Multi-key checks** still clone once and save one of *N* registrations. Since a registration rebuilds the ordering cache and re-runs ordering discovery, that single saved call is worth 9–14% here, rising to −37.8% for `4_keys` at 32 classes. ### Reproducing The benchmark is included in this PR, so reverting just the one source file gives you the baseline: ```bash # baseline: this PR's parent version of the file, with the new benchmark kept git checkout HEAD^ -- datafusion/physical-expr/src/equivalence/properties/mod.rs cargo bench -p datafusion-physical-expr --bench equivalence_properties -- --save-baseline before # with the change git checkout HEAD -- datafusion/physical-expr/src/equivalence/properties/mod.rs cargo bench -p datafusion-physical-expr --bench equivalence_properties -- --baseline before ``` The second run prints criterion's own `change: [...] (p = ...)` line per benchmark. ## Are these changes tested? No new correctness tests: this does not change what any of these functions return, so existing coverage is the right check. Covered by the `equivalence` unit tests in `datafusion/physical-expr` and, for plan-shape regressions, by sqllogictest — these functions decide whether a `SortExec` can be removed, so a behavior change would surface as a diff in an `EXPLAIN` plan. Full workspace suite (`--features avro,json,backtrace,extended_tests,recursive_protection,parquet_encryption`): 10,981 passed, 0 failed, and all 505 sqllogictest files pass. `./dev/rust_lint.sh` is clean. ## Are there any user-facing changes? No. No public API or behavior changes — planning is just faster. --- datafusion/physical-expr/Cargo.toml | 4 + .../benches/equivalence_properties.rs | 169 ++++++++++++++++++ .../src/equivalence/properties/mod.rs | 31 +++- 3 files changed, 199 insertions(+), 5 deletions(-) create mode 100644 datafusion/physical-expr/benches/equivalence_properties.rs diff --git a/datafusion/physical-expr/Cargo.toml b/datafusion/physical-expr/Cargo.toml index 0588a777230fb..9b0ae96c4f60f 100644 --- a/datafusion/physical-expr/Cargo.toml +++ b/datafusion/physical-expr/Cargo.toml @@ -105,3 +105,7 @@ name = "string_concat" [package.metadata.cargo-machete] ignored = ["half"] + +[[bench]] +harness = false +name = "equivalence_properties" diff --git a/datafusion/physical-expr/benches/equivalence_properties.rs b/datafusion/physical-expr/benches/equivalence_properties.rs new file mode 100644 index 0000000000000..6993b1056dedc --- /dev/null +++ b/datafusion/physical-expr/benches/equivalence_properties.rs @@ -0,0 +1,169 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Benchmarks for the ordering satisfaction checks on [`EquivalenceProperties`]. +//! +//! These are called repeatedly during physical optimization (sort removal, +//! `EnforceSorting`, `EnforceDistribution`, and the requirement checks for +//! windows, joins and aggregates), so their cost shows up directly in planning +//! time. +//! +//! # Scope +//! +//! These measure the satisfaction check itself, not the cost of assembling its +//! arguments. The sort expressions, requirements and orderings are built once, +//! up front. Because the checks take their input by value, each iteration gets a +//! fresh copy from the untimed setup step of `iter_batched`; only the call is +//! timed. Any copying the check does internally is part of what is measured. +//! +//! The benchmarks are parameterized by the number of equivalence classes, since +//! that -- not the schema width, which sits behind an `Arc` -- is what these +//! checks carry around. + +use std::sync::Arc; + +use arrow::compute::SortOptions; +use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use criterion::{BatchSize, BenchmarkId, Criterion, criterion_group, criterion_main}; +use datafusion_physical_expr::expressions::Column; +use datafusion_physical_expr::{ + EquivalenceProperties, LexOrdering, PhysicalExpr, PhysicalSortExpr, + PhysicalSortRequirement, +}; + +fn schema(n_cols: usize) -> SchemaRef { + Arc::new(Schema::new( + (0..n_cols) + .map(|i| Field::new(format!("c{i}"), DataType::Int32, true)) + .collect::>(), + )) +} + +fn col(i: usize) -> Arc { + Arc::new(Column::new(&format!("c{i}"), i)) +} + +fn asc(i: usize) -> PhysicalSortExpr { + PhysicalSortExpr::new(col(i), SortOptions::default()) +} + +/// Properties with three equivalent orderings and `n_classes` equivalence +/// classes, i.e. roughly what a scan feeding a join and a window function looks +/// like. Columns `c0..c7` carry the orderings; the equivalence classes are built +/// from the columns above them. +fn properties(n_classes: usize) -> EquivalenceProperties { + let schema = schema(8 + 2 * n_classes); + let mut props = EquivalenceProperties::new(schema); + props.add_orderings([ + vec![asc(0), asc(1), asc(2), asc(3)], + vec![asc(4), asc(5)], + vec![asc(6)], + ]); + for i in 0..n_classes { + props + .add_equal_conditions(col(8 + 2 * i), col(9 + 2 * i)) + .unwrap(); + } + props +} + +fn bench_ordering_satisfaction(c: &mut Criterion) { + let mut group = c.benchmark_group("equivalence_properties"); + + // Built once; see the "Scope" note at the top of this file. + let one_key = vec![asc(0)]; + // `c7` leads none of the orderings, so this exits on the first key. + let one_key_unsatisfied = vec![asc(7)]; + let four_keys = (0..4).map(asc).collect::>(); + let four_key_ordering = LexOrdering::new(four_keys.clone()).unwrap(); + let one_req = vec![PhysicalSortRequirement::new(col(0), None)]; + let four_reqs = (0..4) + .map(|i| PhysicalSortRequirement::new(col(i), None)) + .collect::>(); + + for n_classes in [2, 8, 32] { + let props = properties(n_classes); + + // A single sort key: the most common shape by far. + group.bench_function( + BenchmarkId::new("ordering_satisfy/1_key", n_classes), + |b| { + b.iter_batched( + || one_key.clone(), + |keys| props.ordering_satisfy(keys).unwrap(), + BatchSize::SmallInput, + ) + }, + ); + group.bench_function( + BenchmarkId::new("ordering_satisfy/1_key_unsatisfied", n_classes), + |b| { + b.iter_batched( + || one_key_unsatisfied.clone(), + |keys| props.ordering_satisfy(keys).unwrap(), + BatchSize::SmallInput, + ) + }, + ); + // Four sort keys: exercises the per-key constant registration. + group.bench_function( + BenchmarkId::new("ordering_satisfy/4_keys", n_classes), + |b| { + b.iter_batched( + || four_keys.clone(), + |keys| props.ordering_satisfy(keys).unwrap(), + BatchSize::SmallInput, + ) + }, + ); + group.bench_function( + BenchmarkId::new("ordering_satisfy_requirement/1_key", n_classes), + |b| { + b.iter_batched( + || one_req.clone(), + |reqs| props.ordering_satisfy_requirement(reqs).unwrap(), + BatchSize::SmallInput, + ) + }, + ); + group.bench_function( + BenchmarkId::new("ordering_satisfy_requirement/4_keys", n_classes), + |b| { + b.iter_batched( + || four_reqs.clone(), + |reqs| props.ordering_satisfy_requirement(reqs).unwrap(), + BatchSize::SmallInput, + ) + }, + ); + group.bench_function( + BenchmarkId::new("extract_common_sort_prefix/4_keys", n_classes), + |b| { + b.iter_batched( + || four_key_ordering.clone(), + |ordering| props.extract_common_sort_prefix(ordering).unwrap(), + BatchSize::SmallInput, + ) + }, + ); + } + + group.finish(); +} + +criterion_group!(benches, bench_ordering_satisfaction); +criterion_main!(benches); diff --git a/datafusion/physical-expr/src/equivalence/properties/mod.rs b/datafusion/physical-expr/src/equivalence/properties/mod.rs index 08c05efe0ccc0..c68157fecbd8c 100644 --- a/datafusion/physical-expr/src/equivalence/properties/mod.rs +++ b/datafusion/physical-expr/src/equivalence/properties/mod.rs @@ -640,8 +640,13 @@ impl EquivalenceProperties { return Ok(true); } let schema = self.schema(); - let mut eq_properties = self.clone(); - for element in normal_reqs { + // Registering satisfied keys as constants mutates the state, so it + // needs an owned copy -- but only from the second requirement onwards. + // Single-element requirements (the common case) never pay for the clone. + let last_idx = normal_reqs.len() - 1; + let mut owned = None::; + for (idx, element) in normal_reqs.into_iter().enumerate() { + let eq_properties = owned.as_ref().unwrap_or(self); // Check whether given requirement is satisfied: let ExprProperties { sort_properties, .. @@ -658,10 +663,16 @@ impl EquivalenceProperties { if !satisfy { return Ok(false); } + if idx == last_idx { + // Nothing left to check, so no need to update the state: + break; + } // Treat satisfied keys (and the sub-expressions they pin down) as // constants in subsequent iterations. See // [`Self::add_satisfied_key_constants`] for the rationale. - eq_properties.add_satisfied_key_constants(element.expr)?; + owned + .get_or_insert_with(|| self.clone()) + .add_satisfied_key_constants(element.expr)?; } Ok(true) } @@ -718,8 +729,12 @@ impl EquivalenceProperties { return Ok(full_length); } let schema = self.schema(); - let mut eq_properties = self.clone(); + // Registering satisfied keys as constants mutates the state, so it + // needs an owned copy -- but only from the second sort expression + // onwards. Single-element orderings never pay for the clone. + let mut owned = None::; for (idx, element) in normal_ordering.into_iter().enumerate() { + let eq_properties = owned.as_ref().unwrap_or(self); // Check whether given ordering is satisfied: let ExprProperties { sort_properties, .. @@ -739,10 +754,16 @@ impl EquivalenceProperties { // many we've satisfied so far: return Ok(idx); } + if idx + 1 == full_length { + // Nothing left to check, so no need to update the state: + break; + } // Treat satisfied keys (and the sub-expressions they pin down) as // constants in subsequent iterations. See // [`Self::add_satisfied_key_constants`] for the rationale. - eq_properties.add_satisfied_key_constants(Arc::clone(&element.expr))?; + owned + .get_or_insert_with(|| self.clone()) + .add_satisfied_key_constants(Arc::clone(&element.expr))?; } // All sort expressions are satisfied, return full length: Ok(full_length) From 864b6665710d517f5ce43af25d2948395329f204 Mon Sep 17 00:00:00 2001 From: Roshan Ramani <154859727+rawsun007@users.noreply.github.com> Date: Wed, 2 Sep 2026 15:01:20 +0000 Subject: [PATCH 15/37] docs: fix three dead links to datafusion.apache.org (#24872) ## Which issue does this PR close? - Closes #24871. ## Rationale for this change Three links in the docs point at pages on our own site that return 404, so a reader following them from the readings list or the GSoC project ideas page lands on an error. ## What changes are included in this PR? - `concepts-readings-events.md`: `blog/2026/03/20/multi-layer-pruning/` becomes `blog/2026/03/20/limit-pruning/`, and `blog/2026/02/02/case-expression/` becomes `blog/2026/02/02/datafusion_case/`. Those are the slugs the two posts were actually published under. - `gsoc/gsoc_project_ideas_2025.md`: `contributor-guide/gsoc_application_guidelines.html` becomes `contributor-guide/gsoc/gsoc_application_guidelines_2025.html`, matching the file the GSoC `index.rst` toctree includes. ## What is the testing strategy for this PR? No tests; this is three URL strings in markdown. Each old URL was requested and returns 404, each new one returns 200. ## Are there any user-facing changes? Docs only. Three links that were broken now resolve. Disclosure: this was written with AI assistance (Claude, via Claude Code). I verified every URL by request rather than inferring it, and the issue lists the three further dead links I deliberately did not touch because the right replacement is your call. --- docs/source/contributor-guide/gsoc/gsoc_project_ideas_2025.md | 2 +- docs/source/user-guide/concepts-readings-events.md | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/docs/source/contributor-guide/gsoc/gsoc_project_ideas_2025.md b/docs/source/contributor-guide/gsoc/gsoc_project_ideas_2025.md index d81d9eb9adab5..5366080ffe051 100644 --- a/docs/source/contributor-guide/gsoc/gsoc_project_ideas_2025.md +++ b/docs/source/contributor-guide/gsoc/gsoc_project_ideas_2025.md @@ -2,7 +2,7 @@ ## Introduction -Welcome to the Apache DataFusion Google Summer of Code (GSoC) 2025 project ideas list. Below you can find information about the projects. Please refer to [this page](https://datafusion.apache.org/contributor-guide/gsoc_application_guidelines.html) for application guidelines. +Welcome to the Apache DataFusion Google Summer of Code (GSoC) 2025 project ideas list. Below you can find information about the projects. Please refer to [this page](https://datafusion.apache.org/contributor-guide/gsoc/gsoc_application_guidelines_2025.html) for application guidelines. ## Projects diff --git a/docs/source/user-guide/concepts-readings-events.md b/docs/source/user-guide/concepts-readings-events.md index 366defa35eb5f..9edd83b33a4b7 100644 --- a/docs/source/user-guide/concepts-readings-events.md +++ b/docs/source/user-guide/concepts-readings-events.md @@ -53,13 +53,13 @@ This is a list of DataFusion related blog posts, articles, and other resources. - **2026-03-31** [Blog: Writing Custom Table Providers in Apache DataFusion](https://datafusion.apache.org/blog/2026/03/31/writing-table-providers/) -- **2026-03-20** [Blog: Turning LIMIT into an I/O Optimization: Inside DataFusion’s Multi-Layer Pruning Stack](https://datafusion.apache.org/blog/2026/03/20/multi-layer-pruning/) +- **2026-03-20** [Blog: Turning LIMIT into an I/O Optimization: Inside DataFusion’s Multi-Layer Pruning Stack](https://datafusion.apache.org/blog/2026/03/20/limit-pruning/) - **2026-02-23** [Blog: Apache DataFusion: A Data Engineer's Guide to the Query Engine Reshaping How We Build Data Systems](https://andrewmadson.substack.com/p/apache-datafusion-a-data-engineers) - **2026-02-09** [Blog: Vector search using only Parquet and DataFusion](https://blog.xiangpeng.systems/posts/vector-search-with-parquet-datafusion/) -- **2026-02-02** [Blog: Optimizing SQL CASE Expression Evaluation](https://datafusion.apache.org/blog/2026/02/02/case-expression/) +- **2026-02-02** [Blog: Optimizing SQL CASE Expression Evaluation](https://datafusion.apache.org/blog/2026/02/02/datafusion_case/) - **2026-01-12** [Blog: Extending SQL in DataFusion: from ->> to TABLESAMPLE](https://datafusion.apache.org/blog/2026/01/12/extending-sql) From d7b8e4fc1ea0a8d946886ade3847df172c85b234 Mon Sep 17 00:00:00 2001 From: MsfPablo <129399053+MsfPablo@users.noreply.github.com> Date: Wed, 2 Sep 2026 15:17:03 +0000 Subject: [PATCH 16/37] docs: fix 'sqllite' typo in sqllogictest README (#24825) ## Which issue does this PR close? N/A - trivial documentation fix. ## Rationale for this change `datafusion/sqllogictest/README.md` misspells "sqlite" as "sqllite" in the instructions for regenerating expected answers. The script it references is `regenerate_sqlite_files.sh`, so the surrounding prose should match. ## What changes are included in this PR? One-word spelling correction in `datafusion/sqllogictest/README.md`. ## Are these changes tested? No code changes; documentation only. ## Are there any user-facing changes? No. --- datafusion/sqllogictest/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/datafusion/sqllogictest/README.md b/datafusion/sqllogictest/README.md index 76c28b5ecd2ec..ca00fff7f6544 100644 --- a/datafusion/sqllogictest/README.md +++ b/datafusion/sqllogictest/README.md @@ -349,7 +349,7 @@ export RUST_MIN_STACK=30485760; PG_COMPAT=true INCLUDE_SQLITE=true cargo test --features=postgres --test sqllogictests ``` -To update the sqllite expected answers use the `datafusion/sqllogictest/regenerate_sqlite_files.sh` script. +To update the sqlite expected answers use the `datafusion/sqllogictest/regenerate_sqlite_files.sh` script. Note this must be run with an empty postgres instance. For example From d2b626cc93616b5bf80b7ca2a079e9859d992e32 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Wed, 2 Sep 2026 16:14:28 +0000 Subject: [PATCH 17/37] chore: update chacha20 dependency (#24880) This is a *very* minor update to the Cargo.lock file to resolve a warning when running `cargo audit` as in our CI suite. v0.10.0 was yanked from crates.io so this bumps to 0.10.2. --- Cargo.lock | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 3077074e862fa..293ac45c51e38 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1244,9 +1244,9 @@ checksum = "613afe47fcd5fac7ccf1db93babcb082c5994d996f20b8b159f2ad1658eb5724" [[package]] name = "chacha20" -version = "0.10.0" +version = "0.10.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6f8d983286843e49675a4b7a2d174efe136dc93a18d69130dd18198a6c167601" +checksum = "65c35e4b699c7e15ccbe7ee35c005e4fc0a278d22238a2857e6ce2dadeda1b06" dependencies = [ "cfg-if", "cpufeatures 0.3.0", From dcfb144974814aba31a71b733ff5720fa110c287 Mon Sep 17 00:00:00 2001 From: Stefan Wang <1fannnw@gmail.com> Date: Wed, 2 Sep 2026 20:38:23 +0000 Subject: [PATCH 18/37] Prevent array_position panic for minimum start positions (#24838) ## Which issue does this PR close? - Closes https://github.com/apache/datafusion/issues/22220. ## Rationale for this change `array_position` accepts an optional one-based start position. Passing `-9223372036854775808` currently aborts evaluation with: ```text datafusion/functions-nested/src/position.rs:210:21: attempt to subtract with overflow ``` It now returns the normal execution error `start_from out of bounds: -9223372036854775808`. ## What changes are included in this PR? The one-based-to-zero-based conversion now uses checked subtraction in both optimized and generic execution paths. ## What is the testing strategy for this PR? SQL regressions exercise a scalar start position, a column start position with a scalar needle, and column start and needle values.
Raw results ```console $ git checkout upstream/main -- datafusion/functions-nested/src/position.rs $ cargo test -p datafusion-sqllogictest --test sqllogictests -- array_position thread 'tokio-rt-worker' panicked at datafusion/functions-nested/src/position.rs:210:21: attempt to subtract with overflow Error: Execution("1 failures") $ git checkout HEAD -- datafusion/functions-nested/src/position.rs $ cargo test -p datafusion-sqllogictest --test sqllogictests -- array_position Running with 12 test threads (available parallelism: 12) Progress: 1/1 files completed (100%) $ RUST_BACKTRACE=1 cargo test --profile ci \ --exclude datafusion-examples --exclude datafusion-benchmarks --exclude datafusion-cli \ --workspace --lib --tests --bins \ --features avro,json,backtrace,extended_tests,recursive_protection,parquet_encryption passed=10997 ignored=8 ```
## Are there any user-facing changes? Yes. An invalid minimum `Int64` start position returns an execution error instead of panicking. --------- Signed-off-by: 1fanwang <1fannnw@gmail.com> --- datafusion/functions-nested/src/position.rs | 28 ++++++++++++++----- .../test_files/array/array_position.slt | 11 ++++++++ 2 files changed, 32 insertions(+), 7 deletions(-) diff --git a/datafusion/functions-nested/src/position.rs b/datafusion/functions-nested/src/position.rs index f677cfb4979d3..9ef8142fa48dd 100644 --- a/datafusion/functions-nested/src/position.rs +++ b/datafusion/functions-nested/src/position.rs @@ -40,7 +40,9 @@ use arrow::array::{ use datafusion_common::cast::{ as_generic_list_array, as_int64_array, as_large_list_array, as_list_array, }; -use datafusion_common::{Result, exec_err, utils::take_function_args}; +use datafusion_common::{ + Result, exec_datafusion_err, exec_err, utils::take_function_args, +}; use itertools::Itertools; use crate::utils::{compare_element_to_list, make_scalar_function}; @@ -198,6 +200,16 @@ fn array_position_inner(args: &[ArrayRef]) -> Result { } } +fn resolve_zero_based_start_from(start_from: i64) -> Result { + start_from.checked_sub(1).ok_or_else(|| { + exec_datafusion_err!( + "start_from out of bounds: {start_from}, expected {} to {}", + i64::MIN + 1, + i64::MAX + ) + }) +} + /// Resolves the optional `start_from` argument into a `Vec` of /// 0-indexed starting positions. fn resolve_start_from( @@ -207,14 +219,16 @@ fn resolve_start_from( match third_arg { None => Ok(vec![0i64; num_rows]), Some(ColumnarValue::Scalar(ScalarValue::Int64(Some(v)))) => { - Ok(vec![v - 1; num_rows]) + Ok(vec![resolve_zero_based_start_from(*v)?; num_rows]) } Some(ColumnarValue::Scalar(s)) => { exec_err!("array_position expected Int64 for start_from, got {s}") } - Some(ColumnarValue::Array(a)) => { - Ok(as_int64_array(a)?.values().iter().map(|&x| x - 1).collect()) - } + Some(ColumnarValue::Array(a)) => as_int64_array(a)? + .values() + .iter() + .map(|&x| resolve_zero_based_start_from(x)) + .collect(), } } @@ -309,8 +323,8 @@ fn general_position_dispatch(args: &[ArrayRef]) -> Result>() + .map(|&x| resolve_zero_based_start_from(x)) + .collect::>>()? } else { vec![0; haystack.len()] }; diff --git a/datafusion/sqllogictest/test_files/array/array_position.slt b/datafusion/sqllogictest/test_files/array/array_position.slt index e3dd830dfb77a..8fe1826619431 100644 --- a/datafusion/sqllogictest/test_files/array/array_position.slt +++ b/datafusion/sqllogictest/test_files/array/array_position.slt @@ -282,6 +282,17 @@ select array_position([1, 2, 3], 3, 4), array_position([1], 1, 2); ---- NULL NULL +query error start_from out of bounds: -9223372036854775808, expected -9223372036854775807 to 9223372036854775807 +select array_position([1], 1, -9223372036854775808); + +query error start_from out of bounds: -9223372036854775808, expected -9223372036854775807 to 9223372036854775807 +select array_position([1], 1, start_from) +from (values (-9223372036854775808)) as t(start_from); + +query error start_from out of bounds: -9223372036854775808, expected -9223372036854775807 to 9223372036854775807 +select array_position([1], needle, start_from) +from (values (1, -9223372036854775808)) as t(needle, start_from); + # array_position with empty array in various contexts query II select array_position(arrow_cast(make_array(), 'List(Int64)'), 1), array_position(arrow_cast(make_array(), 'LargeList(Int64)'), 1); From d02bd786de49f2ad3dc15cb947608fb0f01fe142 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Wed, 2 Sep 2026 21:06:58 +0000 Subject: [PATCH 19/37] chore(deps-dev): bump fast-uri from 3.1.5 to 3.1.7 in /datafusion/wasmtest/datafusion-wasm-app (#24882) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Bumps [fast-uri](https://github.com/fastify/fast-uri) from 3.1.5 to 3.1.7.
Release notes

Sourced from fast-uri's releases.

v3.1.7

⚠️ Security Warning

This is a security release that fixes the following high-severity security advisories:

Users of the v3.x release line should upgrade to v3.1.7.

Full Changelog: https://github.com/fastify/fast-uri/compare/v3.1.6...v3.1.7

v3.1.6

⚠️ Security Warning

This release addresses the following high-severity security advisories:

Users of the v3.x release line should upgrade to v3.1.6.

Full Changelog: https://github.com/fastify/fast-uri/compare/v3.1.5...v3.1.6

Commits
  • 412e40a Bumped v3.1.7
  • 9f4c943 fix: backport port and IP-literal validation to v3.x (#216)
  • 1eb3ce4 fix: treat unterminated bracket hosts as reg-names again (#214)
  • 6f970b2 Bumped v3.1.6
  • d941579 fix: never run IDN canonicalization on bracketed IP literals
  • c0f0279 test: adapt decoded-scheme handler assertion to 3.x (no mailto scheme)
  • 37f3417 Merge commit from fork
  • 607bfbe Merge commit from fork
  • ae92a4c Merge commit from fork
  • 444ecda Merge commit from fork
  • Additional commits viewable in compare view

[![Dependabot compatibility score](https://dependabot-badges.githubapp.com/badges/compatibility_score?dependency-name=fast-uri&package-manager=npm_and_yarn&previous-version=3.1.5&new-version=3.1.7)](https://docs.github.com/en/github/managing-security-vulnerabilities/about-dependabot-security-updates#about-compatibility-scores) Dependabot will resolve any conflicts with this PR as long as you don't alter it yourself. You can also trigger a rebase manually by commenting `@dependabot rebase`. [//]: # (dependabot-automerge-start) [//]: # (dependabot-automerge-end) ---
Dependabot commands and options
You can trigger Dependabot actions by commenting on this PR: - `@dependabot rebase` will rebase this PR - `@dependabot recreate` will recreate this PR, overwriting any edits that have been made to it - `@dependabot show ignore conditions` will show all of the ignore conditions of the specified dependency - `@dependabot ignore this major version` will close this PR and stop Dependabot creating any more for this major version (unless you reopen the PR or upgrade to it yourself) - `@dependabot ignore this minor version` will close this PR and stop Dependabot creating any more for this minor version (unless you reopen the PR or upgrade to it yourself) - `@dependabot ignore this dependency` will close this PR and stop Dependabot creating any more for this dependency (unless you reopen the PR or upgrade to it yourself) You can disable automated security fix PRs for this repo from the [Security Alerts page](https://github.com/apache/datafusion/network/alerts).
Signed-off-by: dependabot[bot] Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> --- .../datafusion-wasm-app/package-lock.json | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/datafusion/wasmtest/datafusion-wasm-app/package-lock.json b/datafusion/wasmtest/datafusion-wasm-app/package-lock.json index c0e1d2e6abf14..09422c7001a55 100644 --- a/datafusion/wasmtest/datafusion-wasm-app/package-lock.json +++ b/datafusion/wasmtest/datafusion-wasm-app/package-lock.json @@ -1482,9 +1482,9 @@ "dev": true }, "node_modules/fast-uri": { - "version": "3.1.5", - "resolved": "https://registry.npmjs.org/fast-uri/-/fast-uri-3.1.5.tgz", - "integrity": "sha512-gHwA1O9LDIcKunMKhObS/HimwtehO1nPUECKAu5TpKgaO19fcWEl4bliWe1jWxVFvIXztJjjQ4L8XQ1EU9f7Jw==", + "version": "3.1.7", + "resolved": "https://registry.npmjs.org/fast-uri/-/fast-uri-3.1.7.tgz", + "integrity": "sha512-dOvZVzjdZdz7phd9v6jCbwxrBW3fK6n8Rc0CtdmM4bumzMnxywBYhuph6J819RRw/ku+rLbelwfMunktuzVVHg==", "dev": true, "funding": [ { @@ -1495,7 +1495,8 @@ "type": "opencollective", "url": "https://opencollective.com/fastify" } - ] + ], + "license": "BSD-3-Clause" }, "node_modules/fill-range": { "version": "7.1.1", @@ -5216,9 +5217,9 @@ "dev": true }, "fast-uri": { - "version": "3.1.5", - "resolved": "https://registry.npmjs.org/fast-uri/-/fast-uri-3.1.5.tgz", - "integrity": "sha512-gHwA1O9LDIcKunMKhObS/HimwtehO1nPUECKAu5TpKgaO19fcWEl4bliWe1jWxVFvIXztJjjQ4L8XQ1EU9f7Jw==", + "version": "3.1.7", + "resolved": "https://registry.npmjs.org/fast-uri/-/fast-uri-3.1.7.tgz", + "integrity": "sha512-dOvZVzjdZdz7phd9v6jCbwxrBW3fK6n8Rc0CtdmM4bumzMnxywBYhuph6J819RRw/ku+rLbelwfMunktuzVVHg==", "dev": true }, "fill-range": { From 7a1f468643f5fcc33166b05a6cf2210675cab449 Mon Sep 17 00:00:00 2001 From: discord9 Date: Thu, 3 Sep 2026 03:02:13 +0000 Subject: [PATCH 20/37] fix: preserve nested projection expressions (#24686) ## Which issue does this PR close? - No issue has been filed. ## Rationale for this change Anonymous nested projections that reuse the same output name can return wrong results. For example, each `i + 1 AS i` layer must be evaluated independently, but the projection optimizer could drop one layer when two consecutive projection expression vectors were structurally equal. Structural equality does not imply that a projection is safe to elide: repeated computations such as `i + 1 AS i` have the same expression shape but must still be evaluated twice. ## What changes are included in this PR? - Remove the structural-equality fast path that directly elided one of two consecutive projections. - Keep the existing iterative whole-chain merge, so deep projection chains still collapse within one optimizer rule invocation. - Continue using the normal projection rewrite path, which composes repeated expressions and preserves aliases and field metadata. - Add focused regressions for repeated non-idempotent projections, one-pass collapse of a 12-level chain, and metadata-bearing aliases. - Add an execution-level SQLLogicTest with six anonymous `i + 1 AS i` layers under both `max_passes = 1` and the default optimizer configuration. ## What is the testing strategy for this PR? The focused unit tests verify that: - two structurally equal `i + 1 AS i` projections retain both additions; - a 12-level chain preserves all 12 additions and collapses to one `Projection` with `max_passes = 1`; - a metadata-bearing `Alias(Column)` is merged without losing field metadata. The SQLLogicTest executes a six-level anonymous projection chain against a temporary table with both one optimizer pass and the default pass count. Verified with: ```text cargo fmt --all --check cargo test -p datafusion-optimizer optimize_projections # 58 passed; 0 failed cargo test -p datafusion-optimizer --test optimizer_integration # 26 passed; 0 failed cargo test --profile ci -p datafusion-sqllogictest --test sqllogictests -- projection.slt # 1/1 files completed; 0 failures cargo clippy -p datafusion-optimizer --all-targets --all-features -- -D warnings # passed git diff --check # passed ``` ## Are there any user-facing changes? Yes. Deep anonymous nested projections that reuse an output name now preserve every projection expression and return the correct result. There are no public API or configuration changes. --------- Signed-off-by: discord9 Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- .../optimizer/src/optimize_projections/mod.rs | 74 ++++++++++++++++--- .../sqllogictest/test_files/projection.slt | 68 +++++++++++++++++ 2 files changed, 131 insertions(+), 11 deletions(-) diff --git a/datafusion/optimizer/src/optimize_projections/mod.rs b/datafusion/optimizer/src/optimize_projections/mod.rs index b0543a871f52b..9604bf2f188c0 100644 --- a/datafusion/optimizer/src/optimize_projections/mod.rs +++ b/datafusion/optimizer/src/optimize_projections/mod.rs @@ -556,17 +556,6 @@ fn merge_consecutive_projections_one_level( return Projection::try_new_with_schema(expr, input, schema).map(Transformed::no); }; - // A fast path: if the previous projection is same as the current projection - // we can directly remove the current projection and return child projection. - if prev_projection.expr == expr { - return Projection::try_new_with_schema( - expr, - Arc::clone(&prev_projection.input), - schema, - ) - .map(Transformed::yes); - } - // Count usages (referrals) of each projection expression in its input fields: let mut column_referral_map = HashMap::<&Column, usize>::new(); for expr in &expr { @@ -1210,6 +1199,69 @@ mod tests { ) } + #[test] + fn merge_structurally_equal_non_idempotent_projections() -> Result<()> { + let schema = Schema::new(vec![Field::new("i", DataType::Int32, false)]); + let projection = || col("i").add(lit(1)).alias("i"); + let plan = table_scan(TableReference::none(), &schema, None)? + .project(vec![projection()])? + .project(vec![projection()])? + .build()?; + + assert_optimized_plan_equal!( + plan, + @r" + Projection: ?table?.i + Int32(1) + Int32(1) AS i + TableScan: ?table? projection=[i] + " + ) + } + + #[test] + fn merge_deep_projection_chain_in_one_pass() -> Result<()> { + let schema = Schema::new(vec![Field::new("i", DataType::Int32, false)]); + let mut plan = table_scan(TableReference::none(), &schema, None)? + .project(vec![col("i").add(lit(1)).alias("i")])? + .build()?; + for _ in 1..12 { + plan = LogicalPlanBuilder::from(plan) + .project(vec![col("i").add(lit(1)).alias("i")])? + .build()?; + } + + let optimizer = Optimizer::with_rules(vec![Arc::new(OptimizeProjections::new())]); + let optimized = optimizer.optimize( + plan, + &OptimizerContext::new().with_max_passes(1), + observe, + )?; + let plan_string = format!("{optimized}"); + assert_eq!(12, plan_string.matches("Int32(1)").count()); + assert_eq!(1, plan_string.matches("Projection:").count()); + Ok(()) + } + + #[test] + fn merge_columns_and_metadata_alias() -> Result<()> { + let metadata = + datafusion_common::metadata::FieldMetadata::from(HashMap::from([( + "key".to_string(), + "value".to_string(), + )])); + let plan = LogicalPlanBuilder::from(test_table_scan()?) + .project(vec![col("a")])? + .project(vec![col("a").alias_with_metadata("a", Some(metadata))])? + .build()?; + + let optimized = optimize(plan)?; + assert_eq!( + "value", + optimized.schema().field(0).metadata().get("key").unwrap() + ); + assert_eq!(1, format!("{optimized}").matches("Projection:").count()); + Ok(()) + } + #[test] fn merge_alias() -> Result<()> { let table_scan = test_table_scan()?; diff --git a/datafusion/sqllogictest/test_files/projection.slt b/datafusion/sqllogictest/test_files/projection.slt index e18114bc51ca8..bac94670f4ffd 100644 --- a/datafusion/sqllogictest/test_files/projection.slt +++ b/datafusion/sqllogictest/test_files/projection.slt @@ -217,6 +217,74 @@ SELECT column1 as a from (values (1), (2)) f where f.column1 = 2; ---- 2 +# Regression: preserve deep anonymous nested projections with one and default passes. +statement ok +CREATE TABLE nested_projection(i INT); + +statement ok +INSERT INTO nested_projection VALUES (3), (4), (5); + +statement ok +SET datafusion.optimizer.max_passes = 1; + +query I rowsort +SELECT i +FROM ( + SELECT i + 1 AS i + FROM ( + SELECT i + 1 AS i + FROM ( + SELECT i + 1 AS i + FROM ( + SELECT i + 1 AS i + FROM ( + SELECT i + 1 AS i + FROM ( + SELECT i + 1 AS i + FROM nested_projection + ) + ) + ) + ) + ) +) +---- +10 +11 +9 + +statement ok +RESET datafusion.optimizer.max_passes; + +query I rowsort +SELECT i +FROM ( + SELECT i + 1 AS i + FROM ( + SELECT i + 1 AS i + FROM ( + SELECT i + 1 AS i + FROM ( + SELECT i + 1 AS i + FROM ( + SELECT i + 1 AS i + FROM ( + SELECT i + 1 AS i + FROM nested_projection + ) + ) + ) + ) + ) +) +---- +10 +11 +9 + +statement ok +DROP TABLE nested_projection; + # clean data statement ok DROP TABLE aggregate_simple; From 6ffc45d0815f5f1e0b912cf5f42b34a5ab236dce Mon Sep 17 00:00:00 2001 From: Bruno Volpato Date: Thu, 3 Sep 2026 03:10:08 +0000 Subject: [PATCH 21/37] fix(substrait): preserve grouping set output order (#23468) ## Which issue does this PR close? - Closes #17910. ## Rationale for this change For multiple grouping sets, Substrait orders aggregate output as grouping expressions, measures, then grouping-set ID. DataFusion puts its internal `__grouping_id` before measures. The consumer applied `RelCommon.emit` to DataFusion's order, so a mapping such as `[0, 1, 2]` returned `__grouping_id` where the plan requested its first measure. See the [Substrait AggregateRel output mapping specification](https://substrait.io/relations/logical_relations/#aggregate-operation). ## What changes are included in this PR? - Reorder consumer output to Substrait's direct order before applying emit mappings. - Emit a producer mapping back to DataFusion's aggregate order. - Deduplicate producer grouping expressions used by grouping references. - Test both issue reproducer and two-measure roundtrip. ## Are these changes tested? - `cargo test -p datafusion-substrait --features protoc` - `cargo clippy --all-targets --all-features -- -D warnings` ## Are there any user-facing changes? Substrait plans with multiple grouping sets now return emitted columns in specified order. No public API change. --- .../consumer/rel/aggregate_rel.rs | 48 ++++++- .../producer/rel/aggregate_rel.rs | 31 ++++- .../tests/cases/aggregation_tests.rs | 38 ++++++ .../tests/cases/roundtrip_logical_plan.rs | 55 +++++++- .../multiple_groupings.json | 127 ++++++++++++++++++ 5 files changed, 289 insertions(+), 10 deletions(-) create mode 100644 datafusion/substrait/tests/testdata/test_plans/aggregate_groupings/multiple_groupings.json diff --git a/datafusion/substrait/src/logical_plan/consumer/rel/aggregate_rel.rs b/datafusion/substrait/src/logical_plan/consumer/rel/aggregate_rel.rs index 8c0114b90ee13..982a87d6d5e83 100644 --- a/datafusion/substrait/src/logical_plan/consumer/rel/aggregate_rel.rs +++ b/datafusion/substrait/src/logical_plan/consumer/rel/aggregate_rel.rs @@ -17,8 +17,11 @@ use crate::logical_plan::consumer::{NameTracker, SubstraitConsumer}; use crate::logical_plan::consumer::{from_substrait_agg_func, from_substrait_sorts}; -use datafusion::common::{DFSchemaRef, not_impl_err}; -use datafusion::logical_expr::{Expr, GroupingSet, LogicalPlan, LogicalPlanBuilder}; +use datafusion::common::{Column, DFSchemaRef, internal_err, not_impl_err}; +use datafusion::logical_expr::builder::project; +use datafusion::logical_expr::{ + Aggregate, Expr, GroupingSet, LogicalPlan, LogicalPlanBuilder, +}; use substrait::proto::AggregateRel; use substrait::proto::aggregate_function::AggregationInvocation; use substrait::proto::aggregate_rel::Grouping; @@ -122,12 +125,51 @@ pub async fn from_aggregate_rel( .map(|e| name_tracker.get_uniquely_named_expr(e)) .collect::, _>>()?; - input.aggregate(group_exprs, aggr_exprs)?.build() + let plan = input.aggregate(group_exprs, aggr_exprs)?.build()?; + if agg.groupings.len() > 1 { + reorder_grouping_set_output(plan, agg.measures.len()) + } else { + Ok(plan) + } } else { not_impl_err!("Aggregate without an input is not valid") } } +/// Reorders DataFusion's `[groups, grouping_id, measures]` aggregate schema to +/// Substrait's direct output order of `[groups, measures, grouping_id]`. +fn reorder_grouping_set_output( + plan: LogicalPlan, + measure_count: usize, +) -> datafusion::common::Result { + let exprs: Vec = { + let schema = plan.schema(); + let Some(grouping_id_index) = + schema.index_of_column_by_name(None, Aggregate::INTERNAL_GROUPING_ID) + else { + return internal_err!( + "Grouping set aggregate schema is missing {}", + Aggregate::INTERNAL_GROUPING_ID + ); + }; + if grouping_id_index + measure_count + 1 != schema.fields().len() { + return internal_err!( + "Grouping set aggregate schema has {} fields after {}, expected {} measures", + schema.fields().len() - grouping_id_index - 1, + Aggregate::INTERNAL_GROUPING_ID, + measure_count + ); + } + + (0..grouping_id_index) + .chain(grouping_id_index + 1..schema.fields().len()) + .chain(std::iter::once(grouping_id_index)) + .map(|index| Expr::Column(Column::from(schema.qualified_field(index)))) + .collect() + }; + project(plan, exprs) +} + #[expect(deprecated)] async fn from_substrait_grouping( consumer: &impl SubstraitConsumer, diff --git a/datafusion/substrait/src/logical_plan/producer/rel/aggregate_rel.rs b/datafusion/substrait/src/logical_plan/producer/rel/aggregate_rel.rs index dec94b0422257..7b6c113ccec0f 100644 --- a/datafusion/substrait/src/logical_plan/producer/rel/aggregate_rel.rs +++ b/datafusion/substrait/src/logical_plan/producer/rel/aggregate_rel.rs @@ -24,7 +24,8 @@ use datafusion::logical_expr::utils::powerset; use datafusion::logical_expr::{Aggregate, Distinct, Expr, GroupingSet}; use substrait::proto::aggregate_rel::{Grouping, Measure}; use substrait::proto::rel::RelType; -use substrait::proto::{AggregateRel, Expression, Rel}; +use substrait::proto::rel_common::EmitKind; +use substrait::proto::{AggregateRel, Expression, Rel, RelCommon, rel_common}; pub fn from_aggregate( producer: &mut impl SubstraitProducer, @@ -38,10 +39,12 @@ pub fn from_aggregate( .iter() .map(|e| to_substrait_agg_measure(producer, e, agg.input.schema())) .collect::>>()?; + let common = (groupings.len() > 1) + .then(|| grouping_set_output_mapping(grouping_expressions.len(), measures.len())); Ok(Box::new(Rel { rel_type: Some(RelType::Aggregate(Box::new(AggregateRel { - common: None, + common, input: Some(input), grouping_expressions, groupings, @@ -51,6 +54,22 @@ pub fn from_aggregate( })) } +/// Maps Substrait's `[groups, measures, grouping_id]` direct output to +/// DataFusion's `[groups, grouping_id, measures]` aggregate schema. +fn grouping_set_output_mapping(grouping_count: usize, measure_count: usize) -> RelCommon { + let grouping_id_index = grouping_count + measure_count; + let output_mapping = (0..grouping_count) + .chain(std::iter::once(grouping_id_index)) + .chain(grouping_count..grouping_id_index) + .map(|index| index as i32) + .collect(); + RelCommon { + emit_kind: Some(EmitKind::Emit(rel_common::Emit { output_mapping })), + hint: None, + advanced_extension: None, + } +} + pub fn from_distinct( producer: &mut impl SubstraitProducer, distinct: &Distinct, @@ -165,8 +184,12 @@ pub fn parse_flat_grouping_exprs( for e in exprs { let rex = producer.handle_expr(e, schema)?; grouping_expressions.push(rex.clone()); - ref_group_exprs.push(rex); - expression_references.push((ref_group_exprs.len() - 1) as u32); + let reference = ref_group_exprs.iter().position(|existing| existing == &rex); + let reference = reference.unwrap_or_else(|| { + ref_group_exprs.push(rex); + ref_group_exprs.len() - 1 + }); + expression_references.push(reference as u32); } #[expect(deprecated)] Ok(Grouping { diff --git a/datafusion/substrait/tests/cases/aggregation_tests.rs b/datafusion/substrait/tests/cases/aggregation_tests.rs index 92a41850b208d..e572023f17a92 100644 --- a/datafusion/substrait/tests/cases/aggregation_tests.rs +++ b/datafusion/substrait/tests/cases/aggregation_tests.rs @@ -68,4 +68,42 @@ mod tests { Ok(()) } + + #[tokio::test] + async fn multiple_grouping_sets_follow_substrait_output_order() -> Result<()> { + let proto_plan = read_json( + "tests/testdata/test_plans/aggregate_groupings/multiple_groupings.json", + ); + let ctx = add_plan_schemas_to_ctx(SessionContext::new(), &proto_plan)?; + let plan = from_substrait_plan(&ctx.state(), &proto_plan).await?; + + assert_snapshot!( + plan, + @r" + Projection: c0, c1, sum(c0) AS summation + Aggregate: groupBy=[[GROUPING SETS ((c0), (c1), (c0, c1))]], aggr=[[sum(c0)]] + Values: (Int64(1), Int64(10)), (Int64(1), Int64(20)), (Int64(2), Int64(10)) + " + ); + + let results = DataFrame::new(ctx.state(), plan).collect().await?; + datafusion::assert_batches_sorted_eq!( + [ + "+----+----+-----------+", + "| c0 | c1 | summation |", + "+----+----+-----------+", + "| | 10 | 3 |", + "| | 20 | 1 |", + "| 1 | | 2 |", + "| 1 | 10 | 1 |", + "| 1 | 20 | 1 |", + "| 2 | | 2 |", + "| 2 | 10 | 2 |", + "+----+----+-----------+", + ], + &results + ); + + Ok(()) + } } diff --git a/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs b/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs index c9f874dd9b095..3716b0feba3cc 100644 --- a/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs +++ b/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs @@ -344,10 +344,59 @@ async fn aggregate_multiple_keys() -> Result<()> { #[tokio::test] async fn aggregate_grouping_sets() -> Result<()> { - roundtrip( - "SELECT a, c, d, avg(b) FROM data GROUP BY GROUPING SETS ((a, c), (a), (d), ())", + let ctx = create_context().await?; + let proto = roundtrip_with_ctx( + "SELECT a, c, d, avg(b), sum(e) FROM data GROUP BY GROUPING SETS ((a, c), (a), (d), ())", + ctx.clone(), ) - .await + .await?; + + let Some(plan_rel::RelType::Root(root)) = &proto.relations[0].rel_type else { + panic!("expected root relation"); + }; + let Some(RelType::Project(project)) = &root.input.as_ref().unwrap().rel_type else { + panic!("expected project relation"); + }; + let Some(RelType::Aggregate(aggregate)) = &project.input.as_ref().unwrap().rel_type + else { + panic!("expected aggregate relation"); + }; + + assert_eq!(aggregate.grouping_expressions.len(), 3); + assert_eq!(aggregate.groupings[0].expression_references, [0, 1]); + assert_eq!(aggregate.groupings[1].expression_references, [0]); + assert_eq!(aggregate.groupings[2].expression_references, [2]); + assert!(aggregate.groupings[3].expression_references.is_empty()); + let output_mapping = match aggregate + .common + .as_ref() + .and_then(|common| common.emit_kind.as_ref()) + { + Some(substrait::proto::rel_common::EmitKind::Emit(emit)) => &emit.output_mapping, + _ => panic!("expected aggregate output mapping"), + }; + assert_eq!(output_mapping, &[0, 1, 2, 5, 3, 4]); + + let plan = from_substrait_plan(&ctx.state(), &proto).await?; + let results = DataFrame::new(ctx.state(), plan).collect().await?; + datafusion::assert_batches_sorted_eq!( + [ + "+---+------------+-------+-------------+-------------+", + "| a | c | d | avg(data.b) | sum(data.e) |", + "+---+------------+-------+-------------+-------------+", + "| | | | 3.250000 | 6442450943 |", + "| | | false | 2.000000 | 4294967295 |", + "| | | true | 4.500000 | 2147483648 |", + "| 1 | | | 2.000000 | 4294967295 |", + "| 1 | 2020-01-01 | | 2.000000 | 4294967295 |", + "| 3 | | | 4.500000 | 2147483648 |", + "| 3 | 2020-01-01 | | 4.500000 | 2147483648 |", + "+---+------------+-------+-------------+-------------+", + ], + &results + ); + + Ok(()) } #[tokio::test] diff --git a/datafusion/substrait/tests/testdata/test_plans/aggregate_groupings/multiple_groupings.json b/datafusion/substrait/tests/testdata/test_plans/aggregate_groupings/multiple_groupings.json new file mode 100644 index 0000000000000..3779ee32c6ab8 --- /dev/null +++ b/datafusion/substrait/tests/testdata/test_plans/aggregate_groupings/multiple_groupings.json @@ -0,0 +1,127 @@ +{ + "extensionUris": [ + { + "extensionUriAnchor": 1, + "uri": "https://github.com/substrait-io/substrait/blob/main/extensions/functions_arithmetic.yaml" + } + ], + "extensions": [ + { + "extensionFunction": { + "extensionUriReference": 1, + "functionAnchor": 1, + "name": "sum:i64" + } + } + ], + "relations": [ + { + "root": { + "input": { + "aggregate": { + "common": { + "emit": { + "outputMapping": [0, 1, 2] + } + }, + "input": { + "read": { + "baseSchema": { + "names": ["c0", "c1"], + "struct": { + "nullability": "NULLABILITY_REQUIRED", + "types": [ + { + "i64": { + "nullability": "NULLABILITY_REQUIRED" + } + }, + { + "i64": { + "nullability": "NULLABILITY_REQUIRED" + } + } + ] + } + }, + "common": { + "direct": {} + }, + "virtualTable": { + "expressions": [ + { + "fields": [ + { "literal": { "i64": "1", "nullable": false } }, + { "literal": { "i64": "10", "nullable": false } } + ] + }, + { + "fields": [ + { "literal": { "i64": "1", "nullable": false } }, + { "literal": { "i64": "20", "nullable": false } } + ] + }, + { + "fields": [ + { "literal": { "i64": "2", "nullable": false } }, + { "literal": { "i64": "10", "nullable": false } } + ] + } + ] + } + } + }, + "groupingExpressions": [ + { + "selection": { + "directReference": { "structField": {} }, + "rootReference": {} + } + }, + { + "selection": { + "directReference": { "structField": { "field": 1 } }, + "rootReference": {} + } + } + ], + "groupings": [ + { "expressionReferences": [0] }, + { "expressionReferences": [1] }, + { "expressionReferences": [0, 1] } + ], + "measures": [ + { + "measure": { + "arguments": [ + { + "value": { + "selection": { + "directReference": { "structField": {} }, + "rootReference": {} + } + } + } + ], + "functionReference": 1, + "invocation": "AGGREGATION_INVOCATION_ALL", + "outputType": { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + "phase": "AGGREGATION_PHASE_INITIAL_TO_RESULT" + } + } + ] + } + }, + "names": ["c0", "c1", "summation"] + } + } + ], + "version": { + "minorNumber": 29, + "producer": "substrait-go v4.2.0" + } +} From 17115f368dbdf0f21428c0991c3f8ec11f96c090 Mon Sep 17 00:00:00 2001 From: Subham Singhal Date: Thu, 3 Sep 2026 03:44:55 +0000 Subject: [PATCH 22/37] bench(h2o): window Top-N sweep over a declared-sorted input (#24732) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Which issue does this PR close? Related to #6899 — adds benchmark coverage for the WindowTopN operators. ## Rationale for this change The existing `h2o --subgroup window` Top-N sweep (q13–q29) registers `x` with no declared ordering, so `output_ordering()` is `None`. Any plan that depends on the input being sorted is unreachable from those queries, however the data happens to be laid out on disk. This adds a `window_sorted` subgroup that publishes the same sweep over a `WITH ORDER` table, so ordering-dependent plans can be measured. ## What changes are included in this PR? A `window_sorted` h2o subgroup: 16 queries covering ROW_NUMBER / RANK / DENSE_RANK × 100 / 1K / 10K / 100K partitions, plus heavy-ties variants. - Two `load` directives: the existing `load_window_${SIZE}_${FORMAT}.sql` creates `x`, then a new script writes a sorted copy via `COPY (... ORDER BY pk, ob DESC)` and registers it `WITH ORDER (pk ASC, ob DESC)`. Reusing the existing loader keeps both the `--size` and `--format` axes working with no duplication. `load` is untimed, so the sort stays out of the measurement. - `WITH ORDER` can only name columns, so the partition key (`id3 % N`) and the tie expression are materialized as `pk` and `ob`. - Asserts guard the three ways this could silently measure nothing: both config flags took effect, the sorted copy holds every source row, and `ob` has the cardinality the query name claims. - `expect_plan PartitionedTopKExec` is deliberately the shared substring, so the same file validates the heap operator and any streaming variant that replaces it. No Rust changes; benchmark files only. ## Are these changes tested? `benchmark_runner h2o --subgroup window_sorted` runs all 16 queries green. Row counts are exact: 2 per partition for the distinct-ORDER-BY shapes, and 1,010,812 (RANK) / 2,020,722 (DENSE_RANK) for the tie shapes — the top-1 and top-2 of 10 distinct values over 10M rows. ## Are there any user-facing changes? No. --- benchmarks/bench.sh | 76 ++++++++++++++ .../benchmarks/window_sorted/q01.benchmark | 13 +++ .../benchmarks/window_sorted/q02.benchmark | 11 +++ .../benchmarks/window_sorted/q03.benchmark | 11 +++ .../benchmarks/window_sorted/q04.benchmark | 12 +++ .../benchmarks/window_sorted/q05.benchmark | 15 +++ .../benchmarks/window_sorted/q06.benchmark | 11 +++ .../benchmarks/window_sorted/q07.benchmark | 11 +++ .../benchmarks/window_sorted/q08.benchmark | 11 +++ .../benchmarks/window_sorted/q09.benchmark | 17 ++++ .../benchmarks/window_sorted/q10.benchmark | 11 +++ .../benchmarks/window_sorted/q11.benchmark | 13 +++ .../benchmarks/window_sorted/q12.benchmark | 11 +++ .../benchmarks/window_sorted/q13.benchmark | 11 +++ .../benchmarks/window_sorted/q14.benchmark | 11 +++ .../benchmarks/window_sorted/q15.benchmark | 13 +++ .../benchmarks/window_sorted/q16.benchmark | 11 +++ benchmarks/sql_benchmarks/h2o/h2o.suite | 4 + .../h2o/init/load_window_sorted.sql | 40 ++++++++ .../h2o/init/window_sorted_cleanup.sql | 1 + .../h2o/init/window_sorted_settings.sql | 24 +++++ .../h2o/window_sorted.benchmark.template | 99 +++++++++++++++++++ 22 files changed, 437 insertions(+) create mode 100644 benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q01.benchmark create mode 100644 benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q02.benchmark create mode 100644 benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q03.benchmark create mode 100644 benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q04.benchmark create mode 100644 benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q05.benchmark create mode 100644 benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q06.benchmark create mode 100644 benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q07.benchmark create mode 100644 benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q08.benchmark create mode 100644 benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q09.benchmark create mode 100644 benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q10.benchmark create mode 100644 benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q11.benchmark create mode 100644 benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q12.benchmark create mode 100644 benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q13.benchmark create mode 100644 benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q14.benchmark create mode 100644 benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q15.benchmark create mode 100644 benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q16.benchmark create mode 100644 benchmarks/sql_benchmarks/h2o/init/load_window_sorted.sql create mode 100644 benchmarks/sql_benchmarks/h2o/init/window_sorted_cleanup.sql create mode 100644 benchmarks/sql_benchmarks/h2o/init/window_sorted_settings.sql create mode 100644 benchmarks/sql_benchmarks/h2o/window_sorted.benchmark.template diff --git a/benchmarks/bench.sh b/benchmarks/bench.sh index df9b7f6c94f16..419fd5be3ad2b 100755 --- a/benchmarks/bench.sh +++ b/benchmarks/bench.sh @@ -139,6 +139,12 @@ h2o_big_join: h2oai benchmark with large dataset (1e9 rows) fo h2o_small_window: Extended h2oai benchmark with small dataset (1e7 rows) for window, default file format is csv h2o_medium_window: Extended h2oai benchmark with medium dataset (1e8 rows) for window, default file format is csv h2o_big_window: Extended h2oai benchmark with large dataset (1e9 rows) for window, default file format is csv +h2o_small_window_sorted: Window Top-N over a declared-sorted h2o input, small dataset (1e7 rows), default file format is csv +h2o_medium_window_sorted: Window Top-N over a declared-sorted h2o input, medium dataset (1e8 rows), default file format is csv +h2o_big_window_sorted: Window Top-N over a declared-sorted h2o input, large dataset (1e9 rows), default file format is csv +h2o_small_window_sorted_parquet: Window Top-N over a declared-sorted h2o input, small dataset (1e7 rows), source file format is parquet +h2o_medium_window_sorted_parquet: Window Top-N over a declared-sorted h2o input, medium dataset (1e8 rows), source file format is parquet +h2o_big_window_sorted_parquet: Window Top-N over a declared-sorted h2o input, large dataset (1e9 rows), source file format is parquet h2o_small_parquet: h2oai benchmark with small dataset (1e7 rows) for groupby, file format is parquet h2o_medium_parquet: h2oai benchmark with medium dataset (1e8 rows) for groupby, file format is parquet h2o_big_parquet: h2oai benchmark with large dataset (1e9 rows) for groupby, file format is parquet @@ -309,6 +315,26 @@ main() { h2o_big_window) data_h2o_join "BIG" "CSV" ;; + # the sorted window subgroup derives its data from the same + # source, then sorts it inside its load SQL + h2o_small_window_sorted) + data_h2o_join "SMALL" "CSV" + ;; + h2o_medium_window_sorted) + data_h2o_join "MEDIUM" "CSV" + ;; + h2o_big_window_sorted) + data_h2o_join "BIG" "CSV" + ;; + h2o_small_window_sorted_parquet) + data_h2o_join "SMALL" "PARQUET" + ;; + h2o_medium_window_sorted_parquet) + data_h2o_join "MEDIUM" "PARQUET" + ;; + h2o_big_window_sorted_parquet) + data_h2o_join "BIG" "PARQUET" + ;; h2o_small_parquet) data_h2o "SMALL" "PARQUET" ;; @@ -534,6 +560,24 @@ main() { h2o_big_window) run_h2o_window "BIG" "CSV" "window" ;; + h2o_small_window_sorted) + run_h2o_window_sorted "small" "csv" + ;; + h2o_medium_window_sorted) + run_h2o_window_sorted "medium" "csv" + ;; + h2o_big_window_sorted) + run_h2o_window_sorted "big" "csv" + ;; + h2o_small_window_sorted_parquet) + run_h2o_window_sorted "small" "parquet" + ;; + h2o_medium_window_sorted_parquet) + run_h2o_window_sorted "medium" "parquet" + ;; + h2o_big_window_sorted_parquet) + run_h2o_window_sorted "big" "parquet" + ;; h2o_small_parquet) run_h2o "SMALL" "PARQUET" ;; @@ -1266,6 +1310,38 @@ run_h2o_window() { h2o_runner "$1" "$2" "window" } +# Runs the h2o window_sorted subgroup: window Top-N over an input that declares +# the ordering the window requires. +# +# The `window` subgroup registers `x` with no declared ordering, so +# `output_ordering()` is None and any plan that depends on a declared ordering is +# unreachable from it, however the data happens to sit on disk. This subgroup's +# load SQL writes a sorted copy and registers it `WITH ORDER`, then asserts a +# per-partition top-K operator is actually in the plan so a silent fallback to +# window-plus-filter cannot masquerade as a result. +# +# The sort happens in the untimed `load` step, so it stays out of the +# measurement. Data comes from the same source as the window subgroup, so the +# data step is data_h2o_join. +# +# H2O_FILE_TYPE selects the *source* format and so only affects that untimed +# load: the measured query always reads the sorted Parquet copy, whatever the +# source was. A Parquet source still makes the load markedly cheaper than +# re-parsing a multi-GB CSV, which is why the `_parquet` entries exist; they are +# expected to produce the same measured numbers as their CSV counterparts, not +# different ones. Each entry pairs with the data step that generates its format. +run_h2o_window_sorted() { + SIZE=${1:-"small"} + FILE_TYPE=${2:-"csv"} + echo "Running h2o window_sorted benchmark (size=${SIZE}, source format=${FILE_TYPE})..." + debug_run env BENCH_NAME=h2o \ + BENCH_SUBGROUP=window_sorted \ + H2O_BENCH_SIZE="${SIZE}" \ + H2O_FILE_TYPE="${FILE_TYPE}" \ + ${QUERY:+BENCH_QUERY="${QUERY}"} \ + bash -c "$SQL_CARGO_COMMAND" +} + # Runs the external aggregation benchmark run_external_aggr() { # Use TPC-H SF1 dataset diff --git a/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q01.benchmark b/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q01.benchmark new file mode 100644 index 0000000000000..236220b7595b4 --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q01.benchmark @@ -0,0 +1,13 @@ +# ROW_NUMBER top-2 per partition over the declared-sorted input, swept across +# partition cardinality. This mirrors the `window` subgroup q14-q17 sweep, with +# the partition key materialized as a column so the ordering can be declared. + +subgroup window_sorted + +template sql_benchmarks/h2o/window_sorted.benchmark.template +QPAD=01 +FN=ROW_NUMBER +PK_MOD=100 +OB_EXPR=v2 +OB_BOUND=count(DISTINCT ob) > 1000 +NAME=window_sorted_q01_row_number_100_partitions \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q02.benchmark b/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q02.benchmark new file mode 100644 index 0000000000000..a9c5733d62e43 --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q02.benchmark @@ -0,0 +1,11 @@ +# ROW_NUMBER top-2, ~1K partitions. + +subgroup window_sorted + +template sql_benchmarks/h2o/window_sorted.benchmark.template +QPAD=02 +FN=ROW_NUMBER +PK_MOD=1000 +OB_EXPR=v2 +OB_BOUND=count(DISTINCT ob) > 1000 +NAME=window_sorted_q02_row_number_1k_partitions \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q03.benchmark b/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q03.benchmark new file mode 100644 index 0000000000000..4a000af56bbf6 --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q03.benchmark @@ -0,0 +1,11 @@ +# ROW_NUMBER top-2, ~10K partitions. + +subgroup window_sorted + +template sql_benchmarks/h2o/window_sorted.benchmark.template +QPAD=03 +FN=ROW_NUMBER +PK_MOD=10000 +OB_EXPR=v2 +OB_BOUND=count(DISTINCT ob) > 1000 +NAME=window_sorted_q03_row_number_10k_partitions \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q04.benchmark b/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q04.benchmark new file mode 100644 index 0000000000000..77208a20249ae --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q04.benchmark @@ -0,0 +1,12 @@ +# ROW_NUMBER top-2, ~100K partitions. Fewest rows per partition on the sweep, +# so per-partition bookkeeping is amortized over the least work. + +subgroup window_sorted + +template sql_benchmarks/h2o/window_sorted.benchmark.template +QPAD=04 +FN=ROW_NUMBER +PK_MOD=100000 +OB_EXPR=v2 +OB_BOUND=count(DISTINCT ob) > 1000 +NAME=window_sorted_q04_row_number_100k_partitions \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q05.benchmark b/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q05.benchmark new file mode 100644 index 0000000000000..e06acffd9a440 --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q05.benchmark @@ -0,0 +1,15 @@ +# RANK top-2 per partition, mirroring the ROW_NUMBER sweep above. RANK retains +# boundary ties, so `WHERE rnk <= 2` can keep more than 2 rows per partition and +# the operator has to track the tied set rather than a fixed row count. `v2` is +# near-unique, so these four measure RANK with the tie machinery idle; q09-q10 +# drive it. + +subgroup window_sorted + +template sql_benchmarks/h2o/window_sorted.benchmark.template +QPAD=05 +FN=RANK +PK_MOD=100 +OB_EXPR=v2 +OB_BOUND=count(DISTINCT ob) > 1000 +NAME=window_sorted_q05_rank_100_partitions \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q06.benchmark b/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q06.benchmark new file mode 100644 index 0000000000000..d9421902b412a --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q06.benchmark @@ -0,0 +1,11 @@ +# RANK top-2, ~1K partitions. + +subgroup window_sorted + +template sql_benchmarks/h2o/window_sorted.benchmark.template +QPAD=06 +FN=RANK +PK_MOD=1000 +OB_EXPR=v2 +OB_BOUND=count(DISTINCT ob) > 1000 +NAME=window_sorted_q06_rank_1k_partitions \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q07.benchmark b/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q07.benchmark new file mode 100644 index 0000000000000..7a2a49ebce84e --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q07.benchmark @@ -0,0 +1,11 @@ +# RANK top-2, ~10K partitions. + +subgroup window_sorted + +template sql_benchmarks/h2o/window_sorted.benchmark.template +QPAD=07 +FN=RANK +PK_MOD=10000 +OB_EXPR=v2 +OB_BOUND=count(DISTINCT ob) > 1000 +NAME=window_sorted_q07_rank_10k_partitions \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q08.benchmark b/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q08.benchmark new file mode 100644 index 0000000000000..409ed17bba712 --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q08.benchmark @@ -0,0 +1,11 @@ +# RANK top-2, ~100K partitions. + +subgroup window_sorted + +template sql_benchmarks/h2o/window_sorted.benchmark.template +QPAD=08 +FN=RANK +PK_MOD=100000 +OB_EXPR=v2 +OB_BOUND=count(DISTINCT ob) > 1000 +NAME=window_sorted_q08_rank_100k_partitions \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q09.benchmark b/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q09.benchmark new file mode 100644 index 0000000000000..9728f470819b0 --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q09.benchmark @@ -0,0 +1,17 @@ +# RANK top-2 with heavy ties: 10 distinct ORDER BY values, so a whole partition +# slice ties at rank 1 and every tied row survives `rnk <= 2`. The tied set, not +# the top-K bound, dominates the work. +# +# The cast is load-bearing. `v2` is Float64, so `v2 % 10` is float modulo and +# stays near-unique — it produces no ties whatsoever. Truncating to an integer +# first is what actually collapses the domain to 10 values. + +subgroup window_sorted + +template sql_benchmarks/h2o/window_sorted.benchmark.template +QPAD=09 +FN=RANK +PK_MOD=1000 +OB_EXPR=CAST(v2 AS BIGINT) % 10 +OB_BOUND=count(DISTINCT ob) = 10 +NAME=window_sorted_q09_rank_1k_partitions_heavy_ties \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q10.benchmark b/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q10.benchmark new file mode 100644 index 0000000000000..232e3f1b3b17d --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q10.benchmark @@ -0,0 +1,11 @@ +# RANK top-2, ~10K partitions, heavy ties. + +subgroup window_sorted + +template sql_benchmarks/h2o/window_sorted.benchmark.template +QPAD=10 +FN=RANK +PK_MOD=10000 +OB_EXPR=CAST(v2 AS BIGINT) % 10 +OB_BOUND=count(DISTINCT ob) = 10 +NAME=window_sorted_q10_rank_10k_partitions_heavy_ties \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q11.benchmark b/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q11.benchmark new file mode 100644 index 0000000000000..d5d5d6b256927 --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q11.benchmark @@ -0,0 +1,13 @@ +# DENSE_RANK top-2 per partition, mirroring the RANK sweep. DENSE_RANK keeps +# every row whose ORDER BY value is among the 2 distinct-greatest in the +# partition, so rows kept per partition is unbounded in rows-per-distinct-value. + +subgroup window_sorted + +template sql_benchmarks/h2o/window_sorted.benchmark.template +QPAD=11 +FN=DENSE_RANK +PK_MOD=100 +OB_EXPR=v2 +OB_BOUND=count(DISTINCT ob) > 1000 +NAME=window_sorted_q11_dense_rank_100_partitions \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q12.benchmark b/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q12.benchmark new file mode 100644 index 0000000000000..c8d277ff040d6 --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q12.benchmark @@ -0,0 +1,11 @@ +# DENSE_RANK top-2, ~1K partitions. + +subgroup window_sorted + +template sql_benchmarks/h2o/window_sorted.benchmark.template +QPAD=12 +FN=DENSE_RANK +PK_MOD=1000 +OB_EXPR=v2 +OB_BOUND=count(DISTINCT ob) > 1000 +NAME=window_sorted_q12_dense_rank_1k_partitions \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q13.benchmark b/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q13.benchmark new file mode 100644 index 0000000000000..be7246232cc56 --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q13.benchmark @@ -0,0 +1,11 @@ +# DENSE_RANK top-2, ~10K partitions. + +subgroup window_sorted + +template sql_benchmarks/h2o/window_sorted.benchmark.template +QPAD=13 +FN=DENSE_RANK +PK_MOD=10000 +OB_EXPR=v2 +OB_BOUND=count(DISTINCT ob) > 1000 +NAME=window_sorted_q13_dense_rank_10k_partitions \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q14.benchmark b/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q14.benchmark new file mode 100644 index 0000000000000..9c773a93ea206 --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q14.benchmark @@ -0,0 +1,11 @@ +# DENSE_RANK top-2, ~100K partitions. + +subgroup window_sorted + +template sql_benchmarks/h2o/window_sorted.benchmark.template +QPAD=14 +FN=DENSE_RANK +PK_MOD=100000 +OB_EXPR=v2 +OB_BOUND=count(DISTINCT ob) > 1000 +NAME=window_sorted_q14_dense_rank_100k_partitions \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q15.benchmark b/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q15.benchmark new file mode 100644 index 0000000000000..3a655c0b4d397 --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q15.benchmark @@ -0,0 +1,13 @@ +# DENSE_RANK top-2, ~1K partitions, heavy ties. With 10 distinct ORDER BY values +# the top-2 distinct values cover roughly a fifth of each partition, so this is +# the heaviest surviving-row count on the sweep. See q09 on why the cast matters. + +subgroup window_sorted + +template sql_benchmarks/h2o/window_sorted.benchmark.template +QPAD=15 +FN=DENSE_RANK +PK_MOD=1000 +OB_EXPR=CAST(v2 AS BIGINT) % 10 +OB_BOUND=count(DISTINCT ob) = 10 +NAME=window_sorted_q15_dense_rank_1k_partitions_heavy_ties \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q16.benchmark b/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q16.benchmark new file mode 100644 index 0000000000000..65baaed851b26 --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q16.benchmark @@ -0,0 +1,11 @@ +# DENSE_RANK top-2, ~10K partitions, heavy ties. + +subgroup window_sorted + +template sql_benchmarks/h2o/window_sorted.benchmark.template +QPAD=16 +FN=DENSE_RANK +PK_MOD=10000 +OB_EXPR=CAST(v2 AS BIGINT) % 10 +OB_BOUND=count(DISTINCT ob) = 10 +NAME=window_sorted_q16_dense_rank_10k_partitions_heavy_ties \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/h2o/h2o.suite b/benchmarks/sql_benchmarks/h2o/h2o.suite index 27d83285ba026..cb2314b13c6e4 100644 --- a/benchmarks/sql_benchmarks/h2o/h2o.suite +++ b/benchmarks/sql_benchmarks/h2o/h2o.suite @@ -31,3 +31,7 @@ description = "Run H2O window query 3." [[examples]] command = "cargo run --release --bin benchmark_runner -- h2o --subgroup join --size medium -f parquet" description = "Run the H2O join queries with the medium Parquet dataset." + +[[examples]] +command = "cargo run --release --bin benchmark_runner -- h2o --subgroup window_sorted" +description = "Run the window Top-N sweep over an input that declares its ordering." \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/h2o/init/load_window_sorted.sql b/benchmarks/sql_benchmarks/h2o/init/load_window_sorted.sql new file mode 100644 index 0000000000000..6e8f2a67b79d1 --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/init/load_window_sorted.sql @@ -0,0 +1,40 @@ +-- Materializes a pre-sorted copy of the h2o window table that *declares* its +-- ordering, for the `window_sorted` subgroup. +-- +-- Why this exists: the `window` subgroup partitions by an expression +-- (`id3 % N`) and, in its heavy-ties queries, orders by another (`v2 % 10`). +-- `WITH ORDER` can only name columns, so both are materialized here as `pk` +-- and `ob`. The queries then partition and order by plain columns, and the +-- declared ordering can satisfy the window's requirement. +-- +-- Parameters, supplied by window_sorted.benchmark.template so each benchmark +-- writes only the shape it measures: +-- PK_MOD partition count (100, 1000, 10000, 100000) +-- OB_EXPR ORDER BY value expression (`v2` distinct, integer-cast for ties) +-- +-- Note `v2` is Float64 here, so a plain `v2 % 10` is float modulo and stays +-- near-unique. Producing genuine ties needs an integer cast first; the template +-- asserts ob's resulting cardinality so a query cannot get this wrong silently. +-- +-- The ORDER BY is load-bearing: the file has to be written in (pk, ob DESC) +-- order for the WITH ORDER declaration below to be true. A declaration that +-- does not match the file is not a planning error, it silently produces wrong +-- results, so this ORDER BY and that WITH ORDER must be edited together. +-- +-- `load` runs before `init` and is not timed, so this sort never enters the +-- measurement. +DROP TABLE IF EXISTS x_sorted; + +COPY ( + SELECT id3 % ${PK_MOD:-1000} AS pk, ${OB_EXPR:-v2} AS ob + FROM x + WHERE v2 IS NOT NULL + ORDER BY pk ASC, ob DESC +) +TO 'sql_benchmarks/h2o/scratch/window_sorted.parquet' +STORED AS PARQUET; + +CREATE EXTERNAL TABLE x_sorted +STORED AS PARQUET +WITH ORDER (pk ASC, ob DESC) +LOCATION 'sql_benchmarks/h2o/scratch/window_sorted.parquet'; \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/h2o/init/window_sorted_cleanup.sql b/benchmarks/sql_benchmarks/h2o/init/window_sorted_cleanup.sql new file mode 100644 index 0000000000000..9601519d48f7c --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/init/window_sorted_cleanup.sql @@ -0,0 +1 @@ +DROP TABLE IF EXISTS x_sorted; \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/h2o/init/window_sorted_settings.sql b/benchmarks/sql_benchmarks/h2o/init/window_sorted_settings.sql new file mode 100644 index 0000000000000..7ed87226149c9 --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/init/window_sorted_settings.sql @@ -0,0 +1,24 @@ +-- Session settings for the h2o window_sorted subgroup. `init` runs after +-- `load`, so these apply to the measured query and not to the COPY that writes +-- the sorted file. +-- +-- enable_window_topn turns on the WindowTopN rewrite, making the baseline the +-- existing per-partition top-K operator rather than an unoptimized +-- window-plus-filter plan. The flag defaults to false, so without this the +-- subgroup would measure a different plan shape than its name claims. +-- +-- information_schema is enabled so the template can assert the flags actually +-- took effect. +-- +-- prefer_existing_sort decides *how* a declared input ordering is honored once +-- some operator requires it. The heap top-K operator requires no input ordering, +-- so on its own the ordering is simply dropped: the hash repartition that +-- co-locates each partition key is inserted with preserve_order: false, and that +-- is not a defect. For an operator that does require the ordering, the planner +-- then has two ways to satisfy it — an order-preserving repartition, or a +-- SortExec that re-sorts data already sorted on disk. This setting picks the +-- former. Without it such an operator would measure a full sort and the sorted +-- input would be pointless. +set datafusion.catalog.information_schema = true; +set datafusion.optimizer.enable_window_topn = true; +set datafusion.optimizer.prefer_existing_sort = true; \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/h2o/window_sorted.benchmark.template b/benchmarks/sql_benchmarks/h2o/window_sorted.benchmark.template new file mode 100644 index 0000000000000..9170511492ea5 --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/window_sorted.benchmark.template @@ -0,0 +1,99 @@ +# Shared template for the h2o `window_sorted` subgroup: WindowTopN over an +# input that *declares* the ordering the window requires. +# +# Parameters set by each qNN.benchmark: +# QPAD zero-padded query id +# FN window function (ROW_NUMBER, RANK, DENSE_RANK) +# PK_MOD partition count (100, 1000, 10000, 100000) +# OB_EXPR ORDER BY value expression (`v2` distinct, integer-cast for ties) +# OB_BOUND SQL predicate pinning ob's cardinality, asserted below +# NAME display name +# +# Why this subgroup exists: the `window` subgroup registers `x` with no declared +# ordering, so `output_ordering()` is None and the planner cannot know the input +# is sorted. Any operator that requires a declared ordering is therefore +# unreachable from those queries, however the data happens to be laid out on +# disk. This subgroup republishes the same sweep over a `WITH ORDER` table so +# the ordering-dependent plans are reachable. +# +# The first `load` is the existing unsorted loader, which creates `x`; the +# second reads `x` and writes the sorted `x_sorted`. Reusing the existing loader +# keeps both the `--size` and `--format` axes working with no duplication. + +load sql_benchmarks/h2o/init/load_window_${H2O_BENCH_SIZE:-small}_${H2O_FILE_TYPE:-csv}.sql + +load sql_benchmarks/h2o/init/load_window_sorted.sql + +init sql_benchmarks/h2o/init/window_sorted_settings.sql + +name ${NAME} +group h2o + +echo Loading ${H2O_BENCH_SIZE:-small} window ${H2O_FILE_TYPE:-csv} h2o data, sorted by (pk, ob DESC) over ${PK_MOD} partitions + +# Guard against a silent config no-op: with the rewrite off the plan is a plain +# window plus filter, and the subgroup would measure something other than what +# its name claims. +assert I +SELECT value = 'true' FROM information_schema.df_settings WHERE name = 'datafusion.optimizer.enable_window_topn'; +---- +true + +# Same guard for the setting that decides an ordering-requiring operator gets an +# order-preserving repartition instead of a redundant SortExec. +assert I +SELECT value = 'true' FROM information_schema.df_settings WHERE name = 'datafusion.optimizer.prefer_existing_sort'; +---- +true + +# The sorted copy must hold every non-null-v2 row of the source. This catches a +# COPY that silently wrote a partial file, which would otherwise show up only as +# a suspiciously fast run. +assert I +SELECT count(*) = (SELECT count(*) FROM x WHERE v2 IS NOT NULL) FROM x_sorted; +---- +true + +# The ORDER BY column really has the cardinality this query's name claims. A +# tie-shape query whose `ob` turns out near-unique measures the no-ties case +# under a heavy-ties name, and nothing else in the run would show it: the timings +# stay plausible and only the output row count quietly collapses to K per +# partition. `v2` is Float64, so this is a live trap — `v2 % 10` is float modulo +# and yields no ties at all. +assert I +SELECT ${OB_BOUND} FROM x_sorted; +---- +true + +# The per-partition top-K operator has to be in the plan, otherwise the +# declared ordering was not picked up and this is just a window-plus-filter +# measurement wearing the subgroup's name. +# +# `PartitionedTopKExec` is deliberately the string asserted rather than a +# fully-qualified operator name: it is also a substring of the streaming +# variant's name, so this same check holds on a branch where a streaming +# operator replaces the heap one for declared-sorted input. That is the point of +# the subgroup — the comparison is this file's numbers across two branches, so +# the assertion must not be the thing that differs between them. +expect_plan PartitionedTopKExec + +# `PartitionedTopKExec` alone only proves the plan has an operator that +# distributes by the partition key; that operator requires no input ordering, +# so it appears whether or not the scan below it declares one. Assert the +# scan's own declared ordering too, so a `WITH ORDER` that silently stops +# reaching `DataSourceExec` (e.g. a rewrite that drops it, or a future change +# to `load_window_sorted.sql`) fails loudly here instead of just being absent +# from a plan that still "passes" the check above. +expect_plan output_ordering=[pk@0 ASC NULLS LAST, ob@1 DESC] + +run +-- Top-2 per partition by ${FN} over the declared-sorted input. +SELECT pk, ob FROM ( + SELECT pk, ob, + ${FN}() OVER (PARTITION BY pk ORDER BY ob DESC) AS rnk + FROM x_sorted +) sub_query WHERE rnk <= 2; + +result sql_benchmarks/h2o/results/window_sorted/${H2O_BENCH_SIZE:-small}/q${QPAD}.csv + +cleanup sql_benchmarks/h2o/init/window_sorted_cleanup.sql \ No newline at end of file From 27e81e9a45317afeacd7d7dc2b77f10637ecd5d5 Mon Sep 17 00:00:00 2001 From: Jaideep Pyne Date: Thu, 3 Sep 2026 03:50:58 +0000 Subject: [PATCH 23/37] fix: do not duplicate volatile expressions when extracting leaf expressions (#24720) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Which issue does this PR close? - Closes #24678. ## Rationale for this change A column alias denotes one value per row, and `WHERE p` may only return rows for which `p` held for that row. Today the leaf-expression extraction passes break both: ```sql SELECT s, s['a'] AS field FROM (SELECT named_struct('a', random()) AS s FROM generate_series(1, 3)); ``` | `s['a']` | `field` | |---|---| | 0.5159112071865757 | 0.4514238291986653 | | 0.0029104680074608646 | 0.28979983332288195 | | 0.48542729227457915 | 0.04499392663566881 | `field` is defined as `s['a']` but differs from it on every row, because the plan is ``` Projection: named_struct(Utf8("a"), random()) AS s, random() AS field ``` The `Filter` form is worse — it returns rows that fail their own predicate: ```sql SELECT bool_and(s['a'] > 0.5) FROM (SELECT s FROM (SELECT named_struct('a', random()) AS s FROM generate_series(1, 1000)) WHERE s['a'] > 0.5); -- false; the predicate tested a different draw than the one in the returned `s` ``` Setting `datafusion.optimizer.enable_leaf_expression_pushdown = false` returns the correct answer in both cases, so the rewrite alone changes the meaning of the query. **Root cause.** `build_extraction_projection_impl` merges an extraction into the input projection by resolving column references through `build_projection_replace_map`, i.e. by inlining each referenced column's *defining* expression. Inlining a volatile definition produces a second, independent evaluation. There was no volatility check in the file. This is the same invariant `FileScanConfig::try_swapping_with_projection` already enforces for the physical projection-pushdown path via `would_duplicate_costly_exprs` (#23220) — the logical extraction path was missing it. ## What changes are included in this PR? - `volatile_output_columns()` — a projection's output columns whose definition is volatile. - `would_duplicate_volatile()` — true when an extraction references one of them. - The guard is applied at the three places that can merge into an input projection: `extract_from_plan` (pass 1: Filter/Sort/Limit/Aggregate/Join), `split_and_push_projection` (pass 2), and `try_push_into_inputs` (multi-input/Union routing). Each already had a "leave the plan alone" return path. The guard is targeted rather than blanket: for `ORDER BY s['a']` the extraction still happens, stacked above the volatile projection instead of merged into it, so the optimization is kept and the result is correct. ### Relationship to #23691 @fornwall wondered on the issue whether #23691 already covers this. I checked out that branch and ran both shapes against it: `SELECT s, s['a']` is incidentally fixed there, but `WHERE s['a'] > 0.5` still duplicates `random()` and still returns rows failing the predicate. #23691 guards `KeepInPlace` *compute cost* in `split_and_push_projection` only; volatility is a separate concern (one duplication is already wrong, regardless of cost or placement) and pass 1 is a different code path. The two changes look independent to me and I believe they compose, but I'd appreciate a second opinion on that from whoever reviews #23691. ## Are these changes tested? Yes. - `datafusion/sqllogictest/test_files/projection_pushdown.slt` — a new section beside the existing #23220 volatile section: two `EXPLAIN`s pinning `random()` to a single occurrence, and two deterministic `bool_and(...)` correctness queries. All four fail on `main` and pass here. - Two rule-level snapshot tests in `extract_leaf_expressions.rs` covering the pass-2 projection merge and the pass-1 `Filter` extraction, using a new test-only `PlacementTestUDF::with_volatility()`. `cargo test -p datafusion-optimizer` (796), `cargo test -p datafusion --lib --tests` (2049) and the full 504-file sqllogictest suite all pass; `cargo fmt --all` and `cargo clippy --all-targets --all-features -- -D warnings` are clean. ## Are there any user-facing changes? Queries that were silently returning wrong results now return correct ones. No API change. In the affected shapes the extraction is skipped, which can cost a small amount of column pruning — only when the referenced column is defined by a volatile expression. --- Per the ASF generative-tooling policy and DataFusion's AI-assisted contribution guidance: this patch was prepared with AI assistance. The core idea is the one described above — the merge path inlines a referenced column's defining expression, which duplicates a volatile definition — and the open question about how this composes with #23691 is flagged deliberately rather than glossed over. --- .../optimizer/src/extract_leaf_expressions.rs | 156 +++++++++++++++++- datafusion/optimizer/src/test/udfs.rs | 11 +- .../test_files/projection_pushdown.slt | 77 +++++++++ 3 files changed, 240 insertions(+), 4 deletions(-) diff --git a/datafusion/optimizer/src/extract_leaf_expressions.rs b/datafusion/optimizer/src/extract_leaf_expressions.rs index 107758cf6b6e7..dfaf93bc1e8fa 100644 --- a/datafusion/optimizer/src/extract_leaf_expressions.rs +++ b/datafusion/optimizer/src/extract_leaf_expressions.rs @@ -57,6 +57,58 @@ fn has_extractable_expr(exprs: &[Expr]) -> bool { }) } +/// Returns the flat names of `plan`'s output columns whose defining expression +/// is volatile (e.g. `random()`). +/// +/// Only a [`LogicalPlan::Projection`] can define such a column: anywhere else +/// the value has already been materialized by the projection that produced it, +/// so referencing the column again does not re-evaluate anything. +fn volatile_output_columns(plan: &LogicalPlan) -> BTreeSet { + let LogicalPlan::Projection(projection) = plan else { + return BTreeSet::new(); + }; + projection + .schema + .iter() + .zip(projection.expr.iter()) + .filter(|(_, expr)| expr.is_volatile()) + .map(|((qualifier, field), _)| Column::from((qualifier, field)).flat_name()) + .collect() +} + +/// Returns `true` if building an extraction projection for `exprs` on top of +/// `input` would duplicate a volatile computation. +/// +/// When `input` is already a projection, [`build_extraction_projection_impl`] +/// *merges* into it: every column reference in an extracted expression is +/// replaced by that column's defining expression (see +/// [`build_projection_replace_map`]). Inlining a volatile definition makes the +/// merged projection evaluate it a second, independent time, so the extracted +/// value no longer matches the column it was derived from: +/// +/// ```text +/// Projection: s, get_field(s, 'a') AS field +/// Projection: named_struct('a', random()) AS s +/// ``` +/// +/// would merge into a single projection computing `random()` twice, and +/// `field` would then differ from `s['a']` on every row. Callers skip the +/// extraction instead. +fn would_duplicate_volatile<'a>( + exprs: impl IntoIterator, + input: &LogicalPlan, +) -> bool { + let volatile = volatile_output_columns(input); + if volatile.is_empty() { + return false; + } + exprs.into_iter().any(|expr| { + expr.column_refs() + .iter() + .any(|col| volatile.contains(&col.flat_name())) + }) +} + /// Extracts `MoveTowardsLeafNodes` sub-expressions from non-projection nodes /// into **extraction projections** (pass 1 of 2). /// @@ -195,7 +247,18 @@ fn extract_from_plan( } // Fast pre-check: skip all allocations if no extractable expressions exist - if !has_extractable_expr(&plan.expressions()) { + let node_exprs = plan.expressions(); + if !has_extractable_expr(&node_exprs) { + return Ok(Transformed::no(plan)); + } + + // The extraction projection is merged into an input that is already a + // projection, which inlines the referenced columns' definitions. Skip the + // extraction when that would duplicate a volatile computation. + if inputs + .iter() + .any(|input| would_duplicate_volatile(node_exprs.iter(), input)) + { return Ok(Transformed::no(plan)); } @@ -917,6 +980,16 @@ fn split_and_push_projection( return Ok(None); } + // Pushing into an input that is already a projection merges into it and + // inlines the referenced columns' definitions. Leave the projection alone + // when that would duplicate a volatile computation. + if would_duplicate_volatile( + extraction_pairs.iter().map(|(expr, _)| expr), + input.as_ref(), + ) { + return Ok(None); + } + // ── Phase 2: Push down ────────────────────────────────────────────── let proj_input = Arc::clone(&proj.input); let pushed = push_extraction_pairs( @@ -1215,6 +1288,15 @@ fn try_push_into_inputs( if per_input[idx].pairs.is_empty() { new_inputs.push(input.clone()); } else { + // Merging into an input projection inlines the referenced columns' + // definitions; bail out when that would duplicate a volatile + // computation. + if would_duplicate_volatile( + per_input[idx].pairs.iter().map(|(expr, _)| expr), + input, + ) { + return Ok(None); + } let input_arc = Arc::new(input.clone()); let target_schema = Arc::clone(input.schema()); let proj = build_extraction_projection_impl( @@ -1270,7 +1352,7 @@ mod tests { use crate::{Optimizer, OptimizerContext}; use datafusion_expr::expr::ScalarFunction; use datafusion_expr::{ - ScalarUDF, col, lit, logical_plan::builder::LogicalPlanBuilder, + ScalarUDF, Volatility, col, lit, logical_plan::builder::LogicalPlanBuilder, }; fn leaf_udf(expr: Expr, name: &str) -> Expr { @@ -1283,6 +1365,19 @@ mod tests { )) } + /// A stand-in for `random()`: a volatile expression that must be evaluated + /// exactly once per row. + fn volatile_udf(expr: Expr) -> Expr { + Expr::ScalarFunction(ScalarFunction::new_udf( + Arc::new(ScalarUDF::new_from_impl( + PlacementTestUDF::new() + .with_placement(ExpressionPlacement::KeepInPlace) + .with_volatility(Volatility::Volatile), + )), + vec![expr], + )) + } + // ========================================================================= // Combined optimization stage formatter // ========================================================================= @@ -1960,6 +2055,63 @@ mod tests { "#) } + /// Merging an extraction into an input projection inlines the definition of + /// every column the extraction references. When that definition is volatile + /// the inlined copy is an independent evaluation, so the extraction must be + /// skipped and the plan left alone. + #[test] + fn test_no_merge_into_volatile_projection() -> Result<()> { + let table_scan = test_table_scan()?; + let plan = LogicalPlanBuilder::from(table_scan) + .project(vec![volatile_udf(col("a")).alias("v")])? + .project(vec![col("v"), leaf_udf(col("v"), "x")])? + .build()?; + + assert_stages!(plan, @r#" + ## Original Plan + Projection: v, leaf_udf(v, Utf8("x")) + Projection: keep_in_place_udf(test.a) AS v + TableScan: test projection=[a] + + ## After Extraction + (same as original) + + ## After Pushdown + (same as after extraction) + + ## Optimized + (same as after pushdown) + "#) + } + + /// Same guard for pass 1: extracting out of a `Filter` whose input + /// projection defines the referenced column volatilely would make the + /// predicate test a second, independent evaluation. + #[test] + fn test_no_extraction_from_filter_over_volatile_projection() -> Result<()> { + let table_scan = test_table_scan()?; + let plan = LogicalPlanBuilder::from(table_scan) + .project(vec![volatile_udf(col("a")).alias("v")])? + .filter(leaf_udf(col("v"), "x").eq(lit(1u32)))? + .build()?; + + assert_stages!(plan, @r#" + ## Original Plan + Filter: leaf_udf(v, Utf8("x")) = UInt32(1) + Projection: keep_in_place_udf(test.a) AS v + TableScan: test projection=[a] + + ## After Extraction + (same as original) + + ## After Pushdown + (same as after extraction) + + ## Optimized + (same as after pushdown) + "#) + } + /// Projections with aliased columns (nothing to extract) return unchanged. #[test] fn test_projection_early_return_no_extraction() -> Result<()> { diff --git a/datafusion/optimizer/src/test/udfs.rs b/datafusion/optimizer/src/test/udfs.rs index ba71b6a04a7a2..72b1a63f6720d 100644 --- a/datafusion/optimizer/src/test/udfs.rs +++ b/datafusion/optimizer/src/test/udfs.rs @@ -19,7 +19,7 @@ use arrow::datatypes::DataType; use datafusion_common::Result; use datafusion_expr::{ ColumnarValue, Expr, ExpressionPlacement, ScalarFunctionArgs, ScalarUDF, - ScalarUDFImpl, Signature, TypeSignature, + ScalarUDFImpl, Signature, TypeSignature, Volatility, }; /// A configurable test UDF for optimizer tests. @@ -44,7 +44,7 @@ impl PlacementTestUDF { // The actual types don't matter since this UDF is not intended for execution. signature: Signature::new( TypeSignature::OneOf(vec![TypeSignature::Any(1), TypeSignature::Any(2)]), - datafusion_expr::Volatility::Immutable, + Volatility::Immutable, ), placement: ExpressionPlacement::MoveTowardsLeafNodes, id: 0, @@ -64,6 +64,13 @@ impl PlacementTestUDF { self.id = id; self } + + /// Set the volatility of the UDF, so that rules which must not duplicate a + /// volatile computation (e.g. `random()`) can be exercised. + pub fn with_volatility(mut self, volatility: Volatility) -> Self { + self.signature.volatility = volatility; + self + } } impl ScalarUDFImpl for PlacementTestUDF { diff --git a/datafusion/sqllogictest/test_files/projection_pushdown.slt b/datafusion/sqllogictest/test_files/projection_pushdown.slt index 3c7b6f4cb1127..1f9176f7137bd 100644 --- a/datafusion/sqllogictest/test_files/projection_pushdown.slt +++ b/datafusion/sqllogictest/test_files/projection_pushdown.slt @@ -2144,6 +2144,83 @@ true true true +##################### +# Section: leaf expression extraction does not duplicate volatile expressions +# +# Regression test for #24678. `extract_leaf_expressions` / `push_down_leaf_projections` +# build an extraction projection by *merging* into the input projection, which +# inlines the definition of every column the extracted expression references. +# When that definition is volatile the inlined copy is an independent +# evaluation, so `random()` is drawn twice and the extracted value no longer +# matches the struct it was derived from. +# +# Before the fix the plan below read +# Projection: named_struct(Utf8("a"), random()) AS s, random() AS field +# and `field` differed from `s['a']` on every row. This is the same invariant +# `FileScanConfig::try_swapping_with_projection` already enforces for the +# physical pushdown path (#23220). +##################### + +statement ok +set datafusion.explain.logical_plan_only = true; + +# `random()` must appear once; `field` reads the struct materialized below it. +query TT +EXPLAIN SELECT s, s['a'] AS field +FROM (SELECT named_struct('a', random()) AS s FROM generate_series(1, 3)); +---- +logical_plan +01)Projection: s, get_field(s, Utf8("a")) AS field +02)--Projection: named_struct(Utf8("a"), random()) AS s +03)----TableScan: generate_series() projection=[] + +# Same for an extraction out of a Filter: the predicate must test the struct +# that is returned, not a second draw. +query TT +EXPLAIN SELECT s +FROM (SELECT named_struct('a', random()) AS s FROM generate_series(1, 3)) +WHERE s['a'] > 0.5; +---- +logical_plan +01)Filter: get_field(s, Utf8("a")) > Float64(0.5) +02)--Projection: named_struct(Utf8("a"), random()) AS s +03)----TableScan: generate_series() projection=[] + +statement ok +set datafusion.explain.logical_plan_only = false; + +# Every row's `field` is the `a` field of that row's own struct. +query B +SELECT bool_and(field = s['a']) +FROM ( + SELECT s, s['a'] AS field + FROM (SELECT named_struct('a', random()) AS s FROM generate_series(1, 1000)) +); +---- +true + +# Same invariant when the extraction is routed into one side of a join. +query B +SELECT bool_and(field = s['a']) +FROM ( + SELECT l.s, l.s['a'] AS field + FROM (SELECT named_struct('a', random()) AS s FROM generate_series(1, 1000)) AS l + INNER JOIN (SELECT 1 AS k) AS r ON true +); +---- +true + +# Every surviving row satisfies the predicate it was filtered on. +query B +SELECT bool_and(s['a'] > 0.5) +FROM ( + SELECT s + FROM (SELECT named_struct('a', random()) AS s FROM generate_series(1, 1000)) + WHERE s['a'] > 0.5 +); +---- +true + ##################### # Section: expensive expressions are not re-inlined by projection pushdown # From e36c5f842aaedfc1afb8c69572f1f1a323f0cdde Mon Sep 17 00:00:00 2001 From: Sergey Zhukov <62326549+cj-zhukov@users.noreply.github.com> Date: Thu, 3 Sep 2026 04:34:09 +0000 Subject: [PATCH 24/37] feat(dataframe): improve DataFrame::from_columns input types (#24630) (#24633) ## Which issue does this PR close? - Closes #https://github.com/apache/datafusion/issues/24630. ## Rationale for this change Improve the `DataFrame::from_columns` API and give users more flexibility when constructing a `DataFrame` from columns. ## What changes are included in this PR? - Generalize `DataFrame::from_columns` to accept `IntoIterator` of columns. - This allows users to pass both `arrays` and `Vecs` of columns. - Update tests to cover both input forms. ## Are these changes tested? Yes. Tests cover both `array` and `Vec` inputs and verify the resulting schema, data types, row count, and values. ## Are there any user-facing changes? Yes. This is a breaking API change. `DataFrame::from_columns` now accepts an `IntoIterator` of columns instead of specifically accepting a `Vec`. Existing `Vec` usage continues to work, while users can also pass `arrays` and other compatible iterators. Users relying on the exact non-generic function signature may need to update their code to account for the new generic API. --- datafusion/core/src/dataframe/mod.rs | 23 +++-- datafusion/core/tests/dataframe/mod.rs | 91 +++++++++++++++++-- .../library-user-guide/upgrading/56.0.0.md | 23 ++++- 3 files changed, 116 insertions(+), 21 deletions(-) diff --git a/datafusion/core/src/dataframe/mod.rs b/datafusion/core/src/dataframe/mod.rs index 1299e5fb7bd65..ed3dc5ea838b9 100644 --- a/datafusion/core/src/dataframe/mod.rs +++ b/datafusion/core/src/dataframe/mod.rs @@ -2614,7 +2614,7 @@ impl DataFrame { /// # async fn main() -> Result<()> { /// let id: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3])); /// let name: ArrayRef = Arc::new(StringArray::from(vec!["foo", "bar", "baz"])); - /// let df = DataFrame::from_columns(vec![("id", id), ("name", name)])?; + /// let df = DataFrame::from_columns([("id", id), ("name", name)])?; /// let expected = vec![ /// "+----+------+", /// "| id | name |", @@ -2628,17 +2628,16 @@ impl DataFrame { /// # Ok(()) /// # } /// ``` - pub fn from_columns(columns: Vec<(&str, ArrayRef)>) -> Result { - let fields = columns - .iter() - .map(|(name, array)| Field::new(*name, array.data_type().clone(), true)) - .collect::>(); - - let arrays = columns + pub fn from_columns<'a, I>(columns: I) -> Result + where + I: IntoIterator, + { + let (fields, arrays): (Vec<_>, Vec<_>) = columns .into_iter() - .map(|(_, array)| array) - .collect::>(); - + .map(|(name, array)| { + (Field::new(name, array.data_type().clone(), true), array) + }) + .unzip(); let schema = Arc::new(Schema::new(fields)); let batch = RecordBatch::try_new(schema, arrays)?; let ctx = SessionContext::new(); @@ -2695,7 +2694,7 @@ macro_rules! dataframe { use datafusion::prelude::DataFrame; use datafusion::common::test_util::IntoArrayRef; - let columns = vec![ + let columns = [ $( ($name, $data.into_array_ref()), )+ diff --git a/datafusion/core/tests/dataframe/mod.rs b/datafusion/core/tests/dataframe/mod.rs index 1676a69de86d3..527c9567495e9 100644 --- a/datafusion/core/tests/dataframe/mod.rs +++ b/datafusion/core/tests/dataframe/mod.rs @@ -6968,7 +6968,7 @@ async fn test_dataframe_from_columns() -> Result<()> { let strings: ArrayRef = Arc::new(StringArray::from(vec![Some("foo"), Some("bar"), None])); - let df = DataFrame::from_columns(vec![ + let columns = [ ("bool", bools), ("i8", i8s), ("i16", i16s), @@ -6982,10 +6982,10 @@ async fn test_dataframe_from_columns() -> Result<()> { ("f32", f32s), ("f64", f64s), ("str", strings), - ])?; + ]; - assert_eq!(df.schema().fields().len(), 13); - assert_eq!(df.clone().count().await?, 3); + let df1 = DataFrame::from_columns(columns.clone())?; + let df2 = DataFrame::from_columns(columns.to_vec())?; let expected_types = [ ("bool", DataType::Boolean), @@ -7003,14 +7003,89 @@ async fn test_dataframe_from_columns() -> Result<()> { ("str", DataType::Utf8), ]; - let schema = df.schema(); + for df in [df1, df2] { + assert_eq!(df.schema().fields().len(), expected_types.len()); + assert_eq!(df.clone().count().await?, 3); - for (name, data_type) in expected_types { - assert_eq!(schema.field_with_name(None, name)?.data_type(), &data_type); + let schema = df.schema(); + + for (name, data_type) in &expected_types { + assert_eq!(schema.field_with_name(None, name)?.data_type(), data_type); + } + + let rows = df.sort(vec![col("i32").sort(true, true)])?; + + assert_batches_eq!( + &[ + "+-------+----+-----+-----+-----+----+-----+-----+-----+-----+-----+-----+-----+", + "| bool | i8 | i16 | i32 | i64 | u8 | u16 | u32 | u64 | f16 | f32 | f64 | str |", + "+-------+----+-----+-----+-----+----+-----+-----+-----+-----+-----+-----+-----+", + "| true | -1 | -1 | -1 | -1 | 0 | 0 | 0 | 0 | 1 | 1.0 | 1.0 | foo |", + "| false | 0 | 0 | 0 | 0 | 1 | 1 | 1 | 1 | 2 | 2.0 | 2.0 | bar |", + "| true | 1 | 1 | 1 | 1 | 2 | 2 | 2 | 2 | 3 | 3.0 | 3.0 | |", + "+-------+----+-----+-----+-----+----+-----+-----+-----+-----+-----+-----+-----+", + ], + &rows.collect().await? + ); } - let rows = df.sort(vec![col("i32").sort(true, true)])?; + Ok(()) +} + +#[test] +fn test_dataframe_from_columns_empty() { + let result = DataFrame::from_columns(vec![]); + assert!(result.is_err()); + + let result = DataFrame::from_columns([]); + assert!(result.is_err()); +} + +#[tokio::test] +async fn test_dataframe_from_columns_with_iterator() -> Result<()> { + let bools: ArrayRef = Arc::new(BooleanArray::from(vec![true, false, true])); + let i8s: ArrayRef = Arc::new(Int8Array::from(vec![-1, 0, 1])); + let i16s: ArrayRef = Arc::new(Int16Array::from(vec![-1, 0, 1])); + let i32s: ArrayRef = Arc::new(Int32Array::from(vec![-1, 0, 1])); + let i64s: ArrayRef = Arc::new(Int64Array::from(vec![-1, 0, 1])); + + let u8s: ArrayRef = Arc::new(UInt8Array::from(vec![0, 1, 2])); + let u16s: ArrayRef = Arc::new(UInt16Array::from(vec![0, 1, 2])); + let u32s: ArrayRef = Arc::new(UInt32Array::from(vec![0, 1, 2])); + let u64s: ArrayRef = Arc::new(UInt64Array::from(vec![0, 1, 2])); + + let f16s: ArrayRef = Arc::new(Float16Array::from(vec![ + half::f16::from_f64(1.0), + half::f16::from_f64(2.0), + half::f16::from_f64(3.0), + ])); + let f32s: ArrayRef = Arc::new(Float32Array::from(vec![1.0, 2.0, 3.0])); + let f64s: ArrayRef = Arc::new(Float64Array::from(vec![1.0, 2.0, 3.0])); + + let strings: ArrayRef = + Arc::new(StringArray::from(vec![Some("foo"), Some("bar"), None])); + + let columns = [ + ("bool", bools), + ("i8", i8s), + ("i16", i16s), + ("i32", i32s), + ("i64", i64s), + ("u8", u8s), + ("u16", u16s), + ("u32", u32s), + ("u64", u64s), + ("f16", f16s), + ("f32", f32s), + ("f64", f64s), + ("str", strings), + ]; + let df = DataFrame::from_columns(columns.into_iter())?; + + assert_eq!(df.schema().fields().len(), 13); + assert_eq!(df.clone().count().await?, 3); + let rows = df.sort(vec![col("i32").sort(true, true)])?; assert_batches_eq!( &[ "+-------+----+-----+-----+-----+----+-----+-----+-----+-----+-----+-----+-----+", diff --git a/docs/source/library-user-guide/upgrading/56.0.0.md b/docs/source/library-user-guide/upgrading/56.0.0.md index 97ac11a0169a0..881aee0c6308f 100644 --- a/docs/source/library-user-guide/upgrading/56.0.0.md +++ b/docs/source/library-user-guide/upgrading/56.0.0.md @@ -97,4 +97,25 @@ The output type of the `floor` and `ceil` UDFs has been changed from the exact i Change the expected type or wrap the expression in `CAST`. It's recommended to avoid relying on decimal's exact precision and scale. -[#24703]: https://github.com/apache/datafusion/pull/24703 +### `DataFrame::from_columns` accepts `IntoIterator` + +`DataFrame::from_columns` now accepts any `IntoIterator` +instead of specifically accepting a `Vec<(&str, ArrayRef)>`. + +```rust,ignore +// Existing Vec usage continues to work +let df = DataFrame::from_columns(vec![ + ("id", id), + ("name", name), +])?; + +// Arrays can now be used directly +let df = DataFrame::from_columns([ + ("id", id), + ("name", name), +])?; +``` + +Most existing call sites using `Vec` require no changes. Code that relies on +the exact non-generic function signature of `DataFrame::from_columns` may need +to be updated to account for the new generic API. From 69a4c72264a3abd7f491f5d13013b0d57ab14c68 Mon Sep 17 00:00:00 2001 From: Hasnaat hussain <110020083+Hasnaathussain@users.noreply.github.com> Date: Thu, 3 Sep 2026 04:40:09 +0000 Subject: [PATCH 25/37] Fix CASE evaluation for custom column expressions (#24484) ## Which issue does this PR address? - Part of #21231. ## Rationale for this change `CaseBody::project` discovers input columns by downcasting expression-tree nodes to DataFusion's concrete `Column` or `LambdaVariable`. A third-party column-like `PhysicalExpr` can evaluate an input column without exposing either node. The projection then omits that dependency and CASE evaluates the expression against a batch with the wrong schema. ## What changes are included in this PR? CASE projection now falls back to the original input batch when it encounters an unknown leaf expression. Built-in `ScalarFunctionExpr` leaves remain eligible for projection because they receive the row count and do not read input columns through the batch. The projection decision is centralized in `ProjectedCaseBody::projection_for`, with the invariant documented next to the representation. The regression coverage exercises searched CASE, base-expression CASE, and the single-WHEN expression path with a custom column-like expression. It also checks that a nullary `random()` scalar function keeps the projection fast path. This is a conservative correctness fix for the custom-leaf case. The broader dependency-reporting design discussed in #21231 remains separate work. ## Are these changes tested? - `cargo fmt --all -- --check` - `cargo clippy -p datafusion-physical-expr --all-targets --all-features -- -D warnings` - `cargo test -p datafusion-physical-expr expressions::case::tests` - `cargo test -p datafusion-physical-expr` ## Are there any user-facing changes? Custom physical expressions can now be evaluated correctly inside searched and base-expression CASE expressions. There is no public API change. --------- Signed-off-by: Hasnaat Hussain --- .../physical-expr/src/expressions/case.rs | 260 +++++++++++++++--- 1 file changed, 227 insertions(+), 33 deletions(-) diff --git a/datafusion/physical-expr/src/expressions/case.rs b/datafusion/physical-expr/src/expressions/case.rs index e95e9e570cf1b..dd98029c739a8 100644 --- a/datafusion/physical-expr/src/expressions/case.rs +++ b/datafusion/physical-expr/src/expressions/case.rs @@ -18,10 +18,10 @@ mod literal_lookup_table; use super::{Column, Literal}; -use crate::PhysicalExpr; use crate::expressions::{ CastExpr, LambdaVariable, NegativeExpr, NotExpr, lit, try_cast, }; +use crate::{PhysicalExpr, ScalarFunctionExpr}; use arrow::array::*; use arrow::compute::kernels::zip::zip; use arrow::compute::{ @@ -133,6 +133,7 @@ impl CaseBody { // Determine the set of columns that are used in all the expressions of the case body. // Use an ordered set so lambda variables continue to be positioned after columns let mut used_column_indices = BTreeSet::::new(); + let mut supports_projection = true; let mut collect_column_indices = |expr: &Arc| { expr.apply(|expr| { if let Some(column) = expr.downcast_ref::() { @@ -141,6 +142,12 @@ impl CaseBody { expr.downcast_ref::() { used_column_indices.insert(lambda_variable.index()); + } else if expr.downcast_ref::().is_none() + && expr.downcast_ref::().is_none() + && expr.children().is_empty() + { + // Unknown leaves may read input columns without exposing a Column child. + supports_projection = false; } Ok(TreeNodeRecursion::Continue) }) @@ -216,6 +223,7 @@ impl CaseBody { Ok(ProjectedCaseBody { projection, body: projected_body, + supports_projection, }) } } @@ -247,10 +255,32 @@ impl CaseBody { /// /// The projection vector and the rewritten expression (which only differs from the original in /// column reference indices) are held in a `ProjectedCaseBody`. +/// +/// When `supports_projection` is false, evaluation uses the original body and full batch because +/// the expression may depend on input columns that are not visible in its children. #[derive(Debug, Hash, PartialEq, Eq)] struct ProjectedCaseBody { projection: Vec, body: CaseBody, + /// Whether `body` is safe to evaluate against a projected batch. + /// + /// When false, `projection` and the rewritten `body` must not be used because an + /// expression may read input columns without exposing them through its children. + supports_projection: bool, +} + +impl ProjectedCaseBody { + /// Returns a projection when evaluating the derived body can avoid unused input columns. + fn projection_for(&self, batch: &RecordBatch) -> Option> { + let projection = self + .projection + .iter() + .copied() + .filter(|index| *index < batch.num_columns()) + .collect::>(); + (self.supports_projection && projection.len() < batch.num_columns()) + .then_some(projection) + } } /// The CASE expression is similar to a series of nested if/else and there are two forms that @@ -1049,14 +1079,7 @@ impl CaseExpr { projected: &ProjectedCaseBody, ) -> Result { let return_type = self.data_type(&batch.schema())?; - // projected.projection may include indexes of lambda variables not available on this batch - let projection = projected - .projection - .iter() - .copied() - .filter(|index| *index < batch.num_columns()) - .collect::>(); - if projection.len() < batch.num_columns() { + if let Some(projection) = projected.projection_for(batch) { let projected_batch = batch.project(&projection)?; projected .body @@ -1079,14 +1102,7 @@ impl CaseExpr { projected: &ProjectedCaseBody, ) -> Result { let return_type = self.data_type(&batch.schema())?; - // projected.projection may include indexes of lambda variables not available on this batch - let projection = projected - .projection - .iter() - .copied() - .filter(|index| *index < batch.num_columns()) - .collect::>(); - if projection.len() < batch.num_columns() { + if let Some(projection) = projected.projection_for(batch) { let projected_batch = batch.project(&projection)?; projected .body @@ -1204,23 +1220,14 @@ impl CaseExpr { )?)) } } + } else if let Some(projection) = projected.projection_for(batch) { + // The case expressions do not use all the columns of the input batch. + // Project first to reduce time spent filtering. + let projected_batch = batch.project(&projection)?; + projected.body.expr_or_expr(&projected_batch, when_value) } else { - // projected.projection may include indexes of lambda variables not available on this batch - let projection = projected - .projection - .iter() - .copied() - .filter(|index| *index < batch.num_columns()) - .collect::>(); - if projection.len() < batch.num_columns() { - // The case expressions do not use all the columns of the input batch. - // Project first to reduce time spent filtering. - let projected_batch = batch.project(&projection)?; - projected.body.expr_or_expr(&projected_batch, when_value) - } else { - // All columns are used in the case expressions, so there is no need to project. - self.body.expr_or_expr(batch, when_value) - } + // All columns are used in the case expressions, so there is no need to project. + self.body.expr_or_expr(batch, when_value) } } @@ -1592,6 +1599,7 @@ mod tests { use arrow::datatypes::DataType::Float64; use arrow::datatypes::Field; use datafusion_common::cast::{as_float64_array, as_int32_array}; + use datafusion_common::config::ConfigOptions; use datafusion_common::plan_err; use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode}; use datafusion_expr::type_coercion::binary::type_union_coercion; @@ -1599,6 +1607,55 @@ mod tests { use datafusion_physical_expr_common::physical_expr::fmt_sql; use half::f16; + #[derive(Debug, Hash, PartialEq, Eq)] + struct CustomColumn { + inner: Column, + } + + impl CustomColumn { + fn new(name: &str, index: usize) -> Self { + Self { + inner: Column::new(name, index), + } + } + } + + impl std::fmt::Display for CustomColumn { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + std::fmt::Display::fmt(&self.inner, f) + } + } + + impl PhysicalExpr for CustomColumn { + fn data_type(&self, input_schema: &Schema) -> Result { + self.inner.data_type(input_schema) + } + + fn nullable(&self, input_schema: &Schema) -> Result { + self.inner.nullable(input_schema) + } + + fn evaluate(&self, batch: &RecordBatch) -> Result { + self.inner.evaluate(batch) + } + + fn children(&self) -> Vec<&Arc> { + vec![] + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + assert!(children.is_empty()); + Ok(self) + } + + fn fmt_sql(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + self.inner.fmt_sql(f) + } + } + #[test] fn case_with_expr() -> Result<()> { let batch = case_test_batch()?; @@ -1894,6 +1951,127 @@ mod tests { Ok(()) } + #[test] + fn case_without_expr_with_custom_column() -> Result<()> { + let batch = case_test_batch()?; + let schema = batch.schema(); + + let when1 = binary( + Arc::new(CustomColumn::new("a", 0)), + Operator::Eq, + lit("foo"), + &schema, + )?; + let when2 = binary( + Arc::new(CustomColumn::new("a", 0)), + Operator::Eq, + lit("bar"), + &schema, + )?; + let expr = generate_case_when_with_type_coercion( + None, + vec![(when1, lit(123i32)), (when2, lit(456i32))], + None, + schema.as_ref(), + )?; + + let result = expr.evaluate(&batch)?.into_array(batch.num_rows())?; + let result = as_int32_array(&result)?; + let expected = Int32Array::from(vec![Some(123), None, None, Some(456)]); + + assert_eq!(&expected, result); + Ok(()) + } + + #[test] + fn case_with_expr_with_custom_column() -> Result<()> { + let batch = case_test_batch_with_extra_columns()?; + let schema = batch.schema(); + + let expr = CaseExpr::try_new( + Some(Arc::new(CustomColumn::new("a", 0))), + vec![(lit("foo"), col("b", &schema)?)], + Some(col("c", &schema)?), + )?; + + match &expr.eval_method { + EvalMethod::WithExpression(projected) => { + assert!(!projected.supports_projection); + } + method => panic!("expected WithExpression, got {method:?}"), + } + + let result = expr.evaluate(&batch)?.into_array(batch.num_rows())?; + let result = as_int32_array(&result)?; + let expected = Int32Array::from(vec![Some(1), Some(20), Some(30), Some(40)]); + + assert_eq!(&expected, result); + Ok(()) + } + + #[test] + fn case_expr_or_expr_with_custom_column() -> Result<()> { + let batch = case_test_batch_with_extra_columns()?; + let schema = batch.schema(); + let when = binary( + Arc::new(CustomColumn::new("a", 0)), + Operator::Eq, + lit("foo"), + &schema, + )?; + + let expr = CaseExpr::try_new( + None, + vec![(when, col("b", &schema)?)], + Some(col("c", &schema)?), + )?; + + match &expr.eval_method { + EvalMethod::ExpressionOrExpression(projected) => { + assert!(!projected.supports_projection); + } + method => panic!("expected ExpressionOrExpression, got {method:?}"), + } + + let result = expr.evaluate(&batch)?.into_array(batch.num_rows())?; + let result = as_int32_array(&result)?; + let expected = Int32Array::from(vec![Some(1), Some(20), Some(30), Some(40)]); + + assert_eq!(&expected, result); + Ok(()) + } + + #[test] + fn case_expr_or_expr_keeps_nullary_scalar_functions_projectable() -> Result<()> { + let batch = case_test_batch1()?; + let schema = batch.schema(); + let random: Arc = Arc::new(ScalarFunctionExpr::try_new( + datafusion_functions::math::random(), + vec![], + &schema, + Arc::new(ConfigOptions::default()), + )?); + let when = binary(random, Operator::Gt, lit(0.5f64), &schema)?; + + let expr = CaseExpr::try_new( + None, + vec![(when, col("b", &schema)?)], + Some(col("c", &schema)?), + )?; + + match &expr.eval_method { + EvalMethod::ExpressionOrExpression(projected) => { + assert!(projected.supports_projection); + assert_eq!(projected.projection, vec![1, 2]); + } + method => panic!("expected ExpressionOrExpression, got {method:?}"), + } + + let result = expr.evaluate(&batch)?.into_array(batch.num_rows())?; + assert_eq!(result.len(), batch.num_rows()); + Ok(()) + } + #[test] fn case_with_expr_when_null() -> Result<()> { let batch = case_test_batch()?; @@ -1961,6 +2139,22 @@ mod tests { Ok(()) } + fn case_test_batch_with_extra_columns() -> Result { + let schema = Schema::new(vec![ + Field::new("a", DataType::Utf8, true), + Field::new("b", DataType::Int32, true), + Field::new("c", DataType::Int32, true), + ]); + let a = StringArray::from(vec![Some("foo"), Some("baz"), None, Some("bar")]); + let b = Int32Array::from(vec![Some(1), Some(2), Some(3), Some(4)]); + let c = Int32Array::from(vec![Some(10), Some(20), Some(30), Some(40)]); + RecordBatch::try_new( + Arc::new(schema), + vec![Arc::new(a), Arc::new(b), Arc::new(c)], + ) + .map_err(Into::into) + } + fn case_test_batch1() -> Result { let schema = Schema::new(vec![ Field::new("a", DataType::Int32, true), From 7e5561dff96a8103fc8a28a0723deb6b252a6b49 Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Thu, 3 Sep 2026 06:03:33 +0000 Subject: [PATCH 26/37] fix: clear memory after emit all in count distinct (#24888) ## Which issue does this PR close? N/A ## Rationale for this change after emit all, the memory is still being held in count distinct causing OOM issues ## What changes are included in this PR? release memory in emit all and added tests ## What is the testing strategy for this PR? integration test ## Are there any user-facing changes? no ---- Founded while running: - #24881 --- datafusion/core/tests/memory_limit/mod.rs | 84 ++++++++++++++++++- .../src/aggregate/count_distinct/groups.rs | 9 +- 2 files changed, 90 insertions(+), 3 deletions(-) diff --git a/datafusion/core/tests/memory_limit/mod.rs b/datafusion/core/tests/memory_limit/mod.rs index 15b224d200bf4..39be33587fb90 100644 --- a/datafusion/core/tests/memory_limit/mod.rs +++ b/datafusion/core/tests/memory_limit/mod.rs @@ -26,11 +26,14 @@ mod nlj_spill_unmatched; mod repartition_mem_limit; mod union_nullable_spill; mod view_spill_compaction; -use arrow::array::{ArrayRef, DictionaryArray, Int32Array, RecordBatch, StringViewArray}; +use arrow::array::{ + ArrayRef, DictionaryArray, Int32Array, Int64Array, RecordBatch, StringViewArray, +}; use arrow::compute::SortOptions; use arrow::datatypes::{Int32Type, SchemaRef}; use arrow_schema::{DataType, Field, Schema}; use datafusion::assert_batches_eq; +use datafusion::assert_batches_sorted_eq; use datafusion::config::SpillCompression; use datafusion::datasource::memory::MemorySourceConfig; use datafusion::datasource::source::DataSourceExec; @@ -125,6 +128,85 @@ async fn group_by_hash() { .await } +/// `count(distinct)` over integers under a memory limit. +/// +/// The integer distinct-count groups accumulator reports the capacity of its +/// buffers in `size()`. After an aggregate stream emits all groups, either to +/// emit partial state early or to spill, it resizes its reservation to the +/// table's reported size and expects it to have shrunk. If the accumulator +/// keeps its capacity, that resize is a grow against an exhausted pool and the +/// query fails although everything was already written out. +const COUNT_DISTINCT_ROWS: usize = 200_000; +const COUNT_DISTINCT_GROUPS: i64 = 64; +const COUNT_DISTINCT_BATCH_ROWS: usize = 8_192; + +/// Far below the distinct sets (200k values, several megabytes across the +/// partial and final tables), far above the fixed cost of the stages. With the +/// accumulator releasing its buffers the query passes from 2 MB upwards; +/// without, it fails up to 4 MB with "Decreasing allocation after spilling +/// should succeed" in the final stage or a failed emit in the partial stage. +const COUNT_DISTINCT_MEMORY_LIMIT: usize = 4 * 1024 * 1024; + +/// `g` has 64 groups, `v` is unique, so every group holds 3125 distinct values. +fn count_distinct_table() -> MemTable { + let schema = Arc::new(Schema::new(vec![ + Field::new("g", DataType::Int64, false), + Field::new("v", DataType::Int64, false), + ])); + let batches = (0..COUNT_DISTINCT_ROWS) + .step_by(COUNT_DISTINCT_BATCH_ROWS) + .map(|start| { + let rows = + start..(start + COUNT_DISTINCT_BATCH_ROWS).min(COUNT_DISTINCT_ROWS); + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int64Array::from_iter_values( + rows.clone().map(|row| row as i64 % COUNT_DISTINCT_GROUPS), + )), + Arc::new(Int64Array::from_iter_values(rows.map(|row| row as i64))), + ], + ) + .unwrap() + }) + .collect(); + MemTable::try_new(schema, vec![batches]).unwrap() +} + +/// Four partial stages emit their state early and four hash-partitioned final +/// stages spill; every one of them must see the accumulator memory drop after +/// emitting all groups. +#[tokio::test] +async fn count_distinct_releases_memory_after_emitting_all() { + let runtime = RuntimeEnvBuilder::new() + .with_memory_limit(COUNT_DISTINCT_MEMORY_LIMIT, 1.0) + .with_disk_manager_builder(DiskManagerBuilder::default()) + .build_arc() + .unwrap(); + let config = SessionConfig::new().with_target_partitions(4); + let ctx = SessionContext::new_with_config_rt(config, runtime); + ctx.register_table("t", Arc::new(count_distinct_table())) + .unwrap(); + + let batches = ctx + .sql("select count(distinct v) as d, count(*) as n from t group by g") + .await + .unwrap() + .collect() + .await + .unwrap_or_else(|error| panic!("query failed under the memory limit: {error}")); + + let per_group = (COUNT_DISTINCT_ROWS as i64 / COUNT_DISTINCT_GROUPS).to_string(); + let row = format!("| {per_group} | {per_group} |"); + let mut expected = vec!["+------+------+", "| d | n |", "+------+------+"]; + expected.extend(std::iter::repeat_n( + row.as_str(), + COUNT_DISTINCT_GROUPS as usize, + )); + expected.push("+------+------+"); + assert_batches_sorted_eq!(expected, &batches); +} + #[tokio::test] async fn join_by_key_multiple_partitions() { let config = SessionConfig::new().with_target_partitions(2); diff --git a/datafusion/functions-aggregate-common/src/aggregate/count_distinct/groups.rs b/datafusion/functions-aggregate-common/src/aggregate/count_distinct/groups.rs index 6e3e3b91a74f7..38f6d26200ec9 100644 --- a/datafusion/functions-aggregate-common/src/aggregate/count_distinct/groups.rs +++ b/datafusion/functions-aggregate-common/src/aggregate/count_distinct/groups.rs @@ -89,7 +89,10 @@ where match emit_to { EmitTo::All => { - self.seen.clear(); + // Release the capacity, not just the entries: `size()` reports + // capacity, and the aggregate streams rely on it dropping after + // emitting everything. + self.seen = HashSet::default(); } EmitTo::First(n) => { let mut remaining = HashSet::default(); @@ -145,7 +148,9 @@ where all_values[pos] = value; cursors[group_idx] += 1; } - self.counts.clear(); + // Release the capacity, see `evaluate`. + self.seen = HashSet::default(); + self.counts = Vec::new(); } else { let mut remaining = HashSet::default(); for (group_idx, value) in self.seen.drain() { From a6f1dce70c6cb1faac48ec51f85f6d6e53829a4a Mon Sep 17 00:00:00 2001 From: blinding-pixels Date: Thu, 3 Sep 2026 06:38:48 +0000 Subject: [PATCH 27/37] fix: avoid invalid qualifiers in unparsed subqueries (#24808) ## Which issue does this PR close? - Closes #13027. ## Rationale for this change I noticed that this issue had an earlier PR from almost two years ago. That implementation was closed because it used broad identifier cleanup, and maintainers were concerned it could introduce subtle bugs. Since then, DataFusion has added much of the machinery needed for the approaches recommended in that review. This solution uses that newer machinery and combines both approaches. It removes the qualifier when the derived table does not need an alias, and rewrites the reference to the derived alias when the dialect requires one. ## What changes are included in this PR? This PR detects when a projection enters a new derived-table scope and prevents its outer expressions from referring to an inner table alias that is no longer visible. For dialects that do not require a derived-table alias, the invalid inner qualifier is removed. For dialects such as MySQL that require an alias, the outer reference is rewritten to the generated derived-table alias. Explicitly named subquery scopes remain unchanged. ## What is the testing strategy for this PR? The SQL unparser round-trip tests cover the issue's original query with both the generic and MySQL dialects. They also cover filtered and distinct derived inputs, and update existing limit and nested-projection cases to assert valid outer references. The required formatting, Clippy, and extended workspace test suite all pass, including all 505 SQL logic test files. ## Are there any user-facing changes? Yes. SQL produced by the unparser no longer contains qualifiers that refer to tables outside their visible scope. Dialects that require derived-table aliases now qualify the outer reference with the generated alias. There are no public API changes. --------- Co-authored-by: blinding-pixels <281499151+blinding-pixels@users.noreply.github.com> --- datafusion/sql/src/unparser/ast.rs | 16 +++ datafusion/sql/src/unparser/plan.rs | 160 ++++++++++++++++++++-- datafusion/sql/tests/cases/plan_to_sql.rs | 87 +++++++++++- 3 files changed, 251 insertions(+), 12 deletions(-) diff --git a/datafusion/sql/src/unparser/ast.rs b/datafusion/sql/src/unparser/ast.rs index 7418d0b5b7605..c335d3ee4c16d 100644 --- a/datafusion/sql/src/unparser/ast.rs +++ b/datafusion/sql/src/unparser/ast.rs @@ -167,6 +167,8 @@ pub struct SelectBuilder { /// Table aliases that correspond to LATERAL FLATTEN relations. /// Column references into these aliases must use `VALUE` as the column name. flatten_table_aliases: Vec, + /// Depth of explicitly named subqueries currently being rendered. + subquery_alias_depth: usize, } /// Prefix used for auto-generated LATERAL FLATTEN table aliases. @@ -195,6 +197,19 @@ impl SelectBuilder { self.flatten_table_aliases.iter().any(|a| a == alias) } + pub(super) fn enter_subquery_alias(&mut self) { + self.subquery_alias_depth += 1; + } + + pub(super) fn exit_subquery_alias(&mut self) { + debug_assert!(self.subquery_alias_depth > 0); + self.subquery_alias_depth -= 1; + } + + pub(super) fn inside_subquery_alias(&self) -> bool { + self.subquery_alias_depth > 0 + } + /// Returns the most recently generated flatten alias, or `None` if /// `next_flatten_alias` has not been called yet. pub fn current_flatten_alias(&self) -> Option { @@ -419,6 +434,7 @@ impl SelectBuilder { flavor: Some(SelectFlavor::Standard), flatten_alias_counter: 0, flatten_table_aliases: Vec::new(), + subquery_alias_depth: 0, } } } diff --git a/datafusion/sql/src/unparser/plan.rs b/datafusion/sql/src/unparser/plan.rs index 30320acbb24dc..1a4a094e96ed6 100644 --- a/datafusion/sql/src/unparser/plan.rs +++ b/datafusion/sql/src/unparser/plan.rs @@ -54,7 +54,7 @@ use datafusion_expr::{ TableScan, Unnest, UserDefinedLogicalNode, Window, expr::Alias, }; use sqlparser::ast::{self, Ident, OrderByKind, SetExpr, TableAliasColumnDef}; -use std::{sync::Arc, vec}; +use std::{collections::HashSet, sync::Arc, vec}; /// Convert a DataFusion [`LogicalPlan`] to [`ast::Statement`] /// @@ -164,6 +164,12 @@ impl<'a> UnparserAggScope<'a> { } } +#[derive(Clone, Copy)] +struct DerivedInputScope<'a> { + alias: &'static str, + schema: &'a DFSchema, +} + impl Unparser<'_> { pub fn plan_to_sql(&self, plan: &LogicalPlan) -> Result { let mut plan = normalize_union_schema(plan)?; @@ -451,6 +457,102 @@ impl Unparser<'_> { } } + /// Return the alias recursion would assign when `plan` must become a + /// derived relation below an already rendered projection. + fn derived_input_alias(plan: &LogicalPlan) -> Option<&'static str> { + match plan { + LogicalPlan::Projection(_) => Some("derived_projection"), + LogicalPlan::Limit(_) => Some("derived_limit"), + LogicalPlan::Sort(_) => Some("derived_sort"), + LogicalPlan::Distinct(_) => Some("derived_distinct"), + LogicalPlan::Filter(filter) => { + Self::derived_input_alias(filter.input.as_ref()) + } + LogicalPlan::Repartition(repartition) => { + Self::derived_input_alias(repartition.input.as_ref()) + } + _ => None, + } + } + + fn derived_input_scope<'a>( + plan: &'a LogicalPlan, + select: &SelectBuilder, + ) -> Option> { + if select.inside_subquery_alias() { + return None; + } + + if select.already_projected() + && find_unnest_node_within_select(plan).is_none() + && let Some(alias) = Self::derived_input_alias(plan) + { + return Some(DerivedInputScope { + alias, + schema: plan.schema().as_ref(), + }); + } + + match plan { + LogicalPlan::Projection(projection) => { + let alias = Self::derived_input_alias(projection.input.as_ref())?; + let qualified_projection = projection.expr.iter().any(|expr| { + expr.column_refs() + .iter() + .any(|column| column.relation.is_some()) + }); + let mut input_names = HashSet::new(); + let unique_input_names = projection + .input + .schema() + .fields() + .iter() + .all(|field| input_names.insert(field.name())); + + (qualified_projection + && unique_input_names + && find_unnest_node_within_select(plan).is_none()) + .then_some(DerivedInputScope { + alias, + schema: projection.input.schema().as_ref(), + }) + } + LogicalPlan::Filter(filter) => { + Self::derived_input_scope(filter.input.as_ref(), select) + } + LogicalPlan::Limit(limit) => { + Self::derived_input_scope(limit.input.as_ref(), select) + } + LogicalPlan::Sort(sort) => { + Self::derived_input_scope(sort.input.as_ref(), select) + } + LogicalPlan::Repartition(repartition) => { + Self::derived_input_scope(repartition.input.as_ref(), select) + } + _ => None, + } + } + + fn rebase_derived_input_expr( + &self, + expr: Expr, + scope: Option>, + ) -> Result { + let Some(scope) = scope else { + return Ok(expr); + }; + if self.dialect.requires_derived_table_alias() { + let mut alias_rewriter = TableAliasRewriter { + table_schema: scope.schema, + alias_name: TableReference::bare(scope.alias), + rewrite_unqualified: false, + }; + expr.rewrite(&mut alias_rewriter).data() + } else { + Self::strip_column_qualifiers_for_schema(expr, scope.schema) + } + } + fn contains_aggregate_before_relation(plan: &LogicalPlan) -> bool { match plan { LogicalPlan::Aggregate(_) => true, @@ -829,6 +931,26 @@ impl Unparser<'_> { columns, ); } + + if let Some(scope) = Self::derived_input_scope(plan, select) { + // The input is about to enter a new SQL scope. Preserve that + // boundary explicitly and make the outer expressions resolve + // against the relation that will actually be visible there. + let requires_alias = self.dialect.requires_derived_table_alias(); + let alias = requires_alias + .then(|| self.new_table_alias(scope.alias.to_string(), vec![])); + self.derive(p.input.as_ref(), relation, alias, false)?; + + let items = p + .expr + .iter() + .cloned() + .map(|expr| self.rebase_derived_input_expr(expr, Some(scope))) + .map(|expr| self.select_item_to_sql(&expr?)) + .collect::>>()?; + select.projection(items); + return Ok(()); + } // For Snowflake FLATTEN: when the outer Projection has // UNNEST(...) display-name columns (from SELECT * / SELECT // UNNEST(...)), generate a flatten alias now so that @@ -1072,6 +1194,7 @@ impl Unparser<'_> { self.select_to_sql_recursively(cur, query, select, relation) } LogicalPlan::Filter(filter) => { + let derived_input_scope = Self::derived_input_scope(plan, select); let window = find_window_nodes_within_select( plan, None, @@ -1088,15 +1211,23 @@ impl Unparser<'_> { unprojected = UnparserAggScope::new(agg).prepare(unprojected, None)?; } + unprojected = + self.rebase_derived_input_expr(unprojected, derived_input_scope)?; let filter_expr = self.expr_to_sql(&unprojected)?; select.qualify(Some(filter_expr)); } else if let Some(agg) = agg { - let unprojected = UnparserAggScope::new(agg) + let mut unprojected = UnparserAggScope::new(agg) .prepare(filter.predicate.clone(), None)?; + unprojected = + self.rebase_derived_input_expr(unprojected, derived_input_scope)?; let filter_expr = self.expr_to_sql(&unprojected)?; select.having(Some(filter_expr)); } else { - let filter_expr = self.expr_to_sql(&filter.predicate)?; + let predicate = self.rebase_derived_input_expr( + filter.predicate.clone(), + derived_input_scope, + )?; + let filter_expr = self.expr_to_sql(&predicate)?; select.selection(Some(filter_expr)); } @@ -1172,17 +1303,25 @@ impl Unparser<'_> { )))); } + let derived_input_scope = Self::derived_input_scope(plan, select); let agg = find_agg_node_within_select(plan, select.already_projected()); // unproject sort expressions let sort_exprs: Vec = sort .expr .iter() .map(|sort_expr| { - Self::unproject_sort_expr_in_scope( + let sort_expr = Self::unproject_sort_expr_in_scope( sort_expr.clone(), agg, sort.input.as_ref(), - ) + )?; + Ok(SortExpr { + expr: self.rebase_derived_input_expr( + sort_expr.expr, + derived_input_scope, + )?, + ..sort_expr + }) }) .collect::>>()?; @@ -1537,7 +1676,8 @@ impl Unparser<'_> { )]); } let plan = unparsed_table_scan.unwrap_or_else(|| plan.clone()); - if !columns.is_empty() + select.enter_subquery_alias(); + let recursive_result = if !columns.is_empty() && !self.dialect.supports_column_alias_in_table_alias() { // Instead of specifying column aliases as part of the outer table, inject them directly into the inner projection @@ -1558,10 +1698,12 @@ impl Unparser<'_> { query, select, relation, - )?; + ) } else { - self.select_to_sql_recursively(&plan, query, select, relation)?; - } + self.select_to_sql_recursively(&plan, query, select, relation) + }; + select.exit_subquery_alias(); + recursive_result?; relation.alias(Some( self.new_table_alias(plan_alias.alias.table().to_string(), columns), diff --git a/datafusion/sql/tests/cases/plan_to_sql.rs b/datafusion/sql/tests/cases/plan_to_sql.rs index d6c31570bf1b0..32dfade056dd2 100644 --- a/datafusion/sql/tests/cases/plan_to_sql.rs +++ b/datafusion/sql/tests/cases/plan_to_sql.rs @@ -392,13 +392,94 @@ fn roundtrip_statement_with_dialect_4() -> Result<(), DataFusionError> { Ok(()) } +#[test] +fn roundtrip_rebases_derived_projection_references() -> Result<(), DataFusionError> { + roundtrip_statement_with_dialect_helper!( + sql: "select j1_id from (select ta.j1_id as j1_id from j1 ta);", + parser_dialect: GenericDialect {}, + unparser_dialect: UnparserDefaultDialect {}, + expected: @"SELECT j1_id FROM (SELECT ta.j1_id FROM j1 AS ta)", + ); + roundtrip_statement_with_dialect_helper!( + sql: "select j1_id from (select ta.j1_id as j1_id from j1 ta);", + parser_dialect: MySqlDialect {}, + unparser_dialect: UnparserMySqlDialect {}, + expected: @"SELECT `derived_projection`.`j1_id` FROM (SELECT `ta`.`j1_id` FROM `j1` AS `ta`) AS `derived_projection`", + ); + roundtrip_statement_with_dialect_helper!( + sql: "select j1_id from (select ta.j1_id as j1_id from j1 ta) order by j1_id;", + parser_dialect: GenericDialect {}, + unparser_dialect: UnparserDefaultDialect {}, + expected: @"SELECT j1_id FROM (SELECT ta.j1_id FROM j1 AS ta) ORDER BY j1_id ASC NULLS LAST", + ); + roundtrip_statement_with_dialect_helper!( + sql: "select j1_id from (select ta.j1_id as j1_id from j1 ta) order by j1_id;", + parser_dialect: MySqlDialect {}, + unparser_dialect: UnparserMySqlDialect {}, + expected: @"SELECT `derived_projection`.`j1_id` FROM (SELECT `ta`.`j1_id` FROM `j1` AS `ta`) AS `derived_projection` ORDER BY `derived_projection`.`j1_id` ASC", + ); + roundtrip_statement_with_dialect_helper!( + sql: "select j1_id from (select ta.j1_id as j1_id from j1 ta) where j1_id > 1;", + parser_dialect: GenericDialect {}, + unparser_dialect: UnparserDefaultDialect {}, + expected: @"SELECT j1_id FROM (SELECT ta.j1_id FROM j1 AS ta WHERE (ta.j1_id > 1))", + ); + roundtrip_statement_with_dialect_helper!( + sql: "select j1_id from (select ta.j1_id as j1_id from j1 ta) where j1_id > 1;", + parser_dialect: MySqlDialect {}, + unparser_dialect: UnparserMySqlDialect {}, + expected: @"SELECT `derived_projection`.`j1_id` FROM (SELECT `ta`.`j1_id` FROM `j1` AS `ta` WHERE (`ta`.`j1_id` > 1)) AS `derived_projection`", + ); + roundtrip_statement_with_dialect_helper!( + sql: "select j1_id from (select distinct ta.j1_id as j1_id from j1 ta);", + parser_dialect: GenericDialect {}, + unparser_dialect: UnparserDefaultDialect {}, + expected: @"SELECT j1_id FROM (SELECT DISTINCT ta.j1_id FROM j1 AS ta)", + ); + roundtrip_statement_with_dialect_helper!( + sql: "select j1_id from (select distinct ta.j1_id as j1_id from j1 ta);", + parser_dialect: MySqlDialect {}, + unparser_dialect: UnparserMySqlDialect {}, + expected: @"SELECT `derived_distinct`.`j1_id` FROM (SELECT DISTINCT `ta`.`j1_id` FROM `j1` AS `ta`) AS `derived_distinct`", + ); + + let statement = Parser::new(&GenericDialect {}) + .try_with_sql("select j1_id from (select ta.j1_id as j1_id from j1 ta)")? + .parse_statement()?; + let context = MockContextProvider { + state: MockSessionState::default(), + }; + let plan = SqlToRel::new(&context).sql_statement_to_plan(statement)?; + let plan = LogicalPlanBuilder::from(plan) + .filter(col("ta.j1_id").gt(lit(1)))? + .build()?; + let unparser = Unparser::new(&UnparserMySqlDialect {}); + assert_snapshot!( + unparser.plan_to_sql(&plan)?, + @"SELECT `derived_projection`.`j1_id` FROM (SELECT `ta`.`j1_id` FROM `j1` AS `ta`) AS `derived_projection` WHERE (`derived_projection`.`j1_id` > 1)" + ); + + let schema = Schema::new(vec![Field::new("j1_id", DataType::Int32, false)]); + let plan = table_scan(Some("j1"), &schema, None)? + .alias("ta")? + .project(vec![col("ta.j1_id")])? + .filter(col("ta.j1_id").gt(lit(0)))? + .project(vec![lit(1)])? + .build()?; + assert_snapshot!( + unparser.plan_to_sql(&plan)?, + @"SELECT 1 FROM (SELECT `ta`.`j1_id` FROM `j1` AS `ta`) AS `derived_projection` WHERE (`derived_projection`.`j1_id` > 0)" + ); + Ok(()) +} + #[test] fn roundtrip_statement_with_dialect_5() -> Result<(), DataFusionError> { roundtrip_statement_with_dialect_helper!( sql: "select j1_id from (select j1_id from j1 limit 10);", parser_dialect: MySqlDialect {}, unparser_dialect: UnparserMySqlDialect {}, - expected: @"SELECT `j1`.`j1_id` FROM (SELECT `j1`.`j1_id` FROM `j1` LIMIT 10) AS `derived_limit`", + expected: @"SELECT `derived_limit`.`j1_id` FROM (SELECT `j1`.`j1_id` FROM `j1` LIMIT 10) AS `derived_limit`", ); Ok(()) } @@ -1614,7 +1695,7 @@ fn test_table_scan_pushdown() -> Result<()> { plan_to_sql(&query_from_table_scan_with_two_projections)?; assert_snapshot!( query_from_table_scan_with_two_projections, - @"SELECT t1.id, t1.age FROM (SELECT t1.id, t1.age FROM t1)" + @"SELECT id, age FROM (SELECT t1.id, t1.age FROM t1)" ); let table_scan_with_filter = table_scan_with_filters( @@ -1791,7 +1872,7 @@ fn test_sort_with_scalar_fn_and_push_down_fetch() -> Result<()> { let sql = plan_to_sql(&plan)?; assert_snapshot!( sql, - @"SELECT t1.search_phrase FROM (SELECT t1.search_phrase, t1.event_time FROM t1 WHERE (t1.search_phrase <> '') ORDER BY substr(t1.event_time, 1, 5) ASC NULLS FIRST LIMIT 10)" + @"SELECT search_phrase FROM (SELECT t1.search_phrase, t1.event_time FROM t1 WHERE (t1.search_phrase <> '') ORDER BY substr(t1.event_time, 1, 5) ASC NULLS FIRST LIMIT 10)" ); Ok(()) } From 4a93adee064f63ae4d78fa92a39edf68c5dfb9f8 Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Thu, 3 Sep 2026 09:43:06 +0000 Subject: [PATCH 28/37] test: skip count distinct spill memory test under forced hash collisions (#24918) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Which issue does this PR close? Follow-up to #24888, which broke the `cargo test hash collisions` CI job on main. ## Rationale for this change The memory limit test added in #24888 fails when built with `force_hash_collisions`. Every key hashes to the same value there, so the hash repartition sends all 64 groups to one final stage. That single table needs 5.3 MB against the test's 4 MB pool, and it has nothing reserved yet, so there is nothing to spill. It fails no matter how well the accumulator releases memory, which is what the test is actually about. I tried a few ways to keep it running under the feature first: - **Bigger limit for the collision build.** Below 5.3 MB it dies on that one state batch; at 6 MB and up nothing spills, so the unfixed accumulator passes too and the test asserts nothing. Nothing in between. - **Single partition, no repartition at all.** Same wall. With 64 groups the whole distinct state lives in 64 rows, so total state and one batch are the same 5.3 MB. Also 82s instead of 0.18s. - **More groups, to spread the state over more batches.** With every key in one hash bucket, interning goes quadratic: 4096 groups did not finish in 400s. - **More rows (800k), to make total state exceed one batch.** Fails even with the fix. They all hit the same thing: under forced collisions the total state and a single batch are the same size, and the pool would have to sit above one and below the other. ## What changes are included in this PR? The test and its helpers move into a module gated on `not(feature = "force_hash_collisions")`. ## What is the testing strategy for this PR? `cargo test -p datafusion --features force_hash_collisions --test core_integration count_distinct_releases` runs 0 tests. Without the feature it still runs and passes. ## Are there any user-facing changes? No. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --- datafusion/core/tests/memory_limit/mod.rs | 173 ++++++++++++---------- 1 file changed, 93 insertions(+), 80 deletions(-) diff --git a/datafusion/core/tests/memory_limit/mod.rs b/datafusion/core/tests/memory_limit/mod.rs index 39be33587fb90..9172e34a87c79 100644 --- a/datafusion/core/tests/memory_limit/mod.rs +++ b/datafusion/core/tests/memory_limit/mod.rs @@ -26,14 +26,11 @@ mod nlj_spill_unmatched; mod repartition_mem_limit; mod union_nullable_spill; mod view_spill_compaction; -use arrow::array::{ - ArrayRef, DictionaryArray, Int32Array, Int64Array, RecordBatch, StringViewArray, -}; +use arrow::array::{ArrayRef, DictionaryArray, Int32Array, RecordBatch, StringViewArray}; use arrow::compute::SortOptions; use arrow::datatypes::{Int32Type, SchemaRef}; use arrow_schema::{DataType, Field, Schema}; use datafusion::assert_batches_eq; -use datafusion::assert_batches_sorted_eq; use datafusion::config::SpillCompression; use datafusion::datasource::memory::MemorySourceConfig; use datafusion::datasource::source::DataSourceExec; @@ -128,83 +125,99 @@ async fn group_by_hash() { .await } -/// `count(distinct)` over integers under a memory limit. -/// -/// The integer distinct-count groups accumulator reports the capacity of its -/// buffers in `size()`. After an aggregate stream emits all groups, either to -/// emit partial state early or to spill, it resizes its reservation to the -/// table's reported size and expects it to have shrunk. If the accumulator -/// keeps its capacity, that resize is a grow against an exhausted pool and the -/// query fails although everything was already written out. -const COUNT_DISTINCT_ROWS: usize = 200_000; -const COUNT_DISTINCT_GROUPS: i64 = 64; -const COUNT_DISTINCT_BATCH_ROWS: usize = 8_192; - -/// Far below the distinct sets (200k values, several megabytes across the -/// partial and final tables), far above the fixed cost of the stages. With the -/// accumulator releasing its buffers the query passes from 2 MB upwards; -/// without, it fails up to 4 MB with "Decreasing allocation after spilling -/// should succeed" in the final stage or a failed emit in the partial stage. -const COUNT_DISTINCT_MEMORY_LIMIT: usize = 4 * 1024 * 1024; - -/// `g` has 64 groups, `v` is unique, so every group holds 3125 distinct values. -fn count_distinct_table() -> MemTable { - let schema = Arc::new(Schema::new(vec![ - Field::new("g", DataType::Int64, false), - Field::new("v", DataType::Int64, false), - ])); - let batches = (0..COUNT_DISTINCT_ROWS) - .step_by(COUNT_DISTINCT_BATCH_ROWS) - .map(|start| { - let rows = - start..(start + COUNT_DISTINCT_BATCH_ROWS).min(COUNT_DISTINCT_ROWS); - RecordBatch::try_new( - Arc::clone(&schema), - vec![ - Arc::new(Int64Array::from_iter_values( - rows.clone().map(|row| row as i64 % COUNT_DISTINCT_GROUPS), - )), - Arc::new(Int64Array::from_iter_values(rows.map(|row| row as i64))), - ], - ) - .unwrap() - }) - .collect(); - MemTable::try_new(schema, vec![batches]).unwrap() -} - -/// Four partial stages emit their state early and four hash-partitioned final -/// stages spill; every one of them must see the accumulator memory drop after -/// emitting all groups. -#[tokio::test] -async fn count_distinct_releases_memory_after_emitting_all() { - let runtime = RuntimeEnvBuilder::new() - .with_memory_limit(COUNT_DISTINCT_MEMORY_LIMIT, 1.0) - .with_disk_manager_builder(DiskManagerBuilder::default()) - .build_arc() - .unwrap(); - let config = SessionConfig::new().with_target_partitions(4); - let ctx = SessionContext::new_with_config_rt(config, runtime); - ctx.register_table("t", Arc::new(count_distinct_table())) - .unwrap(); +/// With `force_hash_collisions` every key hashes alike, so the hash +/// repartitioning sends all groups to a single final stage, whose table then +/// does not fit the limit however well memory is released. The limit is sized +/// for the real distribution across four final stages, so this test is skipped +/// under that feature. +#[cfg(not(feature = "force_hash_collisions"))] +mod count_distinct_spill { + use super::*; + use arrow::array::Int64Array; + use datafusion::assert_batches_sorted_eq; + + /// `count(distinct)` over integers under a memory limit. + /// + /// The integer distinct-count groups accumulator reports the capacity of its + /// buffers in `size()`. After an aggregate stream emits all groups, either to + /// emit partial state early or to spill, it resizes its reservation to the + /// table's reported size and expects it to have shrunk. If the accumulator + /// keeps its capacity, that resize is a grow against an exhausted pool and the + /// query fails although everything was already written out. + const COUNT_DISTINCT_ROWS: usize = 200_000; + const COUNT_DISTINCT_GROUPS: i64 = 64; + const COUNT_DISTINCT_BATCH_ROWS: usize = 8_192; + + /// Far below the distinct sets (200k values, several megabytes across the + /// partial and final tables), far above the fixed cost of the stages. With the + /// accumulator releasing its buffers the query passes from 2 MB upwards; + /// without, it fails up to 4 MB with "Decreasing allocation after spilling + /// should succeed" in the final stage or a failed emit in the partial stage. + const COUNT_DISTINCT_MEMORY_LIMIT: usize = 4 * 1024 * 1024; + + /// `g` has 64 groups, `v` is unique, so every group holds 3125 distinct values. + fn count_distinct_table() -> MemTable { + let schema = Arc::new(Schema::new(vec![ + Field::new("g", DataType::Int64, false), + Field::new("v", DataType::Int64, false), + ])); + let batches = (0..COUNT_DISTINCT_ROWS) + .step_by(COUNT_DISTINCT_BATCH_ROWS) + .map(|start| { + let rows = + start..(start + COUNT_DISTINCT_BATCH_ROWS).min(COUNT_DISTINCT_ROWS); + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int64Array::from_iter_values( + rows.clone().map(|row| row as i64 % COUNT_DISTINCT_GROUPS), + )), + Arc::new(Int64Array::from_iter_values( + rows.map(|row| row as i64), + )), + ], + ) + .unwrap() + }) + .collect(); + MemTable::try_new(schema, vec![batches]).unwrap() + } - let batches = ctx - .sql("select count(distinct v) as d, count(*) as n from t group by g") - .await - .unwrap() - .collect() - .await - .unwrap_or_else(|error| panic!("query failed under the memory limit: {error}")); - - let per_group = (COUNT_DISTINCT_ROWS as i64 / COUNT_DISTINCT_GROUPS).to_string(); - let row = format!("| {per_group} | {per_group} |"); - let mut expected = vec!["+------+------+", "| d | n |", "+------+------+"]; - expected.extend(std::iter::repeat_n( - row.as_str(), - COUNT_DISTINCT_GROUPS as usize, - )); - expected.push("+------+------+"); - assert_batches_sorted_eq!(expected, &batches); + /// Four partial stages emit their state early and four hash-partitioned final + /// stages spill; every one of them must see the accumulator memory drop after + /// emitting all groups. + #[tokio::test] + async fn count_distinct_releases_memory_after_emitting_all() { + let runtime = RuntimeEnvBuilder::new() + .with_memory_limit(COUNT_DISTINCT_MEMORY_LIMIT, 1.0) + .with_disk_manager_builder(DiskManagerBuilder::default()) + .build_arc() + .unwrap(); + let config = SessionConfig::new().with_target_partitions(4); + let ctx = SessionContext::new_with_config_rt(config, runtime); + ctx.register_table("t", Arc::new(count_distinct_table())) + .unwrap(); + + let batches = ctx + .sql("select count(distinct v) as d, count(*) as n from t group by g") + .await + .unwrap() + .collect() + .await + .unwrap_or_else(|error| { + panic!("query failed under the memory limit: {error}") + }); + + let per_group = (COUNT_DISTINCT_ROWS as i64 / COUNT_DISTINCT_GROUPS).to_string(); + let row = format!("| {per_group} | {per_group} |"); + let mut expected = vec!["+------+------+", "| d | n |", "+------+------+"]; + expected.extend(std::iter::repeat_n( + row.as_str(), + COUNT_DISTINCT_GROUPS as usize, + )); + expected.push("+------+------+"); + assert_batches_sorted_eq!(expected, &batches); + } } #[tokio::test] From f0e0ac4a36093ef79dc406334a4c29e528f0444e Mon Sep 17 00:00:00 2001 From: Goutam Adwant <8672451+goutamadwant@users.noreply.github.com> Date: Thu, 3 Sep 2026 12:10:37 +0000 Subject: [PATCH 29/37] fix: exclude child compute from NestedLoopJoinExec metrics (#24491) ## Which issue does this PR close? - Part of #24459. ## Rationale for this change `NestedLoopJoinExec` included time spent polling its build-side and probe-side inputs in `build_time` and `join_time`. Because child operators report their own compute time, this double-counted child work in plan-level `elapsed_compute` metrics. ## What changes are included in this PR? - Time build-side bookkeeping and materialization only after each left input batch is ready. - Start probe-side timing only after the right input returns a ready result. - Apply the same accounting to the memory-limited spill and replay paths. - Add regressions for standard and spill execution that verify child polling is excluded while join-owned work remains timed. ## Are these changes tested? Yes. - `cargo test -p datafusion-physical-plan joins::nested_loop_join::tests --all-features` - `cargo test -p datafusion-physical-plan --all-features` - `cargo clippy --all-targets --all-features -- -D warnings` - `RUST_BACKTRACE=1 cargo test --profile ci --exclude datafusion-examples --exclude datafusion-benchmarks --exclude datafusion-cli --workspace --lib --tests --bins --features avro,json,backtrace,extended_tests,recursive_protection,parquet_encryption` The four new timing regressions were also ablated by restoring the previous outer timer scopes; each failed because the injected child delay was included. ## Are there any user-facing changes? `NestedLoopJoinExec` metrics now exclude child input polling. Query results and public APIs are unchanged. --- .../src/joins/nested_loop_join.rs | 433 ++++++++++++++++-- 1 file changed, 385 insertions(+), 48 deletions(-) diff --git a/datafusion/physical-plan/src/joins/nested_loop_join.rs b/datafusion/physical-plan/src/joins/nested_loop_join.rs index f63418f63e5f0..8cd603e91eea5 100644 --- a/datafusion/physical-plan/src/joins/nested_loop_join.rs +++ b/datafusion/physical-plan/src/joins/nested_loop_join.rs @@ -1078,6 +1078,7 @@ async fn collect_left_input( while let Some(batch) = stream.next().await { let batch = batch?; + let build_timer = metrics.build_time.timer(); let batch_size = batch.get_array_memory_size(); match reservation.try_grow(batch_size) { Ok(()) => { @@ -1087,6 +1088,9 @@ async fn collect_left_input( batches.push(batch); } Err(e) if is_spillable_oom(&e, spill_manager.as_ref()) => { + // Do not keep the operator timer running while the spill path + // drains the child stream. + build_timer.done(); let spill_manager = spill_manager.expect("checked by is_spillable_oom"); let spilled = spill_left_input( spill_manager, @@ -1109,6 +1113,10 @@ async fn collect_left_input( } } + // Only time the build-side materialization performed by this operator, not + // polling the child stream above. + let build_timer = metrics.build_time.timer(); + let merged_batch = concat_batches(&schema, &batches)?; // Reserve memory for visited_left_side bitmap if required by join type @@ -1118,6 +1126,9 @@ async fn collect_left_input( match reservation.try_grow(buffer_size) { Ok(()) => {} Err(e) if is_spillable_oom(&e, spill_manager.as_ref()) => { + // `spill_left_input` owns its timing and polls the input stream + // outside that timer. + build_timer.done(); let spill_manager = spill_manager.expect("checked by is_spillable_oom"); drop(batches); let spilled = spill_left_input( @@ -1198,6 +1209,7 @@ async fn spill_left_input( metrics: BuildProbeJoinMetrics, reservation: &MemoryReservation, ) -> Result> { + let build_timer = metrics.build_time.timer(); let mut spill_file = spill_manager.create_in_progress_file("NestedLoopJoin left spill")?; @@ -1217,9 +1229,11 @@ async fn spill_left_input( spill_file.append_batch(&batch)?; } } + build_timer.done(); while let Some(batch) = stream.next().await { let batch = batch?; + let _build_timer = metrics.build_time.timer(); if batch.num_rows() > 0 { metrics.build_input_batches.add(1); metrics.build_input_rows.add(batch.num_rows()); @@ -1227,6 +1241,7 @@ async fn spill_left_input( } } + let _build_timer = metrics.build_time.timer(); Ok(spill_file.finish()?.map(|file| LeftSpillData { spill_manager, spill_file: file, @@ -1524,12 +1539,6 @@ impl Stream for NestedLoopJoinStream { // side batch, before start joining. NLJState::BufferingLeft => { debug!("[NLJState] Entering: {:?}", self.state); - // inside `collect_left_input` (the routine to buffer build - // -side batches), related metrics except build time will be - // updated. - // stop on drop - let build_metric = self.metrics.join_metrics.build_time.clone(); - let _build_timer = build_metric.timer(); match self.handle_buffering_left(cx) { ControlFlow::Continue(()) => {} @@ -1562,9 +1571,6 @@ impl Stream for NestedLoopJoinStream { // handling (e.g., in cases like left join). NLJState::FetchingRight => { debug!("[NLJState] Entering: {:?}", self.state); - // stop on drop - let join_metric = self.metrics.join_metrics.join_time.clone(); - let _join_timer = join_metric.timer(); match self.handle_fetching_right(cx) { ControlFlow::Continue(()) => {} @@ -1682,9 +1688,6 @@ impl Stream for NestedLoopJoinStream { NLJState::EmitGlobalRightUnmatched => { debug!("[NLJState] Entering: {:?}", self.state); - let join_metric = self.metrics.join_metrics.join_time.clone(); - let _join_timer = join_metric.timer(); - match self.handle_emit_global_right_unmatched(cx) { ControlFlow::Continue(()) => {} ControlFlow::Break(poll) => { @@ -1879,6 +1882,7 @@ impl NestedLoopJoinStream { &mut self, cx: &mut std::task::Context<'_>, ) -> ControlFlow>>> { + let build_metric = self.metrics.join_metrics.build_time.clone(); let SpillState::Active(active) = &mut self.spill_state else { unreachable!( "handle_buffering_left_memory_limited called without Active spill state" @@ -1889,6 +1893,7 @@ impl NestedLoopJoinStream { // stream was consumed, open a fresh stream over the left spill file. if active.left_stream.is_none() { let spill_data = Arc::clone(&active.left_spill); + let _build_timer = build_metric.timer(); match spill_data .spill_manager .read_spill_as_stream(Arc::clone(&spill_data.spill_file), None) @@ -1914,6 +1919,7 @@ impl NestedLoopJoinStream { loop { match left_stream.poll_next_unpin(cx) { Poll::Ready(Some(Ok(batch))) => { + let _build_timer = build_metric.timer(); if batch.num_rows() == 0 { continue; } @@ -1954,6 +1960,8 @@ impl NestedLoopJoinStream { } } + let _build_timer = build_metric.timer(); + // If the left stream is fully exhausted, release its resources so the // upstream pipeline can be torn down before we move on to probing. if self.left_exhausted { @@ -2046,47 +2054,52 @@ impl NestedLoopJoinStream { &mut self, cx: &mut std::task::Context<'_>, ) -> ControlFlow>>> { - match self + let result = match self .right_data .as_mut() .expect("right_data must be present while fetching right") .poll_next_unpin(cx) { - Poll::Ready(result) => match result { - Some(Ok(right_batch)) => { - // Update metrics - let right_batch_rows = right_batch.num_rows(); - self.metrics.join_metrics.input_rows.add(right_batch_rows); - self.metrics.join_metrics.input_batches.add(1); - - // Skip the empty batch - if right_batch_rows == 0 { - return ControlFlow::Continue(()); - } + Poll::Ready(result) => result, + Poll::Pending => return ControlFlow::Break(Poll::Pending), + }; - self.current_right_batch = Some(right_batch); + let join_metric = self.metrics.join_metrics.join_time.clone(); + let _join_timer = join_metric.timer(); - // Prepare right bitmap - if self.should_track_unmatched_right { - let zeroed_buf = BooleanBuffer::new_unset(right_batch_rows); - self.current_right_batch_matched = - Some(BooleanArray::new(zeroed_buf, None)); - } + match result { + Some(Ok(right_batch)) => { + // Update metrics + let right_batch_rows = right_batch.num_rows(); + self.metrics.join_metrics.input_rows.add(right_batch_rows); + self.metrics.join_metrics.input_batches.add(1); - self.left_probe_idx = 0; - self.state = NLJState::ProbeRight; - ControlFlow::Continue(()) + // Skip the empty batch + if right_batch_rows == 0 { + return ControlFlow::Continue(()); } - Some(Err(e)) => ControlFlow::Break(Poll::Ready(Some(Err(e)))), - None => { - // Right side exhausted: probing for the current left chunk - // is finished. `ProbeEnd` reports probe completion before - // emitting unmatched-left rows. - self.state = NLJState::ProbeEnd; - ControlFlow::Continue(()) + + self.current_right_batch = Some(right_batch); + + // Prepare right bitmap + if self.should_track_unmatched_right { + let zeroed_buf = BooleanBuffer::new_unset(right_batch_rows); + self.current_right_batch_matched = + Some(BooleanArray::new(zeroed_buf, None)); } - }, - Poll::Pending => ControlFlow::Break(Poll::Pending), + + self.left_probe_idx = 0; + self.state = NLJState::ProbeRight; + ControlFlow::Continue(()) + } + Some(Err(e)) => ControlFlow::Break(Poll::Ready(Some(Err(e)))), + None => { + // Right side exhausted: probing for the current left chunk + // is finished. `ProbeEnd` reports probe completion before + // emitting unmatched-left rows. + self.state = NLJState::ProbeEnd; + ControlFlow::Continue(()) + } } } @@ -2303,6 +2316,8 @@ impl NestedLoopJoinStream { // On first entry, open a new replay pass on the right input if self.right_data.is_none() { + let join_metric = self.metrics.join_metrics.join_time.clone(); + let _join_timer = join_metric.timer(); let SpillState::Active(ref mut active) = self.spill_state else { unreachable!("EmitGlobalRightUnmatched without Active spill state"); }; @@ -2318,13 +2333,20 @@ impl NestedLoopJoinStream { } // Poll the replay stream for the next right batch - match self + let result = match self .right_data .as_mut() .expect("right_data must be present") .poll_next_unpin(cx) { - Poll::Ready(Some(Ok(right_batch))) => { + Poll::Ready(result) => result, + Poll::Pending => return ControlFlow::Break(Poll::Pending), + }; + + let join_metric = self.metrics.join_metrics.join_time.clone(); + let _join_timer = join_metric.timer(); + match result { + Some(Ok(right_batch)) => { if right_batch.num_rows() == 0 { return ControlFlow::Continue(()); } @@ -2370,8 +2392,8 @@ impl NestedLoopJoinStream { Err(e) => ControlFlow::Break(Poll::Ready(Some(Err(e)))), } } - Poll::Ready(Some(Err(e))) => ControlFlow::Break(Poll::Ready(Some(Err(e)))), - Poll::Ready(None) => { + Some(Err(e)) => ControlFlow::Break(Poll::Ready(Some(Err(e)))), + None => { // All right batches replayed match self.output_buffer.finish_buffered_batch() { Ok(()) => { @@ -2381,7 +2403,6 @@ impl NestedLoopJoinStream { Err(e) => ControlFlow::Break(Poll::Ready(Some(arrow_err!(e)))), } } - Poll::Pending => ControlFlow::Break(Poll::Pending), } } @@ -3378,6 +3399,9 @@ fn build_unmatched_batch( #[cfg(test)] pub(crate) mod tests { + use std::pin::Pin; + use std::time::Duration; + use super::*; use crate::statistics::{StatisticsArgs, StatisticsContext}; use crate::test::{TestMemoryExec, assert_join_metrics}; @@ -3387,9 +3411,15 @@ pub(crate) mod tests { use arrow::compute::SortOptions; use arrow::datatypes::{DataType, Field}; + use bytes::Bytes; use datafusion_common::assert_contains; + use datafusion_common::instant::Instant; use datafusion_common::test_util::batches_to_sort_string; + use datafusion_execution::disk_manager::{ + DiskManager, DiskManagerBuilder, DiskManagerMode, + }; use datafusion_execution::runtime_env::RuntimeEnvBuilder; + use datafusion_execution::spill_file::{SpillFile, SpillWriter, TempFileFactory}; use datafusion_expr::Operator; use datafusion_physical_expr::expressions::{BinaryExpr, Literal}; use datafusion_physical_expr::{Partitioning, PhysicalExpr}; @@ -3399,6 +3429,73 @@ pub(crate) mod tests { use insta::assert_snapshot; use rstest::rstest; + fn delayed_stream(batch: RecordBatch, delay: Duration) -> SendableRecordBatchStream { + let schema = batch.schema(); + Box::pin(crate::stream::RecordBatchStreamAdapter::new( + schema, + futures::stream::once(async move { + std::thread::sleep(delay); + Ok(batch) + }), + )) + } + + /// Delays the first item while the spill stream is polled, making an + /// incorrectly scoped operator timer include the delay. + struct DelayedReadSpillFile { + inner: Arc, + delay: Duration, + read_count: Arc, + } + + impl SpillFile for DelayedReadSpillFile { + fn path(&self) -> Option<&std::path::Path> { + self.inner.path() + } + + fn size(&self) -> Option { + self.inner.size() + } + + fn read_stream( + &self, + ) -> Result> + Send>>> { + let delay = self.delay; + let read_count = Arc::clone(&self.read_count); + let mut delay_first_item = true; + let stream = self.inner.read_stream()?.map(move |item| { + if delay_first_item { + delay_first_item = false; + read_count.fetch_add(1, Ordering::Relaxed); + std::thread::sleep(delay); + } + item + }); + Ok(Box::pin(stream)) + } + + fn open_writer(&self) -> Result> { + self.inner.open_writer() + } + } + + /// Wraps local spill files so replay reads can be delayed deterministically. + struct DelayedReadTempFileFactory { + inner: Arc, + delay: Duration, + read_count: Arc, + } + + impl TempFileFactory for DelayedReadTempFileFactory { + fn create_temp_file(&self, description: &str) -> Result> { + Ok(Arc::new(DelayedReadSpillFile { + inner: self.inner.create_tmp_file(description)?, + delay: self.delay, + read_count: Arc::clone(&self.read_count), + })) + } + } + fn build_table( a: (&str, &Vec), b: (&str, &Vec), @@ -3554,6 +3651,246 @@ pub(crate) mod tests { ) } + async fn run_join_with_child_poll_delays( + left_delay: Duration, + right_delay: Duration, + memory_limited: bool, + ) -> Result<(Duration, Duration, Duration)> { + run_join_with_poll_delays( + left_delay, + right_delay, + memory_limited, + JoinType::Inner, + None, + ) + .await + } + + async fn run_join_with_poll_delays( + left_delay: Duration, + right_delay: Duration, + memory_limited: bool, + join_type: JoinType, + spill_read_delay: Option<(Duration, Arc)>, + ) -> Result<(Duration, Duration, Duration)> { + let left_batch = + build_table_i32(("a1", &vec![1]), ("b1", &vec![2]), ("c1", &vec![3])); + let right_batch = + build_table_i32(("a2", &vec![4]), ("b2", &vec![5]), ("c2", &vec![6])); + let left_schema = left_batch.schema(); + let right_schema = right_batch.schema(); + + let left_stream = delayed_stream(left_batch.clone(), left_delay); + let right_stream = delayed_stream(right_batch, right_delay); + + let task_ctx = if let Some((delay, read_count)) = spill_read_delay { + let inner = Arc::new( + DiskManagerBuilder::default() + .with_mode(DiskManagerMode::OsTmpDirectory) + .build()?, + ); + let runtime = RuntimeEnvBuilder::new() + .with_disk_manager_builder(DiskManagerBuilder::default().with_mode( + DiskManagerMode::Custom(Arc::new(DelayedReadTempFileFactory { + inner, + delay, + read_count, + })), + )) + .build_arc()?; + Arc::new(TaskContext::default().with_runtime(runtime)) + } else { + Arc::new(TaskContext::default()) + }; + let metrics_set = ExecutionPlanMetricsSet::new(); + let metrics = NestedLoopJoinMetrics::new(&metrics_set, 0); + let build_time = metrics.join_metrics.build_time.clone(); + let join_time = metrics.join_metrics.join_time.clone(); + let (left_data, right_stream, spill_state) = if memory_limited { + let reservation = MemoryConsumer::new("NestedLoopJoinLoad[test]".to_string()) + .with_can_spill(true) + .register(task_ctx.memory_pool()); + let global_right_bitmaps_reservation = + MemoryConsumer::new("NestedLoopJoinGlobalRightBitmaps[test]".to_string()) + .register(task_ctx.memory_pool()); + let spill_manager = SpillManager::new( + task_ctx.runtime_env(), + metrics.spill_metrics.clone(), + Arc::clone(&right_schema), + ); + let left_spill_manager = SpillManager::new( + task_ctx.runtime_env(), + metrics.spill_metrics.clone(), + Arc::clone(&left_schema), + ); + let mut left_spill_file = + left_spill_manager.create_in_progress_file("test left spill")?; + left_spill_file.append_batch(&left_batch)?; + let left_spill_file = left_spill_file + .finish()? + .expect("the test left spill contains one batch"); + let active = SpillStateActive { + left_spill: Arc::new(LeftSpillData { + spill_manager: left_spill_manager, + spill_file: left_spill_file, + schema: Arc::clone(&left_schema), + }), + left_stream: Some(left_stream), + left_schema: Some(Arc::clone(&left_schema)), + reservation, + pending_batches: Vec::new(), + right_input: ReplayableStreamSource::new( + right_stream, + spill_manager, + "test right spill", + ), + global_right_bitmaps: Vec::new(), + global_right_bitmaps_reservation, + right_batch_index: 0, + }; + ( + OnceFut::new(async { internal_err!("unused left data was polled") }), + Box::pin(crate::EmptyRecordBatchStream::new(Arc::clone( + &right_schema, + ))) as SendableRecordBatchStream, + SpillState::Active(Box::new(active)), + ) + } else { + let reservation = MemoryConsumer::new("NestedLoopJoinLoad[test]".to_string()) + .register(task_ctx.memory_pool()); + ( + OnceFut::new(collect_left_input( + left_stream, + metrics.join_metrics.clone(), + reservation, + false, + 1, + None, + )), + right_stream, + SpillState::Disabled, + ) + }; + let (output_schema, column_indices) = + build_join_schema(&left_schema, &right_schema, &join_type); + let stream = NestedLoopJoinStream::new( + Arc::new(output_schema), + None, + join_type, + right_stream, + left_data, + column_indices, + metrics, + 1024, + spill_state, + ); + + let start = Instant::now(); + let batches = common::collect(Box::pin(stream)).await?; + let wall_time = start.elapsed(); + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 1); + + Ok(( + Duration::from_nanos(build_time.value() as u64), + Duration::from_nanos(join_time.value() as u64), + wall_time, + )) + } + + async fn check_child_poll_time_excluded(mut run: F) -> Result<()> + where + F: FnMut(Duration) -> Fut, + Fut: Future>, + { + // Escalating the delay filters out fixed-size scheduler preemption without + // masking the bug: an incorrectly scoped timer grows with every delay. + let mut delay = Duration::from_millis(50); + for attempt in 0..3 { + let (operator_time, wall_time) = run(delay).await?; + assert!( + !operator_time.is_zero(), + "operator work should still be timed" + ); + assert!( + wall_time >= delay, + "child poll delay should dominate wall time: {wall_time:?} < {delay:?}" + ); + if operator_time < delay { + return Ok(()); + } + assert!( + attempt < 2, + "operator time ({operator_time:?}) included the child poll delay ({delay:?})" + ); + delay *= 4; + } + unreachable!() + } + + #[tokio::test] + async fn build_time_excludes_left_child_poll() -> Result<()> { + check_child_poll_time_excluded(|delay| async move { + let (build_time, _, wall_time) = + run_join_with_child_poll_delays(delay, Duration::ZERO, false).await?; + Ok((build_time, wall_time)) + }) + .await + } + + #[tokio::test] + async fn join_time_excludes_right_child_poll() -> Result<()> { + check_child_poll_time_excluded(|delay| async move { + let (_, join_time, wall_time) = + run_join_with_child_poll_delays(Duration::ZERO, delay, false).await?; + Ok((join_time, wall_time)) + }) + .await + } + + #[tokio::test] + async fn build_time_excludes_spill_stream_poll() -> Result<()> { + check_child_poll_time_excluded(|delay| async move { + let (build_time, _, wall_time) = + run_join_with_child_poll_delays(delay, Duration::ZERO, true).await?; + Ok((build_time, wall_time)) + }) + .await + } + + #[tokio::test] + async fn join_time_excludes_replayable_input_poll() -> Result<()> { + check_child_poll_time_excluded(|delay| async move { + let (_, join_time, wall_time) = + run_join_with_child_poll_delays(Duration::ZERO, delay, true).await?; + Ok((join_time, wall_time)) + }) + .await + } + + #[tokio::test(flavor = "current_thread")] + async fn join_time_excludes_global_right_unmatched_replay_poll() -> Result<()> { + check_child_poll_time_excluded(|delay| async move { + let read_count = Arc::new(AtomicUsize::new(0)); + // The only left chunk is supplied directly, so the sole spill read + // is the final right replay in EmitGlobalRightUnmatched. + let (_, join_time, wall_time) = run_join_with_poll_delays( + Duration::ZERO, + Duration::ZERO, + true, + JoinType::Right, + Some((delay, Arc::clone(&read_count))), + ) + .await?; + assert_eq!( + read_count.load(Ordering::Relaxed), + 1, + "EmitGlobalRightUnmatched should replay the right spill exactly once" + ); + Ok((join_time, wall_time)) + }) + .await + } + fn prepare_join_filter() -> JoinFilter { let column_indices = vec![ ColumnIndex { From 20d1c5676118e0bca955ccf92c0d2cf665384d0d Mon Sep 17 00:00:00 2001 From: Bharadwaj Pendyala Date: Thu, 3 Sep 2026 12:15:10 +0000 Subject: [PATCH 30/37] fix: literal-on-left equality no longer collapses the filter to false (#24763) ## Which issue does this PR close? No separate issue. I found this while reading `simplify_predicates`. ## Rationale for this change `WHERE s = 'a' AND 'a' = s` returns no rows, where one row is expected: ```sql > CREATE TABLE t(s VARCHAR) AS VALUES ('a'), ('b'); > SELECT * FROM t WHERE s = 'a' AND 'a' = s; 0 row(s) fetched. ``` Either half on its own returns `a`, and the same query against an INT column returns the row. On `main` (4d3e79e), `EXPLAIN VERBOSE` shows the filter turning into a constant between two rules: ``` logical_plan after simplify_expressions Filter: t.s = Utf8View("a") AND Utf8View("a") = t.s logical_plan after push_down_filter Filter: Boolean(false) logical_plan after eliminate_filter EmptyRelation: rows=0 ``` `PushDownFilter` splits the conjuncts and calls `simplify_predicates`. It accepts both ` ` and ` `, but `simplify_column_predicates` compares whole `Expr`s. `t.s = Utf8View("a")` and `Utf8View("a") = t.s` aren't structurally equal, so the two equalities read as a contradiction and the conjunction becomes `false`. The INT version survives because the `Canonicalizer` reorders it first. It can't do that here: it runs once at `expr_simplifier.rs:203`, ahead of the const-evaluation loop, so it sees `CAST(Utf8("a") AS Utf8View)` rather than a `Literal` and its `(Literal, Column)` arm doesn't match. The cast folds to a literal afterwards. Canonicalization is skipped entirely for `Join` (`simplify_exprs.rs:130`), so `simplify_predicates` can't assume canonical input either way. The same gap costs a strict bound. Given `a >= 5` and `5 < a`, `find_most_restrictive_predicate` breaks the tie on `op == Gt`, doesn't count `Lt` with the literal on the left as strict, keeps `a >= 5`, and lets `a = 5` through. ## What changes are included in this PR? `simplify_predicates` now normalizes the literal to the right with `op.swap()`, at the point where it already distinguishes the two orientations. `simplify_column_predicates` can then match on the operator alone. No signature changes. ## Are these changes tested? Two unit tests in `simplify_predicates.rs` and four cases in `simplify_predicates.slt`. All six fail before the fix. With only `simplify_predicates.rs` reverted the SLT reports `EmptyRelation: rows=0` where `Filter: test_data.str_col = Utf8View("apple")` is expected, and the `apple` row goes missing. `datafusion-optimizer` is green (765 lib, 26 integration, 5 doc) and clippy with `-D warnings` is clean. The full `sqllogictests` run passes except `window_limits.slt`, which fails identically on an unmodified `main`. `SELECT * FROM t WHERE s = 'a' AND 'b' = s` stays `EmptyRelation: rows=0` before and after, and that's pinned in the SLT. ## Are there any user-facing changes? Affected queries return the right rows instead of none. Predicates reaching `simplify_predicates` with the literal on the left now come back with it on the right, so a plan can show `a > 5` where it used to show `5 < a`. Nothing in the test suite depended on that, but the function is public. Equalities whose literals are equal in value but differ in `ScalarValue` representation still collapse to `false`. On `main`, `[a = 5i32, a = 5i64]` in the same orientation already returns `Boolean(false)`, so that predates this change and isn't orientation related. This PR was written with AI assistance. --- .../simplify_predicates.rs | 92 ++++++++++++------- .../test_files/simplify_predicates.slt | 26 ++++++ 2 files changed, 87 insertions(+), 31 deletions(-) diff --git a/datafusion/optimizer/src/simplify_expressions/simplify_predicates.rs b/datafusion/optimizer/src/simplify_expressions/simplify_predicates.rs index 356f2711b708e..e7edc34cfe4e6 100644 --- a/datafusion/optimizer/src/simplify_expressions/simplify_predicates.rs +++ b/datafusion/optimizer/src/simplify_expressions/simplify_predicates.rs @@ -52,27 +52,45 @@ pub fn simplify_predicates(predicates: Vec) -> Result> { let mut other_predicates = Vec::new(); for pred in predicates { - match &pred { - Expr::BinaryExpr(BinaryExpr { - left, - op: + match pred { + Expr::BinaryExpr(BinaryExpr { left, op, right }) + if matches!( + op, Operator::Gt - | Operator::GtEq - | Operator::Lt - | Operator::LtEq - | Operator::Eq, - right, - }) => { + | Operator::GtEq + | Operator::Lt + | Operator::LtEq + | Operator::Eq + ) => + { if let (Some(col), Some(_)) = - (extract_column_from_expr(left), right.as_literal()) - { - column_predicates.entry(col).or_default().push(pred); - } else if let (Some(_), Some(col)) = - (left.as_literal(), extract_column_from_expr(right)) + (extract_column_from_expr(&left), right.as_literal()) { - column_predicates.entry(col).or_default().push(pred); + column_predicates + .entry(col) + .or_default() + .push(Expr::BinaryExpr(BinaryExpr { left, op, right })); + } else if let (Some(_), Some(col), Some(swapped_op)) = ( + left.as_literal(), + extract_column_from_expr(&right), + op.swap(), + ) { + // Put the literal on the right so that predicates differing only in + // operand order compare equal below + column_predicates + .entry(col) + .or_default() + .push(Expr::BinaryExpr(BinaryExpr { + left: right, + op: swapped_op, + right: left, + })); } else { - other_predicates.push(pred); + other_predicates.push(Expr::BinaryExpr(BinaryExpr { + left, + op, + right, + })); } } _ => other_predicates.push(pred), @@ -114,20 +132,12 @@ fn simplify_column_predicates(predicates: Vec) -> Result> { for pred in predicates { match &pred { - Expr::BinaryExpr(BinaryExpr { left: _, op, right }) => { - match (op, right.as_literal().is_some()) { - (Operator::Gt, true) - | (Operator::Lt, false) - | (Operator::GtEq, true) - | (Operator::LtEq, false) => greater_predicates.push(pred), - (Operator::Lt, true) - | (Operator::Gt, false) - | (Operator::LtEq, true) - | (Operator::GtEq, false) => less_predicates.push(pred), - (Operator::Eq, _) => eq_predicates.push(pred), - _ => unreachable!("Unexpected operator: {}", op), - } - } + Expr::BinaryExpr(BinaryExpr { op, .. }) => match op { + Operator::Gt | Operator::GtEq => greater_predicates.push(pred), + Operator::Lt | Operator::LtEq => less_predicates.push(pred), + Operator::Eq => eq_predicates.push(pred), + _ => unreachable!("Unexpected operator: {}", op), + }, _ => unreachable!("Unexpected predicate {}", pred.to_string()), } } @@ -305,6 +315,26 @@ mod tests { assert_eq!(extract_column_from_expr(&col_expr), Some(Column::from("a"))); } + #[test] + fn test_eq_predicates_are_matched_regardless_of_operand_order() { + // a = 5 AND 5 = a is a single condition, not a contradiction + let predicates = vec![col("a").eq(lit(5i32)), lit(5i32).eq(col("a"))]; + + let result = simplify_predicates(predicates).unwrap(); + + assert_eq!(result, vec![col("a").eq(lit(5i32))]); + } + + #[test] + fn test_strict_bound_wins_when_literal_is_on_the_left() { + // a >= 5 AND 5 < a is a > 5; keeping a >= 5 would let a = 5 through + let predicates = vec![col("a").gt_eq(lit(5i32)), lit(5i32).lt(col("a"))]; + + let result = simplify_predicates(predicates).unwrap(); + + assert_eq!(result, vec![col("a").gt(lit(5i32))]); + } + #[test] fn test_simplify_predicates_direct_columns_only() { // Test that only predicates on direct columns are simplified together diff --git a/datafusion/sqllogictest/test_files/simplify_predicates.slt b/datafusion/sqllogictest/test_files/simplify_predicates.slt index 44fdedc9c8e1d..c36b0c864d592 100644 --- a/datafusion/sqllogictest/test_files/simplify_predicates.slt +++ b/datafusion/sqllogictest/test_files/simplify_predicates.slt @@ -242,5 +242,31 @@ logical_plan 02)--TableScan: test_data projection=[int_col, float_col, str_col, date_col, bool_col] +# x = 'apple' AND 'apple' = x is one condition, not a contradiction +query TT +EXPLAIN SELECT * FROM test_data WHERE str_col = 'apple' AND 'apple' = str_col; +---- +logical_plan +01)Filter: test_data.str_col = Utf8View("apple") +02)--TableScan: test_data projection=[int_col, float_col, str_col, date_col, bool_col] + +# x = 'apple' AND 'banana' = x is still impossible +query TT +EXPLAIN SELECT * FROM test_data WHERE str_col = 'apple' AND 'banana' = str_col; +---- +logical_plan EmptyRelation: rows=0 + statement ok set datafusion.explain.logical_plan_only=false; + +statement ok +CREATE TABLE fruit (name VARCHAR) AS VALUES ('apple'), ('banana'); + +query T +SELECT name FROM fruit WHERE name = 'apple' AND 'apple' = name; +---- +apple + +query T +SELECT name FROM fruit WHERE name = 'apple' AND 'banana' = name; +---- From ca7a86a42bfe3f9922a591694ae5736146ecc9ab Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Tue, 1 Sep 2026 01:14:49 -0500 Subject: [PATCH 31/37] fix: report the real hashbrown allocation from ArrowBytesMap and ArrowBytesViewMap Both maps tracked their hash table footprint in a `map_size` field that was only ever incremented by `HashTableAllocExt::insert_accounted`, which charges `capacity * size_of::()` on growth and nothing else. That undercounts in two ways. `ArrowBytesViewMap::new` seeded `map_size` with `capacity() * size_of::>()`, which ignores the control bytes and the trailing group that hashbrown allocates alongside the entry array, so the reported size was roughly half the real allocation. `ArrowBytesMap::new` seeded `map_size` with 0 despite pre-allocating a table for 128 entries. Since `insert_accounted` only charges when the table grows, any map holding fewer entries than the pre-allocated capacity reported its hash table as free forever. Drop the field and ask hashbrown for the exact figure with `HashTable::allocation_size`, which covers entries, control bytes and the trailing group. It is a constant time layout calculation, so `size()` stays cheap, and it cannot drift out of sync with the table the way an incrementally maintained counter can. --- .../physical-expr-common/src/binary_map.rs | 26 +++++-------- .../src/binary_view_map.rs | 38 +++++++++++-------- 2 files changed, 31 insertions(+), 33 deletions(-) diff --git a/datafusion/physical-expr-common/src/binary_map.rs b/datafusion/physical-expr-common/src/binary_map.rs index 7543e6b297329..ceae72d4ccafe 100644 --- a/datafusion/physical-expr-common/src/binary_map.rs +++ b/datafusion/physical-expr-common/src/binary_map.rs @@ -28,7 +28,7 @@ use arrow::buffer::{Buffer, NullBuffer, OffsetBuffer, ScalarBuffer}; use arrow::datatypes::DataType; use datafusion_common::hash_utils::RandomState; use datafusion_common::hash_utils::create_hashes; -use datafusion_common::utils::proxy::{HashTableAllocExt, VecAllocExt}; +use datafusion_common::utils::proxy::VecAllocExt; use datafusion_common::{Result, exec_err}; use std::any::type_name; use std::fmt::Debug; @@ -217,8 +217,6 @@ where output_type: OutputType, /// Underlying hash set for each distinct value map: hashbrown::hash_table::HashTable>, - /// Total size of the map in bytes - map_size: usize, /// In progress buffer containing all values buffer: Vec, /// Offsets into `buffer` for each distinct value. These offsets as used @@ -248,7 +246,6 @@ where Self { output_type, map: hashbrown::hash_table::HashTable::with_capacity(INITIAL_MAP_CAPACITY), - map_size: 0, buffer: Vec::with_capacity(INITIAL_BUFFER_CAPACITY), offsets: vec![O::default()], // first offset is always 0 random_state: RandomState::default(), @@ -415,11 +412,7 @@ where offset_or_inline: inline, payload, }; - self.map.insert_accounted( - new_header, - |header| header.hash, - &mut self.map_size, - ); + self.map.insert_unique(hash, new_header, |header| header.hash); payload } } @@ -457,11 +450,7 @@ where offset_or_inline: offset, payload, }; - self.map.insert_accounted( - new_header, - |header| header.hash, - &mut self.map_size, - ); + self.map.insert_unique(hash, new_header, |header| header.hash); payload } }; @@ -486,7 +475,6 @@ where let Self { output_type, map: _, - map_size: _, offsets, buffer, random_state: _, @@ -591,7 +579,11 @@ where /// Return the total size, in bytes, of memory used to store the data in /// this set, not including `self` pub fn size(&self) -> usize { - self.map_size + // `HashTable::allocation_size` reports the whole hashbrown allocation, + // which is the entry array plus the control bytes plus the trailing + // group, so it is larger than `capacity() * size_of::>()`. + // It is a constant time layout calculation, not a walk of the table. + self.map.allocation_size() + self.buffer.capacity() * size_of::() + self.offsets.allocated_size() + self.hashes_buffer.allocated_size() @@ -615,7 +607,7 @@ where fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("ArrowBytesMap") .field("map", &"") - .field("map_size", &self.map_size) + .field("map_allocation_size", &self.map.allocation_size()) .field("buffer", &self.buffer) .field("random_state", &self.random_state) .field("hashes_buffer", &self.hashes_buffer) diff --git a/datafusion/physical-expr-common/src/binary_view_map.rs b/datafusion/physical-expr-common/src/binary_view_map.rs index 0457825decb96..304b321870467 100644 --- a/datafusion/physical-expr-common/src/binary_view_map.rs +++ b/datafusion/physical-expr-common/src/binary_view_map.rs @@ -28,10 +28,9 @@ use arrow::buffer::{Buffer, ScalarBuffer}; use arrow::datatypes::{BinaryViewType, ByteViewType, DataType, StringViewType}; use datafusion_common::hash_utils::RandomState; use datafusion_common::hash_utils::create_hashes; -use datafusion_common::utils::proxy::{HashTableAllocExt, VecAllocExt}; +use datafusion_common::utils::proxy::VecAllocExt; use datafusion_common::{Result, exec_err}; use std::fmt::Debug; -use std::mem::size_of; use std::sync::Arc; /// HashSet optimized for storing string or binary values that can produce that @@ -129,8 +128,6 @@ where output_type: OutputType, /// Underlying hash set for each distinct value map: hashbrown::hash_table::HashTable>, - /// Total size of the map in bytes - map_size: usize, /// Views for all stored values (in insertion order) views: Vec, @@ -159,13 +156,9 @@ where V: Debug + PartialEq + Eq + Clone + Copy + Default, { pub fn new(output_type: OutputType) -> Self { - let map = hashbrown::hash_table::HashTable::with_capacity(INITIAL_MAP_CAPACITY); - let map_size = map.capacity() * size_of::>(); - Self { output_type, - map, - map_size, + map: hashbrown::hash_table::HashTable::with_capacity(INITIAL_MAP_CAPACITY), views: Vec::new(), in_progress: Vec::new(), completed: Vec::new(), @@ -374,8 +367,7 @@ where payload, }; - self.map - .insert_accounted(new_header, |h| h.hash, &mut self.map_size); + self.map.insert_unique(hash, new_header, |h| h.hash); payload }; observe_payload_fn(payload); @@ -540,13 +532,18 @@ where pub fn size(&self) -> usize { // All fields below own their allocations. Count retained capacity rather // than used length because this value drives memory accounting. + // + // `HashTable::allocation_size` reports the whole hashbrown allocation, + // which is the entry array plus the control bytes plus the trailing + // group, so it is larger than `capacity() * size_of::>()`. It + // is a constant time layout calculation, not a walk of the table. let views_size = self.views.allocated_size(); let in_progress_size = self.in_progress.allocated_size(); let completed_size = self.completed.allocated_size() + self.completed.iter().map(Buffer::capacity).sum::(); let nulls_size = self.nulls.allocated_size(); - self.map_size + self.map.allocation_size() + views_size + in_progress_size + completed_size @@ -562,7 +559,7 @@ where fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("ArrowBytesMap") .field("map", &"") - .field("map_size", &self.map_size) + .field("map_allocation_size", &self.map.allocation_size()) .field("views_len", &self.views.len()) .field("completed_buffers", &self.completed.len()) .field("random_state", &self.random_state) @@ -597,6 +594,7 @@ where mod tests { use arrow::array::{GenericByteViewArray, StringViewArray}; use datafusion_common::HashMap; + use std::mem::size_of; use super::*; @@ -789,7 +787,15 @@ mod tests { fn test_size_counts_initial_hash_table_capacity() { let map = ArrowBytesViewMap::<()>::new(OutputType::Utf8View); - assert_eq!(map.size(), map.map.capacity() * size_of::>()); + assert_eq!(map.size(), map.map.allocation_size()); + // The reported size covers the control bytes as well as the entries, so + // it is strictly larger than the entry array on its own. + assert!( + map.size() > map.map.capacity() * size_of::>(), + "expected {} to exceed {}", + map.size(), + map.map.capacity() * size_of::>() + ); } #[test] @@ -822,7 +828,7 @@ mod tests { .any(|buffer| buffer.capacity() > buffer.len()) ); - let expected_size = map.map_size + let expected_size = map.map.allocation_size() + map.views.allocated_size() + map.in_progress.allocated_size() + map.completed.allocated_size() @@ -832,7 +838,7 @@ mod tests { assert_eq!(map.size(), expected_size); // Verify the retained-capacity delta independently from the production formula. - let legacy_size = map.map_size + let legacy_size = map.map.allocation_size() + map.views.len() * size_of::() + map.in_progress.capacity() + map.completed.iter().map(Buffer::len).sum::() From 8b2f71217fea8fc3b3425f93a9b70524f1e1bcae Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Tue, 1 Sep 2026 01:17:54 -0500 Subject: [PATCH 32/37] perf: stop pre-allocating a hash table per COUNT(DISTINCT) group `ArrowBytesMap` and `ArrowBytesViewMap` always pre-allocated their hash table, and `ArrowBytesMap` also pre-allocated an 8 KiB value buffer. That is the right trade for the single map that backs a `GROUP BY` on one string column, which goes on to hold every group value in the query. It is the wrong trade for `BytesDistinctCountAccumulator` and `BytesViewDistinctCountAccumulator`, because `GroupsAccumulatorAdapter` creates one accumulator per group: a grouped `COUNT(DISTINCT)` over a high cardinality key holds hundreds of thousands of them at once, and most see only a handful of values, so the pre-allocation dwarfs the data. Split the constructors. `new` no longer allocates anything, and `with_capacity` keeps the previous behavior for the callers that want it. The capacity is stored so `take` re-creates the map the way it was built. The `GroupValuesBytes` and `GroupValuesBytesView` call sites move to `with_capacity`; the two distinct-count accumulators stay on `new`. The `arrow_bytes_map` benchmark also moves to `with_capacity`: its `long_low_cardinality` case is defined by the distinct values fitting inside the pre-allocated buffer. --- .../src/aggregate/count_distinct/bytes.rs | 6 + .../benches/arrow_bytes_map.rs | 12 +- .../physical-expr-common/src/binary_map.rs | 134 ++++++++++++++++-- .../src/binary_view_map.rs | 103 ++++++++++++-- .../group_values/single_group_by/bytes.rs | 8 +- .../single_group_by/bytes_view.rs | 8 +- 6 files changed, 247 insertions(+), 24 deletions(-) diff --git a/datafusion/functions-aggregate-common/src/aggregate/count_distinct/bytes.rs b/datafusion/functions-aggregate-common/src/aggregate/count_distinct/bytes.rs index f6df4182a879b..3aa60f6f3b6ad 100644 --- a/datafusion/functions-aggregate-common/src/aggregate/count_distinct/bytes.rs +++ b/datafusion/functions-aggregate-common/src/aggregate/count_distinct/bytes.rs @@ -39,6 +39,10 @@ use std::mem::size_of_val; pub struct BytesDistinctCountAccumulator(ArrowBytesSet); impl BytesDistinctCountAccumulator { + /// The set deliberately does not pre-allocate. `GroupsAccumulatorAdapter` + /// creates one accumulator per group, so a grouped `COUNT(DISTINCT)` over a + /// high cardinality key holds hundreds of thousands of these at once and + /// most of them see only a handful of values. pub fn new(output_type: OutputType) -> Self { Self(ArrowBytesSet::new(output_type)) } @@ -100,6 +104,8 @@ impl Accumulator for BytesDistinctCountAccumulator { pub struct BytesViewDistinctCountAccumulator(ArrowBytesViewSet); impl BytesViewDistinctCountAccumulator { + /// See [`BytesDistinctCountAccumulator::new`] for why the set does not + /// pre-allocate. pub fn new(output_type: OutputType) -> Self { Self(ArrowBytesViewSet::new(output_type)) } diff --git a/datafusion/physical-expr-common/benches/arrow_bytes_map.rs b/datafusion/physical-expr-common/benches/arrow_bytes_map.rs index 7c8cdc3b4c50e..68351a839554a 100644 --- a/datafusion/physical-expr-common/benches/arrow_bytes_map.rs +++ b/datafusion/physical-expr-common/benches/arrow_bytes_map.rs @@ -17,7 +17,9 @@ use arrow::array::{ArrayRef, StringArray}; use criterion::{Criterion, Throughput, criterion_group, criterion_main}; -use datafusion_physical_expr_common::binary_map::{ArrowBytesMap, OutputType}; +use datafusion_physical_expr_common::binary_map::{ + ArrowBytesMap, INITIAL_MAP_CAPACITY, OutputType, +}; use std::hint::black_box; use std::sync::Arc; @@ -57,7 +59,13 @@ fn bench_arrow_bytes_map(c: &mut Criterion) { for (name, values) in cases { group.bench_function(name, |b| { b.iter(|| { - let mut map = ArrowBytesMap::::new(OutputType::Utf8); + // The `long_low_cardinality` case is defined by the distinct + // values fitting in the pre-allocated buffer, so this benchmark + // measures the pre-allocating constructor. + let mut map = ArrowBytesMap::::with_capacity( + OutputType::Utf8, + INITIAL_MAP_CAPACITY, + ); let mut next_payload = 0; map.insert_if_new( &values, diff --git a/datafusion/physical-expr-common/src/binary_map.rs b/datafusion/physical-expr-common/src/binary_map.rs index ceae72d4ccafe..61b1e6e008c7c 100644 --- a/datafusion/physical-expr-common/src/binary_map.rs +++ b/datafusion/physical-expr-common/src/binary_map.rs @@ -55,10 +55,21 @@ pub enum OutputType { pub struct ArrowBytesSet(ArrowBytesMap); impl ArrowBytesSet { + /// Creates a set that does not pre-allocate its hash table or value buffer. + /// + /// See [`ArrowBytesMap::new`] for when to prefer this over + /// [`Self::with_capacity`]. pub fn new(output_type: OutputType) -> Self { Self(ArrowBytesMap::new(output_type)) } + /// Creates a set with room for `map_capacity` entries. + /// + /// See [`ArrowBytesMap::with_capacity`]. + pub fn with_capacity(output_type: OutputType, map_capacity: usize) -> Self { + Self(ArrowBytesMap::with_capacity(output_type, map_capacity)) + } + /// Return the contents of this set and replace it with a new empty /// set with the same output type pub fn take(&mut self) -> Self { @@ -217,6 +228,13 @@ where output_type: OutputType, /// Underlying hash set for each distinct value map: hashbrown::hash_table::HashTable>, + /// Hash table capacity to re-create the map with in [`Self::take`], so a + /// map built with [`Self::with_capacity`] keeps its pre-allocation when it + /// is emptied and reused + initial_map_capacity: usize, + /// Value buffer capacity to re-create the buffer with in [`Self::take`], + /// for the same reason as `initial_map_capacity` + initial_buffer_capacity: usize, /// In progress buffer containing all values buffer: Vec, /// Offsets into `buffer` for each distinct value. These offsets as used @@ -234,19 +252,49 @@ where null: Option<(V, usize)>, } -/// The size, in number of entries, of the initial hash table -const INITIAL_MAP_CAPACITY: usize = 128; -/// The initial size, in bytes, of the string data +/// The size, in number of entries, of the hash table pre-allocated by +/// [`ArrowBytesMap::with_capacity`]. It is a warm up size for maps that go on +/// to hold many values, not a bound on what the map can hold. +pub const INITIAL_MAP_CAPACITY: usize = 128; +/// The size, in bytes, of the string data buffer pre-allocated by +/// [`ArrowBytesMap::with_capacity`] pub const INITIAL_BUFFER_CAPACITY: usize = 8 * 1024; impl ArrowBytesMap where V: Debug + PartialEq + Eq + Clone + Copy + Default, { + /// Creates a map that does not pre-allocate its hash table or value buffer. + /// + /// Use this when maps are created in large numbers and most of them stay + /// small, such as the per group `COUNT(DISTINCT)` accumulators that + /// `GroupsAccumulatorAdapter` creates one of per group. There the + /// pre-allocation dwarfs the values the map actually holds. pub fn new(output_type: OutputType) -> Self { + Self::new_inner(output_type, 0, 0) + } + + /// Creates a map whose hash table is pre-allocated for `map_capacity` + /// entries and whose value buffer is pre-allocated with + /// [`INITIAL_BUFFER_CAPACITY`] bytes. + /// + /// Use this for the few long lived maps that are each expected to hold many + /// values, such as the single map backing a `GROUP BY` on one string + /// column. The capacities are preserved across [`Self::take`]. + pub fn with_capacity(output_type: OutputType, map_capacity: usize) -> Self { + Self::new_inner(output_type, map_capacity, INITIAL_BUFFER_CAPACITY) + } + + fn new_inner( + output_type: OutputType, + map_capacity: usize, + buffer_capacity: usize, + ) -> Self { Self { output_type, - map: hashbrown::hash_table::HashTable::with_capacity(INITIAL_MAP_CAPACITY), - buffer: Vec::with_capacity(INITIAL_BUFFER_CAPACITY), + map: hashbrown::hash_table::HashTable::with_capacity(map_capacity), + initial_map_capacity: map_capacity, + initial_buffer_capacity: buffer_capacity, + buffer: Vec::with_capacity(buffer_capacity), offsets: vec![O::default()], // first offset is always 0 random_state: RandomState::default(), hashes_buffer: vec![], @@ -257,7 +305,11 @@ where /// Return the contents of this map and replace it with a new empty map with /// the same output type pub fn take(&mut self) -> Self { - let mut new_self = Self::new(self.output_type); + let mut new_self = Self::new_inner( + self.output_type, + self.initial_map_capacity, + self.initial_buffer_capacity, + ); swap(self, &mut new_self); new_self } @@ -412,7 +464,8 @@ where offset_or_inline: inline, payload, }; - self.map.insert_unique(hash, new_header, |header| header.hash); + self.map + .insert_unique(hash, new_header, |header| header.hash); payload } } @@ -450,7 +503,8 @@ where offset_or_inline: offset, payload, }; - self.map.insert_unique(hash, new_header, |header| header.hash); + self.map + .insert_unique(hash, new_header, |header| header.hash); payload } }; @@ -475,6 +529,8 @@ where let Self { output_type, map: _, + initial_map_capacity: _, + initial_buffer_capacity: _, offsets, buffer, random_state: _, @@ -655,6 +711,68 @@ mod tests { use arrow::array::{BinaryArray, LargeBinaryArray, StringArray}; use std::collections::HashMap; + /// The bytes a hashbrown table of `buckets` buckets must allocate for + /// entries of type `T`, ignoring the alignment padding and the trailing + /// group. Derived independently of the production accounting so it can + /// bracket it. + fn min_table_bytes(buckets: usize) -> usize { + // One entry slot plus one control byte per bucket. + buckets * (size_of::() + 1) + } + + #[test] + fn map_new_does_not_allocate() { + let map = ArrowBytesMap::::new(OutputType::Utf8); + + assert_eq!(map.map.capacity(), 0); + assert_eq!(map.map.allocation_size(), 0); + assert_eq!(map.buffer.capacity(), 0); + // Only the single leading zero offset is allocated. + assert!(map.size() < 128, "expected {} to be tiny", map.size()); + } + + #[test] + fn map_with_capacity_reports_the_real_hash_table_allocation() { + let map = ArrowBytesMap::::with_capacity( + OutputType::Utf8, + INITIAL_MAP_CAPACITY, + ); + + assert!(map.map.capacity() >= INITIAL_MAP_CAPACITY); + assert_eq!(map.buffer.capacity(), INITIAL_BUFFER_CAPACITY); + + // Before this accounting was corrected the map reported its hash table + // as costing zero bytes until the table grew past its pre-allocation. + let table_bytes = map.map.allocation_size(); + let lower_bound = min_table_bytes::>(map.map.capacity()); + assert!( + table_bytes >= lower_bound, + "expected {table_bytes} to be at least {lower_bound}" + ); + assert!( + table_bytes <= 2 * lower_bound + 64, + "expected {table_bytes} to be within a small factor of {lower_bound}" + ); + assert!(map.size() >= table_bytes + INITIAL_BUFFER_CAPACITY); + } + + #[test] + fn take_preserves_the_capacity_the_map_was_built_with() { + let mut preallocated = ArrowBytesMap::::with_capacity( + OutputType::Utf8, + INITIAL_MAP_CAPACITY, + ); + let capacity = preallocated.map.capacity(); + preallocated.take(); + assert_eq!(preallocated.map.capacity(), capacity); + assert_eq!(preallocated.buffer.capacity(), INITIAL_BUFFER_CAPACITY); + + let mut lazy = ArrowBytesMap::::new(OutputType::Utf8); + lazy.take(); + assert_eq!(lazy.map.capacity(), 0); + assert_eq!(lazy.buffer.capacity(), 0); + } + #[test] fn string_set_empty() { let mut set = ArrowBytesSet::::new(OutputType::Utf8); diff --git a/datafusion/physical-expr-common/src/binary_view_map.rs b/datafusion/physical-expr-common/src/binary_view_map.rs index 304b321870467..9351d92092984 100644 --- a/datafusion/physical-expr-common/src/binary_view_map.rs +++ b/datafusion/physical-expr-common/src/binary_view_map.rs @@ -39,10 +39,21 @@ use std::sync::Arc; pub struct ArrowBytesViewSet(ArrowBytesViewMap<()>); impl ArrowBytesViewSet { + /// Creates a set that does not pre-allocate its hash table. + /// + /// See [`ArrowBytesViewMap::new`] for when to prefer this over + /// [`Self::with_capacity`]. pub fn new(output_type: OutputType) -> Self { Self(ArrowBytesViewMap::new(output_type)) } + /// Creates a set with room for `map_capacity` entries. + /// + /// See [`ArrowBytesViewMap::with_capacity`]. + pub fn with_capacity(output_type: OutputType, map_capacity: usize) -> Self { + Self(ArrowBytesViewMap::with_capacity(output_type, map_capacity)) + } + /// Inserts each value from `values` into the set pub fn insert(&mut self, values: &ArrayRef) { fn make_payload_fn(_value: Option<&[u8]>) {} @@ -54,9 +65,7 @@ impl ArrowBytesViewSet { /// Return the contents of this map and replace it with a new empty map with /// the same output type pub fn take(&mut self) -> Self { - let mut new_self = Self::new(self.0.output_type); - std::mem::swap(self, &mut new_self); - new_self + Self(self.0.take()) } /// Converts this set into a `StringViewArray` or `BinaryViewArray` @@ -128,6 +137,10 @@ where output_type: OutputType, /// Underlying hash set for each distinct value map: hashbrown::hash_table::HashTable>, + /// Hash table capacity to re-create the map with in [`Self::take`], so a + /// map built with [`Self::with_capacity`] keeps its pre-allocation when it + /// is emptied and reused + initial_map_capacity: usize, /// Views for all stored values (in insertion order) views: Vec, @@ -148,17 +161,36 @@ where null: Option<(V, usize)>, } -/// The size, in number of entries, of the initial hash table -const INITIAL_MAP_CAPACITY: usize = 512; +/// The size, in number of entries, of the hash table pre-allocated by +/// [`ArrowBytesViewMap::with_capacity`]. It is a warm up size for maps that go +/// on to hold many values, not a bound on what the map can hold. +pub const INITIAL_MAP_CAPACITY: usize = 512; impl ArrowBytesViewMap where V: Debug + PartialEq + Eq + Clone + Copy + Default, { + /// Creates a map that does not pre-allocate its hash table. + /// + /// Use this when maps are created in large numbers and most of them stay + /// small, such as the per group `COUNT(DISTINCT)` accumulators that + /// `GroupsAccumulatorAdapter` creates one of per group. There the + /// pre-allocation dwarfs the values the map actually holds. pub fn new(output_type: OutputType) -> Self { + Self::with_capacity(output_type, 0) + } + + /// Creates a map whose hash table is pre-allocated for `map_capacity` + /// entries. + /// + /// Use this for the few long lived maps that are each expected to hold many + /// values, such as the single map backing a `GROUP BY` on one string + /// column. The capacity is preserved across [`Self::take`]. + pub fn with_capacity(output_type: OutputType, map_capacity: usize) -> Self { Self { output_type, - map: hashbrown::hash_table::HashTable::with_capacity(INITIAL_MAP_CAPACITY), + map: hashbrown::hash_table::HashTable::with_capacity(map_capacity), + initial_map_capacity: map_capacity, views: Vec::new(), in_progress: Vec::new(), completed: Vec::new(), @@ -172,7 +204,8 @@ where /// Return the contents of this map and replace it with a new empty map with /// the same output type pub fn take(&mut self) -> Self { - let mut new_self = Self::new(self.output_type); + let mut new_self = + Self::with_capacity(self.output_type, self.initial_map_capacity); std::mem::swap(self, &mut new_self); new_self } @@ -783,13 +816,48 @@ mod tests { assert_eq!(set.len(), 10); } + /// The bytes a hashbrown table of `buckets` buckets must allocate for + /// entries of type `T`, ignoring the alignment padding and the trailing + /// group. Derived independently of the production accounting so it can + /// bracket it. + fn min_table_bytes(buckets: usize) -> usize { + // One entry slot plus one control byte per bucket. + buckets * (size_of::() + 1) + } + #[test] - fn test_size_counts_initial_hash_table_capacity() { + fn map_new_does_not_allocate() { let map = ArrowBytesViewMap::<()>::new(OutputType::Utf8View); + assert_eq!(map.map.capacity(), 0); + assert_eq!(map.map.allocation_size(), 0); + assert_eq!(map.size(), 0); + } + + #[test] + fn test_size_counts_initial_hash_table_capacity() { + let map = ArrowBytesViewMap::<()>::with_capacity( + OutputType::Utf8View, + INITIAL_MAP_CAPACITY, + ); + + assert!(map.map.capacity() >= INITIAL_MAP_CAPACITY); assert_eq!(map.size(), map.map.allocation_size()); - // The reported size covers the control bytes as well as the entries, so - // it is strictly larger than the entry array on its own. + + // Before this accounting was corrected the map reported exactly + // `capacity() * size_of::>()`, which leaves out the control + // bytes and undercounts the real allocation by roughly half. + let lower_bound = min_table_bytes::>(map.map.capacity()); + assert!( + map.size() >= lower_bound, + "expected {} to be at least {lower_bound}", + map.size() + ); + assert!( + map.size() <= 2 * lower_bound + 64, + "expected {} to be within a small factor of {lower_bound}", + map.size() + ); assert!( map.size() > map.map.capacity() * size_of::>(), "expected {} to exceed {}", @@ -798,6 +866,21 @@ mod tests { ); } + #[test] + fn take_preserves_the_capacity_the_map_was_built_with() { + let mut preallocated = ArrowBytesViewMap::<()>::with_capacity( + OutputType::Utf8View, + INITIAL_MAP_CAPACITY, + ); + let capacity = preallocated.map.capacity(); + preallocated.take(); + assert_eq!(preallocated.map.capacity(), capacity); + + let mut lazy = ArrowBytesViewMap::<()>::new(OutputType::Utf8View); + lazy.take(); + assert_eq!(lazy.map.capacity(), 0); + } + #[test] fn test_size_counts_retained_buffer_capacities() { let first = "a".repeat(BYTE_VIEW_MAX_BLOCK_SIZE / 2 + 1); diff --git a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes.rs b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes.rs index 34ec36be31d2e..87eeba489ecc9 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes.rs @@ -22,7 +22,9 @@ use crate::aggregates::group_values::GroupValues; use arrow::array::{Array, ArrayRef, OffsetSizeTrait}; use datafusion_common::Result; use datafusion_expr::{EmitTo, GroupSelection}; -use datafusion_physical_expr_common::binary_map::{ArrowBytesMap, OutputType}; +use datafusion_physical_expr_common::binary_map::{ + ArrowBytesMap, INITIAL_MAP_CAPACITY, OutputType, +}; /// A [`GroupValues`] storing single column of Utf8/LargeUtf8/Binary/LargeBinary values /// @@ -38,7 +40,9 @@ pub struct GroupValuesBytes { impl GroupValuesBytes { pub fn new(output_type: OutputType) -> Self { Self { - map: ArrowBytesMap::new(output_type), + // One map holds every group value for the whole query, so it is + // worth pre-allocating the hash table and the value buffer. + map: ArrowBytesMap::with_capacity(output_type, INITIAL_MAP_CAPACITY), num_groups: 0, } } diff --git a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes_view.rs b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes_view.rs index 997a7ce166a71..f5183ba329b30 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes_view.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes_view.rs @@ -19,7 +19,9 @@ use crate::aggregates::group_values::GroupValues; use arrow::array::{Array, ArrayRef}; use datafusion_expr::{EmitTo, GroupSelection}; use datafusion_physical_expr::binary_map::OutputType; -use datafusion_physical_expr_common::binary_view_map::ArrowBytesViewMap; +use datafusion_physical_expr_common::binary_view_map::{ + ArrowBytesViewMap, INITIAL_MAP_CAPACITY, +}; use std::mem::size_of; /// A [`GroupValues`] storing single column of Utf8View/BinaryView values @@ -36,7 +38,9 @@ pub struct GroupValuesBytesView { impl GroupValuesBytesView { pub fn new(output_type: OutputType) -> Self { Self { - map: ArrowBytesViewMap::new(output_type), + // One map holds every group value for the whole query, so it is + // worth pre-allocating the hash table. + map: ArrowBytesViewMap::with_capacity(output_type, INITIAL_MAP_CAPACITY), num_groups: 0, } } From 7cccdd244f2e3e62b34e93a1e4699b2970284b26 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Tue, 1 Sep 2026 01:46:45 -0500 Subject: [PATCH 33/37] refactor: name the hash table term in ArrowBytesViewMap::size Keep the comment about what `HashTable::allocation_size` covers next to the value it describes, and say what the test helper's lower bound is derived from. --- datafusion/physical-expr-common/src/binary_map.rs | 12 +++++------- .../physical-expr-common/src/binary_view_map.rs | 15 +++++++-------- 2 files changed, 12 insertions(+), 15 deletions(-) diff --git a/datafusion/physical-expr-common/src/binary_map.rs b/datafusion/physical-expr-common/src/binary_map.rs index 61b1e6e008c7c..0b59bcf5937bc 100644 --- a/datafusion/physical-expr-common/src/binary_map.rs +++ b/datafusion/physical-expr-common/src/binary_map.rs @@ -711,13 +711,11 @@ mod tests { use arrow::array::{BinaryArray, LargeBinaryArray, StringArray}; use std::collections::HashMap; - /// The bytes a hashbrown table of `buckets` buckets must allocate for - /// entries of type `T`, ignoring the alignment padding and the trailing - /// group. Derived independently of the production accounting so it can - /// bracket it. - fn min_table_bytes(buckets: usize) -> usize { - // One entry slot plus one control byte per bucket. - buckets * (size_of::() + 1) + /// A lower bound on the bytes a hashbrown table holding `entries` entries + /// of type `T` must allocate: one entry slot and one control byte each. + /// Derived independently of the production accounting so it can bracket it. + fn min_table_bytes(entries: usize) -> usize { + entries * (size_of::() + 1) } #[test] diff --git a/datafusion/physical-expr-common/src/binary_view_map.rs b/datafusion/physical-expr-common/src/binary_view_map.rs index 9351d92092984..35f45e8e9a12e 100644 --- a/datafusion/physical-expr-common/src/binary_view_map.rs +++ b/datafusion/physical-expr-common/src/binary_view_map.rs @@ -570,13 +570,14 @@ where // which is the entry array plus the control bytes plus the trailing // group, so it is larger than `capacity() * size_of::>()`. It // is a constant time layout calculation, not a walk of the table. + let map_size = self.map.allocation_size(); let views_size = self.views.allocated_size(); let in_progress_size = self.in_progress.allocated_size(); let completed_size = self.completed.allocated_size() + self.completed.iter().map(Buffer::capacity).sum::(); let nulls_size = self.nulls.allocated_size(); - self.map.allocation_size() + map_size + views_size + in_progress_size + completed_size @@ -816,13 +817,11 @@ mod tests { assert_eq!(set.len(), 10); } - /// The bytes a hashbrown table of `buckets` buckets must allocate for - /// entries of type `T`, ignoring the alignment padding and the trailing - /// group. Derived independently of the production accounting so it can - /// bracket it. - fn min_table_bytes(buckets: usize) -> usize { - // One entry slot plus one control byte per bucket. - buckets * (size_of::() + 1) + /// A lower bound on the bytes a hashbrown table holding `entries` entries + /// of type `T` must allocate: one entry slot and one control byte each. + /// Derived independently of the production accounting so it can bracket it. + fn min_table_bytes(entries: usize) -> usize { + entries * (size_of::() + 1) } #[test] From a4f43dc83369f238730ee0202c54fcc821313d4f Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Tue, 1 Sep 2026 12:06:54 -0500 Subject: [PATCH 34/37] fix: release the group values map allocations in `clear_shrink` `GroupValuesBytes::clear_shrink` and `GroupValuesBytesView::clear_shrink` reset their map with `take()`, which restores the capacity the map was configured with so the emptied map stays warm. That is what the emit path wants, but `clear_shrink` exists to hand memory back before spilling and before the spilled batch is sorted, so it left roughly 16 KiB (string and binary) and 34 KiB (view) reserved instead of releasing it. Add `clear_and_release` to `ArrowBytesMap` and `ArrowBytesViewMap`, which empties the map and drops its allocations while remembering the configured capacities so a later `take()` still warms the map up, and call it from the two `clear_shrink` implementations. The pre-allocation stays at construction, where the hot single column string `GROUP BY` path earns it. --- .../physical-expr-common/src/binary_map.rs | 55 ++++++++++++++++ .../src/binary_view_map.rs | 52 +++++++++++++++ .../group_values/single_group_by/bytes.rs | 65 ++++++++++++++++++- .../single_group_by/bytes_view.rs | 64 +++++++++++++++++- 4 files changed, 230 insertions(+), 6 deletions(-) diff --git a/datafusion/physical-expr-common/src/binary_map.rs b/datafusion/physical-expr-common/src/binary_map.rs index 0b59bcf5937bc..bc7d987173770 100644 --- a/datafusion/physical-expr-common/src/binary_map.rs +++ b/datafusion/physical-expr-common/src/binary_map.rs @@ -314,6 +314,22 @@ where new_self } + /// Empties this map and releases every allocation it holds, so + /// [`Self::size`] drops to approximately zero. + /// + /// This is the difference from [`Self::take`]: `take` restores the + /// capacities the map was configured with so the emptied map stays warm for + /// continued use, which is what emitting wants, whereas here the point is + /// to hand the memory back, as before spilling or before a downstream sort. + /// The configured capacities are remembered, so the next [`Self::take`] + /// warms the map up again. + pub fn clear_and_release(&mut self) { + let mut released = Self::new_inner(self.output_type, 0, 0); + released.initial_map_capacity = self.initial_map_capacity; + released.initial_buffer_capacity = self.initial_buffer_capacity; + *self = released; + } + /// Inserts each value from `values` into the map, invoking `payload_fn` for /// each value if *not* already present, deferring the allocation of the /// payload until it is needed. @@ -771,6 +787,45 @@ mod tests { assert_eq!(lazy.buffer.capacity(), 0); } + #[test] + fn clear_and_release_frees_the_preallocation_that_take_keeps() { + let mut map = ArrowBytesMap::::with_capacity( + OutputType::Utf8, + INITIAL_MAP_CAPACITY, + ); + let values: ArrayRef = Arc::new(StringArray::from_iter_values( + (0..1_000).map(|i| format!("distinct value number {i}")), + )); + map.insert_if_new(&values, |_| (), |_| ()); + + let populated_size = map.size(); + assert!(populated_size > INITIAL_BUFFER_CAPACITY); + + // `take` deliberately keeps the map warm, so it does not release the + // configured capacities. + map.take(); + let taken_size = map.size(); + assert!( + taken_size > INITIAL_BUFFER_CAPACITY, + "expected take to retain the warm up allocations, got {taken_size}" + ); + + map.clear_and_release(); + let released_size = map.size(); + assert_eq!(map.map.allocation_size(), 0); + assert_eq!(map.buffer.capacity(), 0); + assert!( + released_size < 128, + "expected the released map to report approximately zero bytes, got {released_size}" + ); + + // The configured capacities survive, so the map warms back up when it + // is emitted from again. + map.take(); + assert!(map.map.capacity() >= INITIAL_MAP_CAPACITY); + assert_eq!(map.buffer.capacity(), INITIAL_BUFFER_CAPACITY); + } + #[test] fn string_set_empty() { let mut set = ArrowBytesSet::::new(OutputType::Utf8); diff --git a/datafusion/physical-expr-common/src/binary_view_map.rs b/datafusion/physical-expr-common/src/binary_view_map.rs index 35f45e8e9a12e..29f4014c5f9a4 100644 --- a/datafusion/physical-expr-common/src/binary_view_map.rs +++ b/datafusion/physical-expr-common/src/binary_view_map.rs @@ -210,6 +210,21 @@ where new_self } + /// Empties this map and releases every allocation it holds, so + /// [`Self::size`] drops to approximately zero. + /// + /// This is the difference from [`Self::take`]: `take` restores the capacity + /// the map was configured with so the emptied map stays warm for continued + /// use, which is what emitting wants, whereas here the point is to hand the + /// memory back, as before spilling or before a downstream sort. The + /// configured capacity is remembered, so the next [`Self::take`] warms the + /// map up again. + pub fn clear_and_release(&mut self) { + let mut released = Self::with_capacity(self.output_type, 0); + released.initial_map_capacity = self.initial_map_capacity; + *self = released; + } + /// Inserts each value from `values` into the map, invoking `payload_fn` for /// each value if *not* already present, deferring the allocation of the /// payload until it is needed. @@ -880,6 +895,43 @@ mod tests { assert_eq!(lazy.map.capacity(), 0); } + #[test] + fn clear_and_release_frees_the_preallocation_that_take_keeps() { + let mut map = ArrowBytesViewMap::<()>::with_capacity( + OutputType::Utf8View, + INITIAL_MAP_CAPACITY, + ); + let values: ArrayRef = Arc::new(StringViewArray::from_iter_values( + (0..1_000).map(|i| format!("distinct value number {i}")), + )); + map.insert_if_new(&values, |_| (), |_| ()); + + let warm_size = map.map.allocation_size(); + assert!(warm_size > 0); + + // `take` deliberately keeps the map warm, so it does not release the + // configured capacity. + map.take(); + let taken_size = map.size(); + assert!( + taken_size >= min_table_bytes::>(INITIAL_MAP_CAPACITY), + "expected take to retain the warm up allocation, got {taken_size}" + ); + + map.clear_and_release(); + assert_eq!(map.map.allocation_size(), 0); + let released_size = map.size(); + assert!( + released_size < 128, + "expected the released map to report approximately zero bytes, got {released_size}" + ); + + // The configured capacity survives, so the map warms back up when it is + // emitted from again. + map.take(); + assert!(map.map.capacity() >= INITIAL_MAP_CAPACITY); + } + #[test] fn test_size_counts_retained_buffer_capacities() { let first = "a".repeat(BYTE_VIEW_MAX_BLOCK_SIZE / 2 + 1); diff --git a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes.rs b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes.rs index 87eeba489ecc9..9f7b4b4e91cba 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes.rs @@ -137,8 +137,67 @@ impl GroupValues for GroupValuesBytes { } fn clear_shrink(&mut self, _num_rows: usize) { - // in theory we could potentially avoid this reallocation and clear the - // contents of the maps, but for now we just reset the map from the beginning - self.map.take(); + // Callers use this to hand memory back before spilling or sorting, so + // release the map's allocations rather than restoring the warm up + // capacities that `take` keeps for the emit path. + self.map.clear_and_release(); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + use std::sync::Arc; + + use arrow::array::StringArray; + use datafusion_physical_expr_common::binary_map::INITIAL_BUFFER_CAPACITY; + + /// `clear_shrink` is how the aggregate stream hands memory back before it + /// spills and before the spilled batch is sorted, so the memory it releases + /// has to actually show up in the size it reports afterwards. + #[test] + fn clear_shrink_releases_the_reported_memory() { + let mut group_values = GroupValuesBytes::::new(OutputType::Utf8); + let empty = size_of::>(); + + // The map is pre-allocated at construction, so it is already well above + // its own struct size before a single row is interned. + let warm_size = group_values.size(); + assert!( + warm_size > empty + INITIAL_BUFFER_CAPACITY, + "expected the pre-allocated map to report more than {} bytes, got {warm_size}", + empty + INITIAL_BUFFER_CAPACITY + ); + + let values: ArrayRef = Arc::new(StringArray::from_iter_values( + (0..1_000).map(|i| format!("group value number {i}")), + )); + let mut groups = vec![]; + group_values + .intern(&[Arc::clone(&values)], &mut groups) + .unwrap(); + let populated_size = group_values.size(); + assert!(populated_size > warm_size); + + group_values.clear_shrink(0); + + // Everything the map held is gone: what remains is the struct itself + // plus the single leading zero offset. + let released_size = group_values.size(); + assert!( + released_size < empty + 128, + "expected clear_shrink to release the map, got {released_size} with a struct size of {empty}" + ); + assert!( + released_size * 10 < populated_size, + "expected {released_size} to be far below {populated_size}" + ); + + // The map still works, and warms back up on the next emit. + group_values.intern(&[values], &mut groups).unwrap(); + assert!(group_values.size() > released_size); + group_values.emit(EmitTo::All).unwrap(); + assert!(group_values.size() > empty + INITIAL_BUFFER_CAPACITY); } } diff --git a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes_view.rs b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes_view.rs index f5183ba329b30..23ea4e7ed3f88 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes_view.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes_view.rs @@ -139,8 +139,66 @@ impl GroupValues for GroupValuesBytesView { } fn clear_shrink(&mut self, _num_rows: usize) { - // in theory we could potentially avoid this reallocation and clear the - // contents of the maps, but for now we just reset the map from the beginning - self.map.take(); + // Callers use this to hand memory back before spilling or sorting, so + // release the map's allocations rather than restoring the warm up + // capacity that `take` keeps for the emit path. + self.map.clear_and_release(); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + use std::sync::Arc; + + use arrow::array::StringViewArray; + + /// `clear_shrink` is how the aggregate stream hands memory back before it + /// spills and before the spilled batch is sorted, so the memory it releases + /// has to actually show up in the size it reports afterwards. + #[test] + fn clear_shrink_releases_the_reported_memory() { + let mut group_values = GroupValuesBytesView::new(OutputType::Utf8View); + let empty = size_of::(); + + // The hash table is pre-allocated at construction, so the map is + // already well above its own struct size before a single row is + // interned. + let warm_size = group_values.size(); + assert!( + warm_size > empty + INITIAL_MAP_CAPACITY, + "expected the pre-allocated map to report more than {} bytes, got {warm_size}", + empty + INITIAL_MAP_CAPACITY + ); + + let values: ArrayRef = Arc::new(StringViewArray::from_iter_values( + (0..1_000).map(|i| format!("group value number {i}")), + )); + let mut groups = vec![]; + group_values + .intern(&[Arc::clone(&values)], &mut groups) + .unwrap(); + let populated_size = group_values.size(); + assert!(populated_size > warm_size); + + group_values.clear_shrink(0); + + // Everything the map held is gone: what remains is the struct itself. + let released_size = group_values.size(); + assert!( + released_size < empty + 128, + "expected clear_shrink to release the map, got {released_size} with a struct size of {empty}" + ); + assert!( + released_size * 10 < populated_size, + "expected {released_size} to be far below {populated_size}" + ); + + // The map still works, and warms back up on the next emit. + group_values.intern(&[values], &mut groups).unwrap(); + assert!(group_values.size() > released_size); + group_values.emit(EmitTo::All).unwrap(); + assert!(group_values.size() > empty + INITIAL_MAP_CAPACITY); } } From b0cd6fba5927059a58cb0609c26d3479c2eae3d6 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Tue, 1 Sep 2026 16:09:57 -0500 Subject: [PATCH 35/37] test: memory limit test for grouped `COUNT(DISTINCT )` A grouped `COUNT(DISTINCT )` gets one accumulator per group, and each of those owns a hash set of the distinct values it has seen. Those sets were created pre-allocated, so the query's memory use tracked the number of groups rather than the amount of data. Add two `memory_limit` tests that turn that into a binary observable, one for `Utf8` and one for `Utf8View`, over a new scenario of 4,000 groups holding 2 distinct values each. Measured against this branch's base commit with spilling disabled and `target_partitions` pinned to 1: | value column | budget needed before | budget needed after | | ------------ | -------------------- | ------------------- | | `Utf8` | ~35.5 MB | ~1.9 MB | | `Utf8View` | ~123 MB | ~2.7 MB | The tests run at 8 MB and 16 MB respectively, so each sits at least 4x above what the branch needs and at least 4x below what the base needs. Both fail on the base commit with `Resources exhausted` and pass here. --- datafusion/core/tests/memory_limit/mod.rs | 125 +++++++++++++++++++++- 1 file changed, 124 insertions(+), 1 deletion(-) diff --git a/datafusion/core/tests/memory_limit/mod.rs b/datafusion/core/tests/memory_limit/mod.rs index 9172e34a87c79..274ace182bf13 100644 --- a/datafusion/core/tests/memory_limit/mod.rs +++ b/datafusion/core/tests/memory_limit/mod.rs @@ -26,7 +26,9 @@ mod nlj_spill_unmatched; mod repartition_mem_limit; mod union_nullable_spill; mod view_spill_compaction; -use arrow::array::{ArrayRef, DictionaryArray, Int32Array, RecordBatch, StringViewArray}; +use arrow::array::{ + ArrayRef, DictionaryArray, Int32Array, RecordBatch, StringArray, StringViewArray, +}; use arrow::compute::SortOptions; use arrow::datatypes::{Int32Type, SchemaRef}; use arrow_schema::{DataType, Field, Schema}; @@ -220,6 +222,60 @@ mod count_distinct_spill { } } +/// A grouped `COUNT(DISTINCT )` gets one accumulator per group, and +/// each of those owns a hash set of the distinct values it has seen. Those +/// sets used to be created pre-allocated, which costs far more than the +/// handful of values a group typically holds, so the query's memory use +/// tracked the number of groups rather than the amount of data. +/// +/// With 4,000 groups holding 2 distinct values each, this query needed about +/// 35.5 MB of budget before the per group pre-allocation was removed and +/// about 1.9 MB after, so an 8 MB limit is a failure before the change and a +/// success after it. Spilling is disabled, so completing means the query +/// genuinely fit in the budget. +/// +/// The `count(*)` is load bearing: without it +/// `single_distinct_aggregation_to_group_by` rewrites the distinct aggregate +/// into a plain two stage `GROUP BY`, which does not use these accumulators +/// at all. +#[tokio::test] +async fn group_by_count_distinct_utf8() { + TestCase::new() + .with_query( + "select group_key, count(distinct value), count(*) from t group by group_key", + ) + .with_scenario(Scenario::GroupedDistinctStrings { + groups: 4_000, + string_view: false, + }) + .with_config(SessionConfig::new().with_target_partitions(1)) + .with_memory_limit(8_000_000) + .with_expected_success() + .run() + .await +} + +/// The `Utf8View` counterpart of [`group_by_count_distinct_utf8`], covering +/// the separate view flavoured hash set. The same query over a `Utf8View` +/// column needed about 123 MB of budget before the change and about 2.7 MB +/// after, so 16 MB separates the two. +#[tokio::test] +async fn group_by_count_distinct_utf8_view() { + TestCase::new() + .with_query( + "select group_key, count(distinct value), count(*) from t group by group_key", + ) + .with_scenario(Scenario::GroupedDistinctStrings { + groups: 4_000, + string_view: true, + }) + .with_config(SessionConfig::new().with_target_partitions(1)) + .with_memory_limit(16_000_000) + .with_expected_success() + .run() + .await +} + #[tokio::test] async fn join_by_key_multiple_partitions() { let config = SessionConfig::new().with_target_partitions(2); @@ -1077,6 +1133,14 @@ enum Scenario { /// If true, splits all input batches into 1 row each single_row_batches: bool, }, + + /// `groups` distinct integer keys paired with a short string value, for + /// grouped aggregates that build one accumulator per group. + GroupedDistinctStrings { + groups: usize, + /// If true, the value column is `Utf8View` rather than `Utf8` + string_view: bool, + }, } impl Scenario { @@ -1151,6 +1215,15 @@ impl Scenario { let table = SortedTableProvider::new(batches, sort_information); Arc::new(table) } + Self::GroupedDistinctStrings { + groups, + string_view, + } => { + let batches = grouped_distinct_string_batches(*groups, *string_view); + let table = + MemTable::try_new(batches[0].schema(), vec![batches]).unwrap(); + Arc::new(table) + } } } @@ -1174,10 +1247,60 @@ impl Scenario { // Use default rules None } + Self::GroupedDistinctStrings { .. } => { + // Disable the rules that would add a repartition, so the test + // measures the aggregate's budget rather than a repartition's + Some(vec![Arc::new(JoinSelection::new())]) + } } } } +/// Number of distinct string values held by every group produced by +/// [`grouped_distinct_string_batches`] +const DISTINCT_VALUES_PER_GROUP: usize = 2; + +/// Returns batches of 1024 rows with `groups` distinct keys in `group_key`, +/// each key paired with [`DISTINCT_VALUES_PER_GROUP`] distinct short strings +/// in `value`. The values are `Utf8View` if `string_view` is set, `Utf8` +/// otherwise. +fn grouped_distinct_string_batches(groups: usize, string_view: bool) -> Vec { + let value_type = if string_view { + DataType::Utf8View + } else { + DataType::Utf8 + }; + let schema = Arc::new(Schema::new(vec![ + Field::new("group_key", DataType::Int32, false), + Field::new("value", value_type, false), + ])); + + const ROWS_PER_BATCH: usize = 1024; + + let rows = groups * DISTINCT_VALUES_PER_GROUP; + let mut keys = Vec::with_capacity(rows); + let mut values = Vec::with_capacity(rows); + for value in 0..DISTINCT_VALUES_PER_GROUP { + for group in 0..groups { + keys.push(group as i32); + values.push(format!("value-{value}")); + } + } + + keys.chunks(ROWS_PER_BATCH) + .zip(values.chunks(ROWS_PER_BATCH)) + .map(|(keys, values)| { + let keys: ArrayRef = Arc::new(Int32Array::from(keys.to_vec())); + let values: ArrayRef = if string_view { + Arc::new(StringViewArray::from_iter_values(values)) + } else { + Arc::new(StringArray::from_iter_values(values)) + }; + RecordBatch::try_new(Arc::clone(&schema), vec![keys, values]).unwrap() + }) + .collect() +} + fn access_log_batches() -> Vec { AccessLogGenerator::new() .with_row_limit(1000) From e55a3aa15343bd4083bea7e5c73e86286222a5aa Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Tue, 1 Sep 2026 16:40:23 -0500 Subject: [PATCH 36/37] test: keep the memory limit tests off the single-distinct rewrite The two grouped `COUNT(DISTINCT )` memory limit tests only reach the per group accumulators while `single_distinct_aggregation_to_group_by` declines to rewrite the query. They leant on `count(*)` for that, which the rule rejects only because `count` is missing from the `sum`/`min`/`max` allow list. apache/datafusion#24859 proposes adding `count` to that list, which would rewrite the query, remove the accumulators, and leave both tests passing at any memory limit while still looking like they test something. Aggregate `avg(payload)` over a new `Int64` column instead. `avg` cannot be added to that list: the rule re-aggregates its own partial results over the deduplicated inner group by, and averaging per group averages of different sizes gives the wrong answer. That is why ClickBench Q9 keeps its distinct aggregate under #24859. Verified from the physical plan with #24859 cherry-picked on top of this branch: the `avg` query still plans as `aggr=[count(DISTINCT t.value), avg(t.payload)]`, while the `count(*)` query becomes `aggr=[count(alias1), sum(alias2)]` over an inner `GROUP BY group_key, value`, and drops from needing ~1.9 MB to ~0.9 MB. Re-swept both thresholds against the base commit. `Utf8` needs ~35.5 MB before and ~1.9 MB after; `Utf8View` needs ~123 MB before and ~2.5 MB after, so the 8 MB and 16 MB limits keep at least 4x margin on each side and are unchanged. --- datafusion/core/tests/memory_limit/mod.rs | 35 ++++++++++++++++------- 1 file changed, 24 insertions(+), 11 deletions(-) diff --git a/datafusion/core/tests/memory_limit/mod.rs b/datafusion/core/tests/memory_limit/mod.rs index 274ace182bf13..f873f956010e1 100644 --- a/datafusion/core/tests/memory_limit/mod.rs +++ b/datafusion/core/tests/memory_limit/mod.rs @@ -27,7 +27,8 @@ mod repartition_mem_limit; mod union_nullable_spill; mod view_spill_compaction; use arrow::array::{ - ArrayRef, DictionaryArray, Int32Array, RecordBatch, StringArray, StringViewArray, + ArrayRef, DictionaryArray, Int32Array, Int64Array, RecordBatch, StringArray, + StringViewArray, }; use arrow::compute::SortOptions; use arrow::datatypes::{Int32Type, SchemaRef}; @@ -234,15 +235,23 @@ mod count_distinct_spill { /// success after it. Spilling is disabled, so completing means the query /// genuinely fit in the budget. /// -/// The `count(*)` is load bearing: without it -/// `single_distinct_aggregation_to_group_by` rewrites the distinct aggregate -/// into a plain two stage `GROUP BY`, which does not use these accumulators -/// at all. +/// The `avg(payload)` is load bearing, and `avg` specifically. Without a +/// second aggregate, `single_distinct_aggregation_to_group_by` rewrites the +/// distinct aggregate into a plain two stage `GROUP BY` that does not use +/// these accumulators at all. That rule tolerates a non-distinct `sum`, `min` +/// or `max` beside the distinct aggregate, because it re-aggregates its own +/// partial results over the deduplicated inner group by, and those three +/// compose with themselves. `avg` does not: averaging per group averages of +/// different sizes gives the wrong answer, so the rule can never accept it. +/// That is why ClickBench Q9 keeps its distinct aggregate. Do not replace +/// this with `count(*)`: `count` is only incidentally rejected today, and +/// proposes accepting it, +/// which would rewrite the query and leave this test passing by construction. #[tokio::test] async fn group_by_count_distinct_utf8() { TestCase::new() .with_query( - "select group_key, count(distinct value), count(*) from t group by group_key", + "select group_key, count(distinct value), avg(payload) from t group by group_key", ) .with_scenario(Scenario::GroupedDistinctStrings { groups: 4_000, @@ -257,13 +266,13 @@ async fn group_by_count_distinct_utf8() { /// The `Utf8View` counterpart of [`group_by_count_distinct_utf8`], covering /// the separate view flavoured hash set. The same query over a `Utf8View` -/// column needed about 123 MB of budget before the change and about 2.7 MB +/// column needed about 123 MB of budget before the change and about 2.5 MB /// after, so 16 MB separates the two. #[tokio::test] async fn group_by_count_distinct_utf8_view() { TestCase::new() .with_query( - "select group_key, count(distinct value), count(*) from t group by group_key", + "select group_key, count(distinct value), avg(payload) from t group by group_key", ) .with_scenario(Scenario::GroupedDistinctStrings { groups: 4_000, @@ -1262,8 +1271,8 @@ const DISTINCT_VALUES_PER_GROUP: usize = 2; /// Returns batches of 1024 rows with `groups` distinct keys in `group_key`, /// each key paired with [`DISTINCT_VALUES_PER_GROUP`] distinct short strings -/// in `value`. The values are `Utf8View` if `string_view` is set, `Utf8` -/// otherwise. +/// in `value` and an `Int64` `payload` to aggregate over. The values are +/// `Utf8View` if `string_view` is set, `Utf8` otherwise. fn grouped_distinct_string_batches(groups: usize, string_view: bool) -> Vec { let value_type = if string_view { DataType::Utf8View @@ -1273,6 +1282,7 @@ fn grouped_distinct_string_batches(groups: usize, string_view: bool) -> Vec Vec Date: Wed, 2 Sep 2026 17:54:06 -0500 Subject: [PATCH 37/37] fix: grow the value buffer on a power of two ladder `ArrowBytesMap::new` starts its value buffer empty and `ArrowBytesMap::with_capacity` starts it at `INITIAL_BUFFER_CAPACITY`. `Vec` then doubles from wherever its first allocation landed, so the two sit on different ladders and can hold the same values at capacities differing by up to 2x, in either direction depending on the value lengths. Measured on 500,000 distinct 28 byte values, the lazily grown map reported 52,494,344 bytes against 45,154,312 for a pre-allocated one, 16% more for identical contents. That matters because the ungrouped `COUNT(DISTINCT )` accumulator is the caller that had a use for the warm up: it builds one map and grows it to hold every distinct value in the input. Rounding every buffer growth up to a power of two puts both constructors on one ladder, so a lazily allocated map is never larger than a pre-allocated one holding the same values. Growth stays geometric, so appending is still amortized constant time. `ArrowBytesViewMap` has no such buffer and is unaffected. Two new tests cover the ungrouped path, which had none: `ungrouped_utf8_accumulator_is_never_worse_than_a_pre_allocated_set` and its `Utf8View` counterpart drive an accumulator to 0 through 500,000 distinct values and assert it is strictly cheaper than a pre-allocated set at per group cardinalities and exactly equal at ungrouped ones. The `Utf8` one fails without this change, at 1,000 distinct values, with the lazy set reporting 110,408 bytes against 96,072. Two map level tests pin the ladder itself. --- .../src/aggregate/count_distinct/bytes.rs | 149 ++++++++++++++++++ .../physical-expr-common/src/binary_map.rs | 76 ++++++++- 2 files changed, 223 insertions(+), 2 deletions(-) diff --git a/datafusion/functions-aggregate-common/src/aggregate/count_distinct/bytes.rs b/datafusion/functions-aggregate-common/src/aggregate/count_distinct/bytes.rs index 3aa60f6f3b6ad..d955d343ad629 100644 --- a/datafusion/functions-aggregate-common/src/aggregate/count_distinct/bytes.rs +++ b/datafusion/functions-aggregate-common/src/aggregate/count_distinct/bytes.rs @@ -43,6 +43,14 @@ impl BytesDistinctCountAccumulator { /// creates one accumulator per group, so a grouped `COUNT(DISTINCT)` over a /// high cardinality key holds hundreds of thousands of these at once and /// most of them see only a handful of values. + /// + /// The ungrouped path builds one of these and grows it to hold every + /// distinct value in the input, so it is the caller that had a use for the + /// warm up. It loses nothing here: the set grows into exactly the + /// capacities a pre-allocated one reaches, which + /// `ungrouped_utf8_accumulator_is_never_worse_than_a_pre_allocated_set` + /// pins. That is why this constructor needs no signal distinguishing the + /// two callers. pub fn new(output_type: OutputType) -> Self { Self(ArrowBytesSet::new(output_type)) } @@ -157,3 +165,144 @@ impl Accumulator for BytesViewDistinctCountAccumulator { size_of_val(self) + self.0.size() } } + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::{StringArray, StringViewArray}; + use datafusion_physical_expr_common::binary_map::INITIAL_MAP_CAPACITY; + use datafusion_physical_expr_common::binary_view_map::INITIAL_MAP_CAPACITY as INITIAL_VIEW_MAP_CAPACITY; + use std::sync::Arc; + + /// The batch size the aggregate stream drives an ungrouped accumulator with. + const BATCH_SIZE: usize = 8192; + + /// Distinct value counts spanning both sides of the warm up capacity, up to + /// ones where the two constructors have converged. + const CARDINALITIES: [usize; 7] = [0, 1, 100, 1_000, 10_000, 100_000, 500_000]; + + /// Cardinalities small enough that the warm up dominates what the set + /// holds. This is the per group population, where `GroupsAccumulatorAdapter` + /// holds one accumulator per group and most see a handful of values. + const PER_GROUP_SCALE: usize = 100; + + /// Cardinalities at which a lazily built set has grown into exactly the + /// capacities a pre-allocated one reaches. This is the ungrouped + /// population, one set holding every distinct value in the input. + const UNGROUPED_SCALE: usize = 10_000; + + /// Longer than the map's inline value length, so the value lands in the + /// value buffer rather than inside the hash table entry. + fn distinct_value(i: usize) -> String { + format!("distinct value number {i}") + } + + fn batches(distinct_values: usize, view: bool) -> Vec { + (0..distinct_values) + .step_by(BATCH_SIZE) + .map(|start| { + let values = (start..(start + BATCH_SIZE).min(distinct_values)) + .map(distinct_value); + if view { + Arc::new(StringViewArray::from_iter_values(values)) as ArrayRef + } else { + Arc::new(StringArray::from_iter_values(values)) as ArrayRef + } + }) + .collect() + } + + /// The property that decides whether the ungrouped path can afford to share + /// the lazy constructor with the per group path. + fn assert_lazy_is_not_worse( + distinct_values: usize, + lazy_size: usize, + pre_allocated_size: usize, + ) { + // The guarantee that lets both paths share one lazy constructor. + assert!( + lazy_size <= pre_allocated_size, + "at {distinct_values} distinct values the lazy set reported \ + {lazy_size} bytes against the {pre_allocated_size} bytes the \ + pre-allocated one reported" + ); + + if distinct_values <= PER_GROUP_SCALE { + // The warm up is pure overhead here, and removing it is the whole + // point of the change. + assert!( + lazy_size < pre_allocated_size, + "at {distinct_values} distinct values the lazy set should be \ + strictly cheaper, but reported {lazy_size} bytes against \ + {pre_allocated_size}" + ); + } + + if distinct_values >= UNGROUPED_SCALE { + // The hash table's bucket count is a power of two fixed by the + // number of entries, and the value buffer grows on a power of two + // ladder, so a set that starts empty lands on exactly the + // capacities a pre-allocated one reaches. The ungrouped path gives + // up nothing by starting empty, which is why these accumulators + // need no signal telling them apart from the per group ones. + assert_eq!( + lazy_size, pre_allocated_size, + "at {distinct_values} distinct values the lazy and pre-allocated \ + sets should have converged" + ); + } + } + + /// An ungrouped `COUNT(DISTINCT )` builds a single accumulator that + /// grows to hold every distinct value in its input, which is the population + /// the warm up existed for. It must be no more expensive without one. + #[test] + fn ungrouped_utf8_accumulator_is_never_worse_than_a_pre_allocated_set() { + for distinct_values in CARDINALITIES { + let mut accumulator = + BytesDistinctCountAccumulator::::new(OutputType::Utf8); + let mut pre_allocated = ArrowBytesSet::::with_capacity( + OutputType::Utf8, + INITIAL_MAP_CAPACITY, + ); + + for batch in batches(distinct_values, false) { + accumulator.update_batch(&[Arc::clone(&batch)]).unwrap(); + pre_allocated.insert(&batch); + } + + assert_eq!(accumulator.0.non_null_len(), distinct_values); + assert_lazy_is_not_worse( + distinct_values, + accumulator.0.size(), + pre_allocated.size(), + ); + } + } + + /// The `Utf8View` counterpart of + /// [`ungrouped_utf8_accumulator_is_never_worse_than_a_pre_allocated_set`]. + #[test] + fn ungrouped_utf8_view_accumulator_is_never_worse_than_a_pre_allocated_set() { + for distinct_values in CARDINALITIES { + let mut accumulator = + BytesViewDistinctCountAccumulator::new(OutputType::Utf8View); + let mut pre_allocated = ArrowBytesViewSet::with_capacity( + OutputType::Utf8View, + INITIAL_VIEW_MAP_CAPACITY, + ); + + for batch in batches(distinct_values, true) { + accumulator.update_batch(&[Arc::clone(&batch)]).unwrap(); + pre_allocated.insert(&batch); + } + + assert_eq!(accumulator.0.non_null_len(), distinct_values); + assert_lazy_is_not_worse( + distinct_values, + accumulator.0.size(), + pre_allocated.size(), + ); + } + } +} diff --git a/datafusion/physical-expr-common/src/binary_map.rs b/datafusion/physical-expr-common/src/binary_map.rs index bc7d987173770..4028520c776d0 100644 --- a/datafusion/physical-expr-common/src/binary_map.rs +++ b/datafusion/physical-expr-common/src/binary_map.rs @@ -259,6 +259,28 @@ pub const INITIAL_MAP_CAPACITY: usize = 128; /// The size, in bytes, of the string data buffer pre-allocated by /// [`ArrowBytesMap::with_capacity`] pub const INITIAL_BUFFER_CAPACITY: usize = 8 * 1024; + +/// Appends `value` to a map's value buffer, growing the buffer on a power of +/// two ladder. +/// +/// `Vec` doubles from wherever its first allocation landed, so a buffer started +/// empty by [`ArrowBytesMap::new`] and one started at +/// [`INITIAL_BUFFER_CAPACITY`] by [`ArrowBytesMap::with_capacity`] sit on +/// different ladders, and can hold the same values at capacities differing by +/// up to 2x in either direction depending on the value lengths. Rounding every +/// growth up to a power of two puts both on one ladder, which is what makes a +/// lazily allocated map never larger than a pre-allocated one holding the same +/// values. Growth stays geometric, so appending is still amortized constant +/// time. +fn push_value_bytes(buffer: &mut Vec, value: &[u8]) { + let required = buffer.len() + value.len(); + if required > buffer.capacity() { + let target = required.checked_next_power_of_two().unwrap_or(required); + buffer.reserve_exact(target - buffer.len()); + } + buffer.extend_from_slice(value); +} + impl ArrowBytesMap where V: Debug + PartialEq + Eq + Clone + Copy + Default, @@ -471,7 +493,7 @@ where // Put the small values into buffer and offsets so it appears // the output array, but store the actual bytes inline for // comparison - self.buffer.extend_from_slice(value); + push_value_bytes(&mut self.buffer, value); self.offsets.push(O::usize_as(self.buffer.len())); let payload = make_payload_fn(Some(value)); let new_header = Entry { @@ -509,7 +531,7 @@ where // appears the output array, and store that offset // so the bytes can be compared if needed let offset = self.buffer.len(); // offset of start for data - self.buffer.extend_from_slice(value); + push_value_bytes(&mut self.buffer, value); self.offsets.push(O::usize_as(self.buffer.len())); let payload = make_payload_fn(Some(value)); @@ -826,6 +848,56 @@ mod tests { assert_eq!(map.buffer.capacity(), INITIAL_BUFFER_CAPACITY); } + #[test] + fn lazy_and_pre_allocated_buffers_grow_on_the_same_ladder() { + // Value lengths chosen so the buffer requirement lands between powers + // of two, which is where the two ladders used to diverge. + for value_len in [9usize, 13, 24, 37] { + let mut lazy = ArrowBytesMap::::new(OutputType::Utf8); + let mut pre_allocated = ArrowBytesMap::::with_capacity( + OutputType::Utf8, + INITIAL_MAP_CAPACITY, + ); + + for batch in 0..8 { + let values: ArrayRef = + Arc::new(StringArray::from_iter_values((0..1_000).map(|i| { + let value = format!("{}:{i}", batch * 1_000 + i); + format!("{value:value_len$}") + }))); + lazy.insert_if_new(&values, |_| (), |_| ()); + pre_allocated.insert_if_new(&values, |_| (), |_| ()); + + assert_eq!(lazy.buffer.len(), pre_allocated.buffer.len()); + assert_eq!( + lazy.buffer.capacity(), + pre_allocated.buffer.capacity(), + "value length {value_len}, batch {batch}: a buffer that \ + started empty reached {} bytes of capacity against {} for \ + one that started at INITIAL_BUFFER_CAPACITY", + lazy.buffer.capacity(), + pre_allocated.buffer.capacity(), + ); + } + } + } + + #[test] + fn a_lazy_buffer_stays_below_the_pre_allocated_floor_while_it_is_small() { + let mut lazy = ArrowBytesMap::::new(OutputType::Utf8); + let values: ArrayRef = Arc::new(StringArray::from_iter_values( + (0..10).map(|i| format!("distinct value number {i}")), + )); + lazy.insert_if_new(&values, |_| (), |_| ()); + + assert!( + lazy.buffer.capacity() < INITIAL_BUFFER_CAPACITY, + "expected a small lazy buffer to stay under the {INITIAL_BUFFER_CAPACITY} \ + byte pre-allocation, got {}", + lazy.buffer.capacity() + ); + } + #[test] fn string_set_empty() { let mut set = ArrowBytesSet::::new(OutputType::Utf8);