8069: Support writing RunEndEncoded as Parquet

arrow-array/src/array/run_array.rs

@@ -32,12 +32,12 @@ use crate::{
 
 /// An array of [run-end encoded values](https://arrow.apache.org/docs/format/Columnar.html#run-end-encoded-layout)
 ///
-/// This encoding is variation on [run-length encoding (RLE)](https://en.wikipedia.org/wiki/Run-length_encoding)
+/// This encoding is a variation on [run-length encoding (RLE)](https://en.wikipedia.org/wiki/Run-length_encoding)
 /// and is good for representing data containing same values repeated consecutively.
 ///
 /// [`RunArray`] contains `run_ends` array and `values` array of same length.
 /// The `run_ends` array stores the indexes at which the run ends. The `values` array
-/// stores the value of each run. Below example illustrates how a logical array is represented in
+/// stores the value of each run. The below example illustrates how a logical array is represented in
 /// [`RunArray`]
 ///
 ///

arrow-schema/src/datatype.rs

@@ -353,7 +353,7 @@ pub enum DataType {
     /// that contain many repeated values using less memory, but with
     /// a higher CPU overhead for some operations.
     ///
-    /// This type mostly used to represent low cardinality string
+    /// This type is mostly used to represent low cardinality string
     /// arrays or a limited set of primitive types as integers.
     Dictionary(Box<DataType>, Box<DataType>),
     /// Exact 32-bit width decimal value with precision and scale

parquet/src/arrow/array_reader/byte_array.rs

@@ -67,6 +67,28 @@ pub fn make_byte_array_reader(
                 pages, data_type, reader,
             )))
         }
+        // TODO eventually add a dedicated [`ArrayReader`] for REE
+        ArrowType::RunEndEncoded(_, ref val_field) => match val_field.data_type() {
+            ArrowType::Binary
+            | ArrowType::Utf8
+            | ArrowType::Decimal128(_, _)
+            | ArrowType::Decimal256(_, _) => {
+                let reader = GenericRecordReader::new(column_desc);
+                Ok(Box::new(ByteArrayReader::<i32>::new(
+                    pages, data_type, reader,
+                )))
+            }
+            ArrowType::LargeUtf8 | ArrowType::LargeBinary => {
+                let reader = GenericRecordReader::new(column_desc);
+                Ok(Box::new(ByteArrayReader::<i64>::new(
+                    pages, data_type, reader,
+                )))
+            }
+            _ => Err(general_err!(
+                "invalid run end encoded value type for byte array reader - {}",
+                data_type
+            )),
+        },
         _ => Err(general_err!(
             "invalid data type for byte array reader - {}",
             data_type
@@ -147,6 +169,11 @@ impl<I: OffsetSizeTrait> ArrayReader for ByteArrayReader<I> {
                 .with_precision_and_scale(p, s)?;
                 Arc::new(decimal)
             }
+            // TODO eventually add a dedicated [`ArrayReader`] for REE
+            ArrowType::RunEndEncoded(_, ref val_field) => {
+                let array = buffer.into_array(null_buffer, val_field.data_type().clone());
+                arrow_cast::cast(&array, &self.data_type)?
+            }
             _ => buffer.into_array(null_buffer, self.data_type.clone()),
         };
 

parquet/src/arrow/arrow_writer/byte_array.rs

@@ -30,7 +30,7 @@ use crate::util::bit_util::num_required_bits;
 use crate::util::interner::{Interner, Storage};
 use arrow_array::{
     Array, ArrayAccessor, BinaryArray, BinaryViewArray, DictionaryArray, FixedSizeBinaryArray,
-    LargeBinaryArray, LargeStringArray, StringArray, StringViewArray,
+    LargeBinaryArray, LargeStringArray, RunArray, StringArray, StringViewArray,
 };
 use arrow_schema::DataType;
 
@@ -61,6 +61,28 @@ macro_rules! downcast_dict_op {
     };
 }
 
+macro_rules! downcast_ree_impl {
+    ($array:ident, $key:ident, $val:ident, $op:expr $(, $arg:expr)*) => {{
+        $op($array
+            .as_any()
+            .downcast_ref::<RunArray<arrow_array::types::$key>>()
+            .unwrap()
+            .downcast::<$val>()
+            .unwrap()$(, $arg)*)
+    }};
+}
+
+macro_rules! downcast_ree_op {
+    ($run_end_field:expr, $val:ident, $array:ident, $op:expr $(, $arg:expr)*) => {
+        match $run_end_field.data_type() {
+            DataType::Int16 => downcast_ree_impl!($array, Int16Type, $val, $op$(, $arg)*),
+            DataType::Int32 => downcast_ree_impl!($array, Int32Type, $val, $op$(, $arg)*),
+            DataType::Int64 => downcast_ree_impl!($array, Int64Type, $val, $op$(, $arg)*),
+            _ => unreachable!(),
+        }
+    };
+}
+
 macro_rules! downcast_op {
     ($data_type:expr, $array:ident, $op:expr $(, $arg:expr)*) => {
         match $data_type {
@@ -92,6 +114,20 @@ macro_rules! downcast_op {
                 }
                 d => unreachable!("cannot downcast {} dictionary value to byte array", d),
             },
+            DataType::RunEndEncoded(run_end, value) => match value.data_type() {
+                DataType::Utf8 => downcast_ree_op!(run_end, StringArray, $array, $op$(, $arg)*),
+                DataType::LargeUtf8 => {
+                    downcast_ree_op!(run_end, LargeStringArray, $array, $op$(, $arg)*)
+                }
+                DataType::Binary => downcast_ree_op!(run_end, BinaryArray, $array, $op$(, $arg)*),
+                DataType::LargeBinary => {
+                    downcast_ree_op!(run_end, LargeBinaryArray, $array, $op$(, $arg)*)
+                }
+                DataType::FixedSizeBinary(_) => {
+                    downcast_ree_op!(run_end, FixedSizeBinaryArray, $array, $op$(, $arg)*)
+                }
+                d => unreachable!("cannot downcast {} run end encoded value to byte array", d),
+            },
             d => unreachable!("cannot downcast {} to byte array", d),
         }
     };

parquet/src/arrow/arrow_writer/levels.rs

@@ -222,6 +222,10 @@ impl LevelInfoBuilder {
                     _ => unreachable!(),
                 })
             }
+            DataType::RunEndEncoded(_, v) if is_leaf(v.data_type()) => {
+                let levels = ArrayLevels::new(parent_ctx, is_nullable, array.clone());
+                Ok(Self::Primitive(levels))
+            }
             d => Err(nyi_err!("Datatype {} is not yet supported", d)),
         }
     }

parquet/src/arrow/arrow_writer/mod.rs

@@ -1122,6 +1122,17 @@ impl ArrowColumnWriterFactory {
                 ArrowDataType::FixedSizeBinary(_) => out.push(bytes(leaves.next().unwrap())?),
                 _ => out.push(col(leaves.next().unwrap())?),
             },
+            ArrowDataType::RunEndEncoded(_, value_type) => match value_type.data_type() {
+                ArrowDataType::Utf8
+                | ArrowDataType::LargeUtf8
+                | ArrowDataType::Binary
+                | ArrowDataType::LargeBinary => out.push(bytes(leaves.next().unwrap())?),
+                ArrowDataType::Utf8View | ArrowDataType::BinaryView => {
+                    out.push(bytes(leaves.next().unwrap())?)
+                }
+                ArrowDataType::FixedSizeBinary(_) => out.push(bytes(leaves.next().unwrap())?),
+                _ => out.push(col(leaves.next().unwrap())?),
+            },
             _ => {
                 return Err(ParquetError::NYI(format!(
                     "Attempting to write an Arrow type {data_type} to parquet that is not yet implemented"
@@ -1215,6 +1226,41 @@ fn write_leaf(writer: &mut ColumnWriter<'_>, levels: &ArrayLevels) -> Result<usi
                         write_primitive(typed, array.values(), levels)
                     }
                 },
+                ArrowDataType::RunEndEncoded(_, value_type) => match value_type.data_type() {
+                    ArrowDataType::Decimal32(_, _) => {
+                        let array = arrow_cast::cast(column, value_type.data_type())?;
+                        let array = array
+                            .as_primitive::<Decimal32Type>()
+                            .unary::<_, Int32Type>(|v| v);
+                        write_primitive(typed, array.values(), levels)
+                    }
+                    ArrowDataType::Decimal64(_, _) => {
+                        let array = arrow_cast::cast(column, value_type.data_type())?;
+                        let array = array
+                            .as_primitive::<Decimal64Type>()
+                            .unary::<_, Int32Type>(|v| v as i32);
+                        write_primitive(typed, array.values(), levels)
+                    }
+                    ArrowDataType::Decimal128(_, _) => {
+                        let array = arrow_cast::cast(column, value_type.data_type())?;
+                        let array = array
+                            .as_primitive::<Decimal128Type>()
+                            .unary::<_, Int32Type>(|v| v as i32);
+                        write_primitive(typed, array.values(), levels)
+                    }
+                    ArrowDataType::Decimal256(_, _) => {
+                        let array = arrow_cast::cast(column, value_type.data_type())?;
+                        let array = array
+                            .as_primitive::<Decimal256Type>()
+                            .unary::<_, Int32Type>(|v| v.as_i128() as i32);
+                        write_primitive(typed, array.values(), levels)
+                    }
+                    _ => {
+                        let array = arrow_cast::cast(column, &ArrowDataType::Int32)?;
+                        let array = array.as_primitive::<Int32Type>();
+                        write_primitive(typed, array.values(), levels)
+                    }
+                },
                 _ => {
                     let array = arrow_cast::cast(column, &ArrowDataType::Int32)?;
                     let array = array.as_primitive::<Int32Type>();
@@ -1297,6 +1343,12 @@ fn write_leaf(writer: &mut ColumnWriter<'_>, levels: &ArrayLevels) -> Result<usi
                         write_primitive(typed, array.values(), levels)
                     }
                 },
+                ArrowDataType::RunEndEncoded(_run_ends, _values) => {
+                    Err(ParquetError::NYI(
+                        "Int64ColumnWriter: Attempting to write an Arrow REE type that is not yet implemented"
+                            .to_string(),
+                    ))
+                }
                 _ => {
                     let array = arrow_cast::cast(column, &ArrowDataType::Int64)?;
                     let array = array.as_primitive::<Int64Type>();
@@ -1371,6 +1423,12 @@ fn write_leaf(writer: &mut ColumnWriter<'_>, levels: &ArrayLevels) -> Result<usi
                     let array = column.as_primitive::<Float16Type>();
                     get_float_16_array_slice(array, indices)
                 }
+                ArrowDataType::RunEndEncoded(_run_ends, _values) => {
+                    return Err(ParquetError::NYI(
+                        "FixedLenByteArrayColumnWriter: Attempting to write an Arrow REE type that is not yet implemented"
+                            .to_string(),
+                    ));
+                }
                 _ => {
                     return Err(ParquetError::NYI(
                         "Attempting to write an Arrow type that is not yet implemented".to_string(),
@@ -4481,4 +4539,153 @@ mod tests {
         assert_eq!(get_dict_page_size(col0_meta), 1024 * 1024);
         assert_eq!(get_dict_page_size(col1_meta), 1024 * 1024 * 4);
     }
+
+    #[test]
+    fn arrow_writer_run_end_encoded_string() {
+        // Create a run array of strings
+        let mut builder = StringRunBuilder::<Int32Type>::new();
+        builder.extend(
+            vec![Some("alpha"); 100000]
+                .into_iter()
+                .chain(vec![Some("beta"); 100000]),
+        );
+        let run_array: RunArray<Int32Type> = builder.finish();
+        let schema = Arc::new(Schema::new(vec![Field::new(
+            "ree",
+            run_array.data_type().clone(),
+            run_array.is_nullable(),
+        )]));
+
+        // Write to parquet
+        let mut parquet_bytes: Vec<u8> = Vec::new();
+        let mut writer = ArrowWriter::try_new(&mut parquet_bytes, schema.clone(), None).unwrap();
+        let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(run_array)]).unwrap();
+        writer.write(&batch).unwrap();
+        writer.close().unwrap();
+
+        // Read back and verify
+        let bytes = Bytes::from(parquet_bytes);
+        let reader = ParquetRecordBatchReaderBuilder::try_new(bytes).unwrap();
+
+        // Check if dictionary was used by examining the metadata
+        let metadata = reader.metadata();
+        let row_group = &metadata.row_groups()[0];
+        let col_meta = &row_group.columns()[0];
+
+        // If dictionary encoding worked, we should see RLE_DICTIONARY encoding
+        // and have a dictionary page offset
+        let has_dict_encoding = col_meta.encodings().any(|e| e == Encoding::RLE_DICTIONARY);
+        let has_dict_page = col_meta.dictionary_page_offset().is_some();
+
+        // Verify the schema is REE encoded when we read it back
+        let expected_schema = Arc::new(Schema::new(vec![Field::new(
+            "ree",
+            DataType::RunEndEncoded(
+                Arc::new(Field::new("run_ends", arrow_schema::DataType::Int32, false)),
+                Arc::new(Field::new("values", arrow_schema::DataType::Utf8, true)),
+            ),
+            false,
+        )]));
+        assert_eq!(&expected_schema, reader.schema());
+
+        // Read the data back
+        let batches: Vec<_> = reader
+            .build()
+            .unwrap()
+            .collect::<ArrowResult<Vec<_>>>()
+            .unwrap();
+        assert_eq!(batches.len(), 196);
+        // Count rows in total
+        let total_rows = batches.iter().map(|b| b.num_rows()).sum::<usize>();
+        assert_eq!(total_rows, 200000);
+
+        // Ensure dictionary encoding
+        assert!(has_dict_encoding, "RunArray should be dictionary encoded");
+        assert!(has_dict_page, "RunArray should have dictionary page");
+    }
+
+    #[test]
+    fn arrow_writer_run_end_encoded_int() {
+        // Create a run array of strings
+        let mut builder = PrimitiveRunBuilder::<Int32Type, Int32Type>::new();
+        builder.extend(
+            vec![Some(1); 100000]
+                .into_iter()
+                .chain(vec![Some(2); 100000]),
+        );
+        let run_array: RunArray<Int32Type> = builder.finish();
+        let schema = Arc::new(Schema::new(vec![Field::new(
+            "ree",
+            run_array.data_type().clone(),
+            run_array.is_nullable(),
+        )]));
+
+        // Write to parquet
+        let mut parquet_bytes: Vec<u8> = Vec::new();
+        let mut writer = ArrowWriter::try_new(&mut parquet_bytes, schema.clone(), None).unwrap();
+        let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(run_array)]).unwrap();
+        writer.write(&batch).unwrap();
+        writer.close().unwrap();
+
+        // Read back and verify
+        let bytes = Bytes::from(parquet_bytes);
+        let reader = ParquetRecordBatchReaderBuilder::try_new(bytes).unwrap();
+
+        // Check if dictionary was used by examining the metadata
+        let metadata = reader.metadata();
+        let row_group = &metadata.row_groups()[0];
+        let col_meta = &row_group.columns()[0];
+        let has_dict_encoding = col_meta.encodings().any(|e| e == Encoding::RLE_DICTIONARY);
+
+        // If dictionary encoding worked, we should see RLE_DICTIONARY encoding
+        // and have a dictionary page offset
+        // let has_dict_encoding = col_meta.encodings().contains(&Encoding::RLE_DICTIONARY);
+        let has_dict_page = col_meta.dictionary_page_offset().is_some();
+
+        // Verify the schema is REE encoded when we read it back
+        let expected_schema = Arc::new(Schema::new(vec![Field::new(
+            "ree",
+            DataType::RunEndEncoded(
+                Arc::new(Field::new("run_ends", arrow_schema::DataType::Int32, false)),
+                Arc::new(Field::new("values", arrow_schema::DataType::Int32, true)),
+            ),
+            false,
+        )]));
+        assert_eq!(&expected_schema, reader.schema());
+
+        // Read the data back
+        let batches: Vec<_> = reader
+            .build()
+            .unwrap()
+            .collect::<ArrowResult<Vec<_>>>()
+            .unwrap();
+        assert_eq!(batches.len(), 196);
+        // Count rows in total
+        let total_rows = batches.iter().map(|b| b.num_rows()).sum::<usize>();
+        assert_eq!(total_rows, 200000);
+
+        // Ensure dictionary encoding
+        assert!(has_dict_encoding, "RunArray should be dictionary encoded");
+        assert!(has_dict_page, "RunArray should have dictionary page");
+    }
+
+    #[test]
+    fn arrow_writer_round_trip_run_end_encoded_string() {
+        // Create a run array of strings (cannot have more than 1024 values per record batch)
+        let mut builder = StringRunBuilder::<Int32Type>::new();
+        builder.extend(
+            vec![Some("alpha"); 512]
+                .into_iter()
+                .chain(vec![Some("beta"); 512]),
+        );
+        let run_array: RunArray<Int32Type> = builder.finish();
+        let schema = Arc::new(Schema::new(vec![Field::new(
+            "ree",
+            run_array.data_type().clone(),
+            run_array.is_nullable(),
+        )]));
+
+        let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(run_array)]).unwrap();
+        roundtrip(batch, None);
+    }
 }