datafusion_physical_expr/
physical_expr.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 std::sync::Arc;
19
20use crate::expressions::{self, Column};
21use crate::{LexOrdering, PhysicalSortExpr, create_physical_expr};
22
23use arrow::compute::SortOptions;
24use arrow::datatypes::{DataType, Schema, SchemaRef};
25use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode};
26use datafusion_common::{DFSchema, HashMap, ScalarValue, SplitPoint};
27use datafusion_common::{Result, plan_err};
28use datafusion_expr::execution_props::ExecutionProps;
29use datafusion_expr::physical_planning_context::PhysicalPlanningContext;
30use datafusion_expr::{Expr, Partitioning as LogicalPartitioning, SortExpr};
31use datafusion_expr_common::casts::try_cast_literal_to_type;
32
33use itertools::izip;
34// Exports:
35use crate::{Partitioning, RangePartitioning};
36pub(crate) use datafusion_physical_expr_common::physical_expr::PhysicalExpr;
37
38/// Adds the `offset` value to `Column` indices inside `expr`. This function is
39/// generally used during the update of the right table schema in join operations.
40pub fn add_offset_to_expr(
41    expr: Arc<dyn PhysicalExpr>,
42    offset: isize,
43) -> Result<Arc<dyn PhysicalExpr>> {
44    expr.transform_down(|e| match e.downcast_ref::<Column>() {
45        Some(col) => {
46            let Some(idx) = col.index().checked_add_signed(offset) else {
47                return plan_err!("Column index overflow");
48            };
49            Ok(Transformed::yes(Arc::new(Column::new(col.name(), idx))))
50        }
51        None => Ok(Transformed::no(e)),
52    })
53    .data()
54}
55
56/// This function is similar to the `contains` method of `Vec`. It finds
57/// whether `expr` is among `physical_exprs`.
58pub fn physical_exprs_contains(
59    physical_exprs: &[Arc<dyn PhysicalExpr>],
60    expr: &Arc<dyn PhysicalExpr>,
61) -> bool {
62    physical_exprs
63        .iter()
64        .any(|physical_expr| physical_expr.as_ref().eq(expr.as_ref()))
65}
66
67/// Checks whether the given physical expression slices are equal.
68pub fn physical_exprs_equal(
69    lhs: &[Arc<dyn PhysicalExpr>],
70    rhs: &[Arc<dyn PhysicalExpr>],
71) -> bool {
72    lhs.len() == rhs.len()
73        && izip!(lhs, rhs).all(|(lhs, rhs)| lhs.as_ref().eq(rhs.as_ref()))
74}
75
76/// Checks whether the given physical expression slices are equal in the sense
77/// of bags (multi-sets), disregarding their orderings.
78pub fn physical_exprs_bag_equal(
79    lhs: &[Arc<dyn PhysicalExpr>],
80    rhs: &[Arc<dyn PhysicalExpr>],
81) -> bool {
82    let mut multi_set_lhs: HashMap<_, usize> = HashMap::new();
83    let mut multi_set_rhs: HashMap<_, usize> = HashMap::new();
84    for expr in lhs {
85        *multi_set_lhs.entry(expr).or_insert(0) += 1;
86    }
87    for expr in rhs {
88        *multi_set_rhs.entry(expr).or_insert(0) += 1;
89    }
90    multi_set_lhs == multi_set_rhs
91}
92
93/// Converts logical sort expressions to physical sort expressions.
94///
95/// This function transforms a collection of logical sort expressions into their
96/// physical representation that can be used during query execution.
97///
98/// # Arguments
99///
100/// * `schema` - The schema containing column definitions.
101/// * `sort_order` - A collection of logical sort expressions grouped into
102///   lexicographic orderings.
103///
104/// # Returns
105///
106/// A vector of lexicographic orderings for physical execution, or an error if
107/// the transformation fails.
108///
109/// # Examples
110///
111/// ```
112/// // Create orderings from columns "id" and "name"
113/// # use arrow::datatypes::{Schema, Field, DataType};
114/// # use datafusion_physical_expr::create_ordering;
115/// # use datafusion_common::Column;
116/// # use datafusion_expr::{Expr, SortExpr};
117/// #
118/// // Create a schema with two fields
119/// let schema = Schema::new(vec![
120///     Field::new("id", DataType::Int32, false),
121///     Field::new("name", DataType::Utf8, false),
122/// ]);
123///
124/// let sort_exprs = vec![
125///     vec![SortExpr {
126///         expr: Expr::Column(Column::new(Some("t"), "id")),
127///         asc: true,
128///         nulls_first: false,
129///     }],
130///     vec![SortExpr {
131///         expr: Expr::Column(Column::new(Some("t"), "name")),
132///         asc: false,
133///         nulls_first: true,
134///     }],
135/// ];
136/// let result = create_ordering(&schema, &sort_exprs).unwrap();
137/// ```
138pub fn create_ordering(
139    schema: &Schema,
140    sort_order: &[Vec<SortExpr>],
141) -> Result<Vec<LexOrdering>> {
142    let mut all_sort_orders = vec![];
143
144    for (group_idx, exprs) in sort_order.iter().enumerate() {
145        // Construct PhysicalSortExpr objects from Expr objects:
146        let mut sort_exprs = vec![];
147        for (expr_idx, sort) in exprs.iter().enumerate() {
148            match &sort.expr {
149                Expr::Column(col) => match expressions::col(&col.name, schema) {
150                    Ok(expr) => {
151                        let opts = SortOptions::new(!sort.asc, sort.nulls_first);
152                        sort_exprs.push(PhysicalSortExpr::new(expr, opts));
153                    }
154                    // Cannot find expression in the projected_schema, stop iterating
155                    // since rest of the orderings are violated
156                    Err(_) => break,
157                },
158                expr => {
159                    return plan_err!(
160                        "Expected single column reference in sort_order[{}][{}], got {}",
161                        group_idx,
162                        expr_idx,
163                        expr
164                    );
165                }
166            }
167        }
168        all_sort_orders.extend(LexOrdering::new(sort_exprs));
169    }
170    Ok(all_sort_orders)
171}
172
173/// Creates a vector of [LexOrdering] from a vector of logical expression
174pub fn create_lex_ordering(
175    schema: &SchemaRef,
176    sort_order: &[Vec<SortExpr>],
177    execution_props: &ExecutionProps,
178) -> Result<Vec<LexOrdering>> {
179    // Try the fast path that only supports column references first
180    // This avoids creating a DFSchema
181    if let Ok(ordering) = create_ordering(schema, sort_order) {
182        return Ok(ordering);
183    }
184
185    let df_schema = DFSchema::try_from(Arc::clone(schema))?;
186
187    let mut all_sort_orders = vec![];
188
189    for exprs in sort_order.iter() {
190        all_sort_orders.extend(LexOrdering::new(create_physical_sort_exprs(
191            exprs,
192            &df_schema,
193            execution_props,
194            &PhysicalPlanningContext::default(),
195        )?));
196    }
197    Ok(all_sort_orders)
198}
199
200/// Create a physical sort expression from a logical expression
201///
202/// See [`create_physical_expr`] for details on the `planning_ctx` argument.
203pub fn create_physical_sort_expr(
204    e: &SortExpr,
205    input_dfschema: &DFSchema,
206    execution_props: &ExecutionProps,
207    planning_ctx: &PhysicalPlanningContext,
208) -> Result<PhysicalSortExpr> {
209    create_physical_expr(&e.expr, input_dfschema, execution_props, planning_ctx).map(
210        |expr| {
211            let options = SortOptions::new(!e.asc, e.nulls_first);
212            PhysicalSortExpr::new(expr, options)
213        },
214    )
215}
216
217/// Create vector of physical sort expression from a vector of logical expression
218///
219/// See [`create_physical_expr`] for details on the `planning_ctx` argument.
220pub fn create_physical_sort_exprs(
221    exprs: &[SortExpr],
222    input_dfschema: &DFSchema,
223    execution_props: &ExecutionProps,
224    planning_ctx: &PhysicalPlanningContext,
225) -> Result<Vec<PhysicalSortExpr>> {
226    exprs
227        .iter()
228        .map(|e| {
229            create_physical_sort_expr(e, input_dfschema, execution_props, planning_ctx)
230        })
231        .collect()
232}
233
234/// Create physical partitioning from logical partitioning.
235///
236/// See [`create_physical_expr`] for details on the `planning_ctx` argument.
237pub fn create_physical_partitioning(
238    partitioning: &LogicalPartitioning,
239    input_dfschema: &DFSchema,
240    execution_props: &ExecutionProps,
241    planning_ctx: &PhysicalPlanningContext,
242) -> Result<Partitioning> {
243    match partitioning {
244        LogicalPartitioning::RoundRobinBatch(n) => Ok(Partitioning::RoundRobinBatch(*n)),
245        LogicalPartitioning::Hash(exprs, partition_count) => {
246            let exprs = exprs
247                .iter()
248                .map(|expr| {
249                    create_physical_expr(
250                        expr,
251                        input_dfschema,
252                        execution_props,
253                        planning_ctx,
254                    )
255                })
256                .collect::<Result<Vec<_>>>()?;
257            Ok(Partitioning::Hash(exprs, *partition_count))
258        }
259        LogicalPartitioning::Range(range) => {
260            let ordering = create_physical_sort_exprs(
261                range.ordering(),
262                input_dfschema,
263                execution_props,
264                planning_ctx,
265            )?;
266            let Some(ordering) = LexOrdering::new(ordering) else {
267                return plan_err!("Range partitioning requires non-empty ordering");
268            };
269            let split_points = normalize_range_split_points(
270                &ordering,
271                range.split_points(),
272                input_dfschema.as_arrow(),
273            )?;
274            let range = RangePartitioning::try_new(ordering, split_points)?;
275            Ok(Partitioning::Range(range))
276        }
277        LogicalPartitioning::DistributeBy(_) => {
278            datafusion_common::not_impl_err!(
279                "Physical plan does not support DistributeBy partitioning"
280            )
281        }
282    }
283}
284
285fn normalize_range_split_points(
286    ordering: &LexOrdering,
287    split_points: &[SplitPoint],
288    schema: &Schema,
289) -> Result<Vec<SplitPoint>> {
290    split_points
291        .iter()
292        .enumerate()
293        .map(|(split_idx, split_point)| {
294            let values = split_point
295                .values()
296                .iter()
297                .zip(ordering.iter())
298                .enumerate()
299                .map(|(value_idx, (value, sort_expr))| {
300                    let target_type = sort_expr.expr.data_type(schema)?;
301                    normalize_range_split_point_value(
302                        value,
303                        &target_type,
304                        split_idx,
305                        value_idx,
306                    )
307                })
308                .collect::<Result<Vec<_>>>()?;
309            Ok(SplitPoint::new(values))
310        })
311        .collect()
312}
313
314fn normalize_range_split_point_value(
315    value: &ScalarValue,
316    target_type: &DataType,
317    split_idx: usize,
318    value_idx: usize,
319) -> Result<ScalarValue> {
320    let value_type = value.data_type();
321    if &value_type == target_type {
322        return Ok(value.clone());
323    }
324
325    if let Some(casted) = try_cast_literal_to_type(value, target_type) {
326        // Split points define physical partition boundaries, so normalization
327        // must reject casts that would change the advertised boundary.
328        if try_cast_literal_to_type(&casted, &value_type).as_ref() == Some(value) {
329            return Ok(casted);
330        }
331    }
332
333    plan_err!(
334        "Range output partitioning split point {split_idx} value {value_idx} with type {value_type} cannot be represented exactly as ordering expression type {target_type}"
335    )
336}
337
338pub fn add_offset_to_physical_sort_exprs(
339    sort_exprs: impl IntoIterator<Item = PhysicalSortExpr>,
340    offset: isize,
341) -> Result<Vec<PhysicalSortExpr>> {
342    sort_exprs
343        .into_iter()
344        .map(|mut sort_expr| {
345            sort_expr.expr = add_offset_to_expr(sort_expr.expr, offset)?;
346            Ok(sort_expr)
347        })
348        .collect()
349}
350
351#[cfg(test)]
352mod tests {
353    use super::*;
354
355    use crate::expressions::{BinaryExpr, Literal, UnKnownColumn};
356    use crate::physical_expr::{
357        physical_exprs_bag_equal, physical_exprs_contains, physical_exprs_equal,
358    };
359    use datafusion_physical_expr_common::physical_expr::is_volatile;
360
361    use arrow::datatypes::DataType;
362    use arrow::record_batch::RecordBatch;
363    use datafusion_common::ScalarValue;
364    use datafusion_expr::ColumnarValue;
365    use datafusion_expr::Operator;
366    use std::fmt;
367
368    #[test]
369    fn test_physical_exprs_contains() {
370        let lit_true = Arc::new(Literal::new(ScalarValue::Boolean(Some(true))))
371            as Arc<dyn PhysicalExpr>;
372        let lit_false = Arc::new(Literal::new(ScalarValue::Boolean(Some(false))))
373            as Arc<dyn PhysicalExpr>;
374        let lit4 =
375            Arc::new(Literal::new(ScalarValue::Int32(Some(4)))) as Arc<dyn PhysicalExpr>;
376        let lit2 =
377            Arc::new(Literal::new(ScalarValue::Int32(Some(2)))) as Arc<dyn PhysicalExpr>;
378        let lit1 =
379            Arc::new(Literal::new(ScalarValue::Int32(Some(1)))) as Arc<dyn PhysicalExpr>;
380        let col_a_expr = Arc::new(Column::new("a", 0)) as Arc<dyn PhysicalExpr>;
381        let col_b_expr = Arc::new(Column::new("b", 1)) as Arc<dyn PhysicalExpr>;
382        let col_c_expr = Arc::new(Column::new("c", 2)) as Arc<dyn PhysicalExpr>;
383
384        // lit(true), lit(false), lit(4), lit(2), Col(a), Col(b)
385        let physical_exprs: Vec<Arc<dyn PhysicalExpr>> = vec![
386            Arc::clone(&lit_true),
387            Arc::clone(&lit_false),
388            Arc::clone(&lit4),
389            Arc::clone(&lit2),
390            Arc::clone(&col_a_expr),
391            Arc::clone(&col_b_expr),
392        ];
393        // below expressions are inside physical_exprs
394        assert!(physical_exprs_contains(&physical_exprs, &lit_true));
395        assert!(physical_exprs_contains(&physical_exprs, &lit2));
396        assert!(physical_exprs_contains(&physical_exprs, &col_b_expr));
397
398        // below expressions are not inside physical_exprs
399        assert!(!physical_exprs_contains(&physical_exprs, &col_c_expr));
400        assert!(!physical_exprs_contains(&physical_exprs, &lit1));
401
402        let unknown = Arc::new(UnKnownColumn::new("unknown")) as Arc<dyn PhysicalExpr>;
403        assert!(!physical_exprs_contains(
404            std::slice::from_ref(&unknown),
405            &unknown
406        ));
407    }
408
409    #[test]
410    fn test_physical_exprs_equal() {
411        let lit_true = Arc::new(Literal::new(ScalarValue::Boolean(Some(true))))
412            as Arc<dyn PhysicalExpr>;
413        let lit_false = Arc::new(Literal::new(ScalarValue::Boolean(Some(false))))
414            as Arc<dyn PhysicalExpr>;
415        let lit1 =
416            Arc::new(Literal::new(ScalarValue::Int32(Some(1)))) as Arc<dyn PhysicalExpr>;
417        let lit2 =
418            Arc::new(Literal::new(ScalarValue::Int32(Some(2)))) as Arc<dyn PhysicalExpr>;
419        let col_b_expr = Arc::new(Column::new("b", 1)) as Arc<dyn PhysicalExpr>;
420
421        let vec1 = vec![Arc::clone(&lit_true), Arc::clone(&lit_false)];
422        let vec2 = vec![Arc::clone(&lit_true), Arc::clone(&col_b_expr)];
423        let vec3 = vec![Arc::clone(&lit2), Arc::clone(&lit1)];
424        let vec4 = vec![Arc::clone(&lit_true), Arc::clone(&lit_false)];
425
426        // these vectors are same
427        assert!(physical_exprs_equal(&vec1, &vec1));
428        assert!(physical_exprs_equal(&vec1, &vec4));
429        assert!(physical_exprs_bag_equal(&vec1, &vec1));
430        assert!(physical_exprs_bag_equal(&vec1, &vec4));
431
432        // these vectors are different
433        assert!(!physical_exprs_equal(&vec1, &vec2));
434        assert!(!physical_exprs_equal(&vec1, &vec3));
435        assert!(!physical_exprs_bag_equal(&vec1, &vec2));
436        assert!(!physical_exprs_bag_equal(&vec1, &vec3));
437
438        let unknown = Arc::new(UnKnownColumn::new("unknown")) as Arc<dyn PhysicalExpr>;
439        assert!(!physical_exprs_equal(
440            std::slice::from_ref(&unknown),
441            std::slice::from_ref(&unknown)
442        ));
443    }
444
445    #[test]
446    fn test_physical_exprs_set_equal() {
447        let list1: Vec<Arc<dyn PhysicalExpr>> = vec![
448            Arc::new(Column::new("a", 0)),
449            Arc::new(Column::new("a", 0)),
450            Arc::new(Column::new("b", 1)),
451        ];
452        let list2: Vec<Arc<dyn PhysicalExpr>> = vec![
453            Arc::new(Column::new("b", 1)),
454            Arc::new(Column::new("b", 1)),
455            Arc::new(Column::new("a", 0)),
456        ];
457        assert!(!physical_exprs_bag_equal(
458            list1.as_slice(),
459            list2.as_slice()
460        ));
461        assert!(!physical_exprs_bag_equal(
462            list2.as_slice(),
463            list1.as_slice()
464        ));
465        assert!(!physical_exprs_equal(list1.as_slice(), list2.as_slice()));
466        assert!(!physical_exprs_equal(list2.as_slice(), list1.as_slice()));
467
468        let list3: Vec<Arc<dyn PhysicalExpr>> = vec![
469            Arc::new(Column::new("a", 0)),
470            Arc::new(Column::new("b", 1)),
471            Arc::new(Column::new("c", 2)),
472            Arc::new(Column::new("a", 0)),
473            Arc::new(Column::new("b", 1)),
474        ];
475        let list4: Vec<Arc<dyn PhysicalExpr>> = vec![
476            Arc::new(Column::new("b", 1)),
477            Arc::new(Column::new("b", 1)),
478            Arc::new(Column::new("a", 0)),
479            Arc::new(Column::new("c", 2)),
480            Arc::new(Column::new("a", 0)),
481        ];
482        assert!(physical_exprs_bag_equal(list3.as_slice(), list4.as_slice()));
483        assert!(physical_exprs_bag_equal(list4.as_slice(), list3.as_slice()));
484        assert!(physical_exprs_bag_equal(list3.as_slice(), list3.as_slice()));
485        assert!(physical_exprs_bag_equal(list4.as_slice(), list4.as_slice()));
486        assert!(!physical_exprs_equal(list3.as_slice(), list4.as_slice()));
487        assert!(!physical_exprs_equal(list4.as_slice(), list3.as_slice()));
488        assert!(physical_exprs_bag_equal(list3.as_slice(), list3.as_slice()));
489        assert!(physical_exprs_bag_equal(list4.as_slice(), list4.as_slice()));
490    }
491
492    #[test]
493    fn test_is_volatile_default_behavior() {
494        // Test that default PhysicalExpr implementations are not volatile
495        let literal =
496            Arc::new(Literal::new(ScalarValue::Int32(Some(42)))) as Arc<dyn PhysicalExpr>;
497        let column = Arc::new(Column::new("test", 0)) as Arc<dyn PhysicalExpr>;
498
499        // Test is_volatile_node() - should return false by default
500        assert!(!literal.is_volatile_node());
501        assert!(!column.is_volatile_node());
502
503        // Test is_volatile() - should return false for non-volatile expressions
504        assert!(!is_volatile(&literal));
505        assert!(!is_volatile(&column));
506    }
507
508    /// Mock volatile PhysicalExpr for testing purposes
509    #[derive(Debug, Clone, PartialEq, Eq, Hash)]
510    struct MockVolatileExpr {
511        volatile: bool,
512    }
513
514    impl MockVolatileExpr {
515        fn new(volatile: bool) -> Self {
516            Self { volatile }
517        }
518    }
519
520    impl fmt::Display for MockVolatileExpr {
521        fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
522            write!(f, "MockVolatile({})", self.volatile)
523        }
524    }
525
526    impl PhysicalExpr for MockVolatileExpr {
527        fn data_type(&self, _input_schema: &Schema) -> Result<DataType> {
528            Ok(DataType::Boolean)
529        }
530
531        fn nullable(&self, _input_schema: &Schema) -> Result<bool> {
532            Ok(false)
533        }
534
535        fn evaluate(&self, _batch: &RecordBatch) -> Result<ColumnarValue> {
536            Ok(ColumnarValue::Scalar(ScalarValue::Boolean(Some(
537                self.volatile,
538            ))))
539        }
540
541        fn children(&self) -> Vec<&Arc<dyn PhysicalExpr>> {
542            vec![]
543        }
544
545        fn with_new_children(
546            self: Arc<Self>,
547            _children: Vec<Arc<dyn PhysicalExpr>>,
548        ) -> Result<Arc<dyn PhysicalExpr>> {
549            Ok(self)
550        }
551
552        fn is_volatile_node(&self) -> bool {
553            self.volatile
554        }
555
556        fn fmt_sql(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
557            write!(f, "mock_volatile({})", self.volatile)
558        }
559    }
560
561    #[test]
562    fn test_nested_expression_volatility() {
563        // Test that is_volatile() recursively detects volatility in expression trees
564
565        // Create a volatile mock expression
566        let volatile_expr =
567            Arc::new(MockVolatileExpr::new(true)) as Arc<dyn PhysicalExpr>;
568        assert!(volatile_expr.is_volatile_node());
569        assert!(is_volatile(&volatile_expr));
570
571        // Create a non-volatile mock expression
572        let stable_expr = Arc::new(MockVolatileExpr::new(false)) as Arc<dyn PhysicalExpr>;
573        assert!(!stable_expr.is_volatile_node());
574        assert!(!is_volatile(&stable_expr));
575
576        // Create a literal (non-volatile)
577        let literal =
578            Arc::new(Literal::new(ScalarValue::Int32(Some(42)))) as Arc<dyn PhysicalExpr>;
579        assert!(!literal.is_volatile_node());
580        assert!(!is_volatile(&literal));
581
582        // Test composite expression: volatile_expr AND literal
583        // The BinaryExpr itself is not volatile, but contains a volatile child
584        let composite_expr = Arc::new(BinaryExpr::new(
585            Arc::clone(&volatile_expr),
586            Operator::And,
587            Arc::clone(&literal),
588        )) as Arc<dyn PhysicalExpr>;
589
590        assert!(!composite_expr.is_volatile_node()); // BinaryExpr itself is not volatile
591        assert!(is_volatile(&composite_expr)); // But it contains a volatile child
592
593        // Test composite expression with all non-volatile children
594        let stable_composite = Arc::new(BinaryExpr::new(
595            Arc::clone(&stable_expr),
596            Operator::And,
597            Arc::clone(&literal),
598        )) as Arc<dyn PhysicalExpr>;
599
600        assert!(!stable_composite.is_volatile_node());
601        assert!(!is_volatile(&stable_composite)); // No volatile children
602    }
603}
datafusion_physical_expr/physical_expr.rs

datafusion_physical_expr/
physical_expr.rs
