Gross hack during type conversion

Author: alamb

parquet/src/arrow/schema/complex.rs

@@ -25,7 +25,7 @@ use crate::errors::ParquetError;
 use crate::errors::Result;
 use crate::schema::types::{SchemaDescriptor, Type, TypePtr};
 use arrow_schema::{DataType, Field, Fields, SchemaBuilder};
-use crate::arrow::schema::extension::add_extension_type;
+use crate::arrow::schema::extension::{add_extension_type, binary_to_binary_view};
 
 fn get_repetition(t: &Type) -> Repetition {
     let info = t.get_basic_info();
@@ -173,7 +173,7 @@ impl Visitor {
 
         let parquet_fields = struct_type.get_fields();
 
-        // Extract the arrow fields
+        // Extract any arrow fields from the hints
         let arrow_fields = match &context.data_type {
             Some(DataType::Struct(fields)) => {
                 if fields.len() != parquet_fields.len() {
@@ -221,10 +221,10 @@ impl Visitor {
                 data_type,
             };
 
-            if let Some(child) = self.dispatch(parquet_field, child_ctx)? {
+            if let Some(mut child) = self.dispatch(parquet_field, child_ctx)? {
                 // The child type returned may be different from what is encoded in the arrow
                 // schema in the event of a mismatch or a projection
-                child_fields.push(convert_field(parquet_field, &child, arrow_field));
+                child_fields.push(convert_field(parquet_field, &mut child, arrow_field));
                 children.push(child);
             }
         }
@@ -353,13 +353,13 @@ impl Visitor {
 
         // Need both columns to be projected
         match (maybe_key, maybe_value) {
-            (Some(key), Some(value)) => {
+            (Some(mut key), Some(mut value)) => {
                 let key_field = Arc::new(
-                    convert_field(map_key, &key, arrow_key)
+                    convert_field(map_key, &mut key, arrow_key)
                         // The key is always non-nullable (#5630)
                         .with_nullable(false),
                 );
-                let value_field = Arc::new(convert_field(map_value, &value, arrow_value));
+                let value_field = Arc::new(convert_field(map_value, &mut value, arrow_value));
                 let field_metadata = match arrow_map {
                     Some(field) => field.metadata().clone(),
                     _ => HashMap::default(),
@@ -496,8 +496,8 @@ impl Visitor {
         };
 
         match self.dispatch(item_type, new_context) {
-            Ok(Some(item)) => {
-                let item_field = Arc::new(convert_field(item_type, &item, arrow_field));
+            Ok(Some(mut item)) => {
+                let item_field = Arc::new(convert_field(item_type, &mut item, arrow_field));
 
                 // Use arrow type as hint for index size
                 let arrow_type = match context.data_type {
@@ -545,7 +545,7 @@ impl Visitor {
 ///
 /// The resulting Arrow [`Field`] will have the type dictated by the Parquet `field`, a name
 /// dictated by the `parquet_type`, and any metadata from `arrow_hint`
-fn convert_field(parquet_type: &Type, field: &ParquetField, arrow_hint: Option<&Field>) -> Field {
+fn convert_field(parquet_type: &Type, field: &mut ParquetField, arrow_hint: Option<&Field>) -> Field {
     let name = parquet_type.name();
     let data_type = field.arrow_type.clone();
     let nullable = field.nullable;
@@ -576,7 +576,9 @@ fn convert_field(parquet_type: &Type, field: &ParquetField, arrow_hint: Option<&
                 );
                 ret.set_metadata(meta);
             }
-            add_extension_type(ret, &parquet_type)
+            let ret = add_extension_type(ret, &parquet_type);
+            field.arrow_type = binary_to_binary_view(&field.arrow_type);
+            ret
         }
     }
 }

parquet/src/arrow/schema/extension.rs

@@ -17,7 +17,8 @@
 
 //! Arrow Extension Type Support
 
-use arrow_schema::Field;
+use std::sync::Arc;
+use arrow_schema::{DataType, Field, Fields};
 use crate::basic::LogicalType;
 use crate::schema::types::Type;
 
@@ -28,11 +29,50 @@ use crate::schema::types::Type;
 /// Arrow DataType, and instead are represented by an Arrow ExtensionType.
 /// Extension types are attached to Arrow Fields via metadata.
 pub(crate) fn add_extension_type(arrow_field: Field, parquet_type: &Type) -> Field {
-    match parquet_type.get_basic_info().logical_type() {
+    println!("Adding extension to field: {:#?}", arrow_field);
+    let result = match parquet_type.get_basic_info().logical_type() {
         #[cfg(feature = "variant_experimental")]
         Some(LogicalType::Variant) => {
-            arrow_field.with_extension_type(crate::variant::VariantType)
+            let new_data_type = binary_to_binary_view(arrow_field.data_type());
+            println!("New variant data type: {:?}", new_data_type);
+            arrow_field
+                .with_data_type(new_data_type)
+                .with_extension_type(crate::variant::VariantType)
         }
         _ => arrow_field
+    };
+    println!("new field: {:#?}", result);
+    result
+}
+
+/// replaces all instances of Binary with BinaryView in a DataType
+///
+/// The `VariantArray` implementation uses BinaryView rather than Binary to
+/// avoid unnecessary copies. This function recursively traverses the DataType
+/// and replaces Binary with BinaryView, including within nested List and Struct
+/// types.
+#[cfg(feature = "variant_experimental")]
+pub(crate) fn binary_to_binary_view(data_type: &DataType) -> DataType {
+    // TODO avoid the clones below
+    match data_type {
+        DataType::Binary => DataType::BinaryView,
+        DataType::List(field) => {
+            let new_field = field
+                .as_ref()
+                .clone()
+                .with_data_type(binary_to_binary_view(field.data_type()));
+            DataType::List(Arc::new(new_field))
+        }
+        DataType::Struct(fields) => {
+            let new_fields: Fields = fields
+                .iter()
+                .map(|f| {
+                    let new_field = f.as_ref().clone().with_data_type(binary_to_binary_view(f.data_type()));
+                    Arc::new(new_field)
+                })
+                .collect();
+            DataType::Struct(new_fields)
+        }
+        _ => data_type.clone(),
     }
 }

parquet/src/variant.rs

@@ -123,6 +123,7 @@ mod tests {
     use crate::arrow::ArrowWriter;
     use crate::file::reader::ChunkReader;
     use arrow::util::test_util::parquet_test_data;
+    use arrow_array::RecordBatchReader;
     use arrow_array::{Array, ArrayRef, RecordBatch};
     use arrow_schema::{Field, Fields, Schema};
     use bytes::Bytes;
@@ -165,7 +166,7 @@ mod tests {
         assert_eq!(var_array.len(), 1);
         assert!(var_array.is_valid(0));
         let var_value = var_array.value(0);
-        assert_eq!(var_value, Variant::String("iceberg - UPDATE"));
+        assert_eq!(var_value, Variant::from("iceberg"));
     }
 
     /// Writes a variant to a parquet file and ensures the parquet logical type
@@ -210,6 +211,7 @@ mod tests {
             .unwrap()
             .build()
             .unwrap();
+        println!("Reader shcema: {:?}", reader.schema());
         let mut batches: Vec<RecordBatch> = reader.collect::<Result<Vec<_>, _>>().unwrap();
         assert_eq!(batches.len(), 1);
         batches.swap_remove(0)