jayzhan211 opened a new issue, #25199: URL: https://github.com/apache/datafusion/issues/25199
### Is your feature request related to a problem or challenge? Follow-up from #24392, which rewrote `OrderSensitiveArrayAggAccumulator` to retain payloads as Arrow arrays instead of `ScalarValue`s. That is a large win when `update_batch` is called with reasonably sized batches, but the accumulator pushes one `ArrayRef` per *call*, so the fixed per-array cost (`ArrayData` + a 64-byte-minimum buffer + the `Arc`) is paid per call rather than per row: https://github.com/apache/datafusion/blob/main/datafusion/functions-aggregate/src/array_agg.rs#L1488 Ordered `ARRAY_AGG` cannot use a `GroupsAccumulator` (`groups_accumulator_supported` requires `order_bys.is_empty()`), so grouped queries go through `GroupsAccumulatorAdapter`, which calls `update_batch` once per group per input batch. For a high-cardinality `GROUP BY` that is typically 1–2 rows per call, which is the worst case for this layout. Measured with the accumulator's own `size()`, 2048 rows, `Int64` payload + `Int64` ordering key: | rows per `update_batch` | after #24392 | before #24392 | |---|---|---| | 1 | 217 B/row | 152 B/row | | 8 | 56 B/row | 152 B/row | | 64 | 43 B/row | 152 B/row | So `GROUP BY high_cardinality_col` with `array_agg(x ORDER BY y)` — roughly the shape of #20788 — uses about 1.4x more memory per row than it did before, while everything else got 3–4x better. ### Describe the solution you'd like Coalesce small inputs instead of retaining each as its own batch: below a row threshold, `concat` into the tail batch and rewrite the affected `entries` rather than pushing a new `ArrayRef`. That keeps the large-batch win and removes the per-call fixed cost. Suggested fix: coalesce small inputs rather than retaining each as its own batch, e.g. ```rust const COALESCE_ROWS: usize = 64; // in store_batch, after compaction: let row_count = values.len(); // ... let start = self.entries.len(); let batch_idx = match self.batches.last() { Some(last) if last.len() + row_count <= COALESCE_ROWS => { let merged = arrow::compute::concat(&[last.as_ref(), values.as_ref()])?; let idx = self.batches.len() - 1; // rows already recorded for this batch keep their row_idx; the new rows // start at the old length let offset = self.batches[idx].len(); self.batches[idx] = merged; self.entries.extend( (0..row_count).map(|row_idx| OrderedArrayAggEntry { batch_idx: idx, row_idx: offset + row_idx, }), ); return Ok(Some(start..self.entries.len())); } _ => { self.batches.push(values); self.batches.len() - 1 } }; ``` ### Describe alternatives you've considered - A real `GroupsAccumulator` for ordered `ARRAY_AGG`, which would side-step `GroupsAccumulatorAdapter`'s per-group slicing entirely. Bigger change, probably the right long-term answer. - Leaving it as is — the regression is confined to very small `update_batch` calls. ### Additional context `benches/array_agg.rs` currently only exercises the unordered `ArrayAggAccumulator::merge_batch`, so neither the win nor this regression is visible in CI benchmarks. A bench over the ordered path covering both regimes (few rows per call vs. full batches, pre-ordered vs. random input) would be worth adding alongside any fix. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
