You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by ag...@apache.org on 2026/08/07 02:14:04 UTC
(datafusion-comet) branch main updated: fix: normalize nested field nullability in ShuffleScanExec and ExpandExec (#5138)
This is an automated email from the ASF dual-hosted git repository.
andygrove pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
The following commit(s) were added to refs/heads/main by this push:
new cd4d0e25d6 fix: normalize nested field nullability in ShuffleScanExec and ExpandExec (#5138)
cd4d0e25d6 is described below
commit cd4d0e25d68510acb83cd637dd4d12ed2ca28123
Author: Andy Grove <ag...@apache.org>
AuthorDate: Thu Aug 6 20:13:59 2026 -0600
fix: normalize nested field nullability in ShuffleScanExec and ExpandExec (#5138)
* fix: normalize nested field nullability in ShuffleScanExec and ExpandExec
Arrow treats nested field nullability as part of `DataType` identity, so
`RecordBatch::try_new_with_options` rejects a column whose nested `nullable`
flags are narrower than the declared type even though the data is fine — a
non-null child is a strict subset of a nullable one.
Most Comet boundaries already normalize instead of asserting (`ScanExec` casts
at the FFI boundary, `SchemaAlignExec` casts the shuffle writer's input,
`CometLocalTableScanExec` widens with `.asNullable`), which is why this drift is
usually invisible. Two stamp points still asserted and aborted the task:
- `ShuffleScanExec` stamped the catalyst schema onto the decoded shuffle block
with no cast. This is the direct-read shuffle path, enabled by default.
- `ExpandExec` derived its schema from `projections[0]` only and stamped it onto
the output of every projection, so it failed if any two projections disagreed
on a nested nullable flag for the same output column.
Add `cast_and_stamp_schema` to `datafusion-comet-common` and use it at both
sites, so the drift is absorbed the way the other boundaries absorb it. Also
derive `ExpandExec`'s schema by widening nested nullability across all
projections rather than reading it off the first, so the declared schema is
valid for every projection.
When a stamp genuinely cannot be reconciled, the error now names the operator
and the dotted path of the offending column instead of only `at column index N`,
which was very hard to act on for wide schemas of deeply nested structs.
Upstream nested-nullability drift itself is tracked by #4515.
Closes #5137
* refactor: address /simplify review of the nested-nullability fix
Reuse:
- Route `SchemaAlignExec::align` through the new shared `cast_and_stamp_schema`
instead of keeping its own copy of the same cast-then-stamp loop. This deletes
the `ColumnAction` enum and its plan-time precomputation and plumbing; the
per-batch type comparison it avoided is negligible against batch-sized work,
and the operator gets the improved error messages for free.
- Share one nested-nullability test fixture between the `expand` and
`shuffle_scan` tests rather than defining the same list-of-struct builder in
each.
- Use `datafusion::physical_plan::collect` and
`MemorySourceConfig::try_new_exec` instead of hand-rolled equivalents.
Simplification:
- Drop the `residual_mismatch` tracking from `cast_and_stamp_schema`. It guarded
against `cast_with_options` returning a type other than the one requested,
which it never does — it returns that type or an error — so after the cast loop
the stamp can only fail on row counts.
- Drop the `ListView`/`LargeListView`/`Dictionary` arms and the nested-metadata
comparison. The reachable nested set is Struct/List/LargeList/FixedSizeList/Map
— the same set `make_all_fields_nullable` walks — and a metadata-only
difference always casts cleanly, so the describer never sees one.
- Make `describe_type_mismatch` private; nothing outside the module used it.
- Drop the planner's `assert!(!projections.is_empty())`, now that
`build_schema` returns a proper error for the same condition.
Efficiency:
- Take `columns` by value so `ExpandStream::expand` moves its vector in as it did
before this change, instead of the helper allocating a second one per
projection per batch.
Documentation:
- Record that reconciliation follows the declared schema in both directions and
will narrow as well as widen, and point callers that must not narrow at
`widen_nested_nullability`. Cross-reference the pre-existing
`make_all_fields_nullable` and arrow's `Field::try_merge`, noting why neither
substitutes for it.
* docs: explain why Expand's output field names stay col_{idx}
Addresses review feedback on #5138: the naming is inherited verbatim from
the planner, and parent operators bind to these columns by index, so the
name is positional either way.
---
native/common/src/lib.rs | 2 +
native/common/src/schema.rs | 378 +++++++++++++++++++++
native/core/src/execution/operators/expand.rs | 199 ++++++++++-
native/core/src/execution/operators/mod.rs | 49 +++
.../core/src/execution/operators/shuffle_scan.rs | 94 ++++-
native/core/src/execution/planner.rs | 16 +-
native/shuffle/src/schema_align.rs | 68 ++--
7 files changed, 731 insertions(+), 75 deletions(-)
diff --git a/native/common/src/lib.rs b/native/common/src/lib.rs
index b2d4a57431..c4153606a1 100644
--- a/native/common/src/lib.rs
+++ b/native/common/src/lib.rs
@@ -17,9 +17,11 @@
mod error;
mod query_context;
+mod schema;
pub mod tracing;
mod utils;
pub use error::{decimal_overflow_error, SparkError, SparkErrorWithContext, SparkResult};
pub use query_context::{create_query_context_map, QueryContext, QueryContextMap};
+pub use schema::{cast_and_stamp_schema, widen_nested_nullability};
pub use utils::{bytes_to_i128, decode_utf8_spark_lossy};
diff --git a/native/common/src/schema.rs b/native/common/src/schema.rs
new file mode 100644
index 0000000000..2b5995c8c2
--- /dev/null
+++ b/native/common/src/schema.rs
@@ -0,0 +1,378 @@
+// 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.
+
+//! Helpers for reconciling the Arrow type an operator actually produced with the type Spark
+//! catalyst declared.
+//!
+//! Arrow treats nested field nullability as part of `DataType` identity, so a column whose nested
+//! `nullable` flags are narrower than the declared type is rejected by
+//! `RecordBatch::try_new_with_options` even though the data is fine — a non-null child is a strict
+//! subset of a nullable one. Every Comet boundary that stamps a declared schema onto a runtime
+//! array therefore normalizes first instead of asserting. See
+//! <https://github.com/apache/datafusion-comet/issues/5137> for the boundaries and
+//! <https://github.com/apache/datafusion-comet/issues/4515> for the upstream drift itself.
+
+use arrow::array::{ArrayRef, RecordBatch, RecordBatchOptions};
+use arrow::compute::{cast_with_options, CastOptions};
+use arrow::datatypes::{DataType, FieldRef, SchemaRef};
+use datafusion::common::DataFusionError;
+use std::sync::Arc;
+
+/// Builds a `RecordBatch` with `schema` from `columns`, casting any column whose Arrow type
+/// differs from the declared field type. `operator` names the caller and is used only in error
+/// messages.
+///
+/// This is the normalizing counterpart to stamping `schema` on directly: it absorbs the
+/// return-type drift — most commonly a nested `nullable` flag — that native kernels and non-Parquet
+/// sources introduce, the same way `ScanExec` absorbs it at the FFI boundary.
+///
+/// Reconciliation follows `schema` in both directions, so it will also narrow a nullable nested
+/// child to non-null when that is what `schema` declares. That is not silently lossy: arrow's
+/// `StructArray::try_new` rejects unmasked nulls under a non-nullable field, so data that cannot
+/// survive the narrowing errors here rather than producing an array that misreports itself.
+/// Callers that must not narrow should widen `schema` first — see [`widen_nested_nullability`],
+/// which is what `SchemaAlignExec` and `ExpandExec` do.
+pub fn cast_and_stamp_schema(
+ operator: &str,
+ schema: &SchemaRef,
+ mut columns: Vec<ArrayRef>,
+ num_rows: usize,
+) -> Result<RecordBatch, DataFusionError> {
+ if columns.len() != schema.fields().len() {
+ return Err(DataFusionError::Internal(format!(
+ "{operator} produced {} columns but its schema declares {}",
+ columns.len(),
+ schema.fields().len()
+ )));
+ }
+
+ for (idx, (column, field)) in columns.iter_mut().zip(schema.fields()).enumerate() {
+ if column.data_type() != field.data_type() {
+ *column = cast_with_options(column, field.data_type(), &CastOptions::default())
+ .map_err(|e| cast_error(operator, schema, idx, column.data_type(), e))?;
+ }
+ }
+
+ // Every column's type now matches its declared field — `cast_with_options` returns either an
+ // error or an array of exactly the requested type — so the stamp can only fail on row counts.
+ let options = RecordBatchOptions::new().with_row_count(Some(num_rows));
+ RecordBatch::try_new_with_options(Arc::clone(schema), columns, &options).map_err(|e| {
+ DataFusionError::Context(
+ format!("{operator} cannot build a batch of {num_rows} rows with its declared schema"),
+ Box::new(DataFusionError::from(e)),
+ )
+ })
+}
+
+/// Names the operator and the dotted path of the column that could not be reconciled, since
+/// arrow's own message reports only `at column index N` and the two printed types may differ by a
+/// single flag hundreds of characters in.
+fn cast_error(
+ operator: &str,
+ schema: &SchemaRef,
+ idx: usize,
+ actual: &DataType,
+ source: arrow::error::ArrowError,
+) -> DataFusionError {
+ let field = schema.field(idx);
+ let detail = describe_type_mismatch(field.name(), field.data_type(), actual)
+ .unwrap_or_else(|| format!("{}: expected {}", field.name(), field.data_type()));
+ DataFusionError::Context(
+ format!("{operator} cannot reconcile col[{idx}] with its declared schema at {detail}"),
+ Box::new(DataFusionError::from(source)),
+ )
+}
+
+/// Describes where `expected` and `actual` first diverge, as a dotted path rooted at `path`.
+/// Returns `None` when the two types are equal.
+///
+/// The nested arms cover the shapes Comet actually builds — see `make_all_fields_nullable` in the
+/// planner and `to_arrow_datatype` in the serde layer, which walk the same set.
+fn describe_type_mismatch(path: &str, expected: &DataType, actual: &DataType) -> Option<String> {
+ if expected == actual {
+ return None;
+ }
+ match (expected, actual) {
+ (DataType::List(e), DataType::List(a))
+ | (DataType::LargeList(e), DataType::LargeList(a)) => {
+ describe_field_mismatch(&format!("{path}.element"), e, a)
+ }
+ (DataType::FixedSizeList(e, e_len), DataType::FixedSizeList(a, a_len))
+ if e_len == a_len =>
+ {
+ describe_field_mismatch(&format!("{path}.element"), e, a)
+ }
+ (DataType::Map(e, e_sorted), DataType::Map(a, a_sorted)) if e_sorted == a_sorted => {
+ describe_field_mismatch(&format!("{path}.entries"), e, a)
+ }
+ (DataType::Struct(e), DataType::Struct(a)) if e.len() == a.len() => e
+ .iter()
+ .zip(a.iter())
+ .find_map(|(e, a)| describe_field_mismatch(&format!("{path}.{}", e.name()), e, a)),
+ _ => Some(format!("{path}: expected {expected}, found {actual}")),
+ }
+}
+
+fn describe_field_mismatch(path: &str, expected: &FieldRef, actual: &FieldRef) -> Option<String> {
+ if expected.name() != actual.name() {
+ return Some(format!(
+ "{path}: expected field name '{}', found '{}'",
+ expected.name(),
+ actual.name()
+ ));
+ }
+ if expected.is_nullable() != actual.is_nullable() {
+ return Some(format!(
+ "{path}: expected {}, found {}",
+ nullability(expected),
+ nullability(actual)
+ ));
+ }
+ describe_type_mismatch(path, expected.data_type(), actual.data_type())
+}
+
+fn nullability(field: &FieldRef) -> String {
+ let qualifier = if field.is_nullable() {
+ "nullable"
+ } else {
+ "non-null"
+ };
+ format!("{qualifier} {}", field.data_type())
+}
+
+/// Returns `base` with nested field nullability widened to also cover `other`, so that arrays of
+/// either type can be stamped with the result. Shapes that do not line up (different struct field
+/// counts, different list flavours, ...) are left as `base`; those are real type differences and
+/// are handled by the cast in [`cast_and_stamp_schema`].
+///
+/// Related but not interchangeable: `make_all_fields_nullable` in the planner widens one type
+/// unconditionally, and arrow's `Field::try_merge` unions two but errors on a leaf type difference
+/// instead of tolerating it and does not recurse into maps or fixed-size lists.
+pub fn widen_nested_nullability(base: &DataType, other: &DataType) -> DataType {
+ match (base, other) {
+ (DataType::List(b), DataType::List(o)) => DataType::List(widen_field(b, o)),
+ (DataType::LargeList(b), DataType::LargeList(o)) => DataType::LargeList(widen_field(b, o)),
+ (DataType::FixedSizeList(b, b_len), DataType::FixedSizeList(o, o_len))
+ if b_len == o_len =>
+ {
+ DataType::FixedSizeList(widen_field(b, o), *b_len)
+ }
+ (DataType::Map(b, b_sorted), DataType::Map(o, o_sorted)) if b_sorted == o_sorted => {
+ DataType::Map(widen_field(b, o), *b_sorted)
+ }
+ (DataType::Struct(b), DataType::Struct(o)) if b.len() == o.len() => DataType::Struct(
+ b.iter()
+ .zip(o.iter())
+ .map(|(b, o)| widen_field(b, o))
+ .collect(),
+ ),
+ _ => base.clone(),
+ }
+}
+
+/// Widens a single field, keeping `base`'s name and metadata. A map's `entries` field must stay
+/// non-nullable in Arrow, so only the nullability of fields that are already nullable on either
+/// side is propagated — which is exactly `base.nullable || other.nullable`.
+fn widen_field(base: &FieldRef, other: &FieldRef) -> FieldRef {
+ let data_type = widen_nested_nullability(base.data_type(), other.data_type());
+ let nullable = base.is_nullable() || other.is_nullable();
+ if nullable == base.is_nullable() && &data_type == base.data_type() {
+ return Arc::clone(base);
+ }
+ Arc::new(
+ base.as_ref()
+ .clone()
+ .with_data_type(data_type)
+ .with_nullable(nullable),
+ )
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use arrow::array::{Array, Int32Array, Int64Array, ListArray, StringArray, StructArray};
+ use arrow::buffer::OffsetBuffer;
+ use arrow::datatypes::{Field, Fields, Schema};
+
+ fn struct_field(nullable_child: bool) -> FieldRef {
+ Arc::new(Field::new_list_field(
+ DataType::Struct(Fields::from(vec![
+ Field::new("id", DataType::Int64, true),
+ Field::new("flag", DataType::Boolean, nullable_child),
+ ])),
+ true,
+ ))
+ }
+
+ /// The drift from the issue: `List(Struct(..non-null Boolean))` where catalyst declared the
+ /// child nullable. Only the child's `nullable` flag differs, so the path must pin it down.
+ #[test]
+ fn describes_nested_nullability_path() {
+ let expected = DataType::List(struct_field(true));
+ let actual = DataType::List(struct_field(false));
+ assert_eq!(
+ describe_type_mismatch("c0", &expected, &actual).unwrap(),
+ "c0.element.flag: expected nullable Boolean, found non-null Boolean"
+ );
+ }
+
+ #[test]
+ fn describes_leaf_type_difference() {
+ let expected = DataType::List(Arc::new(Field::new_list_field(DataType::Int64, true)));
+ let actual = DataType::List(Arc::new(Field::new_list_field(DataType::Int32, true)));
+ assert_eq!(
+ describe_type_mismatch("c0", &expected, &actual).unwrap(),
+ "c0.element: expected Int64, found Int32"
+ );
+ }
+
+ #[test]
+ fn describes_nothing_for_equal_types() {
+ let dt = DataType::List(struct_field(true));
+ assert_eq!(describe_type_mismatch("c0", &dt, &dt), None);
+ }
+
+ #[test]
+ fn widens_nested_nullability_in_both_directions() {
+ let nullable_child = DataType::List(struct_field(true));
+ let non_null_child = DataType::List(struct_field(false));
+ assert_eq!(
+ widen_nested_nullability(&non_null_child, &nullable_child),
+ nullable_child
+ );
+ assert_eq!(
+ widen_nested_nullability(&nullable_child, &non_null_child),
+ nullable_child
+ );
+ }
+
+ #[test]
+ fn widening_leaves_real_type_differences_alone() {
+ let int64 = DataType::List(Arc::new(Field::new_list_field(DataType::Int64, false)));
+ let int32 = DataType::List(Arc::new(Field::new_list_field(DataType::Int32, true)));
+ // The element nullability still widens; the leaf type stays `base`'s for the cast to fix.
+ assert_eq!(
+ widen_nested_nullability(&int64, &int32),
+ DataType::List(Arc::new(Field::new_list_field(DataType::Int64, true)))
+ );
+ }
+
+ /// A `List(Struct(non-null Boolean))` array stamped with a schema declaring the child nullable
+ /// must be absorbed rather than rejected, and the values must survive unchanged.
+ #[test]
+ fn stamps_nested_nullability_drift() {
+ let actual = drifting_list_of_struct();
+ let schema = Arc::new(Schema::new(vec![Field::new(
+ "c0",
+ DataType::List(struct_field(true)),
+ true,
+ )]));
+
+ // Pin the reason this helper exists: stamping the declared schema on directly rejects the
+ // batch on the child's `nullable` flag alone. If arrow ever relaxes that, this assertion
+ // flips and every `cast_and_stamp_schema` call site can go back to a plain stamp.
+ let options = RecordBatchOptions::new().with_row_count(Some(2));
+ assert!(
+ RecordBatch::try_new_with_options(
+ Arc::clone(&schema),
+ vec![Arc::clone(&actual)],
+ &options
+ )
+ .is_err(),
+ "arrow no longer treats nested field nullability as part of DataType identity"
+ );
+
+ let batch = cast_and_stamp_schema("TestExec", &schema, vec![actual], 2).unwrap();
+ assert_eq!(batch.num_rows(), 2);
+ assert_eq!(batch.schema(), schema);
+ let list = batch
+ .column(0)
+ .as_any()
+ .downcast_ref::<ListArray>()
+ .unwrap();
+ assert_eq!(list.value(0).len(), 2);
+ assert_eq!(list.value(1).len(), 1);
+ let entries = list
+ .values()
+ .as_any()
+ .downcast_ref::<StructArray>()
+ .unwrap();
+ let ids = entries
+ .column(0)
+ .as_any()
+ .downcast_ref::<Int64Array>()
+ .unwrap();
+ assert_eq!(ids.values(), &[1, 2, 3]);
+ }
+
+ /// `[[{1,true},{2,false}], [{3,true}]]` with a non-null `flag` child.
+ fn drifting_list_of_struct() -> ArrayRef {
+ let entries = StructArray::new(
+ Fields::from(vec![
+ Field::new("id", DataType::Int64, true),
+ Field::new("flag", DataType::Boolean, false),
+ ]),
+ vec![
+ Arc::new(Int64Array::from(vec![1, 2, 3])) as ArrayRef,
+ Arc::new(arrow::array::BooleanArray::from(vec![true, false, true])),
+ ],
+ None,
+ );
+ Arc::new(ListArray::new(
+ struct_field(false),
+ OffsetBuffer::new(vec![0, 2, 3].into()),
+ Arc::new(entries),
+ None,
+ ))
+ }
+
+ #[test]
+ fn stamps_equal_types_without_copying() {
+ let schema = Arc::new(Schema::new(vec![Field::new("c0", DataType::Int32, true)]));
+ let column: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3]));
+ let batch =
+ cast_and_stamp_schema("TestExec", &schema, vec![Arc::clone(&column)], 3).unwrap();
+ assert!(Arc::ptr_eq(batch.column(0), &column));
+ }
+
+ /// An unreconcilable difference must still name the operator and the column, not just an index.
+ #[test]
+ fn error_names_operator_and_column() {
+ let schema = Arc::new(Schema::new(vec![Field::new(
+ "payload",
+ DataType::Struct(Fields::from(vec![Field::new("id", DataType::Int64, true)])),
+ true,
+ )]));
+ let column: ArrayRef = Arc::new(StringArray::from(vec!["a", "b"]));
+ let err = cast_and_stamp_schema("TestExec", &schema, vec![column], 2).unwrap_err();
+ let msg = err.to_string();
+ assert!(msg.contains("TestExec"), "{msg}");
+ assert!(msg.contains("col[0]"), "{msg}");
+ assert!(msg.contains("payload: expected Struct"), "{msg}");
+ }
+
+ #[test]
+ fn error_on_column_count_mismatch() {
+ let schema = Arc::new(Schema::new(vec![Field::new("c0", DataType::Int32, true)]));
+ let err = cast_and_stamp_schema("TestExec", &schema, vec![], 0).unwrap_err();
+ assert!(
+ err.to_string()
+ .contains("produced 0 columns but its schema declares 1"),
+ "{err}"
+ );
+ }
+}
diff --git a/native/core/src/execution/operators/expand.rs b/native/core/src/execution/operators/expand.rs
index 8edc8c4d50..28d0c70225 100644
--- a/native/core/src/execution/operators/expand.rs
+++ b/native/core/src/execution/operators/expand.rs
@@ -15,8 +15,8 @@
// specific language governing permissions and limitations
// under the License.
-use arrow::array::{RecordBatch, RecordBatchOptions};
-use arrow::datatypes::SchemaRef;
+use arrow::array::RecordBatch;
+use arrow::datatypes::{Field, Schema, SchemaRef};
use datafusion::common::DataFusionError;
use datafusion::physical_expr::{EquivalenceProperties, PhysicalExpr};
use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType};
@@ -27,6 +27,7 @@ use datafusion::{
RecordBatchStream, SendableRecordBatchStream,
},
};
+use datafusion_comet_common::{cast_and_stamp_schema, widen_nested_nullability};
use futures::{Stream, StreamExt};
use std::{
pin::Pin,
@@ -65,6 +66,51 @@ impl ExpandExec {
cache,
}
}
+
+ /// Derives the operator's output schema from *every* projection, not just the first.
+ ///
+ /// Each projection produces the same output columns, so their types agree except where a
+ /// native kernel drifts on nested field nullability. Widening each column's nested `nullable`
+ /// flags across all projections gives a schema that is valid for all of them; deriving it from
+ /// `projections[0]` alone yields a schema that the remaining projections cannot be stamped
+ /// with. See <https://github.com/apache/datafusion-comet/issues/5137>.
+ pub fn build_schema(
+ projections: &[Vec<Arc<dyn PhysicalExpr>>],
+ input_schema: &SchemaRef,
+ ) -> Result<SchemaRef, DataFusionError> {
+ let mut data_types = projections
+ .first()
+ .ok_or_else(|| {
+ DataFusionError::Internal("Expand should have at least one projection".to_string())
+ })?
+ .iter()
+ .map(|expr| expr.data_type(input_schema))
+ .collect::<Result<Vec<_>, _>>()?;
+
+ for projection in &projections[1..] {
+ if projection.len() != data_types.len() {
+ return Err(DataFusionError::Internal(format!(
+ "Expand projections produce differing column counts: {} and {}",
+ data_types.len(),
+ projection.len()
+ )));
+ }
+ for (data_type, expr) in data_types.iter_mut().zip(projection) {
+ *data_type = widen_nested_nullability(data_type, &expr.data_type(input_schema)?);
+ }
+ }
+
+ // `col_{idx}` is the name the planner used before this logic moved here, kept so that only
+ // the nullability derivation changes. A projection is an arbitrary expression with no
+ // natural output name, and parent operators bind to these columns by index (a bound
+ // reference takes its name from the child schema), so the name is positional either way.
+ let fields: Vec<Field> = data_types
+ .into_iter()
+ .enumerate()
+ .map(|(idx, dt)| Field::new(format!("col_{idx}"), dt, true))
+ .collect();
+ Ok(Arc::new(Schema::new(fields)))
+ }
}
impl DisplayAs for ExpandExec {
@@ -174,9 +220,9 @@ impl ExpandStream {
Ok::<(), DataFusionError>(())
})?;
- let options = RecordBatchOptions::new().with_row_count(Some(batch.num_rows()));
- RecordBatch::try_new_with_options(Arc::clone(&self.schema), columns, &options)
- .map_err(|e| e.into())
+ // A projection whose nested nullability is narrower than the operator's declared schema is
+ // reconciled here rather than rejected by the stamp.
+ cast_and_stamp_schema("CometExpandExec", &self.schema, columns, batch.num_rows())
}
}
@@ -214,3 +260,146 @@ impl RecordBatchStream for ExpandStream {
Arc::clone(&self.schema)
}
}
+
+#[cfg(test)]
+mod tests {
+ //! `ExpandExec` stamps its declared schema onto the output of every projection. Nested field
+ //! nullability is part of arrow's `DataType` identity, so a projection whose nested `nullable`
+ //! flags differ from the declared schema used to abort the task. These tests pin both halves of
+ //! the fix independently: `build_schema` widens the declared schema across all projections, and
+ //! the stamp reconciles anything still narrower. See
+ //! <https://github.com/apache/datafusion-comet/issues/5137>.
+
+ use super::*;
+ use crate::execution::operators::nested_nullability_fixture::{
+ list_of_struct, list_of_struct_type,
+ };
+ use arrow::array::{Array, ListArray};
+ use datafusion::datasource::memory::MemorySourceConfig;
+ use datafusion::physical_expr::expressions::Column;
+ use datafusion::physical_plan::collect;
+ use datafusion::prelude::SessionContext;
+
+ /// Child with two columns of the same logical `array<struct<id,flag>>` type that disagree on
+ /// the `flag` field's nullability, so that projecting one per Expand projection reproduces the
+ /// drift.
+ fn drifting_child() -> Arc<dyn ExecutionPlan> {
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("a", list_of_struct_type(true), true),
+ Field::new("b", list_of_struct_type(false), true),
+ ]));
+ let batch = RecordBatch::try_new(
+ Arc::clone(&schema),
+ vec![list_of_struct(true), list_of_struct(false)],
+ )
+ .unwrap();
+ MemorySourceConfig::try_new_exec(&[vec![batch]], schema, None).unwrap()
+ }
+
+ fn drifting_projections() -> Vec<Vec<Arc<dyn PhysicalExpr>>> {
+ vec![
+ vec![Arc::new(Column::new("a", 0)) as Arc<dyn PhysicalExpr>],
+ vec![Arc::new(Column::new("b", 1)) as Arc<dyn PhysicalExpr>],
+ ]
+ }
+
+ async fn expand_all(plan: Arc<dyn ExecutionPlan>) -> Vec<RecordBatch> {
+ collect(plan, SessionContext::new().task_ctx())
+ .await
+ .unwrap()
+ }
+
+ /// Each output batch must carry the operator's declared schema and the values the projection
+ /// produced, unchanged.
+ fn assert_expanded(batches: &[RecordBatch], schema: &SchemaRef) {
+ assert_eq!(batches.len(), 2, "one batch per projection");
+ for batch in batches {
+ assert_eq!(&batch.schema(), schema);
+ assert_eq!(batch.num_rows(), 2);
+ let list = batch
+ .column(0)
+ .as_any()
+ .downcast_ref::<ListArray>()
+ .unwrap();
+ assert_eq!(list.value(0).len(), 2);
+ assert_eq!(list.value(1).len(), 1);
+ }
+ }
+
+ /// `build_schema` must widen the nested `nullable` flags across all projections, not read them
+ /// off `projections[0]`.
+ #[test]
+ fn build_schema_widens_nested_nullability_across_projections() {
+ let child = drifting_child();
+ // `projections[0]` alone would declare the widened type here, so drive the widening from
+ // the narrow side to prove it is not just projections[0] being echoed back.
+ let projections = vec![
+ vec![Arc::new(Column::new("b", 1)) as Arc<dyn PhysicalExpr>],
+ vec![Arc::new(Column::new("a", 0)) as Arc<dyn PhysicalExpr>],
+ ];
+ let schema = ExpandExec::build_schema(&projections, &child.schema()).unwrap();
+ assert_eq!(
+ schema.field(0).data_type(),
+ &list_of_struct_type(true),
+ "the flag field must be nullable because projection[1] produces it nullable"
+ );
+ }
+
+ #[test]
+ fn build_schema_rejects_projections_of_differing_width() {
+ let child = drifting_child();
+ let projections = vec![
+ vec![Arc::new(Column::new("a", 0)) as Arc<dyn PhysicalExpr>],
+ vec![
+ Arc::new(Column::new("a", 0)) as Arc<dyn PhysicalExpr>,
+ Arc::new(Column::new("b", 1)) as Arc<dyn PhysicalExpr>,
+ ],
+ ];
+ let err = ExpandExec::build_schema(&projections, &child.schema()).unwrap_err();
+ assert!(err.to_string().contains("differing column counts"), "{err}");
+ }
+
+ /// End-to-end: projections that disagree on nested nullability both expand successfully under
+ /// the widened schema.
+ #[tokio::test]
+ async fn expands_projections_with_divergent_nested_nullability() {
+ let child = drifting_child();
+ let projections = drifting_projections();
+ let schema = ExpandExec::build_schema(&projections, &child.schema()).unwrap();
+ let expand = Arc::new(ExpandExec::new(projections, child, Arc::clone(&schema)));
+ assert_expanded(&expand_all(expand).await, &schema);
+ }
+
+ /// The stamp reconciles on its own: even given the narrow `projections[0]`-derived schema that
+ /// the planner used to build, the projection producing a non-null child no longer aborts.
+ #[tokio::test]
+ async fn stamp_reconciles_a_schema_narrower_than_a_projection() {
+ let child = drifting_child();
+ let narrow = Arc::new(Schema::new(vec![Field::new(
+ "col_0",
+ list_of_struct_type(false),
+ true,
+ )]));
+ let expand = Arc::new(ExpandExec::new(
+ drifting_projections(),
+ child,
+ Arc::clone(&narrow),
+ ));
+ assert_expanded(&expand_all(expand).await, &narrow);
+ }
+
+ /// When no projection drifts, the schema is exactly what `projections[0]` yields, and the stamp
+ /// is a no-op passthrough.
+ #[tokio::test]
+ async fn no_drift_is_unchanged() {
+ let child = drifting_child();
+ let projections = vec![
+ vec![Arc::new(Column::new("a", 0)) as Arc<dyn PhysicalExpr>],
+ vec![Arc::new(Column::new("a", 0)) as Arc<dyn PhysicalExpr>],
+ ];
+ let schema = ExpandExec::build_schema(&projections, &child.schema()).unwrap();
+ assert_eq!(schema.field(0).data_type(), &list_of_struct_type(true));
+ let expand = Arc::new(ExpandExec::new(projections, child, Arc::clone(&schema)));
+ assert_expanded(&expand_all(expand).await, &schema);
+ }
+}
diff --git a/native/core/src/execution/operators/mod.rs b/native/core/src/execution/operators/mod.rs
index 6fdc3b0486..b9b2b0fbd7 100644
--- a/native/core/src/execution/operators/mod.rs
+++ b/native/core/src/execution/operators/mod.rs
@@ -39,3 +39,52 @@ mod scan;
mod shuffle_scan;
pub use csv_scan::init_csv_datasource_exec;
pub use shuffle_scan::ShuffleScanExec;
+
+/// Fixtures for the nested-nullability drift from
+/// <https://github.com/apache/datafusion-comet/issues/5137>, shared by the `expand` and
+/// `shuffle_scan` tests so the two boundaries are pinned against the same array.
+#[cfg(test)]
+pub(crate) mod nested_nullability_fixture {
+ use arrow::array::{ArrayRef, BooleanArray, Int64Array, ListArray, StructArray};
+ use arrow::buffer::OffsetBuffer;
+ use arrow::datatypes::{DataType, Field, FieldRef, Fields};
+ use std::sync::Arc;
+
+ pub fn struct_fields(flag_nullable: bool) -> Fields {
+ Fields::from(vec![
+ Field::new("id", DataType::Int64, true),
+ Field::new("flag", DataType::Boolean, flag_nullable),
+ ])
+ }
+
+ /// The element field of an `array<struct<id,flag>>`, with `flag`'s nullability as given.
+ pub fn element_field(flag_nullable: bool) -> FieldRef {
+ Arc::new(Field::new_list_field(
+ DataType::Struct(struct_fields(flag_nullable)),
+ true,
+ ))
+ }
+
+ pub fn list_of_struct_type(flag_nullable: bool) -> DataType {
+ DataType::List(element_field(flag_nullable))
+ }
+
+ /// `[[{1,true},{2,false}], [{3,true}]]`. Both variants hold identical values; only the `flag`
+ /// field's `nullable` flag differs, which is the whole of the drift being absorbed.
+ pub fn list_of_struct(flag_nullable: bool) -> ArrayRef {
+ let entries = StructArray::new(
+ struct_fields(flag_nullable),
+ vec![
+ Arc::new(Int64Array::from(vec![1, 2, 3])) as ArrayRef,
+ Arc::new(BooleanArray::from(vec![true, false, true])),
+ ],
+ None,
+ );
+ Arc::new(ListArray::new(
+ element_field(flag_nullable),
+ OffsetBuffer::new(vec![0, 2, 3].into()),
+ Arc::new(entries),
+ None,
+ ))
+ }
+}
diff --git a/native/core/src/execution/operators/shuffle_scan.rs b/native/core/src/execution/operators/shuffle_scan.rs
index f3209b0c1b..71cfc29ae3 100644
--- a/native/core/src/execution/operators/shuffle_scan.rs
+++ b/native/core/src/execution/operators/shuffle_scan.rs
@@ -24,7 +24,7 @@ use crate::{
};
use arrow::array::ArrayRef;
use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
-use datafusion::common::{arrow_datafusion_err, Result as DataFusionResult};
+use datafusion::common::Result as DataFusionResult;
use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType};
use datafusion::physical_plan::metrics::{
BaselineMetrics, ExecutionPlanMetricsSet, MetricBuilder, MetricsSet, Time,
@@ -34,6 +34,7 @@ use datafusion::{
physical_expr::*,
physical_plan::{ExecutionPlan, *},
};
+use datafusion_comet_common::cast_and_stamp_schema;
use futures::Stream;
use jni::objects::{Global, JByteBuffer, JObject};
use std::{
@@ -318,14 +319,16 @@ impl Stream for ShuffleScanStream {
InputBatch::EOF => Poll::Ready(None),
InputBatch::Batch(columns, num_rows) => {
self.baseline_metrics.record_output(*num_rows);
- let options =
- arrow::array::RecordBatchOptions::new().with_row_count(Some(*num_rows));
- let maybe_batch = arrow::array::RecordBatch::try_new_with_options(
- self.shuffle_scan.schema(),
+ // Reconcile the decoded block with the catalyst-declared schema rather than
+ // stamping it on, so that nested field nullability drift is absorbed here the way
+ // `ScanExec` absorbs it at the FFI boundary.
+ // See https://github.com/apache/datafusion-comet/issues/5137.
+ let maybe_batch = cast_and_stamp_schema(
+ self.shuffle_scan.name(),
+ &self.shuffle_scan.schema,
columns.clone(),
- &options,
- )
- .map_err(|e| arrow_datafusion_err!(e));
+ *num_rows,
+ );
Poll::Ready(Some(maybe_batch))
}
};
@@ -499,4 +502,79 @@ mod tests {
assert_eq!(col1.value(2), "hello");
});
}
+
+ /// A decoded shuffle block whose nested field nullability is narrower than the catalyst-declared
+ /// type must be reconciled, not rejected. `ShuffleScanExec` used to stamp the declared schema
+ /// straight onto the block, which aborted the task on a single nested `nullable` flag even
+ /// though a non-null child is a strict subset of a nullable one.
+ /// See <https://github.com/apache/datafusion-comet/issues/5137>.
+ #[test]
+ #[cfg_attr(miri, ignore)]
+ fn test_nested_nullability_drift_is_reconciled() {
+ use super::*;
+ use crate::execution::operators::nested_nullability_fixture::{
+ list_of_struct, list_of_struct_type,
+ };
+ use arrow::array::Array;
+ use datafusion::physical_plan::ExecutionPlan;
+ use futures::StreamExt;
+
+ // The block carries `List(Struct("id": Int64, "flag": non-null Boolean))` while catalyst
+ // declared the `flag` child nullable.
+ let block_column = list_of_struct(false);
+ let declared = list_of_struct_type(true);
+ let mut scan = ShuffleScanExec::new(
+ super::super::super::planner::TEST_EXEC_CONTEXT_ID,
+ None,
+ vec![declared.clone()],
+ )
+ .unwrap();
+ scan.set_input_batch(InputBatch::new(vec![block_column], Some(2)));
+
+ let rt = tokio::runtime::Runtime::new().unwrap();
+ rt.block_on(async {
+ let ctx = Arc::new(TaskContext::default());
+ let mut stream = scan.execute(0, ctx).unwrap();
+ let batch = stream.next().await.unwrap().unwrap();
+
+ assert_eq!(batch.schema().field(0).data_type(), &declared);
+ assert_eq!(batch.num_rows(), 2);
+ // The values must survive the reconciliation untouched. `ArrayData` equality is
+ // logical, so this holds regardless of whether the cast materialized an all-valid
+ // null buffer for the widened child.
+ assert_eq!(batch.column(0).to_data(), list_of_struct(true).to_data());
+ });
+ }
+
+ /// An unreconcilable column must name the operator and the column, since arrow's own message
+ /// reports only `at column index N`.
+ #[test]
+ #[cfg_attr(miri, ignore)]
+ fn test_unreconcilable_column_error_names_operator() {
+ use super::*;
+ use arrow::datatypes::Fields;
+ use datafusion::physical_plan::ExecutionPlan;
+ use futures::StreamExt;
+
+ let declared =
+ DataType::Struct(Fields::from(vec![Field::new("id", DataType::Int64, true)]));
+ let mut scan = ShuffleScanExec::new(
+ super::super::super::planner::TEST_EXEC_CONTEXT_ID,
+ None,
+ vec![declared],
+ )
+ .unwrap();
+ let column: ArrayRef = Arc::new(StringArray::from(vec!["a", "b"]));
+ scan.set_input_batch(InputBatch::new(vec![column], Some(2)));
+
+ let rt = tokio::runtime::Runtime::new().unwrap();
+ rt.block_on(async {
+ let ctx = Arc::new(TaskContext::default());
+ let mut stream = scan.execute(0, ctx).unwrap();
+ let err = stream.next().await.unwrap().unwrap_err().to_string();
+ assert!(err.contains("ShuffleScanExec"), "{err}");
+ assert!(err.contains("col[0]"), "{err}");
+ assert!(err.contains("col_0: expected Struct"), "{err}");
+ });
+ }
}
diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs
index 85851b17ca..43f066e76c 100644
--- a/native/core/src/execution/planner.rs
+++ b/native/core/src/execution/planner.rs
@@ -1880,21 +1880,7 @@ impl PhysicalPlanner {
Ok::<(), ExecutionError>(())
})?;
- assert!(
- !projections.is_empty(),
- "Expand should have at least one projection"
- );
-
- let datatypes = projections[0]
- .iter()
- .map(|expr| expr.data_type(&child.schema()))
- .collect::<Result<Vec<DataType>, _>>()?;
- let fields: Vec<Field> = datatypes
- .iter()
- .enumerate()
- .map(|(idx, dt)| Field::new(format!("col_{idx}"), dt.clone(), true))
- .collect();
- let schema = Arc::new(Schema::new(fields));
+ let schema = ExpandExec::build_schema(&projections, &child.schema())?;
// `Expand` operator keeps the input batch and expands it to multiple output
// batches. However, `ScanExec` will reuse input arrays for the next
diff --git a/native/shuffle/src/schema_align.rs b/native/shuffle/src/schema_align.rs
index 6c8d0bb97b..33ce490b69 100644
--- a/native/shuffle/src/schema_align.rs
+++ b/native/shuffle/src/schema_align.rs
@@ -22,21 +22,19 @@
//! return-type drift from DataFusion / `datafusion-spark` is self-healing. When a native plan's
//! output crosses back to the JVM and feeds another native plan, the consuming `ScanExec` casts
//! every imported column to the catalyst-declared type, so a wrong Arrow type never survives the
-//! boundary. Shuffle is the lone exception, on two counts:
+//! boundary. Shuffle is the lone exception because the writer hash-partitions on these columns, and
+//! Spark's hash differs by type (e.g. `Int32` vs `Int64`), so a drifted type would route rows to
+//! the wrong partition. A read-side cast cannot undo a wrong partition assignment, so the type must
+//! be corrected before partitioning — which forces the alignment onto the writer input.
//!
-//! 1. The writer hash-partitions on these columns, and Spark's hash differs by type (e.g. `Int32`
-//! vs `Int64`), so a drifted type would route rows to the wrong partition. A read-side cast
-//! cannot undo a wrong partition assignment, so the type must be corrected before partitioning.
-//! 2. The shuffle read path (`ShuffleScanExec`) does not cast; it stamps the catalyst schema onto
-//! the decoded block and errors on any mismatch. The schema is serialized into the block on
-//! write and trusted on read.
+//! The read path (`ShuffleScanExec`) casts too, so a drift that only affects the batch's Arrow type
+//! and not its partition assignment is absorbed there as well. See
+//! <https://github.com/apache/datafusion-comet/issues/5137>.
//!
-//! Both force the alignment to happen on the writer input. See
-//! <https://github.com/apache/datafusion-comet/issues/4515> for the running list of mismatched
+//! See <https://github.com/apache/datafusion-comet/issues/4515> for the running list of mismatched
//! functions.
-use arrow::array::{ArrayRef, RecordBatch, RecordBatchOptions};
-use arrow::compute::{cast_with_options, CastOptions};
+use arrow::array::RecordBatch;
use arrow::datatypes::{Field, Schema, SchemaRef};
use datafusion::common::DataFusionError;
use datafusion::physical_expr::EquivalenceProperties;
@@ -48,6 +46,7 @@ use datafusion::{
RecordBatchStream, SendableRecordBatchStream,
},
};
+use datafusion_comet_common::cast_and_stamp_schema;
use futures::{Stream, StreamExt};
use std::{
collections::HashSet,
@@ -72,19 +71,9 @@ fn warn_dedup() -> &'static Mutex<HashSet<String>> {
pub struct SchemaAlignExec {
child: Arc<dyn ExecutionPlan>,
target_schema: SchemaRef,
- column_actions: Arc<Vec<ColumnAction>>,
cache: Arc<PlanProperties>,
}
-#[derive(Debug, Clone)]
-enum ColumnAction {
- /// Pass the input column through unchanged. Any nullability/metadata difference is
- /// absorbed when the batch is re-stamped via `RecordBatch::try_new_with_options`.
- Passthrough,
- /// Cast the input column to the target data_type.
- Cast,
-}
-
impl SchemaAlignExec {
/// Build a SchemaAlignExec that aligns `child`'s output to `expected`. Returns
/// `Ok(child)` unchanged when no per-column reshape is needed; otherwise wraps `child`
@@ -103,7 +92,6 @@ impl SchemaAlignExec {
)));
}
let mut needs_alignment = false;
- let mut actions = Vec::with_capacity(actual.fields().len());
let mut target_fields = Vec::with_capacity(actual.fields().len());
for (idx, (actual_field, expected_field)) in actual
.fields()
@@ -111,8 +99,8 @@ impl SchemaAlignExec {
.zip(expected.fields().iter())
.enumerate()
{
- let action = if actual_field.data_type() == expected_field.data_type() {
- ColumnAction::Passthrough
+ let needs_cast = if actual_field.data_type() == expected_field.data_type() {
+ false
} else {
let signature = format!(
"{}|{:?}|{:?}",
@@ -130,10 +118,10 @@ impl SchemaAlignExec {
expected_field.data_type()
);
}
- ColumnAction::Cast
+ true
};
let target_nullable = actual_field.is_nullable() || expected_field.is_nullable();
- let field_changed = !matches!(action, ColumnAction::Passthrough)
+ let field_changed = needs_cast
|| target_nullable != actual_field.is_nullable()
|| expected_field.metadata() != actual_field.metadata()
|| expected_field.name() != actual_field.name();
@@ -148,7 +136,6 @@ impl SchemaAlignExec {
)
.with_metadata(expected_field.metadata().clone()),
);
- actions.push(action);
}
if !needs_alignment {
return Ok(child);
@@ -163,7 +150,6 @@ impl SchemaAlignExec {
Ok(Arc::new(Self {
child,
target_schema,
- column_actions: Arc::new(actions),
cache,
}))
}
@@ -204,7 +190,6 @@ impl ExecutionPlan for SchemaAlignExec {
Ok(Arc::new(Self {
child: new_child,
target_schema: Arc::clone(&self.target_schema),
- column_actions: Arc::clone(&self.column_actions),
cache,
}))
}
@@ -218,7 +203,6 @@ impl ExecutionPlan for SchemaAlignExec {
Ok(Box::pin(SchemaAlignStream {
child_stream,
target_schema: Arc::clone(&self.target_schema),
- column_actions: Arc::clone(&self.column_actions),
}))
}
@@ -234,27 +218,17 @@ impl ExecutionPlan for SchemaAlignExec {
struct SchemaAlignStream {
child_stream: SendableRecordBatchStream,
target_schema: SchemaRef,
- column_actions: Arc<Vec<ColumnAction>>,
}
impl SchemaAlignStream {
fn align(&self, batch: RecordBatch) -> Result<RecordBatch, DataFusionError> {
- let mut columns: Vec<ArrayRef> = Vec::with_capacity(batch.num_columns());
- for (idx, action) in self.column_actions.iter().enumerate() {
- let column = batch.column(idx);
- let aligned = match action {
- ColumnAction::Passthrough => Arc::clone(column),
- ColumnAction::Cast => cast_with_options(
- column,
- self.target_schema.field(idx).data_type(),
- &CastOptions::default(),
- )?,
- };
- columns.push(aligned);
- }
- let options = RecordBatchOptions::new().with_row_count(Some(batch.num_rows()));
- RecordBatch::try_new_with_options(Arc::clone(&self.target_schema), columns, &options)
- .map_err(DataFusionError::from)
+ let num_rows = batch.num_rows();
+ cast_and_stamp_schema(
+ "CometSchemaAlignExec",
+ &self.target_schema,
+ batch.columns().to_vec(),
+ num_rows,
+ )
}
}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org