[do not merge] list transform array with templatearray instantiation

vortex-array/src/arrays/list_transform/array.rs

@@ -0,0 +1,163 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use vortex_error::VortexExpect;
+use vortex_error::VortexResult;
+use vortex_error::vortex_bail;
+use vortex_error::vortex_ensure;
+
+use crate::ArrayRef;
+use crate::ArraySlots;
+use crate::array::Array;
+use crate::array::ArrayParts;
+use crate::array::EmptyArrayData;
+use crate::array::TypedArrayRef;
+use crate::arrays::ListTransform;
+use crate::arrays::list_transform::template::build_template;
+use crate::dtype::DType;
+use crate::expr::BoundLambda;
+
+/// A lazy list transformation has structural children only:
+/// `list`, a zero-length lambda `body`, then outer-row capture arrays.
+pub trait ListTransformArrayExt: TypedArrayRef<ListTransform> {
+    fn list(&self) -> &ArrayRef {
+        self.as_ref().slots()[0]
+            .as_ref()
+            .vortex_expect("validated ListTransformArray list slot")
+    }
+
+    fn body(&self) -> &ArrayRef {
+        self.as_ref().slots()[1]
+            .as_ref()
+            .vortex_expect("validated ListTransformArray body slot")
+    }
+
+    fn captures(&self) -> impl Iterator<Item = &ArrayRef> + '_ {
+        self.as_ref().slots()[2..].iter().map(|capture| {
+            capture
+                .as_ref()
+                .vortex_expect("validated ListTransformArray capture slot")
+        })
+    }
+
+    fn capture_count(&self) -> usize {
+        self.as_ref().slots().len() - 2
+    }
+}
+impl<T: TypedArrayRef<ListTransform>> ListTransformArrayExt for T {}
+
+impl Array<ListTransform> {
+    /// Build a structural list transform from a bound lambda and outer-row captures.
+    pub fn try_new(
+        list: ArrayRef,
+        lambda: BoundLambda,
+        captures: impl IntoIterator<Item = ArrayRef>,
+    ) -> VortexResult<Self> {
+        let captures = captures.into_iter().collect::<Vec<_>>();
+        validate_lambda(&list, &lambda, &captures)?;
+        let body = build_template(&lambda)?;
+        Self::try_new_from_parts(list, body, captures)
+    }
+
+    /// Rebuild a transform after substituting its outer template inputs.
+    pub(crate) fn try_new_from_parts(
+        list: ArrayRef,
+        body: ArrayRef,
+        captures: impl IntoIterator<Item = ArrayRef>,
+    ) -> VortexResult<Self> {
+        let captures = captures.into_iter().collect::<Vec<_>>();
+        let dtype = output_dtype(list.dtype(), body.dtype())?;
+        vortex_ensure!(
+            body.is_empty(),
+            "ListTransformArray body must be a zero-length template, got {}",
+            body.len()
+        );
+        vortex_ensure!(
+            captures.iter().all(|capture| capture.len() == list.len()),
+            "ListTransformArray captures must have the outer list length {}",
+            list.len()
+        );
+        let len = list.len();
+        let slots = std::iter::once(list)
+            .chain(std::iter::once(body))
+            .chain(captures)
+            .map(Some)
+            .collect::<ArraySlots>();
+        Array::try_from_parts(
+            ArrayParts::new(ListTransform, dtype, len, EmptyArrayData).with_slots(slots),
+        )
+    }
+}
+
+pub(crate) fn output_dtype(list: &DType, body: &DType) -> VortexResult<DType> {
+    match list {
+        DType::List(_, nullability) => Ok(DType::List(body.clone().into(), *nullability)),
+        DType::FixedSizeList(_, size, nullability) => Ok(DType::FixedSizeList(
+            body.clone().into(),
+            *size,
+            *nullability,
+        )),
+        _ => vortex_bail!("list_transform() requires List, ListView, or FixedSizeList, got {list}"),
+    }
+}
+
+fn validate_lambda(
+    list: &ArrayRef,
+    lambda: &BoundLambda,
+    captures: &[ArrayRef],
+) -> VortexResult<()> {
+    let element_dtype = match list.dtype() {
+        DType::List(element, _) | DType::FixedSizeList(element, ..) => element.as_ref(),
+        _ => vortex_bail!(
+            "list_transform() requires List, ListView, or FixedSizeList, got {}",
+            list.dtype()
+        ),
+    };
+    vortex_ensure!(
+        matches!(lambda.param_dtypes().len(), 1 | 2),
+        "list_transform() lambda must take one or two parameters, got {}",
+        lambda.param_dtypes().len()
+    );
+    vortex_ensure!(
+        lambda.param_dtypes()[0] == *element_dtype,
+        "list_transform() element parameter expects dtype {}, got {}",
+        lambda.param_dtypes()[0],
+        element_dtype
+    );
+    if lambda.param_dtypes().len() == 2 {
+        let index = DType::Primitive(
+            crate::dtype::PType::U64,
+            crate::dtype::Nullability::NonNullable,
+        );
+        vortex_ensure!(
+            lambda.param_dtypes()[1] == index,
+            "list_transform() index parameter expects dtype {index}, got {}",
+            lambda.param_dtypes()[1]
+        );
+    }
+    vortex_ensure!(
+        lambda.captures().len() == captures.len(),
+        "list_transform() lambda requires {} captures, got {}",
+        lambda.captures().len(),
+        captures.len()
+    );
+    for (index, (capture, array)) in lambda.captures().iter().zip(captures).enumerate() {
+        vortex_ensure!(
+            capture.dtype() == array.dtype(),
+            "list_transform() capture {index} expects dtype {}, got {}",
+            capture.dtype(),
+            array.dtype()
+        );
+        vortex_ensure!(
+            array.len() == list.len(),
+            "list_transform() capture {index} has length {}, expected {}",
+            array.len(),
+            list.len()
+        );
+    }
+    vortex_ensure!(
+        lambda.body().is_root_bound_to(element_dtype),
+        "list_transform() lambda root expects a different dtype than {element_dtype}"
+    );
+    Ok(())
+}

vortex-array/src/arrays/list_transform/mod.rs

@@ -0,0 +1,14 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+pub(crate) mod array;
+mod rules;
+mod template;
+mod vtable;
+
+pub use array::ListTransformArrayExt;
+pub use vtable::ListTransform;
+pub use vtable::ListTransformArray;
+
+#[cfg(test)]
+mod tests;

vortex-array/src/arrays/list_transform/rules.rs

@@ -0,0 +1,93 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use vortex_error::VortexResult;
+
+use crate::ArrayRef;
+use crate::IntoArray;
+use crate::array::ArrayView;
+use crate::arrays::Filter;
+use crate::arrays::ListTransform;
+use crate::arrays::ListTransformArray;
+use crate::arrays::ListTransformArrayExt;
+use crate::arrays::Slice;
+use crate::arrays::dict::TakeReduce;
+use crate::arrays::dict::TakeReduceAdaptor;
+use crate::optimizer::rules::ArrayParentReduceRule;
+use crate::optimizer::rules::ParentRuleSet;
+
+pub(super) const PARENT_RULES: ParentRuleSet<ListTransform> = ParentRuleSet::new(&[
+    ParentRuleSet::lift(&ListTransformSlicePushDown),
+    ParentRuleSet::lift(&ListTransformFilterPushDown),
+    ParentRuleSet::lift(&TakeReduceAdaptor(ListTransform)),
+]);
+
+#[derive(Debug)]
+struct ListTransformSlicePushDown;
+
+impl ArrayParentReduceRule<ListTransform> for ListTransformSlicePushDown {
+    type Parent = Slice;
+
+    fn reduce_parent(
+        &self,
+        transform: ArrayView<'_, ListTransform>,
+        parent: ArrayView<'_, Slice>,
+        _child_idx: usize,
+    ) -> VortexResult<Option<ArrayRef>> {
+        let range = parent.slice_range().clone();
+        Ok(Some(
+            ListTransformArray::try_new_from_parts(
+                transform.list().slice(range.clone())?,
+                transform.body().clone(),
+                transform
+                    .captures()
+                    .map(|capture| capture.slice(range.clone()))
+                    .collect::<VortexResult<Vec<_>>>()?,
+            )?
+            .into_array(),
+        ))
+    }
+}
+
+#[derive(Debug)]
+struct ListTransformFilterPushDown;
+
+impl ArrayParentReduceRule<ListTransform> for ListTransformFilterPushDown {
+    type Parent = Filter;
+
+    fn reduce_parent(
+        &self,
+        transform: ArrayView<'_, ListTransform>,
+        parent: ArrayView<'_, Filter>,
+        _child_idx: usize,
+    ) -> VortexResult<Option<ArrayRef>> {
+        let mask = parent.filter_mask().clone();
+        Ok(Some(
+            ListTransformArray::try_new_from_parts(
+                transform.list().filter(mask.clone())?,
+                transform.body().clone(),
+                transform
+                    .captures()
+                    .map(|capture| capture.filter(mask.clone()))
+                    .collect::<VortexResult<Vec<_>>>()?,
+            )?
+            .into_array(),
+        ))
+    }
+}
+
+impl TakeReduce for ListTransform {
+    fn take(transform: ArrayView<'_, Self>, indices: &ArrayRef) -> VortexResult<Option<ArrayRef>> {
+        Ok(Some(
+            ListTransformArray::try_new_from_parts(
+                transform.list().take(indices.clone())?,
+                transform.body().clone(),
+                transform
+                    .captures()
+                    .map(|capture| capture.take(indices.clone()))
+                    .collect::<VortexResult<Vec<_>>>()?,
+            )?
+            .into_array(),
+        ))
+    }
+}

vortex-array/src/arrays/list_transform/template.rs

@@ -0,0 +1,112 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use itertools::Itertools;
+use vortex_error::VortexResult;
+use vortex_error::vortex_bail;
+use vortex_error::vortex_ensure;
+
+use crate::ArrayRef;
+use crate::IntoArray;
+use crate::arrays::ConstantArray;
+use crate::arrays::ScalarFnArray;
+use crate::arrays::TemplateInputArray;
+use crate::arrays::template::TemplateInputArrayExt;
+use crate::arrays::template::TemplateScope;
+use crate::expr::BoundExpression;
+use crate::expr::BoundLambda;
+use crate::scalar_fn::fns::literal::Literal;
+
+/// Convert one bound lambda into its zero-length, lexically scoped body template.
+pub(super) fn build_template(lambda: &BoundLambda) -> VortexResult<ArrayRef> {
+    let scope = TemplateScope::fresh();
+    // Slots 0 and 1 are permanently reserved for the element and optional local index. Captures
+    // always start at 2, which keeps a one-parameter lambda's first capture distinct from the
+    // index slot after the bound lambda itself has been discarded.
+    let mut inputs = vec![None; 2 + lambda.captures().len()];
+    inputs[0] =
+        Some(TemplateInputArray::new(scope, 0, lambda.param_dtypes()[0].clone()).into_array());
+    if lambda.param_dtypes().len() == 2 {
+        inputs[1] =
+            Some(TemplateInputArray::new(scope, 1, lambda.param_dtypes()[1].clone()).into_array());
+    }
+    for (index, capture) in lambda.captures().iter().enumerate() {
+        inputs[index + 2] =
+            Some(TemplateInputArray::new(scope, index + 2, capture.dtype().clone()).into_array());
+    }
+    build_expression(lambda.body(), lambda, scope, &inputs)
+}
+
+fn build_expression(
+    expression: &BoundExpression,
+    lambda: &BoundLambda,
+    scope: TemplateScope,
+    inputs: &[Option<ArrayRef>],
+) -> VortexResult<ArrayRef> {
+    match expression {
+        BoundExpression::Root { .. } => template_input(inputs, 0),
+        BoundExpression::Variable(variable) => {
+            let slot = lambda
+                .param_refs()
+                .iter()
+                .position(|reference| *reference == variable.variable_ref())
+                .or_else(|| {
+                    lambda
+                        .captures()
+                        .iter()
+                        .position(|capture| capture.variable_ref() == variable.variable_ref())
+                        .map(|index| 2 + index)
+                })
+                .ok_or_else(|| {
+                    vortex_error::vortex_err!(
+                        "variable '{}' is unresolved while building a template",
+                        variable
+                    )
+                })?;
+            let input = template_input(inputs, slot)?;
+            vortex_ensure!(
+                input.as_::<crate::arrays::TemplateInput>().scope() == scope,
+                "template builder mixed scopes"
+            );
+            Ok(input)
+        }
+        BoundExpression::Scalar {
+            scalar_fn,
+            children,
+            ..
+        } => {
+            if let Some(value) = scalar_fn.as_opt::<Literal>() {
+                return Ok(ConstantArray::new(value.clone(), 0).into_array());
+            }
+            let children = children
+                .iter()
+                .map(|child| build_expression(child, lambda, scope, inputs))
+                .try_collect()?;
+            Ok(ScalarFnArray::try_new_with_len(scalar_fn.clone(), children, 0)?.into_array())
+        }
+        BoundExpression::Lambda(_) => {
+            vortex_bail!("a detached lambda cannot appear in a template body")
+        }
+        BoundExpression::ListTransform {
+            lambda: nested_lambda,
+            children,
+            ..
+        } => {
+            let list = build_expression(&children[0], lambda, scope, inputs)?;
+            let captures = children[1..]
+                .iter()
+                .map(|capture| build_expression(capture, lambda, scope, inputs))
+                .collect::<VortexResult<Vec<_>>>()?;
+            crate::arrays::ListTransformArray::try_new(list, nested_lambda.clone(), captures)
+                .map(IntoArray::into_array)
+        }
+    }
+}
+
+fn template_input(inputs: &[Option<ArrayRef>], slot: usize) -> VortexResult<ArrayRef> {
+    inputs
+        .get(slot)
+        .and_then(Option::as_ref)
+        .cloned()
+        .ok_or_else(|| vortex_error::vortex_err!("template input slot {slot} is not bound"))
+}