switch to using "project" to get schemas

Author: nicklan

derive-macros/Cargo.toml

@@ -14,7 +14,7 @@ proc-macro = true
 
 [dependencies]
 proc-macro2 = "1"
-syn = "2.0"
+syn = { version = "2.0", features = ["extra-traits"] }
 quote = "1.0"
 
 

derive-macros/src/lib.rs

@@ -1,66 +1,15 @@
-use proc_macro2::{Ident, Spacing, TokenStream, TokenTree};
+use proc_macro2::{Ident, TokenStream};
 use quote::{quote, quote_spanned};
 use syn::spanned::Spanned;
-use syn::{
-    parse_macro_input, Attribute, Data, DataStruct, DeriveInput, Fields, Meta, PathArguments, Type,
-};
-
-static SCHEMA_ERR_STR: &str = "schema(...) only supports schema(name = name)";
-
-// Return the ident to use as the schema name if it's been specified in the attributes of the struct
-fn get_schema_name_from_attr<'a>(attrs: impl Iterator<Item = &'a Attribute>) -> Option<Ident> {
-    for attr in attrs {
-        if let Meta::List(list) = &attr.meta {
-            if let Some(attr_name) = list.path.segments.iter().last() {
-                if attr_name.ident == "schema" {
-                    // We have some schema(...) attribute, see if we've specified a different name
-                    let tokens: Vec<TokenTree> = list.tokens.clone().into_iter().collect();
-                    match tokens[..] {
-                        // we only support `name = name` style
-                        [TokenTree::Ident(ref name_ident), TokenTree::Punct(ref punct), TokenTree::Ident(ref schema_ident)] =>
-                        {
-                            assert!(name_ident == "name", "{}", SCHEMA_ERR_STR);
-                            assert!(punct.as_char() == '=', "{}", SCHEMA_ERR_STR);
-                            assert!(punct.spacing() == Spacing::Alone, "{}", SCHEMA_ERR_STR);
-                            return Some(schema_ident.clone());
-                        }
-                        _ => panic!("{}", SCHEMA_ERR_STR),
-                    }
-                } else {
-                    panic!("Schema only accepts `schema` as an extra attribute")
-                }
-            }
-        }
-    }
-    None
-}
+use syn::{parse_macro_input, Data, DataStruct, DeriveInput, Fields, PathArguments, Type};
 
 #[proc_macro_derive(Schema, attributes(schema))]
 pub fn derive_schema(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
     let input = parse_macro_input!(input as DeriveInput);
     let struct_ident = input.ident;
-    let schema_name = get_schema_name_from_attr(input.attrs.iter()).unwrap_or_else(|| {
-        // default to the struct name, but lowercased
-        Ident::new(
-            &struct_ident.to_string().to_lowercase(),
-            struct_ident.span(),
-        )
-    });
 
     let schema_fields = gen_schema_fields(&input.data);
     let output = quote! {
-        impl crate::actions::schemas::GetSchema for #struct_ident {
-            fn get_schema() -> crate::schema::SchemaRef {
-                use crate::actions::schemas::GetField;
-                static SCHEMA_LOCK: std::sync::OnceLock<crate::schema::SchemaRef> = std::sync::OnceLock::new();
-                SCHEMA_LOCK.get_or_init(|| {
-                    std::sync::Arc::new(crate::schema::StructType::new(vec![
-                        Self::get_field(stringify!(#schema_name))
-                    ]))
-                }).clone() // cheap clone, it's an Arc
-            }
-        }
-
         impl crate::actions::schemas::GetField for #struct_ident {
             fn get_field(name: impl Into<String>) -> crate::schema::StructField {
                 use crate::actions::schemas::GetField;

kernel/src/actions/mod.rs

@@ -4,12 +4,39 @@ pub(crate) mod schemas;
 pub(crate) mod visitors;
 
 use derive_macros::Schema;
-use std::collections::HashMap;
+use lazy_static::lazy_static;
 use visitors::{AddVisitor, MetadataVisitor, ProtocolVisitor};
 
+use self::deletion_vector::DeletionVectorDescriptor;
+use crate::actions::schemas::GetField;
 use crate::{schema::StructType, DeltaResult, EngineData};
 
-use self::{deletion_vector::DeletionVectorDescriptor, schemas::GetSchema};
+use std::collections::HashMap;
+
+lazy_static! {
+    static ref LOG_SCHEMA: StructType = StructType::new(
+        vec![
+            Option::<Add>::get_field("add"),
+            Option::<Remove>::get_field("remove"),
+            Option::<Metadata>::get_field("metaData"),
+            Option::<Protocol>::get_field("protocol"),
+            // We don't support the following actions yet
+            //Option<Cdc>::get_field("cdc"),
+            //Option<CommitInfo>::get_field("commitInfo"),
+            //Option<DomainMetadata>::get_field("domainMetadata"),
+            //Option<Transaction>::get_field("txn"),
+        ]
+    );
+}
+
+pub(crate) static ADD_NAME: &str = "add";
+pub(crate) static REMOVE_NAME: &str = "remove";
+pub(crate) static METADATA_NAME: &str = "metaData";
+pub(crate) static PROTOCOL_NAME: &str = "protocol";
+
+pub(crate) fn get_log_schema() -> &'static StructType {
+    &LOG_SCHEMA
+}
 
 #[derive(Debug, Clone, PartialEq, Eq, Schema)]
 pub struct Format {
@@ -52,7 +79,10 @@ pub struct Metadata {
 impl Metadata {
     pub fn try_new_from_data(data: &dyn EngineData) -> DeltaResult<Option<Metadata>> {
         let mut visitor = MetadataVisitor::default();
-        data.extract(Metadata::get_schema(), &mut visitor)?;
+        data.extract(
+            get_log_schema().project_as_schema(&[METADATA_NAME])?,
+            &mut visitor,
+        )?;
         Ok(visitor.metadata)
     }
 
@@ -80,7 +110,10 @@ pub struct Protocol {
 impl Protocol {
     pub fn try_new_from_data(data: &dyn EngineData) -> DeltaResult<Option<Protocol>> {
         let mut visitor = ProtocolVisitor::default();
-        data.extract(Protocol::get_schema(), &mut visitor)?;
+        data.extract(
+            get_log_schema().project_as_schema(&[PROTOCOL_NAME])?,
+            &mut visitor,
+        )?;
         Ok(visitor.protocol)
     }
 }
@@ -134,7 +167,10 @@ impl Add {
     /// Since we always want to parse multiple adds from data, we return a `Vec<Add>`
     pub fn parse_from_data(data: &dyn EngineData) -> DeltaResult<Vec<Add>> {
         let mut visitor = AddVisitor::default();
-        data.extract(Add::get_schema(), &mut visitor)?;
+        data.extract(
+            get_log_schema().project_as_schema(&[ADD_NAME])?,
+            &mut visitor,
+        )?;
         Ok(visitor.adds)
     }
 
@@ -189,43 +225,18 @@ impl Remove {
     }
 }
 
-use crate::actions::schemas::GetField;
-use lazy_static::lazy_static;
-
-lazy_static! {
-    static ref LOG_SCHEMA: StructType = StructType::new(
-        vec![
-            Option::<Add>::get_field("add"),
-            Option::<Remove>::get_field("remove"),
-            Option::<Metadata>::get_field("metaData"),
-            Option::<Protocol>::get_field("protocol"),
-            // We don't support the following actions yet
-            //Option<Cdc>::get_field("cdc"),
-            //Option<CommitInfo>::get_field("commitInfo"),
-            //Option<DomainMetadata>::get_field("domainMetadata"),
-            //Option<Transaction>::get_field("txn"),
-        ]
-    );
-}
-
-#[cfg(test)]
-pub(crate) fn get_log_schema() -> &'static StructType {
-    &LOG_SCHEMA
-}
-
 #[cfg(test)]
 mod tests {
     use std::sync::Arc;
 
     use super::*;
-    use crate::{
-        actions::schemas::GetSchema,
-        schema::{ArrayType, DataType, MapType, StructField},
-    };
+    use crate::schema::{ArrayType, DataType, MapType, StructField};
 
     #[test]
     fn test_metadata_schema() {
-        let schema = Metadata::get_schema();
+        let schema = get_log_schema()
+            .project_as_schema(&["metaData"])
+            .expect("Couldn't get metaData field");
 
         let expected = Arc::new(StructType::new(vec![StructField::new(
             "metaData",
@@ -258,7 +269,7 @@ mod tests {
                     false,
                 ),
             ]),
-            false,
+            true,
         )]));
         assert_eq!(schema, expected);
     }
@@ -295,7 +306,9 @@ mod tests {
 
     #[test]
     fn test_remove_schema() {
-        let schema = Remove::get_schema();
+        let schema = get_log_schema()
+            .project_as_schema(&["remove"])
+            .expect("Couldn't get remove field");
         let expected = Arc::new(StructType::new(vec![StructField::new(
             "remove",
             StructType::new(vec![
@@ -310,7 +323,7 @@ mod tests {
                 StructField::new("baseRowId", DataType::LONG, true),
                 StructField::new("defaultRowCommitVersion", DataType::LONG, true),
             ]),
-            false,
+            true,
         )]));
         assert_eq!(schema, expected);
     }

kernel/src/actions/schemas.rs

@@ -2,12 +2,7 @@
 
 use std::collections::HashMap;
 
-use crate::schema::{ArrayType, DataType, MapType, SchemaRef, StructField};
-
-/// A trait that says you can ask for the [`Schema`] of the implementor
-pub(crate) trait GetSchema {
-    fn get_schema() -> SchemaRef;
-}
+use crate::schema::{ArrayType, DataType, MapType, StructField};
 
 /// A trait that allows getting a `StructField` based on the provided name and nullability
 pub(crate) trait GetField {

kernel/src/actions/visitors.rs

@@ -262,8 +262,7 @@ mod tests {
 
     use super::*;
     use crate::{
-        actions::get_log_schema,
-        actions::schemas::GetSchema,
+        actions::{get_log_schema, ADD_NAME},
         simple_client::{data::SimpleData, json::SimpleJsonHandler, SimpleClient},
         EngineData, EngineInterface, JsonHandler,
     };
@@ -356,7 +355,9 @@ mod tests {
         let batch = json_handler
             .parse_json(string_array_to_engine_data(json_strings), output_schema)
             .unwrap();
-        let add_schema = Add::get_schema();
+        let add_schema = get_log_schema()
+            .project_as_schema(&[ADD_NAME])
+            .expect("Can't get add schema");
         let mut add_visitor = AddVisitor::default();
         batch.extract(add_schema, &mut add_visitor).unwrap();
         let add1 = Add {

kernel/src/scan/file_stream.rs

@@ -1,15 +1,14 @@
 use std::collections::HashSet;
-use std::sync::Arc;
 
 use either::Either;
 use tracing::debug;
 
 use super::data_skipping::DataSkippingFilter;
-use crate::actions::schemas::{GetField, GetSchema};
+use crate::actions::{get_log_schema, ADD_NAME, REMOVE_NAME};
 use crate::actions::{visitors::AddVisitor, visitors::RemoveVisitor, Add, Remove};
 use crate::engine_data::{GetData, TypedGetData};
 use crate::expressions::Expression;
-use crate::schema::{SchemaRef, StructType};
+use crate::schema::SchemaRef;
 use crate::{DataVisitor, DeltaResult, EngineData, EngineInterface};
 
 struct LogReplayScanner {
@@ -82,14 +81,11 @@ impl LogReplayScanner {
         };
 
         let schema_to_use = if is_log_batch {
-            Arc::new(StructType::new(vec![
-                Option::<Add>::get_field("add"),
-                Option::<Remove>::get_field("remove"),
-            ]))
+            get_log_schema().project_as_schema(&[ADD_NAME, REMOVE_NAME])?
         } else {
             // All checkpoint actions are already reconciled and Remove actions in checkpoint files
             // only serve as tombstones for vacuum jobs. So no need to load them here.
-            Add::get_schema()
+            get_log_schema().project_as_schema(&[ADD_NAME])?
         };
         let mut visitor = AddRemoveVisitor::default();
         actions.extract(schema_to_use, &mut visitor)?;

kernel/src/scan/mod.rs

@@ -3,8 +3,7 @@ use std::sync::Arc;
 use itertools::Itertools;
 
 use self::file_stream::log_replay_iter;
-use crate::actions::schemas::GetField;
-use crate::actions::{Add, Remove};
+use crate::actions::{get_log_schema, Add, ADD_NAME, REMOVE_NAME};
 use crate::expressions::{Expression, Scalar};
 use crate::schema::{DataType, SchemaRef, StructType};
 use crate::snapshot::Snapshot;
@@ -129,10 +128,7 @@ impl Scan {
         &self,
         engine_interface: &dyn EngineInterface,
     ) -> DeltaResult<impl Iterator<Item = DeltaResult<Add>>> {
-        let action_schema = Arc::new(StructType::new(vec![
-            Option::<Add>::get_field("add"),
-            Option::<Remove>::get_field("remove"),
-        ]));
+        let action_schema = get_log_schema().project_as_schema(&[ADD_NAME, REMOVE_NAME])?;
 
         let log_iter = self.snapshot.log_segment.replay(
             engine_interface,

kernel/src/schema.rs

@@ -3,8 +3,11 @@ use std::sync::Arc;
 use std::{collections::HashMap, fmt::Display};
 
 use indexmap::IndexMap;
+use itertools::Itertools;
 use serde::{Deserialize, Serialize};
 
+use crate::{DeltaResult, Error};
+
 pub type Schema = StructType;
 pub type SchemaRef = Arc<StructType>;
 
@@ -140,6 +143,38 @@ impl StructType {
         }
     }
 
+    /// Get a [`StructType`] containing [`StructField`]s of the given names, preserving the original
+    /// order of fields. Returns an Err if a specified field doesn't exist
+    pub fn project(&self, names: &[impl AsRef<str>]) -> DeltaResult<StructType> {
+        let mut indexes: Vec<usize> = names
+            .iter()
+            .map(|name| {
+                self.fields
+                    .get_index_of(name.as_ref())
+                    .ok_or_else(|| Error::missing_column(name.as_ref()))
+            })
+            .try_collect()?;
+        indexes.sort(); // keep schema order
+        let fields: Vec<StructField> = indexes
+            .iter()
+            .map(|index| {
+                self.fields
+                    .get_index(*index)
+                    .expect("get_index_of returned non-existant index")
+                    .1
+                    .clone()
+            })
+            .collect();
+        Ok(Self::new(fields))
+    }
+
+    /// Get a [`SchemaRef`] containing [`StructField`]s of the given names, preserving the original
+    /// order of fields. Returns an Err if a specified field doesn't exist
+    pub fn project_as_schema(&self, names: &[impl AsRef<str>]) -> DeltaResult<SchemaRef> {
+        let struct_type = self.project(names)?;
+        Ok(Arc::new(struct_type))
+    }
+
     pub fn field(&self, name: impl AsRef<str>) -> Option<&StructField> {
         self.fields.get(name.as_ref())
     }