fix reduction/execute cycle

Signed-off-by: Connor Tsui <connor.tsui20@gmail.com>

Author: connortsui20

vortex-array/src/arrays/constant/vtable/canonical.rs

@@ -10,6 +10,7 @@ use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
 
 use crate::Canonical;
+use crate::ExecutionCtx;
 use crate::IntoArray;
 use crate::array::ArrayView;
 use crate::arrays::BoolArray;
@@ -36,7 +37,10 @@ use crate::scalar::Scalar;
 use crate::validity::Validity;
 
 /// Shared implementation for both `canonicalize` and `execute` methods.
-pub(crate) fn constant_canonicalize(array: ArrayView<'_, Constant>) -> VortexResult<Canonical> {
+pub(crate) fn constant_canonicalize(
+    array: ArrayView<'_, Constant>,
+    ctx: &mut ExecutionCtx,
+) -> VortexResult<Canonical> {
     let scalar = array.scalar();
 
     let validity = match array.dtype().nullability() {
@@ -163,7 +167,16 @@ pub(crate) fn constant_canonicalize(array: ArrayView<'_, Constant>) -> VortexRes
             let s = scalar.as_extension();
 
             let storage_scalar = s.to_storage_scalar();
-            let storage_self = ConstantArray::new(storage_scalar, array.len()).into_array();
+
+            // NB: We need to execute the constant array to be canonical because there is a
+            // reduction rule that turns `Extension(Constant(..))` into `Constant(Extension(..))`,
+            // and if we don't do this we create an infinite cycle.
+            // See `ExtensionConstantRule` for more details.
+            let storage_self = ConstantArray::new(storage_scalar, array.len())
+                .into_array()
+                .execute::<Canonical>(ctx)?
+                .into_array();
+
             Canonical::Extension(ExtensionArray::new(ext_dtype.clone(), storage_self))
         }
         DType::Variant(_) => {

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

@@ -160,9 +160,10 @@ impl VTable for Constant {
         PARENT_RULES.evaluate(array, parent, child_idx)
     }
 
-    fn execute(array: Array<Self>, _ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
+    fn execute(array: Array<Self>, ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
         Ok(ExecutionResult::done(constant_canonicalize(
             array.as_view(),
+            ctx,
         )?))
     }
 
@@ -268,10 +269,11 @@ fn append_value_or_nulls<B: ArrayBuilder + 'static>(
 #[cfg(test)]
 mod tests {
     use rstest::rstest;
+    use vortex_error::VortexResult;
     use vortex_session::VortexSession;
 
-    use crate::ExecutionCtx;
     use crate::IntoArray;
+    use crate::VortexSessionExecute;
     use crate::arrays::ConstantArray;
     use crate::arrays::constant::vtable::canonical::constant_canonicalize;
     use crate::assert_arrays_eq;
@@ -282,42 +284,37 @@ mod tests {
     use crate::dtype::StructFields;
     use crate::scalar::Scalar;
 
-    fn ctx() -> ExecutionCtx {
-        ExecutionCtx::new(VortexSession::empty())
-    }
-
     /// Appends `array` into a fresh builder and asserts the result matches `constant_canonicalize`.
-    fn assert_append_matches_canonical(array: ConstantArray) -> vortex_error::VortexResult<()> {
-        let expected = constant_canonicalize(array.as_view())?.into_array();
+    fn assert_append_matches_canonical(array: ConstantArray) -> VortexResult<()> {
+        let mut ctx = VortexSession::empty().create_execution_ctx();
+
+        let expected = constant_canonicalize(array.as_view(), &mut ctx)?.into_array();
         let mut builder = builder_with_capacity(array.dtype(), array.len());
         array
             .into_array()
-            .append_to_builder(builder.as_mut(), &mut ctx())?;
+            .append_to_builder(builder.as_mut(), &mut ctx)?;
         let result = builder.finish();
         assert_arrays_eq!(&result, &expected);
         Ok(())
     }
 
     #[test]
-    fn test_null_constant_append() -> vortex_error::VortexResult<()> {
+    fn test_null_constant_append() -> VortexResult<()> {
         assert_append_matches_canonical(ConstantArray::new(Scalar::null(DType::Null), 5))
     }
 
     #[rstest]
     #[case::bool_true(true, 5)]
     #[case::bool_false(false, 3)]
-    fn test_bool_constant_append(
-        #[case] value: bool,
-        #[case] n: usize,
-    ) -> vortex_error::VortexResult<()> {
+    fn test_bool_constant_append(#[case] value: bool, #[case] n: usize) -> VortexResult<()> {
         assert_append_matches_canonical(ConstantArray::new(
             Scalar::bool(value, Nullability::NonNullable),
             n,
         ))
     }
 
     #[test]
-    fn test_bool_null_constant_append() -> vortex_error::VortexResult<()> {
+    fn test_bool_null_constant_append() -> VortexResult<()> {
         assert_append_matches_canonical(ConstantArray::new(
             Scalar::null(DType::Bool(Nullability::Nullable)),
             4,
@@ -332,7 +329,7 @@ mod tests {
     fn test_primitive_constant_append(
         #[case] scalar: Scalar,
         #[case] n: usize,
-    ) -> vortex_error::VortexResult<()> {
+    ) -> VortexResult<()> {
         assert_append_matches_canonical(ConstantArray::new(scalar, n))
     }
 
@@ -341,18 +338,15 @@ mod tests {
     #[case::utf8_noninline("hello world!!", 5)] // >12 bytes: requires buffer block
     #[case::utf8_empty("", 3)]
     #[case::utf8_n_zero("hello world!!", 0)] // n=0 with non-inline: must not write orphaned bytes
-    fn test_utf8_constant_append(
-        #[case] value: &str,
-        #[case] n: usize,
-    ) -> vortex_error::VortexResult<()> {
+    fn test_utf8_constant_append(#[case] value: &str, #[case] n: usize) -> VortexResult<()> {
         assert_append_matches_canonical(ConstantArray::new(
             Scalar::utf8(value, Nullability::NonNullable),
             n,
         ))
     }
 
     #[test]
-    fn test_utf8_null_constant_append() -> vortex_error::VortexResult<()> {
+    fn test_utf8_null_constant_append() -> VortexResult<()> {
         assert_append_matches_canonical(ConstantArray::new(
             Scalar::null(DType::Utf8(Nullability::Nullable)),
             4,
@@ -362,26 +356,23 @@ mod tests {
     #[rstest]
     #[case::binary_inline(vec![1u8, 2, 3], 5)] // ≤12 bytes: inlined
     #[case::binary_noninline(vec![0u8; 13], 5)] // >12 bytes: buffer block
-    fn test_binary_constant_append(
-        #[case] value: Vec<u8>,
-        #[case] n: usize,
-    ) -> vortex_error::VortexResult<()> {
+    fn test_binary_constant_append(#[case] value: Vec<u8>, #[case] n: usize) -> VortexResult<()> {
         assert_append_matches_canonical(ConstantArray::new(
             Scalar::binary(value, Nullability::NonNullable),
             n,
         ))
     }
 
     #[test]
-    fn test_binary_null_constant_append() -> vortex_error::VortexResult<()> {
+    fn test_binary_null_constant_append() -> VortexResult<()> {
         assert_append_matches_canonical(ConstantArray::new(
             Scalar::null(DType::Binary(Nullability::Nullable)),
             4,
         ))
     }
 
     #[test]
-    fn test_struct_constant_append() -> vortex_error::VortexResult<()> {
+    fn test_struct_constant_append() -> VortexResult<()> {
         let fields = StructFields::new(
             ["x", "y"].into(),
             vec![
@@ -400,7 +391,7 @@ mod tests {
     }
 
     #[test]
-    fn test_null_struct_constant_append() -> vortex_error::VortexResult<()> {
+    fn test_null_struct_constant_append() -> VortexResult<()> {
         let fields = StructFields::new(
             ["x"].into(),
             vec![DType::Primitive(PType::I32, Nullability::Nullable)],

vortex-array/src/arrays/extension/vtable/canonical.rs

@@ -1,9 +1,5 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
-mod canonical;
-mod kernel;
-mod operations;
-mod validity;
 
 use std::hash::Hasher;
 
@@ -36,6 +32,10 @@ use crate::buffer::BufferHandle;
 use crate::dtype::DType;
 use crate::serde::ArrayChildren;
 
+mod kernel;
+mod operations;
+mod validity;
+
 #[derive(Clone, Debug)]
 pub struct Extension;