/Users/andrewlamb/Software/datafusion/datafusion/physical-expr/src/utils/mod.rs
Line | Count | Source (jump to first uncovered line) |
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 | | mod guarantee; |
19 | | pub use guarantee::{Guarantee, LiteralGuarantee}; |
20 | | use hashbrown::HashSet; |
21 | | |
22 | | use std::borrow::Borrow; |
23 | | use std::collections::HashMap; |
24 | | use std::sync::Arc; |
25 | | |
26 | | use crate::expressions::{BinaryExpr, Column}; |
27 | | use crate::tree_node::ExprContext; |
28 | | use crate::PhysicalExpr; |
29 | | use crate::PhysicalSortExpr; |
30 | | |
31 | | use arrow::datatypes::SchemaRef; |
32 | | use datafusion_common::tree_node::{ |
33 | | Transformed, TransformedResult, TreeNode, TreeNodeRecursion, |
34 | | }; |
35 | | use datafusion_common::Result; |
36 | | use datafusion_expr::Operator; |
37 | | |
38 | | use itertools::Itertools; |
39 | | use petgraph::graph::NodeIndex; |
40 | | use petgraph::stable_graph::StableGraph; |
41 | | |
42 | | /// Assume the predicate is in the form of CNF, split the predicate to a Vec of PhysicalExprs. |
43 | | /// |
44 | | /// For example, split "a1 = a2 AND b1 <= b2 AND c1 != c2" into ["a1 = a2", "b1 <= b2", "c1 != c2"] |
45 | 33 | pub fn split_conjunction( |
46 | 33 | predicate: &Arc<dyn PhysicalExpr>, |
47 | 33 | ) -> Vec<&Arc<dyn PhysicalExpr>> { |
48 | 33 | split_impl(Operator::And, predicate, vec![]) |
49 | 33 | } |
50 | | |
51 | | /// Assume the predicate is in the form of DNF, split the predicate to a Vec of PhysicalExprs. |
52 | | /// |
53 | | /// For example, split "a1 = a2 OR b1 <= b2 OR c1 != c2" into ["a1 = a2", "b1 <= b2", "c1 != c2"] |
54 | 0 | pub fn split_disjunction( |
55 | 0 | predicate: &Arc<dyn PhysicalExpr>, |
56 | 0 | ) -> Vec<&Arc<dyn PhysicalExpr>> { |
57 | 0 | split_impl(Operator::Or, predicate, vec![]) |
58 | 0 | } |
59 | | |
60 | 67 | fn split_impl<'a>( |
61 | 67 | operator: Operator, |
62 | 67 | predicate: &'a Arc<dyn PhysicalExpr>, |
63 | 67 | mut exprs: Vec<&'a Arc<dyn PhysicalExpr>>, |
64 | 67 | ) -> Vec<&'a Arc<dyn PhysicalExpr>> { |
65 | 67 | match predicate.as_any().downcast_ref::<BinaryExpr>() { |
66 | 67 | Some(binary17 ) if binary.op() == &operator => { |
67 | 17 | let exprs = split_impl(operator, binary.left(), exprs); |
68 | 17 | split_impl(operator, binary.right(), exprs) |
69 | | } |
70 | | Some(_) | None => { |
71 | 50 | exprs.push(predicate); |
72 | 50 | exprs |
73 | | } |
74 | | } |
75 | 67 | } |
76 | | |
77 | | /// This function maps back requirement after ProjectionExec |
78 | | /// to the Executor for its input. |
79 | | // Specifically, `ProjectionExec` changes index of `Column`s in the schema of its input executor. |
80 | | // This function changes requirement given according to ProjectionExec schema to the requirement |
81 | | // according to schema of input executor to the ProjectionExec. |
82 | | // For instance, Column{"a", 0} would turn to Column{"a", 1}. Please note that this function assumes that |
83 | | // name of the Column is unique. If we have a requirement such that Column{"a", 0}, Column{"a", 1}. |
84 | | // This function will produce incorrect result (It will only emit single Column as a result). |
85 | 0 | pub fn map_columns_before_projection( |
86 | 0 | parent_required: &[Arc<dyn PhysicalExpr>], |
87 | 0 | proj_exprs: &[(Arc<dyn PhysicalExpr>, String)], |
88 | 0 | ) -> Vec<Arc<dyn PhysicalExpr>> { |
89 | 0 | let column_mapping = proj_exprs |
90 | 0 | .iter() |
91 | 0 | .filter_map(|(expr, name)| { |
92 | 0 | expr.as_any() |
93 | 0 | .downcast_ref::<Column>() |
94 | 0 | .map(|column| (name.clone(), column.clone())) |
95 | 0 | }) |
96 | 0 | .collect::<HashMap<_, _>>(); |
97 | 0 | parent_required |
98 | 0 | .iter() |
99 | 0 | .filter_map(|r| { |
100 | 0 | r.as_any() |
101 | 0 | .downcast_ref::<Column>() |
102 | 0 | .and_then(|c| column_mapping.get(c.name())) |
103 | 0 | }) |
104 | 0 | .map(|e| Arc::new(e.clone()) as _) |
105 | 0 | .collect() |
106 | 0 | } |
107 | | |
108 | | /// This function returns all `Arc<dyn PhysicalExpr>`s inside the given |
109 | | /// `PhysicalSortExpr` sequence. |
110 | 0 | pub fn convert_to_expr<T: Borrow<PhysicalSortExpr>>( |
111 | 0 | sequence: impl IntoIterator<Item = T>, |
112 | 0 | ) -> Vec<Arc<dyn PhysicalExpr>> { |
113 | 0 | sequence |
114 | 0 | .into_iter() |
115 | 0 | .map(|elem| Arc::clone(&elem.borrow().expr)) |
116 | 0 | .collect() |
117 | 0 | } |
118 | | |
119 | | /// This function finds the indices of `targets` within `items` using strict |
120 | | /// equality. |
121 | 0 | pub fn get_indices_of_exprs_strict<T: Borrow<Arc<dyn PhysicalExpr>>>( |
122 | 0 | targets: impl IntoIterator<Item = T>, |
123 | 0 | items: &[Arc<dyn PhysicalExpr>], |
124 | 0 | ) -> Vec<usize> { |
125 | 0 | targets |
126 | 0 | .into_iter() |
127 | 0 | .filter_map(|target| items.iter().position(|e| e.eq(target.borrow()))) |
128 | 0 | .collect() |
129 | 0 | } |
130 | | |
131 | | pub type ExprTreeNode<T> = ExprContext<Option<T>>; |
132 | | |
133 | | /// This struct is used to convert a [`PhysicalExpr`] tree into a DAEG (i.e. an expression |
134 | | /// DAG) by collecting identical expressions in one node. Caller specifies the node type |
135 | | /// in the DAEG via the `constructor` argument, which constructs nodes in the DAEG from |
136 | | /// the [`ExprTreeNode`] ancillary object. |
137 | | struct PhysicalExprDAEGBuilder<'a, T, F: Fn(&ExprTreeNode<NodeIndex>) -> Result<T>> { |
138 | | // The resulting DAEG (expression DAG). |
139 | | graph: StableGraph<T, usize>, |
140 | | // A vector of visited expression nodes and their corresponding node indices. |
141 | | visited_plans: Vec<(Arc<dyn PhysicalExpr>, NodeIndex)>, |
142 | | // A function to convert an input expression node to T. |
143 | | constructor: &'a F, |
144 | | } |
145 | | |
146 | | impl<'a, T, F: Fn(&ExprTreeNode<NodeIndex>) -> Result<T>> |
147 | | PhysicalExprDAEGBuilder<'a, T, F> |
148 | | { |
149 | | // This method mutates an expression node by transforming it to a physical expression |
150 | | // and adding it to the graph. The method returns the mutated expression node. |
151 | 17.0k | fn mutate( |
152 | 17.0k | &mut self, |
153 | 17.0k | mut node: ExprTreeNode<NodeIndex>, |
154 | 17.0k | ) -> Result<Transformed<ExprTreeNode<NodeIndex>>> { |
155 | 17.0k | // Get the expression associated with the input expression node. |
156 | 17.0k | let expr = &node.expr; |
157 | | |
158 | | // Check if the expression has already been visited. |
159 | 92.2k | let node_idx17.0k = match self.visited_plans.iter().find(17.0k |(e, _)| expr.eq(e))17.0k { |
160 | | // If the expression has been visited, return the corresponding node index. |
161 | 2.64k | Some((_, idx)) => *idx, |
162 | | // If the expression has not been visited, add a new node to the graph and |
163 | | // add edges to its child nodes. Add the visited expression to the vector |
164 | | // of visited expressions and return the newly created node index. |
165 | | None => { |
166 | 14.4k | let node_idx = self.graph.add_node((self.constructor)(&node)?0 ); |
167 | 15.6k | for expr_node in node.children.iter()14.4k { |
168 | 15.6k | self.graph.add_edge(node_idx, expr_node.data.unwrap(), 0); |
169 | 15.6k | } |
170 | 14.4k | self.visited_plans.push((Arc::clone(expr), node_idx)); |
171 | 14.4k | node_idx |
172 | | } |
173 | | }; |
174 | | // Set the data field of the input expression node to the corresponding node index. |
175 | 17.0k | node.data = Some(node_idx); |
176 | 17.0k | // Return the mutated expression node. |
177 | 17.0k | Ok(Transformed::yes(node)) |
178 | 17.0k | } |
179 | | } |
180 | | |
181 | | // A function that builds a directed acyclic graph of physical expression trees. |
182 | 1.14k | pub fn build_dag<T, F>( |
183 | 1.14k | expr: Arc<dyn PhysicalExpr>, |
184 | 1.14k | constructor: &F, |
185 | 1.14k | ) -> Result<(NodeIndex, StableGraph<T, usize>)> |
186 | 1.14k | where |
187 | 1.14k | F: Fn(&ExprTreeNode<NodeIndex>) -> Result<T>, |
188 | 1.14k | { |
189 | 1.14k | // Create a new expression tree node from the input expression. |
190 | 1.14k | let init = ExprTreeNode::new_default(expr); |
191 | 1.14k | // Create a new `PhysicalExprDAEGBuilder` instance. |
192 | 1.14k | let mut builder = PhysicalExprDAEGBuilder { |
193 | 1.14k | graph: StableGraph::<T, usize>::new(), |
194 | 1.14k | visited_plans: Vec::<(Arc<dyn PhysicalExpr>, NodeIndex)>::new(), |
195 | 1.14k | constructor, |
196 | 1.14k | }; |
197 | | // Use the builder to transform the expression tree node into a DAG. |
198 | 17.0k | let root1.14k = init.transform_up(1.14k |node| builder.mutate(node)).data()1.14k ?0 ; |
199 | | // Return a tuple containing the root node index and the DAG. |
200 | 1.14k | Ok((root.data.unwrap(), builder.graph)) |
201 | 1.14k | } |
202 | | |
203 | | /// Recursively extract referenced [`Column`]s within a [`PhysicalExpr`]. |
204 | 4.50k | pub fn collect_columns(expr: &Arc<dyn PhysicalExpr>) -> HashSet<Column> { |
205 | 4.50k | let mut columns = HashSet::<Column>::new(); |
206 | 4.90k | expr.apply(|expr| { |
207 | 4.90k | if let Some(column4.59k ) = expr.as_any().downcast_ref::<Column>() { |
208 | 4.59k | columns.get_or_insert_owned(column); |
209 | 4.59k | }307 |
210 | 4.90k | Ok(TreeNodeRecursion::Continue) |
211 | 4.90k | }) |
212 | 4.50k | // pre_visit always returns OK, so this will always too |
213 | 4.50k | .expect("no way to return error during recursion"); |
214 | 4.50k | columns |
215 | 4.50k | } |
216 | | |
217 | | /// Re-assign column indices referenced in predicate according to given schema. |
218 | | /// This may be helpful when dealing with projections. |
219 | 0 | pub fn reassign_predicate_columns( |
220 | 0 | pred: Arc<dyn PhysicalExpr>, |
221 | 0 | schema: &SchemaRef, |
222 | 0 | ignore_not_found: bool, |
223 | 0 | ) -> Result<Arc<dyn PhysicalExpr>> { |
224 | 0 | pred.transform_down(|expr| { |
225 | 0 | let expr_any = expr.as_any(); |
226 | | |
227 | 0 | if let Some(column) = expr_any.downcast_ref::<Column>() { |
228 | 0 | let index = match schema.index_of(column.name()) { |
229 | 0 | Ok(idx) => idx, |
230 | 0 | Err(_) if ignore_not_found => usize::MAX, |
231 | 0 | Err(e) => return Err(e.into()), |
232 | | }; |
233 | 0 | return Ok(Transformed::yes(Arc::new(Column::new( |
234 | 0 | column.name(), |
235 | 0 | index, |
236 | 0 | )))); |
237 | 0 | } |
238 | 0 | Ok(Transformed::no(expr)) |
239 | 0 | }) |
240 | 0 | .data() |
241 | 0 | } |
242 | | |
243 | | /// Merge left and right sort expressions, checking for duplicates. |
244 | 2 | pub fn merge_vectors( |
245 | 2 | left: &[PhysicalSortExpr], |
246 | 2 | right: &[PhysicalSortExpr], |
247 | 2 | ) -> Vec<PhysicalSortExpr> { |
248 | 2 | left.iter() |
249 | 2 | .cloned() |
250 | 2 | .chain(right.iter().cloned()) |
251 | 2 | .unique() |
252 | 2 | .collect() |
253 | 2 | } |
254 | | |
255 | | #[cfg(test)] |
256 | | pub(crate) mod tests { |
257 | | use std::any::Any; |
258 | | use std::fmt::{Display, Formatter}; |
259 | | |
260 | | use super::*; |
261 | | use crate::expressions::{binary, cast, col, in_list, lit, Literal}; |
262 | | |
263 | | use arrow_array::{ArrayRef, Float32Array, Float64Array}; |
264 | | use arrow_schema::{DataType, Field, Schema}; |
265 | | use datafusion_common::{exec_err, DataFusionError, ScalarValue}; |
266 | | use datafusion_expr::sort_properties::{ExprProperties, SortProperties}; |
267 | | use datafusion_expr::{ColumnarValue, ScalarUDFImpl, Signature, Volatility}; |
268 | | |
269 | | use petgraph::visit::Bfs; |
270 | | |
271 | | #[derive(Debug, Clone)] |
272 | | pub struct TestScalarUDF { |
273 | | pub(crate) signature: Signature, |
274 | | } |
275 | | |
276 | | impl TestScalarUDF { |
277 | | pub fn new() -> Self { |
278 | | use DataType::*; |
279 | | Self { |
280 | | signature: Signature::uniform( |
281 | | 1, |
282 | | vec![Float64, Float32], |
283 | | Volatility::Immutable, |
284 | | ), |
285 | | } |
286 | | } |
287 | | } |
288 | | |
289 | | impl ScalarUDFImpl for TestScalarUDF { |
290 | | fn as_any(&self) -> &dyn Any { |
291 | | self |
292 | | } |
293 | | fn name(&self) -> &str { |
294 | | "test-scalar-udf" |
295 | | } |
296 | | |
297 | | fn signature(&self) -> &Signature { |
298 | | &self.signature |
299 | | } |
300 | | |
301 | | fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> { |
302 | | let arg_type = &arg_types[0]; |
303 | | |
304 | | match arg_type { |
305 | | DataType::Float32 => Ok(DataType::Float32), |
306 | | _ => Ok(DataType::Float64), |
307 | | } |
308 | | } |
309 | | |
310 | | fn output_ordering(&self, input: &[ExprProperties]) -> Result<SortProperties> { |
311 | | Ok(input[0].sort_properties) |
312 | | } |
313 | | |
314 | | fn invoke(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> { |
315 | | let args = ColumnarValue::values_to_arrays(args)?; |
316 | | |
317 | | let arr: ArrayRef = match args[0].data_type() { |
318 | | DataType::Float64 => Arc::new({ |
319 | | let arg = &args[0] |
320 | | .as_any() |
321 | | .downcast_ref::<Float64Array>() |
322 | | .ok_or_else(|| { |
323 | | DataFusionError::Internal(format!( |
324 | | "could not cast {} to {}", |
325 | | self.name(), |
326 | | std::any::type_name::<Float64Array>() |
327 | | )) |
328 | | })?; |
329 | | |
330 | | arg.iter() |
331 | | .map(|a| a.map(f64::floor)) |
332 | | .collect::<Float64Array>() |
333 | | }), |
334 | | DataType::Float32 => Arc::new({ |
335 | | let arg = &args[0] |
336 | | .as_any() |
337 | | .downcast_ref::<Float32Array>() |
338 | | .ok_or_else(|| { |
339 | | DataFusionError::Internal(format!( |
340 | | "could not cast {} to {}", |
341 | | self.name(), |
342 | | std::any::type_name::<Float32Array>() |
343 | | )) |
344 | | })?; |
345 | | |
346 | | arg.iter() |
347 | | .map(|a| a.map(f32::floor)) |
348 | | .collect::<Float32Array>() |
349 | | }), |
350 | | other => { |
351 | | return exec_err!( |
352 | | "Unsupported data type {other:?} for function {}", |
353 | | self.name() |
354 | | ); |
355 | | } |
356 | | }; |
357 | | Ok(ColumnarValue::Array(arr)) |
358 | | } |
359 | | } |
360 | | |
361 | | #[derive(Clone)] |
362 | | struct DummyProperty { |
363 | | expr_type: String, |
364 | | } |
365 | | |
366 | | /// This is a dummy node in the DAEG; it stores a reference to the actual |
367 | | /// [PhysicalExpr] as well as a dummy property. |
368 | | #[derive(Clone)] |
369 | | struct PhysicalExprDummyNode { |
370 | | pub expr: Arc<dyn PhysicalExpr>, |
371 | | pub property: DummyProperty, |
372 | | } |
373 | | |
374 | | impl Display for PhysicalExprDummyNode { |
375 | | fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { |
376 | | write!(f, "{}", self.expr) |
377 | | } |
378 | | } |
379 | | |
380 | | fn make_dummy_node(node: &ExprTreeNode<NodeIndex>) -> Result<PhysicalExprDummyNode> { |
381 | | let expr = Arc::clone(&node.expr); |
382 | | let dummy_property = if expr.as_any().is::<BinaryExpr>() { |
383 | | "Binary" |
384 | | } else if expr.as_any().is::<Column>() { |
385 | | "Column" |
386 | | } else if expr.as_any().is::<Literal>() { |
387 | | "Literal" |
388 | | } else { |
389 | | "Other" |
390 | | } |
391 | | .to_owned(); |
392 | | Ok(PhysicalExprDummyNode { |
393 | | expr, |
394 | | property: DummyProperty { |
395 | | expr_type: dummy_property, |
396 | | }, |
397 | | }) |
398 | | } |
399 | | |
400 | | #[test] |
401 | | fn test_build_dag() -> Result<()> { |
402 | | let schema = Schema::new(vec![ |
403 | | Field::new("0", DataType::Int32, true), |
404 | | Field::new("1", DataType::Int32, true), |
405 | | Field::new("2", DataType::Int32, true), |
406 | | ]); |
407 | | let expr = binary( |
408 | | cast( |
409 | | binary( |
410 | | col("0", &schema)?, |
411 | | Operator::Plus, |
412 | | col("1", &schema)?, |
413 | | &schema, |
414 | | )?, |
415 | | &schema, |
416 | | DataType::Int64, |
417 | | )?, |
418 | | Operator::Gt, |
419 | | binary( |
420 | | cast(col("2", &schema)?, &schema, DataType::Int64)?, |
421 | | Operator::Plus, |
422 | | lit(ScalarValue::Int64(Some(10))), |
423 | | &schema, |
424 | | )?, |
425 | | &schema, |
426 | | )?; |
427 | | let mut vector_dummy_props = vec![]; |
428 | | let (root, graph) = build_dag(expr, &make_dummy_node)?; |
429 | | let mut bfs = Bfs::new(&graph, root); |
430 | | while let Some(node_index) = bfs.next(&graph) { |
431 | | let node = &graph[node_index]; |
432 | | vector_dummy_props.push(node.property.clone()); |
433 | | } |
434 | | |
435 | | assert_eq!( |
436 | | vector_dummy_props |
437 | | .iter() |
438 | | .filter(|property| property.expr_type == "Binary") |
439 | | .count(), |
440 | | 3 |
441 | | ); |
442 | | assert_eq!( |
443 | | vector_dummy_props |
444 | | .iter() |
445 | | .filter(|property| property.expr_type == "Column") |
446 | | .count(), |
447 | | 3 |
448 | | ); |
449 | | assert_eq!( |
450 | | vector_dummy_props |
451 | | .iter() |
452 | | .filter(|property| property.expr_type == "Literal") |
453 | | .count(), |
454 | | 1 |
455 | | ); |
456 | | assert_eq!( |
457 | | vector_dummy_props |
458 | | .iter() |
459 | | .filter(|property| property.expr_type == "Other") |
460 | | .count(), |
461 | | 2 |
462 | | ); |
463 | | Ok(()) |
464 | | } |
465 | | |
466 | | #[test] |
467 | | fn test_convert_to_expr() -> Result<()> { |
468 | | let schema = Schema::new(vec![Field::new("a", DataType::UInt64, false)]); |
469 | | let sort_expr = vec![PhysicalSortExpr { |
470 | | expr: col("a", &schema)?, |
471 | | options: Default::default(), |
472 | | }]; |
473 | | assert!(convert_to_expr(&sort_expr)[0].eq(&sort_expr[0].expr)); |
474 | | Ok(()) |
475 | | } |
476 | | |
477 | | #[test] |
478 | | fn test_get_indices_of_exprs_strict() { |
479 | | let list1: Vec<Arc<dyn PhysicalExpr>> = vec![ |
480 | | Arc::new(Column::new("a", 0)), |
481 | | Arc::new(Column::new("b", 1)), |
482 | | Arc::new(Column::new("c", 2)), |
483 | | Arc::new(Column::new("d", 3)), |
484 | | ]; |
485 | | let list2: Vec<Arc<dyn PhysicalExpr>> = vec![ |
486 | | Arc::new(Column::new("b", 1)), |
487 | | Arc::new(Column::new("c", 2)), |
488 | | Arc::new(Column::new("a", 0)), |
489 | | ]; |
490 | | assert_eq!(get_indices_of_exprs_strict(&list1, &list2), vec![2, 0, 1]); |
491 | | assert_eq!(get_indices_of_exprs_strict(&list2, &list1), vec![1, 2, 0]); |
492 | | } |
493 | | |
494 | | #[test] |
495 | | fn test_reassign_predicate_columns_in_list() { |
496 | | let int_field = Field::new("should_not_matter", DataType::Int64, true); |
497 | | let dict_field = Field::new( |
498 | | "id", |
499 | | DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)), |
500 | | true, |
501 | | ); |
502 | | let schema_small = Arc::new(Schema::new(vec![dict_field.clone()])); |
503 | | let schema_big = Arc::new(Schema::new(vec![int_field, dict_field])); |
504 | | let pred = in_list( |
505 | | Arc::new(Column::new_with_schema("id", &schema_big).unwrap()), |
506 | | vec![lit(ScalarValue::Dictionary( |
507 | | Box::new(DataType::Int32), |
508 | | Box::new(ScalarValue::from("2")), |
509 | | ))], |
510 | | &false, |
511 | | &schema_big, |
512 | | ) |
513 | | .unwrap(); |
514 | | |
515 | | let actual = reassign_predicate_columns(pred, &schema_small, false).unwrap(); |
516 | | |
517 | | let expected = in_list( |
518 | | Arc::new(Column::new_with_schema("id", &schema_small).unwrap()), |
519 | | vec![lit(ScalarValue::Dictionary( |
520 | | Box::new(DataType::Int32), |
521 | | Box::new(ScalarValue::from("2")), |
522 | | ))], |
523 | | &false, |
524 | | &schema_small, |
525 | | ) |
526 | | .unwrap(); |
527 | | |
528 | | assert_eq!(actual.as_ref(), expected.as_any()); |
529 | | } |
530 | | |
531 | | #[test] |
532 | | fn test_collect_columns() -> Result<()> { |
533 | | let expr1 = Arc::new(Column::new("col1", 2)) as _; |
534 | | let mut expected = HashSet::new(); |
535 | | expected.insert(Column::new("col1", 2)); |
536 | | assert_eq!(collect_columns(&expr1), expected); |
537 | | |
538 | | let expr2 = Arc::new(Column::new("col2", 5)) as _; |
539 | | let mut expected = HashSet::new(); |
540 | | expected.insert(Column::new("col2", 5)); |
541 | | assert_eq!(collect_columns(&expr2), expected); |
542 | | |
543 | | let expr3 = Arc::new(BinaryExpr::new(expr1, Operator::Plus, expr2)) as _; |
544 | | let mut expected = HashSet::new(); |
545 | | expected.insert(Column::new("col1", 2)); |
546 | | expected.insert(Column::new("col2", 5)); |
547 | | assert_eq!(collect_columns(&expr3), expected); |
548 | | Ok(()) |
549 | | } |
550 | | } |