Skip to content

Commit 7b36f7f

Browse files
authored
fix(array): preserve constants when zip selects one branch (#11373)
When zip's mask is uniformly true or false (including null masks, which select the false branch), executing the selected ConstantArray materializes it unnecessarily. Scalar access then logs `Scalar function vortex.zip() returned non-constant array from execution over all scalar inputs`. Return the selected branch cast to the output dtype without executing it again. Removing the zip node already makes execution progress; the caller can preserve constants with `Columnar` or continue materializing with `Canonical`. Added a parameterized regression test for true, false, and null masks, with matching branch types and with nullable/non-nullable branches. It checks constant encoding, values, and output nullability. Validation: reviewed the patch and execution paths. Tests, formatting, and lints were not run locally, following AGENTS.md's instruction to leave verification execution to the user. Implemented and opened by Codex at Daniel King's request. Daniel King (@danking) is the responsible human. Signed-off-by: Daniel King <[email protected]>
1 parent 2e1b05b commit 7b36f7f

1 file changed

Lines changed: 28 additions & 2 deletions

File tree

  • vortex-array/src/scalar_fn/fns/zip

‎vortex-array/src/scalar_fn/fns/zip/mod.rs‎

Lines changed: 28 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -141,11 +141,11 @@ impl ScalarFnVTable for Zip {
141141
let return_dtype = zip_return_dtype(if_true.dtype(), if_false.dtype())?;
142142

143143
if mask.all_true() {
144-
return if_true.cast(return_dtype)?.execute(ctx);
144+
return if_true.cast(return_dtype);
145145
}
146146

147147
if mask.all_false() {
148-
return if_false.cast(return_dtype)?.execute(ctx);
148+
return if_false.cast(return_dtype);
149149
}
150150

151151
if !if_true.is_canonical() || !if_false.is_canonical() {
@@ -345,6 +345,7 @@ pub(crate) fn zip_validity(
345345
mod tests {
346346
use std::sync::LazyLock;
347347

348+
use rstest::rstest;
348349
use vortex_buffer::buffer;
349350
use vortex_error::VortexResult;
350351
use vortex_mask::Mask;
@@ -353,6 +354,7 @@ mod tests {
353354
use crate::ArrayRef;
354355
use crate::IntoArray;
355356
use crate::VortexSessionExecute;
357+
use crate::arrays::Constant;
356358
use crate::arrays::ConstantArray;
357359
use crate::arrays::PrimitiveArray;
358360
use crate::arrays::Struct;
@@ -437,6 +439,30 @@ mod tests {
437439
assert_eq!(result.dtype(), if_true.dtype());
438440
}
439441

442+
#[rstest]
443+
#[case::all_true(Some(true), 10)]
444+
#[case::all_false(Some(false), 20)]
445+
#[case::null_mask(None, 20)]
446+
fn test_zip_constant_mask_preserves_constant(
447+
#[case] mask: Option<bool>,
448+
#[case] expected: i32,
449+
#[values(Nullability::NonNullable, Nullability::Nullable)] nullability: Nullability,
450+
) -> VortexResult<()> {
451+
let mut ctx = SESSION.create_execution_ctx();
452+
let if_true = ConstantArray::new(10i32, 4).into_array();
453+
let if_false = ConstantArray::new(Scalar::primitive(20i32, nullability), 4).into_array();
454+
let result = ConstantArray::new(mask, 4)
455+
.into_array()
456+
.zip(if_true, if_false)?
457+
.execute::<Columnar>(&mut ctx)?
458+
.into_array();
459+
460+
assert!(result.is::<Constant>());
461+
let expected = ConstantArray::new(Scalar::primitive(expected, nullability), 4);
462+
assert_arrays_eq!(result, expected, &mut ctx);
463+
Ok(())
464+
}
465+
440466
#[test]
441467
fn test_invalid_lengths() {
442468
let mask = Mask::new_false(4);

0 commit comments

Comments
 (0)