You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@arrow.apache.org by GitBox <gi...@apache.org> on 2022/06/22 13:08:34 UTC
[GitHub] [arrow-rs] alamb commented on a diff in pull request #1855: support compression for IPC
alamb commented on code in PR #1855:
URL: https://github.com/apache/arrow-rs/pull/1855#discussion_r903701041
##########
arrow/Cargo.toml:
##########
@@ -60,7 +63,7 @@ multiversion = { version = "0.6.1", default-features = false }
bitflags = { version = "1.2.1", default-features = false }
[features]
-default = ["csv", "ipc", "test_utils"]
+default = ["csv", "ipc", "test_utils", "zstd", "lz4"]
csv = ["csv_crate"]
ipc = ["flatbuffers"]
Review Comment:
I also then think we would need to add an entry for `byteorder` under `features`
So if we want `ipc` to always support compression (seems reasonable to me?) it would be like
```suggestion
ipc = ["flatbuffers", "zstd", "lz4", "byte_order"]
```
If we want another specific feature to control ipc compression (e.g. `ipc_compression`) perhaps we could do something like
```toml
[features]
ipc_compression = ["zstd", "lz4", "byte_order"]
```
?
##########
arrow/src/ipc/compression/ipc_compression.rs:
##########
@@ -0,0 +1,124 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::ipc::CompressionType;
+
+#[derive(Debug, Clone, Copy, PartialEq)]
+pub enum CompressionCodecType {
+ NoCompression,
+ Lz4Frame,
+ Zstd,
+}
+
+impl From<CompressionType> for CompressionCodecType {
+ fn from(compression_type: CompressionType) -> Self {
+ match compression_type {
+ CompressionType::ZSTD => CompressionCodecType::Zstd,
+ CompressionType::LZ4_FRAME => CompressionCodecType::Lz4Frame,
+ _ => CompressionCodecType::NoCompression,
+ }
+ }
+}
+
+impl From<CompressionCodecType> for Option<CompressionType> {
+ fn from(codec: CompressionCodecType) -> Self {
+ match codec {
+ CompressionCodecType::NoCompression => None,
+ CompressionCodecType::Lz4Frame => Some(CompressionType::LZ4_FRAME),
+ CompressionCodecType::Zstd => Some(CompressionType::ZSTD),
+ }
+ }
+}
+
+#[cfg(any(feature = "zstd,lz4", test))]
+mod compression_function {
+ use crate::error::Result;
+ use crate::ipc::compression::ipc_compression::CompressionCodecType;
+ use std::io::{Read, Write};
+
+ impl CompressionCodecType {
+ pub fn compress(&self, input: &[u8], output: &mut Vec<u8>) -> Result<()> {
+ match self {
+ CompressionCodecType::Lz4Frame => {
+ let mut encoder = lz4::EncoderBuilder::new().build(output).unwrap();
+ encoder.write_all(input).unwrap();
Review Comment:
I suggest passing the errors back out of here directly (as `compress` returns a `Result`) rather than `unwrap()` which will `panic` on error
##########
arrow/src/ipc/compression/ipc_compression.rs:
##########
@@ -0,0 +1,124 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::ipc::CompressionType;
+
+#[derive(Debug, Clone, Copy, PartialEq)]
+pub enum CompressionCodecType {
+ NoCompression,
+ Lz4Frame,
+ Zstd,
+}
+
+impl From<CompressionType> for CompressionCodecType {
+ fn from(compression_type: CompressionType) -> Self {
+ match compression_type {
+ CompressionType::ZSTD => CompressionCodecType::Zstd,
+ CompressionType::LZ4_FRAME => CompressionCodecType::Lz4Frame,
+ _ => CompressionCodecType::NoCompression,
+ }
+ }
+}
+
+impl From<CompressionCodecType> for Option<CompressionType> {
+ fn from(codec: CompressionCodecType) -> Self {
+ match codec {
+ CompressionCodecType::NoCompression => None,
+ CompressionCodecType::Lz4Frame => Some(CompressionType::LZ4_FRAME),
+ CompressionCodecType::Zstd => Some(CompressionType::ZSTD),
+ }
+ }
+}
+
+#[cfg(any(feature = "zstd,lz4", test))]
Review Comment:
I wonder if you can put this guard on the entire `mod ipc_compression` statement so that the entire module (including the test) is not compiled unless that feature is active
##########
arrow/src/ipc/writer.rs:
##########
@@ -55,9 +61,50 @@ pub struct IpcWriteOptions {
/// version 2.0.0: V4, with legacy format enabled
/// version 4.0.0: V5
metadata_version: ipc::MetadataVersion,
+ batch_compression_type: CompressionCodecType,
}
impl IpcWriteOptions {
+ pub fn try_new_with_compression(
+ alignment: usize,
+ write_legacy_ipc_format: bool,
+ metadata_version: ipc::MetadataVersion,
+ batch_compression_type: CompressionCodecType,
+ ) -> Result<Self> {
+ if alignment == 0 || alignment % 8 != 0 {
+ return Err(ArrowError::InvalidArgumentError(
+ "Alignment should be greater than 0 and be a multiple of 8".to_string(),
+ ));
+ }
+ match batch_compression_type {
+ CompressionCodecType::NoCompression => {}
+ _ => {
+ if metadata_version != ipc::MetadataVersion::V5 {
+ return Err(ArrowError::InvalidArgumentError(
+ "Compress buffer just support from metadata v5".to_string(),
Review Comment:
```suggestion
"Compression only supported in metadata v5 and above".to_string(),
```
##########
arrow/src/ipc/writer.rs:
##########
@@ -55,9 +61,50 @@ pub struct IpcWriteOptions {
/// version 2.0.0: V4, with legacy format enabled
/// version 4.0.0: V5
metadata_version: ipc::MetadataVersion,
+ batch_compression_type: CompressionCodecType,
}
impl IpcWriteOptions {
+ pub fn try_new_with_compression(
+ alignment: usize,
+ write_legacy_ipc_format: bool,
+ metadata_version: ipc::MetadataVersion,
+ batch_compression_type: CompressionCodecType,
+ ) -> Result<Self> {
+ if alignment == 0 || alignment % 8 != 0 {
+ return Err(ArrowError::InvalidArgumentError(
+ "Alignment should be greater than 0 and be a multiple of 8".to_string(),
+ ));
+ }
+ match batch_compression_type {
+ CompressionCodecType::NoCompression => {}
+ _ => {
+ if metadata_version != ipc::MetadataVersion::V5 {
Review Comment:
should this check be `<` instead of `!=` to cover future versions?
```suggestion
if metadata_version < ipc::MetadataVersion::V5 {
```
##########
arrow/src/ipc/writer.rs:
##########
@@ -922,20 +1021,67 @@ fn write_array_data(
}
/// Write a buffer to a vector of bytes, and add its ipc::Buffer to a vector
+/// From <https://github.com/apache/arrow/blob/6a936c4ff5007045e86f65f1a6b6c3c955ad5103/format/Message.fbs#L58>
+/// Each constituent buffer is first compressed with the indicated
+/// compressor, and then written with the uncompressed length in the first 8
+/// bytes as a 64-bit little-endian signed integer followed by the compressed
+/// buffer bytes (and then padding as required by the protocol). The
+/// uncompressed length may be set to -1 to indicate that the data that
+/// follows is not compressed, which can be useful for cases where
+/// compression does not yield appreciable savings.
fn write_buffer(
buffer: &Buffer,
buffers: &mut Vec<ipc::Buffer>,
arrow_data: &mut Vec<u8>,
offset: i64,
+ compression_codec: &CompressionCodecType,
) -> i64 {
- let len = buffer.len();
- let pad_len = pad_to_8(len as u32);
- let total_len: i64 = (len + pad_len) as i64;
- // assert_eq!(len % 8, 0, "Buffer width not a multiple of 8 bytes");
- buffers.push(ipc::Buffer::new(offset, total_len));
- arrow_data.extend_from_slice(buffer.as_slice());
- arrow_data.extend_from_slice(&vec![0u8; pad_len][..]);
- offset + total_len
+ let origin_buffer_len = buffer.len();
+ let mut _compression_buffer = Vec::<u8>::new();
+ let (data, uncompression_buffer_len) = match compression_codec {
+ CompressionCodecType::NoCompression => {
+ // this buffer_len will not used in the following logic
+ // If we don't use the compression, just write the data in the array
+ (buffer.as_slice(), origin_buffer_len as i64)
+ }
+ CompressionCodecType::Lz4Frame | CompressionCodecType::Zstd => {
+ if (origin_buffer_len as i64) == LENGTH_EMPTY_COMPRESSED_DATA {
Review Comment:
It seems like this code should throw an "Not supported" error (or `panic` if it is encountered without support compiled in)
##########
arrow/src/ipc/reader.rs:
##########
@@ -1414,6 +1493,53 @@ mod tests {
});
}
+ #[test]
+ fn read_generated_streams_200() {
+ let testdata = crate::util::test_util::arrow_test_data();
+ let version = "2.0.0-compression";
+
+ // the test is repetitive, thus we can read all supported files at once
+ let paths = vec!["generated_lz4", "generated_zstd"];
+ paths.iter().for_each(|path| {
+ let file = File::open(format!(
+ "{}/arrow-ipc-stream/integration/{}/{}.stream",
+ testdata, version, path
+ ))
+ .unwrap();
+
+ let mut reader = StreamReader::try_new(file, None).unwrap();
+
+ // read expected JSON output
+ let arrow_json = read_gzip_json(version, path);
+ assert!(arrow_json.equals_reader(&mut reader));
+ // the next batch must be empty
+ assert!(reader.next().is_none());
+ // the stream must indicate that it's finished
+ assert!(reader.is_finished());
+ });
+ }
+
+ #[test]
+ fn read_generated_files_200() {
Review Comment:
nice!
##########
arrow/src/ipc/writer.rs:
##########
@@ -55,9 +61,50 @@ pub struct IpcWriteOptions {
/// version 2.0.0: V4, with legacy format enabled
/// version 4.0.0: V5
metadata_version: ipc::MetadataVersion,
+ batch_compression_type: CompressionCodecType,
}
impl IpcWriteOptions {
+ pub fn try_new_with_compression(
Review Comment:
I wonder if you could avoid the duplication here with more of a `Builder` style:
```rust
impl IpcWriteOptions {
pub fn with_compression(mut self, batch_compression_type: CompressionCodecType) -> Result<Self> {
.. // do checks here
self.batch_compresson_type = batch_compression_type;
Ok(self)
}
...
}
```
Then one could use it like:
```rust
let options = IpcWriteOptions::try_new(8, false, ipc::MetadataVersion::v5)?
.with_compression(CompressionCodecType::LZ4)?;
...
```
##########
arrow/Cargo.toml:
##########
@@ -60,7 +63,7 @@ multiversion = { version = "0.6.1", default-features = false }
bitflags = { version = "1.2.1", default-features = false }
[features]
-default = ["csv", "ipc", "test_utils"]
+default = ["csv", "ipc", "test_utils", "zstd", "lz4"]
Review Comment:
I agree -- I suggest something like
```suggestion
default = ["csv", "ipc", "test_utils"]
```
##########
arrow/src/ipc/mod.rs:
##########
@@ -22,6 +22,7 @@ pub mod convert;
pub mod reader;
pub mod writer;
+mod compression;
Review Comment:
I recommend guarding this entire thing via `#cfg(...)` rather than more granular guards within the module
##########
arrow/src/ipc/writer.rs:
##########
@@ -55,9 +61,50 @@ pub struct IpcWriteOptions {
/// version 2.0.0: V4, with legacy format enabled
/// version 4.0.0: V5
metadata_version: ipc::MetadataVersion,
+ batch_compression_type: CompressionCodecType,
}
impl IpcWriteOptions {
+ pub fn try_new_with_compression(
+ alignment: usize,
+ write_legacy_ipc_format: bool,
+ metadata_version: ipc::MetadataVersion,
+ batch_compression_type: CompressionCodecType,
+ ) -> Result<Self> {
+ if alignment == 0 || alignment % 8 != 0 {
+ return Err(ArrowError::InvalidArgumentError(
+ "Alignment should be greater than 0 and be a multiple of 8".to_string(),
+ ));
+ }
+ match batch_compression_type {
+ CompressionCodecType::NoCompression => {}
+ _ => {
+ if metadata_version != ipc::MetadataVersion::V5 {
+ return Err(ArrowError::InvalidArgumentError(
+ "Compress buffer just support from metadata v5".to_string(),
+ ));
+ }
+ }
+ };
+ match metadata_version {
+ ipc::MetadataVersion::V5 => {
+ if write_legacy_ipc_format {
+ Err(ArrowError::InvalidArgumentError(
+ "Legacy IPC format only supported on metadata version 4"
+ .to_string(),
+ ))
+ } else {
+ Ok(Self {
+ alignment,
+ write_legacy_ipc_format,
+ metadata_version,
+ batch_compression_type,
+ })
+ }
+ }
+ z => panic!("Unsupported ipc::MetadataVersion {:?}", z),
Review Comment:
Also given that this function returns a `Result` it seems like we could return a proper error here rather than `panic`ing
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org