apply a schema to fix column names

kernel/examples/read-table-single-threaded/src/main.rs

@@ -41,6 +41,10 @@ struct Cli {
     /// to the aws metadata server, which will fail unless you're on an ec2 instance.
     #[arg(long)]
     public: bool,
+
+    /// Only print the schema of the table
+    #[arg(long)]
+    schema_only: bool,
 }
 
 #[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, ValueEnum)]
@@ -90,6 +94,11 @@ fn try_main() -> DeltaResult<()> {
 
     let snapshot = table.snapshot(engine.as_ref(), None)?;
 
+    if cli.schema_only {
+        println!("{:#?}", snapshot.schema());
+        return Ok(());
+    }
+
     let read_schema_opt = cli
         .columns
         .map(|cols| -> DeltaResult<_> {

kernel/src/engine/arrow_expression.rs

@@ -1,10 +1,11 @@
 //! Expression handling based on arrow-rs compute kernels.
+use std::collections::HashMap;
 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::types::*;
+use arrow_array::{types::*, MapArray};
 use arrow_array::{
     Array, ArrayRef, BinaryArray, BooleanArray, Date32Array, Datum, Decimal128Array, Float32Array,
     Float64Array, Int16Array, Int32Array, Int64Array, Int8Array, ListArray, RecordBatch,
@@ -21,20 +22,21 @@ use arrow_select::concat::concat;
 use itertools::Itertools;
 
 use super::arrow_conversion::LIST_ARRAY_ROOT;
+use super::arrow_utils::make_arrow_error;
 use crate::engine::arrow_data::ArrowEngineData;
 use crate::engine::arrow_utils::ensure_data_types;
 use crate::engine::arrow_utils::prim_array_cmp;
 use crate::error::{DeltaResult, Error};
 use crate::expressions::{BinaryOperator, Expression, Scalar, UnaryOperator, VariadicOperator};
-use crate::schema::{DataType, PrimitiveType, SchemaRef};
+use crate::schema::{ArrayType, DataType, MapType, PrimitiveType, Schema, SchemaRef, StructField};
 use crate::{EngineData, ExpressionEvaluator, ExpressionHandler};
 
 // TODO leverage scalars / Datum
 
 fn downcast_to_bool(arr: &dyn Array) -> DeltaResult<&BooleanArray> {
     arr.as_any()
         .downcast_ref::<BooleanArray>()
-        .ok_or(Error::generic("expected boolean array"))
+        .ok_or_else(|| Error::generic("expected boolean array"))
 }
 
 impl Scalar {
@@ -154,10 +156,7 @@ fn extract_column<'array, 'path>(
 ) -> Result<&'array Arc<dyn Array>, ArrowError> {
     let child = array
         .column_by_name(path_step)
-        .ok_or(ArrowError::SchemaError(format!(
-            "No such field: {}",
-            path_step,
-        )))?;
+        .ok_or_else(|| ArrowError::SchemaError(format!("No such field: {}", path_step,)))?;
     if let Some(next_path_step) = remaining_path_steps.next() {
         // This is not the last path step. Drill deeper.
         extract_column(
@@ -176,10 +175,10 @@ fn column_as_struct<'a>(
     column: &Option<&'a Arc<dyn Array>>,
 ) -> Result<&'a StructArray, ArrowError> {
     column
-        .ok_or(ArrowError::SchemaError(format!("No such column: {}", name)))?
+        .ok_or_else(|| ArrowError::SchemaError(format!("No such column: {}", name)))?
         .as_any()
         .downcast_ref::<StructArray>()
-        .ok_or(ArrowError::SchemaError(format!("{} is not a struct", name)))
+        .ok_or_else(|| ArrowError::SchemaError(format!("{} is not a struct", name)))
 }
 
 fn evaluate_expression(
@@ -201,7 +200,7 @@ fn evaluate_expression(
             } else {
                 batch
                     .column_by_name(name)
-                    .ok_or(Error::missing_column(name))
+                    .ok_or_else(|| Error::missing_column(name))
                     .cloned()
             }
         }
@@ -215,7 +214,6 @@ fn evaluate_expression(
                 .iter()
                 .zip(output_schema.fields())
                 .map(|(output_col, output_field)| -> DeltaResult<_> {
-                    ensure_data_types(output_field.data_type(), output_col.data_type())?;
                     Ok(ArrowField::new(
                         output_field.name(),
                         output_col.data_type().clone(),
@@ -366,6 +364,193 @@ fn evaluate_expression(
     }
 }
 
+// Apply a schema to an array. The array _must_ be a `StructArray`. Returns a `RecordBatch where the
+// names of fields, nullable, and metadata in the struct have been transformed to match those in
+// schema specified by `schema`
+fn apply_schema(array: &dyn Array, schema: &DataType) -> DeltaResult<RecordBatch> {
+    let DataType::Struct(struct_schema) = schema else {
+        return Err(Error::generic(
+            "apply_schema at top-level must be passed a struct schema",
+        ))?;
+    };
+    let applied = apply_schema_to_struct(array, struct_schema)?;
+    Ok(applied.into())
+}
+
+// helper to transform an arrow field+col into the specified target type. If `rename` is specified
+// the field will be renamed to the contained `str`.
+fn transform_field_and_col(
+    arrow_field: &Arc<ArrowField>,
+    arrow_col: Arc<dyn Array>,
+    target_type: &DataType,
+    nullable: bool,
+    rename: Option<&str>,
+    metadata: Option<HashMap<String, String>>,
+) -> DeltaResult<(ArrowField, Arc<dyn Array>)> {
+    let transformed_col = apply_schema_to(&arrow_col, target_type)?.unwrap_or(arrow_col);
+    let transformed_field = arrow_field
+        .as_ref()
+        .clone()
+        .with_nullable(nullable)
+        .with_data_type(transformed_col.data_type().clone());
+    let transformed_field = match rename {
+        Some(name) => transformed_field.with_name(name),
+        None => transformed_field,
+    };
+    let transformed_field = match metadata {
+        Some(metadata) => transformed_field.with_metadata(metadata),
+        None => transformed_field,
+    };
+    Ok((transformed_field, transformed_col))
+}
+
+// A helper that is a wrapper over `transform_field_and_col`. This will take apart the passed struct
+// and use that method to transform each column and then put the struct back together. Target types
+// and names for each column should be passed in `target_types_and_names`. The number of elements in
+// the `target_types_and_names` iterator _must_ be the same as the number of columns in
+// `struct_array`. The transformation is ordinal. That is, the order of fields in `target_fields`
+// _must_ match the order of the columns in `struct_array`.
+fn transform_struct<'a>(
+    struct_array: &StructArray,
+    target_fields: impl Iterator<Item = &'a StructField>,
+) -> DeltaResult<StructArray> {
+    let (arrow_fields, arrow_cols, nulls) = struct_array.clone().into_parts();
+    let input_col_count = arrow_cols.len();
+    let result_iter = arrow_fields
+        .into_iter()
+        .zip(arrow_cols)
+        .zip(target_fields)
+        .map(|((sa_field, sa_col), target_field)| {
+            transform_field_and_col(
+                sa_field,
+                sa_col,
+                target_field.data_type(),
+                target_field.nullable,
+                Some(target_field.name.as_str()),
+                Some(target_field.metadata_as_string()),
+            )
+        });
+    let (transformed_fields, transformed_cols): (Vec<ArrowField>, Vec<Arc<dyn Array>>) =
+        result_iter.process_results(|iter| iter.unzip())?;
+    if transformed_cols.len() != input_col_count {
+        return Err(Error::InternalError(format!(
+            "Passed struct had {input_col_count} columns, but transformed column has {}",
+            transformed_cols.len()
+        )));
+    }
+    Ok(StructArray::try_new(
+        transformed_fields.into(),
+        transformed_cols,
+        nulls,
+    )?)
+}
+
+// Transform a struct array. The data is in `sa`, the current fields are in `arrow_fields`, and the
+// target fields are in `kernel_fields`.
+fn apply_schema_to_struct(array: &dyn Array, kernel_fields: &Schema) -> DeltaResult<StructArray> {
+    let Some(sa) = array.as_struct_opt() else {
+        return Err(make_arrow_error(
+            "Arrow claimed to be a struct but isn't a StructArray",
+        ));
+    };
+    transform_struct(sa, kernel_fields.fields())
+}
+
+// deconstruct the array, then rebuild the mapped version
+fn apply_schema_to_list(
+    array: &dyn Array,
+    target_inner_type: &ArrayType,
+) -> DeltaResult<ListArray> {
+    let Some(la) = array.as_list_opt() else {
+        return Err(make_arrow_error(
+            "Arrow claimed to be a list but isn't a ListArray",
+        ));
+    };
+    let (field, offset_buffer, values, nulls) = la.clone().into_parts();
+    let (transformed_field, transformed_values) = transform_field_and_col(
+        &field,
+        values,
+        &target_inner_type.element_type,
+        target_inner_type.contains_null,
+        None,
+        None,
+    )?;
+    Ok(ListArray::try_new(
+        Arc::new(transformed_field),
+        offset_buffer,
+        transformed_values,
+        nulls,
+    )?)
+}
+
+// deconstruct a map, and rebuild it with the specified target kernel type
+fn apply_schema_to_map(array: &dyn Array, kernel_map_type: &MapType) -> DeltaResult<MapArray> {
+    let Some(ma) = array.as_map_opt() else {
+        return Err(make_arrow_error(
+            "Arrow claimed to be a map but isn't a MapArray",
+        ));
+    };
+
+    let ArrowDataType::Map(arrow_map_type, _) = array.data_type() else {
+        return Err(make_arrow_error(
+            "Arrow claimed to be a map but doesn't have map DataType",
+        ));
+    };
+
+    let (map_field, offset_buffer, map_struct_array, nulls, ordered) = ma.clone().into_parts();
+
+    let ArrowDataType::Struct(_) = arrow_map_type.data_type() else {
+        return Err(make_arrow_error("Arrow map type wasn't a struct."));
+    };
+
+    let target_fields: Vec<StructField> = map_struct_array
+        .fields()
+        .iter()
+        .zip([
+            (&kernel_map_type.key_type, false),
+            (
+                &kernel_map_type.value_type,
+                kernel_map_type.value_contains_null,
+            ),
+        ])
+        .map(|(arrow_field, (target_type, nullable))| {
+            StructField::new(arrow_field.name(), target_type.clone(), nullable)
+        })
+        .collect();
+
+    // Arrow puts the key type/val as the first field/col and the value type/val as the second. So
+    // we just transform like a 'normal' struct, but we know there are two fields/cols and we
+    // specify the key/value types as the target type iterator.
+    let transformed_map_struct_array = transform_struct(&map_struct_array, target_fields.iter())?;
+
+    let transformed_map_field = Arc::new(
+        map_field
+            .as_ref()
+            .clone()
+            .with_data_type(transformed_map_struct_array.data_type().clone()),
+    );
+    Ok(MapArray::try_new(
+        transformed_map_field,
+        offset_buffer,
+        transformed_map_struct_array,
+        nulls,
+        ordered,
+    )?)
+}
+
+// make column `col` with type `arrow_type` look like `kernel_type`. For now this only handles name
+// transforms. if the actual data types don't match, this will return an error
+fn apply_schema_to(array: &dyn Array, schema: &DataType) -> DeltaResult<Option<ArrayRef>> {
+    use DataType::*;
+    let array: ArrayRef = match schema {
+        Struct(stype) => Arc::new(apply_schema_to_struct(array, stype)?),
+        Array(atype) => Arc::new(apply_schema_to_list(array, atype)?),
+        Map(mtype) => Arc::new(apply_schema_to_map(array, mtype)?),
+        _ => return ensure_data_types(schema, array.data_type(), true).map(|_| None),
+    };
+    Ok(Some(array))
+}
+
 #[derive(Debug)]
 pub struct ArrowExpressionHandler;
 
@@ -396,7 +581,7 @@ impl ExpressionEvaluator for DefaultExpressionEvaluator {
         let batch = batch
             .as_any()
             .downcast_ref::<ArrowEngineData>()
-            .ok_or(Error::engine_data_type("ArrowEngineData"))?
+            .ok_or_else(|| Error::engine_data_type("ArrowEngineData"))?
             .record_batch();
         let _input_schema: ArrowSchema = self.input_schema.as_ref().try_into()?;
         // TODO: make sure we have matching schemas for validation
@@ -408,13 +593,14 @@ impl ExpressionEvaluator for DefaultExpressionEvaluator {
         //     )));
         // };
         let array_ref = evaluate_expression(&self.expression, batch, Some(&self.output_type))?;
-        let arrow_type: ArrowDataType = ArrowDataType::try_from(&self.output_type)?;
         let batch: RecordBatch = if let DataType::Struct(_) = self.output_type {
-            array_ref
-                .as_struct_opt()
-                .ok_or(Error::unexpected_column_type("Expected a struct array"))?
-                .into()
+            apply_schema(&array_ref, &self.output_type)?
         } else {
+            let array_ref = match apply_schema_to(&array_ref, &self.output_type)? {
+                Some(transformed) => transformed,
+                None => array_ref, // Were a primitive type, we just validated
+            };
+            let arrow_type: ArrowDataType = ArrowDataType::try_from(&self.output_type)?;
             let schema = ArrowSchema::new(vec![ArrowField::new("output", arrow_type, true)]);
             RecordBatch::try_new(Arc::new(schema), vec![array_ref])?
         };