datafusion_physical_expr/expressions/
lambda_variable.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
18//! Physical lambda variable reference: [`LambdaVariable`]
19
20use std::hash::Hash;
21use std::sync::Arc;
22
23use crate::physical_expr::PhysicalExpr;
24use arrow::datatypes::FieldRef;
25use arrow::{
26    datatypes::{DataType, Schema},
27    record_batch::RecordBatch,
28};
29
30use datafusion_common::{Result, exec_err, internal_err};
31use datafusion_expr::ColumnarValue;
32
33/// Represents the lambda variable with a given index and field
34#[derive(Debug, Clone)]
35pub struct LambdaVariable {
36    index: usize,
37    field: FieldRef,
38}
39
40impl Eq for LambdaVariable {}
41
42impl PartialEq for LambdaVariable {
43    fn eq(&self, other: &Self) -> bool {
44        self.index == other.index && self.field == other.field
45    }
46}
47
48impl Hash for LambdaVariable {
49    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
50        self.index.hash(state);
51        self.field.hash(state);
52    }
53}
54
55impl LambdaVariable {
56    /// Create a new lambda variable expression
57    pub fn new(index: usize, field: FieldRef) -> Self {
58        Self { index, field }
59    }
60
61    /// Get the variable's name
62    pub fn name(&self) -> &str {
63        self.field.name()
64    }
65
66    /// Get the variable's index
67    pub fn index(&self) -> usize {
68        self.index
69    }
70
71    /// Get the variable's field
72    pub fn field(&self) -> &FieldRef {
73        &self.field
74    }
75
76    #[cfg(feature = "proto")]
77    /// Reconstruct a [`LambdaVariable`] from a proto node.
78    pub fn try_from_proto(
79        node: &datafusion_proto_models::protobuf::PhysicalExprNode,
80        _ctx: &datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx<'_>,
81    ) -> Result<Arc<dyn PhysicalExpr>> {
82        use datafusion_physical_expr_common::{
83            expect_expr_variant, physical_expr::proto_decode::require_proto_field,
84        };
85        use datafusion_proto_models::protobuf;
86
87        let var = expect_expr_variant!(
88            node,
89            protobuf::physical_expr_node::ExprType::LambdaVariable,
90            "LambdaVariable",
91        );
92
93        Ok(Arc::new(LambdaVariable::new(
94            var.index as usize,
95            Arc::new(
96                require_proto_field(var.field.as_ref(), "LambdaVariable", "field")?
97                    .try_into()?,
98            ),
99        )))
100    }
101}
102
103impl std::fmt::Display for LambdaVariable {
104    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
105        write!(f, "{}@{}", self.name(), self.index)
106    }
107}
108
109impl PhysicalExpr for LambdaVariable {
110    fn data_type(&self, _input_schema: &Schema) -> Result<DataType> {
111        Ok(self.field.data_type().clone())
112    }
113
114    fn nullable(&self, _input_schema: &Schema) -> Result<bool> {
115        Ok(self.field.is_nullable())
116    }
117
118    fn evaluate(&self, batch: &RecordBatch) -> Result<ColumnarValue> {
119        if self.index >= batch.num_columns() {
120            return internal_err!(
121                "PhysicalExpr LambdaVariable references column '{}' at index {} (zero-based) but batch only has {} columns: {:?}",
122                self.name(),
123                self.index,
124                batch.num_columns(),
125                batch
126                    .schema_ref()
127                    .fields()
128                    .iter()
129                    .map(|f| f.name())
130                    .collect::<Vec<_>>()
131            );
132        }
133
134        if self.field.as_ref() != batch.schema_ref().field(self.index) {
135            return exec_err!(
136                "Field of physical LambdaVariable with index {} doesn't match batch field during evaluation {} != {}",
137                self.index,
138                self.field,
139                batch.schema_ref().field(self.index)
140            );
141        }
142
143        Ok(ColumnarValue::Array(Arc::clone(batch.column(self.index))))
144    }
145
146    fn return_field(&self, _input_schema: &Schema) -> Result<FieldRef> {
147        Ok(Arc::clone(&self.field))
148    }
149
150    fn children(&self) -> Vec<&Arc<dyn PhysicalExpr>> {
151        vec![]
152    }
153
154    fn with_new_children(
155        self: Arc<Self>,
156        _children: Vec<Arc<dyn PhysicalExpr>>,
157    ) -> Result<Arc<dyn PhysicalExpr>> {
158        Ok(self)
159    }
160
161    fn fmt_sql(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
162        write!(f, "{}@{}", self.name(), self.index)
163    }
164
165    #[cfg(feature = "proto")]
166    fn try_to_proto(
167        &self,
168        _ctx: &datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx<'_>,
169    ) -> Result<Option<datafusion_proto_models::protobuf::PhysicalExprNode>> {
170        use datafusion_proto_models::protobuf;
171
172        Ok(Some(protobuf::PhysicalExprNode {
173            expr_id: None,
174            expr_type: Some(protobuf::physical_expr_node::ExprType::LambdaVariable(
175                protobuf::PhysicalLambdaVariableExprNode {
176                    index: self.index() as u32,
177                    field: Some(self.field().as_ref().try_into()?),
178                },
179            )),
180        }))
181    }
182}
183
184/// Create a lambda variable expression
185pub fn lambda_variable(name: &str, schema: &Schema) -> Result<Arc<dyn PhysicalExpr>> {
186    let index = schema.index_of(name)?;
187    let field = Arc::clone(&schema.fields()[index]);
188
189    Ok(Arc::new(LambdaVariable::new(index, field)))
190}
datafusion_physical_expr/expressions/lambda_variable.rs

datafusion_physical_expr/expressions/
lambda_variable.rs
