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")