ailake_query/
schema_filler.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2//! Schema filler — inject missing columns at read time (Phase G).
3//!
4//! When the table schema has been evolved by adding new columns, old data files
5//! do not contain those columns. `SchemaFiller::fill` detects absent columns and
6//! appends them to the `RecordBatch`, filled with the field's `initial_default`
7//! (or null when no default is set). This implements schema evolution without
8//! rewriting data files, equivalent to Iceberg V2/V3 §4.1.1.
9
10use std::sync::Arc;
11
12use ailake_catalog::SchemaField;
13use ailake_core::{AilakeError, AilakeResult};
14use arrow_array::{
15    Array, ArrayRef, BooleanArray, Date32Array, Float32Array, Float64Array, Int32Array, Int64Array,
16    StringArray, TimestampMicrosecondArray,
17};
18use arrow_schema::{DataType, Field, Schema, TimeUnit};
19
20pub struct SchemaFiller;
21
22impl SchemaFiller {
23    /// Inject columns present in `schema_fields` but absent from `batch`.
24    ///
25    /// Added columns appear after the existing columns, filled with `initial_default`
26    /// (or null). Columns already in the batch are left untouched.
27    ///
28    /// Returns `batch` unchanged if `schema_fields` is empty or no columns are missing.
29    pub fn fill(
30        batch: arrow_array::RecordBatch,
31        schema_fields: &[SchemaField],
32    ) -> AilakeResult<arrow_array::RecordBatch> {
33        if schema_fields.is_empty() {
34            return Ok(batch);
35        }
36
37        let batch_schema = batch.schema();
38        let existing: std::collections::HashSet<&str> = batch_schema
39            .fields()
40            .iter()
41            .map(|f| f.name().as_str())
42            .collect();
43
44        let missing: Vec<&SchemaField> = schema_fields
45            .iter()
46            .filter(|sf| !existing.contains(sf.name.as_str()))
47            .collect();
48
49        if missing.is_empty() {
50            return Ok(batch);
51        }
52
53        let n = batch.num_rows();
54        let mut new_fields: Vec<arrow_schema::FieldRef> =
55            batch.schema().fields().iter().cloned().collect();
56        let mut new_cols: Vec<Arc<dyn Array>> = batch.columns().to_vec();
57
58        for sf in missing {
59            let dtype = iceberg_type_to_arrow(&sf.iceberg_type);
60            let arr = make_default_array(&dtype, sf.initial_default.as_ref(), n)?;
61            new_fields.push(Arc::new(Field::new(sf.name.clone(), dtype, !sf.required)));
62            new_cols.push(arr);
63        }
64
65        let new_schema = Arc::new(Schema::new(new_fields));
66        arrow_array::RecordBatch::try_new(new_schema, new_cols)
67            .map_err(|e| AilakeError::Arrow(e.to_string()))
68    }
69}
70
71/// Map an Iceberg type string to an Arrow `DataType`.
72///
73/// Complex types (`list<…>`, `map<…>`, `struct<…>`) are mapped to `Utf8`
74/// as a conservative fallback — the raw JSON is stored as a string.
75pub fn iceberg_type_to_arrow(typ: &str) -> DataType {
76    match typ.trim() {
77        "boolean" => DataType::Boolean,
78        "int" | "integer" => DataType::Int32,
79        "long" => DataType::Int64,
80        "float" => DataType::Float32,
81        "double" => DataType::Float64,
82        "date" => DataType::Date32,
83        "time" => DataType::Time64(TimeUnit::Microsecond),
84        "timestamp" => DataType::Timestamp(TimeUnit::Microsecond, None),
85        "timestamptz" => DataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into())),
86        "string" | "uuid" => DataType::Utf8,
87        "binary" | "fixed" => DataType::Binary,
88        _ => DataType::Utf8,
89    }
90}
91
92/// Build an array of `n` rows filled with `default_val` (or all-null if `None`).
93fn make_default_array(
94    dtype: &DataType,
95    default_val: Option<&serde_json::Value>,
96    n: usize,
97) -> AilakeResult<ArrayRef> {
98    use serde_json::Value;
99
100    Ok(match dtype {
101        DataType::Boolean => {
102            let v = default_val.and_then(Value::as_bool);
103            Arc::new(BooleanArray::from(vec![v; n]))
104        }
105        DataType::Int32 => {
106            let v = default_val.and_then(Value::as_i64).map(|i| i as i32);
107            Arc::new(Int32Array::from(vec![v; n]))
108        }
109        DataType::Int64 => {
110            let v = default_val.and_then(Value::as_i64);
111            Arc::new(Int64Array::from(vec![v; n]))
112        }
113        DataType::Float32 => {
114            let v = default_val.and_then(Value::as_f64).map(|f| f as f32);
115            Arc::new(Float32Array::from(vec![v; n]))
116        }
117        DataType::Float64 => {
118            let v = default_val.and_then(Value::as_f64);
119            Arc::new(Float64Array::from(vec![v; n]))
120        }
121        DataType::Date32 => {
122            // Iceberg date default is an integer (days since epoch).
123            let v = default_val.and_then(Value::as_i64).map(|d| d as i32);
124            Arc::new(Date32Array::from(vec![v; n]))
125        }
126        DataType::Timestamp(TimeUnit::Microsecond, tz) => {
127            // Iceberg timestamp default is µs since epoch (i64).
128            let v = default_val.and_then(Value::as_i64);
129            let arr = TimestampMicrosecondArray::from(vec![v; n]);
130            Arc::new(if tz.is_some() {
131                arr.with_timezone("UTC")
132            } else {
133                arr
134            })
135        }
136        DataType::Utf8 => {
137            let v: Option<&str> = default_val.and_then(Value::as_str);
138            Arc::new(StringArray::from(vec![v; n]))
139        }
140        _ => {
141            // Binary, complex types, unknowns — inject null Utf8.
142            Arc::new(StringArray::from(vec![None::<&str>; n]))
143        }
144    })
145}
146
147#[cfg(test)]
148mod tests {
149    use std::sync::Arc;
150
151    use ailake_catalog::SchemaField;
152    use arrow_array::{Array, Float32Array, Int32Array, RecordBatch, StringArray};
153    use arrow_schema::{DataType, Field, Schema};
154
155    use super::SchemaFiller;
156
157    fn make_base_batch() -> RecordBatch {
158        let schema = Arc::new(Schema::new(vec![
159            Field::new("id", DataType::Int32, false),
160            Field::new("text", DataType::Utf8, true),
161        ]));
162        RecordBatch::try_new(
163            schema,
164            vec![
165                Arc::new(Int32Array::from(vec![1, 2, 3])),
166                Arc::new(StringArray::from(vec!["a", "b", "c"])),
167            ],
168        )
169        .unwrap()
170    }
171
172    #[test]
173    fn no_op_when_no_schema_fields() {
174        let batch = make_base_batch();
175        let filled = SchemaFiller::fill(batch.clone(), &[]).unwrap();
176        assert_eq!(filled.num_columns(), batch.num_columns());
177        assert_eq!(filled.num_rows(), batch.num_rows());
178    }
179
180    #[test]
181    fn no_op_when_all_columns_present() {
182        let batch = make_base_batch();
183        let fields = vec![
184            SchemaField {
185                id: 1,
186                name: "id".into(),
187                required: true,
188                iceberg_type: "int".into(),
189                initial_default: None,
190                write_default: None,
191            },
192            SchemaField {
193                id: 2,
194                name: "text".into(),
195                required: false,
196                iceberg_type: "string".into(),
197                initial_default: None,
198                write_default: None,
199            },
200        ];
201        let filled = SchemaFiller::fill(batch.clone(), &fields).unwrap();
202        assert_eq!(filled.num_columns(), 2);
203    }
204
205    #[test]
206    fn injects_missing_column_with_null_default() {
207        let batch = make_base_batch();
208        let fields = vec![SchemaField {
209            id: 3,
210            name: "score".into(),
211            required: false,
212            iceberg_type: "float".into(),
213            initial_default: None,
214            write_default: None,
215        }];
216        let filled = SchemaFiller::fill(batch, &fields).unwrap();
217        assert_eq!(filled.num_columns(), 3);
218        assert_eq!(filled.num_rows(), 3);
219        let score_col = filled
220            .column_by_name("score")
221            .unwrap()
222            .as_any()
223            .downcast_ref::<Float32Array>()
224            .unwrap();
225        // All nulls because initial_default is None.
226        assert!(!score_col.is_valid(0));
227        assert!(!score_col.is_valid(1));
228    }
229
230    #[test]
231    fn injects_missing_column_with_value_default() {
232        let batch = make_base_batch();
233        let fields = vec![SchemaField {
234            id: 4,
235            name: "score".into(),
236            required: false,
237            iceberg_type: "float".into(),
238            initial_default: Some(serde_json::json!(0.5)),
239            write_default: None,
240        }];
241        let filled = SchemaFiller::fill(batch, &fields).unwrap();
242        let score_col = filled
243            .column_by_name("score")
244            .unwrap()
245            .as_any()
246            .downcast_ref::<Float32Array>()
247            .unwrap();
248        assert!((score_col.value(0) - 0.5).abs() < 1e-6);
249        assert!((score_col.value(2) - 0.5).abs() < 1e-6);
250    }
251
252    #[test]
253    fn injects_string_column_with_default() {
254        let batch = make_base_batch();
255        let fields = vec![SchemaField {
256            id: 5,
257            name: "category".into(),
258            required: false,
259            iceberg_type: "string".into(),
260            initial_default: Some(serde_json::json!("uncategorized")),
261            write_default: None,
262        }];
263        let filled = SchemaFiller::fill(batch, &fields).unwrap();
264        let cat = filled
265            .column_by_name("category")
266            .unwrap()
267            .as_any()
268            .downcast_ref::<StringArray>()
269            .unwrap();
270        assert_eq!(cat.value(0), "uncategorized");
271        assert_eq!(cat.value(2), "uncategorized");
272    }
273}
ailake_query/schema_filler.rs

ailake_query/
schema_filler.rs
