datafusion_functions/core/
planner.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 arrow::datatypes::Field;
19use datafusion_common::Result;
20use datafusion_common::{Column, DFSchema, ScalarValue, TableReference};
21use datafusion_expr::expr::ScalarFunction;
22use datafusion_expr::planner::{ExprPlanner, PlannerResult, RawDictionaryExpr};
23use datafusion_expr::{lit, Expr};
24
25use super::named_struct;
26
27#[derive(Default, Debug)]
28pub struct CoreFunctionPlanner {}
29
30impl ExprPlanner for CoreFunctionPlanner {
31    fn plan_dictionary_literal(
32        &self,
33        expr: RawDictionaryExpr,
34        _schema: &DFSchema,
35    ) -> Result<PlannerResult<RawDictionaryExpr>> {
36        let mut args = vec![];
37        for (k, v) in expr.keys.into_iter().zip(expr.values.into_iter()) {
38            args.push(k);
39            args.push(v);
40        }
41        Ok(PlannerResult::Planned(named_struct().call(args)))
42    }
43
44    fn plan_struct_literal(
45        &self,
46        args: Vec<Expr>,
47        is_named_struct: bool,
48    ) -> Result<PlannerResult<Vec<Expr>>> {
49        Ok(PlannerResult::Planned(Expr::ScalarFunction(
50            ScalarFunction::new_udf(
51                if is_named_struct {
52                    named_struct()
53                } else {
54                    crate::core::r#struct()
55                },
56                args,
57            ),
58        )))
59    }
60
61    fn plan_overlay(&self, args: Vec<Expr>) -> Result<PlannerResult<Vec<Expr>>> {
62        Ok(PlannerResult::Planned(Expr::ScalarFunction(
63            ScalarFunction::new_udf(crate::string::overlay(), args),
64        )))
65    }
66
67    fn plan_compound_identifier(
68        &self,
69        field: &Field,
70        qualifier: Option<&TableReference>,
71        nested_names: &[String],
72    ) -> Result<PlannerResult<Vec<Expr>>> {
73        let col = Expr::Column(Column::from((qualifier, field)));
74
75        // Start with the base column expression
76        let mut expr = col;
77
78        // Iterate over nested_names and create nested get_field expressions
79        for nested_name in nested_names {
80            let get_field_args = vec![expr, lit(ScalarValue::from(nested_name.clone()))];
81            expr = Expr::ScalarFunction(ScalarFunction::new_udf(
82                crate::core::get_field(),
83                get_field_args,
84            ));
85        }
86
87        Ok(PlannerResult::Planned(expr))
88    }
89}
datafusion_functions/core/planner.rs

datafusion_functions/core/
planner.rs
