Fix: Correctly resolve renamed columns using Parquet field_id → DuckLake column_id mapping (#24)

Author: shefeek-jinnah

src/column_rename.rs

@@ -0,0 +1,229 @@
+//! 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::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 (Parquet) column names to new (DuckLake) column names
+    name_mapping: HashMap<String, String>,
+    /// Reverse mapping: new name -> old name, for looking up input columns
+    reverse_mapping: Arc<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(Arc::clone(&output_schema));
+        let properties = PlanProperties::new(
+            eq_props,
+            input.output_partitioning().clone(),
+            input.pipeline_behavior(),
+            Boundedness::Bounded,
+        );
+
+        // Pre-compute reverse mapping once (new_name -> old_name)
+        let reverse_mapping: HashMap<String, String> = name_mapping
+            .iter()
+            .map(|(old, new)| (new.clone(), old.clone()))
+            .collect();
+
+        Self {
+            input,
+            output_schema,
+            name_mapping,
+            reverse_mapping: Arc::new(reverse_mapping),
+            properties,
+        }
+    }
+}
+
+impl DisplayAs for ColumnRenameExec {
+    fn fmt_as(&self, _t: DisplayFormatType, f: &mut std::fmt::Formatter) -> std::fmt::Result {
+        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(),
+            ));
+        }
+
+        // Must call new() to rebuild properties from new child's partitioning
+        Ok(Arc::new(ColumnRenameExec::new(
+            Arc::clone(&children[0]),
+            Arc::clone(&self.output_schema),
+            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: Arc::clone(&self.output_schema),
+            reverse_mapping: Arc::clone(&self.reverse_mapping),
+        }))
+    }
+}
+
+/// Stream that renames columns in output batches
+struct ColumnRenameStream {
+    input: SendableRecordBatchStream,
+    output_schema: SchemaRef,
+    /// Mapping from output column name -> input column name (for renamed columns only)
+    reverse_mapping: Arc<HashMap<String, String>>,
+}
+
+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(
+                        Arc::clone(&self.output_schema),
+                        vec![],
+                        &options,
+                    )
+                } else {
+                    // Build columns by looking up each output field in the input batch
+                    let input_schema = batch.schema();
+                    let columns: Result<Vec<_>, _> = self
+                        .output_schema
+                        .fields()
+                        .iter()
+                        .map(|output_field| {
+                            // Check if this column was renamed (new_name -> old_name)
+                            let input_name = self
+                                .reverse_mapping
+                                .get(output_field.name())
+                                .map(|s| s.as_str())
+                                .unwrap_or_else(|| output_field.name().as_str());
+
+                            input_schema
+                                .index_of(input_name)
+                                .map(|idx| batch.column(idx).clone())
+                        })
+                        .collect();
+
+                    match columns {
+                        Ok(cols) => RecordBatch::try_new(Arc::clone(&self.output_schema), cols),
+                        Err(e) => Err(e),
+                    }
+                };
+
+                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 {
+        Arc::clone(&self.output_schema)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use arrow::datatypes::{DataType, Field, Schema};
+    use datafusion::physical_plan::EmptyRecordBatchStream;
+
+    #[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 mut reverse_mapping = HashMap::new();
+        reverse_mapping.insert("new_col".to_string(), "old_col".to_string());
+
+        let stream = ColumnRenameStream {
+            input: Box::pin(EmptyRecordBatchStream::new(input_schema)),
+            output_schema: Arc::clone(&output_schema),
+            reverse_mapping: Arc::new(reverse_mapping),
+        };
+
+        // The stream should report the output schema
+        assert_eq!(stream.schema().field(0).name(), "new_col");
+    }
+}

src/lib.rs

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

src/metadata_provider.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";
@@ -118,7 +118,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
@@ -236,14 +237,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_duckdb.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<_>, _>>()?;
@@ -306,10 +308,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_postgres.rs

@@ -174,7 +174,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",
@@ -185,10 +185,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()
@@ -390,7 +392,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
@@ -411,10 +413,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,