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

iceberg/arrow/
record_batch_projector.rs
