sunchao commented on code in PR #4791:
URL: https://github.com/apache/datafusion-comet/pull/4791#discussion_r4104772666
##########
native/core/src/execution/planner.rs:
##########
@@ -1325,17 +1350,39 @@ impl PhysicalPlanner {
let (scans, shuffle_scans, child) =
self.create_plan(&children[0], inputs, partition_count)?;
- let group_exprs: PhyExprResult = agg
+ // When `spark.comet.exec.aggregation.useLargeDataTypes` is
on, wrap Utf8/Binary
+ // group keys in a Cast to LargeUtf8/LargeBinary so DataFusion
dispatches
+ // to ByteGroupValueBuilder::<i64> and the per-task group-key
byte buffer
+ // is no longer capped at i32::MAX (2 GiB). The promotion is
reverted at
+ // the aggregate's output by a Projection below, so
LargeUtf8/LargeBinary
+ // never leaves this operator -- keeps the FFI, JVM shuffle,
and Spark
+ // consumer paths untouched.
+ let use_large = agg.use_large_data_types;
+ let child_schema_ref = child.schema();
+ let child_schema = child_schema_ref.as_ref();
+ // Per group column: `Some(original_dt)` when the column was
promoted to a
+ // Large* variant, `None` when it was passed through as-is.
Populated in
+ // lockstep with `group_exprs` and consumed below to build the
revert
+ // projection.
+ let mut group_reverts: Vec<Option<DataType>> =
+ Vec::with_capacity(agg.grouping_exprs.len());
+ let group_exprs: Vec<(Arc<dyn PhysicalExpr>, String)> = agg
.grouping_exprs
.iter()
.enumerate()
.map(|(idx, expr)| {
- self.create_expr(expr, child.schema())
- .map(|r| (r, format!("col_{idx}")))
+ let raw = self.create_expr(expr,
Arc::clone(&child_schema_ref))?;
+ let (wrapped, revert) = if use_large {
+ promote_byte_group_key(raw, child_schema)?
Review Comment:
[P1] Keep the final aggregate’s input and spill schemas consistent with
promoted keys. This promotion also runs in `Final` mode, while its child still
produces `Utf8`/`Binary`. DataFusion 55.1.0 uses `agg.input().schema()` as the
final aggregate’s state schema. When memory pressure triggers spilling, it
stamps that small-offset schema onto promoted group arrays and fails with
`expected Utf8 but found LargeUtf8` (likewise for binary). The output alignment
cannot fix an error inside the aggregate. With the default enabled, ordinary
spilling `GROUP BY`/`DISTINCT` queries now fail. Promote the final input
through a matching projection, or otherwise make its spill state schema agree
with the promoted keys.
Evidence: Bounded reproduction in `/tmp/comet-4791-validation/src/main.rs`:
native Final distinct aggregation over 80 batches of 256 rows, 10,000 distinct
128-byte keys, batch size 256, and a 1 MiB FairSpillPool. Without promotion,
both Utf8 and Binary return 10,000 keys after five spills. With the PR’s
grouping-expression promotion, both fail before returning rows with the
corresponding small/large type mismatch. Promoting the input through
ProjectionExec restores 10,000 results and five spills. DataFusion’s
`aggregate_hash_table/final_table.rs` takes the state schema from
`agg.input().schema()`, and `common.rs::take_state_batch` constructs the
failing batch.
##########
native/core/src/execution/planner.rs:
##########
@@ -1451,6 +1498,36 @@ impl PhysicalPlanner {
)?,
);
+ // Cast promoted group columns back to their original
Utf8/Binary type
+ // so LargeUtf8/LargeBinary never crosses the FFI boundary,
JVM columnar
+ // shuffle, or the Spark consumer path (all of which reject
Large*).
+ // Uses `SchemaAlignExec` (not a plain `ProjectionExec` +
`CastExpr`)
+ // because arrow's cast kernel rejects any single Large* array
whose
+ // value bytes exceed `i32::MAX` -- which is precisely the
regime this
+ // flag is used in. `SchemaAlignExec` splits each batch by row
ranges
+ // so every emitted small-offset chunk fits under 2 GiB.
+ let aggregate = if use_large &&
group_reverts.iter().any(Option::is_some) {
+ let agg_schema = aggregate.schema();
+ let target_fields: Vec<Field> = agg_schema
+ .fields()
+ .iter()
+ .enumerate()
+ .map(|(idx, f)| {
+ let target_dt = group_reverts
+ .get(idx)
+ .and_then(|r| r.clone())
+ .unwrap_or_else(|| f.data_type().clone());
+ Field::new(f.name(), target_dt, f.is_nullable())
+ .with_metadata(f.metadata().clone())
+ })
+ .collect();
+ let target_schema: SchemaRef =
Arc::new(Schema::new(target_fields));
+ SchemaAlignExec::try_new_or_passthrough(aggregate,
&target_schema)
Review Comment:
[P2] Preserve aggregate metrics when installing this wrapper.
`SchemaAlignExec` has no `metrics()` implementation, but it replaces
`AggregateExec` as the native plan passed to `SparkPlan::new`.
`to_native_metric_node` consequently reads no metrics and never visits the
wrapped aggregate. Every promoted string/binary aggregation loses output-row,
execution-time, and spill reporting, including the spill information propagated
into Spark task metrics. Forward the aggregate’s metrics through the wrapper or
explicitly retain it as the metric source.
Evidence: The exact-source harness executes string and binary aggregates
containing duplicates, nulls, and empty keys. Both return four groups, and the
underlying AggregateExec reports `output_rows=4` and nonzero `elapsed_compute`,
while the wrapper returns `metrics=None`. `SparkPlan::new` leaves
`additional_native_plans` empty, and
`native/core/src/execution/metrics/utils.rs::to_native_metric_node` only reads
the root metrics in that case.
##########
native/shuffle/src/schema_align.rs:
##########
@@ -228,27 +261,241 @@ impl ExecutionPlan for SchemaAlignExec {
struct SchemaAlignStream {
child_stream: SendableRecordBatchStream,
target_schema: SchemaRef,
+ column_actions: Arc<Vec<ColumnAction>>,
+ /// Sub-batches produced by the last input batch and not yet yielded. Used
when a
+ /// `CastLargeStringToString` / `CastLargeBinaryToBinary` column would
overflow the
+ /// destination `i32` offsets, so the input is split into multiple
Utf8/Binary outputs.
+ pending: VecDeque<RecordBatch>,
}
+/// `i32::MAX` bytes — the cap on a Utf8/Binary values buffer (its offsets are
`i32`).
+const I32_BYTE_CAP: i64 = i32::MAX as i64;
+
impl SchemaAlignStream {
- fn align(&self, batch: RecordBatch) -> Result<RecordBatch,
DataFusionError> {
+ /// Apply the per-column actions to `batch` and push the resulting
(possibly multiple)
+ /// aligned batches into `out`. Splits the input by row ranges when any
+ /// `CastLargeStringToString` / `CastLargeBinaryToBinary` column would
otherwise emit a
+ /// values buffer larger than `i32::MAX`.
+ fn align_into(
+ &self,
+ batch: RecordBatch,
+ out: &mut VecDeque<RecordBatch>,
+ ) -> Result<(), DataFusionError> {
+ let ranges = self.compute_row_ranges(&batch)?;
+ for (start, length) in ranges {
+ let slice = if start == 0 && length == batch.num_rows() {
+ batch.clone()
+ } else {
+ batch.slice(start, length)
+ };
+ out.push_back(self.align_slice(slice)?);
+ }
+ Ok(())
+ }
+
+ /// Apply `column_actions` to a single row range that is already known to
fit each
+ /// shrinking-cast column's destination offset width.
+ fn align_slice(&self, batch: RecordBatch) -> Result<RecordBatch,
DataFusionError> {
+ let mut columns: Vec<ArrayRef> =
Vec::with_capacity(batch.num_columns());
+ for (idx, action) in self.column_actions.iter().enumerate() {
+ let column = batch.column(idx);
+ let aligned = match action {
+ ColumnAction::Passthrough => Arc::clone(column),
+ ColumnAction::Cast => cast_with_options(
+ column,
+ self.target_schema.field(idx).data_type(),
+ &CastOptions::default(),
+ )?,
+ // Build a fresh Utf8/Binary array from the slice rather than
calling
+ // arrow's cast kernel. `cast_byte_container` reads the
underlying
+ // offsets buffer in full and verifies every absolute offset
fits the
+ // destination offset type — slicing the source array does not
rebase
+ // the offsets, so a slice that is logically small can still
trip the
+ // i32::MAX check if its offsets sit far into the values
buffer. We
+ // copy values explicitly so the new offsets start at 0.
+ ColumnAction::CastLargeStringToString => {
+ let arr = column
+ .as_any()
+ .downcast_ref::<LargeStringArray>()
+ .ok_or_else(|| {
+ DataFusionError::Internal(format!(
+ "SchemaAlignExec: column[{idx}] expected
LargeStringArray, \
+ got {:?}",
+ column.data_type()
+ ))
+ })?;
+ // Pre-size the values buffer with the exact byte total
for this slice so
+ // the underlying Vec never has to grow-and-memcpy while
we replay rows.
+ let offsets = arr.value_offsets();
+ let values_bytes = (offsets[arr.len()] - offsets[0]) as
usize;
+ let mut builder = StringBuilder::with_capacity(arr.len(),
values_bytes);
Review Comment:
[P2] Avoid copying value buffers when the offsets already fit. This builder
path runs for every promoted string key, with the equivalent copy for binary,
even when no splitting or rebasing is needed. Since promotion defaults to
enabled, each partial/final aggregation adds a full allocation and copy of its
emitted key bytes. Arrow can narrow ordinary batches while sharing those bytes.
Use that fast path when offsets fit, reserving rebuilding for overflow slices,
or rebase offsets against a shared value-buffer slice.
Evidence: The exact SchemaAlignExec source copies a 4 MiB value buffer for
8,192 keys of 512 bytes, verified by buffer-pointer comparison. Arrow 59.3.0’s
cast shares it and produces identical values. An optimized benchmark of the
PR’s exact builder branch measured median conversion times of 168.1 µs versus
8.4 µs for that batch, and 7.81 ms versus 8.3 µs for 8,192 keys of 8 KiB. These
are narrowing-only measurements, not whole-query timings. Reproduction:
`/tmp/comet-4791-cast-bench/src/main.rs`.
--
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]
