fix: handle renamed columns by matching Parquet field_id to DuckLake column_id (#24)

Author: Shefeek Jinnah

Cargo.lock

@@ -0,0 +1,253 @@
+//! Custom execution plan for renaming columns
+//!
+//! This module implements a DataFusion execution plan that wraps a scan
+//! and renames columns from their original Parquet names to current DuckLake names.
+//! This is needed when columns have been renamed in DuckLake metadata but the
+//! Parquet files still have the original column names.
+
+use std::any::Any;
+use std::collections::HashMap;
+use std::pin::Pin;
+use std::sync::Arc;
+use std::task::{Context, Poll};
+
+use arrow::datatypes::{Field, Schema, SchemaRef};
+use arrow::record_batch::RecordBatch;
+use datafusion::error::{DataFusionError, Result as DataFusionResult};
+use datafusion::execution::{RecordBatchStream, SendableRecordBatchStream, TaskContext};
+use datafusion::physical_expr::EquivalenceProperties;
+use datafusion::physical_plan::execution_plan::Boundedness;
+use datafusion::physical_plan::{
+    DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, PlanProperties,
+};
+use futures::Stream;
+
+/// Custom execution plan that renames columns from Parquet file names to current DuckLake names
+#[derive(Debug)]
+pub struct ColumnRenameExec {
+    /// The input execution plan (typically ParquetExec)
+    input: Arc<dyn ExecutionPlan>,
+    /// Output schema with renamed columns
+    output_schema: SchemaRef,
+    /// Mapping from old column names to new column names (for display purposes)
+    name_mapping: HashMap<String, String>,
+    /// Cached plan properties with updated schema
+    properties: PlanProperties,
+}
+
+impl ColumnRenameExec {
+    pub fn new(
+        input: Arc<dyn ExecutionPlan>,
+        output_schema: SchemaRef,
+        name_mapping: HashMap<String, String>,
+    ) -> Self {
+        // PlanProperties must use output schema for DataFusion schema validation
+        let eq_props = EquivalenceProperties::new(output_schema.clone());
+        let properties = PlanProperties::new(
+            eq_props,
+            input.output_partitioning().clone(),
+            input.pipeline_behavior(),
+            Boundedness::Bounded,
+        );
+
+        Self {
+            input,
+            output_schema,
+            name_mapping,
+            properties,
+        }
+    }
+}
+
+impl DisplayAs for ColumnRenameExec {
+    fn fmt_as(&self, t: DisplayFormatType, f: &mut std::fmt::Formatter) -> std::fmt::Result {
+        match t {
+            DisplayFormatType::Default | DisplayFormatType::Verbose => {
+                write!(f, "ColumnRenameExec: renames={}", self.name_mapping.len())
+            },
+            DisplayFormatType::TreeRender => {
+                write!(f, "ColumnRenameExec: renames={}", self.name_mapping.len())
+            },
+        }
+    }
+}
+
+impl ExecutionPlan for ColumnRenameExec {
+    fn name(&self) -> &str {
+        "ColumnRenameExec"
+    }
+
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+
+    fn properties(&self) -> &PlanProperties {
+        &self.properties
+    }
+
+    fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
+        vec![&self.input]
+    }
+
+    fn with_new_children(
+        self: Arc<Self>,
+        children: Vec<Arc<dyn ExecutionPlan>>,
+    ) -> DataFusionResult<Arc<dyn ExecutionPlan>> {
+        if children.len() != 1 {
+            return Err(DataFusionError::Internal(
+                "ColumnRenameExec expects exactly one child".into(),
+            ));
+        }
+
+        Ok(Arc::new(ColumnRenameExec::new(
+            children[0].clone(),
+            self.output_schema.clone(),
+            self.name_mapping.clone(),
+        )))
+    }
+
+    fn execute(
+        &self,
+        partition: usize,
+        context: Arc<TaskContext>,
+    ) -> DataFusionResult<SendableRecordBatchStream> {
+        let input_stream = self.input.execute(partition, context)?;
+
+        Ok(Box::pin(ColumnRenameStream {
+            input: input_stream,
+            output_schema: self.output_schema.clone(),
+        }))
+    }
+}
+
+/// Stream that renames columns in output batches
+struct ColumnRenameStream {
+    input: SendableRecordBatchStream,
+    output_schema: SchemaRef,
+}
+
+impl Stream for ColumnRenameStream {
+    type Item = DataFusionResult<RecordBatch>;
+
+    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        match Pin::new(&mut self.input).poll_next(cx) {
+            Poll::Ready(Some(Ok(batch))) => {
+                let result = if batch.num_columns() == 0 {
+                    // COUNT(*) case: preserve row count with empty schema
+                    use arrow::record_batch::RecordBatchOptions;
+                    let options = RecordBatchOptions::new().with_row_count(Some(batch.num_rows()));
+                    RecordBatch::try_new_with_options(self.output_schema.clone(), vec![], &options)
+                } else {
+                    RecordBatch::try_new(self.output_schema.clone(), batch.columns().to_vec())
+                };
+
+                match result {
+                    Ok(renamed_batch) => Poll::Ready(Some(Ok(renamed_batch))),
+                    Err(e) => {
+                        Poll::Ready(Some(Err(DataFusionError::ArrowError(Box::new(e), None))))
+                    },
+                }
+            },
+            Poll::Ready(Some(Err(e))) => Poll::Ready(Some(Err(e))),
+            Poll::Ready(None) => Poll::Ready(None),
+            Poll::Pending => Poll::Pending,
+        }
+    }
+}
+
+impl RecordBatchStream for ColumnRenameStream {
+    fn schema(&self) -> SchemaRef {
+        self.output_schema.clone()
+    }
+}
+
+/// Build a schema with renamed fields, preserving data types and metadata.
+pub fn build_renamed_schema(
+    input_schema: &Schema,
+    name_mapping: &HashMap<String, String>,
+) -> Schema {
+    let renamed_fields: Vec<Field> = input_schema
+        .fields()
+        .iter()
+        .map(|field| {
+            let new_name = name_mapping
+                .get(field.name())
+                .cloned()
+                .unwrap_or_else(|| field.name().clone());
+            Field::new(new_name, field.data_type().clone(), field.is_nullable())
+                .with_metadata(field.metadata().clone())
+        })
+        .collect();
+
+    Schema::new(renamed_fields).with_metadata(input_schema.metadata().clone())
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use arrow::datatypes::DataType;
+    use datafusion::physical_plan::EmptyRecordBatchStream;
+
+    #[test]
+    fn test_build_renamed_schema() {
+        let input_schema = Schema::new(vec![
+            Field::new("old_id", DataType::Int32, false),
+            Field::new("old_name", DataType::Utf8, true),
+        ]);
+
+        let mut name_mapping = HashMap::new();
+        name_mapping.insert("old_id".to_string(), "new_id".to_string());
+        name_mapping.insert("old_name".to_string(), "new_name".to_string());
+
+        let renamed_schema = build_renamed_schema(&input_schema, &name_mapping);
+
+        assert_eq!(renamed_schema.fields().len(), 2);
+        assert_eq!(renamed_schema.field(0).name(), "new_id");
+        assert_eq!(renamed_schema.field(1).name(), "new_name");
+        assert_eq!(renamed_schema.field(0).data_type(), &DataType::Int32);
+        assert_eq!(renamed_schema.field(1).data_type(), &DataType::Utf8);
+    }
+
+    #[test]
+    fn test_build_renamed_schema_partial_mapping() {
+        // Test when only some columns are renamed
+        let input_schema = Schema::new(vec![
+            Field::new("col1", DataType::Int32, false),
+            Field::new("col2", DataType::Utf8, true),
+            Field::new("col3", DataType::Float64, false),
+        ]);
+
+        let mut name_mapping = HashMap::new();
+        name_mapping.insert("col1".to_string(), "renamed_col1".to_string());
+        // col2 and col3 are not renamed
+
+        let renamed_schema = build_renamed_schema(&input_schema, &name_mapping);
+
+        assert_eq!(renamed_schema.field(0).name(), "renamed_col1");
+        assert_eq!(renamed_schema.field(1).name(), "col2"); // unchanged
+        assert_eq!(renamed_schema.field(2).name(), "col3"); // unchanged
+    }
+
+    #[test]
+    fn test_column_rename_stream_schema() {
+        let input_schema = Arc::new(Schema::new(vec![Field::new(
+            "old_col",
+            DataType::Int32,
+            false,
+        )]));
+
+        let output_schema = Arc::new(Schema::new(vec![Field::new(
+            "new_col",
+            DataType::Int32,
+            false,
+        )]));
+
+        let stream = ColumnRenameStream {
+            input: Box::pin(EmptyRecordBatchStream::new(input_schema)),
+            output_schema: output_schema.clone(),
+        };
+
+        // The stream should report the output schema
+        assert_eq!(stream.schema().field(0).name(), "new_col");
+    }
+}

src/column_rename.rs

@@ -36,6 +36,7 @@
 //! ```
 
 pub mod catalog;
+pub mod column_rename;
 pub mod delete_filter;
 pub mod encryption;
 pub mod error;

src/lib.rs

@@ -17,7 +17,7 @@ pub const SQL_LIST_TABLES: &str =
        AND ? >= begin_snapshot
        AND (? < end_snapshot OR end_snapshot IS NULL)";
 
-pub const SQL_GET_TABLE_COLUMNS: &str = "SELECT column_id, column_name, column_type
+pub const SQL_GET_TABLE_COLUMNS: &str = "SELECT column_id, column_name, column_type, nulls_allowed
      FROM ducklake_column
      WHERE table_id = ?
      ORDER BY column_order";
@@ -116,7 +116,8 @@ pub const SQL_LIST_ALL_COLUMNS: &str = "
         t.table_name,
         c.column_id,
         c.column_name,
-        c.column_type
+        c.column_type,
+        c.nulls_allowed
     FROM ducklake_schema s
     JOIN ducklake_table t ON s.schema_id = t.schema_id
     JOIN ducklake_column c ON t.table_id = c.table_id
@@ -232,14 +233,22 @@ pub struct DuckLakeTableColumn {
     pub column_name: String,
     /// DuckLake type string (e.g., "varchar", "int64", "decimal(10,2)")
     pub column_type: String,
+    /// Whether this column allows NULL values
+    pub is_nullable: bool,
 }
 
 impl DuckLakeTableColumn {
-    pub fn new(column_id: i64, column_name: String, column_type: String) -> Self {
+    pub fn new(
+        column_id: i64,
+        column_name: String,
+        column_type: String,
+        is_nullable: bool,
+    ) -> Self {
         Self {
             column_id,
             column_name,
             column_type,
+            is_nullable,
         }
     }
 }

src/metadata_provider.rs

@@ -145,10 +145,12 @@ impl MetadataProvider for DuckdbMetadataProvider {
                 let column_id: i64 = row.get(0)?;
                 let column_name: String = row.get(1)?;
                 let column_type: String = row.get(2)?;
+                let nulls_allowed: Option<bool> = row.get(3)?;
                 Ok(DuckLakeTableColumn::new(
                     column_id,
                     column_name,
                     column_type,
+                    nulls_allowed.unwrap_or(true),
                 ))
             })?
             .collect::<Result<Vec<_>, _>>()?;
@@ -303,10 +305,12 @@ impl MetadataProvider for DuckdbMetadataProvider {
                 |row| {
                     let schema_name: String = row.get(0)?;
                     let table_name: String = row.get(1)?;
+                    let nulls_allowed: Option<bool> = row.get(5)?;
                     let column = DuckLakeTableColumn {
                         column_id: row.get(2)?,
                         column_name: row.get(3)?,
                         column_type: row.get(4)?,
+                        is_nullable: nulls_allowed.unwrap_or(true),
                     };
                     Ok(ColumnWithTable {
                         schema_name,

src/metadata_provider_duckdb.rs

@@ -163,7 +163,7 @@ impl MetadataProvider for PostgresMetadataProvider {
     fn get_table_structure(&self, table_id: i64) -> Result<Vec<DuckLakeTableColumn>> {
         block_on(async {
             let rows = sqlx::query(
-                "SELECT column_id, column_name, column_type
+                "SELECT column_id, column_name, column_type, nulls_allowed
                  FROM ducklake_column
                  WHERE table_id = $1
                  ORDER BY column_order",
@@ -174,10 +174,12 @@ impl MetadataProvider for PostgresMetadataProvider {
 
             rows.into_iter()
                 .map(|row| {
+                    let nulls_allowed: Option<bool> = row.try_get(3)?;
                     Ok(DuckLakeTableColumn {
                         column_id: row.try_get(0)?,
                         column_name: row.try_get(1)?,
                         column_type: row.try_get(2)?,
+                        is_nullable: nulls_allowed.unwrap_or(true),
                     })
                 })
                 .collect()
@@ -375,7 +377,7 @@ impl MetadataProvider for PostgresMetadataProvider {
     fn list_all_columns(&self, snapshot_id: i64) -> Result<Vec<ColumnWithTable>> {
         block_on(async {
             let rows = sqlx::query(
-                "SELECT s.schema_name, t.table_name, c.column_id, c.column_name, c.column_type
+                "SELECT s.schema_name, t.table_name, c.column_id, c.column_name, c.column_type, c.nulls_allowed
                  FROM ducklake_schema s
                  JOIN ducklake_table t ON s.schema_id = t.schema_id
                  JOIN ducklake_column c ON t.table_id = c.table_id
@@ -396,10 +398,12 @@ impl MetadataProvider for PostgresMetadataProvider {
                 .map(|row| {
                     let schema_name: String = row.try_get(0)?;
                     let table_name: String = row.try_get(1)?;
+                    let nulls_allowed: Option<bool> = row.try_get(5)?;
                     let column = DuckLakeTableColumn {
                         column_id: row.try_get(2)?,
                         column_name: row.try_get(3)?,
                         column_type: row.try_get(4)?,
+                        is_nullable: nulls_allowed.unwrap_or(true),
                     };
                     Ok(ColumnWithTable {
                         schema_name,