ziting-openai commented on code in PR #5650:
URL: https://github.com/apache/datafusion-comet/pull/5650#discussion_r3919172892
##########
native/shuffle/src/remote_schema.rs:
##########
@@ -16,18 +16,73 @@
// under the License.
use arrow::array::RecordBatch;
-use arrow::datatypes::DataType;
+use arrow::datatypes::{DataType, Schema};
use datafusion::common::DataFusionError;
use datafusion::error::Result;
+use datafusion_comet_common::cast_and_stamp_schema;
+use std::sync::Arc;
+
+/// Decode a remote shuffle batch and reconcile its encoding with Spark's
declared logical types.
+/// Validate buffers and logical types before casting, so a corrupt frame
cannot be made to look
+/// compatible by a value-changing cast. Dictionary keys are an encoding
detail, including inside
+/// containers: decode them here before either native execution or the JVM
Arrow importer sees them.
+pub fn decode_remote_shuffle_batch(
+ bytes: &[u8],
+ expected_types: &[DataType],
+) -> Result<RecordBatch> {
+ let batch = crate::read_ipc_compressed_validated(bytes)?;
+ validate_remote_schema(&batch, expected_types)?;
+ if batch
+ .columns()
+ .iter()
+ .zip(expected_types)
+ .all(|(column, expected)| column.data_type() == expected)
+ {
+ return Ok(batch);
+ }
+
+ let fields: Vec<_> = batch
+ .schema()
+ .fields()
+ .iter()
+ .zip(expected_types)
+ .map(|(field, expected)| {
+ // Non-null dictionary keys can reference null values. The
declared types do not
+ // constrain top-level nullability; match ShuffleScanExec's
nullable output fields.
+ field
+ .as_ref()
+ .clone()
+ .with_data_type(expected.clone())
+ .with_nullable(true)
+ })
+ .collect();
+ let schema = Arc::new(Schema::new_with_metadata(
+ fields,
+ batch.schema().metadata().clone(),
+ ));
+ // Match the local reader's unpack-then-normalize order. Taking the
dictionary keys first
+ // removes unused values, which may contain nulls that cannot satisfy the
expected nested
+ // nullability even though every referenced value does.
+ let columns = batch
+ .columns()
+ .iter()
+ .map(|column| match column.data_type() {
Review Comment:
[P2] Decode dictionaries inside containers before narrowing nullability
`same_logical_type` now accepts a `List<Dictionary<Int32, Struct<value:
nullable Int32>>>` when Spark expects `List<Struct<value: non-null Int32>>`,
but this pre-pass unwraps a dictionary only when it is the top-level column.
With keys `[0, 0]` and dictionary values `[{7}, {null}]`, the logical list
contains only valid `{7}` values. `cast_and_stamp_schema` then asks Arrow to
cast the list, and Arrow 58.4 casts the entire dictionary-values table to the
non-null struct before applying `take(keys)`, so the unused `{null}` fails with
`Found unmasked nulls for non-nullable StructArray field`. That turns a valid
writer-emitted frame into a remote decode/fetch failure. Please recursively
compact dictionaries inside list/struct/map containers before nullability
reconciliation, and add this nested unused-value regression.
--
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]
