substrait_expr_macros/
lib.rs

1use proc_macro::TokenStream;
2use quote::quote;
3use syn::parse::{Parse, ParseStream};
4
5// Convert rust code to a field in a names-only schema
6// Example Input:
7//  a: { b: {}, c: { d: {} } }
8//
9// Example Output:
10//  NamesOnlySchemaNode { name: "a".to_string(), children: vec![
11//    NamesOnlySchemaNode { name: "b".to_string(), children: vec![] },
12//    NamesOnlySchemaNode { name: "c".to_string(), children: vec![
13//      NamesOnlySchemaNode { name: "d".to_string(), children: vec![] }
14//    ]}
15//  ]}
16fn rust_field_to_names_field(field: &Field) -> proc_macro2::TokenStream {
17    let name = field.name.to_string();
18    let children = rust_to_names_fields(&field.ty);
19    quote! {substrait_expr::helpers::schema::NamesOnlySchemaNode {
20        name: #name.to_string(),
21        children: #children,
22    }}
23}
24
25// Convert rust code to a vector of NamesOnlySchemaNode
26//
27// Example input:
28//  { a: { b: {}, c: { d: {} } } }
29//
30// Example Output:
31// vec![
32//   NamesOnlySchemaNode { name: "a".to_string(), children: vec![
33//     NamesOnlySchemaNode { name: "b".to_string(), children: vec![] },
34//     NamesOnlySchemaNode { name: "c".to_string(), children: vec![
35//       NamesOnlySchemaNode { name: "d".to_string(), children: vec![] }
36//     ]}
37//   ]}
38// ]
39fn rust_to_names_fields(schema: &NestedType) -> proc_macro2::TokenStream {
40    let parsed_fields = schema
41        .fields
42        .iter()
43        .map(|field| rust_field_to_names_field(field))
44        .collect::<Vec<_>>();
45    quote! {vec![#(#parsed_fields),*]}
46}
47
48// Convert rust code to a root NamesOnlySchemaNode (that has the empty string for a name)
49//
50// Example input:
51//  { a: { b: {}, c: { d: {} } } }
52//
53// Example Output:
54// NamesOnlySchemaNode {
55//   name: "".to_string(),
56//   vec![
57//     NamesOnlySchemaNode { name: "a".to_string(), children: vec![
58//       NamesOnlySchemaNode { name: "b".to_string(), children: vec![] },
59//       NamesOnlySchemaNode { name: "c".to_string(), children: vec![
60//         NamesOnlySchemaNode { name: "d".to_string(), children: vec![] }
61//       ]}
62//     ]}
63//   ]
64// }
65fn rust_to_names_schema(schema: &NestedType) -> proc_macro2::TokenStream {
66    let children = rust_to_names_fields(schema);
67    quote! {
68        substrait_expr::helpers::schema::SchemaInfo::Names(substrait_expr::helpers::schema::NamesOnlySchema::new(#children))
69    }
70}
71
72// New rust syntax for a field in a names only schema
73//
74// Examples:
75//  foo: {}
76//  blah: { x: {}, y: { z: {} } }
77struct Field {
78    name: syn::Ident,
79    _colon_token: syn::Token![:],
80    ty: NestedType,
81}
82
83impl Parse for Field {
84    fn parse(input: ParseStream) -> syn::Result<Self> {
85        Ok(Field {
86            name: input.parse()?,
87            _colon_token: input.parse()?,
88            ty: input.parse()?,
89        })
90    }
91}
92
93// New rust syntax for a nested names-only type ({field, field, field})
94//
95// Examples:
96//  { foo: {} }
97//  { blah: { x: {}, y: { z: {} } } }
98struct NestedType {
99    _brace_token: syn::token::Brace,
100    fields: syn::punctuated::Punctuated<Field, syn::Token![,]>,
101}
102
103impl Parse for NestedType {
104    fn parse(input: ParseStream) -> syn::Result<Self> {
105        let content;
106        Ok(Self {
107            _brace_token: syn::braced!(content in input),
108            fields: content.parse_terminated(Field::parse, syn::Token![,])?,
109        })
110    }
111}
112
113fn names_schema2(input: proc_macro2::TokenStream) -> syn::Result<proc_macro2::TokenStream> {
114    let struct_type: NestedType = syn::parse2(input)?;
115    Ok(rust_to_names_schema(&struct_type))
116}
117
118/// A macro to create names-only schemas from a dictionary-like rust syntax
119///
120/// # Examples
121/// ```ignore
122/// use substrait_expr::macros::names_schema;
123///
124/// let schema = names_schema!({
125///   vector: {},
126///   metadata: {
127///     caption: {},
128///     user_score: {}
129///  }
130/// });
131/// ```
132#[proc_macro]
133pub fn names_schema(input: TokenStream) -> TokenStream {
134    let input = proc_macro2::TokenStream::from(input);
135    names_schema2(input).unwrap().into()
136}
substrait_expr_macros/lib.rs

substrait_expr_macros/
lib.rs
