support compression for IPC

arrow/Cargo.toml

@@ -38,6 +38,8 @@ path = "src/lib.rs"
 bench = false
 
 [dependencies]
+lz4 = { version = "1.23", optional = true }
+zstd = { version = "0.11.1", optional = true, default-features = false }
 serde = { version = "1.0" }
 serde_derive = "1.0"
 serde_json = { version = "1.0", features = ["preserve_order"] }

arrow/src/ipc/compression/compression.rs

@@ -0,0 +1,16 @@
+
+pub fn lz4_compress() {
+
+}
+
+pub fn lz4_decompress() {
+
+}
+
+pub fn zstd_compress() {
+
+}
+
+pub fn zstd_decompress() {
+
+}

arrow/src/ipc/compression/mod.rs

@@ -0,0 +1 @@
+mod compression;

arrow/src/ipc/mod.rs

@@ -29,6 +29,7 @@ pub mod writer;
 #[allow(clippy::redundant_field_names)]
 #[allow(non_camel_case_types)]
 pub mod gen;
+pub mod compression;
 
 pub use self::gen::File::*;
 pub use self::gen::Message::*;

arrow/src/ipc/writer.rs

@@ -37,6 +37,7 @@ use crate::record_batch::RecordBatch;
 use crate::util::bit_util;
 
 use ipc::CONTINUATION_MARKER;
+use crate::ipc::{BodyCompressionMethod, CompressionType};
 
 /// IPC write options used to control the behaviour of the writer
 #[derive(Debug, Clone)]
@@ -55,6 +56,8 @@ pub struct IpcWriteOptions {
     /// version 2.0.0: V4, with legacy format enabled
     /// version 4.0.0: V5
     metadata_version: ipc::MetadataVersion,
+    /// TODO comments
+    batch_compression_type: CompressionType
 }
 
 impl IpcWriteOptions {
@@ -91,6 +94,7 @@ impl IpcWriteOptions {
                         alignment,
                         write_legacy_ipc_format,
                         metadata_version,
+                        batch_compression_type: CompressionType::NO_COMPRESSION,
                     })
                 }
             }
@@ -105,6 +109,7 @@ impl Default for IpcWriteOptions {
             alignment: 8,
             write_legacy_ipc_format: false,
             metadata_version: ipc::MetadataVersion::V5,
+            batch_compression_type: CompressionType::NO_COMPRESSION,
         }
     }
 }
@@ -326,6 +331,7 @@ impl IpcDataGenerator {
         let mut buffers: Vec<ipc::Buffer> = vec![];
         let mut arrow_data: Vec<u8> = vec![];
         let mut offset = 0;
+        // TODO: add compression type
         for array in batch.columns() {
             let array_data = array.data();
             offset = write_array_data(
@@ -342,12 +348,18 @@ impl IpcDataGenerator {
         // write data
         let buffers = fbb.create_vector(&buffers);
         let nodes = fbb.create_vector(&nodes);
+        // TODO check the compression type
+        let mut bodyCompressionBuilder = ipc::BodyCompressionBuilder::new(&mut fbb);
+        bodyCompressionBuilder.add_method(BodyCompressionMethod::BUFFER);
+        bodyCompressionBuilder.add_codec(CompressionType::LZ4_FRAME);
 
         let root = {
             let mut batch_builder = ipc::RecordBatchBuilder::new(&mut fbb);
             batch_builder.add_length(batch.num_rows() as i64);
             batch_builder.add_nodes(nodes);
             batch_builder.add_buffers(buffers);
+            // TODO check the compression type
+            batch_builder.add_compression(bodyCompressionBuilder.finish());
             let b = batch_builder.finish();
             b.as_union_value()
         };