feat: add Struct Accessors to BoundReferences

sdd · sdd · commit b531d52cb887 · 2024-04-06T08:19:36.000+01:00
diff --git a/crates/iceberg/src/expr/accessor.rs b/crates/iceberg/src/expr/accessor.rs
@@ -0,0 +1,120 @@
+// 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::{Datum, Literal, PrimitiveType, Struct};
+use crate::{Error, ErrorKind};
+use serde_derive::{Deserialize, Serialize};
+use std::sync::Arc;
+
+#[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Eq)]
+pub struct StructAccessor {
+    position: usize,
+    r#type: PrimitiveType,
+    inner: Option<Box<StructAccessor>>,
+}
+
+pub(crate) type StructAccessorRef = Arc<StructAccessor>;
+
+impl StructAccessor {
+    pub(crate) fn new(position: usize, r#type: PrimitiveType) -> Self {
+        StructAccessor {
+            position,
+            r#type,
+            inner: None,
+        }
+    }
+
+    pub(crate) fn wrap(position: usize, inner: Box<StructAccessor>) -> Self {
+        StructAccessor {
+            position,
+            r#type: inner.r#type().clone(),
+            inner: Some(inner),
+        }
+    }
+
+    pub(crate) fn position(&self) -> usize {
+        self.position
+    }
+
+    pub(crate) fn r#type(&self) -> &PrimitiveType {
+        &self.r#type
+    }
+
+    pub(crate) fn get<'a>(&'a self, container: &'a Struct) -> crate::Result<Datum> {
+        match &self.inner {
+            None => {
+                if let Literal::Primitive(literal) = &container[self.position] {
+                    Ok(Datum::new(self.r#type().clone(), literal.clone()))
+                } else {
+                    Err(Error::new(
+                        ErrorKind::Unexpected,
+                        "Expected Literal to be Primitive",
+                    ))
+                }
+            }
+            Some(inner) => {
+                if let Literal::Struct(wrapped) = &container[self.position] {
+                    inner.get(wrapped)
+                } else {
+                    Err(Error::new(
+                        ErrorKind::Unexpected,
+                        "Nested accessor should only be wrapping a Struct",
+                    ))
+                }
+            }
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::expr::accessor::StructAccessor;
+    use crate::spec::{Datum, Literal, PrimitiveType, Struct};
+
+    #[test]
+    fn test_single_level_accessor() {
+        let accessor = StructAccessor::new(1, PrimitiveType::Boolean);
+
+        assert_eq!(accessor.r#type(), &PrimitiveType::Boolean);
+        assert_eq!(accessor.position(), 1);
+
+        let test_struct =
+            Struct::from_iter(vec![Some(Literal::bool(false)), Some(Literal::bool(true))]);
+
+        assert_eq!(accessor.get(&test_struct).unwrap(), Datum::bool(true));
+    }
+
+    #[test]
+    fn test_nested_accessor() {
+        let nested_accessor = StructAccessor::new(1, PrimitiveType::Boolean);
+        let accessor = StructAccessor::wrap(2, Box::new(nested_accessor));
+
+        assert_eq!(accessor.r#type(), &PrimitiveType::Boolean);
+        //assert_eq!(accessor.position(), 1);
+
+        let nested_test_struct =
+            Struct::from_iter(vec![Some(Literal::bool(false)), Some(Literal::bool(true))]);
+
+        let test_struct = Struct::from_iter(vec![
+            Some(Literal::bool(false)),
+            Some(Literal::bool(false)),
+            Some(Literal::Struct(nested_test_struct)),
+        ]);
+
+        assert_eq!(accessor.get(&test_struct).unwrap(), Datum::bool(true));
+    }
+}
diff --git a/crates/iceberg/src/expr/mod.rs b/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;
diff --git a/crates/iceberg/src/expr/term.rs b/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, 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, 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);
diff --git a/crates/iceberg/src/spec/schema.rs b/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,61 @@ 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() {
+            if let Type::Primitive(prim_type) = field.field_type.as_ref() {
+                // add an accessor for this field
+                let accessor = Arc::new(StructAccessor::new(pos, prim_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, accessor));
+                        map.insert(field_id, new_accessor.clone());
+                    }
+                }
+            }
+        }
+
+        map
+    }
+
+    fn build_accessors_nested(fields: &[NestedFieldRef]) -> Vec<(i32, Box<StructAccessor>)> {
+        let mut results = vec![];
+        for (pos, field) in fields.iter().enumerate() {
+            match field.field_type.as_ref() {
+                Type::Primitive(prim_type) => {
+                    let accessor = Box::new(StructAccessor::new(pos, prim_type.clone()));
+                    results.push((field.id, accessor));
+                }
+                Type::Struct(nested) => {
+                    let nested_accessors = Self::build_accessors_nested(nested.fields());
+
+                    let wrapped_nested_accessors =
+                        nested_accessors.into_iter().map(|(id, accessor)| {
+                            let new_accessor = Box::new(StructAccessor::wrap(pos, accessor));
+                            (id, new_accessor.clone())
+                        });
+
+                    results.extend(wrapped_nested_accessors);
+                }
+                _ => {
+                    todo!()
+                }
+            }
+        }
+
+        results
+    }
+
     fn validate_identifier_ids(
         r#struct: &StructType,
         id_to_field: &HashMap<i32, NestedFieldRef>,
@@ -264,6 +321,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 +443,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>);
 
diff --git a/crates/iceberg/src/spec/values.rs b/crates/iceberg/src/spec/values.rs
@@ -20,6 +20,7 @@
  */
 
 use std::fmt::{Display, Formatter};
+use std::ops::Index;
 use std::str::FromStr;
 use std::{any::Any, collections::BTreeMap};
 
@@ -141,6 +142,11 @@ impl From<Datum> for Literal {
 }
 
 impl Datum {
+    /// Creates a `Datum` from a `PrimitiveType` and a `PrimitiveLiteral`
+    pub(crate) fn new(r#type: PrimitiveType, literal: PrimitiveLiteral) -> Self {
+        Datum { r#type, literal }
+    }
+
     /// Creates a boolean value.
     ///
     /// Example:
@@ -1143,6 +1149,14 @@ impl Struct {
     }
 }
 
+impl Index<usize> for Struct {
+    type Output = Literal;
+
+    fn index(&self, idx: usize) -> &Self::Output {
+        &self.fields[idx]
+    }
+}
+
 /// An iterator that moves out of a struct.
 pub struct StructValueIntoIter {
     null_bitmap: bitvec::boxed::IntoIter,
diff --git a/crates/iceberg/src/transform/mod.rs b/crates/iceberg/src/transform/mod.rs
