implement DictArray batched execute

Signed-off-by: Connor Tsui <[email protected]>

Author: connortsui20

vortex-array/src/arrays/dict/vtable/mod.rs

@@ -2,18 +2,21 @@
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
 use vortex_buffer::BufferHandle;
+use vortex_compute::take::Take;
 use vortex_dtype::DType;
 use vortex_dtype::Nullability;
 use vortex_dtype::PType;
 use vortex_error::VortexResult;
 use vortex_error::vortex_bail;
 use vortex_error::vortex_err;
+use vortex_vector::Vector;
 
 use super::DictArray;
 use super::DictMetadata;
 use crate::DeserializeMetadata;
 use crate::ProstMetadata;
 use crate::SerializeMetadata;
+use crate::execution::ExecutionCtx;
 use crate::serde::ArrayChildren;
 use crate::vtable;
 use crate::vtable::ArrayId;
@@ -31,6 +34,9 @@ mod visitor;
 
 vtable!(Dict);
 
+#[cfg(test)]
+mod tests;
+
 #[derive(Debug)]
 pub struct DictVTable;
 
@@ -108,4 +114,10 @@ impl VTable for DictVTable {
             DictArray::new_unchecked(codes, values).set_all_values_referenced(all_values_referenced)
         })
     }
+
+    fn batch_execute(array: &DictArray, ctx: &mut ExecutionCtx) -> VortexResult<Vector> {
+        let values = array.values().batch_execute(ctx)?;
+        let codes = array.codes().batch_execute(ctx)?.into_primitive();
+        Ok(values.take(&codes))
+    }
 }

vortex-array/src/arrays/dict/vtable/tests.rs

@@ -0,0 +1,284 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! Tests for the optimized boolean take in dictionary batch execution.
+
+use vortex_buffer::buffer;
+use vortex_session::VortexSession;
+use vortex_vector::VectorOps;
+use vortex_vector::bool::BoolVector;
+
+use crate::IntoArray;
+use crate::arrays::BoolArray;
+use crate::arrays::PrimitiveArray;
+use crate::arrays::dict::DictArray;
+
+/// Helper to create a boolean dict array and execute it via `batch_execute`.
+fn bool_dict_execute(codes: impl IntoArray, values: impl IntoArray) -> BoolVector {
+    let dict = DictArray::try_new(codes.into_array(), values.into_array()).unwrap();
+    let session = VortexSession::empty();
+    dict.execute(&session).unwrap().into_bool()
+}
+
+/// Helper to assert two BoolVectors are equal by comparing element-by-element.
+fn assert_bool_vectors_eq(result: &BoolVector, expected: &[Option<bool>]) {
+    assert_eq!(result.len(), expected.len(), "length mismatch");
+    for (i, expected_val) in expected.iter().enumerate() {
+        let result_val = result.scalar_at(i);
+        assert_eq!(
+            result_val.value(),
+            *expected_val,
+            "mismatch at index {i}: result={:?}, expected={expected_val:?}",
+            result_val.value()
+        );
+    }
+}
+
+/// Test case: single true, single false (the common case).
+/// values = [false, true], codes = [0, 1, 0, 1, 1]
+/// expected = [false, true, false, true, true]
+#[test]
+fn test_bool_dict_single_true_single_false() {
+    let values = BoolArray::from_iter([false, true]);
+    let codes = buffer![0u8, 1, 0, 1, 1];
+
+    let result = bool_dict_execute(codes, values);
+    assert_bool_vectors_eq(
+        &result,
+        &[Some(false), Some(true), Some(false), Some(true), Some(true)],
+    );
+}
+
+/// Test case: all values are true (single true, no false).
+/// values = [true], codes = [0, 0, 0]
+/// expected = [true, true, true]
+#[test]
+fn test_bool_dict_all_true() {
+    let values = BoolArray::from_iter([true]);
+    let codes = buffer![0u8, 0, 0];
+
+    let result = bool_dict_execute(codes, values);
+    assert_bool_vectors_eq(&result, &[Some(true), Some(true), Some(true)]);
+}
+
+/// Test case: all values are false (single false, no true).
+/// values = [false], codes = [0, 0, 0]
+/// expected = [false, false, false]
+#[test]
+fn test_bool_dict_all_false() {
+    let values = BoolArray::from_iter([false]);
+    let codes = buffer![0u8, 0, 0];
+
+    let result = bool_dict_execute(codes, values);
+    assert_bool_vectors_eq(&result, &[Some(false), Some(false), Some(false)]);
+}
+
+/// Test case: multiple true values, no false values.
+/// values = [true, true], codes = [0, 1, 0, 1]
+/// expected = [true, true, true, true]
+#[test]
+fn test_bool_dict_multiple_true_no_false() {
+    let values = BoolArray::from_iter([true, true]);
+    let codes = buffer![0u8, 1, 0, 1];
+
+    let result = bool_dict_execute(codes, values);
+    assert_bool_vectors_eq(&result, &[Some(true), Some(true), Some(true), Some(true)]);
+}
+
+/// Test case: no true values, multiple false values.
+/// values = [false, false], codes = [0, 1, 0, 1]
+/// expected = [false, false, false, false]
+#[test]
+fn test_bool_dict_no_true_multiple_false() {
+    let values = BoolArray::from_iter([false, false]);
+    let codes = buffer![0u8, 1, 0, 1];
+
+    let result = bool_dict_execute(codes, values);
+    assert_bool_vectors_eq(
+        &result,
+        &[Some(false), Some(false), Some(false), Some(false)],
+    );
+}
+
+/// Test case: single true, multiple false values.
+/// values = [true, false, false], codes = [0, 1, 2, 0, 1]
+/// expected = [true, false, false, true, false]
+#[test]
+fn test_bool_dict_single_true_multiple_false() {
+    let values = BoolArray::from_iter([true, false, false]);
+    let codes = buffer![0u8, 1, 2, 0, 1];
+
+    let result = bool_dict_execute(codes, values);
+    assert_bool_vectors_eq(
+        &result,
+        &[
+            Some(true),
+            Some(false),
+            Some(false),
+            Some(true),
+            Some(false),
+        ],
+    );
+}
+
+/// Test case: multiple true values, single false.
+/// values = [true, true, false], codes = [0, 1, 2, 0, 2]
+/// expected = [true, true, false, true, false]
+#[test]
+fn test_bool_dict_multiple_true_single_false() {
+    let values = BoolArray::from_iter([true, true, false]);
+    let codes = buffer![0u8, 1, 2, 0, 2];
+
+    let result = bool_dict_execute(codes, values);
+    assert_bool_vectors_eq(
+        &result,
+        &[Some(true), Some(true), Some(false), Some(true), Some(false)],
+    );
+}
+
+/// Test case: multiple of both true and false (fallback case).
+/// values = [true, false, true, false], codes = [0, 1, 2, 3, 0]
+/// expected = [true, false, true, false, true]
+#[test]
+fn test_bool_dict_multiple_of_both() {
+    let values = BoolArray::from_iter([true, false, true, false]);
+    let codes = buffer![0u8, 1, 2, 3, 0];
+
+    let result = bool_dict_execute(codes, values);
+    assert_bool_vectors_eq(
+        &result,
+        &[Some(true), Some(false), Some(true), Some(false), Some(true)],
+    );
+}
+
+/// Test case: nullable codes (null code produces null output).
+/// values = [false, true], codes = [0, null, 1]
+/// expected = [false, null, true]
+#[test]
+fn test_bool_dict_nullable_codes() {
+    let values = BoolArray::from_iter([false, true]);
+    let codes = PrimitiveArray::from_option_iter([Some(0u8), None, Some(1)]);
+
+    let result = bool_dict_execute(codes, values);
+    assert_bool_vectors_eq(&result, &[Some(false), None, Some(true)]);
+}
+
+/// Test case: nullable values (code pointing to null produces null output).
+/// values = [null, true], codes = [0, 1, 0]
+/// expected = [null, true, null]
+#[test]
+fn test_bool_dict_nullable_values() {
+    let values = BoolArray::from_iter([None, Some(true)]);
+    let codes = buffer![0u8, 1, 0];
+
+    let result = bool_dict_execute(codes, values);
+    assert_bool_vectors_eq(&result, &[None, Some(true), None]);
+}
+
+/// Test case: all values are null (no valid true or false).
+/// values = [null, null], codes = [0, 1, 0]
+/// expected = [null, null, null]
+#[test]
+fn test_bool_dict_all_values_null() {
+    let values = BoolArray::from_iter([None::<bool>, None]);
+    let codes = buffer![0u8, 1, 0];
+
+    let result = bool_dict_execute(codes, values);
+    assert_bool_vectors_eq(&result, &[None, None, None]);
+}
+
+/// Test case: single non-null true with null values.
+/// values = [null, true], codes = [0, 1, 1, 0]
+/// expected = [null, true, true, null]
+#[test]
+fn test_bool_dict_single_true_with_nulls() {
+    let values = BoolArray::from_iter([None, Some(true)]);
+    let codes = buffer![0u8, 1, 1, 0];
+
+    let result = bool_dict_execute(codes, values);
+    assert_bool_vectors_eq(&result, &[None, Some(true), Some(true), None]);
+}
+
+/// Test case: single non-null false with null values.
+/// values = [null, false], codes = [0, 1, 1, 0]
+/// expected = [null, false, false, null]
+#[test]
+fn test_bool_dict_single_false_with_nulls() {
+    let values = BoolArray::from_iter([None, Some(false)]);
+    let codes = buffer![0u8, 1, 1, 0];
+
+    let result = bool_dict_execute(codes, values);
+    assert_bool_vectors_eq(&result, &[None, Some(false), Some(false), None]);
+}
+
+/// Test case: both nullable codes and nullable values.
+/// values = [null, true, false], codes = [0, null, 1, 2]
+/// expected = [null, null, true, false]
+#[test]
+fn test_bool_dict_nullable_codes_and_values() {
+    let values = BoolArray::from_iter([None, Some(true), Some(false)]);
+    let codes = PrimitiveArray::from_option_iter([Some(0u8), None, Some(1), Some(2)]);
+
+    let result = bool_dict_execute(codes, values);
+    assert_bool_vectors_eq(&result, &[None, None, Some(true), Some(false)]);
+}
+
+/// Test case: empty codes.
+/// values = [true, false], codes = []
+/// expected = []
+#[test]
+fn test_bool_dict_empty_codes() {
+    let values = BoolArray::from_iter([true, false]);
+    let codes = buffer![0u8; 0];
+
+    let result = bool_dict_execute(codes, values);
+    assert_bool_vectors_eq(&result, &[]);
+}
+
+/// Test case: larger index types (u32).
+/// values = [false, true], codes = [0u32, 1, 0, 1]
+/// expected = [false, true, false, true]
+#[test]
+fn test_bool_dict_u32_codes() {
+    let values = BoolArray::from_iter([false, true]);
+    let codes = buffer![0u32, 1, 0, 1];
+
+    let result = bool_dict_execute(codes, values);
+    assert_bool_vectors_eq(&result, &[Some(false), Some(true), Some(false), Some(true)]);
+}
+
+/// Test case: single true at non-zero index.
+/// values = [false, false, true], codes = [0, 1, 2, 2, 0]
+/// expected = [false, false, true, true, false]
+#[test]
+fn test_bool_dict_single_true_at_nonzero_index() {
+    let values = BoolArray::from_iter([false, false, true]);
+    let codes = buffer![0u8, 1, 2, 2, 0];
+
+    let result = bool_dict_execute(codes, values);
+    assert_bool_vectors_eq(
+        &result,
+        &[
+            Some(false),
+            Some(false),
+            Some(true),
+            Some(true),
+            Some(false),
+        ],
+    );
+}
+
+/// Test case: single false at non-zero index.
+/// values = [true, true, false], codes = [0, 1, 2, 2, 0]
+/// expected = [true, true, false, false, true]
+#[test]
+fn test_bool_dict_single_false_at_nonzero_index() {
+    let values = BoolArray::from_iter([true, true, false]);
+    let codes = buffer![0u8, 1, 2, 2, 0];
+
+    let result = bool_dict_execute(codes, values);
+    assert_bool_vectors_eq(
+        &result,
+        &[Some(true), Some(true), Some(false), Some(false), Some(true)],
+    );
+}