This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-24869-3767e88ec4499d3f23971813082886f00c8dfe33 in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit 4a1fe81dba3f0e6cec3fa78292d545159b42019a Author: linfeng <[email protected]> AuthorDate: Thu Sep 10 09:12:29 2026 +0000 perf: compact inputs for grouped aggregate filter (#24869) ## Which issue does this PR close? <!-- We generally require a GitHub issue to be filed for all bug fixes and enhancements and this helps us generate change logs for our releases. You can link an issue to this PR using the GitHub syntax. For example `Closes #123` indicates that this PR will close issue #123. --> - Closes #. ## Rationale for this change <!-- Why are you proposing this change? If this is already explained clearly in the issue then this section is not needed. Explaining clearly why changes are proposed helps reviewers understand your changes and offer better suggestions for fixes. Please explain the problem you are trying to solve in terms of the user-visible behavior, rather than the implementation. For example, "The code in `foo.rs` doesn't handle nulls" is a symptom of the implementation. "COUNT(DISTINCT) returns wrong results when the column contains nulls" is the user-visible problem. --> Grouped aggregate `FILTER` currently scatters evaluated arguments back to their original row positions. The accumulator then receives the full-length arguments and group IDs and scans the filter again. This introduces unnecessary scatter and filtering work. ## What changes are included in this PR? <!-- There is no need to duplicate the description in the issue here, but it is sometimes worth providing a summary of the individual changes in this PR. --> - For regular `Partial` and `Single` aggregation, evaluate aggregate arguments using only rows selected by the filter. - Compact the corresponding group IDs before updating each accumulator. - Reuse the original group IDs when every row is selected. - Use bitmap indices or contiguous ranges depending on filter selectivity when compacting group IDs. - Preserve group registration using the complete input batch. - Handle sparse groups correctly in accumulator null-state tracking. - Preserve row-aligned argument evaluation for skip-partial aggregation. - Preserve ordered aggregation and empty-grouping-set behavior. ## What is the testing strategy for this PR? <!-- We typically require tests for all PRs in order to: 1. Prevent the code from being accidentally broken by subsequent changes 2. Serve as another way to document the expected behavior of the code Briefly describe how this PR is tested, and point to the specific tests you added. For example: 'This new feature is covered by the `sqllogictest` cases added in `foo.slt`'. If this PR does not add tests, explain why. For example, if the change is already covered by existing tests, please mention it. You should also check the `codecov` bot reply on this PR to confirm the changed code is exercised. --> Yes. slts cover independent filters, rejected groups, primitive and boolean group accumulators, ordered-input scenarios, and skip-partial aggregation. Unit tests cover group ID compaction and row-aligned skip-partial state conversion. ## Are there any user-facing changes? <!-- If there are user-facing changes then we may require documentation to be updated before approving the PR. If there are any breaking changes to public APIs, please add the `api change` label. --> No. This is an internal performance optimization. ## Benchmark ```text group main optimize ----- ---- -------- grouped_aggregate_filter/one_aggregate/1024_groups/column/100_percent 1.60 291.8±10.49µs ? ?/sec 1.00 182.2±1.95µs ? ?/sec grouped_aggregate_filter/one_aggregate/1024_groups/column/10_percent 1.74 343.9±60.21µs ? ?/sec 1.00 197.9±3.64µs ? ?/sec grouped_aggregate_filter/one_aggregate/1024_groups/column/1_percent 1.56 268.6±33.15µs ? ?/sec 1.00 172.0±3.28µs ? ?/sec grouped_aggregate_filter/one_aggregate/1024_groups/column/50_percent 1.77 639.4±108.17µs ? ?/sec 1.00 362.2±15.07µs ? ?/sec grouped_aggregate_filter/one_aggregate/1024_groups/column/90_percent 1.00 766.6±44.60µs ? ?/sec 1.01 774.5±15.30µs ? ?/sec grouped_aggregate_filter/one_aggregate/1024_groups/column/99_percent 1.13 361.9±9.36µs ? ?/sec 1.00 320.4±3.95µs ? ?/sec grouped_aggregate_filter/one_aggregate/1024_groups/column/unfiltered 1.05 195.7±10.26µs ? ?/sec 1.00 185.7±5.89µs ? ?/sec grouped_aggregate_filter/one_aggregate/1024_groups/multiply/100_percent 1.47 305.1±5.93µs ? ?/sec 1.00 207.1±3.59µs ? ?/sec grouped_aggregate_filter/one_aggregate/1024_groups/multiply/10_percent 1.74 342.4±43.21µs ? ?/sec 1.00 197.2±0.83µs ? ?/sec grouped_aggregate_filter/one_aggregate/1024_groups/multiply/1_percent 1.54 251.8±11.14µs ? ?/sec 1.00 163.9±2.20µs ? ?/sec grouped_aggregate_filter/one_aggregate/1024_groups/multiply/50_percent 1.77 631.5±22.55µs ? ?/sec 1.00 356.8±4.79µs ? ?/sec grouped_aggregate_filter/one_aggregate/1024_groups/multiply/90_percent 1.01 797.8±18.16µs ? ?/sec 1.00 786.3±31.99µs ? ?/sec grouped_aggregate_filter/one_aggregate/1024_groups/multiply/99_percent 1.00 401.4±15.31µs ? ?/sec 1.04 416.2±28.65µs ? ?/sec grouped_aggregate_filter/one_aggregate/1024_groups/multiply/unfiltered 1.02 217.0±20.33µs ? ?/sec 1.00 213.6±11.31µs ? ?/sec grouped_aggregate_filter/one_aggregate/16_groups/column/100_percent 1.61 266.6±7.40µs ? ?/sec 1.00 165.7±1.61µs ? ?/sec grouped_aggregate_filter/one_aggregate/16_groups/column/10_percent 1.52 296.1±7.73µs ? ?/sec 1.00 195.0±7.79µs ? ?/sec grouped_aggregate_filter/one_aggregate/16_groups/column/1_percent 1.36 220.2±3.91µs ? ?/sec 1.00 161.5±8.14µs ? ?/sec grouped_aggregate_filter/one_aggregate/16_groups/column/50_percent 1.62 573.8±13.79µs ? ?/sec 1.00 353.1±6.08µs ? ?/sec grouped_aggregate_filter/one_aggregate/16_groups/column/90_percent 1.10 705.5±10.46µs ? ?/sec 1.00 638.7±12.83µs ? ?/sec grouped_aggregate_filter/one_aggregate/16_groups/column/99_percent 1.33 348.3±47.72µs ? ?/sec 1.00 261.5±25.44µs ? ?/sec grouped_aggregate_filter/one_aggregate/16_groups/column/unfiltered 1.00 163.5±5.24µs ? ?/sec 1.08 177.4±12.99µs ? ?/sec grouped_aggregate_filter/one_aggregate/16_groups/multiply/100_percent 2.15 396.8±182.01µs ? ?/sec 1.00 184.3±2.48µs ? ?/sec grouped_aggregate_filter/one_aggregate/16_groups/multiply/10_percent 1.64 299.1±8.26µs ? ?/sec 1.00 182.9±1.63µs ? ?/sec grouped_aggregate_filter/one_aggregate/16_groups/multiply/1_percent 1.51 224.2±4.13µs ? ?/sec 1.00 148.5±1.32µs ? ?/sec grouped_aggregate_filter/one_aggregate/16_groups/multiply/50_percent 1.65 581.6±17.48µs ? ?/sec 1.00 352.3±8.53µs ? ?/sec grouped_aggregate_filter/one_aggregate/16_groups/multiply/90_percent 1.28 804.6±120.15µs ? ?/sec 1.00 626.6±31.91µs ? ?/sec grouped_aggregate_filter/one_aggregate/16_groups/multiply/99_percent 1.35 382.3±42.57µs ? ?/sec 1.00 283.2±14.58µs ? ?/sec grouped_aggregate_filter/one_aggregate/16_groups/multiply/unfiltered 1.02 188.6±19.17µs ? ?/sec 1.00 184.1±3.40µs ? ?/sec grouped_aggregate_filter/one_aggregate/65536_groups/column/100_percent 1.16 1341.0±53.20µs ? ?/sec 1.00 1152.0±66.17µs ? ?/sec grouped_aggregate_filter/one_aggregate/65536_groups/column/10_percent 1.18 1421.0±52.66µs ? ?/sec 1.00 1205.4±79.31µs ? ?/sec grouped_aggregate_filter/one_aggregate/65536_groups/column/1_percent 1.14 1333.0±48.42µs ? ?/sec 1.00 1171.5±57.48µs ? ?/sec grouped_aggregate_filter/one_aggregate/65536_groups/column/50_percent 1.31 1689.3±44.50µs ? ?/sec 1.00 1294.0±41.36µs ? ?/sec grouped_aggregate_filter/one_aggregate/65536_groups/column/90_percent 1.12 1865.2±90.42µs ? ?/sec 1.00 1663.3±49.81µs ? ?/sec grouped_aggregate_filter/one_aggregate/65536_groups/column/99_percent 1.12 1484.9±154.95µs ? ?/sec 1.00 1325.3±33.43µs ? ?/sec grouped_aggregate_filter/one_aggregate/65536_groups/column/unfiltered 1.08 1260.3±101.90µs ? ?/sec 1.00 1165.3±19.06µs ? ?/sec grouped_aggregate_filter/one_aggregate/65536_groups/multiply/100_percent 1.29 1435.1±158.66µs ? ?/sec 1.00 1110.8±72.64µs ? ?/sec grouped_aggregate_filter/one_aggregate/65536_groups/multiply/10_percent 1.24 1420.1±92.72µs ? ?/sec 1.00 1145.4±28.61µs ? ?/sec grouped_aggregate_filter/one_aggregate/65536_groups/multiply/1_percent 1.17 1310.0±62.02µs ? ?/sec 1.00 1115.0±45.38µs ? ?/sec grouped_aggregate_filter/one_aggregate/65536_groups/multiply/50_percent 1.29 1691.4±86.65µs ? ?/sec 1.00 1306.5±63.15µs ? ?/sec grouped_aggregate_filter/one_aggregate/65536_groups/multiply/90_percent 1.10 1850.0±76.61µs ? ?/sec 1.00 1674.2±55.65µs ? ?/sec grouped_aggregate_filter/one_aggregate/65536_groups/multiply/99_percent 1.17 1453.6±42.35µs ? ?/sec 1.00 1244.7±53.70µs ? ?/sec grouped_aggregate_filter/one_aggregate/65536_groups/multiply/unfiltered 1.05 1260.2±30.72µs ? ?/sec 1.00 1203.8±22.23µs ? ?/sec grouped_aggregate_filter/two_aggregates_distinct_filters/1024_groups/multiply/100_percent 1.82 441.4±14.01µs ? ?/sec 1.00 242.2±4.97µs ? ?/sec grouped_aggregate_filter/two_aggregates_distinct_filters/1024_groups/multiply/10_percent 1.90 490.2±9.58µs ? ?/sec 1.00 258.6±7.60µs ? ?/sec grouped_aggregate_filter/two_aggregates_distinct_filters/1024_groups/multiply/1_percent 1.78 350.7±11.14µs ? ?/sec 1.00 196.8±15.83µs ? ?/sec grouped_aggregate_filter/two_aggregates_distinct_filters/1024_groups/multiply/50_percent 1.91 1077.6±21.84µs ? ?/sec 1.00 563.5±27.23µs ? ?/sec grouped_aggregate_filter/two_aggregates_distinct_filters/1024_groups/multiply/90_percent 1.13 1458.9±80.33µs ? ?/sec 1.00 1290.9±37.66µs ? ?/sec grouped_aggregate_filter/two_aggregates_distinct_filters/1024_groups/multiply/99_percent 1.20 637.5±84.74µs ? ?/sec 1.00 530.8±15.89µs ? ?/sec grouped_aggregate_filter/two_aggregates_distinct_filters/1024_groups/multiply/unfiltered 1.01 260.0±20.62µs ? ?/sec 1.00 257.7±5.63µs ? ?/sec grouped_aggregate_filter/two_aggregates_distinct_filters/16_groups/multiply/100_percent 1.83 442.6±25.65µs ? ?/sec 1.00 242.0±18.15µs ? ?/sec grouped_aggregate_filter/two_aggregates_distinct_filters/16_groups/multiply/10_percent 2.06 491.8±20.88µs ? ?/sec 1.00 238.4±7.67µs ? ?/sec grouped_aggregate_filter/two_aggregates_distinct_filters/16_groups/multiply/1_percent 2.04 335.2±13.63µs ? ?/sec 1.00 164.4±3.26µs ? ?/sec grouped_aggregate_filter/two_aggregates_distinct_filters/16_groups/multiply/50_percent 1.87 1065.6±20.79µs ? ?/sec 1.00 569.3±17.13µs ? ?/sec grouped_aggregate_filter/two_aggregates_distinct_filters/16_groups/multiply/90_percent 1.09 1460.0±145.69µs ? ?/sec 1.00 1338.8±247.67µs ? ?/sec grouped_aggregate_filter/two_aggregates_distinct_filters/16_groups/multiply/99_percent 1.53 655.3±122.20µs ? ?/sec 1.00 428.4±80.90µs ? ?/sec grouped_aggregate_filter/two_aggregates_distinct_filters/16_groups/multiply/unfiltered 1.00 246.5±15.46µs ? ?/sec 1.04 256.1±49.50µs ? ?/sec grouped_aggregate_filter/two_aggregates_distinct_filters/65536_groups/multiply/100_percent 1.22 1479.8±45.71µs ? ?/sec 1.00 1214.9±43.68µs ? ?/sec grouped_aggregate_filter/two_aggregates_distinct_filters/65536_groups/multiply/10_percent 1.29 1582.4±189.79µs ? ?/sec 1.00 1222.4±47.32µs ? ?/sec grouped_aggregate_filter/two_aggregates_distinct_filters/65536_groups/multiply/1_percent 1.26 1378.2±89.04µs ? ?/sec 1.00 1097.3±34.69µs ? ?/sec grouped_aggregate_filter/two_aggregates_distinct_filters/65536_groups/multiply/50_percent 1.62 2.5±0.48ms ? ?/sec 1.00 1519.1±44.74µs ? ?/sec grouped_aggregate_filter/two_aggregates_distinct_filters/65536_groups/multiply/90_percent 1.12 2.5±0.10ms ? ?/sec 1.00 2.2±0.04ms ? ?/sec grouped_aggregate_filter/two_aggregates_distinct_filters/65536_groups/multiply/99_percent 1.27 1884.7±402.52µs ? ?/sec 1.00 1484.2±45.19µs ? ?/sec grouped_aggregate_filter/two_aggregates_distinct_filters/65536_groups/multiply/unfiltered 1.07 1304.4±101.10µs ? ?/sec 1.00 1217.9±53.10µs ? ?/sec grouped_aggregate_filter/two_aggregates_shared_filter/1024_groups/multiply/100_percent 1.96 469.2±45.79µs ? ?/sec 1.00 239.7±4.70µs ? ?/sec grouped_aggregate_filter/two_aggregates_shared_filter/1024_groups/multiply/10_percent 2.40 566.0±94.16µs ? ?/sec 1.00 236.3±5.20µs ? ?/sec grouped_aggregate_filter/two_aggregates_shared_filter/1024_groups/multiply/1_percent 2.02 370.4±44.87µs ? ?/sec 1.00 183.3±3.44µs ? ?/sec grouped_aggregate_filter/two_aggregates_shared_filter/1024_groups/multiply/50_percent 2.07 1103.2±99.60µs ? ?/sec 1.00 532.8±20.23µs ? ?/sec grouped_aggregate_filter/two_aggregates_shared_filter/1024_groups/multiply/90_percent 1.21 1411.1±119.95µs ? ?/sec 1.00 1165.6±38.46µs ? ?/sec grouped_aggregate_filter/two_aggregates_shared_filter/1024_groups/multiply/99_percent 1.36 720.9±194.27µs ? ?/sec 1.00 528.6±21.47µs ? ?/sec grouped_aggregate_filter/two_aggregates_shared_filter/1024_groups/multiply/unfiltered 1.03 245.7±3.38µs ? ?/sec 1.00 237.9±5.10µs ? ?/sec grouped_aggregate_filter/two_aggregates_shared_filter/16_groups/multiply/100_percent 1.85 424.4±20.65µs ? ?/sec 1.00 229.6±12.86µs ? ?/sec grouped_aggregate_filter/two_aggregates_shared_filter/16_groups/multiply/10_percent 2.19 495.7±75.88µs ? ?/sec 1.00 226.7±2.49µs ? ?/sec grouped_aggregate_filter/two_aggregates_shared_filter/16_groups/multiply/1_percent 2.10 340.5±45.25µs ? ?/sec 1.00 161.8±2.39µs ? ?/sec grouped_aggregate_filter/two_aggregates_shared_filter/16_groups/multiply/50_percent 2.00 1038.7±143.30µs ? ?/sec 1.00 519.7±9.86µs ? ?/sec grouped_aggregate_filter/two_aggregates_shared_filter/16_groups/multiply/90_percent 1.33 1290.4±48.84µs ? ?/sec 1.00 972.3±29.13µs ? ?/sec grouped_aggregate_filter/two_aggregates_shared_filter/16_groups/multiply/99_percent 1.80 665.6±231.71µs ? ?/sec 1.00 369.5±50.72µs ? ?/sec grouped_aggregate_filter/two_aggregates_shared_filter/16_groups/multiply/unfiltered 1.00 238.8±24.69µs ? ?/sec 1.00 239.1±19.73µs ? ?/sec grouped_aggregate_filter/two_aggregates_shared_filter/65536_groups/multiply/100_percent 1.24 1643.7±92.99µs ? ?/sec 1.00 1325.0±155.28µs ? ?/sec grouped_aggregate_filter/two_aggregates_shared_filter/65536_groups/multiply/10_percent 1.46 1765.9±118.19µs ? ?/sec 1.00 1211.7±56.11µs ? ?/sec grouped_aggregate_filter/two_aggregates_shared_filter/65536_groups/multiply/1_percent 1.36 1529.3±67.92µs ? ?/sec 1.00 1127.5±49.69µs ? ?/sec grouped_aggregate_filter/two_aggregates_shared_filter/65536_groups/multiply/50_percent 1.62 2.4±0.26ms ? ?/sec 1.00 1507.6±57.98µs ? ?/sec grouped_aggregate_filter/two_aggregates_shared_filter/65536_groups/multiply/90_percent 1.11 2.5±0.28ms ? ?/sec 1.00 2.3±0.07ms ? ?/sec grouped_aggregate_filter/two_aggregates_shared_filter/65536_groups/multiply/99_percent 1.10 1699.1±153.04µs ? ?/sec 1.00 1550.3±158.97µs ? ?/sec grouped_aggregate_filter/two_aggregates_shared_filter/65536_groups/multiply/unfiltered 1.19 1476.5±124.38µs ? ?/sec 1.00 1237.4±47.13µs ? ?/sec ``` --- .../src/aggregate/groups_accumulator/accumulate.rs | 57 ++-- .../src/aggregates/aggregate_hash_table/common.rs | 307 +++++++++++++++++---- .../aggregate_hash_table/common_ordered.rs | 2 +- .../aggregate_hash_table/partial_table.rs | 30 +- .../test_files/aggregate_filter_selection.slt | 101 +++++++ 5 files changed, 413 insertions(+), 84 deletions(-) diff --git a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/accumulate.rs b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/accumulate.rs index fb84d5882f..36ba0fba0b 100644 --- a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/accumulate.rs +++ b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/accumulate.rs @@ -31,11 +31,8 @@ use datafusion_expr_common::groups_accumulator::{EmitTo, GroupSelection}; /// handle each input null value specially (e.g. for `SUM` to mark the /// corresponding sum as null) /// -/// If there are filters present, `NullState` tracks if it has seen -/// *any* value for that group (as some values may be filtered -/// out). Without a filter, the accumulator is only passed groups that -/// had at least one value to accumulate so they do not need to track -/// if they have seen values for a particular group. +/// `NullState` tracks if it has seen *any* value for each group when filters or +/// sparse group indices may omit input for a registered group. #[derive(Debug)] pub enum SeenValues { /// All groups seen so far have seen at least one non-null value @@ -86,6 +83,31 @@ impl SeenValues { } } +/// Returns true when all newly registered groups are present in `group_indices`. +/// +/// Group indices are assigned in first-seen order, so an unfiltered batch visits +/// new groups in ascending order. Pre-filtered input can omit a new group, making +/// the indices sparse even though the accumulator no longer receives a filter. +fn new_groups_are_dense( + group_indices: &[usize], + first_new_group: usize, + total_num_groups: usize, +) -> bool { + if first_new_group == total_num_groups { + return true; + } + + let mut next_new_group = first_new_group; + for &group_index in group_indices { + if group_index == next_new_group { + next_new_group += 1; + } else if group_index > next_new_group { + return false; + } + } + next_new_group == total_num_groups +} + /// Track the accumulator null state per row: if any values for that /// group were null and if any values have been seen at all for that group. /// @@ -104,11 +126,8 @@ impl SeenValues { /// handle each input null value specially (e.g. for `SUM` to mark the /// corresponding sum as null) /// -/// If there are filters present, `NullState` tracks if it has seen -/// *any* value for that group (as some values may be filtered -/// out). Without a filter, the accumulator is only passed groups that -/// had at least one value to accumulate so they do not need to track -/// if they have seen values for a particular group. +/// `NullState` tracks if it has seen *any* value for each group when filters or +/// sparse group indices may omit input for a registered group. /// /// [`GroupsAccumulator`]: datafusion_expr_common::groups_accumulator::GroupsAccumulator #[derive(Debug)] @@ -173,10 +192,13 @@ impl NullState { T: ArrowPrimitiveType + Send, F: FnMut(usize, T::Native) + Send, { - // skip null handling if no nulls in input or accumulator - if let SeenValues::All { num_values } = &mut self.seen_values - && opt_filter.is_none() + // Skip per-value null handling when every input value is valid and all + // newly registered groups are represented. Pre-filtered inputs can have + // sparse group indices despite not passing a filter to the accumulator. + if opt_filter.is_none() && values.null_count() == 0 + && let SeenValues::All { num_values } = &mut self.seen_values + && new_groups_are_dense(group_indices, *num_values, total_num_groups) { accumulate(group_indices, values, None, value_fn); *num_values = total_num_groups; @@ -213,10 +235,13 @@ impl NullState { let data = values.values(); assert_eq!(data.len(), group_indices.len()); - // skip null handling if no nulls in input or accumulator - if let SeenValues::All { num_values } = &mut self.seen_values - && opt_filter.is_none() + // Skip per-value null handling when every input value is valid and all + // newly registered groups are represented. Pre-filtered inputs can have + // sparse group indices despite not passing a filter to the accumulator. + if opt_filter.is_none() && values.null_count() == 0 + && let SeenValues::All { num_values } = &mut self.seen_values + && new_groups_are_dense(group_indices, *num_values, total_num_groups) { group_indices .iter() diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs index 75468198f5..5f10195087 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs @@ -15,11 +15,15 @@ // specific language governing permissions and limitations // under the License. +use std::borrow::Cow; use std::collections::HashMap; use std::marker::PhantomData; use std::sync::Arc; -use arrow::array::{ArrayRef, AsArray, BooleanArray, new_null_array}; +use arrow::array::{ + Array, ArrayRef, AsArray, BooleanArray, new_empty_array, new_null_array, +}; +use arrow::compute::{filter_record_batch, prep_null_mask_filter}; use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; use datafusion_common::{Result, internal_err}; @@ -186,7 +190,7 @@ impl<AggrMode> AggregateHashTable<AggrMode> { .enumerate() .map(|(idx, acc)| { self.aggregate_argument_metrics - .time(idx, || acc.evaluate_acc_args(batch)) + .time(idx, || acc.evaluate_compacted_args(batch)) }) .collect::<Result<Vec<_>>>() })?; @@ -221,6 +225,8 @@ impl<AggrMode> AggregateHashTable<AggrMode> { .intern(group_values, &mut state.batch_group_indices) })?; + // Register groups from the full input. Each filtered aggregate compacts + // this row-aligned vector independently immediately before its update. let group_indices = &state.batch_group_indices; let total_num_groups = state.group_values.len(); group_by_metrics.time_aggregation(|| { @@ -400,9 +406,8 @@ impl<AggrMode> AggregateHashTable<AggrMode> { /// GROUP BY GROUPING SETS (()); /// ``` /// - /// The synthetic row is filtered out before accumulator update so aggregates - /// see the same state they would see for an empty input, rather than a real - /// null-valued row. + /// Accumulators receive zero argument rows and zero group IDs, together with the + /// full registered group count, so they produce the same state as empty input. /// /// Only the raw-input tables (partial and single aggregation) call this /// method: grouping sets are expanded while consuming raw rows, so the @@ -454,14 +459,13 @@ impl<AggrMode> AggregateHashTable<AggrMode> { if any_interned { let total_groups = state.group_values.len(); - let false_filter = BooleanArray::from(vec![false]); let values = state .accumulators .iter() .map(|acc| { - Ok(EvaluatedAccumulatorArgs { - arguments: acc.null_arguments(&self.input_schema)?, - filter: Some(Arc::new(false_filter.clone())), + Ok(CompactedAccumulatorArgs { + arguments: acc.null_arguments(&self.input_schema, 0)?, + selection: None, }) }) .collect::<Result<Vec<_>>>()?; @@ -470,7 +474,7 @@ impl<AggrMode> AggregateHashTable<AggrMode> { state.accumulators.iter_mut().zip(values.iter()).enumerate() { accumulator_metrics.time(idx, AccumulatorPhase::Update, || { - acc.update_batch(values, &[0], total_groups) + acc.update_batch(values, &[], total_groups) })?; } Ok::<(), datafusion_common::DataFusionError>(()) @@ -511,13 +515,13 @@ pub(super) type AggregateAccumulator = HashAggregateAccumulator; /// /// Arguments: /// * accumulator to update. -/// * accumulator's evaluated arguments and optional filter. +/// * accumulator's compacted arguments and optional row-aligned selection. /// * one group index per input row, mapping each row to its interned group. /// * total number of groups currently interned in that buffer, including newly /// interned groups. pub(super) type AggregateBatchFn = fn( &mut AggregateAccumulator, - &EvaluatedAccumulatorArgs, + &CompactedAccumulatorArgs, &[usize], usize, ) -> Result<()>; @@ -531,17 +535,21 @@ pub(super) type AggregateBatchFn = fn( pub(super) type MaterializeAccumulatorFn = fn(&mut AggregateAccumulator, EmitTo) -> Result<Vec<ArrayRef>>; -/// Evaluated aggregate arguments and filter for one input batch. -/// -/// For example, `AVG(x + 1) FILTER (WHERE x > 0)` evaluates both `x + 1` -/// and `x > 0`. -/// -/// These arrays can be passed directly to [`GroupsAccumulator`]. -pub(super) struct EvaluatedAccumulatorArgs { - /// Evaluated argument arrays. Some aggregate functions take multiple arguments. +/// Aggregate arguments compacted according to one aggregate's `FILTER`. +pub(super) struct CompactedAccumulatorArgs { + /// Argument arrays containing only selected rows. Some aggregate functions take + /// multiple arguments. pub(super) arguments: Vec<ArrayRef>, - /// Evaluated filter array, `Some` if the aggregate has a `FILTER` expression. - pub(super) filter: Option<ArrayRef>, + /// Original row-aligned selection used only to compact the matching group IDs. + pub(super) selection: Option<BooleanArray>, +} + +/// Evaluated aggregate arguments that preserve one output row per input row. +pub(super) struct RowAlignedAccumulatorArgs { + /// Row-aligned argument arrays. Rejected rows are represented as nulls. + pub(super) arguments: Vec<ArrayRef>, + /// Original row-aligned filter passed through to state conversion. + pub(super) filter: Option<BooleanArray>, } /// Evaluated all group by keys and accumulator args. @@ -553,8 +561,8 @@ pub(super) struct EvaluatedAggregateBatch { /// arrays for the current input batch. pub(super) grouping_set_args: Vec<Vec<ArrayRef>>, - /// Evaluated arguments and filters, one entry per aggregate expression. - pub(super) accumulator_args: Vec<EvaluatedAccumulatorArgs>, + /// Compacted arguments and selections, one entry per aggregate expression. + pub(super) accumulator_args: Vec<CompactedAccumulatorArgs>, } /// Buffer for the aggregate hash table's group keys and accumulator states. @@ -633,6 +641,42 @@ impl MaterializedAggregateOutput { } } +/// Compacts row-aligned group indices using an aggregate filter. +/// +/// Returns `None` when every row is selected so callers can reuse the input +/// slice without allocating. At high selectivity, copying contiguous selected +/// ranges avoids branching once per row. +fn compact_group_indices( + group_indices: &[usize], + filter: &BooleanArray, +) -> Option<Vec<usize>> { + debug_assert_eq!(group_indices.len(), filter.len()); + + let filter = match filter.null_count() { + 0 => Cow::Borrowed(filter), + _ => Cow::Owned(prep_null_mask_filter(filter)), + }; + let mask = filter.values(); + let selected_rows = mask.count_set_bits(); + + if selected_rows == group_indices.len() { + return None; + } + + let mut compacted = Vec::with_capacity(selected_rows); + // Match scatter's strategy: above 80% selectivity, copy contiguous ranges. + if selected_rows * 5 > group_indices.len() * 4 { + for (start, end) in mask.set_slices() { + compacted.extend_from_slice(&group_indices[start..end]); + } + } else { + compacted.extend(mask.set_indices().map(|index| group_indices[index])); + } + debug_assert_eq!(compacted.len(), selected_rows); + + Some(compacted) +} + impl HashAggregateAccumulator { pub(super) fn new( aggregate_expr: Arc<AggregateFunctionExpr>, @@ -663,25 +707,62 @@ impl HashAggregateAccumulator { /// Evaluate aggregate arguments and filter for one input batch. /// /// For example, `AVG(2 / x) FILTER (WHERE x > 0)` evaluates `x > 0` - /// first, then evaluates `2 / x` only for selected rows. - /// Filtered rows will be evaluated to `NULL`, and won't trigger errors - /// such as divide by zero. + /// first, then evaluates `2 / x` against a compact batch containing only + /// selected rows. Filtered rows won't trigger errors such as divide by zero. /// - /// These arrays can be passed directly to [`GroupsAccumulator`] next. - pub(super) fn evaluate_acc_args( + /// Before updating [`GroupsAccumulator`], the retained selection is used to + /// compact the matching group IDs and is not passed through. + pub(super) fn evaluate_compacted_args( &self, batch: &RecordBatch, - ) -> Result<EvaluatedAccumulatorArgs> { - let filter = self - .filter - .as_ref() - .map(|filter| { - filter - .evaluate(batch) - .and_then(|value| value.into_array(batch.num_rows())) + ) -> Result<CompactedAccumulatorArgs> { + let selection = self.evaluate_filter(batch)?; + let selected_rows = selection.as_ref().map(|selection| selection.true_count()); + let filtered_batch = match (selection.as_ref(), selected_rows) { + (Some(selection), Some(selected_rows)) + if selected_rows > 0 && selected_rows < batch.num_rows() => + { + Some(filter_record_batch(batch, selection)?) + } + _ => None, + }; + let argument_batch = match selected_rows { + None => Some(batch), + Some(0) => None, + Some(selected_rows) if selected_rows == batch.num_rows() => Some(batch), + Some(_) => filtered_batch.as_ref(), + }; + let arguments = self + .arguments + .iter() + .map(|expr| { + if let Some(argument_batch) = argument_batch { + expr.evaluate(argument_batch) + .and_then(|value| value.into_array(argument_batch.num_rows())) + } else { + let data_type = expr.data_type(batch.schema_ref().as_ref())?; + Ok(new_empty_array(&data_type)) + } }) - .transpose()?; - let selection = filter.as_ref().map(|filter| filter.as_boolean()); + .collect::<Result<_>>()?; + + Ok(CompactedAccumulatorArgs { + arguments, + selection, + }) + } + + /// Evaluates selected arguments while preserving the input batch row count. + /// + /// Skip-partial conversion produces one state row per input row, so rejected + /// rows remain as null argument values and the filter is passed to + /// [`GroupsAccumulator::convert_to_state`]. + pub(super) fn evaluate_row_aligned_args( + &self, + batch: &RecordBatch, + ) -> Result<RowAlignedAccumulatorArgs> { + let filter = self.evaluate_filter(batch)?; + let selection = filter.as_ref(); let arguments = self .arguments .iter() @@ -695,7 +776,19 @@ impl HashAggregateAccumulator { }) .collect::<Result<_>>()?; - Ok(EvaluatedAccumulatorArgs { arguments, filter }) + Ok(RowAlignedAccumulatorArgs { arguments, filter }) + } + + fn evaluate_filter(&self, batch: &RecordBatch) -> Result<Option<BooleanArray>> { + self.filter + .as_ref() + .map(|filter| { + filter + .evaluate(batch) + .and_then(|value| value.into_array(batch.num_rows())) + .map(|filter| filter.as_boolean().clone()) + }) + .transpose() } pub(super) fn size(&self) -> usize { @@ -704,26 +797,30 @@ impl HashAggregateAccumulator { pub(super) fn update_batch( &mut self, - values: &EvaluatedAccumulatorArgs, + values: &CompactedAccumulatorArgs, group_indices: &[usize], total_num_groups: usize, ) -> Result<()> { - let filter = values.filter.as_ref().map(|filter| filter.as_boolean()); + let filtered_group_indices = values + .selection + .as_ref() + .and_then(|selection| compact_group_indices(group_indices, selection)); + let group_indices = filtered_group_indices.as_deref().unwrap_or(group_indices); self.accumulator.update_batch( &values.arguments, group_indices, - filter, + None, total_num_groups, ) } pub(super) fn merge_batch( &mut self, - values: &EvaluatedAccumulatorArgs, + values: &CompactedAccumulatorArgs, group_indices: &[usize], total_num_groups: usize, ) -> Result<()> { - debug_assert!(values.filter.is_none()); + debug_assert!(values.selection.is_none()); self.accumulator .merge_batch(&values.arguments, group_indices, total_num_groups) } @@ -749,24 +846,25 @@ impl HashAggregateAccumulator { self.accumulator.state(emit_to) } + /// Converts evaluated row-aligned arguments directly to partial state. pub(super) fn convert_to_state( - &mut self, - values: &EvaluatedAccumulatorArgs, + &self, + values: &RowAlignedAccumulatorArgs, ) -> Result<Vec<ArrayRef>> { - let opt_filter = values.filter.as_ref().map(|filter| filter.as_boolean()); self.accumulator - .convert_to_state(&values.arguments, opt_filter) + .convert_to_state(&values.arguments, values.filter.as_ref()) } pub(super) fn null_arguments( &self, input_schema: &SchemaRef, + num_rows: usize, ) -> Result<Vec<ArrayRef>> { self.arguments .iter() .map(|expr| { let data_type = expr.data_type(input_schema)?; - Ok(new_null_array(&data_type, 1)) + Ok(new_null_array(&data_type, num_rows)) }) .collect() } @@ -792,10 +890,35 @@ impl AggregateHashTableState { mod tests { use std::sync::Arc; - use arrow::array::{Array, Int32Array}; + use arrow::array::{Array, BooleanArray, Int32Array, Int64Array}; use arrow::datatypes::{DataType, Field, Schema}; + use datafusion_functions_aggregate::sum::sum_udaf; + use datafusion_physical_expr::aggregate::AggregateExprBuilder; + use datafusion_physical_expr::expressions::Column; use super::*; + use crate::metrics::ExecutionPlanMetricsSet; + + #[test] + fn compact_group_indices_uses_filter_bitmap() { + let group_indices = (0..10).collect::<Vec<_>>(); + let all_true = BooleanArray::from(vec![true; 10]); + assert_eq!(compact_group_indices(&group_indices, &all_true), None); + + let high_selectivity = + BooleanArray::from((0..10).map(|index| index != 4).collect::<Vec<_>>()); + assert_eq!( + compact_group_indices(&group_indices, &high_selectivity), + Some(vec![0, 1, 2, 3, 5, 6, 7, 8, 9]) + ); + + let with_nulls = + BooleanArray::from(vec![Some(true), None, Some(false), Some(true), None]); + assert_eq!( + compact_group_indices(&group_indices[..5], &with_nulls), + Some(vec![0, 3]) + ); + } #[test] fn materialized_aggregate_output_slices_batches_until_exhausted() -> Result<()> { @@ -819,6 +942,88 @@ mod tests { Ok(()) } + #[test] + fn convert_to_state_preserves_rows_and_metrics() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("value", DataType::Int64, false), + Field::new("include", DataType::Boolean, false), + ])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int64Array::from(vec![10, 20, 30, 40])), + Arc::new(BooleanArray::from(vec![true, false, true, false])), + ], + )?; + let accumulator = sum_accumulator(&schema, "include", 1)?; + let metrics = ExecutionPlanMetricsSet::new(); + let group_by_metrics = GroupByMetrics::new(&metrics, 0); + let argument_metrics = AggregateArgumentMetrics::new(&metrics, 0, ["SUM(value)"]); + let accumulator_metrics = AggregateAccumulatorMetrics::new( + &metrics, + 0, + ["SUM(value)"], + &[AccumulatorPhase::ConvertToState], + ); + + let values = group_by_metrics.time_aggregate_arguments(|| { + argument_metrics.time(0, || accumulator.evaluate_row_aligned_args(&batch)) + })?; + let state = + accumulator_metrics.time(0, AccumulatorPhase::ConvertToState, || { + accumulator.convert_to_state(&values) + })?; + + assert_eq!( + int64_options(&state[0]), + vec![Some(10), None, Some(30), None] + ); + let metrics = metrics.clone_inner(); + for metric_name in [ + "aggregate_arguments_time", + "agg_expr_0_arguments_time", + "agg_expr_0_convert_to_state_time", + ] { + assert!( + metrics + .sum_by_name(metric_name) + .is_some_and(|time| { time.as_usize() > 0 }) + ); + } + + Ok(()) + } + + fn sum_accumulator( + schema: &SchemaRef, + filter_name: &str, + filter_index: usize, + ) -> Result<HashAggregateAccumulator> { + let argument: Arc<dyn PhysicalExpr> = Arc::new(Column::new("value", 0)); + let aggregate_expr = Arc::new( + AggregateExprBuilder::new(sum_udaf(), vec![Arc::clone(&argument)]) + .schema(Arc::clone(schema)) + .alias("SUM(value)") + .build()?, + ); + let accumulator = create_group_accumulator(&aggregate_expr)?; + Ok(HashAggregateAccumulator::new( + aggregate_expr, + vec![argument], + Some(Arc::new(Column::new(filter_name, filter_index))), + accumulator, + )) + } + + fn int64_options(array: &ArrayRef) -> Vec<Option<i64>> { + array + .as_any() + .downcast_ref::<Int64Array>() + .unwrap() + .iter() + .collect() + } + fn int32_values(batch: &RecordBatch, column: usize) -> Vec<i32> { let array = batch .column(column) diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs index 418f3f376b..97ef898b51 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs @@ -256,7 +256,7 @@ impl<AggrMode> OrderedAggregateTable<AggrMode> { .enumerate() .map(|(idx, acc)| { self.aggregate_argument_metrics - .time(idx, || acc.evaluate_acc_args(batch)) + .time(idx, || acc.evaluate_compacted_args(batch)) }) .collect::<Result<Vec<_>>>() })?; diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs index bbc51ae666..377abdae71 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs @@ -22,9 +22,9 @@ use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; use datafusion_common::{Result, assert_eq_or_internal_err}; -use crate::aggregates::AggregateExec; use crate::aggregates::group_values::{AccumulatorPhase, new_group_values}; use crate::aggregates::order::GroupOrdering; +use crate::aggregates::{AggregateExec, evaluate_group_by}; use super::common::{ AggregateHashTable, AggregateHashTableBuffer, AggregateHashTableState, @@ -130,31 +130,29 @@ impl AggregateHashTable<PartialSkipMarker> { &mut self, batch: &RecordBatch, ) -> Result<RecordBatch> { - let evaluated_batch = self.evaluate_batch(batch)?; + let state = self.state.building(); + let grouping_set_args = self + .group_by_metrics + .time_group_key_preparation(|| evaluate_group_by(&state.group_by, batch))?; assert_eq_or_internal_err!( - evaluated_batch.grouping_set_args.len(), + grouping_set_args.len(), 1, "group_values expected to have single element" ); - let mut output = evaluated_batch - .grouping_set_args - .into_iter() - .next() - .unwrap_or_default(); + let mut output = grouping_set_args.into_iter().next().unwrap_or_default(); let accumulator_metrics = Arc::clone(&self.aggregate_accumulator_metrics); - let state = self.state.building_mut(); - for (idx, (acc, values)) in state - .accumulators - .iter_mut() - .zip(evaluated_batch.accumulator_args.iter()) - .enumerate() - { + for (idx, acc) in state.accumulators.iter().enumerate() { + let values = self.group_by_metrics.time_aggregate_arguments(|| { + self.aggregate_argument_metrics + .time(idx, || acc.evaluate_row_aligned_args(batch)) + })?; + output.extend(accumulator_metrics.time( idx, AccumulatorPhase::ConvertToState, - || acc.convert_to_state(values), + || acc.convert_to_state(&values), )?); } diff --git a/datafusion/sqllogictest/test_files/aggregate_filter_selection.slt b/datafusion/sqllogictest/test_files/aggregate_filter_selection.slt index ca58378584..d8e4e97f5a 100644 --- a/datafusion/sqllogictest/test_files/aggregate_filter_selection.slt +++ b/datafusion/sqllogictest/test_files/aggregate_filter_selection.slt @@ -44,9 +44,110 @@ ORDER BY g; 3 NULL 4 2 +# Each aggregate must compact group indices using its own FILTER. Cover both +# primitive and boolean group accumulators, including groups rejected by both. +query IIB +SELECT g, + SUM(v) FILTER (WHERE v = 2), + BOOL_AND(v > 0) FILTER (WHERE v = 5) +FROM aggregate_filter_selection +GROUP BY g +ORDER BY g; +---- +1 NULL NULL +2 2 NULL +3 NULL NULL +4 NULL true + +# NULL and false predicates must reject rows before evaluating 10 / v. +# Interleaved groups keep argument and group-ID compaction observable; groups +# with no selected rows must still appear in the output. +statement ok +CREATE TABLE aggregate_filter_selection_nullable (g INT, v BIGINT, keep BOOLEAN); + +statement ok +INSERT INTO aggregate_filter_selection_nullable VALUES + (1, 0, NULL), + (2, 2, true), + (1, 5, true), + (3, 0, NULL), + (2, 0, false), + (4, 0, false); + +query III +SELECT g, SUM(10 / v) FILTER (WHERE keep), COUNT(*) FILTER (WHERE keep) +FROM aggregate_filter_selection_nullable +GROUP BY g +ORDER BY g; +---- +1 2 1 +2 5 1 +3 NULL 0 +4 NULL 0 + +statement ok +DROP TABLE aggregate_filter_selection_nullable; + +statement ok +SET datafusion.execution.target_partitions = 1; + +query I +COPY ( + SELECT * + FROM (VALUES + (1, 10, 0), + (1, 10, 2), + (1, 11, 5), + (2, 20, 0), + (2, 21, 4) + ) AS t(sort_col, group_col, v) + ORDER BY sort_col, group_col +) +TO 'test_files/scratch/aggregate_filter_selection/ordered.parquet' +STORED AS PARQUET; +---- +5 + +statement ok +CREATE EXTERNAL TABLE aggregate_filter_selection_ordered ( + sort_col INT, + group_col INT, + v BIGINT +) +STORED AS PARQUET +WITH ORDER (sort_col) +LOCATION 'test_files/scratch/aggregate_filter_selection'; + +# Ordered single aggregation must also apply FILTER before evaluating arguments. +query III +SELECT sort_col, group_col, SUM(10 / v) FILTER (WHERE v <> 0) +FROM aggregate_filter_selection_ordered +GROUP BY sort_col, group_col +ORDER BY sort_col, group_col; +---- +1 10 5 +1 11 2 +2 20 NULL +2 21 2 + statement ok SET datafusion.execution.target_partitions = 2; +# Repartitioning ordered input selects ordered partial aggregation. +query III +SELECT sort_col, group_col, SUM(10 / v) FILTER (WHERE v <> 0) +FROM aggregate_filter_selection_ordered +GROUP BY sort_col, group_col +ORDER BY sort_col, group_col; +---- +1 10 5 +1 11 2 +2 20 NULL +2 21 2 + +statement ok +DROP TABLE aggregate_filter_selection_ordered; + statement ok SET datafusion.execution.batch_size = 1; --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
