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 <[email protected]>
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: [email protected]
For additional commands, e-mail: [email protected]
