ajsquared commented on code in PR #5650:
URL: https://github.com/apache/datafusion-comet/pull/5650#discussion_r3919110738
##########
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() {
+ DataType::Dictionary(_, values) => {
+ arrow::compute::cast(column,
values).map_err(DataFusionError::from)
Review Comment:
[P2][non-blocking] Unpack nested dictionaries without narrowing the outer
keys
For a valid `Dictionary<UInt16, Dictionary<Int8, Int32>>` whose outer keys
contain `128` and whose inner dictionary has 129 rows all referencing value
`7`, recursive validation accepts the logical `Int32` type. However,
`cast(column, values)` targets `Dictionary<Int8, Int32>`: Arrow 58.4's
`dictionary_to_dictionary_cast` casts the outer UInt16 keys to Int8 and rejects
128, so this valid shuffle block still becomes a fetch failure. This cast does
not unwrap one dictionary layer when the value type is itself a dictionary.
Gather the dictionary values using the outer keys before normalization,
preserving the inner key width, and cover this mixed-width nested case.
--
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]
