Skip to content

Commit d2ddcf8

Browse files
authored
Merge branch 'develop' into alfonso/lower-execute-step
2 parents 0d135ba + 3f55648 commit d2ddcf8

4 files changed

Lines changed: 178 additions & 29 deletions

File tree

‎vortex-array/src/arrays/dict/compute/cast.rs‎

Lines changed: 70 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -16,8 +16,20 @@ use crate::scalar_fn::fns::cast::CastReduce;
1616

1717
impl CastReduce for Dict {
1818
fn cast(array: ArrayView<'_, Dict>, dtype: &DType) -> VortexResult<Option<ArrayRef>> {
19-
// Can have un-reference null values making the cast of values fail without a possible mask.
20-
// TODO(joe): optimize this, could look at accessible values and fill_null not those?
19+
// Casting unused entries must not introduce errors that a cast over the rows would skip.
20+
if !array.has_all_values_referenced() {
21+
let safe_for_unused_values = match (array.dtype(), dtype) {
22+
(DType::Primitive(from, _), DType::Primitive(to, _)) => {
23+
from.casts_losslessly_to(*to)
24+
}
25+
(from, to) => from.as_nullable() == to.as_nullable(),
26+
};
27+
if !safe_for_unused_values {
28+
return Ok(None);
29+
}
30+
}
31+
32+
// An unused null entry must not make a cast to a non-nullable dtype fail.
2133
if !dtype.is_nullable() && !array.values().validity()?.definitely_no_nulls() {
2234
return Ok(None);
2335
}
@@ -51,9 +63,11 @@ mod tests {
5163
use rstest::rstest;
5264
use vortex_buffer::BitBuffer;
5365
use vortex_buffer::buffer;
66+
use vortex_error::VortexResult;
5467
use vortex_session::VortexSession;
5568

5669
use crate::ArrayRef;
70+
use crate::Canonical;
5771
use crate::IntoArray;
5872
use crate::VortexSessionExecute;
5973
use crate::arrays::Dict;
@@ -229,4 +243,58 @@ mod tests {
229243
);
230244
assert_arrays_eq!(casted, PrimitiveArray::from_iter([1.0f64, 3.0, 1.0]), ctx);
231245
}
246+
247+
/// Slicing keeps the whole dictionary, so the value `1000` stays in `values` although no
248+
/// remaining row references it. The cast must only consider the rows: `[1, 1, 1]` fits in i8.
249+
#[test]
250+
fn test_cast_dict_after_slice_ignores_unreferenced_values() -> VortexResult<()> {
251+
let ctx = &mut SESSION.create_execution_ctx();
252+
let dict = dict_encode(&buffer![1000i64, 1, 1, 1].into_array(), ctx)?;
253+
let rows = dict.into_array().slice(1..4)?.optimize_ctx(&SESSION)?;
254+
255+
let casted = rows
256+
.cast(DType::Primitive(PType::I8, Nullability::NonNullable))?
257+
.optimize_ctx(&SESSION)?
258+
.execute::<Canonical>(ctx)?;
259+
260+
assert_arrays_eq!(casted, PrimitiveArray::from_iter([1i8, 1, 1]), ctx);
261+
Ok(())
262+
}
263+
264+
#[test]
265+
fn test_cast_dict_with_unreferenced_values_preserves_nullability_rewrite() -> VortexResult<()> {
266+
let ctx = &mut SESSION.create_execution_ctx();
267+
let dict = DictArray::try_new(
268+
buffer![1u8, 1, 1].into_array(),
269+
buffer![1000i64, 1].into_array(),
270+
)?
271+
.into_array();
272+
273+
let nullable = dict
274+
.cast(DType::Primitive(PType::I64, Nullability::Nullable))?
275+
.optimize_ctx(&SESSION)?;
276+
assert!(nullable.is::<Dict>());
277+
let non_nullable = nullable
278+
.cast(DType::Primitive(PType::I64, Nullability::NonNullable))?
279+
.optimize_ctx(&SESSION)?;
280+
assert!(non_nullable.is::<Dict>());
281+
282+
assert_arrays_eq!(non_nullable, PrimitiveArray::from_iter([1i64, 1, 1]), ctx);
283+
Ok(())
284+
}
285+
286+
#[test]
287+
fn test_cast_dict_rejects_referenced_out_of_range_value() -> VortexResult<()> {
288+
let ctx = &mut SESSION.create_execution_ctx();
289+
let dict = dict_encode(&buffer![1000i64, 1, 1, 1].into_array(), ctx)?;
290+
let casted = dict
291+
.into_array()
292+
.cast(DType::Primitive(PType::I8, Nullability::NonNullable))?
293+
.optimize_ctx(&SESSION)?;
294+
295+
assert!(casted.is::<Dict>());
296+
let error = casted.execute::<Canonical>(ctx).unwrap_err();
297+
assert!(error.to_string().contains("value exceeds target range"));
298+
Ok(())
299+
}
232300
}

‎vortex-array/src/arrays/dict/compute/rules.rs‎

Lines changed: 81 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -268,6 +268,12 @@ impl ArrayParentReduceRule<Dict> for DictionaryScalarFnCodesPullUpRule {
268268
return Ok(None);
269269
}
270270

271+
// Pulling up the codes evaluates every dictionary value, including unreferenced entries
272+
// on which a fallible function could fail even though no row would evaluate them.
273+
if !array.has_all_values_referenced() && !parent.scalar_fn().signature().is_infallible() {
274+
return Ok(None);
275+
}
276+
271277
// Check that all siblings are dictionaries, and have the same number of values as us.
272278
// This is a cheap first loop.
273279
if !parent.iter_children().enumerate().all(|(idx, c)| {
@@ -310,11 +316,13 @@ impl ArrayParentReduceRule<Dict> for DictionaryScalarFnCodesPullUpRule {
310316
mod tests {
311317
use std::sync::LazyLock;
312318

319+
use rstest::rstest;
313320
use vortex_buffer::buffer;
314321
use vortex_error::VortexResult;
315322
use vortex_session::VortexSession;
316323

317324
use crate::ArrayRef;
325+
use crate::Canonical;
318326
use crate::IntoArray;
319327
use crate::arrays::BoolArray;
320328
use crate::arrays::Chunked;
@@ -444,4 +452,77 @@ mod tests {
444452
assert_arrays_eq!(optimized, expected, &mut ctx);
445453
Ok(())
446454
}
455+
456+
/// Both slices share the codes `[1, 1, 1]`, so the codes pull-up divides the two values
457+
/// arrays. The divisor still holds the unreferenced `0`, which no row divides by, so the
458+
/// result must be `[8, 8, 8] / [2, 2, 2]`.
459+
#[test]
460+
fn codes_pull_up_skips_fallible_fn_over_unreferenced_values() -> VortexResult<()> {
461+
let mut ctx = SESSION.create_execution_ctx();
462+
let dividend = DictArray::try_new(
463+
buffer![0u8, 1, 1, 1].into_array(),
464+
buffer![7i32, 8].into_array(),
465+
)?
466+
.into_array()
467+
.slice(1..4)?
468+
.optimize_ctx(&SESSION)?;
469+
let divisor = DictArray::try_new(
470+
buffer![0u8, 1, 1, 1].into_array(),
471+
buffer![0i32, 2].into_array(),
472+
)?
473+
.into_array()
474+
.slice(1..4)?
475+
.optimize_ctx(&SESSION)?;
476+
let expr = Binary::try_new(dividend, divisor, Operator::Div)?.into_array();
477+
478+
let quotient = expr
479+
.optimize_ctx(&SESSION)?
480+
.execute::<Canonical>(&mut ctx)?;
481+
482+
assert_arrays_eq!(quotient, PrimitiveArray::from_iter([4i32, 4, 4]), &mut ctx);
483+
Ok(())
484+
}
485+
486+
#[rstest]
487+
#[case::infallible(false, Operator::Eq, BoolArray::from_iter([false, true, false, true]).into_array())]
488+
#[case::all_values_referenced(true, Operator::Div, buffer![4i32, 1, 4, 1].into_array())]
489+
fn codes_pull_up_preserves_safe_rewrites(
490+
#[case] all_values_referenced: bool,
491+
#[case] operator: Operator,
492+
#[case] expected: ArrayRef,
493+
) -> VortexResult<()> {
494+
let codes = buffer![0u8, 1, 0, 1].into_array();
495+
// SAFETY: Both dictionary entries are referenced by valid codes.
496+
let lhs = unsafe {
497+
DictArray::try_new(codes.clone(), buffer![8i32, 9].into_array())?
498+
.set_all_values_referenced(all_values_referenced)
499+
}
500+
.into_array();
501+
let rhs = DictArray::try_new(codes, buffer![2i32, 9].into_array())?.into_array();
502+
503+
let optimized = Binary::try_new(lhs, rhs, operator)?
504+
.into_array()
505+
.optimize_ctx(&SESSION)?;
506+
assert!(optimized.is::<Dict>());
507+
let mut ctx = SESSION.create_execution_ctx();
508+
let result = optimized.execute::<Canonical>(&mut ctx)?;
509+
assert_arrays_eq!(result, expected, &mut ctx);
510+
Ok(())
511+
}
512+
513+
#[test]
514+
fn codes_pull_up_rejects_referenced_zero_divisor() -> VortexResult<()> {
515+
let codes = buffer![0u8, 1, 0, 1].into_array();
516+
let dividend =
517+
DictArray::try_new(codes.clone(), buffer![7i32, 8].into_array())?.into_array();
518+
let divisor = DictArray::try_new(codes, buffer![0i32, 2].into_array())?.into_array();
519+
let expr = Binary::try_new(dividend, divisor, Operator::Div)?.into_array();
520+
521+
let error = expr
522+
.optimize_ctx(&SESSION)?
523+
.execute::<Canonical>(&mut SESSION.create_execution_ctx())
524+
.unwrap_err();
525+
assert!(error.to_string().contains("division by zero"));
526+
Ok(())
527+
}
447528
}

‎vortex-array/src/arrays/primitive/compute/cast.rs‎

Lines changed: 1 addition & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -519,35 +519,9 @@ where
519519
)
520520
};
521521

522-
// Returns `true` if every value of `from` is representable in `to` without loss.
523-
fn casts_losslessly_to(from: PType, to: PType) -> bool {
524-
if from == to {
525-
return true;
526-
}
527-
if (from.is_unsigned_int() && to.is_unsigned_int())
528-
|| (from.is_signed_int() && to.is_signed_int())
529-
|| (from.is_float() && to.is_float())
530-
{
531-
return from.byte_width() <= to.byte_width();
532-
}
533-
if from.is_unsigned_int() && to.is_signed_int() {
534-
return from.byte_width() < to.byte_width();
535-
}
536-
if from.is_int() && to.is_float() {
537-
let minimum_float_width = match from.byte_width() {
538-
1 => 2,
539-
2 => 4,
540-
4 => 8,
541-
_ => return false,
542-
};
543-
return to.byte_width() >= minimum_float_width;
544-
}
545-
false
546-
}
547-
548522
// Skip the fallible kernel when type widening or (cached) min/max prove every value fits.
549523
let target_dtype = DType::Primitive(T::PTYPE, Nullability::NonNullable);
550-
let infallible = casts_losslessly_to(F::PTYPE, T::PTYPE)
524+
let infallible = F::PTYPE.casts_losslessly_to(T::PTYPE)
551525
|| cached_values_fit_in(array, &target_dtype).unwrap_or(false);
552526

553527
let len = array.len();

‎vortex-array/src/dtype/ptype.rs‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -827,6 +827,32 @@ impl PType {
827827
self.byte_width() * 8
828828
}
829829

830+
/// Returns whether every value of this type is representable in `to` without loss.
831+
pub(crate) fn casts_losslessly_to(self, to: Self) -> bool {
832+
if self == to {
833+
return true;
834+
}
835+
if (self.is_unsigned_int() && to.is_unsigned_int())
836+
|| (self.is_signed_int() && to.is_signed_int())
837+
|| (self.is_float() && to.is_float())
838+
{
839+
return self.byte_width() <= to.byte_width();
840+
}
841+
if self.is_unsigned_int() && to.is_signed_int() {
842+
return self.byte_width() < to.byte_width();
843+
}
844+
if self.is_int() && to.is_float() {
845+
let minimum_float_width = match self.byte_width() {
846+
1 => 2,
847+
2 => 4,
848+
4 => 8,
849+
_ => return false,
850+
};
851+
return to.byte_width() >= minimum_float_width;
852+
}
853+
false
854+
}
855+
830856
/// Returns the maximum value of this PType if it is an integer type
831857
/// Returns `u64::MAX` if the value is too large to fit in a `u64`
832858
#[inline]

0 commit comments

Comments
 (0)