@@ -20,7 +20,7 @@ use crate::{
2020 execution:: {
2121 operators:: ExecutionError ,
2222 planner:: TEST_EXEC_CONTEXT_ID ,
23- shuffle:: { read_ipc_compressed , read_ipc_compressed_validated , validate_remote_schema } ,
23+ shuffle:: { decode_remote_shuffle_batch , read_ipc_compressed } ,
2424 } ,
2525 jvm_bridge:: { jni_call, JVMClasses } ,
2626} ;
@@ -230,11 +230,9 @@ fn decode_shuffle_batch(
230230 requires_validation : bool ,
231231) -> DataFusionResult < RecordBatch > {
232232 if requires_validation {
233- let batch = read_ipc_compressed_validated ( bytes) ?;
234- // Validate before unpack_dictionary or cast_and_stamp_schema can hide an incompatible
235- // wire type by changing values. Keep failures inside get_next's recovery callback.
236- validate_remote_schema ( & batch, expected_types) ?;
237- Ok ( batch)
233+ // Validate logical types before decoding dictionaries or normalizing nested fields.
234+ // Keep both validation and normalization failures inside get_next's recovery callback.
235+ decode_remote_shuffle_batch ( bytes, expected_types)
238236 } else {
239237 check_column_count ( read_ipc_compressed ( bytes) ?, expected_types. len ( ) )
240238 }
@@ -452,7 +450,7 @@ mod tests {
452450 } ;
453451 payload[ signed_flag] = 0 ;
454452
455- let decoded = super :: read_ipc_compressed_validated ( & payload) . unwrap ( ) ;
453+ let decoded = crate :: execution :: shuffle :: read_ipc_compressed_validated ( & payload) . unwrap ( ) ;
456454 assert_eq ! ( decoded. num_columns( ) , 1 ) ;
457455 assert_eq ! ( decoded. column( 0 ) . data_type( ) , & DataType :: UInt32 ) ;
458456 assert_eq ! (
@@ -568,13 +566,13 @@ mod tests {
568566 fn test_dictionary_encoded_shuffle_block_is_unpacked ( ) {
569567 use super :: * ;
570568 use arrow:: array:: StringDictionaryBuilder ;
571- use arrow:: datatypes:: Int32Type ;
569+ use arrow:: datatypes:: Int8Type ;
572570 use datafusion:: physical_plan:: ExecutionPlan ;
573571 use futures:: StreamExt ;
574572
575573 // Build a batch with a dictionary-encoded string column (simulating what
576574 // the native shuffle writer produces for string columns).
577- let mut dict_builder = StringDictionaryBuilder :: < Int32Type > :: new ( ) ;
575+ let mut dict_builder = StringDictionaryBuilder :: < Int8Type > :: new ( ) ;
578576 dict_builder. append_value ( "hello" ) ;
579577 dict_builder. append_value ( "world" ) ;
580578 dict_builder. append_value ( "hello" ) ; // repeated value, good for dictionary
@@ -585,7 +583,7 @@ mod tests {
585583 Field :: new( "id" , DataType :: Int32 , false ) ,
586584 Field :: new(
587585 "name" ,
588- DataType :: Dictionary ( Box :: new( DataType :: Int32 ) , Box :: new( DataType :: Utf8 ) ) ,
586+ DataType :: Dictionary ( Box :: new( DataType :: Int8 ) , Box :: new( DataType :: Utf8 ) ) ,
589587 true ,
590588 ) ,
591589 ] ) ) ;
@@ -614,14 +612,17 @@ mod tests {
614612 let bytes = buf. into_inner ( ) ;
615613 let body = & bytes[ 16 ..] ;
616614
617- // Remote schema validation must preserve the writer's dictionary representation.
615+ // Local decoding preserves the wire encoding for get_next to unpack. Remote decoding
616+ // validates and unpacks first, so the same result can also be safely imported by the JVM.
617+ let local =
618+ super :: decode_shuffle_batch ( body, & [ DataType :: Int32 , DataType :: Utf8 ] , false ) . unwrap ( ) ;
619+ assert ! ( matches!(
620+ local. column( 1 ) . data_type( ) ,
621+ DataType :: Dictionary ( _, _)
622+ ) ) ;
618623 let decoded =
619624 super :: decode_shuffle_batch ( body, & [ DataType :: Int32 , DataType :: Utf8 ] , true ) . unwrap ( ) ;
620- assert ! (
621- matches!( decoded. column( 1 ) . data_type( ) , DataType :: Dictionary ( _, _) ) ,
622- "Expected dictionary-encoded column from IPC, got {:?}" ,
623- decoded. column( 1 ) . data_type( )
624- ) ;
625+ assert_eq ! ( decoded. column( 1 ) . data_type( ) , & DataType :: Utf8 ) ;
625626
626627 // Create ShuffleScanExec with value types (Utf8, not Dictionary) — this is
627628 // what the protobuf schema provides.
0 commit comments