feat: add Struct Accessors to BoundReferences

Author: sdd

crates/iceberg/src/expr/accessor.rs

@@ -0,0 +1,73 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::spec::{Literal, Struct, Type};
+use serde_derive::{Deserialize, Serialize};
+use std::sync::Arc;
+
+#[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Eq)]
+enum InnerOrType {
+    Inner(Arc<StructAccessor>),
+    Type(Type),
+}
+
+#[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Eq)]
+pub struct StructAccessor {
+    position: i32,
+    inner_or_type: InnerOrType,
+}
+
+pub(crate) type StructAccessorRef = Arc<StructAccessor>;
+
+impl StructAccessor {
+    pub(crate) fn new(position: i32, r#type: Type) -> Self {
+        StructAccessor {
+            position,
+            inner_or_type: InnerOrType::Type(r#type),
+        }
+    }
+
+    pub(crate) fn wrap(position: i32, inner: StructAccessorRef) -> Self {
+        StructAccessor {
+            position,
+            inner_or_type: InnerOrType::Inner(inner),
+        }
+    }
+
+    pub fn position(&self) -> i32 {
+        self.position
+    }
+
+    fn r#type(&self) -> &Type {
+        match &self.inner_or_type {
+            InnerOrType::Inner(inner) => inner.r#type(),
+            InnerOrType::Type(r#type) => r#type,
+        }
+    }
+
+    fn get<'a>(&'a self, container: &'a Struct) -> &Literal {
+        match &self.inner_or_type {
+            InnerOrType::Inner(inner) => match container.get(self.position) {
+                Literal::Struct(wrapped) => inner.get(wrapped),
+                _ => {
+                    panic!("Nested accessor should only be wrapping a Struct")
+                }
+            },
+            InnerOrType::Type(_) => container.get(self.position),
+        }
+    }
+}

crates/iceberg/src/expr/mod.rs

@@ -22,6 +22,7 @@ mod term;
 use std::fmt::{Display, Formatter};
 
 pub use term::*;
+pub(crate) mod accessor;
 mod predicate;
 
 use crate::spec::SchemaRef;

crates/iceberg/src/expr/term.rs

@@ -21,6 +21,7 @@ use std::fmt::{Display, Formatter};
 
 use fnv::FnvHashSet;
 
+use crate::expr::accessor::{StructAccessor, StructAccessorRef};
 use crate::expr::Bind;
 use crate::expr::{BinaryExpression, Predicate, PredicateOperator, SetExpression, UnaryExpression};
 use crate::spec::{Datum, NestedField, NestedFieldRef, SchemaRef};
@@ -333,7 +334,19 @@ impl Bind for Reference {
                 format!("Field {} not found in schema", self.name),
             )
         })?;
-        Ok(BoundReference::new(self.name.clone(), field.clone()))
+
+        let accessor = schema.accessor_for_field_id(field.id).ok_or_else(|| {
+            Error::new(
+                ErrorKind::DataInvalid,
+                format!("Accessor for Field {} not found", self.name),
+            )
+        })?;
+
+        Ok(BoundReference::new(
+            self.name.clone(),
+            field.clone(),
+            accessor.clone(),
+        ))
     }
 }
 
@@ -344,21 +357,32 @@ pub struct BoundReference {
     // For example, if the field is `a.b.c`, then `field.name` is `c`, but `original_name` is `a.b.c`.
     column_name: String,
     field: NestedFieldRef,
+    accessor: StructAccessorRef,
 }
 
 impl BoundReference {
     /// Creates a new bound reference.
-    pub fn new(name: impl Into<String>, field: NestedFieldRef) -> Self {
+    pub fn new(
+        name: impl Into<String>,
+        field: NestedFieldRef,
+        accessor: StructAccessorRef,
+    ) -> Self {
         Self {
             column_name: name.into(),
             field,
+            accessor,
         }
     }
 
     /// Return the field of this reference.
     pub fn field(&self) -> &NestedField {
         &self.field
     }
+
+    /// Get this BoundReference's Accessor
+    pub fn accessor(&self) -> &StructAccessor {
+        &self.accessor
+    }
 }
 
 impl Display for BoundReference {
@@ -374,6 +398,7 @@ pub type BoundTerm = BoundReference;
 mod tests {
     use std::sync::Arc;
 
+    use crate::expr::accessor::StructAccessor;
     use crate::expr::{Bind, BoundReference, Reference};
     use crate::spec::{NestedField, PrimitiveType, Schema, SchemaRef, Type};
 
@@ -397,9 +422,11 @@ mod tests {
         let schema = table_schema_simple();
         let reference = Reference::new("bar").bind(schema, true).unwrap();
 
+        let accessor_ref = Arc::new(StructAccessor::new(1, Type::Primitive(PrimitiveType::Int)));
         let expected_ref = BoundReference::new(
             "bar",
             NestedField::required(2, "bar", Type::Primitive(PrimitiveType::Int)).into(),
+            accessor_ref.clone(),
         );
 
         assert_eq!(expected_ref, reference);
@@ -410,9 +437,11 @@ mod tests {
         let schema = table_schema_simple();
         let reference = Reference::new("BAR").bind(schema, false).unwrap();
 
+        let accessor_ref = Arc::new(StructAccessor::new(1, Type::Primitive(PrimitiveType::Int)));
         let expected_ref = BoundReference::new(
             "BAR",
             NestedField::required(2, "bar", Type::Primitive(PrimitiveType::Int)).into(),
+            accessor_ref.clone(),
         );
 
         assert_eq!(expected_ref, reference);

crates/iceberg/src/spec/schema.rs

@@ -18,6 +18,7 @@
 //! This module defines schema in iceberg.
 
 use crate::error::Result;
+use crate::expr::accessor::StructAccessor;
 use crate::spec::datatypes::{
     ListType, MapType, NestedFieldRef, PrimitiveType, StructType, Type, LIST_FILED_NAME,
     MAP_KEY_FIELD_NAME, MAP_VALUE_FIELD_NAME,
@@ -55,6 +56,8 @@ pub struct Schema {
     name_to_id: HashMap<String, i32>,
     lowercase_name_to_id: HashMap<String, i32>,
     id_to_name: HashMap<i32, String>,
+
+    field_id_to_accessor: HashMap<i32, Arc<StructAccessor>>,
 }
 
 impl PartialEq for Schema {
@@ -105,6 +108,8 @@ impl SchemaBuilder {
     pub fn build(self) -> Result<Schema> {
         let highest_field_id = self.fields.iter().map(|f| f.id).max().unwrap_or(0);
 
+        let field_id_to_accessor = self.build_accessors();
+
         let r#struct = StructType::new(self.fields);
         let id_to_field = index_by_id(&r#struct)?;
 
@@ -137,9 +142,55 @@ impl SchemaBuilder {
             name_to_id,
             lowercase_name_to_id,
             id_to_name,
+
+            field_id_to_accessor,
         })
     }
 
+    fn build_accessors(&self) -> HashMap<i32, Arc<StructAccessor>> {
+        let mut map = HashMap::new();
+
+        for (pos, field) in self.fields.iter().enumerate() {
+            // add an accessor for this field
+
+            let accessor = Arc::new(StructAccessor::new(pos as i32, *field.field_type.clone()));
+            map.insert(field.id, accessor.clone());
+
+            if let Type::Struct(nested) = field.field_type.as_ref() {
+                // add accessors for nested fields
+                for (field_id, accessor) in Self::build_accessors_nested(nested.fields()) {
+                    let new_accessor = Arc::new(StructAccessor::wrap(pos as i32, accessor));
+                    map.insert(field_id, new_accessor.clone());
+                }
+            }
+        }
+
+        map
+    }
+
+    fn build_accessors_nested(fields: &[NestedFieldRef]) -> Vec<(i32, Arc<StructAccessor>)> {
+        let mut results = vec![];
+        for (pos, field) in fields.iter().enumerate() {
+            if let Type::Struct(nested) = field.field_type.as_ref() {
+                let nested_accessors = Self::build_accessors_nested(nested.fields());
+
+                let wrapped_nested_accessors =
+                    nested_accessors.into_iter().map(|(id, accessor)| {
+                        let new_accessor = Arc::new(StructAccessor::wrap(pos as i32, accessor));
+                        (id, new_accessor.clone())
+                    });
+
+                results.extend(wrapped_nested_accessors);
+            }
+
+            let accessor = Arc::new(StructAccessor::new(pos as i32, *field.field_type.clone()));
+
+            results.push((field.id, accessor.clone()));
+        }
+
+        results
+    }
+
     fn validate_identifier_ids(
         r#struct: &StructType,
         id_to_field: &HashMap<i32, NestedFieldRef>,
@@ -264,6 +315,11 @@ impl Schema {
     pub fn name_by_field_id(&self, field_id: i32) -> Option<&str> {
         self.id_to_name.get(&field_id).map(String::as_str)
     }
+
+    /// Get an accessor for retrieving data in a struct
+    pub fn accessor_for_field_id(&self, field_id: i32) -> Option<Arc<StructAccessor>> {
+        self.field_id_to_accessor.get(&field_id).cloned()
+    }
 }
 
 impl Display for Schema {
@@ -381,7 +437,7 @@ pub fn visit_schema<V: SchemaVisitor>(schema: &Schema, visitor: &mut V) -> Resul
     visitor.schema(schema, result)
 }
 
-/// Creates an field id to field map.
+/// Creates a field id to field map.
 pub fn index_by_id(r#struct: &StructType) -> Result<HashMap<i32, NestedFieldRef>> {
     struct IndexById(HashMap<i32, NestedFieldRef>);
 

crates/iceberg/src/spec/values.rs

@@ -1141,6 +1141,11 @@ impl Struct {
             },
         )
     }
+
+    /// Gets a ref to the field at `position` within the `Struct`
+    pub fn get(&self, position: i32) -> &Literal {
+        &self.fields[position as usize]
+    }
 }
 
 /// An iterator that moves out of a struct.