From 47b88fef7ee4a8df9a8e65a1a2f37786a671daf8 Mon Sep 17 00:00:00 2001 From: Hyukjin Kwon Date: Sun, 8 Nov 2026 10:24:26 +0900 Subject: [PATCH] [SPARK-61143][CONNECT] Decode the Arrow batches that the server compresses with zstd or lz4 in the Rust client, as PySpark does With spark.sql.execution.arrow.compression.codec set to zstd or lz4, the server compresses the Arrow batches of the results, which PySpark reads with PyArrow. The client built arrow without the ipc_compression feature, so it failed to decode them ("zstd IPC decompression requires the zstd feature"), in collect, toArrow and toPandas, also of the results of pandas UDFs. arrow now has the feature, which adds no crate to the lock file but the zstd dependency of arrow-ipc. Co-authored-by: Isaac --- Cargo.lock | 1 + Cargo.toml | 2 +- crates/spark-connect/src/dataframe.rs | 45 +++++++++++++++++++++++++++ python/tests/test_dropin_e2e.py | 15 +++++++++ 4 files changed, 62 insertions(+), 1 deletion(-) diff --git a/Cargo.lock b/Cargo.lock index 29a30b61..0e6dd6e7 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -288,6 +288,7 @@ dependencies = [ "arrow-select", "flatbuffers", "lz4_flex", + "zstd", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index 68f28243..8595732f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -73,7 +73,7 @@ tokio-stream = "0.1" futures = "0.3" # Arrow (result decoding). -arrow = { version = "58", default-features = false, features = ["ipc", "prettyprint"] } +arrow = { version = "58", default-features = false, features = ["ipc", "ipc_compression", "prettyprint"] } # DataFusion for the optional to_datafusion() conversion. # Pinned to 53 for arrow 58 compatibility. diff --git a/crates/spark-connect/src/dataframe.rs b/crates/spark-connect/src/dataframe.rs index 14496d6b..812630b5 100644 --- a/crates/spark-connect/src/dataframe.rs +++ b/crates/spark-connect/src/dataframe.rs @@ -3040,6 +3040,51 @@ mod nested_decode_tests { row } + #[test] + fn batches_compressed_with_zstd_or_lz4_decode_as_uncompressed_ones() { + // As PyArrow reads them: the server compresses the Arrow batches of the results with + // the codec of spark.sql.execution.arrow.compression.codec, zstd or lz4. + use arrow::ipc::writer::IpcWriteOptions; + use arrow::ipc::CompressionType; + + let ids: ArrayRef = Arc::new(Int64Array::from_iter((0..2000).map(Some))); + let names: ArrayRef = Arc::new(StringArray::from_iter( + (0..2000).map(|i| (i % 3 != 0).then(|| format!("name {}", i % 7))), + )); + let plain = encoded_batches(vec![("id", ids.clone()), ("name", names.clone())], 2, false); + let rows = decode_arrow_batch(&plain).unwrap(); + assert_eq!(rows.len(), 4000); + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int64, true), + Field::new("name", DataType::Utf8, true), + ])); + let batch = RecordBatch::try_new(schema.clone(), vec![ids, names]).unwrap(); + for codec in [CompressionType::ZSTD, CompressionType::LZ4_FRAME] { + let options = IpcWriteOptions::default() + .try_with_compression(Some(codec)) + .unwrap(); + let mut data = vec![]; + { + let mut writer = + StreamWriter::try_new_with_options(&mut data, &schema, options).unwrap(); + writer.write(&batch).unwrap(); + writer.write(&batch).unwrap(); + writer.finish().unwrap(); + } + assert!(data.len() < plain.data.len(), "{codec:?} compresses"); + let compressed = proto::execute_plan_response::ArrowBatch { + data: data.into(), + ..Default::default() + }; + assert_eq!(decode_arrow_batch(&compressed).unwrap(), rows, "{codec:?}"); + assert_eq!( + decode_arrow_record_batches(&compressed).unwrap(), + decode_arrow_record_batches(&plain).unwrap(), + "{codec:?}" + ); + } + } + #[test] fn mixed_nullable_slices_decode_across_empty_and_repeated_batches() { let columns: Vec<(&str, ArrayRef)> = vec![ diff --git a/python/tests/test_dropin_e2e.py b/python/tests/test_dropin_e2e.py index 7c6326fb..2ebee3ce 100644 --- a/python/tests/test_dropin_e2e.py +++ b/python/tests/test_dropin_e2e.py @@ -1911,6 +1911,21 @@ def test_collected_rows_keep_their_values_as_pyspark(session): assert (rows[0].v.toJson(), rows[0].n) == ("1", 2) +def test_results_compressed_by_the_server_decode_as_pyspark(session): + # As PySpark's: the server compresses the Arrow batches of the results with the codec of + # spark.sql.execution.arrow.compression.codec, zstd or lz4, which the client decodes. + df = session.range(0, 20000, 1, 2).selectExpr("id", "CAST(id % 7 AS STRING) AS s") + rows, table, frame = df.collect(), df.toArrow(), df.toPandas() + try: + for codec in ("zstd", "lz4", "none"): + session.conf.set("spark.sql.execution.arrow.compression.codec", codec) + assert df.collect() == rows, codec + assert df.toArrow().equals(table), codec + assert df.toPandas().equals(frame), codec + finally: + session.conf.unset("spark.sql.execution.arrow.compression.codec") + + def test_transpose_takes_a_column_name_as_pyspark(session): # As PySpark's transpose, of a Column or a column name as the index column. df = session.createDataFrame([("a", 1, 2), ("b", 3, 4)], "k string, x int, y int")