Support Parquet Byte Stream Split Encoding

parquet/Cargo.toml

@@ -175,6 +175,11 @@ name = "compression"
 required-features = ["experimental", "default"]
 harness = false
 
+[[bench]]
+name = "encoding"
+required-features = ["experimental", "default"]
+harness = false
+
 
 [[bench]]
 name = "metadata"

parquet/benches/encoding.rs

@@ -0,0 +1,83 @@
+// 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 criterion::*;
+use parquet::basic::Encoding;
+use parquet::data_type::{DataType, DoubleType, FloatType};
+use parquet::decoding::{get_decoder, Decoder};
+use parquet::encoding::get_encoder;
+use parquet::schema::types::{ColumnDescPtr, ColumnDescriptor, ColumnPath, Type};
+use rand::prelude::*;
+use std::sync::Arc;
+
+fn bench_typed<T: DataType>(c: &mut Criterion, values: &[T::T], encoding: Encoding) {
+    let name = format!(
+        "dtype={}, encoding={:?}",
+        std::any::type_name::<T::T>(),
+        encoding
+    );
+    c.bench_function(&format!("encoding: {}", name), |b| {
+        b.iter(|| {
+            let mut encoder = get_encoder::<T>(encoding).unwrap();
+            encoder.put(values).unwrap();
+            encoder.flush_buffer().unwrap();
+        });
+    });
+
+    let mut encoder = get_encoder::<T>(encoding).unwrap();
+    encoder.put(values).unwrap();
+    let encoded = encoder.flush_buffer().unwrap();
+    println!("{} encoded as {} bytes", name, encoded.len(),);
+
+    let mut buffer = vec![T::T::default(); values.len()];
+    let column_desc_ptr = ColumnDescPtr::new(ColumnDescriptor::new(
+        Arc::new(
+            Type::primitive_type_builder("", T::get_physical_type())
+                .build()
+                .unwrap(),
+        ),
+        0,
+        0,
+        ColumnPath::new(vec![]),
+    ));
+    c.bench_function(&format!("decoding: {}", name), |b| {
+        b.iter(|| {
+            let mut decoder: Box<dyn Decoder<T>> =
+                get_decoder(column_desc_ptr.clone(), encoding).unwrap();
+            decoder.set_data(encoded.clone(), values.len()).unwrap();
+            decoder.get(&mut buffer).unwrap();
+        });
+    });
+}
+
+fn criterion_benchmark(c: &mut Criterion) {
+    let mut rng = StdRng::seed_from_u64(0);
+    let n = 16 * 1024;
+
+    let mut f32s = Vec::new();
+    let mut f64s = Vec::new();
+    for _ in 0..n {
+        f32s.push(rng.gen::<f32>());
+        f64s.push(rng.gen::<f64>());
+    }
+
+    bench_typed::<FloatType>(c, &f32s, Encoding::BYTE_STREAM_SPLIT);
+    bench_typed::<DoubleType>(c, &f64s, Encoding::BYTE_STREAM_SPLIT);
+}
+
+criterion_group!(benches, criterion_benchmark);
+criterion_main!(benches);

parquet/src/arrow/arrow_reader/mod.rs

@@ -755,7 +755,7 @@ mod tests {
     use crate::column::reader::decoder::REPETITION_LEVELS_BATCH_SIZE;
     use crate::data_type::{
         BoolType, ByteArray, ByteArrayType, DataType, FixedLenByteArray, FixedLenByteArrayType,
-        Int32Type, Int64Type, Int96Type,
+        FloatType, Int32Type, Int64Type, Int96Type,
     };
     use crate::errors::Result;
     use crate::file::properties::{EnabledStatistics, WriterProperties, WriterVersion};
@@ -861,6 +861,13 @@ mod tests {
                 Encoding::DELTA_BINARY_PACKED,
             ],
         );
+        run_single_column_reader_tests::<FloatType, _, FloatType>(
+            2,
+            ConvertedType::NONE,
+            None,
+            |vals| Arc::new(Float32Array::from_iter(vals.iter().cloned())),
+            &[Encoding::PLAIN, Encoding::BYTE_STREAM_SPLIT],
+        );
     }
 
     #[test]
@@ -1390,6 +1397,43 @@ mod tests {
         assert!(col.value(2).is_nan());
     }
 
+    #[test]
+    fn test_read_float32_float64_byte_stream_split() {
+        let path = format!(
+            "{}/byte_stream_split.zstd.parquet",
+            arrow::util::test_util::parquet_test_data(),
+        );
+        let file = File::open(path).unwrap();
+        let record_reader = ParquetRecordBatchReader::try_new(file, 128).unwrap();
+
+        let mut row_count = 0;
+        for batch in record_reader {
+            let batch = batch.unwrap();
+            row_count += batch.num_rows();
+            let f32_col = batch
+                .column(0)
+                .as_any()
+                .downcast_ref::<Float32Array>()
+                .unwrap();
+            let f64_col = batch
+                .column(1)
+                .as_any()
+                .downcast_ref::<Float64Array>()
+                .unwrap();
+
+            // This file contains floats from a standard normal distribution
+            for x in f32_col {
+                assert!(x.unwrap() > -10.0);
+                assert!(x.unwrap() < 10.0);
+            }
+            for x in f64_col {
+                assert!(x.unwrap() > -10.0);
+                assert!(x.unwrap() < 10.0);
+            }
+        }
+        assert_eq!(row_count, 300);
+    }
+
     /// Parameters for single_column_reader_test
     #[derive(Clone)]
     struct TestOptions {

parquet/src/arrow/arrow_writer/mod.rs

@@ -1579,6 +1579,9 @@ mod tests {
             | DataType::UInt32
             | DataType::UInt16
             | DataType::UInt8 => vec![Encoding::PLAIN, Encoding::DELTA_BINARY_PACKED],
+            DataType::Float32 | DataType::Float64 => {
+                vec![Encoding::PLAIN, Encoding::BYTE_STREAM_SPLIT]
+            }
             _ => vec![Encoding::PLAIN],
         };
 

parquet/src/encodings/decoding.rs

@@ -31,6 +31,8 @@ use crate::errors::{ParquetError, Result};
 use crate::schema::types::ColumnDescPtr;
 use crate::util::bit_util::{self, BitReader};
 
+mod byte_stream_split_decoder;
+
 pub(crate) mod private {
     use super::*;
 
@@ -103,8 +105,32 @@ pub(crate) mod private {
         }
     }
 
-    impl GetDecoder for f32 {}
-    impl GetDecoder for f64 {}
+    impl GetDecoder for f32 {
+        fn get_decoder<T: DataType<T = Self>>(
+            descr: ColumnDescPtr,
+            encoding: Encoding,
+        ) -> Result<Box<dyn Decoder<T>>> {
+            match encoding {
+                Encoding::BYTE_STREAM_SPLIT => Ok(Box::new(
+                    byte_stream_split_decoder::ByteStreamSplitDecoder::new(),
+                )),
+                _ => get_decoder_default(descr, encoding),
+            }
+        }
+    }
+    impl GetDecoder for f64 {
+        fn get_decoder<T: DataType<T = Self>>(
+            descr: ColumnDescPtr,
+            encoding: Encoding,
+        ) -> Result<Box<dyn Decoder<T>>> {
+            match encoding {
+                Encoding::BYTE_STREAM_SPLIT => Ok(Box::new(
+                    byte_stream_split_decoder::ByteStreamSplitDecoder::new(),
+                )),
+                _ => get_decoder_default(descr, encoding),
+            }
+        }
+    }
 
     impl GetDecoder for ByteArray {
         fn get_decoder<T: DataType<T = Self>>(
@@ -1760,6 +1786,29 @@ mod tests {
         test_delta_byte_array_decode(data);
     }
 
+    #[test]
+    fn test_byte_stream_split_multiple() {
+        let data = vec![
+            vec![
+                f32::from_le_bytes([0xAA, 0xBB, 0xCC, 0xDD]),
+                f32::from_le_bytes([0x00, 0x11, 0x22, 0x33]),
+            ],
+            vec![f32::from_le_bytes([0xA3, 0xB4, 0xC5, 0xD6])],
+        ];
+        test_byte_stream_split_decode::<FloatType>(data);
+    }
+
+    #[test]
+    fn test_skip_byte_stream_split() {
+        let block_data = vec![0.3, 0.4, 0.1, 4.10];
+        test_skip::<FloatType>(block_data.clone(), Encoding::BYTE_STREAM_SPLIT, 2);
+        test_skip::<DoubleType>(
+            block_data.into_iter().map(|x| x as f64).collect(),
+            Encoding::BYTE_STREAM_SPLIT,
+            100,
+        );
+    }
+
     fn test_rle_value_decode<T: DataType>(data: Vec<Vec<T::T>>) {
         test_encode_decode::<T>(data, Encoding::RLE);
     }
@@ -1768,6 +1817,10 @@ mod tests {
         test_encode_decode::<T>(data, Encoding::DELTA_BINARY_PACKED);
     }
 
+    fn test_byte_stream_split_decode<T: DataType>(data: Vec<Vec<T::T>>) {
+        test_encode_decode::<T>(data, Encoding::BYTE_STREAM_SPLIT);
+    }
+
     fn test_delta_byte_array_decode(data: Vec<Vec<ByteArray>>) {
         test_encode_decode::<ByteArrayType>(data, Encoding::DELTA_BYTE_ARRAY);
     }

parquet/src/encodings/decoding/byte_stream_split_decoder.rs

@@ -0,0 +1,121 @@
+// 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 std::marker::PhantomData;
+
+use bytes::Bytes;
+
+use crate::basic::Encoding;
+use crate::data_type::{DataType, SliceAsBytes};
+use crate::errors::{ParquetError, Result};
+
+use super::Decoder;
+
+pub struct ByteStreamSplitDecoder<T: DataType> {
+    _phantom: PhantomData<T>,
+    encoded_bytes: Bytes,
+    total_num_values: usize,
+    values_decoded: usize,
+}
+
+impl<T: DataType> ByteStreamSplitDecoder<T> {
+    pub(crate) fn new() -> Self {
+        Self {
+            _phantom: PhantomData,
+            encoded_bytes: Bytes::new(),
+            total_num_values: 0,
+            values_decoded: 0,
+        }
+    }
+}
+
+// Here we assume src contains the full data (which it must, since we're
+// can only know where to split the streams once all data is collected),
+// but dst can be just a slice starting from the given index.
+// We iterate over the output bytes and fill them in from their strided
+// input byte locations.
+fn join_streams_const<const TYPE_SIZE: usize>(
+    src: &[u8],
+    dst: &mut [u8],
+    stride: usize,
+    values_decoded: usize,
+) {
+    let sub_src = &src[values_decoded..];
+    for i in 0..dst.len() / TYPE_SIZE {
+        for j in 0..TYPE_SIZE {
+            dst[i * TYPE_SIZE + j] = sub_src[i + j * stride];
+        }
+    }
+}
+
+impl<T: DataType> Decoder<T> for ByteStreamSplitDecoder<T> {
+    fn set_data(&mut self, data: Bytes, num_values: usize) -> Result<()> {
+        self.encoded_bytes = data;
+        self.total_num_values = num_values;
+        self.values_decoded = 0;
+
+        Ok(())
+    }
+
+    fn get(&mut self, buffer: &mut [<T as DataType>::T]) -> Result<usize> {
+        let total_remaining_values = self.values_left();
+        let num_values = buffer.len().min(total_remaining_values);
+        let buffer = &mut buffer[..num_values];
+
+        // SAFETY: f32 and f64 has no constraints on their internal representation, so we can modify it as we want
+        let raw_out_bytes = unsafe { <T as DataType>::T::slice_as_bytes_mut(buffer) };
+        let type_size = T::get_type_size();
+        let stride = self.encoded_bytes.len() / type_size;
+        match type_size {
+            4 => join_streams_const::<4>(
+                &self.encoded_bytes,
+                raw_out_bytes,
+                stride,
+                self.values_decoded,
+            ),
+            8 => join_streams_const::<8>(
+                &self.encoded_bytes,
+                raw_out_bytes,
+                stride,
+                self.values_decoded,
+            ),
+            _ => {
+                return Err(general_err!(
+                    "byte stream split unsupported for data types of size {} bytes",
+                    type_size
+                ));
+            }
+        }
+        self.values_decoded += num_values;
+
+        Ok(num_values)
+    }
+
+    fn values_left(&self) -> usize {
+        self.total_num_values - self.values_decoded
+    }
+
+    fn encoding(&self) -> Encoding {
+        Encoding::BYTE_STREAM_SPLIT
+    }
+
+    fn skip(&mut self, num_values: usize) -> Result<usize> {
+        let to_skip = usize::min(self.values_left(), num_values);
+        self.values_decoded += to_skip;
+        Ok(to_skip)
+    }
+}