switch to lzf_flex crate

Author: andygrove

native/Cargo.lock

@@ -55,7 +55,7 @@ jni = "0.21"
 snap = "1.1"
 brotli = "3.3"
 flate2 = "1.0"
-lz4 = "1.24"
+lz4_flex = "0.11.3"
 zstd = "0.11"
 rand = { workspace = true}
 num = { workspace = true }

native/core/Cargo.toml

@@ -63,7 +63,12 @@ fn criterion_benchmark(c: &mut Criterion) {
         let mut cursor = Cursor::new(&mut buffer);
         let ipc_time = Time::default();
         b.iter(|| {
-            write_ipc_compressed(&batch, &mut cursor, &CompressionCodec::Lz4Frame, &ipc_time)
+            write_ipc_compressed(
+                &batch,
+                &mut cursor,
+                &CompressionCodec::Lz4Frame(0),
+                &ipc_time,
+            )
         });
     });
     group.bench_function("shuffle_writer: encode and compress (lz4 block)", |b| {
@@ -77,7 +82,6 @@ fn criterion_benchmark(c: &mut Criterion) {
     });
     group.bench_function("shuffle_writer: end to end", |b| {
         let ctx = SessionContext::new();
-        let batch = create_batch(8192, true);
         let exec = create_shuffle_writer_exec(CompressionCodec::Zstd(1));
         b.iter(|| {
             let task_ctx = ctx.task_ctx();
@@ -90,9 +94,10 @@ fn criterion_benchmark(c: &mut Criterion) {
 
 fn create_shuffle_writer_exec(compression_codec: CompressionCodec) -> ShuffleWriterExec {
     let batches = create_batches(8192, 10);
+    let schema = batches[0].schema();
     let partitions = &[batches];
     let exec = ShuffleWriterExec::try_new(
-        Arc::new(MemoryExec::try_new(partitions, batches[0].schema(), None).unwrap()),
+        Arc::new(MemoryExec::try_new(partitions, schema, None).unwrap()),
         Partitioning::Hash(vec![Arc::new(Column::new("a", 0))], 16),
         compression_codec,
         "/tmp/data.out".to_string(),

native/core/benches/shuffle_writer.rs

@@ -53,6 +53,7 @@ use datafusion_physical_expr::EquivalenceProperties;
 use futures::executor::block_on;
 use futures::{lock::Mutex, Stream, StreamExt, TryFutureExt, TryStreamExt};
 use itertools::Itertools;
+use lz4_flex::compress_prepend_size;
 use simd_adler32::Adler32;
 use std::io::Error;
 use std::{
@@ -1603,32 +1604,45 @@ pub fn write_ipc_compressed<W: Write + Seek>(
             arrow_writer.finish()?;
             let ipc_encoded = arrow_writer.into_inner()?;
 
-            // TODO refactor to compress directly to output buffer and avoid a copy
-            let compressed = lz4::block::compress(ipc_encoded, None, true)?;
+            let compressed = compress_prepend_size(ipc_encoded);
             output.write_all(&compressed)?;
-
             output
+
+            // lz4 crate impl
+            // // TODO refactor to compress directly to output buffer and avoid a copy
+            // let compressed = lz4::block::compress(ipc_encoded, None, true)?;
+            // output.write_all(&compressed)?;
+            //
+            // output
         }
-        CompressionCodec::Lz4Frame(level) => {
+        CompressionCodec::Lz4Frame(_level) => {
             // write IPC first without compression
             let mut buffer = vec![];
             let mut arrow_writer = StreamWriter::try_new(&mut buffer, &batch.schema())?;
             arrow_writer.write(batch)?;
             arrow_writer.finish()?;
             let ipc_encoded = arrow_writer.into_inner()?;
 
-            let mut encoder = lz4::EncoderBuilder::new()
-                .content_size(ipc_encoded.len() as u64)
-                // .block_mode(lz4::BlockMode::Independent)
-                // .block_size(BlockSize::Default)
-                // .checksum(ContentChecksum::NoChecksum)
-                // .block_checksum(BlockChecksum::BlockChecksumEnabled)
-                .level(*level as u32)
-                .build(&mut *output)?;
-            encoder.write_all(ipc_encoded.as_slice())?;
-            let (output, result) = encoder.finish();
-            result?;
-            output
+            // compress
+            let mut reader = Cursor::new(ipc_encoded);
+            let mut wtr = lz4_flex::frame::FrameEncoder::new(output);
+            std::io::copy(&mut reader, &mut wtr).expect("I/O operation failed");
+            wtr.finish().unwrap()
+
+
+            // lz4 crate version
+            // let mut encoder = lz4::EncoderBuilder::new()
+            //     .content_size(ipc_encoded.len() as u64)
+            //     // .block_mode(lz4::BlockMode::Independent)
+            //     // .block_size(BlockSize::Default)
+            //     // .checksum(ContentChecksum::NoChecksum)
+            //     // .block_checksum(BlockChecksum::BlockChecksumEnabled)
+            //     .level(*level as u32)
+            //     .build(&mut *output)?;
+            // encoder.write_all(ipc_encoded.as_slice())?;
+            // let (output, result) = encoder.finish();
+            // result?;
+            // output
         }
         CompressionCodec::Zstd(level) => {
             let encoder = zstd::Encoder::new(output, *level)?;