[SPARK-61143][CONNECT] Decode the Arrow batches that the server compresses with zstd or lz4 in the Rust client, as PySpark does

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.

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![

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