You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@datafusion.apache.org by "andygrove (via GitHub)" <gi...@apache.org> on 2026/08/14 21:18:18 UTC
[PR] fix: make CometExplodeExec respect batch size [datafusion-comet]
andygrove opened a new pull request, #5362:
URL: https://github.com/apache/datafusion-comet/pull/5362
## Which issue does this PR close?
Related to #5210, which tracks the unnest-related workarounds in the explode planner.
## Rationale for this change
`CometExplodeExec` plans onto DataFusion's `UnnestExec`, which emits exactly one output batch per input batch however many rows the unnesting produces. It never consults `datafusion.execution.batch_size` — there is no `batch_size` anywhere in upstream `unnest.rs`, and `UnnestStream` calls `build_batch` once per input batch and returns the whole result.
So an 8192-row batch of 100-element arrays comes back as a single 819,200-row batch, and input batch count always equals output batch count.
Two consequences:
1. Downstream operators receive arbitrarily large batches.
2. Peak memory scales with input batch size times array length rather than with `batch_size` — the full expansion of an input batch is materialized at once, which is an OOM risk on high-fanout explode.
## What changes are included in this PR?
Adds `ExplodeExec` (`native/core/src/execution/operators/explode.rs`), a temporary fork of `UnnestExec` that consumes each input batch in chunks, and points the planner at it instead of `UnnestExec`.
`find_longest_length` already computes how many output rows each input row expands into, with null handling accounted for. Prefix-summing it gives *exact* chunk boundaries for the depth-1 unnesting Comet plans, so each `build_batch` call produces at most `batch_size` rows and the oversized intermediate is never materialized. That is what bounds memory, not just output size.
Two cases can't be chunked on the input side and are handled by slicing the built batch instead:
- a single input row whose array is longer than `batch_size`, since one row is never split across output batches
- recursive unnest (`depth > 1`), which Comet does not plan today, but the fallback keeps the operator correct if that changes
### Why a fork
The unnesting kernels (`build_batch` and everything it calls) are private to `datafusion-physical-plan`, so they cannot be called from Comet without copying them. They are copied verbatim from DataFusion 54.1.0 (upstream revision `cc7565be1ee97ba8fa2f5d6da373c5e38d81bb13`) and marked as such, so the eventual deletion is mechanical. All Comet-specific behavior lives in `ExplodeExec` / `ExplodeStream` above the copied section.
Note that 54.1.0 predates upstream's `NullHandling` enum and still uses `UnnestOptions::preserve_nulls`, which is why the planner still wraps empty arrays with `ListEmptyToNullExpr` for `explode_outer` semantics. That is unchanged by this PR.
### Deleting this fork
The fix has been submitted upstream:
- apache/datafusion#24383 (issue)
- apache/datafusion#24384 (PR)
Once Comet moves to a DataFusion release carrying that PR, delete this module and go back to `UnnestExec`. The module docs say so, and #5210 already tracks the related cleanup.
I went with a fork rather than an opt-in upstream API because Comet is pinned to DataFusion 54.1.0, so anything landing in 55 or 56 does nothing here until the upgrade. A fork gets the fix now and deletes cleanly.
## Are these changes tested?
Yes — six new unit tests in the new module:
- `respects_batch_size` — 10 rows x 10 elements at `batch_size=8`, every output batch bounded, values and order preserved
- `chunks_input_rather_than_slicing_output` — pins *how* the limit is met: 3 rows of 3 elements at `batch_size=4` must give `[3, 3, 3]` (input chunked per row), not `[4, 4, 1]` (built whole, then sliced). These are indistinguishable by row counts alone but have very different memory profiles, so without this a refactor could silently regress to build-then-slice.
- `single_row_exceeding_batch_size_is_sliced` — the unavoidable case
- `chunking_preserves_outer_semantics` / `chunking_preserves_non_outer_semantics` — chunked vs unchunked output compared directly for both `preserve_nulls` settings
- `multiple_input_batches` — short tail chunk per input batch boundary
Full native suite: 163 passed, 0 failed. `cargo fmt` and `cargo clippy --all-targets -- -D warnings` clean.
JVM-side suite runs still in progress; I'll report results before marking this ready for review.
## Are there any user-facing changes?
Explode now produces more, smaller batches, bounded by `datafusion.execution.batch_size`, and uses substantially less memory on high-fanout arrays. Query results are unchanged.
`CometExplodeExec` now displays as `CometExplodeExec` rather than `UnnestExec` in native plan output.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix: make CometExplodeExec respect batch size [datafusion-comet]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove merged PR #5362:
URL: https://github.com/apache/datafusion-comet/pull/5362
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix: make CometExplodeExec respect batch size [datafusion-comet]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove commented on code in PR #5362:
URL: https://github.com/apache/datafusion-comet/pull/5362#discussion_r3799039298
##########
native/core/src/execution/operators/explode.rs:
##########
@@ -0,0 +1,1237 @@
+// 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.
+
+//! A temporary fork of DataFusion's `UnnestExec` that respects
+//! `datafusion.execution.batch_size`.
+//!
+//! # Why this fork exists
+//!
+//! DataFusion's `UnnestExec` emits exactly one output batch per input batch, however many
+//! rows the unnesting produces, and never consults `batch_size`. For `explode` this means
+//! an 8192-row batch of 100-element arrays comes back as a single 819,200-row batch, and
+//! peak memory scales with input batch size times array length rather than with
+//! `batch_size`.
+//!
+//! The fix has been submitted upstream:
+//!
+//! * <https://github.com/apache/datafusion/issues/24383>
+//! * <https://github.com/apache/datafusion/pull/24384>
+//!
+//! # Deleting this file
+//!
+//! Once Comet moves to a DataFusion release carrying apache/datafusion#24384, delete this
+//! module and go back to `datafusion::physical_plan::unnest::UnnestExec` in the planner.
+//!
+//! Note that <https://github.com/apache/datafusion-comet/issues/5210> is a *different*
+//! unnest cleanup — it tracks adopting upstream `unnest_outer`
+//! (apache/datafusion#22100) to retire `ListEmptyToNullExpr`. The two upstream PRs can
+//! land in different releases, so closing 5210 is not a signal to delete this fork.
+//!
+//! # What was forked
+//!
+//! The unnesting kernels below (`build_batch` and everything it calls) are copied from
+//! `datafusion/physical-plan/src/unnest.rs` at DataFusion 54.1.0, upstream revision
+//! `cc7565be1ee97ba8fa2f5d6da373c5e38d81bb13`. They are private to
+//! `datafusion-physical-plan`, so they cannot be called from here without copying them.
+//! Leave them semantically unmodified so the eventual deletion is mechanical; the
+//! Comet-specific behavior lives entirely in `ExplodeExec` and `ExplodeStream`.
+//!
+//! They are not byte-identical to upstream: Comet's rustfmt uses `max_width = 100` and
+//! edition 2021, DataFusion's uses `max_width = 90` and edition 2024, so `cargo fmt`
+//! reflows some signatures. To audit for real changes, reformat this region at
+//! `max_width = 90` and diff it against upstream `unnest.rs`; that reduces the difference
+//! to a single cosmetic line wrap in `flatten_struct_cols`. The only deliberate edits are
+//! the `lt` import path noted below and dropping upstream's `ListUnnest` declaration in
+//! favor of importing the public one.
+//!
+//! Note that 54.1.0 predates upstream's `NullHandling` enum and still uses
+//! `UnnestOptions::preserve_nulls`, which is why the planner wraps empty arrays with
+//! `ListEmptyToNullExpr` to get Spark's `explode_outer` semantics.
+
+use arrow::array::{
+ new_null_array, Array, ArrayRef, AsArray, FixedSizeListArray, Int64Array, LargeListArray,
+ LargeListViewArray, ListArray, ListViewArray, PrimitiveArray, Scalar, StructArray,
+};
+use arrow::compute::kernels::length::length;
+use arrow::compute::kernels::zip::zip;
+use arrow::compute::{cast, is_not_null, kernels, sum};
+use arrow::datatypes::{DataType, Int64Type, SchemaRef};
+use arrow::record_batch::RecordBatch;
+// Upstream imports this as `arrow_ord::cmp::lt`; Comet reaches it through `arrow`,
+// which does not have `arrow_ord` as a direct dependency.
+use arrow::compute::kernels::cmp::lt;
+use datafusion::common::{
+ exec_datafusion_err, exec_err, internal_err, HashMap, HashSet, Result, UnnestOptions,
+};
+use datafusion::execution::TaskContext;
+use datafusion::physical_expr::EquivalenceProperties;
+use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType};
+use datafusion::physical_plan::metrics::{
+ BaselineMetrics, Count, ExecutionPlanMetricsSet, MetricBuilder, MetricsSet, RecordOutput,
+};
+// `ListUnnest` is the one item the copied region below does NOT need to duplicate: unlike the
+// kernels, upstream exports it publicly.
+use datafusion::physical_plan::unnest::ListUnnest;
+use datafusion::physical_plan::{
+ DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties, RecordBatchStream,
+ SendableRecordBatchStream,
+};
+use futures::{Stream, StreamExt};
+use std::cmp::{self, Ordering};
+use std::pin::Pin;
+use std::sync::Arc;
+use std::task::{ready, Context, Poll};
+
+/// Comet's explode operator: DataFusion's `UnnestExec` with the input consumed in chunks so
+/// that output batches respect `datafusion.execution.batch_size`.
+#[derive(Debug)]
+pub struct ExplodeExec {
+ child: Arc<dyn ExecutionPlan>,
+ schema: SchemaRef,
+ list_column_indices: Vec<ListUnnest>,
+ struct_column_indices: Vec<usize>,
+ options: UnnestOptions,
+ metrics: ExecutionPlanMetricsSet,
+ cache: Arc<PlanProperties>,
+}
+
+impl ExplodeExec {
+ pub fn new(
+ child: Arc<dyn ExecutionPlan>,
+ list_column_indices: Vec<ListUnnest>,
+ struct_column_indices: Vec<usize>,
+ schema: SchemaRef,
+ options: UnnestOptions,
+ ) -> Self {
+ // Unnesting invalidates the child's orderings and constraints for the unnested
+ // columns, and Comet plans explode on a single partition, so start from empty
+ // equivalences rather than trying to project the child's.
+ let cache = Arc::new(PlanProperties::new(
+ EquivalenceProperties::new(Arc::clone(&schema)),
+ Partitioning::UnknownPartitioning(1),
+ EmissionType::Incremental,
+ Boundedness::Bounded,
+ ));
+
+ Self {
+ child,
+ schema,
+ list_column_indices,
+ struct_column_indices,
+ options,
+ metrics: ExecutionPlanMetricsSet::new(),
+ cache,
+ }
+ }
+}
+
+impl DisplayAs for ExplodeExec {
+ fn fmt_as(&self, t: DisplayFormatType, f: &mut std::fmt::Formatter) -> std::fmt::Result {
+ match t {
+ DisplayFormatType::Default | DisplayFormatType::Verbose => {
+ write!(f, "CometExplodeExec")
+ }
+ DisplayFormatType::TreeRender => unimplemented!(),
+ }
+ }
+}
+
+impl ExecutionPlan for ExplodeExec {
+ fn name(&self) -> &str {
+ "CometExplodeExec"
+ }
+
+ fn schema(&self) -> SchemaRef {
+ Arc::clone(&self.schema)
+ }
+
+ fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
+ vec![&self.child]
+ }
+
+ fn with_new_children(
+ self: Arc<Self>,
+ children: Vec<Arc<dyn ExecutionPlan>>,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ if children.len() != 1 {
+ return internal_err!("ExplodeExec expects exactly one child");
+ }
+ Ok(Arc::new(ExplodeExec::new(
+ Arc::clone(&children[0]),
+ self.list_column_indices.clone(),
+ self.struct_column_indices.clone(),
+ Arc::clone(&self.schema),
+ self.options.clone(),
+ )))
+ }
+
+ fn execute(
+ &self,
+ partition: usize,
+ context: Arc<TaskContext>,
+ ) -> Result<SendableRecordBatchStream> {
+ let batch_size = context.session_config().batch_size();
+ let input = self.child.execute(partition, context)?;
+
+ Ok(Box::pin(ExplodeStream {
+ input,
+ schema: Arc::clone(&self.schema),
+ list_type_columns: self.list_column_indices.clone(),
+ struct_column_indices: self.struct_column_indices.iter().copied().collect(),
+ options: self.options.clone(),
+ baseline_metrics: BaselineMetrics::new(&self.metrics, partition),
+ input_batches: MetricBuilder::new(&self.metrics).counter("input_batches", partition),
+ input_rows: MetricBuilder::new(&self.metrics).counter("input_rows", partition),
+ batch_size,
+ pending_input: None,
+ pending_output: None,
+ }))
+ }
+
+ fn properties(&self) -> &Arc<PlanProperties> {
+ &self.cache
+ }
+
+ fn metrics(&self) -> Option<MetricsSet> {
+ Some(self.metrics.clone_inner())
+ }
+}
+
+/// An input batch being unnested incrementally, a chunk of rows at a time.
+struct PendingInput {
+ /// The full input batch. Rows before `row_offset` have already been unnested.
+ batch: RecordBatch,
+ /// Index of the next input row to unnest.
+ row_offset: usize,
+ /// How many output rows each input row expands into, indexed by input row.
+ ///
+ /// Empty when the expansion factor cannot be predicted from the input alone, in which
+ /// case the whole remaining input is unnested in one call and only the output is split.
+ /// See [`ExplodeStream::predict_output_lens`].
+ output_lens: Vec<usize>,
+}
+
+impl PendingInput {
+ fn remaining_rows(&self) -> usize {
+ self.batch.num_rows() - self.row_offset
+ }
+
+ /// How many input rows to unnest next so the resulting batch holds at most `batch_size`
+ /// rows.
+ ///
+ /// Always returns at least 1 while rows remain, so the stream always makes progress: a
+ /// single input row is never split across output batches, so one row whose array is
+ /// longer than `batch_size` still produces one oversized build, which is then sliced
+ /// down on the way out.
+ fn next_chunk_rows(&self, batch_size: usize) -> usize {
+ let remaining = self.remaining_rows();
+ if self.output_lens.is_empty() {
+ return remaining;
+ }
+
+ let mut rows = 0;
+ let mut output_rows = 0usize;
+ while rows < remaining {
+ let len = self.output_lens[self.row_offset + rows];
+ if rows > 0 && output_rows.saturating_add(len) > batch_size {
+ break;
+ }
+ output_rows += len;
+ rows += 1;
+ }
+ rows
+ }
+}
+
+/// A stream that unnests its input, bounding output batches to `batch_size` rows.
+struct ExplodeStream {
+ input: SendableRecordBatchStream,
+ schema: SchemaRef,
+ list_type_columns: Vec<ListUnnest>,
+ struct_column_indices: HashSet<usize>,
+ options: UnnestOptions,
+ baseline_metrics: BaselineMetrics,
+ input_batches: Count,
+ input_rows: Count,
+ /// Target number of rows per output batch, from `datafusion.execution.batch_size`.
+ batch_size: usize,
+ /// Rows of the current input batch that have not been unnested yet.
+ pending_input: Option<PendingInput>,
+ /// Unnested rows that have been built but not emitted yet.
+ pending_output: Option<RecordBatch>,
+}
+
+impl RecordBatchStream for ExplodeStream {
+ fn schema(&self) -> SchemaRef {
+ Arc::clone(&self.schema)
+ }
+}
+
+impl Stream for ExplodeStream {
+ type Item = Result<RecordBatch>;
+
+ fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+ self.poll_next_impl(cx)
+ }
+}
+
+impl ExplodeStream {
+ fn poll_next_impl(&mut self, cx: &mut Context<'_>) -> Poll<Option<Result<RecordBatch>>> {
+ loop {
+ // Emit already-unnested rows first, at most `batch_size` at a time.
+ if let Some(batch) = self.pending_output.take() {
+ let (emit, rest) = split_off_head(batch, self.batch_size);
+ self.pending_output = rest;
+ (&emit).record_output(&self.baseline_metrics);
+ return Poll::Ready(Some(Ok(emit)));
+ }
+
+ // Unnest the next chunk of the input batch already in hand.
+ if let Some(pending) = self.pending_input.as_mut() {
+ if pending.remaining_rows() == 0 {
+ self.pending_input = None;
+ continue;
+ }
+
+ let timer = self.baseline_metrics.elapsed_compute().timer();
+
+ let rows = pending.next_chunk_rows(self.batch_size);
+ let chunk = pending.batch.slice(pending.row_offset, rows);
+ pending.row_offset += rows;
Review Comment:
Fixed. The loop now computes `drained` right after advancing `row_offset` and clears `pending_input` before the built batch is returned, so the source `RecordBatch`, the predicted lengths, and the `posexplode` position array are all released ahead of downstream and JVM processing rather than surviving until the next poll. That also let the `remaining_rows() == 0` re-entry branch go, since `PendingInput` is now only ever live with rows left in it — there is a `debug_assert` to that effect.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix: make CometExplodeExec respect batch size [datafusion-comet]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove commented on code in PR #5362:
URL: https://github.com/apache/datafusion-comet/pull/5362#discussion_r3799041110
##########
native/core/src/execution/operators/explode.rs:
##########
@@ -0,0 +1,1237 @@
+// 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.
+
+//! A temporary fork of DataFusion's `UnnestExec` that respects
+//! `datafusion.execution.batch_size`.
+//!
+//! # Why this fork exists
+//!
+//! DataFusion's `UnnestExec` emits exactly one output batch per input batch, however many
+//! rows the unnesting produces, and never consults `batch_size`. For `explode` this means
+//! an 8192-row batch of 100-element arrays comes back as a single 819,200-row batch, and
+//! peak memory scales with input batch size times array length rather than with
+//! `batch_size`.
+//!
+//! The fix has been submitted upstream:
+//!
+//! * <https://github.com/apache/datafusion/issues/24383>
+//! * <https://github.com/apache/datafusion/pull/24384>
+//!
+//! # Deleting this file
+//!
+//! Once Comet moves to a DataFusion release carrying apache/datafusion#24384, delete this
+//! module and go back to `datafusion::physical_plan::unnest::UnnestExec` in the planner.
+//!
+//! Note that <https://github.com/apache/datafusion-comet/issues/5210> is a *different*
+//! unnest cleanup — it tracks adopting upstream `unnest_outer`
+//! (apache/datafusion#22100) to retire `ListEmptyToNullExpr`. The two upstream PRs can
+//! land in different releases, so closing 5210 is not a signal to delete this fork.
+//!
+//! # What was forked
+//!
+//! The unnesting kernels below (`build_batch` and everything it calls) are copied from
+//! `datafusion/physical-plan/src/unnest.rs` at DataFusion 54.1.0, upstream revision
+//! `cc7565be1ee97ba8fa2f5d6da373c5e38d81bb13`. They are private to
+//! `datafusion-physical-plan`, so they cannot be called from here without copying them.
+//! Leave them semantically unmodified so the eventual deletion is mechanical; the
+//! Comet-specific behavior lives entirely in `ExplodeExec` and `ExplodeStream`.
+//!
+//! They are not byte-identical to upstream: Comet's rustfmt uses `max_width = 100` and
+//! edition 2021, DataFusion's uses `max_width = 90` and edition 2024, so `cargo fmt`
+//! reflows some signatures. To audit for real changes, reformat this region at
+//! `max_width = 90` and diff it against upstream `unnest.rs`; that reduces the difference
+//! to a single cosmetic line wrap in `flatten_struct_cols`. The only deliberate edits are
+//! the `lt` import path noted below and dropping upstream's `ListUnnest` declaration in
+//! favor of importing the public one.
+//!
+//! Note that 54.1.0 predates upstream's `NullHandling` enum and still uses
+//! `UnnestOptions::preserve_nulls`, which is why the planner wraps empty arrays with
+//! `ListEmptyToNullExpr` to get Spark's `explode_outer` semantics.
+
+use arrow::array::{
+ new_null_array, Array, ArrayRef, AsArray, FixedSizeListArray, Int64Array, LargeListArray,
+ LargeListViewArray, ListArray, ListViewArray, PrimitiveArray, Scalar, StructArray,
+};
+use arrow::compute::kernels::length::length;
+use arrow::compute::kernels::zip::zip;
+use arrow::compute::{cast, is_not_null, kernels, sum};
+use arrow::datatypes::{DataType, Int64Type, SchemaRef};
+use arrow::record_batch::RecordBatch;
+// Upstream imports this as `arrow_ord::cmp::lt`; Comet reaches it through `arrow`,
+// which does not have `arrow_ord` as a direct dependency.
+use arrow::compute::kernels::cmp::lt;
+use datafusion::common::{
+ exec_datafusion_err, exec_err, internal_err, HashMap, HashSet, Result, UnnestOptions,
+};
+use datafusion::execution::TaskContext;
+use datafusion::physical_expr::EquivalenceProperties;
+use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType};
+use datafusion::physical_plan::metrics::{
+ BaselineMetrics, Count, ExecutionPlanMetricsSet, MetricBuilder, MetricsSet, RecordOutput,
+};
+// `ListUnnest` is the one item the copied region below does NOT need to duplicate: unlike the
+// kernels, upstream exports it publicly.
+use datafusion::physical_plan::unnest::ListUnnest;
+use datafusion::physical_plan::{
+ DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties, RecordBatchStream,
+ SendableRecordBatchStream,
+};
+use futures::{Stream, StreamExt};
+use std::cmp::{self, Ordering};
+use std::pin::Pin;
+use std::sync::Arc;
+use std::task::{ready, Context, Poll};
+
+/// Comet's explode operator: DataFusion's `UnnestExec` with the input consumed in chunks so
+/// that output batches respect `datafusion.execution.batch_size`.
+#[derive(Debug)]
+pub struct ExplodeExec {
+ child: Arc<dyn ExecutionPlan>,
+ schema: SchemaRef,
+ list_column_indices: Vec<ListUnnest>,
+ struct_column_indices: Vec<usize>,
+ options: UnnestOptions,
+ metrics: ExecutionPlanMetricsSet,
+ cache: Arc<PlanProperties>,
+}
+
+impl ExplodeExec {
+ pub fn new(
+ child: Arc<dyn ExecutionPlan>,
+ list_column_indices: Vec<ListUnnest>,
+ struct_column_indices: Vec<usize>,
+ schema: SchemaRef,
+ options: UnnestOptions,
+ ) -> Self {
+ // Unnesting invalidates the child's orderings and constraints for the unnested
+ // columns, and Comet plans explode on a single partition, so start from empty
+ // equivalences rather than trying to project the child's.
+ let cache = Arc::new(PlanProperties::new(
+ EquivalenceProperties::new(Arc::clone(&schema)),
+ Partitioning::UnknownPartitioning(1),
+ EmissionType::Incremental,
+ Boundedness::Bounded,
+ ));
+
+ Self {
+ child,
+ schema,
+ list_column_indices,
+ struct_column_indices,
+ options,
+ metrics: ExecutionPlanMetricsSet::new(),
+ cache,
+ }
+ }
+}
+
+impl DisplayAs for ExplodeExec {
+ fn fmt_as(&self, t: DisplayFormatType, f: &mut std::fmt::Formatter) -> std::fmt::Result {
+ match t {
+ DisplayFormatType::Default | DisplayFormatType::Verbose => {
+ write!(f, "CometExplodeExec")
+ }
+ DisplayFormatType::TreeRender => unimplemented!(),
+ }
+ }
+}
+
+impl ExecutionPlan for ExplodeExec {
+ fn name(&self) -> &str {
+ "CometExplodeExec"
+ }
+
+ fn schema(&self) -> SchemaRef {
+ Arc::clone(&self.schema)
+ }
+
+ fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
+ vec![&self.child]
+ }
+
+ fn with_new_children(
+ self: Arc<Self>,
+ children: Vec<Arc<dyn ExecutionPlan>>,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ if children.len() != 1 {
+ return internal_err!("ExplodeExec expects exactly one child");
+ }
+ Ok(Arc::new(ExplodeExec::new(
+ Arc::clone(&children[0]),
+ self.list_column_indices.clone(),
+ self.struct_column_indices.clone(),
+ Arc::clone(&self.schema),
+ self.options.clone(),
+ )))
+ }
+
+ fn execute(
+ &self,
+ partition: usize,
+ context: Arc<TaskContext>,
+ ) -> Result<SendableRecordBatchStream> {
+ let batch_size = context.session_config().batch_size();
+ let input = self.child.execute(partition, context)?;
+
+ Ok(Box::pin(ExplodeStream {
+ input,
+ schema: Arc::clone(&self.schema),
+ list_type_columns: self.list_column_indices.clone(),
+ struct_column_indices: self.struct_column_indices.iter().copied().collect(),
+ options: self.options.clone(),
+ baseline_metrics: BaselineMetrics::new(&self.metrics, partition),
+ input_batches: MetricBuilder::new(&self.metrics).counter("input_batches", partition),
+ input_rows: MetricBuilder::new(&self.metrics).counter("input_rows", partition),
+ batch_size,
+ pending_input: None,
+ pending_output: None,
+ }))
+ }
+
+ fn properties(&self) -> &Arc<PlanProperties> {
+ &self.cache
+ }
+
+ fn metrics(&self) -> Option<MetricsSet> {
+ Some(self.metrics.clone_inner())
+ }
+}
+
+/// An input batch being unnested incrementally, a chunk of rows at a time.
+struct PendingInput {
+ /// The full input batch. Rows before `row_offset` have already been unnested.
+ batch: RecordBatch,
+ /// Index of the next input row to unnest.
+ row_offset: usize,
+ /// How many output rows each input row expands into, indexed by input row.
+ ///
+ /// Empty when the expansion factor cannot be predicted from the input alone, in which
+ /// case the whole remaining input is unnested in one call and only the output is split.
+ /// See [`ExplodeStream::predict_output_lens`].
+ output_lens: Vec<usize>,
+}
+
+impl PendingInput {
+ fn remaining_rows(&self) -> usize {
+ self.batch.num_rows() - self.row_offset
+ }
+
+ /// How many input rows to unnest next so the resulting batch holds at most `batch_size`
+ /// rows.
+ ///
+ /// Always returns at least 1 while rows remain, so the stream always makes progress: a
+ /// single input row is never split across output batches, so one row whose array is
+ /// longer than `batch_size` still produces one oversized build, which is then sliced
+ /// down on the way out.
+ fn next_chunk_rows(&self, batch_size: usize) -> usize {
+ let remaining = self.remaining_rows();
+ if self.output_lens.is_empty() {
+ return remaining;
+ }
+
+ let mut rows = 0;
+ let mut output_rows = 0usize;
+ while rows < remaining {
+ let len = self.output_lens[self.row_offset + rows];
+ if rows > 0 && output_rows.saturating_add(len) > batch_size {
+ break;
+ }
+ output_rows += len;
+ rows += 1;
+ }
+ rows
+ }
+}
+
+/// A stream that unnests its input, bounding output batches to `batch_size` rows.
+struct ExplodeStream {
+ input: SendableRecordBatchStream,
+ schema: SchemaRef,
+ list_type_columns: Vec<ListUnnest>,
+ struct_column_indices: HashSet<usize>,
+ options: UnnestOptions,
+ baseline_metrics: BaselineMetrics,
+ input_batches: Count,
+ input_rows: Count,
+ /// Target number of rows per output batch, from `datafusion.execution.batch_size`.
+ batch_size: usize,
+ /// Rows of the current input batch that have not been unnested yet.
+ pending_input: Option<PendingInput>,
+ /// Unnested rows that have been built but not emitted yet.
+ pending_output: Option<RecordBatch>,
+}
+
+impl RecordBatchStream for ExplodeStream {
+ fn schema(&self) -> SchemaRef {
+ Arc::clone(&self.schema)
+ }
+}
+
+impl Stream for ExplodeStream {
+ type Item = Result<RecordBatch>;
+
+ fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+ self.poll_next_impl(cx)
+ }
+}
+
+impl ExplodeStream {
+ fn poll_next_impl(&mut self, cx: &mut Context<'_>) -> Poll<Option<Result<RecordBatch>>> {
+ loop {
+ // Emit already-unnested rows first, at most `batch_size` at a time.
+ if let Some(batch) = self.pending_output.take() {
+ let (emit, rest) = split_off_head(batch, self.batch_size);
+ self.pending_output = rest;
+ (&emit).record_output(&self.baseline_metrics);
+ return Poll::Ready(Some(Ok(emit)));
+ }
+
+ // Unnest the next chunk of the input batch already in hand.
+ if let Some(pending) = self.pending_input.as_mut() {
+ if pending.remaining_rows() == 0 {
+ self.pending_input = None;
+ continue;
+ }
+
+ let timer = self.baseline_metrics.elapsed_compute().timer();
+
+ let rows = pending.next_chunk_rows(self.batch_size);
+ let chunk = pending.batch.slice(pending.row_offset, rows);
+ pending.row_offset += rows;
+
+ let result = build_batch(
+ &chunk,
+ &self.schema,
+ &self.list_type_columns,
+ &self.struct_column_indices,
+ &self.options,
+ );
+ timer.done();
+
+ // A chunk can legitimately produce no rows at all (for example rows whose
+ // arrays are all empty and `preserve_nulls` is false); move on to the next
+ // chunk rather than emitting an empty batch.
+ self.pending_output = result?.filter(|batch| batch.num_rows() > 0);
+ continue;
+ }
+
+ // Otherwise pull the next input batch.
+ return Poll::Ready(match ready!(self.input.poll_next_unpin(cx)) {
+ Some(Ok(batch)) => {
+ self.input_batches.add(1);
+ self.input_rows.add(batch.num_rows());
+ if batch.num_rows() == 0 {
+ continue;
+ }
+
+ let timer = self.baseline_metrics.elapsed_compute().timer();
+ let output_lens = self.predict_output_lens(&batch);
+ timer.done();
+
+ match output_lens {
+ Ok(output_lens) => {
+ self.pending_input = Some(PendingInput {
+ batch,
+ row_offset: 0,
+ output_lens,
+ });
+ continue;
+ }
+ Err(e) => Some(Err(e)),
+ }
+ }
+ other => other,
+ });
+ }
+ }
+
+ /// Compute how many output rows each input row of `batch` will expand into, so the input
+ /// can be chunked to keep each build bounded.
+ ///
+ /// Returns an empty vec when the count cannot be derived from the input alone, which is
+ /// the signal to unnest the whole batch in one call:
+ ///
+ /// * With no list columns, unnesting only widens structs and leaves the row count alone,
+ /// so the output is already bounded by the input batch size.
+ /// * With recursion (`depth > 1`), a row's expansion depends on the lengths of inner
+ /// lists that only exist after the outer levels have been unnested, so it cannot be
+ /// predicted up front. Comet only plans depth-1 explode today, but the fallback keeps
+ /// this correct if that changes.
+ fn predict_output_lens(&self, batch: &RecordBatch) -> Result<Vec<usize>> {
+ if self.list_type_columns.is_empty()
+ || self
+ .list_type_columns
+ .iter()
+ .any(|unnest| unnest.depth != 1)
+ {
+ return Ok(vec![]);
+ }
+
+ let list_arrays: Vec<ArrayRef> = self
+ .list_type_columns
+ .iter()
+ .map(|unnest| Arc::clone(batch.column(unnest.index_in_input_schema)))
+ .collect();
+
+ // The same per-row length that `list_unnest_at_level` derives when it actually
+ // unnests, so the chunk boundaries it produces are exact.
+ let longest_length = find_longest_length(&list_arrays, &self.options)?;
Review Comment:
Done, and it now matches what the upstream PR does. `predict_output_lens` keeps the result as a `PrimitiveArray<Int64Type>` rather than collecting into `Vec<usize>`, `PendingInput::chunk_lengths` hands out a zero-copy slice of it per chunk, and that slice threads down through `build_batch` into `list_unnest_at_level`, which uses it in place of calling `find_longest_length` again. Reuse is gated on `max_recursion == 1`, since with recursion the deeper levels depend on arrays that do not exist yet — the same reason the prediction is skipped there in the first place.
This does mean two functions in the vendored region gain a parameter, which cuts against the "do not change it semantically" banner, so I called it out explicitly in the module docs as a deliberate edit that is itself part of apache/datafusion#24384 and therefore disappears with the rest of the fork.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix: make CometExplodeExec respect batch size [datafusion-comet]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove commented on PR #5362:
URL: https://github.com/apache/datafusion-comet/pull/5362#issuecomment-5484810069
Merged. Thanks @sunchao @mbutrovich.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix: make CometExplodeExec respect batch size [datafusion-comet]
Posted by "sunchao (via GitHub)" <gi...@apache.org>.
sunchao commented on code in PR #5362:
URL: https://github.com/apache/datafusion-comet/pull/5362#discussion_r3788435135
##########
native/core/src/execution/operators/explode.rs:
##########
@@ -0,0 +1,1237 @@
+// 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.
+
+//! A temporary fork of DataFusion's `UnnestExec` that respects
+//! `datafusion.execution.batch_size`.
+//!
+//! # Why this fork exists
+//!
+//! DataFusion's `UnnestExec` emits exactly one output batch per input batch, however many
+//! rows the unnesting produces, and never consults `batch_size`. For `explode` this means
+//! an 8192-row batch of 100-element arrays comes back as a single 819,200-row batch, and
+//! peak memory scales with input batch size times array length rather than with
+//! `batch_size`.
+//!
+//! The fix has been submitted upstream:
+//!
+//! * <https://github.com/apache/datafusion/issues/24383>
+//! * <https://github.com/apache/datafusion/pull/24384>
+//!
+//! # Deleting this file
+//!
+//! Once Comet moves to a DataFusion release carrying apache/datafusion#24384, delete this
+//! module and go back to `datafusion::physical_plan::unnest::UnnestExec` in the planner.
+//!
+//! Note that <https://github.com/apache/datafusion-comet/issues/5210> is a *different*
+//! unnest cleanup — it tracks adopting upstream `unnest_outer`
+//! (apache/datafusion#22100) to retire `ListEmptyToNullExpr`. The two upstream PRs can
+//! land in different releases, so closing 5210 is not a signal to delete this fork.
+//!
+//! # What was forked
+//!
+//! The unnesting kernels below (`build_batch` and everything it calls) are copied from
+//! `datafusion/physical-plan/src/unnest.rs` at DataFusion 54.1.0, upstream revision
+//! `cc7565be1ee97ba8fa2f5d6da373c5e38d81bb13`. They are private to
+//! `datafusion-physical-plan`, so they cannot be called from here without copying them.
+//! Leave them semantically unmodified so the eventual deletion is mechanical; the
+//! Comet-specific behavior lives entirely in `ExplodeExec` and `ExplodeStream`.
+//!
+//! They are not byte-identical to upstream: Comet's rustfmt uses `max_width = 100` and
+//! edition 2021, DataFusion's uses `max_width = 90` and edition 2024, so `cargo fmt`
+//! reflows some signatures. To audit for real changes, reformat this region at
+//! `max_width = 90` and diff it against upstream `unnest.rs`; that reduces the difference
+//! to a single cosmetic line wrap in `flatten_struct_cols`. The only deliberate edits are
+//! the `lt` import path noted below and dropping upstream's `ListUnnest` declaration in
+//! favor of importing the public one.
+//!
+//! Note that 54.1.0 predates upstream's `NullHandling` enum and still uses
+//! `UnnestOptions::preserve_nulls`, which is why the planner wraps empty arrays with
+//! `ListEmptyToNullExpr` to get Spark's `explode_outer` semantics.
+
+use arrow::array::{
+ new_null_array, Array, ArrayRef, AsArray, FixedSizeListArray, Int64Array, LargeListArray,
+ LargeListViewArray, ListArray, ListViewArray, PrimitiveArray, Scalar, StructArray,
+};
+use arrow::compute::kernels::length::length;
+use arrow::compute::kernels::zip::zip;
+use arrow::compute::{cast, is_not_null, kernels, sum};
+use arrow::datatypes::{DataType, Int64Type, SchemaRef};
+use arrow::record_batch::RecordBatch;
+// Upstream imports this as `arrow_ord::cmp::lt`; Comet reaches it through `arrow`,
+// which does not have `arrow_ord` as a direct dependency.
+use arrow::compute::kernels::cmp::lt;
+use datafusion::common::{
+ exec_datafusion_err, exec_err, internal_err, HashMap, HashSet, Result, UnnestOptions,
+};
+use datafusion::execution::TaskContext;
+use datafusion::physical_expr::EquivalenceProperties;
+use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType};
+use datafusion::physical_plan::metrics::{
+ BaselineMetrics, Count, ExecutionPlanMetricsSet, MetricBuilder, MetricsSet, RecordOutput,
+};
+// `ListUnnest` is the one item the copied region below does NOT need to duplicate: unlike the
+// kernels, upstream exports it publicly.
+use datafusion::physical_plan::unnest::ListUnnest;
+use datafusion::physical_plan::{
+ DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties, RecordBatchStream,
+ SendableRecordBatchStream,
+};
+use futures::{Stream, StreamExt};
+use std::cmp::{self, Ordering};
+use std::pin::Pin;
+use std::sync::Arc;
+use std::task::{ready, Context, Poll};
+
+/// Comet's explode operator: DataFusion's `UnnestExec` with the input consumed in chunks so
+/// that output batches respect `datafusion.execution.batch_size`.
+#[derive(Debug)]
+pub struct ExplodeExec {
+ child: Arc<dyn ExecutionPlan>,
+ schema: SchemaRef,
+ list_column_indices: Vec<ListUnnest>,
+ struct_column_indices: Vec<usize>,
+ options: UnnestOptions,
+ metrics: ExecutionPlanMetricsSet,
+ cache: Arc<PlanProperties>,
+}
+
+impl ExplodeExec {
+ pub fn new(
+ child: Arc<dyn ExecutionPlan>,
+ list_column_indices: Vec<ListUnnest>,
+ struct_column_indices: Vec<usize>,
+ schema: SchemaRef,
+ options: UnnestOptions,
+ ) -> Self {
+ // Unnesting invalidates the child's orderings and constraints for the unnested
+ // columns, and Comet plans explode on a single partition, so start from empty
+ // equivalences rather than trying to project the child's.
+ let cache = Arc::new(PlanProperties::new(
+ EquivalenceProperties::new(Arc::clone(&schema)),
+ Partitioning::UnknownPartitioning(1),
+ EmissionType::Incremental,
+ Boundedness::Bounded,
+ ));
+
+ Self {
+ child,
+ schema,
+ list_column_indices,
+ struct_column_indices,
+ options,
+ metrics: ExecutionPlanMetricsSet::new(),
+ cache,
+ }
+ }
+}
+
+impl DisplayAs for ExplodeExec {
+ fn fmt_as(&self, t: DisplayFormatType, f: &mut std::fmt::Formatter) -> std::fmt::Result {
+ match t {
+ DisplayFormatType::Default | DisplayFormatType::Verbose => {
+ write!(f, "CometExplodeExec")
+ }
+ DisplayFormatType::TreeRender => unimplemented!(),
+ }
+ }
+}
+
+impl ExecutionPlan for ExplodeExec {
+ fn name(&self) -> &str {
+ "CometExplodeExec"
+ }
+
+ fn schema(&self) -> SchemaRef {
+ Arc::clone(&self.schema)
+ }
+
+ fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
+ vec![&self.child]
+ }
+
+ fn with_new_children(
+ self: Arc<Self>,
+ children: Vec<Arc<dyn ExecutionPlan>>,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ if children.len() != 1 {
+ return internal_err!("ExplodeExec expects exactly one child");
+ }
+ Ok(Arc::new(ExplodeExec::new(
+ Arc::clone(&children[0]),
+ self.list_column_indices.clone(),
+ self.struct_column_indices.clone(),
+ Arc::clone(&self.schema),
+ self.options.clone(),
+ )))
+ }
+
+ fn execute(
+ &self,
+ partition: usize,
+ context: Arc<TaskContext>,
+ ) -> Result<SendableRecordBatchStream> {
+ let batch_size = context.session_config().batch_size();
+ let input = self.child.execute(partition, context)?;
+
+ Ok(Box::pin(ExplodeStream {
+ input,
+ schema: Arc::clone(&self.schema),
+ list_type_columns: self.list_column_indices.clone(),
+ struct_column_indices: self.struct_column_indices.iter().copied().collect(),
+ options: self.options.clone(),
+ baseline_metrics: BaselineMetrics::new(&self.metrics, partition),
+ input_batches: MetricBuilder::new(&self.metrics).counter("input_batches", partition),
+ input_rows: MetricBuilder::new(&self.metrics).counter("input_rows", partition),
+ batch_size,
+ pending_input: None,
+ pending_output: None,
+ }))
+ }
+
+ fn properties(&self) -> &Arc<PlanProperties> {
+ &self.cache
+ }
+
+ fn metrics(&self) -> Option<MetricsSet> {
+ Some(self.metrics.clone_inner())
+ }
+}
+
+/// An input batch being unnested incrementally, a chunk of rows at a time.
+struct PendingInput {
+ /// The full input batch. Rows before `row_offset` have already been unnested.
+ batch: RecordBatch,
+ /// Index of the next input row to unnest.
+ row_offset: usize,
+ /// How many output rows each input row expands into, indexed by input row.
+ ///
+ /// Empty when the expansion factor cannot be predicted from the input alone, in which
+ /// case the whole remaining input is unnested in one call and only the output is split.
+ /// See [`ExplodeStream::predict_output_lens`].
+ output_lens: Vec<usize>,
+}
+
+impl PendingInput {
+ fn remaining_rows(&self) -> usize {
+ self.batch.num_rows() - self.row_offset
+ }
+
+ /// How many input rows to unnest next so the resulting batch holds at most `batch_size`
+ /// rows.
+ ///
+ /// Always returns at least 1 while rows remain, so the stream always makes progress: a
+ /// single input row is never split across output batches, so one row whose array is
+ /// longer than `batch_size` still produces one oversized build, which is then sliced
+ /// down on the way out.
+ fn next_chunk_rows(&self, batch_size: usize) -> usize {
+ let remaining = self.remaining_rows();
+ if self.output_lens.is_empty() {
+ return remaining;
+ }
+
+ let mut rows = 0;
+ let mut output_rows = 0usize;
+ while rows < remaining {
+ let len = self.output_lens[self.row_offset + rows];
+ if rows > 0 && output_rows.saturating_add(len) > batch_size {
+ break;
+ }
+ output_rows += len;
+ rows += 1;
+ }
+ rows
+ }
+}
+
+/// A stream that unnests its input, bounding output batches to `batch_size` rows.
+struct ExplodeStream {
+ input: SendableRecordBatchStream,
+ schema: SchemaRef,
+ list_type_columns: Vec<ListUnnest>,
+ struct_column_indices: HashSet<usize>,
+ options: UnnestOptions,
+ baseline_metrics: BaselineMetrics,
+ input_batches: Count,
+ input_rows: Count,
+ /// Target number of rows per output batch, from `datafusion.execution.batch_size`.
+ batch_size: usize,
+ /// Rows of the current input batch that have not been unnested yet.
+ pending_input: Option<PendingInput>,
+ /// Unnested rows that have been built but not emitted yet.
+ pending_output: Option<RecordBatch>,
+}
+
+impl RecordBatchStream for ExplodeStream {
+ fn schema(&self) -> SchemaRef {
+ Arc::clone(&self.schema)
+ }
+}
+
+impl Stream for ExplodeStream {
+ type Item = Result<RecordBatch>;
+
+ fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+ self.poll_next_impl(cx)
+ }
+}
+
+impl ExplodeStream {
+ fn poll_next_impl(&mut self, cx: &mut Context<'_>) -> Poll<Option<Result<RecordBatch>>> {
+ loop {
+ // Emit already-unnested rows first, at most `batch_size` at a time.
+ if let Some(batch) = self.pending_output.take() {
+ let (emit, rest) = split_off_head(batch, self.batch_size);
+ self.pending_output = rest;
+ (&emit).record_output(&self.baseline_metrics);
+ return Poll::Ready(Some(Ok(emit)));
+ }
+
+ // Unnest the next chunk of the input batch already in hand.
+ if let Some(pending) = self.pending_input.as_mut() {
+ if pending.remaining_rows() == 0 {
+ self.pending_input = None;
+ continue;
+ }
+
+ let timer = self.baseline_metrics.elapsed_compute().timer();
+
+ let rows = pending.next_chunk_rows(self.batch_size);
+ let chunk = pending.batch.slice(pending.row_offset, rows);
+ pending.row_offset += rows;
+
+ let result = build_batch(
+ &chunk,
+ &self.schema,
+ &self.list_type_columns,
+ &self.struct_column_indices,
+ &self.options,
+ );
+ timer.done();
+
+ // A chunk can legitimately produce no rows at all (for example rows whose
+ // arrays are all empty and `preserve_nulls` is false); move on to the next
+ // chunk rather than emitting an empty batch.
+ self.pending_output = result?.filter(|batch| batch.num_rows() > 0);
+ continue;
+ }
+
+ // Otherwise pull the next input batch.
+ return Poll::Ready(match ready!(self.input.poll_next_unpin(cx)) {
+ Some(Ok(batch)) => {
+ self.input_batches.add(1);
+ self.input_rows.add(batch.num_rows());
+ if batch.num_rows() == 0 {
+ continue;
+ }
+
+ let timer = self.baseline_metrics.elapsed_compute().timer();
+ let output_lens = self.predict_output_lens(&batch);
+ timer.done();
+
+ match output_lens {
+ Ok(output_lens) => {
+ self.pending_input = Some(PendingInput {
+ batch,
+ row_offset: 0,
+ output_lens,
+ });
+ continue;
+ }
+ Err(e) => Some(Err(e)),
+ }
+ }
+ other => other,
Review Comment:
**[P1] Release the exhausted child before returning EOF**
Could we restore the EOF cleanup from DataFusion's `UnnestStream` here? The original operator replaces `self.input` with `EmptyRecordBatchStream` when the child is exhausted. Without that, a real `HashJoin -> Projection -> Explode -> ShuffleWriter` plan keeps the partitioned hash join's build table and `MemoryReservation` alive while `shuffle_write()` finalizes. `ParquetWriterExec` similarly retains the input through `writer.close()`. This can force avoidable spills or out-of-memory failures when those writers need the memory. A child stream with a drop guard or reservation would make a useful regression test.
##########
native/core/src/execution/operators/explode.rs:
##########
@@ -0,0 +1,1237 @@
+// 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.
+
+//! A temporary fork of DataFusion's `UnnestExec` that respects
+//! `datafusion.execution.batch_size`.
+//!
+//! # Why this fork exists
+//!
+//! DataFusion's `UnnestExec` emits exactly one output batch per input batch, however many
+//! rows the unnesting produces, and never consults `batch_size`. For `explode` this means
+//! an 8192-row batch of 100-element arrays comes back as a single 819,200-row batch, and
+//! peak memory scales with input batch size times array length rather than with
+//! `batch_size`.
+//!
+//! The fix has been submitted upstream:
+//!
+//! * <https://github.com/apache/datafusion/issues/24383>
+//! * <https://github.com/apache/datafusion/pull/24384>
+//!
+//! # Deleting this file
+//!
+//! Once Comet moves to a DataFusion release carrying apache/datafusion#24384, delete this
+//! module and go back to `datafusion::physical_plan::unnest::UnnestExec` in the planner.
+//!
+//! Note that <https://github.com/apache/datafusion-comet/issues/5210> is a *different*
+//! unnest cleanup — it tracks adopting upstream `unnest_outer`
+//! (apache/datafusion#22100) to retire `ListEmptyToNullExpr`. The two upstream PRs can
+//! land in different releases, so closing 5210 is not a signal to delete this fork.
+//!
+//! # What was forked
+//!
+//! The unnesting kernels below (`build_batch` and everything it calls) are copied from
+//! `datafusion/physical-plan/src/unnest.rs` at DataFusion 54.1.0, upstream revision
+//! `cc7565be1ee97ba8fa2f5d6da373c5e38d81bb13`. They are private to
+//! `datafusion-physical-plan`, so they cannot be called from here without copying them.
+//! Leave them semantically unmodified so the eventual deletion is mechanical; the
+//! Comet-specific behavior lives entirely in `ExplodeExec` and `ExplodeStream`.
+//!
+//! They are not byte-identical to upstream: Comet's rustfmt uses `max_width = 100` and
+//! edition 2021, DataFusion's uses `max_width = 90` and edition 2024, so `cargo fmt`
+//! reflows some signatures. To audit for real changes, reformat this region at
+//! `max_width = 90` and diff it against upstream `unnest.rs`; that reduces the difference
+//! to a single cosmetic line wrap in `flatten_struct_cols`. The only deliberate edits are
+//! the `lt` import path noted below and dropping upstream's `ListUnnest` declaration in
+//! favor of importing the public one.
+//!
+//! Note that 54.1.0 predates upstream's `NullHandling` enum and still uses
+//! `UnnestOptions::preserve_nulls`, which is why the planner wraps empty arrays with
+//! `ListEmptyToNullExpr` to get Spark's `explode_outer` semantics.
+
+use arrow::array::{
+ new_null_array, Array, ArrayRef, AsArray, FixedSizeListArray, Int64Array, LargeListArray,
+ LargeListViewArray, ListArray, ListViewArray, PrimitiveArray, Scalar, StructArray,
+};
+use arrow::compute::kernels::length::length;
+use arrow::compute::kernels::zip::zip;
+use arrow::compute::{cast, is_not_null, kernels, sum};
+use arrow::datatypes::{DataType, Int64Type, SchemaRef};
+use arrow::record_batch::RecordBatch;
+// Upstream imports this as `arrow_ord::cmp::lt`; Comet reaches it through `arrow`,
+// which does not have `arrow_ord` as a direct dependency.
+use arrow::compute::kernels::cmp::lt;
+use datafusion::common::{
+ exec_datafusion_err, exec_err, internal_err, HashMap, HashSet, Result, UnnestOptions,
+};
+use datafusion::execution::TaskContext;
+use datafusion::physical_expr::EquivalenceProperties;
+use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType};
+use datafusion::physical_plan::metrics::{
+ BaselineMetrics, Count, ExecutionPlanMetricsSet, MetricBuilder, MetricsSet, RecordOutput,
+};
+// `ListUnnest` is the one item the copied region below does NOT need to duplicate: unlike the
+// kernels, upstream exports it publicly.
+use datafusion::physical_plan::unnest::ListUnnest;
+use datafusion::physical_plan::{
+ DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties, RecordBatchStream,
+ SendableRecordBatchStream,
+};
+use futures::{Stream, StreamExt};
+use std::cmp::{self, Ordering};
+use std::pin::Pin;
+use std::sync::Arc;
+use std::task::{ready, Context, Poll};
+
+/// Comet's explode operator: DataFusion's `UnnestExec` with the input consumed in chunks so
+/// that output batches respect `datafusion.execution.batch_size`.
+#[derive(Debug)]
+pub struct ExplodeExec {
+ child: Arc<dyn ExecutionPlan>,
+ schema: SchemaRef,
+ list_column_indices: Vec<ListUnnest>,
+ struct_column_indices: Vec<usize>,
+ options: UnnestOptions,
+ metrics: ExecutionPlanMetricsSet,
+ cache: Arc<PlanProperties>,
+}
+
+impl ExplodeExec {
+ pub fn new(
+ child: Arc<dyn ExecutionPlan>,
+ list_column_indices: Vec<ListUnnest>,
+ struct_column_indices: Vec<usize>,
+ schema: SchemaRef,
+ options: UnnestOptions,
+ ) -> Self {
+ // Unnesting invalidates the child's orderings and constraints for the unnested
+ // columns, and Comet plans explode on a single partition, so start from empty
+ // equivalences rather than trying to project the child's.
+ let cache = Arc::new(PlanProperties::new(
+ EquivalenceProperties::new(Arc::clone(&schema)),
+ Partitioning::UnknownPartitioning(1),
+ EmissionType::Incremental,
+ Boundedness::Bounded,
+ ));
+
+ Self {
+ child,
+ schema,
+ list_column_indices,
+ struct_column_indices,
+ options,
+ metrics: ExecutionPlanMetricsSet::new(),
+ cache,
+ }
+ }
+}
+
+impl DisplayAs for ExplodeExec {
+ fn fmt_as(&self, t: DisplayFormatType, f: &mut std::fmt::Formatter) -> std::fmt::Result {
+ match t {
+ DisplayFormatType::Default | DisplayFormatType::Verbose => {
+ write!(f, "CometExplodeExec")
+ }
+ DisplayFormatType::TreeRender => unimplemented!(),
+ }
+ }
+}
+
+impl ExecutionPlan for ExplodeExec {
+ fn name(&self) -> &str {
+ "CometExplodeExec"
+ }
+
+ fn schema(&self) -> SchemaRef {
+ Arc::clone(&self.schema)
+ }
+
+ fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
+ vec![&self.child]
+ }
+
+ fn with_new_children(
+ self: Arc<Self>,
+ children: Vec<Arc<dyn ExecutionPlan>>,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ if children.len() != 1 {
+ return internal_err!("ExplodeExec expects exactly one child");
+ }
+ Ok(Arc::new(ExplodeExec::new(
+ Arc::clone(&children[0]),
+ self.list_column_indices.clone(),
+ self.struct_column_indices.clone(),
+ Arc::clone(&self.schema),
+ self.options.clone(),
+ )))
+ }
+
+ fn execute(
+ &self,
+ partition: usize,
+ context: Arc<TaskContext>,
+ ) -> Result<SendableRecordBatchStream> {
+ let batch_size = context.session_config().batch_size();
+ let input = self.child.execute(partition, context)?;
+
+ Ok(Box::pin(ExplodeStream {
+ input,
+ schema: Arc::clone(&self.schema),
+ list_type_columns: self.list_column_indices.clone(),
+ struct_column_indices: self.struct_column_indices.iter().copied().collect(),
+ options: self.options.clone(),
+ baseline_metrics: BaselineMetrics::new(&self.metrics, partition),
+ input_batches: MetricBuilder::new(&self.metrics).counter("input_batches", partition),
+ input_rows: MetricBuilder::new(&self.metrics).counter("input_rows", partition),
+ batch_size,
+ pending_input: None,
+ pending_output: None,
+ }))
+ }
+
+ fn properties(&self) -> &Arc<PlanProperties> {
+ &self.cache
+ }
+
+ fn metrics(&self) -> Option<MetricsSet> {
+ Some(self.metrics.clone_inner())
+ }
+}
+
+/// An input batch being unnested incrementally, a chunk of rows at a time.
+struct PendingInput {
+ /// The full input batch. Rows before `row_offset` have already been unnested.
+ batch: RecordBatch,
+ /// Index of the next input row to unnest.
+ row_offset: usize,
+ /// How many output rows each input row expands into, indexed by input row.
+ ///
+ /// Empty when the expansion factor cannot be predicted from the input alone, in which
+ /// case the whole remaining input is unnested in one call and only the output is split.
+ /// See [`ExplodeStream::predict_output_lens`].
+ output_lens: Vec<usize>,
+}
+
+impl PendingInput {
+ fn remaining_rows(&self) -> usize {
+ self.batch.num_rows() - self.row_offset
+ }
+
+ /// How many input rows to unnest next so the resulting batch holds at most `batch_size`
+ /// rows.
+ ///
+ /// Always returns at least 1 while rows remain, so the stream always makes progress: a
+ /// single input row is never split across output batches, so one row whose array is
+ /// longer than `batch_size` still produces one oversized build, which is then sliced
+ /// down on the way out.
+ fn next_chunk_rows(&self, batch_size: usize) -> usize {
+ let remaining = self.remaining_rows();
+ if self.output_lens.is_empty() {
+ return remaining;
+ }
+
+ let mut rows = 0;
+ let mut output_rows = 0usize;
+ while rows < remaining {
+ let len = self.output_lens[self.row_offset + rows];
+ if rows > 0 && output_rows.saturating_add(len) > batch_size {
+ break;
+ }
+ output_rows += len;
+ rows += 1;
+ }
+ rows
+ }
+}
+
+/// A stream that unnests its input, bounding output batches to `batch_size` rows.
+struct ExplodeStream {
+ input: SendableRecordBatchStream,
+ schema: SchemaRef,
+ list_type_columns: Vec<ListUnnest>,
+ struct_column_indices: HashSet<usize>,
+ options: UnnestOptions,
+ baseline_metrics: BaselineMetrics,
+ input_batches: Count,
+ input_rows: Count,
+ /// Target number of rows per output batch, from `datafusion.execution.batch_size`.
+ batch_size: usize,
+ /// Rows of the current input batch that have not been unnested yet.
+ pending_input: Option<PendingInput>,
+ /// Unnested rows that have been built but not emitted yet.
+ pending_output: Option<RecordBatch>,
+}
+
+impl RecordBatchStream for ExplodeStream {
+ fn schema(&self) -> SchemaRef {
+ Arc::clone(&self.schema)
+ }
+}
+
+impl Stream for ExplodeStream {
+ type Item = Result<RecordBatch>;
+
+ fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+ self.poll_next_impl(cx)
+ }
+}
+
+impl ExplodeStream {
+ fn poll_next_impl(&mut self, cx: &mut Context<'_>) -> Poll<Option<Result<RecordBatch>>> {
+ loop {
+ // Emit already-unnested rows first, at most `batch_size` at a time.
+ if let Some(batch) = self.pending_output.take() {
+ let (emit, rest) = split_off_head(batch, self.batch_size);
+ self.pending_output = rest;
+ (&emit).record_output(&self.baseline_metrics);
+ return Poll::Ready(Some(Ok(emit)));
+ }
+
+ // Unnest the next chunk of the input batch already in hand.
+ if let Some(pending) = self.pending_input.as_mut() {
+ if pending.remaining_rows() == 0 {
+ self.pending_input = None;
+ continue;
+ }
+
+ let timer = self.baseline_metrics.elapsed_compute().timer();
+
+ let rows = pending.next_chunk_rows(self.batch_size);
+ let chunk = pending.batch.slice(pending.row_offset, rows);
+ pending.row_offset += rows;
+
+ let result = build_batch(
+ &chunk,
+ &self.schema,
+ &self.list_type_columns,
+ &self.struct_column_indices,
+ &self.options,
+ );
+ timer.done();
+
+ // A chunk can legitimately produce no rows at all (for example rows whose
+ // arrays are all empty and `preserve_nulls` is false); move on to the next
+ // chunk rather than emitting an empty batch.
+ self.pending_output = result?.filter(|batch| batch.num_rows() > 0);
+ continue;
+ }
+
+ // Otherwise pull the next input batch.
+ return Poll::Ready(match ready!(self.input.poll_next_unpin(cx)) {
+ Some(Ok(batch)) => {
+ self.input_batches.add(1);
+ self.input_rows.add(batch.num_rows());
+ if batch.num_rows() == 0 {
+ continue;
+ }
+
+ let timer = self.baseline_metrics.elapsed_compute().timer();
+ let output_lens = self.predict_output_lens(&batch);
+ timer.done();
+
+ match output_lens {
+ Ok(output_lens) => {
+ self.pending_input = Some(PendingInput {
+ batch,
+ row_offset: 0,
+ output_lens,
+ });
+ continue;
+ }
+ Err(e) => Some(Err(e)),
+ }
+ }
+ other => other,
+ });
+ }
+ }
+
+ /// Compute how many output rows each input row of `batch` will expand into, so the input
+ /// can be chunked to keep each build bounded.
+ ///
+ /// Returns an empty vec when the count cannot be derived from the input alone, which is
+ /// the signal to unnest the whole batch in one call:
+ ///
+ /// * With no list columns, unnesting only widens structs and leaves the row count alone,
+ /// so the output is already bounded by the input batch size.
+ /// * With recursion (`depth > 1`), a row's expansion depends on the lengths of inner
+ /// lists that only exist after the outer levels have been unnested, so it cannot be
+ /// predicted up front. Comet only plans depth-1 explode today, but the fallback keeps
+ /// this correct if that changes.
+ fn predict_output_lens(&self, batch: &RecordBatch) -> Result<Vec<usize>> {
+ if self.list_type_columns.is_empty()
+ || self
+ .list_type_columns
+ .iter()
+ .any(|unnest| unnest.depth != 1)
+ {
+ return Ok(vec![]);
+ }
+
+ let list_arrays: Vec<ArrayRef> = self
+ .list_type_columns
+ .iter()
+ .map(|unnest| Arc::clone(batch.column(unnest.index_in_input_schema)))
+ .collect();
+
+ // The same per-row length that `list_unnest_at_level` derives when it actually
+ // unnests, so the chunk boundaries it produces are exact.
+ let longest_length = find_longest_length(&list_arrays, &self.options)?;
Review Comment:
**[P2] Reuse the list lengths already computed for chunking**
Could we keep the predicted lengths as an Arrow `PrimitiveArray<Int64Type>` and pass zero-copy slices into `build_batch`? `predict_output_lens` already runs `length`, `cast`, `is_not_null`, and `zip` over every input row, but converts the result into `Vec<usize>`. `list_unnest_at_level` then runs the same kernels again for every chunk. This duplicates work even when no split is needed, and `posexplode` repeats it across both list columns. The linked upstream fix already reuses the precomputed chunk lengths.
##########
native/core/src/execution/operators/explode.rs:
##########
@@ -0,0 +1,1237 @@
+// 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.
+
+//! A temporary fork of DataFusion's `UnnestExec` that respects
+//! `datafusion.execution.batch_size`.
+//!
+//! # Why this fork exists
+//!
+//! DataFusion's `UnnestExec` emits exactly one output batch per input batch, however many
+//! rows the unnesting produces, and never consults `batch_size`. For `explode` this means
+//! an 8192-row batch of 100-element arrays comes back as a single 819,200-row batch, and
+//! peak memory scales with input batch size times array length rather than with
+//! `batch_size`.
+//!
+//! The fix has been submitted upstream:
+//!
+//! * <https://github.com/apache/datafusion/issues/24383>
+//! * <https://github.com/apache/datafusion/pull/24384>
+//!
+//! # Deleting this file
+//!
+//! Once Comet moves to a DataFusion release carrying apache/datafusion#24384, delete this
+//! module and go back to `datafusion::physical_plan::unnest::UnnestExec` in the planner.
+//!
+//! Note that <https://github.com/apache/datafusion-comet/issues/5210> is a *different*
+//! unnest cleanup — it tracks adopting upstream `unnest_outer`
+//! (apache/datafusion#22100) to retire `ListEmptyToNullExpr`. The two upstream PRs can
+//! land in different releases, so closing 5210 is not a signal to delete this fork.
+//!
+//! # What was forked
+//!
+//! The unnesting kernels below (`build_batch` and everything it calls) are copied from
+//! `datafusion/physical-plan/src/unnest.rs` at DataFusion 54.1.0, upstream revision
+//! `cc7565be1ee97ba8fa2f5d6da373c5e38d81bb13`. They are private to
+//! `datafusion-physical-plan`, so they cannot be called from here without copying them.
+//! Leave them semantically unmodified so the eventual deletion is mechanical; the
+//! Comet-specific behavior lives entirely in `ExplodeExec` and `ExplodeStream`.
+//!
+//! They are not byte-identical to upstream: Comet's rustfmt uses `max_width = 100` and
+//! edition 2021, DataFusion's uses `max_width = 90` and edition 2024, so `cargo fmt`
+//! reflows some signatures. To audit for real changes, reformat this region at
+//! `max_width = 90` and diff it against upstream `unnest.rs`; that reduces the difference
+//! to a single cosmetic line wrap in `flatten_struct_cols`. The only deliberate edits are
+//! the `lt` import path noted below and dropping upstream's `ListUnnest` declaration in
+//! favor of importing the public one.
+//!
+//! Note that 54.1.0 predates upstream's `NullHandling` enum and still uses
+//! `UnnestOptions::preserve_nulls`, which is why the planner wraps empty arrays with
+//! `ListEmptyToNullExpr` to get Spark's `explode_outer` semantics.
+
+use arrow::array::{
+ new_null_array, Array, ArrayRef, AsArray, FixedSizeListArray, Int64Array, LargeListArray,
+ LargeListViewArray, ListArray, ListViewArray, PrimitiveArray, Scalar, StructArray,
+};
+use arrow::compute::kernels::length::length;
+use arrow::compute::kernels::zip::zip;
+use arrow::compute::{cast, is_not_null, kernels, sum};
+use arrow::datatypes::{DataType, Int64Type, SchemaRef};
+use arrow::record_batch::RecordBatch;
+// Upstream imports this as `arrow_ord::cmp::lt`; Comet reaches it through `arrow`,
+// which does not have `arrow_ord` as a direct dependency.
+use arrow::compute::kernels::cmp::lt;
+use datafusion::common::{
+ exec_datafusion_err, exec_err, internal_err, HashMap, HashSet, Result, UnnestOptions,
+};
+use datafusion::execution::TaskContext;
+use datafusion::physical_expr::EquivalenceProperties;
+use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType};
+use datafusion::physical_plan::metrics::{
+ BaselineMetrics, Count, ExecutionPlanMetricsSet, MetricBuilder, MetricsSet, RecordOutput,
+};
+// `ListUnnest` is the one item the copied region below does NOT need to duplicate: unlike the
+// kernels, upstream exports it publicly.
+use datafusion::physical_plan::unnest::ListUnnest;
+use datafusion::physical_plan::{
+ DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties, RecordBatchStream,
+ SendableRecordBatchStream,
+};
+use futures::{Stream, StreamExt};
+use std::cmp::{self, Ordering};
+use std::pin::Pin;
+use std::sync::Arc;
+use std::task::{ready, Context, Poll};
+
+/// Comet's explode operator: DataFusion's `UnnestExec` with the input consumed in chunks so
+/// that output batches respect `datafusion.execution.batch_size`.
+#[derive(Debug)]
+pub struct ExplodeExec {
+ child: Arc<dyn ExecutionPlan>,
+ schema: SchemaRef,
+ list_column_indices: Vec<ListUnnest>,
+ struct_column_indices: Vec<usize>,
+ options: UnnestOptions,
+ metrics: ExecutionPlanMetricsSet,
+ cache: Arc<PlanProperties>,
+}
+
+impl ExplodeExec {
+ pub fn new(
+ child: Arc<dyn ExecutionPlan>,
+ list_column_indices: Vec<ListUnnest>,
+ struct_column_indices: Vec<usize>,
+ schema: SchemaRef,
+ options: UnnestOptions,
+ ) -> Self {
+ // Unnesting invalidates the child's orderings and constraints for the unnested
+ // columns, and Comet plans explode on a single partition, so start from empty
+ // equivalences rather than trying to project the child's.
+ let cache = Arc::new(PlanProperties::new(
+ EquivalenceProperties::new(Arc::clone(&schema)),
+ Partitioning::UnknownPartitioning(1),
+ EmissionType::Incremental,
+ Boundedness::Bounded,
+ ));
+
+ Self {
+ child,
+ schema,
+ list_column_indices,
+ struct_column_indices,
+ options,
+ metrics: ExecutionPlanMetricsSet::new(),
+ cache,
+ }
+ }
+}
+
+impl DisplayAs for ExplodeExec {
+ fn fmt_as(&self, t: DisplayFormatType, f: &mut std::fmt::Formatter) -> std::fmt::Result {
+ match t {
+ DisplayFormatType::Default | DisplayFormatType::Verbose => {
+ write!(f, "CometExplodeExec")
+ }
+ DisplayFormatType::TreeRender => unimplemented!(),
+ }
+ }
+}
+
+impl ExecutionPlan for ExplodeExec {
+ fn name(&self) -> &str {
+ "CometExplodeExec"
+ }
+
+ fn schema(&self) -> SchemaRef {
+ Arc::clone(&self.schema)
+ }
+
+ fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
+ vec![&self.child]
+ }
+
+ fn with_new_children(
+ self: Arc<Self>,
+ children: Vec<Arc<dyn ExecutionPlan>>,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ if children.len() != 1 {
+ return internal_err!("ExplodeExec expects exactly one child");
+ }
+ Ok(Arc::new(ExplodeExec::new(
+ Arc::clone(&children[0]),
+ self.list_column_indices.clone(),
+ self.struct_column_indices.clone(),
+ Arc::clone(&self.schema),
+ self.options.clone(),
+ )))
+ }
+
+ fn execute(
+ &self,
+ partition: usize,
+ context: Arc<TaskContext>,
+ ) -> Result<SendableRecordBatchStream> {
+ let batch_size = context.session_config().batch_size();
+ let input = self.child.execute(partition, context)?;
+
+ Ok(Box::pin(ExplodeStream {
+ input,
+ schema: Arc::clone(&self.schema),
+ list_type_columns: self.list_column_indices.clone(),
+ struct_column_indices: self.struct_column_indices.iter().copied().collect(),
+ options: self.options.clone(),
+ baseline_metrics: BaselineMetrics::new(&self.metrics, partition),
+ input_batches: MetricBuilder::new(&self.metrics).counter("input_batches", partition),
+ input_rows: MetricBuilder::new(&self.metrics).counter("input_rows", partition),
+ batch_size,
+ pending_input: None,
+ pending_output: None,
+ }))
+ }
+
+ fn properties(&self) -> &Arc<PlanProperties> {
+ &self.cache
+ }
+
+ fn metrics(&self) -> Option<MetricsSet> {
+ Some(self.metrics.clone_inner())
+ }
+}
+
+/// An input batch being unnested incrementally, a chunk of rows at a time.
+struct PendingInput {
+ /// The full input batch. Rows before `row_offset` have already been unnested.
+ batch: RecordBatch,
+ /// Index of the next input row to unnest.
+ row_offset: usize,
+ /// How many output rows each input row expands into, indexed by input row.
+ ///
+ /// Empty when the expansion factor cannot be predicted from the input alone, in which
+ /// case the whole remaining input is unnested in one call and only the output is split.
+ /// See [`ExplodeStream::predict_output_lens`].
+ output_lens: Vec<usize>,
+}
+
+impl PendingInput {
+ fn remaining_rows(&self) -> usize {
+ self.batch.num_rows() - self.row_offset
+ }
+
+ /// How many input rows to unnest next so the resulting batch holds at most `batch_size`
+ /// rows.
+ ///
+ /// Always returns at least 1 while rows remain, so the stream always makes progress: a
+ /// single input row is never split across output batches, so one row whose array is
+ /// longer than `batch_size` still produces one oversized build, which is then sliced
+ /// down on the way out.
+ fn next_chunk_rows(&self, batch_size: usize) -> usize {
+ let remaining = self.remaining_rows();
+ if self.output_lens.is_empty() {
+ return remaining;
+ }
+
+ let mut rows = 0;
+ let mut output_rows = 0usize;
+ while rows < remaining {
+ let len = self.output_lens[self.row_offset + rows];
+ if rows > 0 && output_rows.saturating_add(len) > batch_size {
+ break;
+ }
+ output_rows += len;
+ rows += 1;
+ }
+ rows
+ }
+}
+
+/// A stream that unnests its input, bounding output batches to `batch_size` rows.
+struct ExplodeStream {
+ input: SendableRecordBatchStream,
+ schema: SchemaRef,
+ list_type_columns: Vec<ListUnnest>,
+ struct_column_indices: HashSet<usize>,
+ options: UnnestOptions,
+ baseline_metrics: BaselineMetrics,
+ input_batches: Count,
+ input_rows: Count,
+ /// Target number of rows per output batch, from `datafusion.execution.batch_size`.
+ batch_size: usize,
+ /// Rows of the current input batch that have not been unnested yet.
+ pending_input: Option<PendingInput>,
+ /// Unnested rows that have been built but not emitted yet.
+ pending_output: Option<RecordBatch>,
+}
+
+impl RecordBatchStream for ExplodeStream {
+ fn schema(&self) -> SchemaRef {
+ Arc::clone(&self.schema)
+ }
+}
+
+impl Stream for ExplodeStream {
+ type Item = Result<RecordBatch>;
+
+ fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+ self.poll_next_impl(cx)
+ }
+}
+
+impl ExplodeStream {
+ fn poll_next_impl(&mut self, cx: &mut Context<'_>) -> Poll<Option<Result<RecordBatch>>> {
+ loop {
+ // Emit already-unnested rows first, at most `batch_size` at a time.
+ if let Some(batch) = self.pending_output.take() {
+ let (emit, rest) = split_off_head(batch, self.batch_size);
+ self.pending_output = rest;
+ (&emit).record_output(&self.baseline_metrics);
+ return Poll::Ready(Some(Ok(emit)));
+ }
+
+ // Unnest the next chunk of the input batch already in hand.
+ if let Some(pending) = self.pending_input.as_mut() {
+ if pending.remaining_rows() == 0 {
+ self.pending_input = None;
+ continue;
+ }
+
+ let timer = self.baseline_metrics.elapsed_compute().timer();
+
+ let rows = pending.next_chunk_rows(self.batch_size);
+ let chunk = pending.batch.slice(pending.row_offset, rows);
+ pending.row_offset += rows;
Review Comment:
**[P2] Drop consumed input before yielding its final output batch**
Could we clear `pending_input` immediately after `pending.row_offset` consumes its final row, before yielding output? As written, the final batch is returned while the full source `RecordBatch`, predicted lengths, and any full-batch `posexplode` positions remain live until the next poll. The upstream implementation checks whether the input is drained and drops `pending_input` before emitting. Preserving that order avoids holding the entire input during downstream or JVM processing of the final batch.
##########
native/core/src/execution/operators/explode.rs:
##########
@@ -0,0 +1,1237 @@
+// 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.
+
+//! A temporary fork of DataFusion's `UnnestExec` that respects
+//! `datafusion.execution.batch_size`.
+//!
+//! # Why this fork exists
+//!
+//! DataFusion's `UnnestExec` emits exactly one output batch per input batch, however many
+//! rows the unnesting produces, and never consults `batch_size`. For `explode` this means
+//! an 8192-row batch of 100-element arrays comes back as a single 819,200-row batch, and
+//! peak memory scales with input batch size times array length rather than with
+//! `batch_size`.
+//!
+//! The fix has been submitted upstream:
+//!
+//! * <https://github.com/apache/datafusion/issues/24383>
+//! * <https://github.com/apache/datafusion/pull/24384>
+//!
+//! # Deleting this file
+//!
+//! Once Comet moves to a DataFusion release carrying apache/datafusion#24384, delete this
+//! module and go back to `datafusion::physical_plan::unnest::UnnestExec` in the planner.
+//!
+//! Note that <https://github.com/apache/datafusion-comet/issues/5210> is a *different*
+//! unnest cleanup — it tracks adopting upstream `unnest_outer`
+//! (apache/datafusion#22100) to retire `ListEmptyToNullExpr`. The two upstream PRs can
+//! land in different releases, so closing 5210 is not a signal to delete this fork.
+//!
+//! # What was forked
+//!
+//! The unnesting kernels below (`build_batch` and everything it calls) are copied from
+//! `datafusion/physical-plan/src/unnest.rs` at DataFusion 54.1.0, upstream revision
+//! `cc7565be1ee97ba8fa2f5d6da373c5e38d81bb13`. They are private to
+//! `datafusion-physical-plan`, so they cannot be called from here without copying them.
+//! Leave them semantically unmodified so the eventual deletion is mechanical; the
+//! Comet-specific behavior lives entirely in `ExplodeExec` and `ExplodeStream`.
+//!
+//! They are not byte-identical to upstream: Comet's rustfmt uses `max_width = 100` and
+//! edition 2021, DataFusion's uses `max_width = 90` and edition 2024, so `cargo fmt`
+//! reflows some signatures. To audit for real changes, reformat this region at
+//! `max_width = 90` and diff it against upstream `unnest.rs`; that reduces the difference
+//! to a single cosmetic line wrap in `flatten_struct_cols`. The only deliberate edits are
+//! the `lt` import path noted below and dropping upstream's `ListUnnest` declaration in
+//! favor of importing the public one.
+//!
+//! Note that 54.1.0 predates upstream's `NullHandling` enum and still uses
+//! `UnnestOptions::preserve_nulls`, which is why the planner wraps empty arrays with
+//! `ListEmptyToNullExpr` to get Spark's `explode_outer` semantics.
+
+use arrow::array::{
+ new_null_array, Array, ArrayRef, AsArray, FixedSizeListArray, Int64Array, LargeListArray,
+ LargeListViewArray, ListArray, ListViewArray, PrimitiveArray, Scalar, StructArray,
+};
+use arrow::compute::kernels::length::length;
+use arrow::compute::kernels::zip::zip;
+use arrow::compute::{cast, is_not_null, kernels, sum};
+use arrow::datatypes::{DataType, Int64Type, SchemaRef};
+use arrow::record_batch::RecordBatch;
+// Upstream imports this as `arrow_ord::cmp::lt`; Comet reaches it through `arrow`,
+// which does not have `arrow_ord` as a direct dependency.
+use arrow::compute::kernels::cmp::lt;
+use datafusion::common::{
+ exec_datafusion_err, exec_err, internal_err, HashMap, HashSet, Result, UnnestOptions,
+};
+use datafusion::execution::TaskContext;
+use datafusion::physical_expr::EquivalenceProperties;
+use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType};
+use datafusion::physical_plan::metrics::{
+ BaselineMetrics, Count, ExecutionPlanMetricsSet, MetricBuilder, MetricsSet, RecordOutput,
+};
+// `ListUnnest` is the one item the copied region below does NOT need to duplicate: unlike the
+// kernels, upstream exports it publicly.
+use datafusion::physical_plan::unnest::ListUnnest;
+use datafusion::physical_plan::{
+ DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties, RecordBatchStream,
+ SendableRecordBatchStream,
+};
+use futures::{Stream, StreamExt};
+use std::cmp::{self, Ordering};
+use std::pin::Pin;
+use std::sync::Arc;
+use std::task::{ready, Context, Poll};
+
+/// Comet's explode operator: DataFusion's `UnnestExec` with the input consumed in chunks so
+/// that output batches respect `datafusion.execution.batch_size`.
+#[derive(Debug)]
+pub struct ExplodeExec {
+ child: Arc<dyn ExecutionPlan>,
+ schema: SchemaRef,
+ list_column_indices: Vec<ListUnnest>,
+ struct_column_indices: Vec<usize>,
+ options: UnnestOptions,
+ metrics: ExecutionPlanMetricsSet,
+ cache: Arc<PlanProperties>,
+}
+
+impl ExplodeExec {
+ pub fn new(
+ child: Arc<dyn ExecutionPlan>,
+ list_column_indices: Vec<ListUnnest>,
+ struct_column_indices: Vec<usize>,
+ schema: SchemaRef,
+ options: UnnestOptions,
+ ) -> Self {
+ // Unnesting invalidates the child's orderings and constraints for the unnested
+ // columns, and Comet plans explode on a single partition, so start from empty
+ // equivalences rather than trying to project the child's.
+ let cache = Arc::new(PlanProperties::new(
+ EquivalenceProperties::new(Arc::clone(&schema)),
Review Comment:
**[P2] Preserve physical properties for passthrough columns**
Could we preserve the projected physical properties from the original `UnnestExec::compute_properties`? It retains ordering and equivalences for passthrough columns, but starting from empty `EquivalenceProperties` loses them. For `SortMergeJoin -> Explode -> GROUP BY k`, a sorted passthrough key `k` now makes the downstream `AggregateExec` choose `InputOrderMode::Linear` instead of `Sorted`, so groups accumulate or spill instead of streaming. This increases memory consumption and runtime on existing native plans.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix: make CometExplodeExec respect batch size [datafusion-comet]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove commented on code in PR #5362:
URL: https://github.com/apache/datafusion-comet/pull/5362#discussion_r3799036865
##########
native/core/src/execution/operators/explode.rs:
##########
@@ -0,0 +1,1237 @@
+// 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.
+
+//! A temporary fork of DataFusion's `UnnestExec` that respects
+//! `datafusion.execution.batch_size`.
+//!
+//! # Why this fork exists
+//!
+//! DataFusion's `UnnestExec` emits exactly one output batch per input batch, however many
+//! rows the unnesting produces, and never consults `batch_size`. For `explode` this means
+//! an 8192-row batch of 100-element arrays comes back as a single 819,200-row batch, and
+//! peak memory scales with input batch size times array length rather than with
+//! `batch_size`.
+//!
+//! The fix has been submitted upstream:
+//!
+//! * <https://github.com/apache/datafusion/issues/24383>
+//! * <https://github.com/apache/datafusion/pull/24384>
+//!
+//! # Deleting this file
+//!
+//! Once Comet moves to a DataFusion release carrying apache/datafusion#24384, delete this
+//! module and go back to `datafusion::physical_plan::unnest::UnnestExec` in the planner.
+//!
+//! Note that <https://github.com/apache/datafusion-comet/issues/5210> is a *different*
+//! unnest cleanup — it tracks adopting upstream `unnest_outer`
+//! (apache/datafusion#22100) to retire `ListEmptyToNullExpr`. The two upstream PRs can
+//! land in different releases, so closing 5210 is not a signal to delete this fork.
+//!
+//! # What was forked
+//!
+//! The unnesting kernels below (`build_batch` and everything it calls) are copied from
+//! `datafusion/physical-plan/src/unnest.rs` at DataFusion 54.1.0, upstream revision
+//! `cc7565be1ee97ba8fa2f5d6da373c5e38d81bb13`. They are private to
+//! `datafusion-physical-plan`, so they cannot be called from here without copying them.
+//! Leave them semantically unmodified so the eventual deletion is mechanical; the
+//! Comet-specific behavior lives entirely in `ExplodeExec` and `ExplodeStream`.
+//!
+//! They are not byte-identical to upstream: Comet's rustfmt uses `max_width = 100` and
+//! edition 2021, DataFusion's uses `max_width = 90` and edition 2024, so `cargo fmt`
+//! reflows some signatures. To audit for real changes, reformat this region at
+//! `max_width = 90` and diff it against upstream `unnest.rs`; that reduces the difference
+//! to a single cosmetic line wrap in `flatten_struct_cols`. The only deliberate edits are
+//! the `lt` import path noted below and dropping upstream's `ListUnnest` declaration in
+//! favor of importing the public one.
+//!
+//! Note that 54.1.0 predates upstream's `NullHandling` enum and still uses
+//! `UnnestOptions::preserve_nulls`, which is why the planner wraps empty arrays with
+//! `ListEmptyToNullExpr` to get Spark's `explode_outer` semantics.
+
+use arrow::array::{
+ new_null_array, Array, ArrayRef, AsArray, FixedSizeListArray, Int64Array, LargeListArray,
+ LargeListViewArray, ListArray, ListViewArray, PrimitiveArray, Scalar, StructArray,
+};
+use arrow::compute::kernels::length::length;
+use arrow::compute::kernels::zip::zip;
+use arrow::compute::{cast, is_not_null, kernels, sum};
+use arrow::datatypes::{DataType, Int64Type, SchemaRef};
+use arrow::record_batch::RecordBatch;
+// Upstream imports this as `arrow_ord::cmp::lt`; Comet reaches it through `arrow`,
+// which does not have `arrow_ord` as a direct dependency.
+use arrow::compute::kernels::cmp::lt;
+use datafusion::common::{
+ exec_datafusion_err, exec_err, internal_err, HashMap, HashSet, Result, UnnestOptions,
+};
+use datafusion::execution::TaskContext;
+use datafusion::physical_expr::EquivalenceProperties;
+use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType};
+use datafusion::physical_plan::metrics::{
+ BaselineMetrics, Count, ExecutionPlanMetricsSet, MetricBuilder, MetricsSet, RecordOutput,
+};
+// `ListUnnest` is the one item the copied region below does NOT need to duplicate: unlike the
+// kernels, upstream exports it publicly.
+use datafusion::physical_plan::unnest::ListUnnest;
+use datafusion::physical_plan::{
+ DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties, RecordBatchStream,
+ SendableRecordBatchStream,
+};
+use futures::{Stream, StreamExt};
+use std::cmp::{self, Ordering};
+use std::pin::Pin;
+use std::sync::Arc;
+use std::task::{ready, Context, Poll};
+
+/// Comet's explode operator: DataFusion's `UnnestExec` with the input consumed in chunks so
+/// that output batches respect `datafusion.execution.batch_size`.
+#[derive(Debug)]
+pub struct ExplodeExec {
+ child: Arc<dyn ExecutionPlan>,
+ schema: SchemaRef,
+ list_column_indices: Vec<ListUnnest>,
+ struct_column_indices: Vec<usize>,
+ options: UnnestOptions,
+ metrics: ExecutionPlanMetricsSet,
+ cache: Arc<PlanProperties>,
+}
+
+impl ExplodeExec {
+ pub fn new(
+ child: Arc<dyn ExecutionPlan>,
+ list_column_indices: Vec<ListUnnest>,
+ struct_column_indices: Vec<usize>,
+ schema: SchemaRef,
+ options: UnnestOptions,
+ ) -> Self {
+ // Unnesting invalidates the child's orderings and constraints for the unnested
+ // columns, and Comet plans explode on a single partition, so start from empty
+ // equivalences rather than trying to project the child's.
+ let cache = Arc::new(PlanProperties::new(
+ EquivalenceProperties::new(Arc::clone(&schema)),
+ Partitioning::UnknownPartitioning(1),
+ EmissionType::Incremental,
+ Boundedness::Bounded,
+ ));
+
+ Self {
+ child,
+ schema,
+ list_column_indices,
+ struct_column_indices,
+ options,
+ metrics: ExecutionPlanMetricsSet::new(),
+ cache,
+ }
+ }
+}
+
+impl DisplayAs for ExplodeExec {
+ fn fmt_as(&self, t: DisplayFormatType, f: &mut std::fmt::Formatter) -> std::fmt::Result {
+ match t {
+ DisplayFormatType::Default | DisplayFormatType::Verbose => {
+ write!(f, "CometExplodeExec")
+ }
+ DisplayFormatType::TreeRender => unimplemented!(),
+ }
+ }
+}
+
+impl ExecutionPlan for ExplodeExec {
+ fn name(&self) -> &str {
+ "CometExplodeExec"
+ }
+
+ fn schema(&self) -> SchemaRef {
+ Arc::clone(&self.schema)
+ }
+
+ fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
+ vec![&self.child]
+ }
+
+ fn with_new_children(
+ self: Arc<Self>,
+ children: Vec<Arc<dyn ExecutionPlan>>,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ if children.len() != 1 {
+ return internal_err!("ExplodeExec expects exactly one child");
+ }
+ Ok(Arc::new(ExplodeExec::new(
+ Arc::clone(&children[0]),
+ self.list_column_indices.clone(),
+ self.struct_column_indices.clone(),
+ Arc::clone(&self.schema),
+ self.options.clone(),
+ )))
+ }
+
+ fn execute(
+ &self,
+ partition: usize,
+ context: Arc<TaskContext>,
+ ) -> Result<SendableRecordBatchStream> {
+ let batch_size = context.session_config().batch_size();
+ let input = self.child.execute(partition, context)?;
+
+ Ok(Box::pin(ExplodeStream {
+ input,
+ schema: Arc::clone(&self.schema),
+ list_type_columns: self.list_column_indices.clone(),
+ struct_column_indices: self.struct_column_indices.iter().copied().collect(),
+ options: self.options.clone(),
+ baseline_metrics: BaselineMetrics::new(&self.metrics, partition),
+ input_batches: MetricBuilder::new(&self.metrics).counter("input_batches", partition),
+ input_rows: MetricBuilder::new(&self.metrics).counter("input_rows", partition),
+ batch_size,
+ pending_input: None,
+ pending_output: None,
+ }))
+ }
+
+ fn properties(&self) -> &Arc<PlanProperties> {
+ &self.cache
+ }
+
+ fn metrics(&self) -> Option<MetricsSet> {
+ Some(self.metrics.clone_inner())
+ }
+}
+
+/// An input batch being unnested incrementally, a chunk of rows at a time.
+struct PendingInput {
+ /// The full input batch. Rows before `row_offset` have already been unnested.
+ batch: RecordBatch,
+ /// Index of the next input row to unnest.
+ row_offset: usize,
+ /// How many output rows each input row expands into, indexed by input row.
+ ///
+ /// Empty when the expansion factor cannot be predicted from the input alone, in which
+ /// case the whole remaining input is unnested in one call and only the output is split.
+ /// See [`ExplodeStream::predict_output_lens`].
+ output_lens: Vec<usize>,
+}
+
+impl PendingInput {
+ fn remaining_rows(&self) -> usize {
+ self.batch.num_rows() - self.row_offset
+ }
+
+ /// How many input rows to unnest next so the resulting batch holds at most `batch_size`
+ /// rows.
+ ///
+ /// Always returns at least 1 while rows remain, so the stream always makes progress: a
+ /// single input row is never split across output batches, so one row whose array is
+ /// longer than `batch_size` still produces one oversized build, which is then sliced
+ /// down on the way out.
+ fn next_chunk_rows(&self, batch_size: usize) -> usize {
+ let remaining = self.remaining_rows();
+ if self.output_lens.is_empty() {
+ return remaining;
+ }
+
+ let mut rows = 0;
+ let mut output_rows = 0usize;
+ while rows < remaining {
+ let len = self.output_lens[self.row_offset + rows];
+ if rows > 0 && output_rows.saturating_add(len) > batch_size {
+ break;
+ }
+ output_rows += len;
+ rows += 1;
+ }
+ rows
+ }
+}
+
+/// A stream that unnests its input, bounding output batches to `batch_size` rows.
+struct ExplodeStream {
+ input: SendableRecordBatchStream,
+ schema: SchemaRef,
+ list_type_columns: Vec<ListUnnest>,
+ struct_column_indices: HashSet<usize>,
+ options: UnnestOptions,
+ baseline_metrics: BaselineMetrics,
+ input_batches: Count,
+ input_rows: Count,
+ /// Target number of rows per output batch, from `datafusion.execution.batch_size`.
+ batch_size: usize,
+ /// Rows of the current input batch that have not been unnested yet.
+ pending_input: Option<PendingInput>,
+ /// Unnested rows that have been built but not emitted yet.
+ pending_output: Option<RecordBatch>,
+}
+
+impl RecordBatchStream for ExplodeStream {
+ fn schema(&self) -> SchemaRef {
+ Arc::clone(&self.schema)
+ }
+}
+
+impl Stream for ExplodeStream {
+ type Item = Result<RecordBatch>;
+
+ fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+ self.poll_next_impl(cx)
+ }
+}
+
+impl ExplodeStream {
+ fn poll_next_impl(&mut self, cx: &mut Context<'_>) -> Poll<Option<Result<RecordBatch>>> {
+ loop {
+ // Emit already-unnested rows first, at most `batch_size` at a time.
+ if let Some(batch) = self.pending_output.take() {
+ let (emit, rest) = split_off_head(batch, self.batch_size);
+ self.pending_output = rest;
+ (&emit).record_output(&self.baseline_metrics);
+ return Poll::Ready(Some(Ok(emit)));
+ }
+
+ // Unnest the next chunk of the input batch already in hand.
+ if let Some(pending) = self.pending_input.as_mut() {
+ if pending.remaining_rows() == 0 {
+ self.pending_input = None;
+ continue;
+ }
+
+ let timer = self.baseline_metrics.elapsed_compute().timer();
+
+ let rows = pending.next_chunk_rows(self.batch_size);
+ let chunk = pending.batch.slice(pending.row_offset, rows);
+ pending.row_offset += rows;
+
+ let result = build_batch(
+ &chunk,
+ &self.schema,
+ &self.list_type_columns,
+ &self.struct_column_indices,
+ &self.options,
+ );
+ timer.done();
+
+ // A chunk can legitimately produce no rows at all (for example rows whose
+ // arrays are all empty and `preserve_nulls` is false); move on to the next
+ // chunk rather than emitting an empty batch.
+ self.pending_output = result?.filter(|batch| batch.num_rows() > 0);
+ continue;
+ }
+
+ // Otherwise pull the next input batch.
+ return Poll::Ready(match ready!(self.input.poll_next_unpin(cx)) {
+ Some(Ok(batch)) => {
+ self.input_batches.add(1);
+ self.input_rows.add(batch.num_rows());
+ if batch.num_rows() == 0 {
+ continue;
+ }
+
+ let timer = self.baseline_metrics.elapsed_compute().timer();
+ let output_lens = self.predict_output_lens(&batch);
+ timer.done();
+
+ match output_lens {
+ Ok(output_lens) => {
+ self.pending_input = Some(PendingInput {
+ batch,
+ row_offset: 0,
+ output_lens,
+ });
+ continue;
+ }
+ Err(e) => Some(Err(e)),
+ }
+ }
+ other => other,
Review Comment:
Fixed, and thanks for spotting it — this was the fork drifting from the upstream PR rather than a deliberate choice, which is true of all four of these.
`ExplodeStream` now swaps in `EmptyRecordBatchStream` when the child is depleted, same as `UnnestStream`. Took your suggestion on the test too: `releases_exhausted_child_at_eof` drives a child stream carrying a drop flag and asserts the flag flips on the poll that reports EOF, not when the stream is finally dropped. It fails on the previous code.
Related: the output now goes through DataFusion's public `BatchSplitStream` instead of the hand-rolled `pending_output`/`split_off_head` pair, which is what upstream uses for the same job. That releases its input on EOF as well, so the whole chain unwinds at the same poll.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix: make CometExplodeExec respect batch size [datafusion-comet]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove commented on code in PR #5362:
URL: https://github.com/apache/datafusion-comet/pull/5362#discussion_r3799038124
##########
native/core/src/execution/operators/explode.rs:
##########
@@ -0,0 +1,1237 @@
+// 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.
+
+//! A temporary fork of DataFusion's `UnnestExec` that respects
+//! `datafusion.execution.batch_size`.
+//!
+//! # Why this fork exists
+//!
+//! DataFusion's `UnnestExec` emits exactly one output batch per input batch, however many
+//! rows the unnesting produces, and never consults `batch_size`. For `explode` this means
+//! an 8192-row batch of 100-element arrays comes back as a single 819,200-row batch, and
+//! peak memory scales with input batch size times array length rather than with
+//! `batch_size`.
+//!
+//! The fix has been submitted upstream:
+//!
+//! * <https://github.com/apache/datafusion/issues/24383>
+//! * <https://github.com/apache/datafusion/pull/24384>
+//!
+//! # Deleting this file
+//!
+//! Once Comet moves to a DataFusion release carrying apache/datafusion#24384, delete this
+//! module and go back to `datafusion::physical_plan::unnest::UnnestExec` in the planner.
+//!
+//! Note that <https://github.com/apache/datafusion-comet/issues/5210> is a *different*
+//! unnest cleanup — it tracks adopting upstream `unnest_outer`
+//! (apache/datafusion#22100) to retire `ListEmptyToNullExpr`. The two upstream PRs can
+//! land in different releases, so closing 5210 is not a signal to delete this fork.
+//!
+//! # What was forked
+//!
+//! The unnesting kernels below (`build_batch` and everything it calls) are copied from
+//! `datafusion/physical-plan/src/unnest.rs` at DataFusion 54.1.0, upstream revision
+//! `cc7565be1ee97ba8fa2f5d6da373c5e38d81bb13`. They are private to
+//! `datafusion-physical-plan`, so they cannot be called from here without copying them.
+//! Leave them semantically unmodified so the eventual deletion is mechanical; the
+//! Comet-specific behavior lives entirely in `ExplodeExec` and `ExplodeStream`.
+//!
+//! They are not byte-identical to upstream: Comet's rustfmt uses `max_width = 100` and
+//! edition 2021, DataFusion's uses `max_width = 90` and edition 2024, so `cargo fmt`
+//! reflows some signatures. To audit for real changes, reformat this region at
+//! `max_width = 90` and diff it against upstream `unnest.rs`; that reduces the difference
+//! to a single cosmetic line wrap in `flatten_struct_cols`. The only deliberate edits are
+//! the `lt` import path noted below and dropping upstream's `ListUnnest` declaration in
+//! favor of importing the public one.
+//!
+//! Note that 54.1.0 predates upstream's `NullHandling` enum and still uses
+//! `UnnestOptions::preserve_nulls`, which is why the planner wraps empty arrays with
+//! `ListEmptyToNullExpr` to get Spark's `explode_outer` semantics.
+
+use arrow::array::{
+ new_null_array, Array, ArrayRef, AsArray, FixedSizeListArray, Int64Array, LargeListArray,
+ LargeListViewArray, ListArray, ListViewArray, PrimitiveArray, Scalar, StructArray,
+};
+use arrow::compute::kernels::length::length;
+use arrow::compute::kernels::zip::zip;
+use arrow::compute::{cast, is_not_null, kernels, sum};
+use arrow::datatypes::{DataType, Int64Type, SchemaRef};
+use arrow::record_batch::RecordBatch;
+// Upstream imports this as `arrow_ord::cmp::lt`; Comet reaches it through `arrow`,
+// which does not have `arrow_ord` as a direct dependency.
+use arrow::compute::kernels::cmp::lt;
+use datafusion::common::{
+ exec_datafusion_err, exec_err, internal_err, HashMap, HashSet, Result, UnnestOptions,
+};
+use datafusion::execution::TaskContext;
+use datafusion::physical_expr::EquivalenceProperties;
+use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType};
+use datafusion::physical_plan::metrics::{
+ BaselineMetrics, Count, ExecutionPlanMetricsSet, MetricBuilder, MetricsSet, RecordOutput,
+};
+// `ListUnnest` is the one item the copied region below does NOT need to duplicate: unlike the
+// kernels, upstream exports it publicly.
+use datafusion::physical_plan::unnest::ListUnnest;
+use datafusion::physical_plan::{
+ DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties, RecordBatchStream,
+ SendableRecordBatchStream,
+};
+use futures::{Stream, StreamExt};
+use std::cmp::{self, Ordering};
+use std::pin::Pin;
+use std::sync::Arc;
+use std::task::{ready, Context, Poll};
+
+/// Comet's explode operator: DataFusion's `UnnestExec` with the input consumed in chunks so
+/// that output batches respect `datafusion.execution.batch_size`.
+#[derive(Debug)]
+pub struct ExplodeExec {
+ child: Arc<dyn ExecutionPlan>,
+ schema: SchemaRef,
+ list_column_indices: Vec<ListUnnest>,
+ struct_column_indices: Vec<usize>,
+ options: UnnestOptions,
+ metrics: ExecutionPlanMetricsSet,
+ cache: Arc<PlanProperties>,
+}
+
+impl ExplodeExec {
+ pub fn new(
+ child: Arc<dyn ExecutionPlan>,
+ list_column_indices: Vec<ListUnnest>,
+ struct_column_indices: Vec<usize>,
+ schema: SchemaRef,
+ options: UnnestOptions,
+ ) -> Self {
+ // Unnesting invalidates the child's orderings and constraints for the unnested
+ // columns, and Comet plans explode on a single partition, so start from empty
+ // equivalences rather than trying to project the child's.
+ let cache = Arc::new(PlanProperties::new(
+ EquivalenceProperties::new(Arc::clone(&schema)),
Review Comment:
You are right, and my comment justifying the empty properties was wrong on both counts. Unnesting only rewrites the list and struct columns, so whatever the child guarantees about the rest still holds, and hardcoding `UnknownPartitioning(1)` because Comet happens to plan explode on one partition is not the operator's call to make.
`compute_properties` now mirrors `UnnestExec::compute_properties`: build a `ProjectionMapping` over the non-unnested indices, project the child's equivalences and partitioning through it, and drop only the constraints, since row duplication genuinely does invalidate uniqueness and primary keys. `new()` is fallible now as a consequence.
`preserves_passthrough_orderings` pins it — a sorted passthrough key over a `DataSourceExec` with sort information, asserting the ordering is still there on the `ExplodeExec` above it. It fails if the properties go back to empty.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org