feat(array): add Narrow integer encoding

encodings/fastlanes/tests/narrow.rs

@@ -0,0 +1,69 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! Narrow operations on compressed integer children.
+//!
+//! BitPacked children retain their stored width through selection and comparison. Canonical
+//! execution widens them to Narrow's logical dtype.
+
+#![cfg(test)]
+
+use std::sync::LazyLock;
+
+use vortex_array::IntoArray;
+use vortex_array::VortexSessionExecute;
+use vortex_array::aggregate_fn::fns::sum::sum;
+use vortex_array::arrays::BoolArray;
+use vortex_array::arrays::ConstantArray;
+use vortex_array::arrays::Narrow;
+use vortex_array::arrays::NarrowArray;
+use vortex_array::arrays::PrimitiveArray;
+use vortex_array::arrays::narrow::NarrowArraySlotsExt;
+use vortex_array::assert_arrays_eq;
+use vortex_array::builtins::ArrayBuiltins;
+use vortex_array::dtype::PType;
+use vortex_array::scalar::Scalar;
+use vortex_array::scalar_fn::fns::operators::Operator;
+use vortex_buffer::buffer;
+use vortex_error::VortexResult;
+use vortex_fastlanes::BitPacked;
+use vortex_mask::Mask;
+use vortex_session::VortexSession;
+
+static SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
+    let session = vortex_array::array_session();
+    vortex_fastlanes::initialize(&session);
+    session
+});
+
+#[test]
+fn test_narrow_bitpacked_child() -> VortexResult<()> {
+    let mut ctx = SESSION.create_execution_ctx();
+    let packed = BitPacked::encode(&buffer![0u8, 1, 3, 7].into_array(), 3, &mut ctx)?;
+    let array = NarrowArray::try_new(packed.into_array(), PType::U64.into())?.into_array();
+    let selected = array.filter(Mask::from_iter([true, false, true, true]))?;
+
+    assert!(selected.is::<Narrow>());
+    assert_eq!(selected.as_::<Narrow>().values().dtype(), &PType::U8.into());
+    assert_arrays_eq!(selected, buffer![0u64, 3, 7].into_array(), &mut ctx);
+    assert_arrays_eq!(
+        array.take(buffer![3u32, 1].into_array())?,
+        buffer![7u64, 1].into_array(),
+        &mut ctx
+    );
+    assert_arrays_eq!(
+        array.binary(
+            ConstantArray::new(Scalar::from(3u64), 4).into_array(),
+            Operator::Lt
+        )?,
+        BoolArray::from_iter([true, true, false, false]),
+        &mut ctx
+    );
+    assert_eq!(sum(&array, &mut ctx)?, Scalar::from(11u64).into_nullable());
+    assert_eq!(
+        array.execute::<PrimitiveArray>(&mut ctx)?.as_slice::<u64>(),
+        &[0, 1, 3, 7]
+    );
+
+    Ok(())
+}

vortex-array/src/aggregate_fn/session.rs

@@ -47,6 +47,7 @@ use crate::arrays::chunked::compute::aggregate::ChunkedArrayAggregate;
 use crate::arrays::dict::compute::is_constant::DictIsConstantKernel;
 use crate::arrays::dict::compute::is_sorted::DictIsSortedKernel;
 use crate::arrays::dict::compute::min_max::DictMinMaxKernel;
+use crate::arrays::narrow;
 use crate::dtype::DType;
 
 /// Session state for aggregate functions and encoding-specific aggregate kernels.
@@ -122,6 +123,8 @@ impl Default for AggregateFnSession {
         this.register_aggregate_kernel(Dict.id(), Some(IsConstant.id()), &DictIsConstantKernel);
         this.register_aggregate_kernel(Dict.id(), Some(IsSorted.id()), &DictIsSortedKernel);
 
+        narrow::register_aggregate_kernels(&this);
+
         // Register the built-in grouped aggregate kernels.
         this.register_grouped_kernel(Count.id(), &CountGroupedKernel);
         this.register_grouped_encoding_kernel(

vortex-array/src/arrays/mod.rs

@@ -86,6 +86,10 @@ pub mod masked;
 pub use masked::Masked;
 pub use masked::MaskedArray;
 
+pub mod narrow;
+pub use narrow::Narrow;
+pub use narrow::NarrowArray;
+
 pub mod null;
 pub use null::Null;
 pub use null::NullArray;
@@ -146,6 +150,7 @@ pub(crate) fn initialize(session: &VortexSession) {
     list::initialize(session);
     listview::initialize(session);
     map::initialize(session);
+    narrow::initialize(session);
     patched::initialize(session);
     primitive::initialize(session);
     struct_::initialize(session);

vortex-array/src/arrays/narrow/aggregates.rs

@@ -0,0 +1,90 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! Aggregate Narrow values without widening the input buffer.
+//!
+//! Supported aggregates preserve integer ordering or use wide sum states. Partial states are cast
+//! to the logical dtype so later batches observe the same boundaries as canonical integers.
+
+use std::sync::LazyLock;
+
+use vortex_error::VortexResult;
+
+use super::Narrow;
+use super::NarrowArraySlotsExt;
+use crate::ArrayRef;
+use crate::ArrayVTable;
+use crate::ExecutionCtx;
+use crate::aggregate_fn::AggregateFnId;
+use crate::aggregate_fn::AggregateFnRef;
+use crate::aggregate_fn::AggregateFnVTable;
+use crate::aggregate_fn::fns::all_non_null::AllNonNull;
+use crate::aggregate_fn::fns::all_null::AllNull;
+use crate::aggregate_fn::fns::count::Count;
+use crate::aggregate_fn::fns::first::First;
+use crate::aggregate_fn::fns::is_constant::IsConstant;
+use crate::aggregate_fn::fns::is_sorted::IsSorted;
+use crate::aggregate_fn::fns::last::Last;
+use crate::aggregate_fn::fns::max::Max;
+use crate::aggregate_fn::fns::min::Min;
+use crate::aggregate_fn::fns::min_max::MinMax;
+use crate::aggregate_fn::fns::null_count::NullCount;
+use crate::aggregate_fn::fns::sum::Sum;
+use crate::aggregate_fn::fns::sum_v2::SumV2;
+use crate::aggregate_fn::kernels::DynAggregateKernel;
+use crate::aggregate_fn::session::AggregateFnSession;
+use crate::scalar::Scalar;
+
+pub(crate) fn register_aggregate_kernels(session: &AggregateFnSession) {
+    // One fallback avoids cloning the registry for each supported aggregate.
+    session.register_aggregate_kernel(Narrow.id(), None::<AggregateFnId>, &NarrowAggregateKernel);
+}
+
+static SUPPORTED_AGGREGATES: LazyLock<[AggregateFnId; 13]> = LazyLock::new(|| {
+    [
+        MinMax.id(),
+        Min.id(),
+        Max.id(),
+        Sum.id(),
+        SumV2.id(),
+        IsConstant.id(),
+        IsSorted.id(),
+        First.id(),
+        Last.id(),
+        Count.id(),
+        NullCount.id(),
+        AllNull.id(),
+        AllNonNull.id(),
+    ]
+});
+
+#[derive(Debug)]
+struct NarrowAggregateKernel;
+
+impl DynAggregateKernel for NarrowAggregateKernel {
+    fn aggregate(
+        &self,
+        aggregate_fn: &AggregateFnRef,
+        batch: &ArrayRef,
+        ctx: &mut ExecutionCtx,
+    ) -> VortexResult<Option<Scalar>> {
+        if !SUPPORTED_AGGREGATES.contains(&aggregate_fn.id()) {
+            return Ok(None);
+        }
+
+        let Some(array) = batch.as_opt::<Narrow>() else {
+            return Ok(None);
+        };
+
+        // These aggregates preserve integer ordering and accumulate sums at i64/u64 width.
+        // Cast the partial state too, since it can contain extrema or boundary values.
+        let mut accumulator = aggregate_fn.accumulator(array.values().dtype())?;
+        accumulator.accumulate(array.values(), ctx)?;
+        let partial = accumulator.partial_scalar()?;
+        let Some(dtype) = aggregate_fn.state_dtype(batch.dtype()) else {
+            return Ok(None);
+        };
+
+        Ok(Some(partial.cast(&dtype)?))
+    }
+}

vortex-array/src/arrays/narrow/compare.rs

@@ -0,0 +1,122 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! Compare Narrow integers at a sufficient stored width.
+//!
+//! Two children use their wider stored dtype. Constants outside the stored range yield a known
+//! comparison result while preserving the array validity.
+
+use vortex_buffer::BitBuffer;
+use vortex_error::VortexResult;
+use vortex_session::VortexSession;
+
+use super::Narrow;
+use super::NarrowArraySlotsExt;
+use crate::ArrayRef;
+use crate::ArrayView;
+use crate::ExecutionCtx;
+use crate::IntoArray;
+use crate::arrays::BoolArray;
+use crate::arrays::Constant;
+use crate::arrays::ConstantArray;
+use crate::arrays::Primitive;
+use crate::builtins::ArrayBuiltins;
+use crate::dtype::DType;
+use crate::optimizer::kernels::ArrayKernelsExt;
+use crate::scalar_fn::ScalarFnVTable;
+use crate::scalar_fn::fns::binary::Binary;
+use crate::scalar_fn::fns::binary::CompareExecuteAdaptor;
+use crate::scalar_fn::fns::binary::CompareKernel;
+use crate::scalar_fn::fns::binary::execute_compare;
+use crate::scalar_fn::fns::operators::CompareOperator;
+
+pub(super) fn initialize(session: &VortexSession) {
+    session.kernels().register_execute_parent_kernel(
+        Binary.id(),
+        Narrow,
+        CompareExecuteAdaptor(Narrow),
+    );
+}
+
+impl CompareKernel for Narrow {
+    fn compare(
+        lhs: ArrayView<'_, Self>,
+        rhs: &ArrayRef,
+        operator: CompareOperator,
+        ctx: &mut ExecutionCtx,
+    ) -> VortexResult<Option<ArrayRef>> {
+        if !lhs.dtype().eq_ignore_nullability(rhs.dtype()) {
+            return Ok(None);
+        }
+
+        if let Some(rhs) = rhs.as_opt::<Narrow>() {
+            let ptype = if lhs.values().dtype().as_ptype().byte_width()
+                >= rhs.values().dtype().as_ptype().byte_width()
+            {
+                lhs.values().dtype().as_ptype()
+            } else {
+                rhs.values().dtype().as_ptype()
+            };
+            let lhs_dtype = DType::Primitive(ptype, lhs.dtype().nullability());
+            let rhs_dtype = DType::Primitive(ptype, rhs.dtype().nullability());
+
+            return compare_values(
+                &lhs.values().cast(lhs_dtype)?,
+                &rhs.values().cast(rhs_dtype)?,
+                operator,
+                ctx,
+            )
+            .map(Some);
+        }
+
+        let Some(constant) = rhs.as_constant() else {
+            return Ok(None);
+        };
+
+        let nullability = lhs.dtype().nullability() | rhs.dtype().nullability();
+        let storage_dtype = lhs.values().dtype().with_nullability(nullability);
+        if let Ok(value) = constant.cast(&storage_dtype) {
+            return compare_values(
+                lhs.values(),
+                &ConstantArray::new(value, lhs.len()).into_array(),
+                operator,
+                ctx,
+            )
+            .map(Some);
+        }
+
+        // An integer outside the child's type range is ordered against every stored value.
+        let below = lhs.dtype().is_signed_int() && i64::try_from(&constant)? < 0;
+        let result = match operator {
+            CompareOperator::Eq => false,
+            CompareOperator::NotEq => true,
+            CompareOperator::Lt | CompareOperator::Lte => !below,
+            CompareOperator::Gt | CompareOperator::Gte => below,
+        };
+        let validity = lhs
+            .values()
+            .validity()?
+            .cast_nullability(nullability, lhs.len(), ctx)?;
+
+        Ok(Some(
+            BoolArray::try_new(
+                BitBuffer::full_in(result, lhs.len(), ctx.allocator().clone()),
+                validity,
+            )?
+            .into_array(),
+        ))
+    }
+}
+
+fn compare_values(
+    lhs: &ArrayRef,
+    rhs: &ArrayRef,
+    operator: CompareOperator,
+    ctx: &mut ExecutionCtx,
+) -> VortexResult<ArrayRef> {
+    if lhs.is::<Primitive>() && (rhs.is::<Primitive>() || rhs.is::<Constant>()) {
+        execute_compare(lhs, rhs, operator, ctx)
+    } else {
+        lhs.binary(rhs.clone(), operator.into())
+    }
+}

vortex-array/src/arrays/narrow/encoding.rs

@@ -0,0 +1,69 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! Select the stored width of an integer array.
+//!
+//! Bounds exclude null payloads. Narrowing changes the physical representation while preserving
+//! the logical dtype and validity.
+
+use vortex_error::VortexResult;
+use vortex_error::vortex_ensure;
+
+use super::NarrowArray;
+use crate::ArrayRef;
+use crate::ExecutionCtx;
+use crate::IntoArray;
+use crate::aggregate_fn::NumericalAggregateOpts;
+use crate::aggregate_fn::fns::min_max::min_max;
+use crate::arrays::PrimitiveArray;
+use crate::builtins::ArrayBuiltins;
+use crate::dtype::DType;
+use crate::dtype::PType;
+
+impl NarrowArray {
+    /// Stores an integer array at the smallest fitting width while retaining its logical dtype.
+    ///
+    /// Uses the non-null bounds and preserves signedness and validity. Empty and all-null arrays
+    /// use the smallest width of the same signedness. Returns the input unchanged if no narrower
+    /// type fits. Floating-point arrays are rejected.
+    pub fn encode(array: PrimitiveArray, ctx: &mut ExecutionCtx) -> VortexResult<ArrayRef> {
+        let ptype = array.ptype();
+        vortex_ensure!(
+            ptype.is_int(),
+            "Narrow requires integer values, got {ptype}"
+        );
+
+        if ptype.byte_width() == 1 {
+            return Ok(array.into_array());
+        }
+
+        let bounds = min_max(array.as_ref(), ctx, NumericalAggregateOpts::default())?;
+        let candidates = if ptype.is_signed_int() {
+            [PType::I8, PType::I16, PType::I32]
+        } else {
+            [PType::U8, PType::U16, PType::U32]
+        };
+
+        for candidate in candidates {
+            if candidate.byte_width() >= ptype.byte_width() {
+                break;
+            }
+
+            let storage_dtype = DType::Primitive(candidate, array.dtype().nullability());
+            if bounds.as_ref().is_some_and(|bounds| {
+                bounds.min.cast(&storage_dtype).is_err() || bounds.max.cast(&storage_dtype).is_err()
+            }) {
+                continue;
+            }
+
+            let values = array
+                .as_ref()
+                .cast(storage_dtype)?
+                .execute::<PrimitiveArray>(ctx)?;
+
+            return Ok(Self::try_new(values.into_array(), array.dtype().clone())?.into_array());
+        }
+
+        Ok(array.into_array())
+    }
+}