Fix extract parquet statistics from Decimal256 columns (#10777)

xinlifoobar · alamb · web-flow · commit 76a372009586 · 2024-06-04T09:17:25.000-04:00
* Fix Extract parquet statistics from Decimal256 columns

* Fix comment

---------

Co-authored-by: Andrew Lamb &lt;andrew@nerdnetworks.org&gt;
diff --git a/datafusion/core/src/datasource/physical_plan/parquet/statistics.rs b/datafusion/core/src/datasource/physical_plan/parquet/statistics.rs
@@ -19,7 +19,7 @@
 
 // TODO: potentially move this to arrow-rs: https://github.com/apache/arrow-rs/issues/4328
 
-use arrow::{array::ArrayRef, datatypes::DataType, datatypes::TimeUnit};
+use arrow::{array::ArrayRef, datatypes::i256, datatypes::DataType, datatypes::TimeUnit};
 use arrow_array::{new_empty_array, new_null_array, UInt64Array};
 use arrow_schema::{Field, FieldRef, Schema};
 use datafusion_common::{
@@ -36,7 +36,13 @@ pub(crate) fn from_bytes_to_i128(b: &[u8]) -> i128 {
     // The bytes array are from parquet file and must be the big-endian.
     // The endian is defined by parquet format, and the reference document
     // https://github.com/apache/parquet-format/blob/54e53e5d7794d383529dd30746378f19a12afd58/src/main/thrift/parquet.thrift#L66
-    i128::from_be_bytes(sign_extend_be(b))
+    i128::from_be_bytes(sign_extend_be::<16>(b))
+}
+
+// Convert the bytes array to i256.
+// The endian of the input bytes array must be big-endian.
+pub(crate) fn from_bytes_to_i256(b: &[u8]) -> i256 {
+    i256::from_be_bytes(sign_extend_be::<32>(b))
 }
 
 // Convert the bytes array to f16
@@ -48,13 +54,13 @@ pub(crate) fn from_bytes_to_f16(b: &[u8]) -> Option<f16> {
 }
 
 // Copy from arrow-rs
-// https://github.com/apache/arrow-rs/blob/733b7e7fd1e8c43a404c3ce40ecf741d493c21b4/parquet/src/arrow/buffer/bit_util.rs#L55
-// Convert the byte slice to fixed length byte array with the length of 16
-fn sign_extend_be(b: &[u8]) -> [u8; 16] {
-    assert!(b.len() <= 16, "Array too large, expected less than 16");
+// https://github.com/apache/arrow-rs/blob/198af7a3f4aa20f9bd003209d9f04b0f37bb120e/parquet/src/arrow/buffer/bit_util.rs#L54
+// Convert the byte slice to fixed length byte array with the length of N.
+pub fn sign_extend_be<const N: usize>(b: &[u8]) -> [u8; N] {
+    assert!(b.len() <= N, "Array too large, expected less than {N}");
     let is_negative = (b[0] & 128u8) == 128u8;
-    let mut result = if is_negative { [255u8; 16] } else { [0u8; 16] };
-    for (d, s) in result.iter_mut().skip(16 - b.len()).zip(b) {
+    let mut result = if is_negative { [255u8; N] } else { [0u8; N] };
+    for (d, s) in result.iter_mut().skip(N - b.len()).zip(b) {
         *d = *s;
     }
     result
@@ -83,6 +89,13 @@ macro_rules! get_statistic {
                             *scale,
                         ))
                     }
+                    Some(DataType::Decimal256(precision, scale)) => {
+                        Some(ScalarValue::Decimal256(
+                            Some(i256::from(*s.$func())),
+                            *precision,
+                            *scale,
+                        ))
+                    }
                     Some(DataType::Int8) => {
                         Some(ScalarValue::Int8(Some((*s.$func()).try_into().unwrap())))
                     }
@@ -123,6 +136,13 @@ macro_rules! get_statistic {
                             *scale,
                         ))
                     }
+                    Some(DataType::Decimal256(precision, scale)) => {
+                        Some(ScalarValue::Decimal256(
+                            Some(i256::from(*s.$func())),
+                            *precision,
+                            *scale,
+                        ))
+                    }
                     Some(DataType::UInt64) => {
                         Some(ScalarValue::UInt64(Some((*s.$func()) as u64)))
                     }
@@ -169,6 +189,13 @@ macro_rules! get_statistic {
                             *scale,
                         ))
                     }
+                    Some(DataType::Decimal256(precision, scale)) => {
+                        Some(ScalarValue::Decimal256(
+                            Some(from_bytes_to_i256(s.$bytes_func())),
+                            *precision,
+                            *scale,
+                        ))
+                    }
                     Some(DataType::Binary) => {
                         Some(ScalarValue::Binary(Some(s.$bytes_func().to_vec())))
                     }
@@ -202,6 +229,13 @@ macro_rules! get_statistic {
                             *scale,
                         ))
                     }
+                    Some(DataType::Decimal256(precision, scale)) => {
+                        Some(ScalarValue::Decimal256(
+                            Some(from_bytes_to_i256(s.$bytes_func())),
+                            *precision,
+                            *scale,
+                        ))
+                    }
                     Some(DataType::FixedSizeBinary(size)) => {
                         let value = s.$bytes_func().to_vec();
                         let value = if value.len().try_into() == Ok(*size) {
@@ -438,13 +472,13 @@ impl<'a> StatisticsConverter<'a> {
 mod test {
     use super::*;
     use arrow::compute::kernels::cast_utils::Parser;
-    use arrow::datatypes::{Date32Type, Date64Type};
+    use arrow::datatypes::{i256, Date32Type, Date64Type};
     use arrow_array::{
         new_null_array, Array, BinaryArray, BooleanArray, Date32Array, Date64Array,
-        Decimal128Array, Float32Array, Float64Array, Int16Array, Int32Array, Int64Array,
-        Int8Array, LargeBinaryArray, RecordBatch, StringArray, StructArray,
-        TimestampMicrosecondArray, TimestampMillisecondArray, TimestampNanosecondArray,
-        TimestampSecondArray,
+        Decimal128Array, Decimal256Array, Float32Array, Float64Array, Int16Array,
+        Int32Array, Int64Array, Int8Array, LargeBinaryArray, RecordBatch, StringArray,
+        StructArray, TimestampMicrosecondArray, TimestampMillisecondArray,
+        TimestampNanosecondArray, TimestampSecondArray,
     };
     use arrow_schema::{Field, SchemaRef};
     use bytes::Bytes;
@@ -824,6 +858,42 @@ mod test {
                     .unwrap(),
             ),
         }
+        .run();
+
+        Test {
+            input: Arc::new(
+                Decimal256Array::from(vec![
+                    // row group 1
+                    Some(i256::from(100)),
+                    None,
+                    Some(i256::from(22000)),
+                    // row group 2
+                    Some(i256::MAX),
+                    Some(i256::MIN),
+                    None,
+                    // row group 3
+                    None,
+                    None,
+                    None,
+                ])
+                .with_precision_and_scale(76, 76)
+                .unwrap(),
+            ),
+            expected_min: Arc::new(
+                Decimal256Array::from(vec![Some(i256::from(100)), Some(i256::MIN), None])
+                    .with_precision_and_scale(76, 76)
+                    .unwrap(),
+            ),
+            expected_max: Arc::new(
+                Decimal256Array::from(vec![
+                    Some(i256::from(22000)),
+                    Some(i256::MAX),
+                    None,
+                ])
+                .with_precision_and_scale(76, 76)
+                .unwrap(),
+            ),
+        }
         .run()
     }
 
diff --git a/datafusion/core/tests/parquet/arrow_statistics.rs b/datafusion/core/tests/parquet/arrow_statistics.rs
@@ -24,17 +24,17 @@ use std::sync::Arc;
 use crate::parquet::{struct_array, Scenario};
 use arrow::compute::kernels::cast_utils::Parser;
 use arrow::datatypes::{
-    Date32Type, Date64Type, TimestampMicrosecondType, TimestampMillisecondType,
+    i256, Date32Type, Date64Type, TimestampMicrosecondType, TimestampMillisecondType,
     TimestampNanosecondType, TimestampSecondType,
 };
 use arrow_array::{
     make_array, Array, ArrayRef, BinaryArray, BooleanArray, Date32Array, Date64Array,
-    Decimal128Array, FixedSizeBinaryArray, Float16Array, Float32Array, Float64Array,
-    Int16Array, Int32Array, Int64Array, Int8Array, LargeBinaryArray, LargeStringArray,
-    RecordBatch, StringArray, Time32MillisecondArray, Time32SecondArray,
-    Time64MicrosecondArray, Time64NanosecondArray, TimestampMicrosecondArray,
-    TimestampMillisecondArray, TimestampNanosecondArray, TimestampSecondArray,
-    UInt16Array, UInt32Array, UInt64Array, UInt8Array,
+    Decimal128Array, Decimal256Array, FixedSizeBinaryArray, Float16Array, Float32Array,
+    Float64Array, Int16Array, Int32Array, Int64Array, Int8Array, LargeBinaryArray,
+    LargeStringArray, RecordBatch, StringArray, Time32MillisecondArray,
+    Time32SecondArray, Time64MicrosecondArray, Time64NanosecondArray,
+    TimestampMicrosecondArray, TimestampMillisecondArray, TimestampNanosecondArray,
+    TimestampSecondArray, UInt16Array, UInt32Array, UInt64Array, UInt8Array,
 };
 use arrow_schema::{DataType, Field, Schema};
 use datafusion::datasource::physical_plan::parquet::{
@@ -1356,6 +1356,38 @@ async fn test_decimal() {
         column_name: "decimal_col",
     }
     .run();
+
+    // This creates a parquet file of 1 column "decimal256_col" with decimal data type and precicion 9, scale 2
+    // file has 3 record batches, each has 5 rows. They will be saved into 3 row groups
+    let decimal256_reader = TestReader {
+        scenario: Scenario::Decimal256,
+        row_per_group: 5,
+    };
+    Test {
+        reader: decimal256_reader.build().await,
+        expected_min: Arc::new(
+            Decimal256Array::from(vec![
+                i256::from(100),
+                i256::from(-500),
+                i256::from(2000),
+            ])
+            .with_precision_and_scale(9, 2)
+            .unwrap(),
+        ),
+        expected_max: Arc::new(
+            Decimal256Array::from(vec![
+                i256::from(600),
+                i256::from(600),
+                i256::from(6000),
+            ])
+            .with_precision_and_scale(9, 2)
+            .unwrap(),
+        ),
+        expected_null_counts: UInt64Array::from(vec![0, 0, 0]),
+        expected_row_counts: UInt64Array::from(vec![5, 5, 5]),
+        column_name: "decimal256_col",
+    }
+    .run();
 }
 #[tokio::test]
 async fn test_dictionary() {
diff --git a/datafusion/core/tests/parquet/mod.rs b/datafusion/core/tests/parquet/mod.rs
@@ -17,15 +17,17 @@
 
 //! Parquet integration tests
 use arrow::array::Decimal128Array;
+use arrow::datatypes::i256;
 use arrow::{
     array::{
         make_array, Array, ArrayRef, BinaryArray, BooleanArray, Date32Array, Date64Array,
-        DictionaryArray, FixedSizeBinaryArray, Float16Array, Float32Array, Float64Array,
-        Int16Array, Int32Array, Int64Array, Int8Array, LargeBinaryArray,
-        LargeStringArray, StringArray, StructArray, Time32MillisecondArray,
-        Time32SecondArray, Time64MicrosecondArray, Time64NanosecondArray,
-        TimestampMicrosecondArray, TimestampMillisecondArray, TimestampNanosecondArray,
-        TimestampSecondArray, UInt16Array, UInt32Array, UInt64Array, UInt8Array,
+        Decimal256Array, DictionaryArray, FixedSizeBinaryArray, Float16Array,
+        Float32Array, Float64Array, Int16Array, Int32Array, Int64Array, Int8Array,
+        LargeBinaryArray, LargeStringArray, StringArray, StructArray,
+        Time32MillisecondArray, Time32SecondArray, Time64MicrosecondArray,
+        Time64NanosecondArray, TimestampMicrosecondArray, TimestampMillisecondArray,
+        TimestampNanosecondArray, TimestampSecondArray, UInt16Array, UInt32Array,
+        UInt64Array, UInt8Array,
     },
     datatypes::{DataType, Field, Int32Type, Int8Type, Schema, TimeUnit},
     record_batch::RecordBatch,
@@ -84,6 +86,7 @@ enum Scenario {
     Float16,
     Float64,
     Decimal,
+    Decimal256,
     DecimalBloomFilterInt32,
     DecimalBloomFilterInt64,
     DecimalLargePrecision,
@@ -618,6 +621,24 @@ fn make_decimal_batch(v: Vec<i128>, precision: u8, scale: i8) -> RecordBatch {
     RecordBatch::try_new(schema, vec![array.clone()]).unwrap()
 }
 
+/// Return record batch with decimal256 vector
+///
+/// Columns are named
+/// "decimal256_col" -> Decimal256Array
+fn make_decimal256_batch(v: Vec<i256>, precision: u8, scale: i8) -> RecordBatch {
+    let schema = Arc::new(Schema::new(vec![Field::new(
+        "decimal256_col",
+        DataType::Decimal256(precision, scale),
+        true,
+    )]));
+    let array = Arc::new(
+        Decimal256Array::from(v)
+            .with_precision_and_scale(precision, scale)
+            .unwrap(),
+    ) as ArrayRef;
+    RecordBatch::try_new(schema, vec![array]).unwrap()
+}
+
 /// Return record batch with a few rows of data for all of the supported date
 /// types with the specified offset (in days)
 ///
@@ -1009,6 +1030,44 @@ fn create_data_batch(scenario: Scenario) -> Vec<RecordBatch> {
                 make_decimal_batch(vec![2000, 3000, 3000, 4000, 6000], 9, 2),
             ]
         }
+        Scenario::Decimal256 => {
+            // decimal256 record batch
+            vec![
+                make_decimal256_batch(
+                    vec![
+                        i256::from(100),
+                        i256::from(200),
+                        i256::from(300),
+                        i256::from(400),
+                        i256::from(600),
+                    ],
+                    9,
+                    2,
+                ),
+                make_decimal256_batch(
+                    vec![
+                        i256::from(-500),
+                        i256::from(100),
+                        i256::from(300),
+                        i256::from(400),
+                        i256::from(600),
+                    ],
+                    9,
+                    2,
+                ),
+                make_decimal256_batch(
+                    vec![
+                        i256::from(2000),
+                        i256::from(3000),
+                        i256::from(3000),
+                        i256::from(4000),
+                        i256::from(6000),
+                    ],
+                    9,
+                    2,
+                ),
+            ]
+        }
         Scenario::DecimalBloomFilterInt32 => {
             // decimal record batch
             vec![
