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| {

kernel/src/engine/arrow_expression.rs

@@ -4,7 +4,7 @@ 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,6 +21,7 @@ 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;
@@ -214,10 +215,9 @@ fn evaluate_expression(
             let output_fields: Vec<ArrowField> = output_cols
                 .iter()
                 .zip(schema.fields())
-                .map(|(array, input_field)| -> DeltaResult<_> {
-                    ensure_data_types(input_field.data_type(), array.data_type())?;
+                .map(|(array, output_field)| -> DeltaResult<_> {
                     Ok(ArrowField::new(
-                        input_field.name(),
+                        output_field.name(),
                         array.data_type().clone(),
                         array.is_nullable(),
                     ))
@@ -366,6 +366,153 @@ fn evaluate_expression(
     }
 }
 
+// return a RecordBatch where the names of fields in `sa` have been transformed to match those in
+// schema specified by `output_type`
+fn apply_schema(sa: &StructArray, output_type: &DataType) -> DeltaResult<RecordBatch> {
+    let applied = apply_to_col(sa.data_type(), sa, output_type)?.ok_or(Error::generic(
+        "apply_to_col at top-level should return something",
+    ))?;
+    let applied_sa = applied.as_struct_opt().ok_or(Error::generic(
+        "apply_to_col at top-level should return a struct array",
+    ))?;
+    Ok(applied_sa.into())
+}
+
+// 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_to_col(
+    arrow_type: &ArrowDataType,
+    col: &dyn Array,
+    kernel_type: &DataType,
+) -> DeltaResult<Option<Arc<dyn Array>>> {
+    match (kernel_type, arrow_type) {
+        (DataType::Struct(kernel_fields), ArrowDataType::Struct(arrow_fields)) => {
+            if kernel_fields.fields.len() != arrow_fields.len() {
+                return Err(make_arrow_error(format!(
+                    "Kernel schema had {} fields, but data has {}",
+                    kernel_fields.fields.len(),
+                    arrow_fields.len()
+                )));
+            }
+            let sa = col.as_struct_opt().ok_or(make_arrow_error(
+                "Arrow claimed to be a struct but isn't a StructArray".to_string(),
+            ))?;
+            let (fields, sa_cols, sa_nulls) = sa.clone().into_parts();
+            let result_iter = fields
+                .into_iter()
+                .zip(sa_cols)
+                .zip(kernel_fields.fields())
+                .map(
+                |((sa_field, sa_col), kernel_field)| -> DeltaResult<(ArrowField, Arc<dyn Array>)> {
+                    let transformed_col =
+                        apply_to_col(sa_field.data_type(), &sa_col, kernel_field.data_type())?
+                            .unwrap_or(sa_col);
+                    let transformed_field = sa_field
+                        .as_ref()
+                        .clone()
+                        .with_name(kernel_field.name.clone())
+                        .with_data_type(transformed_col.data_type().clone());
+                    Ok((transformed_field, transformed_col))
+                },
+            );
+            let (transformed_fields, transformed_cols): (Vec<ArrowField>, Vec<Arc<dyn Array>>) =
+                result_iter.process_results(|iter| iter.unzip())?;
+            let transformed_array =
+                StructArray::try_new(transformed_fields.into(), transformed_cols, sa_nulls)?;
+            Ok(Some(Arc::new(transformed_array)))
+        }
+        (DataType::Array(inner_type), ArrowDataType::List(_arrow_list_type)) => {
+            // deconstruct the array, the rebuild the mapped version
+            let la = col.as_list_opt().ok_or(make_arrow_error(
+                "Arrow claimed to be a list but isn't a ListArray".to_string(),
+            ))?;
+            let (field, offset_buffer, values, nulls) = la.clone().into_parts();
+            let transformed_values =
+                apply_to_col(field.data_type(), &values, &inner_type.element_type)?
+                    .unwrap_or(values);
+            let transformed_field = Arc::new(
+                field
+                    .as_ref()
+                    .clone()
+                    .with_data_type(transformed_values.data_type().clone()),
+            );
+            let transformed_array =
+                ListArray::try_new(transformed_field, offset_buffer, transformed_values, nulls)?;
+            Ok(Some(Arc::new(transformed_array)))
+        }
+        (DataType::Map(kernel_map_type), ArrowDataType::Map(arrow_map_type, _)) => {
+            let ma = col.as_map_opt().ok_or(make_arrow_error(
+                "Arrow claimed to be a map but isn't a MapArray".to_string(),
+            ))?;
+            let (map_field, offset_buffer, map_struct_array, nulls, ordered) =
+                ma.clone().into_parts();
+            if let ArrowDataType::Struct(_) = arrow_map_type.data_type() {
+                let (fields, msa_cols, msa_nulls) = map_struct_array.clone().into_parts();
+                let mut fields = fields.into_iter();
+                let key_field = fields.next().ok_or(make_arrow_error(
+                    "Arrow map struct didn't have a key field".to_string(),
+                ))?;
+                let value_field = fields.next().ok_or(make_arrow_error(
+                    "Arrow map struct didn't have a value field".to_string(),
+                ))?;
+                if fields.next().is_some() {
+                    return Err(Error::generic("map fields had more than 2 members"));
+                }
+                let transformed_key = apply_to_col(
+                    key_field.data_type(),
+                    msa_cols[0].as_ref(),
+                    &kernel_map_type.key_type,
+                )?
+                .unwrap_or(msa_cols[0].clone());
+                let transformed_values = apply_to_col(
+                    value_field.data_type(),
+                    msa_cols[1].as_ref(),
+                    &kernel_map_type.value_type,
+                )?
+                .unwrap_or(msa_cols[1].clone());
+                let transformed_struct_fields = vec![
+                    key_field
+                        .as_ref()
+                        .clone()
+                        .with_data_type(transformed_key.data_type().clone()),
+                    value_field
+                        .as_ref()
+                        .clone()
+                        .with_data_type(transformed_values.data_type().clone()),
+                ];
+                let transformed_struct_cols = vec![transformed_key, transformed_values];
+                let transformed_map_struct_array = StructArray::try_new(
+                    transformed_struct_fields.into(),
+                    transformed_struct_cols,
+                    msa_nulls,
+                )?;
+                let transformed_map_field = Arc::new(
+                    map_field
+                        .as_ref()
+                        .clone()
+                        .with_data_type(transformed_map_struct_array.data_type().clone()),
+                );
+                let transformed_map = MapArray::try_new(
+                    transformed_map_field,
+                    offset_buffer,
+                    transformed_map_struct_array,
+                    nulls,
+                    ordered,
+                )?;
+                Ok(Some(Arc::new(transformed_map)))
+            } else {
+                Err(make_arrow_error(
+                    "Arrow map type wasn't a struct.".to_string(),
+                ))
+            }
+        }
+        _ => {
+            ensure_data_types(kernel_type, arrow_type)?;
+            Ok(None)
+        }
+    }
+}
+
 #[derive(Debug)]
 pub struct ArrowExpressionHandler;
 
@@ -410,10 +557,13 @@ 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
+            let sa: &StructArray = array_ref
                 .as_struct_opt()
-                .ok_or(Error::unexpected_column_type("Expected a struct array"))?
-                .into()
+                .ok_or(Error::unexpected_column_type("Expected a struct array"))?;
+            match ensure_data_types(&self.output_type, sa.data_type()) {
+                Ok(_) => sa.into(),
+                Err(_) => apply_schema(sa, &self.output_type)?,
+            }
         } else {
             let schema = ArrowSchema::new(vec![ArrowField::new("output", arrow_type, true)]);
             RecordBatch::try_new(Arc::new(schema), vec![array_ref])?

kernel/src/engine/arrow_utils.rs

@@ -58,8 +58,8 @@ pub(crate) use prim_array_cmp;
 /// returns a tuples of (mask_indicies: Vec<parquet_schema_index>, reorder_indicies:
 /// Vec<requested_index>). `mask_indicies` is used for generating the mask for reading from the
 
-fn make_arrow_error(s: String) -> Error {
-    Error::Arrow(arrow_schema::ArrowError::InvalidArgumentError(s))
+pub(crate) fn make_arrow_error(s: String) -> Error {
+    Error::Arrow(arrow_schema::ArrowError::InvalidArgumentError(s)).with_backtrace()
 }
 
 /// Capture the compatibility between two data-types, as passed to [`ensure_data_types`]
@@ -636,6 +636,7 @@ pub(crate) fn reorder_struct_array(
     input_data: StructArray,
     requested_ordering: &[ReorderIndex],
 ) -> DeltaResult<StructArray> {
+    debug!("Reordering {input_data:?} with ordering: {requested_ordering:?}");
     if !ordering_needs_transform(requested_ordering) {
         // indices is already sorted, meaning we requested in the order that the columns were
         // stored in the parquet
@@ -944,7 +945,6 @@ mod tests {
             ArrowField::new("s", ArrowDataType::Int32, true),
         ]));
         let res = get_requested_indices(&requested_schema, &parquet_schema);
-        println!("{res:#?}");
         assert!(res.is_err());
     }
 

kernel/src/error.rs

@@ -175,7 +175,7 @@ impl Error {
         Self::FileNotFound(path.to_string())
     }
     pub fn missing_column(name: impl ToString) -> Self {
-        Self::MissingColumn(name.to_string())
+        Self::MissingColumn(name.to_string()).with_backtrace()
     }
     pub fn unexpected_column_type(name: impl ToString) -> Self {
         Self::UnexpectedColumnType(name.to_string())

kernel/src/scan/mod.rs

@@ -12,7 +12,7 @@ use crate::actions::{get_log_schema, ADD_NAME, REMOVE_NAME};
 use crate::expressions::{Expression, Scalar};
 use crate::features::ColumnMappingMode;
 use crate::scan::state::{DvInfo, Stats};
-use crate::schema::{DataType, Schema, SchemaRef, StructField, StructType};
+use crate::schema::{ArrayType, DataType, MapType, Schema, SchemaRef, StructField, StructType};
 use crate::snapshot::Snapshot;
 use crate::{DeltaResult, Engine, EngineData, Error, FileMeta};
 
@@ -384,6 +384,78 @@ fn parse_partition_value(raw: Option<&String>, data_type: &DataType) -> DeltaRes
     }
 }
 
+/// Transform a logical field into the physical form. Currently just renames things for 'name'
+/// column mapping.
+fn make_field_physical(
+    logical_field: &StructField,
+    column_mapping_mode: ColumnMappingMode,
+) -> DeltaResult<StructField> {
+    match column_mapping_mode {
+        ColumnMappingMode::None => Ok(logical_field.clone()),
+        ColumnMappingMode::Name => {
+            let physical_name = logical_field.physical_name(column_mapping_mode)?;
+            let field_data_type = logical_field.data_type();
+            let mapped_data_type = make_data_type_physical(field_data_type, column_mapping_mode)?;
+            Ok(StructField {
+                name: physical_name.to_string(),
+                data_type: mapped_data_type,
+                nullable: logical_field.nullable,
+                metadata: logical_field.metadata.clone(),
+            })
+        }
+        ColumnMappingMode::Id => panic!("No id"),
+    }
+}
+
+/// Transform a DataType into the physical form. Currently just renames anything in a nested type
+/// for 'name' column mapping.
+fn make_data_type_physical(
+    logical_dt: &DataType,
+    column_mapping_mode: ColumnMappingMode,
+) -> DeltaResult<DataType> {
+    match column_mapping_mode {
+        ColumnMappingMode::None => Ok(logical_dt.clone()),
+        ColumnMappingMode::Name => {
+            // we don't need to rename at this level, just need to keep the recursion going
+            // because there might be structs below us
+            match logical_dt {
+                DataType::Array(array_type) => {
+                    let new_type =
+                        make_data_type_physical(&array_type.element_type, column_mapping_mode)?;
+                    Ok(DataType::Array(Box::new(ArrayType::new(
+                        new_type,
+                        array_type.contains_null,
+                    ))))
+                }
+                DataType::Map(map_type) => {
+                    let new_key_type =
+                        make_data_type_physical(&map_type.key_type, column_mapping_mode)?;
+                    let new_value_type =
+                        make_data_type_physical(&map_type.value_type, column_mapping_mode)?;
+                    Ok(DataType::Map(Box::new(MapType::new(
+                        new_key_type,
+                        new_value_type,
+                        map_type.value_contains_null,
+                    ))))
+                }
+                DataType::Struct(struct_type) => {
+                    // build up the mapped child fields
+                    let children = struct_type
+                        .fields()
+                        .map(|field| make_field_physical(field, column_mapping_mode))
+                        .try_collect()?;
+                    Ok(DataType::Struct(Box::new(StructType::new(children))))
+                }
+                _ => {
+                    // types with no children don't change
+                    Ok(logical_dt.clone())
+                }
+            }
+        }
+        ColumnMappingMode::Id => panic!("No id"),
+    }
+}
+
 /// Get the state needed to process a scan. In particular this returns a triple of
 /// (all_fields_in_query, fields_to_read_from_parquet, have_partition_cols) where:
 /// - all_fields_in_query - all fields in the query as [`ColumnType`] enums
@@ -413,10 +485,11 @@ fn get_state_info(
             } else {
                 // Add to read schema, store field so we can build a `Column` expression later
                 // if needed (i.e. if we have partition columns)
-                let physical_name = logical_field.physical_name(column_mapping_mode)?;
-                let physical_field = logical_field.with_name(physical_name);
+                let physical_field = make_field_physical(logical_field, column_mapping_mode)?;
+                debug!("\n\n{logical_field:#?}\nAfter mapping: {physical_field:#?}\n\n");
+                let name = physical_field.name.clone();
                 read_fields.push(physical_field);
-                Ok(ColumnType::Selected(physical_name.to_string()))
+                Ok(ColumnType::Selected(name))
             }
         })
         .try_collect()?;