datafusion_substrait/logical_plan/consumer/expr/
nested.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 crate::logical_plan::consumer::SubstraitConsumer;
19use datafusion::common::{DFSchema, not_impl_err, substrait_err};
20use datafusion::execution::FunctionRegistry;
21use datafusion::logical_expr::Expr;
22use substrait::proto::expression::Nested;
23use substrait::proto::expression::nested::NestedType;
24
25/// Converts a Substrait [Nested] expression into a DataFusion [Expr].
26///
27/// Substrait Nested expressions represent complex type constructors (list, struct, map)
28/// where elements are full expressions rather than just literals. This is used by
29/// producers that emit `Nested { list: ... }` for array construction, as opposed to
30/// `Literal { list: ... }` which only supports scalar values.
31pub async fn from_nested(
32    consumer: &impl SubstraitConsumer,
33    nested: &Nested,
34    input_schema: &DFSchema,
35) -> datafusion::common::Result<Expr> {
36    let Some(nested_type) = &nested.nested_type else {
37        return substrait_err!("Nested expression requires a nested_type");
38    };
39
40    match nested_type {
41        NestedType::List(list) => {
42            if list.values.is_empty() {
43                return substrait_err!(
44                    "Empty Nested lists are not supported; use Literal.empty_list instead"
45                );
46            }
47
48            let mut args = Vec::with_capacity(list.values.len());
49            for value in &list.values {
50                args.push(consumer.consume_expression(value, input_schema).await?);
51            }
52
53            let make_array_udf = consumer.get_function_registry().udf("make_array")?;
54            Ok(Expr::ScalarFunction(
55                datafusion::logical_expr::expr::ScalarFunction::new_udf(
56                    make_array_udf,
57                    args,
58                ),
59            ))
60        }
61        NestedType::Struct(_) => {
62            not_impl_err!("Nested struct expressions are not yet supported")
63        }
64        NestedType::Map(_) => {
65            not_impl_err!("Nested map expressions are not yet supported")
66        }
67    }
68}
69
70#[cfg(test)]
71mod tests {
72    use super::*;
73    use crate::logical_plan::consumer::utils::tests::test_consumer;
74    use substrait::proto::expression::Literal;
75    use substrait::proto::expression::nested::List;
76    use substrait::proto::{self, Expression};
77
78    fn make_i64_literal(value: i64) -> Expression {
79        Expression {
80            rex_type: Some(proto::expression::RexType::Literal(Literal {
81                nullable: false,
82                type_variation_reference: 0,
83                literal_type: Some(proto::expression::literal::LiteralType::I64(value)),
84            })),
85        }
86    }
87
88    #[tokio::test]
89    async fn nested_list_with_literals() -> datafusion::common::Result<()> {
90        let consumer = test_consumer();
91        let schema = DFSchema::empty();
92        let nested = Nested {
93            nullable: false,
94            type_variation_reference: 0,
95            nested_type: Some(NestedType::List(List {
96                values: vec![
97                    make_i64_literal(1),
98                    make_i64_literal(2),
99                    make_i64_literal(3),
100                ],
101            })),
102        };
103
104        let expr = from_nested(&consumer, &nested, &schema).await?;
105        assert_eq!(
106            format!("{expr}"),
107            "make_array(Int64(1), Int64(2), Int64(3))"
108        );
109
110        Ok(())
111    }
112
113    #[tokio::test]
114    async fn nested_list_empty_rejected() -> datafusion::common::Result<()> {
115        let consumer = test_consumer();
116        let schema = DFSchema::empty();
117        let nested = Nested {
118            nullable: true,
119            type_variation_reference: 0,
120            nested_type: Some(NestedType::List(List { values: vec![] })),
121        };
122
123        let result = from_nested(&consumer, &nested, &schema).await;
124        assert!(result.is_err());
125        assert!(
126            result
127                .unwrap_err()
128                .to_string()
129                .contains("Empty Nested lists are not supported")
130        );
131
132        Ok(())
133    }
134
135    #[tokio::test]
136    async fn nested_missing_type() -> datafusion::common::Result<()> {
137        let consumer = test_consumer();
138        let schema = DFSchema::empty();
139        let nested = Nested {
140            nullable: false,
141            type_variation_reference: 0,
142            nested_type: None,
143        };
144
145        let result = from_nested(&consumer, &nested, &schema).await;
146        assert!(result.is_err());
147        assert!(result.unwrap_err().to_string().contains("nested_type"));
148
149        Ok(())
150    }
151}
datafusion_substrait/logical_plan/consumer/expr/nested.rs

datafusion_substrait/logical_plan/consumer/expr/
nested.rs
