Convert RunEndEncoded to 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-cast/src/cast/mod.rs

@@ -130,7 +130,8 @@ pub fn can_cast_types(from_type: &DataType, to_type: &DataType) -> bool {
             | FixedSizeList(_, _)
             | Struct(_)
             | Map(_, _)
-            | Dictionary(_, _),
+            | Dictionary(_, _)
+            | RunEndEncoded(_, _),
         ) => true,
         // Dictionary/List conditions should be put in front of others
         (Dictionary(_, from_value_type), Dictionary(_, to_value_type)) => {
@@ -179,6 +180,7 @@ pub fn can_cast_types(from_type: &DataType, to_type: &DataType) -> bool {
                 _ => false,
             }
         }
+        // TODO: RunEndEncoded here?
         // cast one decimal type to another decimal type
         (
             Decimal32(_, _) | Decimal64(_, _) | Decimal128(_, _) | Decimal256(_, _),
@@ -815,6 +817,7 @@ pub fn cast_with_options(
                 "Casting from type {from_type:?} to dictionary type {to_type:?} not supported",
             ))),
         },
+        // TODO: RunEndEncoded here?
         (List(_), List(to)) => cast_list_values::<i32>(array, to, cast_options),
         (LargeList(_), LargeList(to)) => cast_list_values::<i64>(array, to, cast_options),
         (List(_), LargeList(list_to)) => cast_list::<i32, i64>(array, list_to, cast_options),

arrow-schema/src/datatype.rs

@@ -354,7 +354,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

@@ -28,7 +28,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;
 
@@ -59,6 +59,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 {
@@ -90,6 +112,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)),
         }
     }