Skip to content

Commit 3e45a47

Browse files
committed
simplify
Signed-off-by: Robert Kruszewski <github@robertk.io>
1 parent a7c162e commit 3e45a47

5 files changed

Lines changed: 69 additions & 79 deletions

File tree

‎vortex-array/src/expr/exprs.rs‎

Lines changed: 40 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,8 @@
66
use std::sync::Arc;
77

88
use vortex_error::VortexExpect;
9+
use vortex_error::VortexResult;
10+
use vortex_error::vortex_ensure;
911
use vortex_error::vortex_panic;
1012
use vortex_utils::iter::ReduceBalancedIterExt;
1113

@@ -1148,31 +1150,63 @@ pub fn in_list(needle: Expression, list: Expression) -> Expression {
11481150
)
11491151
}
11501152

1151-
/// Creates `needle IN (elements)` for a list of constant elements of `element_dtype`.
1153+
/// Creates `needle IN (elements)` with SQL null semantics for a list of constant elements of
1154+
/// `element_dtype`, ignoring nullability. An element can also be an untyped null.
1155+
///
1156+
/// The nullability of `element_dtype` is not used: the list elements are nullable only when a
1157+
/// null element is present.
11521158
///
11531159
/// Non-null integers that hold every value between their smallest and largest become a range
11541160
/// check, which is cheaper to evaluate and prunes with tighter bounds than a set lookup, as for
11551161
/// the keys of a small dimension table that a join pushes down.
11561162
///
1163+
/// # Errors
1164+
///
1165+
/// Returns an error if an element is not a null and not of `element_dtype`.
1166+
///
11571167
/// ```rust
11581168
/// # use vortex_array::dtype::{DType, Nullability, PType};
11591169
/// # use vortex_array::expr::{in_list_of, root};
11601170
/// let dtype = DType::Primitive(PType::I32, Nullability::NonNullable);
1161-
/// let expr = in_list_of(root(), dtype, vec![3i32.into(), 1i32.into(), 2i32.into()]);
1171+
/// let expr = in_list_of(root(), dtype, vec![3i32.into(), 1i32.into(), 2i32.into()])?;
11621172
/// assert_eq!(expr.to_string(), "(1i32 <= $ <= 3i32)");
1173+
/// # Ok::<(), vortex_error::VortexError>(())
11631174
/// ```
11641175
pub fn in_list_of(
11651176
needle: Expression,
11661177
element_dtype: DType,
1167-
mut elements: Vec<Scalar>,
1168-
) -> Expression {
1178+
elements: Vec<Scalar>,
1179+
) -> VortexResult<Expression> {
1180+
let element_dtype = if elements.iter().any(Scalar::is_null) {
1181+
element_dtype.as_nullable()
1182+
} else {
1183+
element_dtype.as_nonnullable()
1184+
};
1185+
let mut elements = elements
1186+
.iter()
1187+
.map(|element| {
1188+
vortex_ensure!(
1189+
element.dtype() == &DType::Null
1190+
|| element.dtype().eq_ignore_nullability(&element_dtype),
1191+
"IN list element {element} is not of type {element_dtype}"
1192+
);
1193+
element.cast(&element_dtype)
1194+
})
1195+
.collect::<VortexResult<Vec<_>>>()?;
1196+
11691197
if let Some((low, high)) = contiguous_integer_bounds(&elements) {
11701198
let options = BetweenOptions {
11711199
lower_strict: StrictComparison::NonStrict,
11721200
upper_strict: StrictComparison::NonStrict,
11731201
};
1174-
return between(needle, lit(low.clone()), lit(high.clone()), options);
1202+
return Ok(between(
1203+
needle,
1204+
lit(low.clone()),
1205+
lit(high.clone()),
1206+
options,
1207+
));
11751208
}
1209+
11761210
// Ordered elements make equal lists from the same set, so that equal filters deduplicate.
11771211
if elements
11781212
.windows(2)
@@ -1181,7 +1215,7 @@ pub fn in_list_of(
11811215
elements.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
11821216
}
11831217
let list = Scalar::list(Arc::new(element_dtype), elements, Nullability::NonNullable);
1184-
in_list(needle, lit(list))
1218+
Ok(in_list(needle, lit(list)))
11851219
}
11861220

11871221
/// The smallest and largest of non-null integer `scalars` when they hold every integer between

‎vortex-datafusion/src/convert/exprs/mod.rs‎

Lines changed: 20 additions & 61 deletions
Original file line numberDiff line numberDiff line change
@@ -26,9 +26,7 @@ use vortex::VortexSessionDefault;
2626
use vortex::array::VortexSessionExecute;
2727
use vortex::dtype::DType;
2828
use vortex::dtype::Nullability;
29-
use vortex::error::VortexResult;
3029
use vortex::expr::Expression;
31-
use vortex::expr::and_collect;
3230
use vortex::expr::byte_length;
3331
use vortex::expr::cast;
3432
use vortex::expr::get_item;
@@ -39,10 +37,8 @@ use vortex::expr::list_length;
3937
use vortex::expr::lit;
4038
use vortex::expr::nested_case_when;
4139
use vortex::expr::not;
42-
use vortex::expr::or_collect;
4340
use vortex::expr::pack;
4441
use vortex::expr::root;
45-
use vortex::scalar::Scalar;
4642
use vortex::scalar_fn::ScalarFnVTableExt;
4743
use vortex::scalar_fn::fns::binary::Binary;
4844
use vortex::scalar_fn::fns::like::Like;
@@ -440,64 +436,27 @@ impl DefaultExpressionConverter {
440436
return Err(Unconverted::Unsupported);
441437
}
442438
let value_dtype = converted_dtype(&value, input_dtype)?;
443-
let operator = if in_list.negated() {
444-
Operator::NotEq
445-
} else {
446-
Operator::Eq
447-
};
448-
// A positive list of literals becomes one set lookup. An OR of equalities would
449-
// scan the needle once per element, which dynamic join filters make up to 150 times.
450-
// Floats keep the comparisons, whose -0.0 and NaN semantics differ from set membership.
451-
let is_float = matches!(&value_dtype, DType::Primitive(ptype, _) if ptype.is_float());
452-
if !in_list.negated() && !is_float {
453-
let elements = in_list
454-
.list()
455-
.iter()
456-
.map(|element| self.convert_expr(element, schema, input_dtype))
457-
.collect::<Conversion<Vec<_>>>()?;
458-
let scalars = elements
459-
.iter()
460-
.map(|element| element.as_opt::<Literal>().cloned())
461-
.collect::<Option<Vec<_>>>();
462-
// Null elements make a nullable list, whose SQL semantics `in_list` preserves.
463-
let element_dtype = if scalars.iter().flatten().any(Scalar::is_null) {
464-
value_dtype.as_nullable()
465-
} else {
466-
value_dtype.as_nonnullable()
467-
};
468-
if let Some(scalars) = scalars
469-
&& !scalars.is_empty()
470-
&& scalars
471-
.iter()
472-
.all(|scalar| scalar.dtype().eq_ignore_nullability(&element_dtype))
473-
{
474-
let scalars = scalars
475-
.iter()
476-
.map(|scalar| scalar.cast(&element_dtype))
477-
.collect::<VortexResult<Vec<_>>>()
478-
.map_err(|_| Unconverted::Unsupported)?;
479-
return Ok(in_list_of(value, element_dtype, scalars));
480-
}
481-
}
482-
let mut comparisons = Vec::with_capacity(in_list.len());
483-
for element in in_list.list() {
484-
let element = self.convert_expr(element, schema, input_dtype)?;
485-
let element_dtype = converted_dtype(&element, input_dtype)?;
486-
if element_dtype == DType::Null {
487-
comparisons.push(lit(None::<bool>));
488-
} else if value_dtype.eq_ignore_nullability(&element_dtype) {
489-
comparisons.push(Binary.new_expr(operator, [value.clone(), element]));
490-
} else {
491-
return Err(Unconverted::Unsupported);
492-
}
493-
}
494-
// Kleene AND/OR preserve SQL IN/NOT IN nulls; list_contains does not.
495-
let membership = if in_list.negated() {
496-
and_collect(comparisons)
439+
440+
let elements = in_list
441+
.list()
442+
.iter()
443+
.map(|element| {
444+
let element = self.convert_expr(element, schema, input_dtype)?;
445+
element
446+
.as_opt::<Literal>()
447+
.cloned()
448+
.ok_or(Unconverted::Unsupported)
449+
})
450+
.collect::<Conversion<Vec<_>>>()?;
451+
452+
// An element of another type than the value declines the pushdown: no implicit casts.
453+
let membership =
454+
in_list_of(value, value_dtype, elements).map_err(|_| Unconverted::Unsupported)?;
455+
Ok(if in_list.negated() {
456+
not(membership)
497457
} else {
498-
or_collect(comparisons)
499-
};
500-
membership.ok_or(Unconverted::Unsupported)
458+
membership
459+
})
501460
}
502461

503462
fn convert_scalar_function(

‎vortex-datafusion/src/convert/exprs/tests.rs‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -329,13 +329,13 @@ fn test_native_in_list_large(#[values(false, true)] negated: bool) -> anyhow::Re
329329
)
330330
}
331331

332-
#[test]
333-
fn test_in_list_converts_to_set_lookup() -> DFResult<()> {
332+
#[rstest]
333+
fn test_in_list_converts_to_set_lookup(#[values(false, true)] negated: bool) -> DFResult<()> {
334334
let schema = Schema::new(vec![Field::new("a", DataType::Int32, true)]);
335335
let expr = Arc::new(df_expr::InListExpr::try_new(
336336
Arc::new(df_expr::Column::new("a", 0)),
337337
int_literals([Some(1), Some(3), Some(7)]),
338-
false,
338+
negated,
339339
&schema,
340340
)?);
341341
let element_dtype = Arc::new(DType::Primitive(PType::I32, Nullability::NonNullable));
@@ -344,10 +344,10 @@ fn test_in_list_converts_to_set_lookup() -> DFResult<()> {
344344
vec![1i32.into(), 3i32.into(), 7i32.into()],
345345
Nullability::NonNullable,
346346
);
347-
assert_eq!(
348-
convert(expr, &schema)?,
349-
vortex::expr::in_list(get_item("a", root()), lit(list))
350-
);
347+
let membership = vortex::expr::in_list(get_item("a", root()), lit(list));
348+
let expected = if negated { not(membership) } else { membership };
349+
350+
assert_eq!(convert(expr, &schema)?, expected);
351351
Ok(())
352352
}
353353

‎vortex-duckdb/src/convert/expr.rs‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -751,9 +751,8 @@ fn try_from_compare_in(
751751
else {
752752
return Ok(None);
753753
};
754-
// The IN list itself is never null. Null values are elements, and in_list handles them.
755754
let element_dtype = list_elements[0].dtype().clone();
756-
let expr = in_list_of(element, element_dtype, list_elements);
755+
let expr = in_list_of(element, element_dtype, list_elements)?;
757756
Ok(Some(if not_in { not(expr) } else { expr }))
758757
}
759758

‎vortex-duckdb/src/convert/table_filter.rs‎

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -89,9 +89,7 @@ pub fn try_from_table_filter(
8989
"IN filter must have at least one value"
9090
);
9191
let dtype = scalars[0].dtype().clone();
92-
// The IN list itself is never null. Null values are elements, and in_list handles
93-
// them.
94-
in_list_of(col.clone(), dtype, scalars)
92+
in_list_of(col.clone(), dtype, scalars)?
9593
}
9694
TableFilterClass::Dynamic(dynamic) => {
9795
dynamic_filter_expr(dynamic, col.return_dtype(scope_dtype)?, col.clone())?

0 commit comments

Comments
 (0)