datafusion_physical_expr_common/
datum.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::array::BooleanArray;
19use arrow::array::{ArrayRef, Datum, make_comparator};
20use arrow::buffer::{BooleanBuffer, NullBuffer};
21use arrow::compute::kernels::cmp::{
22    distinct, eq, gt, gt_eq, lt, lt_eq, neq, not_distinct,
23};
24use arrow::compute::{SortOptions, ilike, like, nilike, nlike};
25use arrow::error::ArrowError;
26use datafusion_common::utils::{normalize_float_zero, normalize_float_zero_scalar};
27use datafusion_common::{Result, ScalarValue};
28use datafusion_common::{arrow_datafusion_err, assert_or_internal_err, internal_err};
29use datafusion_expr_common::columnar_value::ColumnarValue;
30use datafusion_expr_common::operator::Operator;
31use std::sync::Arc;
32
33/// Applies a binary [`Datum`] kernel `f` to `lhs` and `rhs`
34///
35/// This maps arrow-rs' [`Datum`] kernels to DataFusion's [`ColumnarValue`] abstraction
36pub fn apply(
37    lhs: &ColumnarValue,
38    rhs: &ColumnarValue,
39    f: impl Fn(&dyn Datum, &dyn Datum) -> Result<ArrayRef, ArrowError>,
40) -> Result<ColumnarValue> {
41    match (&lhs, &rhs) {
42        (ColumnarValue::Array(left), ColumnarValue::Array(right)) => {
43            Ok(ColumnarValue::Array(f(&left.as_ref(), &right.as_ref())?))
44        }
45        (ColumnarValue::Scalar(left), ColumnarValue::Array(right)) => Ok(
46            ColumnarValue::Array(f(&left.to_scalar()?, &right.as_ref())?),
47        ),
48        (ColumnarValue::Array(left), ColumnarValue::Scalar(right)) => Ok(
49            ColumnarValue::Array(f(&left.as_ref(), &right.to_scalar()?)?),
50        ),
51        (ColumnarValue::Scalar(left), ColumnarValue::Scalar(right)) => {
52            let array = f(&left.to_scalar()?, &right.to_scalar()?)?;
53            let scalar = ScalarValue::try_from_array(array.as_ref(), 0)?;
54            Ok(ColumnarValue::Scalar(scalar))
55        }
56    }
57}
58
59/// Applies a binary [`Datum`] comparison operator `op` to `lhs` and `rhs`
60pub fn apply_cmp(
61    op: Operator,
62    lhs: &ColumnarValue,
63    rhs: &ColumnarValue,
64) -> Result<ColumnarValue> {
65    if lhs.data_type().is_nested() {
66        apply_cmp_for_nested(op, lhs, rhs)
67    } else {
68        let f = match op {
69            Operator::Eq => eq,
70            Operator::NotEq => neq,
71            Operator::Lt => lt,
72            Operator::LtEq => lt_eq,
73            Operator::Gt => gt,
74            Operator::GtEq => gt_eq,
75            Operator::IsDistinctFrom => distinct,
76            Operator::IsNotDistinctFrom => not_distinct,
77
78            Operator::LikeMatch => like,
79            Operator::ILikeMatch => ilike,
80            Operator::NotLikeMatch => nlike,
81            Operator::NotILikeMatch => nilike,
82
83            _ => {
84                return internal_err!("Invalid compare operator: {}", op);
85            }
86        };
87
88        // Arrow's comparison kernels use IEEE 754 totalOrder semantics for
89        // floats, which treats `-0.0` and `+0.0` as distinct. Normalize float
90        // operands so SQL semantics (`+0.0 == -0.0`) hold. No-op for
91        // non-float types.
92        let lhs = normalize_cmp_input(lhs);
93        let rhs = normalize_cmp_input(rhs);
94        apply(&lhs, &rhs, |l, r| Ok(Arc::new(f(l, r)?)))
95    }
96}
97
98fn normalize_cmp_input(cv: &ColumnarValue) -> ColumnarValue {
99    match cv {
100        ColumnarValue::Array(a) => ColumnarValue::Array(normalize_float_zero(a)),
101        ColumnarValue::Scalar(s) => {
102            ColumnarValue::Scalar(normalize_float_zero_scalar(s.clone()))
103        }
104    }
105}
106
107/// Applies a binary [`Datum`] comparison operator `op` to `lhs` and `rhs` for nested type like
108/// List, FixedSizeList, LargeList, Struct, Union, Map, or a dictionary of a nested type
109pub fn apply_cmp_for_nested(
110    op: Operator,
111    lhs: &ColumnarValue,
112    rhs: &ColumnarValue,
113) -> Result<ColumnarValue> {
114    let left_data_type = lhs.data_type();
115    let right_data_type = rhs.data_type();
116
117    assert_or_internal_err!(
118        matches!(
119            op,
120            Operator::Eq
121                | Operator::NotEq
122                | Operator::Lt
123                | Operator::Gt
124                | Operator::LtEq
125                | Operator::GtEq
126                | Operator::IsDistinctFrom
127                | Operator::IsNotDistinctFrom
128        ) && left_data_type.equals_datatype(&right_data_type),
129        "invalid operator or data type mismatch for nested data, op {op} left {left_data_type}, right {right_data_type}",
130    );
131
132    apply(lhs, rhs, |l, r| {
133        Ok(Arc::new(compare_op_for_nested(op, l, r)?))
134    })
135}
136
137/// Compare with eq with either nested or non-nested
138pub fn compare_with_eq(
139    lhs: &dyn Datum,
140    rhs: &dyn Datum,
141    is_nested: bool,
142) -> Result<BooleanArray> {
143    if is_nested {
144        compare_op_for_nested(Operator::Eq, lhs, rhs)
145    } else {
146        eq(lhs, rhs).map_err(|e| arrow_datafusion_err!(e))
147    }
148}
149
150/// Compare on nested type List, Struct, and so on
151pub fn compare_op_for_nested(
152    op: Operator,
153    lhs: &dyn Datum,
154    rhs: &dyn Datum,
155) -> Result<BooleanArray> {
156    let (l, is_l_scalar) = lhs.get();
157    let (r, is_r_scalar) = rhs.get();
158    let l_len = l.len();
159    let r_len = r.len();
160
161    assert_or_internal_err!(l_len == r_len || is_l_scalar || is_r_scalar, "len mismatch");
162
163    let len = match is_l_scalar {
164        true => r_len,
165        false => l_len,
166    };
167
168    // fast path, if compare with one null and operator is not 'distinct', then we can return null array directly
169    if !matches!(op, Operator::IsDistinctFrom | Operator::IsNotDistinctFrom)
170        && (is_l_scalar && l.null_count() == 1 || is_r_scalar && r.null_count() == 1)
171    {
172        return Ok(BooleanArray::new_null(len));
173    }
174
175    // TODO: make SortOptions configurable
176    // we choose the default behaviour from arrow-rs which has null-first that follow spark's behaviour
177    let cmp = make_comparator(l, r, SortOptions::default())?;
178
179    let cmp_with_op = |i, j| match op {
180        Operator::Eq | Operator::IsNotDistinctFrom => cmp(i, j).is_eq(),
181        Operator::Lt => cmp(i, j).is_lt(),
182        Operator::Gt => cmp(i, j).is_gt(),
183        Operator::LtEq => !cmp(i, j).is_gt(),
184        Operator::GtEq => !cmp(i, j).is_lt(),
185        Operator::NotEq | Operator::IsDistinctFrom => !cmp(i, j).is_eq(),
186        _ => unreachable!("unexpected operator found"),
187    };
188
189    let values = match (is_l_scalar, is_r_scalar) {
190        (false, false) => BooleanBuffer::collect_bool(len, |i| cmp_with_op(i, i)),
191        (true, false) => BooleanBuffer::collect_bool(len, |i| cmp_with_op(0, i)),
192        (false, true) => BooleanBuffer::collect_bool(len, |i| cmp_with_op(i, 0)),
193        (true, true) => std::iter::once(cmp_with_op(0, 0)).collect(),
194    };
195
196    // Distinct understand how to compare with NULL
197    // i.e NULL is distinct from NULL -> false
198    if matches!(op, Operator::IsDistinctFrom | Operator::IsNotDistinctFrom) {
199        Ok(BooleanArray::new(values, None))
200    } else {
201        // If one of the side is NULL, we return NULL
202        // i.e. NULL eq NULL -> NULL
203        // For nested comparisons, we need to ensure the null buffer matches the result length
204        let nulls = match (is_l_scalar, is_r_scalar) {
205            (false, false) | (true, true) => NullBuffer::union(l.nulls(), r.nulls()),
206            (true, false) => {
207                // When left is null-scalar and right is array, expand left nulls to match result length
208                match l.nulls().filter(|nulls| nulls.is_null(0)) {
209                    Some(_) => Some(NullBuffer::new_null(len)), // Left scalar is null
210                    None => r.nulls().cloned(),                 // Left scalar is non-null
211                }
212            }
213            (false, true) => {
214                // When right is null-scalar and left is array, expand right nulls to match result length
215                match r.nulls().filter(|nulls| nulls.is_null(0)) {
216                    Some(_) => Some(NullBuffer::new_null(len)), // Right scalar is null
217                    None => l.nulls().cloned(), // Right scalar is non-null
218                }
219            }
220        };
221        Ok(BooleanArray::new(values, nulls))
222    }
223}
datafusion_physical_expr_common/datum.rs

datafusion_physical_expr_common/
datum.rs
