Google Git
Sign in
apache/datafusion-sandbox/refs/heads/bash-push/./datafusion/physical-expr/src/expressions/case.rs
blob: 7a33aa95c56b5ecc58d687a5aade5d0f121151a3 [file]
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
use super::{Column, Literal};
use crate::expressions::case::ResultState::{Complete, Empty, Partial};
use crate::expressions::try_cast;
use crate::PhysicalExpr;
use arrow::array::*;
use arrow::compute::kernels::zip::zip;
use arrow::compute::{
is_not_null, not, nullif, prep_null_mask_filter, FilterBuilder, FilterPredicate,
SlicesIterator,
};
use arrow::datatypes::{DataType, Schema, UInt32Type, UnionMode};
use arrow::error::ArrowError;
use datafusion_common::cast::as_boolean_array;
use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRecursion};
use datafusion_common::{
exec_err, internal_datafusion_err, internal_err, DataFusionError, HashMap, HashSet,
Result, ScalarValue,
};
use datafusion_expr::ColumnarValue;
use datafusion_physical_expr_common::datum::compare_with_eq;
use itertools::Itertools;
use std::borrow::Cow;
use std::fmt::{Debug, Formatter};
use std::hash::Hash;
use std::{any::Any, sync::Arc};
type WhenThen = (Arc<dyn PhysicalExpr>, Arc<dyn PhysicalExpr>);
#[derive(Debug, Hash, PartialEq, Eq)]
enum EvalMethod {
/// CASE WHEN condition THEN result
/// [WHEN ...]
/// [ELSE result]
/// END
NoExpression(ProjectedCaseBody),
/// CASE expression
/// WHEN value THEN result
/// [WHEN ...]
/// [ELSE result]
/// END
WithExpression(ProjectedCaseBody),
/// This is a specialization for a specific use case where we can take a fast path
/// for expressions that are infallible and can be cheaply computed for the entire
/// record batch rather than just for the rows where the predicate is true.
///
/// CASE WHEN condition THEN column [ELSE NULL] END
InfallibleExprOrNull,
/// This is a specialization for a specific use case where we can take a fast path
/// if there is just one when/then pair and both the `then` and `else` expressions
/// are literal values
/// CASE WHEN condition THEN literal ELSE literal END
ScalarOrScalar,
/// This is a specialization for a specific use case where we can take a fast path
/// if there is just one when/then pair and both the `then` and `else` are expressions
///
/// CASE WHEN condition THEN expression ELSE expression END
ExpressionOrExpression(ProjectedCaseBody),
}
/// The body of a CASE expression which consists of an optional base expression, the "when/then"
/// branches and an optional "else" branch.
#[derive(Debug, Hash, PartialEq, Eq)]
struct CaseBody {
/// Optional base expression that can be compared to literal values in the "when" expressions
expr: Option<Arc<dyn PhysicalExpr>>,
/// One or more when/then expressions
when_then_expr: Vec<WhenThen>,
/// Optional "else" expression
else_expr: Option<Arc<dyn PhysicalExpr>>,
}
impl CaseBody {
/// Derives a [ProjectedCaseBody] from this [CaseBody].
fn project(&self) -> Result<ProjectedCaseBody> {
// Determine the set of columns that are used in all the expressions of the case body.
let mut used_column_indices = HashSet::<usize>::new();
let mut collect_column_indices = |expr: &Arc<dyn PhysicalExpr>| {
expr.apply(|expr| {
if let Some(column) = expr.as_any().downcast_ref::<Column>() {
used_column_indices.insert(column.index());
}
Ok(TreeNodeRecursion::Continue)
})
.expect("Closure cannot fail");
};
if let Some(e) = &self.expr {
collect_column_indices(e);
}
self.when_then_expr.iter().for_each(|(w, t)| {
collect_column_indices(w);
collect_column_indices(t);
});
if let Some(e) = &self.else_expr {
collect_column_indices(e);
}
// Construct a mapping from the original column index to the projected column index.
let column_index_map = used_column_indices
.iter()
.enumerate()
.map(|(projected, original)| (*original, projected))
.collect::<HashMap<usize, usize>>();
// Construct the projected body by rewriting each expression from the original body
// using the column index mapping.
let project = |expr: &Arc<dyn PhysicalExpr>| -> Result<Arc<dyn PhysicalExpr>> {
Arc::clone(expr)
.transform_down(|e| {
if let Some(column) = e.as_any().downcast_ref::<Column>() {
let original = column.index();
let projected = *column_index_map.get(&original).unwrap();
if projected != original {
return Ok(Transformed::yes(Arc::new(Column::new(
column.name(),
projected,
))));
}
}
Ok(Transformed::no(e))
})
.map(|t| t.data)
};
let projected_body = CaseBody {
expr: self.expr.as_ref().map(project).transpose()?,
when_then_expr: self
.when_then_expr
.iter()
.map(|(e, t)| Ok((project(e)?, project(t)?)))
.collect::<Result<Vec<_>>>()?,
else_expr: self.else_expr.as_ref().map(project).transpose()?,
};
// Construct the projection vector
let projection = column_index_map
.iter()
.sorted_by_key(|(_, v)| **v)
.map(|(k, _)| *k)
.collect::<Vec<_>>();
Ok(ProjectedCaseBody {
projection,
body: projected_body,
})
}
}
/// A derived case body that can be used to evaluate a case expression after projecting
/// record batches using a projection vector.
///
/// This is used to avoid filtering columns that are not used in the
/// input `RecordBatch` when progressively evaluating a `CASE` expression's
/// remainder batches. Filtering these columns is wasteful since for a record
/// batch of `n` rows, filtering requires at worst a copy of `n - 1` values
/// per array. If these filtered values will never be accessed, the time spent
/// producing them is better avoided.
///
/// For example, if we are evaluating the following case expression that
/// only references columns B and D:
///
/// ```sql
/// SELECT CASE WHEN B > 10 THEN D ELSE NULL END FROM (VALUES (...)) T(A, B, C, D)
/// ```
///
/// Of the 4 input columns `[A, B, C, D]`, the `CASE` expression only access `B` and `D`.
/// Filtering `A` and `C` would be unnecessary and wasteful.
///
/// If we only retain columns `B` and `D` using `RecordBatch::project` and the projection vector
/// `[1, 3]`, the indices of these two columns will change to `[0, 1]`. To evaluate the
/// case expression, it will need to be rewritten from `CASE WHEN B@1 > 10 THEN D@3 ELSE NULL END`
/// to `CASE WHEN B@0 > 10 THEN D@1 ELSE NULL END`.
///
/// The projection vector and the rewritten expression (which only differs from the original in
/// column reference indices) are held in a `ProjectedCaseBody`.
#[derive(Debug, Hash, PartialEq, Eq)]
struct ProjectedCaseBody {
projection: Vec<usize>,
body: CaseBody,
}
/// The CASE expression is similar to a series of nested if/else and there are two forms that
/// can be used. The first form consists of a series of boolean "when" expressions with
/// corresponding "then" expressions, and an optional "else" expression.
///
/// CASE WHEN condition THEN result
/// [WHEN ...]
/// [ELSE result]
/// END
///
/// The second form uses a base expression and then a series of "when" clauses that match on a
/// literal value.
///
/// CASE expression
/// WHEN value THEN result
/// [WHEN ...]
/// [ELSE result]
/// END
#[derive(Debug, Hash, PartialEq, Eq)]
pub struct CaseExpr {
/// The case expression body
body: CaseBody,
/// Evaluation method to use
eval_method: EvalMethod,
}
impl std::fmt::Display for CaseExpr {
fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
write!(f, "CASE ")?;
if let Some(e) = &self.body.expr {
write!(f, "{e} ")?;
}
for (w, t) in &self.body.when_then_expr {
write!(f, "WHEN {w} THEN {t} ")?;
}
if let Some(e) = &self.body.else_expr {
write!(f, "ELSE {e} ")?;
}
write!(f, "END")
}
}
/// This is a specialization for a specific use case where we can take a fast path
/// for expressions that are infallible and can be cheaply computed for the entire
/// record batch rather than just for the rows where the predicate is true. For now,
/// this is limited to use with Column expressions but could potentially be used for other
/// expressions in the future
fn is_cheap_and_infallible(expr: &Arc<dyn PhysicalExpr>) -> bool {
expr.as_any().is::<Column>()
}
/// Creates a [FilterPredicate] from a boolean array.
fn create_filter(predicate: &BooleanArray, optimize: bool) -> FilterPredicate {
let mut filter_builder = FilterBuilder::new(predicate);
if optimize {
// Always optimize the filter since we use them multiple times.
filter_builder = filter_builder.optimize();
}
filter_builder.build()
}
fn multiple_arrays(data_type: &DataType) -> bool {
match data_type {
DataType::Struct(fields) => {
fields.len() > 1
|| fields.len() == 1 && multiple_arrays(fields[0].data_type())
}
DataType::Union(fields, UnionMode::Sparse) => !fields.is_empty(),
_ => false,
}
}
// This should be removed when https://github.com/apache/arrow-rs/pull/8693
// is merged and becomes available.
fn filter_record_batch(
record_batch: &RecordBatch,
filter: &FilterPredicate,
) -> std::result::Result<RecordBatch, ArrowError> {
let filtered_columns = record_batch
.columns()
.iter()
.map(|a| filter_array(a, filter))
.collect::<std::result::Result<Vec<_>, _>>()?;
// SAFETY: since we start from a valid RecordBatch, there's no need to revalidate the schema
// since the set of columns has not changed.
// The input column arrays all had the same length (since they're coming from a valid RecordBatch)
// and the filtering them with the same filter will produces a new set of arrays with identical
// lengths.
unsafe {
Ok(RecordBatch::new_unchecked(
record_batch.schema(),
filtered_columns,
filter.count(),
))
}
}
// This function exists purely to be able to use the same call style
// for `filter_record_batch` and `filter_array` at the point of use.
// When https://github.com/apache/arrow-rs/pull/8693 is available, replace
// both with method calls on `FilterPredicate`.
#[inline(always)]
fn filter_array(
array: &dyn Array,
filter: &FilterPredicate,
) -> std::result::Result<ArrayRef, ArrowError> {
filter.filter(array)
}
fn merge(
mask: &BooleanArray,
truthy: ColumnarValue,
falsy: ColumnarValue,
) -> std::result::Result<ArrayRef, ArrowError> {
let (truthy, truthy_is_scalar) = match truthy {
ColumnarValue::Array(a) => (a, false),
ColumnarValue::Scalar(s) => (s.to_array()?, true),
};
let (falsy, falsy_is_scalar) = match falsy {
ColumnarValue::Array(a) => (a, false),
ColumnarValue::Scalar(s) => (s.to_array()?, true),
};
if truthy_is_scalar && falsy_is_scalar {
return zip(mask, &Scalar::new(truthy), &Scalar::new(falsy));
}
let falsy = falsy.to_data();
let truthy = truthy.to_data();
let mut mutable = MutableArrayData::new(vec![&truthy, &falsy], false, truthy.len());
// the SlicesIterator slices only the true values. So the gaps left by this iterator we need to
// fill with falsy values
// keep track of how much is filled
let mut filled = 0;
let mut falsy_offset = 0;
let mut truthy_offset = 0;
SlicesIterator::new(mask).for_each(|(start, end)| {
// the gap needs to be filled with falsy values
if start > filled {
if falsy_is_scalar {
for _ in filled..start {
// Copy the first item from the 'falsy' array into the output buffer.
mutable.extend(1, 0, 1);
}
} else {
let falsy_length = start - filled;
let falsy_end = falsy_offset + falsy_length;
mutable.extend(1, falsy_offset, falsy_end);
falsy_offset = falsy_end;
}
}
// fill with truthy values
if truthy_is_scalar {
for _ in start..end {
// Copy the first item from the 'truthy' array into the output buffer.
mutable.extend(0, 0, 1);
}
} else {
let truthy_length = end - start;
let truthy_end = truthy_offset + truthy_length;
mutable.extend(0, truthy_offset, truthy_end);
truthy_offset = truthy_end;
}
filled = end;
});
// the remaining part is falsy
if filled < mask.len() {
if falsy_is_scalar {
for _ in filled..mask.len() {
// Copy the first item from the 'falsy' array into the output buffer.
mutable.extend(1, 0, 1);
}
} else {
let falsy_length = mask.len() - filled;
let falsy_end = falsy_offset + falsy_length;
mutable.extend(1, falsy_offset, falsy_end);
}
}
let data = mutable.freeze();
Ok(make_array(data))
}
/// Merges elements by index from a list of [`ArrayData`], creating a new [`ColumnarValue`] from
/// those values.
///
/// Each element in `indices` is the index of an array in `values`. The `indices` array is processed
/// sequentially. The first occurrence of index value `n` will be mapped to the first
/// value of the array at index `n`. The second occurrence to the second value, and so on.
/// An index value where `PartialResultIndex::is_none` is `true` is used to indicate null values.
///
/// # Implementation notes
///
/// This algorithm is similar in nature to both `zip` and `interleave`, but there are some important
/// differences.
///
/// In contrast to `zip`, this function supports multiple input arrays. Instead of a boolean
/// selection vector, an index array is to take values from the input arrays, and a special marker
/// value is used to indicate null values.
///
/// In contrast to `interleave`, this function does not use pairs of indices. The values in
/// `indices` serve the same purpose as the first value in the pairs passed to `interleave`.
/// The index in the array is implicit and is derived from the number of times a particular array
/// index occurs.
/// The more constrained indexing mechanism used by this algorithm makes it easier to copy values
/// in contiguous slices. In the example below, the two subsequent elements from array `2` can be
/// copied in a single operation from the source array instead of copying them one by one.
/// Long spans of null values are also especially cheap because they do not need to be represented
/// in an input array.
///
/// # Safety
///
/// This function does not check that the number of occurrences of any particular array index matches
/// the length of the corresponding input array. If an array contains more values than required, the
/// spurious values will be ignored. If an array contains fewer values than necessary, this function
/// will panic.
///
/// # Example
///
/// ```text
/// ┌───────────┐ ┌─────────┐ ┌─────────┐
/// │┌─────────┐│ │ None │ │ NULL │
/// ││ A ││ ├─────────┤ ├─────────┤
/// │└─────────┘│ │ 1 │ │ B │
/// │┌─────────┐│ ├─────────┤ ├─────────┤
/// ││ B ││ │ 0 │ merge(values, indices) │ A │
/// │└─────────┘│ ├─────────┤ ─────────────────────────▶ ├─────────┤
/// │┌─────────┐│ │ None │ │ NULL │
/// ││ C ││ ├─────────┤ ├─────────┤
/// │├─────────┤│ │ 2 │ │ C │
/// ││ D ││ ├─────────┤ ├─────────┤
/// │└─────────┘│ │ 2 │ │ D │
/// └───────────┘ └─────────┘ └─────────┘
/// values indices result
/// ```
fn merge_n(values: &[ArrayData], indices: &[PartialResultIndex]) -> Result<ArrayRef> {
#[cfg(debug_assertions)]
for ix in indices {
if let Some(index) = ix.index() {
assert!(
index < values.len(),
"Index out of bounds: {} >= {}",
index,
values.len()
);
}
}
let data_refs = values.iter().collect();
let mut mutable = MutableArrayData::new(data_refs, true, indices.len());
// This loop extends the mutable array by taking slices from the partial results.
//
// take_offsets keeps track of how many values have been taken from each array.
let mut take_offsets = vec![0; values.len() + 1];
let mut start_row_ix = 0;
loop {
let array_ix = indices[start_row_ix];
// Determine the length of the slice to take.
let mut end_row_ix = start_row_ix + 1;
while end_row_ix < indices.len() && indices[end_row_ix] == array_ix {
end_row_ix += 1;
}
let slice_length = end_row_ix - start_row_ix;
// Extend mutable with either nulls or with values from the array.
match array_ix.index() {
None => mutable.extend_nulls(slice_length),
Some(index) => {
let start_offset = take_offsets[index];
let end_offset = start_offset + slice_length;
mutable.extend(index, start_offset, end_offset);
take_offsets[index] = end_offset;
}
}
if end_row_ix == indices.len() {
break;
} else {
// Set the start_row_ix for the next slice.
start_row_ix = end_row_ix;
}
}
Ok(make_array(mutable.freeze()))
}
/// An index into the partial results array that's more compact than `usize`.
///
/// `u32::MAX` is reserved as a special 'none' value. This is used instead of
/// `Option` to keep the array of indices as compact as possible.
#[derive(Copy, Clone, PartialEq, Eq)]
struct PartialResultIndex {
index: u32,
}
const NONE_VALUE: u32 = u32::MAX;
impl PartialResultIndex {
/// Returns the 'none' placeholder value.
fn none() -> Self {
Self { index: NONE_VALUE }
}
fn zero() -> Self {
Self { index: 0 }
}
/// Creates a new partial result index.
///
/// If the provided value is greater than or equal to `u32::MAX`
/// an error will be returned.
fn try_new(index: usize) -> Result<Self> {
let Ok(index) = u32::try_from(index) else {
return internal_err!("Partial result index exceeds limit");
};
if index == NONE_VALUE {
return internal_err!("Partial result index exceeds limit");
}
Ok(Self { index })
}
/// Determines if this index is the 'none' placeholder value or not.
fn is_none(&self) -> bool {
self.index == NONE_VALUE
}
/// Returns `Some(index)` if this value is not the 'none' placeholder, `None` otherwise.
fn index(&self) -> Option<usize> {
if self.is_none() {
None
} else {
Some(self.index as usize)
}
}
}
impl Debug for PartialResultIndex {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
if self.is_none() {
write!(f, "null")
} else {
write!(f, "{}", self.index)
}
}
}
enum ResultState {
/// The final result is an array containing only null values.
Empty,
/// The final result needs to be computed by merging the data in `arrays`.
Partial {
// A `Vec` of partial results that should be merged.
// `partial_result_indices` contains indexes into this vec.
arrays: Vec<ArrayData>,
// Indicates per result row from which array in `partial_results` a value should be taken.
indices: Vec<PartialResultIndex>,
},
/// A single branch matched all input rows. When creating the final result, no further merging
/// of partial results is necessary.
Complete(ColumnarValue),
}
/// A builder for constructing result arrays for CASE expressions.
///
/// Rather than building a monolithic array containing all results, it maintains a set of
/// partial result arrays and a mapping that indicates for each row which partial array
/// contains the result value for that row.
///
/// On finish(), the builder will merge all partial results into a single array if necessary.
/// If all rows evaluated to the same array, that array can be returned directly without
/// any merging overhead.
struct ResultBuilder {
data_type: DataType,
/// The number of rows in the final result.
row_count: usize,
state: ResultState,
}
impl ResultBuilder {
/// Creates a new ResultBuilder that will produce arrays of the given data type.
///
/// The `row_count` parameter indicates the number of rows in the final result.
fn new(data_type: &DataType, row_count: usize) -> Self {
Self {
data_type: data_type.clone(),
row_count,
state: Empty,
}
}
/// Adds a result for one branch of the case expression.
///
/// `row_indices` should be a [UInt32Array] containing [RecordBatch] relative row indices
/// for which `value` contains result values.
///
/// If `value` is a scalar, the scalar value will be used as the value for each row in `row_indices`.
///
/// If `value` is an array, the values from the array and the indices from `row_indices` will be
/// processed pairwise. The lengths of `value` and `row_indices` must match.
///
/// The diagram below shows a situation where a when expression matched rows 1 and 4 of the
/// record batch. The then expression produced the value array `[A, D]`.
/// After adding this result, the result array will have been added to `partial arrays` and
/// `partial indices` will have been updated at indexes `1` and `4`.
///
/// ```text
/// ┌─────────┐ ┌─────────┐┌───────────┐ ┌─────────┐┌───────────┐
/// │ C │ │ 0: None ││┌ 0 ──────┐│ │ 0: None ││┌ 0 ──────┐│
/// ├─────────┤ ├─────────┤││ A ││ ├─────────┤││ A ││
/// │ D │ │ 1: None ││└─────────┘│ │ 1: 2 ││└─────────┘│
/// └─────────┘ ├─────────┤│┌ 1 ──────┐│ add_branch_result( ├─────────┤│┌ 1 ──────┐│
/// matching │ 2: 0 │││ B ││ row indices, │ 2: 0 │││ B ││
/// 'then' values ├─────────┤│└─────────┘│ value ├─────────┤│└─────────┘│
/// │ 3: None ││ │ ) │ 3: None ││┌ 2 ──────┐│
/// ┌─────────┐ ├─────────┤│ │ ─────────────────────────▶ ├─────────┤││ C ││
/// │ 1 │ │ 4: None ││ │ │ 4: 2 ││├─────────┤│
/// ├─────────┤ ├─────────┤│ │ ├─────────┤││ D ││
/// │ 4 │ │ 5: 1 ││ │ │ 5: 1 ││└─────────┘│
/// └─────────┘ └─────────┘└───────────┘ └─────────┘└───────────┘
/// row indices partial partial partial partial
/// indices arrays indices arrays
/// ```
fn add_branch_result(
&mut self,
row_indices: &ArrayRef,
value: ColumnarValue,
) -> Result<()> {
match value {
ColumnarValue::Array(a) => {
if a.len() != row_indices.len() {
internal_err!("Array length must match row indices length")
} else if row_indices.len() == self.row_count {
self.set_complete_result(ColumnarValue::Array(a))
} else {
self.add_partial_result(row_indices, a.to_data())
}
}
ColumnarValue::Scalar(s) => {
if row_indices.len() == self.row_count {
self.set_complete_result(ColumnarValue::Scalar(s))
} else {
self.add_partial_result(
row_indices,
s.to_array_of_size(row_indices.len())?.to_data(),
)
}
}
}
}
/// Adds a partial result array.
///
/// This method adds the given array data as a partial result and updates the index mapping
/// to indicate that the specified rows should take their values from this array.
/// The partial results will be merged into a single array when finish() is called.
fn add_partial_result(
&mut self,
row_indices: &ArrayRef,
row_values: ArrayData,
) -> Result<()> {
if row_indices.null_count() != 0 {
return internal_err!("Row indices must not contain nulls");
}
match &mut self.state {
Empty => {
let array_index = PartialResultIndex::zero();
let mut indices = vec![PartialResultIndex::none(); self.row_count];
for row_ix in row_indices.as_primitive::<UInt32Type>().values().iter() {
indices[*row_ix as usize] = array_index;
}
self.state = Partial {
arrays: vec![row_values],
indices,
};
Ok(())
}
Partial { arrays, indices } => {
let array_index = PartialResultIndex::try_new(arrays.len())?;
arrays.push(row_values);
for row_ix in row_indices.as_primitive::<UInt32Type>().values().iter() {
// This is check is only active for debug config because the callers of this method,
// `case_when_with_expr` and `case_when_no_expr`, already ensure that
// they only calculate a value for each row at most once.
#[cfg(debug_assertions)]
if !indices[*row_ix as usize].is_none() {
return internal_err!("Duplicate value for row {}", *row_ix);
}
indices[*row_ix as usize] = array_index;
}
Ok(())
}
Complete(_) => internal_err!(
"Cannot add a partial result when complete result is already set"
),
}
}
/// Sets a result that applies to all rows.
///
/// This is an optimization for cases where all rows evaluate to the same result.
/// When a complete result is set, the builder will return it directly from finish()
/// without any merging overhead.
fn set_complete_result(&mut self, value: ColumnarValue) -> Result<()> {
match &self.state {
Empty => {
self.state = Complete(value);
Ok(())
}
Partial { .. } => {
internal_err!(
"Cannot set a complete result when there are already partial results"
)
}
Complete(_) => internal_err!("Complete result already set"),
}
}
/// Finishes building the result and returns the final array.
fn finish(self) -> Result<ColumnarValue> {
match self.state {
Empty => {
// No complete result and no partial results.
// This can happen for case expressions with no else branch where no rows
// matched.
Ok(ColumnarValue::Scalar(ScalarValue::try_new_null(
&self.data_type,
)?))
}
Partial { arrays, indices } => {
// Merge partial results into a single array.
Ok(ColumnarValue::Array(merge_n(&arrays, &indices)?))
}
Complete(v) => {
// If we have a complete result, we can just return it.
Ok(v)
}
}
}
}
impl CaseExpr {
/// Create a new CASE WHEN expression
pub fn try_new(
expr: Option<Arc<dyn PhysicalExpr>>,
when_then_expr: Vec<WhenThen>,
else_expr: Option<Arc<dyn PhysicalExpr>>,
) -> Result<Self> {
// normalize null literals to None in the else_expr (this already happens
// during SQL planning, but not necessarily for other use cases)
let else_expr = match &else_expr {
Some(e) => match e.as_any().downcast_ref::<Literal>() {
Some(lit) if lit.value().is_null() => None,
_ => else_expr,
},
_ => else_expr,
};
if when_then_expr.is_empty() {
return exec_err!("There must be at least one WHEN clause");
}
let body = CaseBody {
expr,
when_then_expr,
else_expr,
};
let eval_method = if body.expr.is_some() {
EvalMethod::WithExpression(body.project()?)
} else if body.when_then_expr.len() == 1
&& is_cheap_and_infallible(&(body.when_then_expr[0].1))
&& body.else_expr.is_none()
{
EvalMethod::InfallibleExprOrNull
} else if body.when_then_expr.len() == 1
&& body.when_then_expr[0].1.as_any().is::<Literal>()
&& body.else_expr.is_some()
&& body.else_expr.as_ref().unwrap().as_any().is::<Literal>()
{
EvalMethod::ScalarOrScalar
} else if body.when_then_expr.len() == 1 && body.else_expr.is_some() {
EvalMethod::ExpressionOrExpression(body.project()?)
} else {
EvalMethod::NoExpression(body.project()?)
};
Ok(Self { body, eval_method })
}
/// Optional base expression that can be compared to literal values in the "when" expressions
pub fn expr(&self) -> Option<&Arc<dyn PhysicalExpr>> {
self.body.expr.as_ref()
}
/// One or more when/then expressions
pub fn when_then_expr(&self) -> &[WhenThen] {
&self.body.when_then_expr
}
/// Optional "else" expression
pub fn else_expr(&self) -> Option<&Arc<dyn PhysicalExpr>> {
self.body.else_expr.as_ref()
}
}
impl CaseBody {
fn data_type(&self, input_schema: &Schema) -> Result<DataType> {
// since all then results have the same data type, we can choose any one as the
// return data type except for the null.
let mut data_type = DataType::Null;
for i in 0..self.when_then_expr.len() {
data_type = self.when_then_expr[i].1.data_type(input_schema)?;
if !data_type.equals_datatype(&DataType::Null) {
break;
}
}
// if all then results are null, we use data type of else expr instead if possible.
if data_type.equals_datatype(&DataType::Null) {
if let Some(e) = &self.else_expr {
data_type = e.data_type(input_schema)?;
}
}
Ok(data_type)
}
/// See [CaseExpr::case_when_with_expr].
fn case_when_with_expr(
&self,
batch: &RecordBatch,
return_type: &DataType,
) -> Result<ColumnarValue> {
let mut result_builder = ResultBuilder::new(return_type, batch.num_rows());
// `remainder_rows` contains the indices of the rows that need to be evaluated
let mut remainder_rows: ArrayRef =
Arc::new(UInt32Array::from_iter_values(0..batch.num_rows() as u32));
// `remainder_batch` contains the rows themselves that need to be evaluated
let mut remainder_batch = Cow::Borrowed(batch);
// evaluate the base expression
let mut base_values = self
.expr
.as_ref()
.unwrap()
.evaluate(batch)?
.into_array(batch.num_rows())?;
// Fill in a result value already for rows where the base expression value is null
// Since each when expression is tested against the base expression using the equality
// operator, null base values can never match any when expression. `x = NULL` is falsy,
// for all possible values of `x`.
if base_values.null_count() > 0 {
// Use `is_not_null` since this is a cheap clone of the null buffer from 'base_value'.
// We already checked there are nulls, so we can be sure a new buffer will not be
// created.
let base_not_nulls = is_not_null(base_values.as_ref())?;
let base_all_null = base_values.null_count() == remainder_batch.num_rows();
// If there is an else expression, use that as the default value for the null rows
// Otherwise the default `null` value from the result builder will be used.
if let Some(e) = &self.else_expr {
let expr = try_cast(Arc::clone(e), &batch.schema(), return_type.clone())?;
if base_all_null {
// All base values were null, so no need to filter
let nulls_value = expr.evaluate(&remainder_batch)?;
result_builder.add_branch_result(&remainder_rows, nulls_value)?;
} else {
// Filter out the null rows and evaluate the else expression for those
let nulls_filter = create_filter(&not(&base_not_nulls)?, true);
let nulls_batch =
filter_record_batch(&remainder_batch, &nulls_filter)?;
let nulls_rows = filter_array(&remainder_rows, &nulls_filter)?;
let nulls_value = expr.evaluate(&nulls_batch)?;
result_builder.add_branch_result(&nulls_rows, nulls_value)?;
}
}
// All base values are null, so we can return early
if base_all_null {
return result_builder.finish();
}
// Remove the null rows from the remainder batch
let not_null_filter = create_filter(&base_not_nulls, true);
remainder_batch =
Cow::Owned(filter_record_batch(&remainder_batch, &not_null_filter)?);
remainder_rows = filter_array(&remainder_rows, &not_null_filter)?;
base_values = filter_array(&base_values, &not_null_filter)?;
}
// The types of case and when expressions will be coerced to match.
// We only need to check if the base_value is nested.
let base_value_is_nested = base_values.data_type().is_nested();
for i in 0..self.when_then_expr.len() {
// Evaluate the 'when' predicate for the remainder batch
// This results in a boolean array with the same length as the remaining number of rows
let when_expr = &self.when_then_expr[i].0;
let when_value = match when_expr.evaluate(&remainder_batch)? {
ColumnarValue::Array(a) => {
compare_with_eq(&a, &base_values, base_value_is_nested)
}
ColumnarValue::Scalar(s) => {
compare_with_eq(&s.to_scalar()?, &base_values, base_value_is_nested)
}
}?;
// `true_count` ignores `true` values where the validity bit is not set, so there's
// no need to call `prep_null_mask_filter`.
let when_true_count = when_value.true_count();
// If the 'when' predicate did not match any rows, continue to the next branch immediately
if when_true_count == 0 {
continue;
}
// If the 'when' predicate matched all remaining rows, there is no need to filter
if when_true_count == remainder_batch.num_rows() {
let then_expression = &self.when_then_expr[i].1;
let then_value = then_expression.evaluate(&remainder_batch)?;
result_builder.add_branch_result(&remainder_rows, then_value)?;
return result_builder.finish();
}
// Filter the remainder batch based on the 'when' value
// This results in a batch containing only the rows that need to be evaluated
// for the current branch
// Still no need to call `prep_null_mask_filter` since `create_filter` will already do
// this unconditionally.
let then_filter = create_filter(&when_value, true);
let then_batch = filter_record_batch(&remainder_batch, &then_filter)?;
let then_rows = filter_array(&remainder_rows, &then_filter)?;
let then_expression = &self.when_then_expr[i].1;
let then_value = then_expression.evaluate(&then_batch)?;
result_builder.add_branch_result(&then_rows, then_value)?;
// If this is the last 'when' branch and there is no 'else' expression, there's no
// point in calculating the remaining rows.
if self.else_expr.is_none() && i == self.when_then_expr.len() - 1 {
return result_builder.finish();
}
// Prepare the next when branch (or the else branch)
let next_selection = match when_value.null_count() {
0 => not(&when_value),
_ => {
// `prep_null_mask_filter` is required to ensure the not operation treats nulls
// as false
not(&prep_null_mask_filter(&when_value))
}
}?;
let next_filter = create_filter(&next_selection, true);
remainder_batch =
Cow::Owned(filter_record_batch(&remainder_batch, &next_filter)?);
remainder_rows = filter_array(&remainder_rows, &next_filter)?;
base_values = filter_array(&base_values, &next_filter)?;
}
// If we reached this point, some rows were left unmatched.
// Check if those need to be evaluated using the 'else' expression.
if let Some(e) = &self.else_expr {
// keep `else_expr`'s data type and return type consistent
let expr = try_cast(Arc::clone(e), &batch.schema(), return_type.clone())?;
let else_value = expr.evaluate(&remainder_batch)?;
result_builder.add_branch_result(&remainder_rows, else_value)?;
}
result_builder.finish()
}
/// See [CaseExpr::case_when_no_expr].
fn case_when_no_expr(
&self,
batch: &RecordBatch,
return_type: &DataType,
) -> Result<ColumnarValue> {
let mut result_builder = ResultBuilder::new(return_type, batch.num_rows());
// `remainder_rows` contains the indices of the rows that need to be evaluated
let mut remainder_rows: ArrayRef =
Arc::new(UInt32Array::from_iter(0..batch.num_rows() as u32));
// `remainder_batch` contains the rows themselves that need to be evaluated
let mut remainder_batch = Cow::Borrowed(batch);
for i in 0..self.when_then_expr.len() {
// Evaluate the 'when' predicate for the remainder batch
// This results in a boolean array with the same length as the remaining number of rows
let when_predicate = &self.when_then_expr[i].0;
let when_value = when_predicate
.evaluate(&remainder_batch)?
.into_array(remainder_batch.num_rows())?;
let when_value = as_boolean_array(&when_value).map_err(|_| {
internal_datafusion_err!("WHEN expression did not return a BooleanArray")
})?;
// `true_count` ignores `true` values where the validity bit is not set, so there's
// no need to call `prep_null_mask_filter`.
let when_true_count = when_value.true_count();
// If the 'when' predicate did not match any rows, continue to the next branch immediately
if when_true_count == 0 {
continue;
}
// If the 'when' predicate matched all remaining rows, there is no need to filter
if when_true_count == remainder_batch.num_rows() {
let then_expression = &self.when_then_expr[i].1;
let then_value = then_expression.evaluate(&remainder_batch)?;
result_builder.add_branch_result(&remainder_rows, then_value)?;
return result_builder.finish();
}
// Filter the remainder batch based on the 'when' value
// This results in a batch containing only the rows that need to be evaluated
// for the current branch
// Still no need to call `prep_null_mask_filter` since `create_filter` will already do
// this unconditionally.
let then_filter = create_filter(when_value, true);
let then_batch = filter_record_batch(&remainder_batch, &then_filter)?;
let then_rows = filter_array(&remainder_rows, &then_filter)?;
let then_expression = &self.when_then_expr[i].1;
let then_value = then_expression.evaluate(&then_batch)?;
result_builder.add_branch_result(&then_rows, then_value)?;
// If this is the last 'when' branch and there is no 'else' expression, there's no
// point in calculating the remaining rows.
if self.else_expr.is_none() && i == self.when_then_expr.len() - 1 {
return result_builder.finish();
}
// Prepare the next when branch (or the else branch)
let next_selection = match when_value.null_count() {
0 => not(when_value),
_ => {
// `prep_null_mask_filter` is required to ensure the not operation treats nulls
// as false
not(&prep_null_mask_filter(when_value))
}
}?;
let next_filter = create_filter(&next_selection, true);
remainder_batch =
Cow::Owned(filter_record_batch(&remainder_batch, &next_filter)?);
remainder_rows = filter_array(&remainder_rows, &next_filter)?;
}
// If we reached this point, some rows were left unmatched.
// Check if those need to be evaluated using the 'else' expression.
if let Some(e) = &self.else_expr {
// keep `else_expr`'s data type and return type consistent
let expr = try_cast(Arc::clone(e), &batch.schema(), return_type.clone())?;
let else_value = expr.evaluate(&remainder_batch)?;
result_builder.add_branch_result(&remainder_rows, else_value)?;
}
result_builder.finish()
}
/// See [CaseExpr::expr_or_expr].
fn expr_or_expr(
&self,
batch: &RecordBatch,
when_value: &BooleanArray,
) -> Result<ColumnarValue> {
let when_value = match when_value.null_count() {
0 => Cow::Borrowed(when_value),
_ => {
// `prep_null_mask_filter` is required to ensure null is treated as false
Cow::Owned(prep_null_mask_filter(when_value))
}
};
let optimize_filter = batch.num_columns() > 1
|| (batch.num_columns() == 1 && multiple_arrays(batch.column(0).data_type()));
let when_filter = create_filter(&when_value, optimize_filter);
let then_batch = filter_record_batch(batch, &when_filter)?;
let then_value = self.when_then_expr[0].1.evaluate(&then_batch)?;
let else_selection = not(&when_value)?;
let else_filter = create_filter(&else_selection, optimize_filter);
let else_batch = filter_record_batch(batch, &else_filter)?;
// keep `else_expr`'s data type and return type consistent
let e = self.else_expr.as_ref().unwrap();
let return_type = self.data_type(&batch.schema())?;
let else_expr = try_cast(Arc::clone(e), &batch.schema(), return_type.clone())
.unwrap_or_else(|_| Arc::clone(e));
let else_value = else_expr.evaluate(&else_batch)?;
Ok(ColumnarValue::Array(merge(
&when_value,
then_value,
else_value,
)?))
}
}
impl CaseExpr {
/// This function evaluates the form of CASE that matches an expression to fixed values.
///
/// CASE expression
/// WHEN value THEN result
/// [WHEN ...]
/// [ELSE result]
/// END
fn case_when_with_expr(
&self,
batch: &RecordBatch,
projected: &ProjectedCaseBody,
) -> Result<ColumnarValue> {
let return_type = self.data_type(&batch.schema())?;
if projected.projection.len() < batch.num_columns() {
let projected_batch = batch.project(&projected.projection)?;
projected
.body
.case_when_with_expr(&projected_batch, &return_type)
} else {
self.body.case_when_with_expr(batch, &return_type)
}
}
/// This function evaluates the form of CASE where each WHEN expression is a boolean
/// expression.
///
/// CASE WHEN condition THEN result
/// [WHEN ...]
/// [ELSE result]
/// END
fn case_when_no_expr(
&self,
batch: &RecordBatch,
projected: &ProjectedCaseBody,
) -> Result<ColumnarValue> {
let return_type = self.data_type(&batch.schema())?;
if projected.projection.len() < batch.num_columns() {
let projected_batch = batch.project(&projected.projection)?;
projected
.body
.case_when_no_expr(&projected_batch, &return_type)
} else {
self.body.case_when_no_expr(batch, &return_type)
}
}
/// This function evaluates the specialized case of:
///
/// CASE WHEN condition THEN column
/// [ELSE NULL]
/// END
///
/// Note that this function is only safe to use for "then" expressions
/// that are infallible because the expression will be evaluated for all
/// rows in the input batch.
fn case_column_or_null(&self, batch: &RecordBatch) -> Result<ColumnarValue> {
let when_expr = &self.body.when_then_expr[0].0;
let then_expr = &self.body.when_then_expr[0].1;
match when_expr.evaluate(batch)? {
// WHEN true --> column
ColumnarValue::Scalar(ScalarValue::Boolean(Some(true))) => {
then_expr.evaluate(batch)
}
// WHEN [false | null] --> NULL
ColumnarValue::Scalar(_) => {
// return scalar NULL value
ScalarValue::try_from(self.data_type(&batch.schema())?)
.map(ColumnarValue::Scalar)
}
// WHEN column --> column
ColumnarValue::Array(bit_mask) => {
let bit_mask = bit_mask
.as_any()
.downcast_ref::<BooleanArray>()
.expect("predicate should evaluate to a boolean array");
// invert the bitmask
let bit_mask = match bit_mask.null_count() {
0 => not(bit_mask)?,
_ => not(&prep_null_mask_filter(bit_mask))?,
};
match then_expr.evaluate(batch)? {
ColumnarValue::Array(array) => {
Ok(ColumnarValue::Array(nullif(&array, &bit_mask)?))
}
ColumnarValue::Scalar(_) => {
internal_err!("expression did not evaluate to an array")
}
}
}
}
}
fn scalar_or_scalar(&self, batch: &RecordBatch) -> Result<ColumnarValue> {
let return_type = self.data_type(&batch.schema())?;
// evaluate when expression
let when_value = self.body.when_then_expr[0].0.evaluate(batch)?;
let when_value = when_value.into_array(batch.num_rows())?;
let when_value = as_boolean_array(&when_value).map_err(|_| {
internal_datafusion_err!("WHEN expression did not return a BooleanArray")
})?;
// Treat 'NULL' as false value
let when_value = match when_value.null_count() {
0 => Cow::Borrowed(when_value),
_ => Cow::Owned(prep_null_mask_filter(when_value)),
};
// evaluate then_value
let then_value = self.body.when_then_expr[0].1.evaluate(batch)?;
let then_value = Scalar::new(then_value.into_array(1)?);
let Some(e) = &self.body.else_expr else {
return internal_err!("expression did not evaluate to an array");
};
// keep `else_expr`'s data type and return type consistent
let expr = try_cast(Arc::clone(e), &batch.schema(), return_type)?;
let else_ = Scalar::new(expr.evaluate(batch)?.into_array(1)?);
Ok(ColumnarValue::Array(zip(&when_value, &then_value, &else_)?))
}
fn expr_or_expr(
&self,
batch: &RecordBatch,
projected: &ProjectedCaseBody,
) -> Result<ColumnarValue> {
// evaluate when condition on batch
let when_value = self.body.when_then_expr[0].0.evaluate(batch)?;
// `num_rows == 1` is intentional to avoid expanding scalars.
// If the `when_value` is effectively a scalar, the 'all true' and 'all false' checks
// below will avoid incorrectly using the scalar as a merge/zip mask.
let when_value = when_value.into_array(1)?;
let when_value = as_boolean_array(&when_value).map_err(|e| {
DataFusionError::Context(
"WHEN expression did not return a BooleanArray".to_string(),
Box::new(e),
)
})?;
let true_count = when_value.true_count();
if true_count == when_value.len() {
// All input rows are true, just call the 'then' expression
self.body.when_then_expr[0].1.evaluate(batch)
} else if true_count == 0 {
// All input rows are false/null, just call the 'else' expression
self.body.else_expr.as_ref().unwrap().evaluate(batch)
} else if projected.projection.len() < batch.num_columns() {
// The case expressions do not use all the columns of the input batch.
// Project first to reduce time spent filtering.
let projected_batch = batch.project(&projected.projection)?;
projected.body.expr_or_expr(&projected_batch, when_value)
} else {
// All columns are used in the case expressions, so there is no need to project.
self.body.expr_or_expr(batch, when_value)
}
}
}
impl PhysicalExpr for CaseExpr {
/// Return a reference to Any that can be used for down-casting
fn as_any(&self) -> &dyn Any {
self
}
fn data_type(&self, input_schema: &Schema) -> Result<DataType> {
self.body.data_type(input_schema)
}
fn nullable(&self, input_schema: &Schema) -> Result<bool> {
// this expression is nullable if any of the input expressions are nullable
let then_nullable = self
.body
.when_then_expr
.iter()
.map(|(_, t)| t.nullable(input_schema))
.collect::<Result<Vec<_>>>()?;
if then_nullable.contains(&true) {
Ok(true)
} else if let Some(e) = &self.body.else_expr {
e.nullable(input_schema)
} else {
// CASE produces NULL if there is no `else` expr
// (aka when none of the `when_then_exprs` match)
Ok(true)
}
}
fn evaluate(&self, batch: &RecordBatch) -> Result<ColumnarValue> {
match &self.eval_method {
EvalMethod::WithExpression(p) => {
// this use case evaluates "expr" and then compares the values with the "when"
// values
self.case_when_with_expr(batch, p)
}
EvalMethod::NoExpression(p) => {
// The "when" conditions all evaluate to boolean in this use case and can be
// arbitrary expressions
self.case_when_no_expr(batch, p)
}
EvalMethod::InfallibleExprOrNull => {
// Specialization for CASE WHEN expr THEN column [ELSE NULL] END
self.case_column_or_null(batch)
}
EvalMethod::ScalarOrScalar => self.scalar_or_scalar(batch),
EvalMethod::ExpressionOrExpression(p) => self.expr_or_expr(batch, p),
}
}
fn children(&self) -> Vec<&Arc<dyn PhysicalExpr>> {
let mut children = vec![];
if let Some(expr) = &self.body.expr {
children.push(expr)
}
self.body.when_then_expr.iter().for_each(|(cond, value)| {
children.push(cond);
children.push(value);
});
if let Some(else_expr) = &self.body.else_expr {
children.push(else_expr)
}
children
}
// For physical CaseExpr, we do not allow modifying children size
fn with_new_children(
self: Arc<Self>,
children: Vec<Arc<dyn PhysicalExpr>>,
) -> Result<Arc<dyn PhysicalExpr>> {
if children.len() != self.children().len() {
internal_err!("CaseExpr: Wrong number of children")
} else {
let (expr, when_then_expr, else_expr) =
match (self.expr().is_some(), self.body.else_expr.is_some()) {
(true, true) => (
Some(&children[0]),
&children[1..children.len() - 1],
Some(&children[children.len() - 1]),
),
(true, false) => {
(Some(&children[0]), &children[1..children.len()], None)
}
(false, true) => (
None,
&children[0..children.len() - 1],
Some(&children[children.len() - 1]),
),
(false, false) => (None, &children[0..children.len()], None),
};
Ok(Arc::new(CaseExpr::try_new(
expr.cloned(),
when_then_expr.iter().cloned().tuples().collect(),
else_expr.cloned(),
)?))
}
}
fn fmt_sql(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
write!(f, "CASE ")?;
if let Some(e) = &self.body.expr {
e.fmt_sql(f)?;
write!(f, " ")?;
}
for (w, t) in &self.body.when_then_expr {
write!(f, "WHEN ")?;
w.fmt_sql(f)?;
write!(f, " THEN ")?;
t.fmt_sql(f)?;
write!(f, " ")?;
}
if let Some(e) = &self.body.else_expr {
write!(f, "ELSE ")?;
e.fmt_sql(f)?;
write!(f, " ")?;
}
write!(f, "END")
}
}
/// Create a CASE expression
pub fn case(
expr: Option<Arc<dyn PhysicalExpr>>,
when_thens: Vec<WhenThen>,
else_expr: Option<Arc<dyn PhysicalExpr>>,
) -> Result<Arc<dyn PhysicalExpr>> {
Ok(Arc::new(CaseExpr::try_new(expr, when_thens, else_expr)?))
}
#[cfg(test)]
mod tests {
use super::*;
use crate::expressions::{binary, cast, col, lit, BinaryExpr};
use arrow::buffer::Buffer;
use arrow::datatypes::DataType::Float64;
use arrow::datatypes::Field;
use datafusion_common::cast::{as_float64_array, as_int32_array};
use datafusion_common::plan_err;
use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode};
use datafusion_expr::type_coercion::binary::comparison_coercion;
use datafusion_expr::Operator;
use datafusion_physical_expr_common::physical_expr::fmt_sql;
#[test]
fn case_with_expr() -> Result<()> {
let batch = case_test_batch()?;
let schema = batch.schema();
// CASE a WHEN 'foo' THEN 123 WHEN 'bar' THEN 456 END
let when1 = lit("foo");
let then1 = lit(123i32);
let when2 = lit("bar");
let then2 = lit(456i32);
let expr = generate_case_when_with_type_coercion(
Some(col("a", &schema)?),
vec![(when1, then1), (when2, then2)],
None,
schema.as_ref(),
)?;
let result = expr
.evaluate(&batch)?
.into_array(batch.num_rows())
.expect("Failed to convert to array");
let result = as_int32_array(&result)?;
let expected = &Int32Array::from(vec![Some(123), None, None, Some(456)]);
assert_eq!(expected, result);
Ok(())
}
#[test]
fn case_with_expr_else() -> Result<()> {
let batch = case_test_batch()?;
let schema = batch.schema();
// CASE a WHEN 'foo' THEN 123 WHEN 'bar' THEN 456 ELSE 999 END
let when1 = lit("foo");
let then1 = lit(123i32);
let when2 = lit("bar");
let then2 = lit(456i32);
let else_value = lit(999i32);
let expr = generate_case_when_with_type_coercion(
Some(col("a", &schema)?),
vec![(when1, then1), (when2, then2)],
Some(else_value),
schema.as_ref(),
)?;
let result = expr
.evaluate(&batch)?
.into_array(batch.num_rows())
.expect("Failed to convert to array");
let result = as_int32_array(&result)?;
let expected =
&Int32Array::from(vec![Some(123), Some(999), Some(999), Some(456)]);
assert_eq!(expected, result);
Ok(())
}
#[test]
fn case_with_expr_divide_by_zero() -> Result<()> {
let batch = case_test_batch1()?;
let schema = batch.schema();
// CASE a when 0 THEN float64(null) ELSE 25.0 / cast(a, float64) END
let when1 = lit(0i32);
let then1 = lit(ScalarValue::Float64(None));
let else_value = binary(
lit(25.0f64),
Operator::Divide,
cast(col("a", &schema)?, &batch.schema(), Float64)?,
&batch.schema(),
)?;
let expr = generate_case_when_with_type_coercion(
Some(col("a", &schema)?),
vec![(when1, then1)],
Some(else_value),
schema.as_ref(),
)?;
let result = expr
.evaluate(&batch)?
.into_array(batch.num_rows())
.expect("Failed to convert to array");
let result =
as_float64_array(&result).expect("failed to downcast to Float64Array");
let expected = &Float64Array::from(vec![Some(25.0), None, None, Some(5.0)]);
assert_eq!(expected, result);
Ok(())
}
#[test]
fn case_without_expr() -> Result<()> {
let batch = case_test_batch()?;
let schema = batch.schema();
// CASE WHEN a = 'foo' THEN 123 WHEN a = 'bar' THEN 456 END
let when1 = binary(
col("a", &schema)?,
Operator::Eq,
lit("foo"),
&batch.schema(),
)?;
let then1 = lit(123i32);
let when2 = binary(
col("a", &schema)?,
Operator::Eq,
lit("bar"),
&batch.schema(),
)?;
let then2 = lit(456i32);
let expr = generate_case_when_with_type_coercion(
None,
vec![(when1, then1), (when2, then2)],
None,
schema.as_ref(),
)?;
let result = expr
.evaluate(&batch)?
.into_array(batch.num_rows())
.expect("Failed to convert to array");
let result = as_int32_array(&result)?;
let expected = &Int32Array::from(vec![Some(123), None, None, Some(456)]);
assert_eq!(expected, result);
Ok(())
}
#[test]
fn case_with_expr_when_null() -> Result<()> {
let batch = case_test_batch()?;
let schema = batch.schema();
// CASE a WHEN NULL THEN 0 WHEN a THEN 123 ELSE 999 END
let when1 = lit(ScalarValue::Utf8(None));
let then1 = lit(0i32);
let when2 = col("a", &schema)?;
let then2 = lit(123i32);
let else_value = lit(999i32);
let expr = generate_case_when_with_type_coercion(
Some(col("a", &schema)?),
vec![(when1, then1), (when2, then2)],
Some(else_value),
schema.as_ref(),
)?;
let result = expr
.evaluate(&batch)?
.into_array(batch.num_rows())
.expect("Failed to convert to array");
let result = as_int32_array(&result)?;
let expected =
&Int32Array::from(vec![Some(123), Some(123), Some(999), Some(123)]);
assert_eq!(expected, result);
Ok(())
}
#[test]
fn case_without_expr_divide_by_zero() -> Result<()> {
let batch = case_test_batch1()?;
let schema = batch.schema();
// CASE WHEN a > 0 THEN 25.0 / cast(a, float64) ELSE float64(null) END
let when1 = binary(col("a", &schema)?, Operator::Gt, lit(0i32), &batch.schema())?;
let then1 = binary(
lit(25.0f64),
Operator::Divide,
cast(col("a", &schema)?, &batch.schema(), Float64)?,
&batch.schema(),
)?;
let x = lit(ScalarValue::Float64(None));
let expr = generate_case_when_with_type_coercion(
None,
vec![(when1, then1)],
Some(x),
schema.as_ref(),
)?;
let result = expr
.evaluate(&batch)?
.into_array(batch.num_rows())
.expect("Failed to convert to array");
let result =
as_float64_array(&result).expect("failed to downcast to Float64Array");
let expected = &Float64Array::from(vec![Some(25.0), None, None, Some(5.0)]);
assert_eq!(expected, result);
Ok(())
}
fn case_test_batch1() -> Result<RecordBatch> {
let schema = Schema::new(vec![
Field::new("a", DataType::Int32, true),
Field::new("b", DataType::Int32, true),
Field::new("c", DataType::Int32, true),
]);
let a = Int32Array::from(vec![Some(1), Some(0), None, Some(5)]);
let b = Int32Array::from(vec![Some(3), None, Some(14), Some(7)]);
let c = Int32Array::from(vec![Some(0), Some(-3), Some(777), None]);
let batch = RecordBatch::try_new(
Arc::new(schema),
vec![Arc::new(a), Arc::new(b), Arc::new(c)],
)?;
Ok(batch)
}
#[test]
fn case_without_expr_else() -> Result<()> {
let batch = case_test_batch()?;
let schema = batch.schema();
// CASE WHEN a = 'foo' THEN 123 WHEN a = 'bar' THEN 456 ELSE 999 END
let when1 = binary(
col("a", &schema)?,
Operator::Eq,
lit("foo"),
&batch.schema(),
)?;
let then1 = lit(123i32);
let when2 = binary(
col("a", &schema)?,
Operator::Eq,
lit("bar"),
&batch.schema(),
)?;
let then2 = lit(456i32);
let else_value = lit(999i32);
let expr = generate_case_when_with_type_coercion(
None,
vec![(when1, then1), (when2, then2)],
Some(else_value),
schema.as_ref(),
)?;
let result = expr
.evaluate(&batch)?
.into_array(batch.num_rows())
.expect("Failed to convert to array");
let result = as_int32_array(&result)?;
let expected =
&Int32Array::from(vec![Some(123), Some(999), Some(999), Some(456)]);
assert_eq!(expected, result);
Ok(())
}
#[test]
fn case_with_type_cast() -> Result<()> {
let batch = case_test_batch()?;
let schema = batch.schema();
// CASE WHEN a = 'foo' THEN 123.3 ELSE 999 END
let when = binary(
col("a", &schema)?,
Operator::Eq,
lit("foo"),
&batch.schema(),
)?;
let then = lit(123.3f64);
let else_value = lit(999i32);
let expr = generate_case_when_with_type_coercion(
None,
vec![(when, then)],
Some(else_value),
schema.as_ref(),
)?;
let result = expr
.evaluate(&batch)?
.into_array(batch.num_rows())
.expect("Failed to convert to array");
let result =
as_float64_array(&result).expect("failed to downcast to Float64Array");
let expected =
&Float64Array::from(vec![Some(123.3), Some(999.0), Some(999.0), Some(999.0)]);
assert_eq!(expected, result);
Ok(())
}
#[test]
fn case_with_matches_and_nulls() -> Result<()> {
let batch = case_test_batch_nulls()?;
let schema = batch.schema();
// SELECT CASE WHEN load4 = 1.77 THEN load4 END
let when = binary(
col("load4", &schema)?,
Operator::Eq,
lit(1.77f64),
&batch.schema(),
)?;
let then = col("load4", &schema)?;
let expr = generate_case_when_with_type_coercion(
None,
vec![(when, then)],
None,
schema.as_ref(),
)?;
let result = expr
.evaluate(&batch)?
.into_array(batch.num_rows())
.expect("Failed to convert to array");
let result =
as_float64_array(&result).expect("failed to downcast to Float64Array");
let expected =
&Float64Array::from(vec![Some(1.77), None, None, None, None, Some(1.77)]);
assert_eq!(expected, result);
Ok(())
}
#[test]
fn case_with_scalar_predicate() -> Result<()> {
let batch = case_test_batch_nulls()?;
let schema = batch.schema();
// SELECT CASE WHEN TRUE THEN load4 END
let when = lit(true);
let then = col("load4", &schema)?;
let expr = generate_case_when_with_type_coercion(
None,
vec![(when, then)],
None,
schema.as_ref(),
)?;
// many rows
let result = expr
.evaluate(&batch)?
.into_array(batch.num_rows())
.expect("Failed to convert to array");
let result =
as_float64_array(&result).expect("failed to downcast to Float64Array");
let expected = &Float64Array::from(vec![
Some(1.77),
None,
None,
Some(1.78),
None,
Some(1.77),
]);
assert_eq!(expected, result);
// one row
let expected = Float64Array::from(vec![Some(1.1)]);
let batch =
RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(expected.clone())])?;
let result = expr
.evaluate(&batch)?
.into_array(batch.num_rows())
.expect("Failed to convert to array");
let result =
as_float64_array(&result).expect("failed to downcast to Float64Array");
assert_eq!(&expected, result);
Ok(())
}
#[test]
fn case_expr_matches_and_nulls() -> Result<()> {
let batch = case_test_batch_nulls()?;
let schema = batch.schema();
// SELECT CASE load4 WHEN 1.77 THEN load4 END
let expr = col("load4", &schema)?;
let when = lit(1.77f64);
let then = col("load4", &schema)?;
let expr = generate_case_when_with_type_coercion(
Some(expr),
vec![(when, then)],
None,
schema.as_ref(),
)?;
let result = expr
.evaluate(&batch)?
.into_array(batch.num_rows())
.expect("Failed to convert to array");
let result =
as_float64_array(&result).expect("failed to downcast to Float64Array");
let expected =
&Float64Array::from(vec![Some(1.77), None, None, None, None, Some(1.77)]);
assert_eq!(expected, result);
Ok(())
}
#[test]
fn test_when_null_and_some_cond_else_null() -> Result<()> {
let batch = case_test_batch()?;
let schema = batch.schema();
let when = binary(
Arc::new(Literal::new(ScalarValue::Boolean(None))),
Operator::And,
binary(col("a", &schema)?, Operator::Eq, lit("foo"), &schema)?,
&schema,
)?;
let then = col("a", &schema)?;
// SELECT CASE WHEN (NULL AND a = 'foo') THEN a ELSE NULL END
let expr = Arc::new(CaseExpr::try_new(None, vec![(when, then)], None)?);
let result = expr
.evaluate(&batch)?
.into_array(batch.num_rows())
.expect("Failed to convert to array");
let result = as_string_array(&result);
// all result values should be null
assert_eq!(result.logical_null_count(), batch.num_rows());
Ok(())
}
fn case_test_batch() -> Result<RecordBatch> {
let schema = Schema::new(vec![Field::new("a", DataType::Utf8, true)]);
let a = StringArray::from(vec![Some("foo"), Some("baz"), None, Some("bar")]);
let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)])?;
Ok(batch)
}
// Construct an array that has several NULL values whose
// underlying buffer actually matches the where expr predicate
fn case_test_batch_nulls() -> Result<RecordBatch> {
let load4: Float64Array = vec![
Some(1.77), // 1.77
Some(1.77), // null <-- same value, but will be set to null
Some(1.77), // null <-- same value, but will be set to null
Some(1.78), // 1.78
None, // null
Some(1.77), // 1.77
]
.into_iter()
.collect();
let null_buffer = Buffer::from([0b00101001u8]);
let load4 = load4
.into_data()
.into_builder()
.null_bit_buffer(Some(null_buffer))
.build()
.unwrap();
let load4: Float64Array = load4.into();
let batch =
RecordBatch::try_from_iter(vec![("load4", Arc::new(load4) as ArrayRef)])?;
Ok(batch)
}
#[test]
fn case_test_incompatible() -> Result<()> {
// 1 then is int64
// 2 then is boolean
let batch = case_test_batch()?;
let schema = batch.schema();
// CASE WHEN a = 'foo' THEN 123 WHEN a = 'bar' THEN true END
let when1 = binary(
col("a", &schema)?,
Operator::Eq,
lit("foo"),
&batch.schema(),
)?;
let then1 = lit(123i32);
let when2 = binary(
col("a", &schema)?,
Operator::Eq,
lit("bar"),
&batch.schema(),
)?;
let then2 = lit(true);
let expr = generate_case_when_with_type_coercion(
None,
vec![(when1, then1), (when2, then2)],
None,
schema.as_ref(),
);
assert!(expr.is_err());
// then 1 is int32
// then 2 is int64
// else is float
// CASE WHEN a = 'foo' THEN 123 WHEN a = 'bar' THEN 456 ELSE 1.23 END
let when1 = binary(
col("a", &schema)?,
Operator::Eq,
lit("foo"),
&batch.schema(),
)?;
let then1 = lit(123i32);
let when2 = binary(
col("a", &schema)?,
Operator::Eq,
lit("bar"),
&batch.schema(),
)?;
let then2 = lit(456i64);
let else_expr = lit(1.23f64);
let expr = generate_case_when_with_type_coercion(
None,
vec![(when1, then1), (when2, then2)],
Some(else_expr),
schema.as_ref(),
);
assert!(expr.is_ok());
let result_type = expr.unwrap().data_type(schema.as_ref())?;
assert_eq!(Float64, result_type);
Ok(())
}
#[test]
fn case_eq() -> Result<()> {
let schema = Schema::new(vec![Field::new("a", DataType::Int32, true)]);
let when1 = lit("foo");
let then1 = lit(123i32);
let when2 = lit("bar");
let then2 = lit(456i32);
let else_value = lit(999i32);
let expr1 = generate_case_when_with_type_coercion(
Some(col("a", &schema)?),
vec![
(Arc::clone(&when1), Arc::clone(&then1)),
(Arc::clone(&when2), Arc::clone(&then2)),
],
Some(Arc::clone(&else_value)),
&schema,
)?;
let expr2 = generate_case_when_with_type_coercion(
Some(col("a", &schema)?),
vec![
(Arc::clone(&when1), Arc::clone(&then1)),
(Arc::clone(&when2), Arc::clone(&then2)),
],
Some(Arc::clone(&else_value)),
&schema,
)?;
let expr3 = generate_case_when_with_type_coercion(
Some(col("a", &schema)?),
vec![(Arc::clone(&when1), Arc::clone(&then1)), (when2, then2)],
None,
&schema,
)?;
let expr4 = generate_case_when_with_type_coercion(
Some(col("a", &schema)?),
vec![(when1, then1)],
Some(else_value),
&schema,
)?;
assert!(expr1.eq(&expr2));
assert!(expr2.eq(&expr1));
assert!(expr2.ne(&expr3));
assert!(expr3.ne(&expr2));
assert!(expr1.ne(&expr4));
assert!(expr4.ne(&expr1));
Ok(())
}
#[test]
fn case_transform() -> Result<()> {
let schema = Schema::new(vec![Field::new("a", DataType::Int32, true)]);
let when1 = lit("foo");
let then1 = lit(123i32);
let when2 = lit("bar");
let then2 = lit(456i32);
let else_value = lit(999i32);
let expr = generate_case_when_with_type_coercion(
Some(col("a", &schema)?),
vec![
(Arc::clone(&when1), Arc::clone(&then1)),
(Arc::clone(&when2), Arc::clone(&then2)),
],
Some(Arc::clone(&else_value)),
&schema,
)?;
let expr2 = Arc::clone(&expr)
.transform(|e| {
let transformed = match e.as_any().downcast_ref::<Literal>() {
Some(lit_value) => match lit_value.value() {
ScalarValue::Utf8(Some(str_value)) => {
Some(lit(str_value.to_uppercase()))
}
_ => None,
},
_ => None,
};
Ok(if let Some(transformed) = transformed {
Transformed::yes(transformed)
} else {
Transformed::no(e)
})
})
.data()
.unwrap();
let expr3 = Arc::clone(&expr)
.transform_down(|e| {
let transformed = match e.as_any().downcast_ref::<Literal>() {
Some(lit_value) => match lit_value.value() {
ScalarValue::Utf8(Some(str_value)) => {
Some(lit(str_value.to_uppercase()))
}
_ => None,
},
_ => None,
};
Ok(if let Some(transformed) = transformed {
Transformed::yes(transformed)
} else {
Transformed::no(e)
})
})
.data()
.unwrap();
assert!(expr.ne(&expr2));
assert!(expr2.eq(&expr3));
Ok(())
}
#[test]
fn test_column_or_null_specialization() -> Result<()> {
// create input data
let mut c1 = Int32Builder::new();
let mut c2 = StringBuilder::new();
for i in 0..1000 {
c1.append_value(i);
if i % 7 == 0 {
c2.append_null();
} else {
c2.append_value(format!("string {i}"));
}
}
let c1 = Arc::new(c1.finish());
let c2 = Arc::new(c2.finish());
let schema = Schema::new(vec![
Field::new("c1", DataType::Int32, true),
Field::new("c2", DataType::Utf8, true),
]);
let batch = RecordBatch::try_new(Arc::new(schema), vec![c1, c2]).unwrap();
// CaseWhenExprOrNull should produce same results as CaseExpr
let predicate = Arc::new(BinaryExpr::new(
make_col("c1", 0),
Operator::LtEq,
make_lit_i32(250),
));
let expr = CaseExpr::try_new(None, vec![(predicate, make_col("c2", 1))], None)?;
assert!(matches!(expr.eval_method, EvalMethod::InfallibleExprOrNull));
match expr.evaluate(&batch)? {
ColumnarValue::Array(array) => {
assert_eq!(1000, array.len());
assert_eq!(785, array.null_count());
}
_ => unreachable!(),
}
Ok(())
}
#[test]
fn test_expr_or_expr_specialization() -> Result<()> {
let batch = case_test_batch1()?;
let schema = batch.schema();
let when = binary(
col("a", &schema)?,
Operator::LtEq,
lit(2i32),
&batch.schema(),
)?;
let then = col("b", &schema)?;
let else_expr = col("c", &schema)?;
let expr = CaseExpr::try_new(None, vec![(when, then)], Some(else_expr))?;
assert!(matches!(
expr.eval_method,
EvalMethod::ExpressionOrExpression(_)
));
let result = expr
.evaluate(&batch)?
.into_array(batch.num_rows())
.expect("Failed to convert to array");
let result = as_int32_array(&result).expect("failed to downcast to Int32Array");
let expected = &Int32Array::from(vec![Some(3), None, Some(777), None]);
assert_eq!(expected, result);
Ok(())
}
fn make_col(name: &str, index: usize) -> Arc<dyn PhysicalExpr> {
Arc::new(Column::new(name, index))
}
fn make_lit_i32(n: i32) -> Arc<dyn PhysicalExpr> {
Arc::new(Literal::new(ScalarValue::Int32(Some(n))))
}
fn generate_case_when_with_type_coercion(
expr: Option<Arc<dyn PhysicalExpr>>,
when_thens: Vec<WhenThen>,
else_expr: Option<Arc<dyn PhysicalExpr>>,
input_schema: &Schema,
) -> Result<Arc<dyn PhysicalExpr>> {
let coerce_type =
get_case_common_type(&when_thens, else_expr.clone(), input_schema);
let (when_thens, else_expr) = match coerce_type {
None => plan_err!(
"Can't get a common type for then {when_thens:?} and else {else_expr:?} expression"
),
Some(data_type) => {
// cast then expr
let left = when_thens
.into_iter()
.map(|(when, then)| {
let then = try_cast(then, input_schema, data_type.clone())?;
Ok((when, then))
})
.collect::<Result<Vec<_>>>()?;
let right = match else_expr {
None => None,
Some(expr) => Some(try_cast(expr, input_schema, data_type.clone())?),
};
Ok((left, right))
}
}?;
case(expr, when_thens, else_expr)
}
fn get_case_common_type(
when_thens: &[WhenThen],
else_expr: Option<Arc<dyn PhysicalExpr>>,
input_schema: &Schema,
) -> Option<DataType> {
let thens_type = when_thens
.iter()
.map(|when_then| {
let data_type = &when_then.1.data_type(input_schema).unwrap();
data_type.clone()
})
.collect::<Vec<_>>();
let else_type = match else_expr {
None => {
// case when then exprs must have one then value
thens_type[0].clone()
}
Some(else_phy_expr) => else_phy_expr.data_type(input_schema).unwrap(),
};
thens_type
.iter()
.try_fold(else_type, |left_type, right_type| {
// TODO: now just use the `equal` coercion rule for case when. If find the issue, and
// refactor again.
comparison_coercion(&left_type, right_type)
})
}
#[test]
fn test_fmt_sql() -> Result<()> {
let schema = Schema::new(vec![Field::new("a", DataType::Utf8, true)]);
// CASE WHEN a = 'foo' THEN 123.3 ELSE 999 END
let when = binary(col("a", &schema)?, Operator::Eq, lit("foo"), &schema)?;
let then = lit(123.3f64);
let else_value = lit(999i32);
let expr = generate_case_when_with_type_coercion(
None,
vec![(when, then)],
Some(else_value),
&schema,
)?;
let display_string = expr.to_string();
assert_eq!(
display_string,
"CASE WHEN a@0 = foo THEN 123.3 ELSE TRY_CAST(999 AS Float64) END"
);
let sql_string = fmt_sql(expr.as_ref()).to_string();
assert_eq!(
sql_string,
"CASE WHEN a = foo THEN 123.3 ELSE TRY_CAST(999 AS Float64) END"
);
Ok(())
}
#[test]
fn test_merge_n() {
let a1 = StringArray::from(vec![Some("A")]).to_data();
let a2 = StringArray::from(vec![Some("B")]).to_data();
let a3 = StringArray::from(vec![Some("C"), Some("D")]).to_data();
let indices = vec![
PartialResultIndex::none(),
PartialResultIndex::try_new(1).unwrap(),
PartialResultIndex::try_new(0).unwrap(),
PartialResultIndex::none(),
PartialResultIndex::try_new(2).unwrap(),
PartialResultIndex::try_new(2).unwrap(),
];
let merged = merge_n(&[a1, a2, a3], &indices).unwrap();
let merged = merged.as_string::<i32>();
assert_eq!(merged.len(), indices.len());
assert!(!merged.is_valid(0));
assert!(merged.is_valid(1));
assert_eq!(merged.value(1), "B");
assert!(merged.is_valid(2));
assert_eq!(merged.value(2), "A");
assert!(!merged.is_valid(3));
assert!(merged.is_valid(4));
assert_eq!(merged.value(4), "C");
assert!(merged.is_valid(5));
assert_eq!(merged.value(5), "D");
}
#[test]
fn test_merge() {
let a1 = Arc::new(StringArray::from(vec![Some("A"), Some("C")]));
let a2 = Arc::new(StringArray::from(vec![Some("B")]));
let mask = BooleanArray::from(vec![true, false, true]);
let merged =
merge(&mask, ColumnarValue::Array(a1), ColumnarValue::Array(a2)).unwrap();
let merged = merged.as_string::<i32>();
assert_eq!(merged.len(), mask.len());
assert!(merged.is_valid(0));
assert_eq!(merged.value(0), "A");
assert!(merged.is_valid(1));
assert_eq!(merged.value(1), "B");
assert!(merged.is_valid(2));
assert_eq!(merged.value(2), "C");
}
}