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: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
