Skip to content

Commit eb410e5

Browse files
committed
Merge remote-tracking branch 'origin/main' into perf/fair-pool-lock-free
2 parents c11b0d5 + ef62b46 commit eb410e5

5 files changed

Lines changed: 806 additions & 53 deletions

File tree

‎native/core/src/execution/jni_api.rs‎

Lines changed: 4 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -100,7 +100,7 @@ use tokio::sync::mpsc;
100100
use crate::execution::memory_pools::{create_memory_pool, parse_memory_pool_config};
101101
use crate::execution::operators::{ScanExec, ShuffleScanExec};
102102
use crate::execution::shuffle::{
103-
read_ipc_compressed, read_ipc_compressed_validated, validate_remote_schema, CompressionCodec,
103+
decode_remote_shuffle_batch, read_ipc_compressed, CompressionCodec,
104104
};
105105
use crate::execution::spark_plan::SparkPlan;
106106

@@ -1297,11 +1297,9 @@ fn decode_shuffle_block(
12971297
let length = length as usize;
12981298
let slice: &[u8] = unsafe { std::slice::from_raw_parts(raw_pointer, length) };
12991299
let batch = if let Some(expected_types) = expected_types {
1300-
let batch = read_ipc_compressed_validated(slice)?;
1301-
// Reject incompatible remote schemas before exporting arrays to the JVM. Casting or
1302-
// importing first can silently change values or bypass remote fetch-failure reporting.
1303-
validate_remote_schema(&batch, expected_types)?;
1304-
batch
1300+
// Reject incompatible logical types, then decode dictionaries before JVM import. The
1301+
// JVM importer supports fewer dictionary key/value layouts than the shuffle writer.
1302+
decode_remote_shuffle_batch(slice, expected_types)?
13051303
} else {
13061304
read_ipc_compressed(slice)?
13071305
};

‎native/core/src/execution/operators/shuffle_scan.rs‎

Lines changed: 17 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -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.

‎native/shuffle/src/lib.rs‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,8 @@ pub(crate) mod metrics;
2121
pub(crate) mod partitioners;
2222
mod remote_schema;
2323
#[cfg(test)]
24+
mod remote_schema_tests;
25+
#[cfg(test)]
2426
mod rss_execution_tests;
2527
mod schema_align;
2628
mod shuffle_writer;
@@ -30,7 +32,7 @@ pub(crate) mod writers;
3032

3133
pub use comet_partitioning::CometPartitioning;
3234
pub use ipc::{read_ipc_compressed, read_ipc_compressed_validated};
33-
pub use remote_schema::validate_remote_schema;
35+
pub use remote_schema::{decode_remote_shuffle_batch, validate_remote_schema};
3436
pub use schema_align::SchemaAlignExec;
3537
pub use shuffle_writer::{ShuffleWriterDestination, ShuffleWriterExec};
3638
pub use writers::{CompressionCodec, ShuffleBlockWriter};

0 commit comments

Comments
 (0)