Convert RunEndEncoded field to Parquet

vegarsti · vegarsti · commit ec40c7f30d32 · 2025-08-09T08:41:57.000+02:00
diff --git a/arrow-schema/src/datatype.rs b/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
diff --git a/parquet/src/arrow/arrow_writer/byte_array.rs b/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),
         }
     };
diff --git a/parquet/src/arrow/arrow_writer/levels.rs b/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)),
         }
     }
diff --git a/parquet/src/arrow/arrow_writer/mod.rs b/parquet/src/arrow/arrow_writer/mod.rs
@@ -1033,15 +1033,15 @@ impl ArrowColumnWriterFactory {
 
         match data_type {
             _ if data_type.is_primitive() => out.push(col(leaves.next().unwrap())?),
-            ArrowDataType::FixedSizeBinary(_) | ArrowDataType::Boolean | ArrowDataType::Null => out.push(col(leaves.next().unwrap())?),
+            ArrowDataType::FixedSizeBinary(_) | ArrowDataType::Boolean | ArrowDataType::Null => {
+                out.push(col(leaves.next().unwrap())?)
+            }
             ArrowDataType::LargeBinary
             | ArrowDataType::Binary
             | ArrowDataType::Utf8
             | ArrowDataType::LargeUtf8
             | ArrowDataType::BinaryView
-            | ArrowDataType::Utf8View => {
-                out.push(bytes(leaves.next().unwrap())?)
-            }
+            | ArrowDataType::Utf8View => out.push(bytes(leaves.next().unwrap())?),
             ArrowDataType::List(f)
             | ArrowDataType::LargeList(f)
             | ArrowDataType::FixedSizeList(f, _) => {
@@ -1058,21 +1058,30 @@ impl ArrowColumnWriterFactory {
                     self.get_arrow_column_writer(f[1].data_type(), props, leaves, out)?
                 }
                 _ => unreachable!("invalid map type"),
-            }
+            },
             ArrowDataType::Dictionary(_, value_type) => match value_type.as_ref() {
-                ArrowDataType::Utf8 | ArrowDataType::LargeUtf8 | ArrowDataType::Binary | ArrowDataType::LargeBinary => {
-                    out.push(bytes(leaves.next().unwrap())?)
-                }
+                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(_) => {
+                ArrowDataType::FixedSizeBinary(_) => out.push(bytes(leaves.next().unwrap())?),
+                _ => out.push(col(leaves.next().unwrap())?),
+            },
+            // TODO: Don't know what I'm doing here!
+            ArrowDataType::RunEndEncoded(_run_ends, 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())?)
                 }
-                _ => {
-                    out.push(col(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"
@@ -1166,6 +1175,7 @@ fn write_leaf(writer: &mut ColumnWriter<'_>, levels: &ArrayLevels) -> Result<usi
                         write_primitive(typed, array.values(), levels)
                     }
                 },
+                ArrowDataType::RunEndEncoded(_run_ends, _value_type) => todo!(),
                 _ => {
                     let array = arrow_cast::cast(column, &ArrowDataType::Int32)?;
                     let array = array.as_primitive::<Int32Type>();
@@ -1248,6 +1258,7 @@ fn write_leaf(writer: &mut ColumnWriter<'_>, levels: &ArrayLevels) -> Result<usi
                         write_primitive(typed, array.values(), levels)
                     }
                 },
+                ArrowDataType::RunEndEncoded(_run_ends, _values) => todo!(),
                 _ => {
                     let array = arrow_cast::cast(column, &ArrowDataType::Int64)?;
                     let array = array.as_primitive::<Int64Type>();
@@ -1324,6 +1335,7 @@ 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) => todo!(),
                 _ => {
                     return Err(ParquetError::NYI(
                         "Attempting to write an Arrow type that is not yet implemented".to_string(),
@@ -1494,6 +1506,7 @@ mod tests {
     use arrow::util::pretty::pretty_format_batches;
     use arrow::{array::*, buffer::Buffer};
     use arrow_buffer::{i256, IntervalDayTime, IntervalMonthDayNano, NullBuffer};
+    use arrow_ipc::RunEndEncoded;
     use arrow_schema::Fields;
     use half::f16;
     use num::{FromPrimitive, ToPrimitive};
@@ -4293,4 +4306,35 @@ mod tests {
         assert_eq!(get_dict_page_size(col0_meta), 1024 * 1024);
         assert_eq!(get_dict_page_size(col1_meta), 1024 * 1024 * 4);
     }
+
+    // TODO: Remove. Just added this to compare with arrow_writer_run_end_encoded
+    #[test]
+    fn arrow_writer_string_dictionary_two() {
+        let mut builder = StringDictionaryBuilder::<Int32Type>::new();
+        builder.extend([Some("alpha"), Some("alpha"), Some("beta")]);
+        let dict_array = builder.finish();
+        println!("dict_array type: {:?}", dict_array.data_type());
+        let schema = Arc::new(Schema::new(vec![Field::new(
+            "dict",
+            dict_array.data_type().clone(),
+            dict_array.is_nullable(),
+        )]));
+        one_column_roundtrip_with_schema(Arc::new(dict_array), schema);
+    }
+
+    #[test]
+    fn arrow_writer_run_end_encoded() {
+        // Create a run array of strings
+        let mut builder = StringRunBuilder::<Int32Type>::new();
+        builder.extend([Some("alpha"), Some("alpha"), Some("beta")]);
+        let run_array: RunArray<Int32Type> = builder.finish();
+        println!("run_array type: {:?}", run_array.data_type());
+        let schema = Arc::new(Schema::new(vec![Field::new(
+            "ree",
+            run_array.data_type().clone(),
+            run_array.is_nullable(),
+        )]));
+        // This fails because we don't read it back as RunArray?
+        one_column_roundtrip_with_schema(Arc::new(run_array), schema);
+    }
 }
diff --git a/parquet/src/arrow/schema/mod.rs b/parquet/src/arrow/schema/mod.rs
@@ -129,7 +129,7 @@ pub fn parquet_to_arrow_field_levels(
     match complex::convert_schema(schema, mask, hint)? {
         Some(field) => match &field.arrow_type {
             DataType::Struct(fields) => Ok(FieldLevels {
-                fields: fields.clone(),
+                fields: fields.to_owned(),
                 levels: Some(field),
             }),
             _ => unreachable!(),
@@ -303,7 +303,7 @@ impl<'a> ArrowSchemaConverter<'a> {
     ///
     /// Setting this option to `true` will result in Parquet files that can be
     /// read by more readers, but may lose precision for Arrow types such as
-    /// [`DataType::Date64`] which have no direct [corresponding Parquet type].
+    /// [`DataType::Date64`] which have no direct corresponding Parquet type.
     ///
     /// By default, this converter does not coerce to native Parquet types. Enabling type
     /// coercion allows for meaningful representations that do not require
@@ -771,12 +771,22 @@ fn arrow_to_parquet_type(field: &Field, coerce_types: bool) -> Result<Type> {
         DataType::Union(_, _) => unimplemented!("See ARROW-8817."),
         DataType::Dictionary(_, ref value) => {
             // Dictionary encoding not handled at the schema level
-            let dict_field = field.clone().with_data_type(value.as_ref().clone());
+            let dict_field = field.to_owned().with_data_type(value.as_ref().clone());
+            arrow_to_parquet_type(&dict_field, coerce_types)
+        }
+        DataType::RunEndEncoded(_run_ends, values) => {
+            // REE arrays store run-length encoded data with:
+            // - run_ends: cumulative end indices for each run
+            // - values: the actual values for each run
+            //
+            // For Parquet conversion, we convert REE arrays to Dictionary arrays:
+            // 1. The logical array is materialized as dictionary indices
+            // 2. The values array becomes the dictionary values
+            let dict_field = field
+                .to_owned()
+                .with_data_type(values.data_type().to_owned());
             arrow_to_parquet_type(&dict_field, coerce_types)
         }
-        DataType::RunEndEncoded(_, _) => Err(arrow_err!(
-            "Converting RunEndEncodedType to parquet not supported",
-        )),
     }
 }
 
@@ -2272,4 +2282,22 @@ mod tests {
 
         Ok(())
     }
+
+    #[test]
+    fn test_run_end_encoded_conversion() {
+        use crate::basic::Type;
+        let run_ends_field = Arc::new(Field::new("run_ends", DataType::Int16, false));
+        let values_field = Arc::new(Field::new("values", DataType::Boolean, true));
+        let run_end_encoded_field = Field::new(
+            "run_end_encoded_16",
+            DataType::RunEndEncoded(run_ends_field, values_field),
+            false,
+        );
+
+        let result = arrow_to_parquet_type(&run_end_encoded_field, false).unwrap();
+        // Should convert to the underlying value type (Boolean in this case)
+        assert_eq!(result.get_physical_type(), Type::BOOLEAN);
+        assert_eq!(result.get_basic_info().repetition(), Repetition::REQUIRED); // field is not nullable
+        assert_eq!(result.name(), "run_end_encoded_16");
+    }
 }

Original file line number	Diff line number	Diff line change
`@@ -222,6 +222,10 @@ impl LevelInfoBuilder {`
`222`	`222`	`_ => unreachable!(),`
`223`	`223`	`})`
`224`	`224`	`}`
	`225`	`+ DataType::RunEndEncoded(_, v) if is_leaf(v.data_type()) => {`
	`226`	`+ let levels = ArrayLevels::new(parent_ctx, is_nullable, array.clone());`
	`227`	`+ Ok(Self::Primitive(levels))`
	`228`	`+ }`
`225`	`229`	`d => Err(nyi_err!("Datatype {} is not yet supported", d)),`
`226`	`230`	`}`
`227`	`231`	`}`