iceberg/arrow/
record_batch_projector.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use std::sync::Arc;
19
20use arrow_array::{Array, ArrayRef, RecordBatch, StructArray, make_array};
21use arrow_buffer::NullBuffer;
22use arrow_schema::{DataType, Field, FieldRef, Fields, Schema, SchemaRef};
23use parquet::arrow::PARQUET_FIELD_ID_META_KEY;
24
25use crate::arrow::schema::schema_to_arrow_schema;
26use crate::error::{Result, invalid_data};
27use crate::spec::Schema as IcebergSchema;
28use crate::{Error, ErrorKind};
29
30/// Help to project specific field from `RecordBatch`` according to the fields id.
31#[derive(Clone, Debug, PartialEq, Eq)]
32pub struct RecordBatchProjector {
33    // A vector of vectors, where each inner vector represents the index path to access a specific field in a nested structure.
34    // E.g. [[0], [1, 2]] means the first field is accessed directly from the first column,
35    // while the second field is accessed from the second column and then from its third subcolumn (second column must be a struct column).
36    field_indices: Vec<Vec<usize>>,
37    // The schema reference after projection. This schema is derived from the original schema based on the given field IDs.
38    projected_schema: SchemaRef,
39}
40
41impl RecordBatchProjector {
42    /// Init ArrowFieldProjector
43    ///
44    /// This function will iterate through the field and fetch the field from the original schema according to the field ids.
45    /// The function to fetch the field id from the field is provided by `field_id_fetch_func`, return None if the field need to be skipped.
46    /// This function will iterate through the nested fields if the field is a struct, `searchable_field_func` can be used to control whether
47    /// iterate into the nested fields.
48    pub(crate) fn new<F1, F2>(
49        original_schema: SchemaRef,
50        field_ids: &[i32],
51        field_id_fetch_func: F1,
52        searchable_field_func: F2,
53    ) -> Result<Self>
54    where
55        F1: Fn(&Field) -> Result<Option<i64>>,
56        F2: Fn(&Field) -> bool,
57    {
58        let mut field_indices = Vec::with_capacity(field_ids.len());
59        let mut fields = Vec::with_capacity(field_ids.len());
60        for &id in field_ids {
61            let mut field_index = vec![];
62            let field = Self::fetch_field_index(
63                original_schema.fields(),
64                &mut field_index,
65                id as i64,
66                &field_id_fetch_func,
67                &searchable_field_func,
68            )?
69            .ok_or_else(|| {
70                Error::new(ErrorKind::Unexpected, "Field not found")
71                    .with_context("field_id", id.to_string())
72            })?;
73            fields.push(field.clone());
74            field_indices.push(field_index);
75        }
76        let delete_arrow_schema = Arc::new(Schema::new(fields));
77        Ok(Self {
78            field_indices,
79            projected_schema: delete_arrow_schema,
80        })
81    }
82
83    /// Create RecordBatchProjector using Iceberg schema.
84    ///
85    /// This constructor converts the Iceberg schema to Arrow schema with field ID metadata,
86    /// then uses the standard field ID lookup for projection.
87    ///
88    /// # Arguments
89    /// * `iceberg_schema` - The Iceberg schema for field ID mapping  
90    /// * `target_field_ids` - The field IDs to project
91    pub fn from_iceberg_schema(
92        iceberg_schema: Arc<IcebergSchema>,
93        target_field_ids: &[i32],
94    ) -> Result<Self> {
95        let arrow_schema_with_ids = Arc::new(schema_to_arrow_schema(&iceberg_schema)?);
96
97        let field_id_fetch_func = |field: &Field| -> Result<Option<i64>> {
98            if let Some(value) = field.metadata().get(PARQUET_FIELD_ID_META_KEY) {
99                let field_id = value.parse::<i32>().map_err(|e| {
100                    invalid_data!("Failed to parse field id")
101                        .with_context("value", value)
102                        .with_source(e)
103                })?;
104                Ok(Some(field_id as i64))
105            } else {
106                Ok(None)
107            }
108        };
109
110        let searchable_field_func = |_field: &Field| -> bool { true };
111
112        Self::new(
113            arrow_schema_with_ids,
114            target_field_ids,
115            field_id_fetch_func,
116            searchable_field_func,
117        )
118    }
119
120    fn fetch_field_index<F1, F2>(
121        fields: &Fields,
122        index_vec: &mut Vec<usize>,
123        target_field_id: i64,
124        field_id_fetch_func: &F1,
125        searchable_field_func: &F2,
126    ) -> Result<Option<FieldRef>>
127    where
128        F1: Fn(&Field) -> Result<Option<i64>>,
129        F2: Fn(&Field) -> bool,
130    {
131        for (pos, field) in fields.iter().enumerate() {
132            let id = field_id_fetch_func(field)?;
133            if let Some(id) = id
134                && target_field_id == id
135            {
136                index_vec.push(pos);
137                return Ok(Some(field.clone()));
138            }
139            if let DataType::Struct(inner) = field.data_type()
140                && searchable_field_func(field)
141                && let Some(res) = Self::fetch_field_index(
142                    inner,
143                    index_vec,
144                    target_field_id,
145                    field_id_fetch_func,
146                    searchable_field_func,
147                )?
148            {
149                index_vec.push(pos);
150                return Ok(Some(res));
151            }
152        }
153        Ok(None)
154    }
155
156    /// Return the reference of projected schema
157    pub(crate) fn projected_schema_ref(&self) -> &SchemaRef {
158        &self.projected_schema
159    }
160
161    /// Do projection with record batch
162    pub(crate) fn project_batch(&self, batch: RecordBatch) -> Result<RecordBatch> {
163        RecordBatch::try_new(
164            self.projected_schema.clone(),
165            self.project_column(batch.columns())?,
166        )
167        .map_err(|err| invalid_data!("{err}"))
168    }
169
170    /// Do projection with columns
171    pub fn project_column(&self, batch: &[ArrayRef]) -> Result<Vec<ArrayRef>> {
172        let mut columns = Vec::with_capacity(self.field_indices.len());
173        for index_vec in &self.field_indices {
174            columns.push(Self::get_column_by_field_index(batch, index_vec)?);
175        }
176        Ok(columns)
177    }
178
179    fn get_column_by_field_index(batch: &[ArrayRef], field_index: &[usize]) -> Result<ArrayRef> {
180        let mut rev_iterator = field_index.iter().rev();
181        let mut array = &batch[*rev_iterator.next().unwrap()];
182        let mut ancestor_nulls = None;
183        for idx in rev_iterator {
184            let struct_array = array
185                .as_any()
186                .downcast_ref::<StructArray>()
187                .ok_or_else(|| {
188                    Error::new(ErrorKind::Unexpected, "Cannot convert Array to StructArray")
189                })?;
190            ancestor_nulls = NullBuffer::union(ancestor_nulls.as_ref(), struct_array.nulls());
191            array = struct_array.column(*idx);
192        }
193        let Some(ancestor_nulls) = ancestor_nulls else {
194            // Iceberg leaf types handled here have logical_nulls() == nulls(), so no null
195            // materialization is needed when no ancestor contributes nulls.
196            return Ok(Arc::clone(array));
197        };
198        let null_buffer = NullBuffer::union(Some(&ancestor_nulls), array.logical_nulls().as_ref());
199        Ok(make_array(
200            array.to_data().into_builder().nulls(null_buffer).build()?,
201        ))
202    }
203}
204
205#[cfg(test)]
206mod test {
207    use std::sync::Arc;
208
209    use arrow_array::{Array, ArrayRef, Int32Array, RecordBatch, StringArray, StructArray};
210    use arrow_buffer::NullBuffer;
211    use arrow_schema::{DataType, Field, Fields, Schema};
212
213    use crate::arrow::record_batch_projector::RecordBatchProjector;
214    use crate::spec::{NestedField, PrimitiveType, Schema as IcebergSchema, Type};
215    use crate::{Error, ErrorKind};
216
217    fn nested_projector() -> (RecordBatchProjector, Field, Field) {
218        let leaf_field = Field::new("leaf", DataType::Int32, true);
219        let inner_field = Field::new(
220            "inner",
221            DataType::Struct(Fields::from(vec![leaf_field.clone()])),
222            true,
223        );
224        let outer_field = Field::new(
225            "outer",
226            DataType::Struct(Fields::from(vec![inner_field.clone()])),
227            true,
228        );
229        let schema = Arc::new(Schema::new(vec![outer_field]));
230        let field_id_fetch_func = |field: &Field| -> crate::Result<Option<i64>> {
231            match field.name().as_str() {
232                "outer" => Ok(Some(1)),
233                "inner" => Ok(Some(2)),
234                "leaf" => Ok(Some(3)),
235                _ => Ok(None),
236            }
237        };
238        let projector =
239            RecordBatchProjector::new(schema, &[3], field_id_fetch_func, |_| true).unwrap();
240        (projector, inner_field, leaf_field)
241    }
242
243    #[test]
244    fn test_record_batch_projector_nested_level() {
245        let inner_fields = vec![
246            Field::new("inner_field1", DataType::Int32, false),
247            Field::new("inner_field2", DataType::Utf8, false),
248        ];
249        let fields = vec![
250            Field::new("field1", DataType::Int32, false),
251            Field::new(
252                "field2",
253                DataType::Struct(Fields::from(inner_fields.clone())),
254                false,
255            ),
256        ];
257        let schema = Arc::new(Schema::new(fields));
258
259        let field_id_fetch_func = |field: &Field| match field.name().as_str() {
260            "field1" => Ok(Some(1)),
261            "field2" => Ok(Some(2)),
262            "inner_field1" => Ok(Some(3)),
263            "inner_field2" => Ok(Some(4)),
264            _ => Err(Error::new(ErrorKind::Unexpected, "Field id not found")),
265        };
266        let projector =
267            RecordBatchProjector::new(schema.clone(), &[1, 3], field_id_fetch_func, |_| true)
268                .unwrap();
269
270        assert_eq!(projector.field_indices.len(), 2);
271        assert_eq!(projector.field_indices[0], vec![0]);
272        assert_eq!(projector.field_indices[1], vec![0, 1]);
273
274        let int_array = Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef;
275        let inner_int_array = Arc::new(Int32Array::from(vec![4, 5, 6])) as ArrayRef;
276        let inner_string_array = Arc::new(StringArray::from(vec!["x", "y", "z"])) as ArrayRef;
277        let struct_array = Arc::new(StructArray::from(vec![
278            (
279                Arc::new(inner_fields[0].clone()),
280                inner_int_array as ArrayRef,
281            ),
282            (
283                Arc::new(inner_fields[1].clone()),
284                inner_string_array as ArrayRef,
285            ),
286        ])) as ArrayRef;
287        let batch = RecordBatch::try_new(schema, vec![int_array, struct_array]).unwrap();
288
289        let projected_batch = projector.project_batch(batch).unwrap();
290        assert_eq!(projected_batch.num_columns(), 2);
291        let projected_int_array = projected_batch
292            .column(0)
293            .as_any()
294            .downcast_ref::<Int32Array>()
295            .unwrap();
296        let projected_inner_int_array = projected_batch
297            .column(1)
298            .as_any()
299            .downcast_ref::<Int32Array>()
300            .unwrap();
301
302        assert_eq!(projected_int_array.values(), &[1, 2, 3]);
303        assert_eq!(projected_inner_int_array.values(), &[4, 5, 6]);
304    }
305
306    #[test]
307    fn test_record_batch_projector_top_level_nullable_column() {
308        let iceberg_schema = IcebergSchema::builder()
309            .with_schema_id(0)
310            .with_fields(vec![
311                NestedField::optional(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
312            ])
313            .build()
314            .unwrap();
315        let projector =
316            RecordBatchProjector::from_iceberg_schema(Arc::new(iceberg_schema), &[1]).unwrap();
317        let input = Arc::new(Int32Array::from(vec![Some(10), None, Some(30)])) as ArrayRef;
318
319        let projected = projector
320            .project_column(std::slice::from_ref(&input))
321            .unwrap();
322        assert!(Arc::ptr_eq(&projected[0], &input));
323        let projected_array = projected[0].as_any().downcast_ref::<Int32Array>().unwrap();
324
325        assert_eq!(projected_array.value(0), 10);
326        assert_eq!(projected_array.null_count(), 1);
327        assert!(projected_array.is_null(1));
328        assert_eq!(projected_array.value(2), 30);
329    }
330
331    #[test]
332    fn test_record_batch_projector_nested_nullable_leaf_without_parent_nulls() {
333        let (projector, inner_field, leaf_field) = nested_projector();
334        let leaf = Arc::new(Int32Array::from(vec![Some(10), None, Some(30)])) as ArrayRef;
335        let inner = Arc::new(StructArray::new(
336            Fields::from(vec![leaf_field]),
337            vec![leaf.clone()],
338            None,
339        )) as ArrayRef;
340        let outer = Arc::new(StructArray::new(
341            Fields::from(vec![inner_field]),
342            vec![inner],
343            None,
344        )) as ArrayRef;
345
346        let projected = projector.project_column(&[outer]).unwrap();
347        assert!(Arc::ptr_eq(&projected[0], &leaf));
348    }
349
350    #[test]
351    fn test_record_batch_projector_nested_all_valid_slice_reuses_leaf() {
352        let (projector, inner_field, leaf_field) = nested_projector();
353        let leaf = Arc::new(Int32Array::from(vec![Some(0), Some(10), None, Some(30)])) as ArrayRef;
354        let inner = Arc::new(StructArray::new(
355            Fields::from(vec![leaf_field]),
356            vec![leaf],
357            None,
358        )) as ArrayRef;
359        let outer = StructArray::new(
360            Fields::from(vec![inner_field]),
361            vec![inner],
362            Some(NullBuffer::from(vec![false, true, true, true])),
363        )
364        .slice(1, 3);
365        assert_eq!(outer.null_count(), 0);
366        assert!(outer.nulls().is_some());
367
368        let sliced_inner = outer
369            .column(0)
370            .as_any()
371            .downcast_ref::<StructArray>()
372            .unwrap();
373        let sliced_leaf = sliced_inner.column(0).clone();
374        let projected = projector
375            .project_column(&[Arc::new(outer) as ArrayRef])
376            .unwrap();
377        assert!(Arc::ptr_eq(&projected[0], &sliced_leaf));
378    }
379
380    #[test]
381    fn test_record_batch_projector_propagates_nested_parent_nulls() {
382        let (projector, inner_field, leaf_field) = nested_projector();
383        let leaf = Arc::new(Int32Array::from(vec![Some(10), Some(20), Some(30), None])) as ArrayRef;
384        let inner = Arc::new(StructArray::new(
385            Fields::from(vec![leaf_field]),
386            vec![leaf],
387            Some(NullBuffer::from(vec![true, true, false, true])),
388        )) as ArrayRef;
389        let outer = Arc::new(StructArray::new(
390            Fields::from(vec![inner_field]),
391            vec![inner],
392            Some(NullBuffer::from(vec![true, false, true, true])),
393        )) as ArrayRef;
394
395        let projected = projector.project_column(&[outer]).unwrap();
396        assert_eq!(
397            projected[0].as_ref(),
398            &Int32Array::from(vec![Some(10), None, None, None])
399        );
400    }
401
402    #[test]
403    fn test_field_not_found() {
404        let inner_fields = vec![
405            Field::new("inner_field1", DataType::Int32, false),
406            Field::new("inner_field2", DataType::Utf8, false),
407        ];
408
409        let fields = vec![
410            Field::new("field1", DataType::Int32, false),
411            Field::new(
412                "field2",
413                DataType::Struct(Fields::from(inner_fields.clone())),
414                false,
415            ),
416        ];
417        let schema = Arc::new(Schema::new(fields));
418
419        let field_id_fetch_func = |field: &Field| match field.name().as_str() {
420            "field1" => Ok(Some(1)),
421            "field2" => Ok(Some(2)),
422            "inner_field1" => Ok(Some(3)),
423            "inner_field2" => Ok(Some(4)),
424            _ => Err(Error::new(ErrorKind::Unexpected, "Field id not found")),
425        };
426        let projector =
427            RecordBatchProjector::new(schema.clone(), &[1, 5], field_id_fetch_func, |_| true);
428
429        assert!(projector.is_err());
430    }
431
432    #[test]
433    fn test_field_not_reachable() {
434        let inner_fields = vec![
435            Field::new("inner_field1", DataType::Int32, false),
436            Field::new("inner_field2", DataType::Utf8, false),
437        ];
438
439        let fields = vec![
440            Field::new("field1", DataType::Int32, false),
441            Field::new(
442                "field2",
443                DataType::Struct(Fields::from(inner_fields.clone())),
444                false,
445            ),
446        ];
447        let schema = Arc::new(Schema::new(fields));
448
449        let field_id_fetch_func = |field: &Field| match field.name().as_str() {
450            "field1" => Ok(Some(1)),
451            "field2" => Ok(Some(2)),
452            "inner_field1" => Ok(Some(3)),
453            "inner_field2" => Ok(Some(4)),
454            _ => Err(Error::new(ErrorKind::Unexpected, "Field id not found")),
455        };
456        let projector =
457            RecordBatchProjector::new(schema.clone(), &[3], field_id_fetch_func, |_| false);
458        assert!(projector.is_err());
459
460        let projector =
461            RecordBatchProjector::new(schema.clone(), &[3], field_id_fetch_func, |_| true);
462        assert!(projector.is_ok());
463    }
464
465    #[test]
466    fn test_from_iceberg_schema() {
467        let iceberg_schema = IcebergSchema::builder()
468            .with_schema_id(0)
469            .with_fields(vec![
470                NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
471                NestedField::required(2, "name", Type::Primitive(PrimitiveType::String)).into(),
472                NestedField::optional(3, "age", Type::Primitive(PrimitiveType::Int)).into(),
473            ])
474            .build()
475            .unwrap();
476
477        let projector =
478            RecordBatchProjector::from_iceberg_schema(Arc::new(iceberg_schema), &[1, 3]).unwrap();
479
480        assert_eq!(projector.field_indices.len(), 2);
481        assert_eq!(projector.projected_schema_ref().fields().len(), 2);
482        assert_eq!(projector.projected_schema_ref().field(0).name(), "id");
483        assert_eq!(projector.projected_schema_ref().field(1).name(), "age");
484    }
485}
iceberg/arrow/record_batch_projector.rs

iceberg/arrow/
record_batch_projector.rs
