[WIP][VARIANT] Support reading basic variant

kernel/src/engine/arrow_conversion.rs

@@ -1,10 +1,11 @@
 //! Conversions from kernel types to arrow types
 
+use std::collections::HashMap;
 use std::sync::Arc;
 
 use arrow_schema::{
-    ArrowError, DataType as ArrowDataType, Field as ArrowField, Schema as ArrowSchema,
-    SchemaRef as ArrowSchemaRef, TimeUnit,
+    ArrowError, DataType as ArrowDataType, Field as ArrowField, Fields as ArrowFields,
+    Schema as ArrowSchema, SchemaRef as ArrowSchemaRef, TimeUnit,
 };
 use itertools::Itertools;
 
@@ -26,13 +27,17 @@ impl TryFrom<&StructField> for ArrowField {
     type Error = ArrowError;
 
     fn try_from(f: &StructField) -> Result<Self, ArrowError> {
-        let metadata = f
+        let mut metadata = f
             .metadata()
             .iter()
             .map(|(key, val)| Ok((key.clone(), serde_json::to_string(val)?)))
-            .collect::<Result<_, serde_json::Error>>()
+            .collect::<Result<HashMap<String, String>, serde_json::Error>>()
             .map_err(|err| ArrowError::JsonError(err.to_string()))?;
 
+        if f.data_type == DataType::VARIANT {
+            metadata.insert("isVariant".to_string(), "true".to_string());
+        }
+
         let field = ArrowField::new(
             f.name(),
             ArrowDataType::try_from(f.data_type())?,
@@ -111,7 +116,16 @@ impl TryFrom<&DataType> for ArrowDataType {
                     )),
                     PrimitiveType::TimestampNtz => {
                         Ok(ArrowDataType::Timestamp(TimeUnit::Microsecond, None))
-                    }
+                    },
+                    // The Delta spec does not enforce the field indexes of the "value" and
+                    // "Metadata" fields. However, within kernel defaults, the "variant_coalesce"
+                    // expression will always return variants with this physical representation.
+                    PrimitiveType::Variant => Ok(ArrowDataType::Struct(ArrowFields::from(
+                        vec![
+                            ArrowField::new("value", ArrowDataType::Binary, false),
+                            ArrowField::new("metadata", ArrowDataType::Binary, false),
+                        ]
+                    )))
                 }
             }
             DataType::Struct(s) => Ok(ArrowDataType::Struct(

kernel/src/engine/arrow_expression.rs

@@ -1,12 +1,13 @@
 //! Expression handling based on arrow-rs compute kernels.
+use std::primitive;
 use std::sync::Arc;
 
 use arrow_arith::boolean::{and_kleene, is_null, not, or_kleene};
 use arrow_arith::numeric::{add, div, mul, sub};
 use arrow_array::cast::AsArray;
 use arrow_array::{
     Array, ArrayRef, BinaryArray, BooleanArray, Date32Array, Datum, Decimal128Array, Float32Array,
-    Float64Array, Int16Array, Int32Array, Int64Array, Int8Array, ListArray, RecordBatch,
+    Float64Array, Int16Array, Int32Array, Int64Array, Int8Array, ListArray, MapArray, RecordBatch,
     StringArray, StructArray, TimestampMicrosecondArray,
 };
 use arrow_ord::cmp::{distinct, eq, gt, gt_eq, lt, lt_eq, neq};
@@ -89,6 +90,10 @@ impl Scalar {
                         Decimal128Array::new_null(num_rows)
                             .with_precision_and_scale(*precision, *scale as i8)?,
                     ),
+                    // TODO(r.chen): Fill this out correctly.
+                    PrimitiveType::Variant => {
+                        panic!("UNSUPPORTED - VARIANT DOESNT HAVE AN ARROW TYPE")
+                    }
                 },
                 DataType::Struct(t) => {
                     let fields: Fields = t.fields().map(ArrowField::try_from).try_collect()?;
@@ -228,13 +233,131 @@ fn ensure_data_types(kernel_type: &DataType, arrow_type: &ArrowDataType) -> Delt
             }
             Ok(())
         }
+        (DataType::Primitive(PrimitiveType::Variant), ArrowDataType::Struct(_)) => Ok(()),
         _ => Err(make_arrow_error(format!(
             "Incorrect datatype. Expected {}, got {}",
             kernel_type, arrow_type
         ))),
     }
 }
 
+fn variant_coalesce_impl(arr: Arc<dyn Array>, kernel_dt: &DataType) -> DeltaResult<ArrayRef> {
+    // TODO(r.chen): Figure out how bad these clone() calls are.
+    match kernel_dt {
+        DataType::Struct(kernel_struct_type) => {
+            if let Some(struct_array) = arr.as_any().downcast_ref::<StructArray>() {
+                let (fields, arrays, nulls) = struct_array.clone().into_parts();
+                let new_arrays = arrays
+                    .into_iter()
+                    .zip(kernel_struct_type.fields().into_iter())
+                    .map(|(child_arr, struct_field)| {
+                        variant_coalesce_impl(child_arr, struct_field.data_type())
+                    })
+                    .collect::<Result<Vec<_>, _>>()?;
+                Ok(Arc::new(StructArray::new(
+                    // TODO(r.chen): Should we be setting this with the new isVariant metadata?
+                    // or will that be handled elsewhere?
+                    fields, new_arrays, nulls,
+                )))
+            } else {
+                Err(Error::unexpected_column_type(kernel_dt))
+            }
+        }
+        DataType::Array(kernel_array_type) => {
+            if let Some(list_array) = arr.as_any().downcast_ref::<ListArray>() {
+                let (field, offsets, values, nulls) = list_array.clone().into_parts();
+                let new_values =
+                    variant_coalesce_impl(values.clone(), kernel_array_type.element_type())?;
+                Ok(Arc::<ListArray>::new(ListArray::new(
+                    // TODO(r.chen): Should we be setting this with the new isVariant metadata?
+                    // or will that be handled elsewhere?
+                    field, offsets, new_values, nulls,
+                )))
+            } else {
+                Err(Error::unexpected_column_type(kernel_dt))
+            }
+        }
+        DataType::Map(kernel_map_type) => {
+            if let Some(map_array) = arr.as_any().downcast_ref::<MapArray>() {
+                let (field, offsets, entries, nulls, ordered) = map_array.clone().into_parts();
+                let new_entries = StructArray::new(
+                    entries.fields().clone(),
+                    vec![
+                        entries.column(0).clone(),
+                        // Map keys cannot be of variant type.
+                        variant_coalesce_impl(
+                            entries.column(1).clone(),
+                            kernel_map_type.value_type(),
+                        )?,
+                    ],
+                    entries.nulls().cloned(),
+                );
+                Ok(Arc::<MapArray>::new(MapArray::new(
+                    // TODO(r.chen): Should we be setting this with the new isVariant metadata?
+                    // or will that be handled elsewhere?
+                    field,
+                    offsets,
+                    new_entries,
+                    nulls,
+                    ordered,
+                )))
+            } else {
+                Err(Error::unexpected_column_type(kernel_dt))
+            }
+        }
+        DataType::Primitive(primitive_type) => match primitive_type {
+            PrimitiveType::Variant => {
+                if let Some(struct_array) = arr.as_any().downcast_ref::<StructArray>() {
+                    let (fields, arrays, nulls) = struct_array.clone().into_parts();
+
+                    let value_idx = if let Some((idx, field)) = fields.find("value") {
+                        if field.data_type() == &ArrowDataType::Binary {
+                            Ok::<usize, Error>(idx)
+                        } else {
+                            Err(Error::invalid_variant_representation(
+                                "\"value\" field is not of binary type.",
+                            ))
+                        }
+                    } else {
+                        Err(Error::invalid_variant_representation(
+                            "\"value\" field is not found.",
+                        ))
+                    }?;
+
+                    let metadata_idx = if let Some((idx, field)) = fields.find("metadata") {
+                        if field.data_type() == &ArrowDataType::Binary {
+                            Ok::<usize, Error>(idx)
+                        } else {
+                            Err(Error::invalid_variant_representation(
+                                "\"metadata\" field is not of binary type.",
+                            ))
+                        }
+                    } else {
+                        Err(Error::invalid_variant_representation(
+                            "\"metadata\" field is not found.",
+                        ))
+                    }?;
+
+                    let new_fields = Fields::from(vec![
+                        fields[value_idx].clone(),
+                        fields[metadata_idx].clone(),
+                    ]);
+                    let new_arrays = vec![
+                        Arc::clone(&arrays[value_idx]),
+                        Arc::clone(&arrays[metadata_idx]),
+                    ];
+                    Ok(Arc::new(StructArray::new(new_fields, new_arrays, nulls)))
+                } else {
+                    Err(Error::invalid_variant_representation(
+                        "variants should be represented in storage as structs.",
+                    ))
+                }
+            }
+            _ => Ok(Arc::new(arr)),
+        },
+    }
+}
+
 fn evaluate_expression(
     expression: &Expression,
     batch: &RecordBatch,
@@ -288,6 +411,12 @@ fn evaluate_expression(
             Ok(match op {
                 UnaryOperator::Not => Arc::new(not(downcast_to_bool(&arr)?)?),
                 UnaryOperator::IsNull => Arc::new(is_null(&arr)?),
+                UnaryOperator::VariantCoalesce => {
+                    let result_type = result_type.ok_or(Error::generic(
+                        "Kernel data type must be supplied to the 'variant_coalesce' expression.",
+                    ))?;
+                    variant_coalesce_impl(arr, result_type)?
+                }
             })
         }
         (BinaryOperation { op, left, right }, _) => {

kernel/src/error.rs

@@ -147,6 +147,9 @@ pub enum Error {
     /// Incosistent data passed to struct scalar
     #[error("Invalid struct data: {0}")]
     InvalidStructData(String),
+
+    #[error ("Invalid variant representation: {0}")]
+    InvalidVariantRepresentation(String),
 }
 
 // Convenience constructors for Error types that take a String argument
@@ -192,6 +195,9 @@ impl Error {
     pub fn invalid_struct_data(msg: impl ToString) -> Self {
         Self::InvalidStructData(msg.to_string())
     }
+    pub fn invalid_variant_representation(msg: impl ToString) -> Self {
+        Self::InvalidVariantRepresentation(msg.to_string())
+    }
 
     // Capture a backtrace when the error is constructed.
     #[must_use]

kernel/src/expressions/mod.rs

@@ -70,6 +70,8 @@ pub enum UnaryOperator {
     Not,
     /// Unary Is Null
     IsNull,
+    /// Unary VariantCoalesce
+    VariantCoalesce,
 }
 
 /// A SQL expression.
@@ -135,6 +137,7 @@ impl Display for Expression {
             Self::UnaryOperation { op, expr } => match op {
                 UnaryOperator::Not => write!(f, "NOT {}", expr),
                 UnaryOperator::IsNull => write!(f, "{} IS NULL", expr),
+                UnaryOperator::VariantCoalesce => write!(f, "VARIANT_COALESCE({})", expr),
             },
             Self::VariadicOperation { op, exprs } => match op {
                 VariadicOperator::And => {
@@ -227,6 +230,11 @@ impl Expression {
         Self::unary(UnaryOperator::IsNull, self)
     }
 
+    /// Create a new expression `variant_coalesce(self)`
+    pub fn variant_coalesce(self) -> Self {
+        Self::unary(UnaryOperator::VariantCoalesce, self)
+    }
+
     /// Create a new expression `self == other`
     pub fn eq(self, other: Self) -> Self {
         Self::binary(BinaryOperator::Equal, self, other)

kernel/src/expressions/scalars.rs

@@ -1,3 +1,4 @@
+use core::panic;
 use std::cmp::Ordering;
 use std::fmt::{Display, Formatter};
 
@@ -314,6 +315,8 @@ impl PrimitiveType {
                     _ => unreachable!(),
                 }
             }
+            // TODO(r.chen): handle this
+            Variant => panic!("UNSUPPORTED")
         }
     }
 

kernel/src/scan/mod.rs

@@ -1,5 +1,6 @@
 //! Functionality to create and execute scans (reads) over data stored in a delta table
 
+use std::ptr::null;
 use std::sync::Arc;
 
 use itertools::Itertools;
@@ -269,6 +270,45 @@ impl Scan {
             let read_results =
                 parquet_handler.read_parquet_files(&[meta], self.physical_schema.clone(), None)?;
 
+            // The rest of the kernel expects variants to be reconstructed, so they must be
+            // reconstructed immediately after scanning the parquet files.
+            let variant_coalesce_expr = match &self.snapshot.protocol().reader_features {
+                Some(reader_features) => {
+                    // TODO(r.chen): Only check this if the variantType feature is enabled.
+                    // if reader_features.contains(&"variantType-preview".to_string())
+                    let all_fields_expr = 
+                        self.physical_schema
+                            .fields
+                            .iter()
+                            .map(|(_, field)| {
+                                Expression::variant_coalesce(Expression::column(
+                                    field.name.clone(),
+                                ))
+                            })
+                            .collect::<Vec<_>>();
+                    Some(Expression::Struct(all_fields_expr))
+                }
+                None => None,
+            };
+
+            let mut variant_coalesce_results: Vec<DeltaResult<Box<dyn EngineData>>> = vec![];
+            match variant_coalesce_expr {
+                Some(expr) => {
+                    for result in read_results {
+                        let coalesced_batch = engine
+                            .get_expression_handler()
+                            .get_evaluator(
+                                self.physical_schema.clone(),
+                                expr.clone(),
+                                output_schema.clone(),
+                            )
+                            .evaluate(result?.as_ref());
+                        variant_coalesce_results.push(coalesced_batch);
+                    }
+                }
+                None => (),
+            };
+
             let read_expression = if self.have_partition_cols
                 || self.snapshot.column_mapping_mode != ColumnMappingMode::None
             {
@@ -307,7 +347,7 @@ impl Scan {
 
             let mut dv_mask = dv_treemap.map(treemap_to_bools);
 
-            for read_result in read_results {
+            for read_result in variant_coalesce_results {
                 let len = if let Ok(ref res) = read_result {
                     res.length()
                 } else {

kernel/src/schema.rs

@@ -358,6 +358,7 @@ pub enum PrimitiveType {
     Timestamp,
     #[serde(rename = "timestamp_ntz")]
     TimestampNtz,
+    Variant,
     #[serde(
         serialize_with = "serialize_decimal",
         deserialize_with = "deserialize_decimal",
@@ -418,7 +419,8 @@ impl Display for PrimitiveType {
             PrimitiveType::TimestampNtz => write!(f, "timestamp_ntz"),
             PrimitiveType::Decimal(precision, scale) => {
                 write!(f, "decimal({},{})", precision, scale)
-            }
+            },
+            PrimitiveType::Variant => write!(f, "variant")
         }
     }
 }
@@ -476,6 +478,7 @@ impl DataType {
     pub const DATE: Self = DataType::Primitive(PrimitiveType::Date);
     pub const TIMESTAMP: Self = DataType::Primitive(PrimitiveType::Timestamp);
     pub const TIMESTAMP_NTZ: Self = DataType::Primitive(PrimitiveType::TimestampNtz);
+    pub const VARIANT: Self = DataType::Primitive(PrimitiveType::Variant);
 
     pub fn decimal(precision: u8, scale: u8) -> DeltaResult<Self> {
         PrimitiveType::check_decimal(precision, scale)?;