byte_length() expression

encodings/fsst/src/compute/byte_length.rs

@@ -0,0 +1,37 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+use vortex_array::ArrayRef;
+use vortex_array::ArrayView;
+use vortex_array::ExecutionCtx;
+use vortex_array::IntoArray;
+use vortex_array::ValidityVTable;
+use vortex_array::arrays::ConstantArray;
+use vortex_array::builtins::ArrayBuiltins;
+use vortex_array::dtype::DType;
+use vortex_array::dtype::PType;
+use vortex_array::scalar::Scalar;
+use vortex_array::scalar_fn::fns::byte_length::ByteLengthKernel;
+use vortex_array::validity::Validity;
+use vortex_error::VortexResult;
+
+use crate::FSST;
+use crate::array::FSSTArrayExt;
+
+impl ByteLengthKernel for FSST {
+    fn byte_length(
+        array: ArrayView<'_, Self>,
+        _ctx: &mut ExecutionCtx,
+    ) -> VortexResult<Option<ArrayRef>> {
+        let nullable = array.dtype().nullability();
+        let dtype = DType::Primitive(PType::U64, nullable);
+        // Uncompressed lengths are non-nullable and may be less than u64 each
+        let lengths = array.uncompressed_lengths().cast(dtype.clone())?;
+        Ok(Some(match FSST::validity(array)? {
+            Validity::NonNullable | Validity::AllValid => lengths,
+            Validity::Array(v) => lengths.mask(v)?,
+            Validity::AllInvalid => {
+                ConstantArray::new(Scalar::null(dtype), lengths.len()).into_array()
+            }
+        }))
+    }
+}

encodings/fsst/src/compute/mod.rs

@@ -1,6 +1,7 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+mod byte_length;
 mod cast;
 mod compare;
 mod filter;

encodings/fsst/src/kernel.rs

@@ -5,6 +5,7 @@ use vortex_array::arrays::dict::TakeExecuteAdaptor;
 use vortex_array::arrays::filter::FilterExecuteAdaptor;
 use vortex_array::kernel::ParentKernelSet;
 use vortex_array::scalar_fn::fns::binary::CompareExecuteAdaptor;
+use vortex_array::scalar_fn::fns::byte_length::ByteLengthExecuteAdaptor;
 use vortex_array::scalar_fn::fns::cast::CastExecuteAdaptor;
 use vortex_array::scalar_fn::fns::like::LikeExecuteAdaptor;
 
@@ -16,6 +17,7 @@ pub(super) const PARENT_KERNELS: ParentKernelSet<FSST> = ParentKernelSet::new(&[
     ParentKernelSet::lift(&FilterExecuteAdaptor(FSST)),
     ParentKernelSet::lift(&TakeExecuteAdaptor(FSST)),
     ParentKernelSet::lift(&LikeExecuteAdaptor(FSST)),
+    ParentKernelSet::lift(&ByteLengthExecuteAdaptor(FSST)),
 ]);
 
 #[cfg(test)]
@@ -27,10 +29,13 @@ mod tests {
     use vortex_array::IntoArray;
     use vortex_array::VortexSessionExecute;
     use vortex_array::arrays::FilterArray;
+    use vortex_array::arrays::PrimitiveArray;
     use vortex_array::arrays::varbin::builder::VarBinBuilder;
     use vortex_array::assert_arrays_eq;
     use vortex_array::dtype::DType;
     use vortex_array::dtype::Nullability;
+    use vortex_array::expr::byte_length;
+    use vortex_array::expr::root;
     use vortex_array::session::ArraySession;
     use vortex_error::VortexResult;
     use vortex_mask::Mask;
@@ -205,4 +210,24 @@ mod tests {
         assert_arrays_eq!(result, fsst_array);
         Ok(())
     }
+
+    #[test]
+    fn test_fsst_byte_length() -> VortexResult<()> {
+        let mut builder = VarBinBuilder::<i32>::with_capacity(3);
+        builder.append_value(b"hello");
+        builder.append_value(b"world!!");
+        builder.append_value("Пуховички"); // 9 characters, 18 bytes
+        builder.append_value(b"");
+
+        let varbin = builder.finish(DType::Utf8(Nullability::NonNullable));
+        let compressor = fsst_train_compressor(&varbin);
+        let len = varbin.len();
+        let dtype = varbin.dtype().clone();
+        let mut ctx = SESSION.create_execution_ctx();
+        let fsst = fsst_compress(varbin, len, &dtype, &compressor, &mut ctx).into_array();
+        let result = fsst.apply(&byte_length(root()))?;
+        let expected = PrimitiveArray::from_iter(vec![5u64, 7, 18, 0]);
+        assert_arrays_eq!(result, expected);
+        Ok(())
+    }
 }

vortex-array/src/expr/exprs.rs

@@ -21,6 +21,7 @@ use crate::scalar_fn::ScalarFnVTableExt;
 use crate::scalar_fn::fns::between::Between;
 use crate::scalar_fn::fns::between::BetweenOptions;
 use crate::scalar_fn::fns::binary::Binary;
+use crate::scalar_fn::fns::byte_length::ByteLength;
 use crate::scalar_fn::fns::case_when::CaseWhen;
 use crate::scalar_fn::fns::case_when::CaseWhenOptions;
 use crate::scalar_fn::fns::cast::Cast;
@@ -718,3 +719,16 @@ pub fn dynamic(
 pub fn list_contains(list: Expression, value: Expression) -> Expression {
     ListContains.new_expr(EmptyOptions, [list, value])
 }
+
+// ---- ByteLength ----
+
+/// Creates an expression that computes the byte length of each element.
+/// This is akin to ANSI SQL OCTET_LENGTH(), or DuckDB's strlen().
+///
+/// ```rust
+/// # use vortex_array::expr::{byte_length, root};
+/// let expr = byte_length(root());
+/// ```
+pub fn byte_length(input: Expression) -> Expression {
+    ByteLength.new_expr(EmptyOptions, [input])
+}

vortex-array/src/scalar_fn/fns/byte_length.rs

@@ -0,0 +1,248 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use num_traits::AsPrimitive;
+use vortex_buffer::Buffer;
+use vortex_error::VortexExpect;
+use vortex_error::VortexResult;
+use vortex_error::vortex_bail;
+use vortex_error::vortex_ensure;
+use vortex_session::VortexSession;
+use vortex_session::registry::CachedId;
+
+use crate::ArrayRef;
+use crate::ExecutionCtx;
+use crate::IntoArray;
+use crate::array::ArrayView;
+use crate::array::VTable;
+use crate::arrays::ConstantArray;
+use crate::arrays::PrimitiveArray;
+use crate::arrays::VarBinViewArray;
+use crate::arrays::scalar_fn::ExactScalarFn;
+use crate::arrays::scalar_fn::ScalarFnArrayView;
+use crate::arrays::varbinview::VarBinViewArrayExt;
+use crate::dtype::DType;
+use crate::dtype::Nullability;
+use crate::dtype::PType;
+use crate::kernel::ExecuteParentKernel;
+use crate::scalar::Scalar;
+use crate::scalar_fn::Arity;
+use crate::scalar_fn::ChildName;
+use crate::scalar_fn::EmptyOptions;
+use crate::scalar_fn::ExecutionArgs;
+use crate::scalar_fn::ScalarFnId;
+use crate::scalar_fn::ScalarFnVTable;
+
+pub trait ByteLengthKernel: VTable {
+    fn byte_length(
+        array: ArrayView<'_, Self>,
+        ctx: &mut ExecutionCtx,
+    ) -> VortexResult<Option<ArrayRef>>;
+}
+
+#[derive(Default, Debug)]
+pub struct ByteLengthExecuteAdaptor<V>(pub V);
+
+impl<V: ByteLengthKernel> ExecuteParentKernel<V> for ByteLengthExecuteAdaptor<V> {
+    type Parent = ExactScalarFn<ByteLength>;
+
+    fn execute_parent(
+        &self,
+        array: ArrayView<'_, V>,
+        _parent: ScalarFnArrayView<'_, ByteLength>,
+        child_idx: usize,
+        ctx: &mut ExecutionCtx,
+    ) -> VortexResult<Option<ArrayRef>> {
+        vortex_ensure!(child_idx == 0);
+        V::byte_length(array, ctx)
+    }
+}
+
+/// Byte length of each element in a Utf8 or Binary array.
+#[derive(Clone)]
+pub struct ByteLength;
+
+impl ScalarFnVTable for ByteLength {
+    type Options = EmptyOptions;
+
+    fn id(&self) -> ScalarFnId {
+        static ID: CachedId = CachedId::new("vortex.byte_length");
+        *ID
+    }
+
+    fn serialize(&self, _instance: &Self::Options) -> VortexResult<Option<Vec<u8>>> {
+        Ok(Some(vec![]))
+    }
+
+    fn deserialize(
+        &self,
+        _metadata: &[u8],
+        _session: &VortexSession,
+    ) -> VortexResult<Self::Options> {
+        Ok(EmptyOptions)
+    }
+
+    fn arity(&self, _options: &Self::Options) -> Arity {
+        Arity::Exact(1)
+    }
+
+    fn child_name(&self, _instance: &Self::Options, child_idx: usize) -> ChildName {
+        match child_idx {
+            0 => ChildName::from("input"),
+            _ => unreachable!("Invalid child index {child_idx} for byte_length()"),
+        }
+    }
+
+    fn return_dtype(&self, _options: &Self::Options, arg_dtypes: &[DType]) -> VortexResult<DType> {
+        match &arg_dtypes[0] {
+            DType::Utf8(nullable) | DType::Binary(nullable) => {
+                Ok(DType::Primitive(PType::U64, *nullable))
+            }
+            other => vortex_bail!("byte_length() requires Utf8 or Binary, got {other}"),
+        }
+    }
+
+    fn execute(
+        &self,
+        _options: &Self::Options,
+        args: &dyn ExecutionArgs,
+        ctx: &mut ExecutionCtx,
+    ) -> VortexResult<ArrayRef> {
+        let input = args.get(0)?;
+        let nullability = input.dtype().nullability();
+
+        if let Some(scalar) = input.as_constant() {
+            let len_scalar = scalar_byte_length(&scalar, nullability)?;
+            return Ok(ConstantArray::new(len_scalar, args.row_count()).into_array());
+        }
+
+        match input.dtype() {
+            DType::Utf8(_) | DType::Binary(_) => byte_length(&input, nullability, ctx),
+            other => vortex_bail!("byte_length() requires Utf8 or Binary, got {other}"),
+        }
+    }
+
+    fn is_null_sensitive(&self, _options: &Self::Options) -> bool {
+        false
+    }
+
+    fn is_fallible(&self, _options: &Self::Options) -> bool {
+        false
+    }
+}
+
+fn scalar_byte_length(scalar: &Scalar, nullability: Nullability) -> VortexResult<Scalar> {
+    if scalar.is_null() {
+        let dtype = DType::Primitive(PType::U64, Nullability::Nullable);
+        return Ok(Scalar::null(dtype));
+    }
+    let len = match scalar.dtype() {
+        DType::Utf8(_) => scalar
+            .as_utf8()
+            .value()
+            .vortex_expect("null utf-8 scalar")
+            .len(),
+        DType::Binary(_) => scalar
+            .as_binary()
+            .value()
+            .vortex_expect("null binary scalar")
+            .len(),
+        other => vortex_bail!("byte_length() requires Utf8 or Binary, got {other}"),
+    };
+    let len: u64 = len.as_();
+    Ok(Scalar::primitive(len, nullability))
+}
+
+pub(crate) fn byte_length(
+    array: &ArrayRef,
+    nullability: Nullability,
+    ctx: &mut ExecutionCtx,
+) -> VortexResult<ArrayRef> {
+    let array = array.clone().execute::<VarBinViewArray>(ctx)?;
+    let validity = array.varbinview_validity();
+    let lengths: Buffer<u64> = array.views().iter().map(|v| v.len() as u64).collect();
+    Ok(PrimitiveArray::new(lengths, validity.union_nullability(nullability)).into_array())
+}
+
+#[cfg(test)]
+mod tests {
+    use rstest::rstest;
+    use vortex_error::VortexResult;
+
+    use crate::ArrayRef;
+    use crate::IntoArray;
+    use crate::LEGACY_SESSION;
+    use crate::VortexSessionExecute;
+    use crate::arrays::ConstantArray;
+    use crate::arrays::PrimitiveArray;
+    use crate::arrays::VarBinArray;
+    use crate::arrays::VarBinViewArray;
+    use crate::assert_arrays_eq;
+    use crate::dtype::DType;
+    use crate::dtype::Nullability;
+    use crate::expr::byte_length;
+    use crate::expr::root;
+    use crate::scalar::Scalar;
+
+    #[rstest]
+    #[case(VarBinArray::from_strs(vec!["hello", "world", ""]).into_array(), vec![5u64, 5, 0])]
+    #[case(VarBinArray::from_bytes(vec![b"ab".as_ref(), b"cde"]).into_array(), vec![2u64, 3])]
+    #[case(VarBinArray::from_strs(vec!["Пуховички"]).into_array(), vec![18u64])]
+    #[case(VarBinArray::from_bytes(vec!["Пуховички".as_ref()]).into_array(), vec![18u64])]
+    fn test_bytes_byte_length(
+        #[case] array: ArrayRef,
+        #[case] expected_lens: Vec<u64>,
+    ) -> VortexResult<()> {
+        let result = array.apply(&byte_length(root()))?;
+        let expected = PrimitiveArray::from_iter(expected_lens);
+        assert_arrays_eq!(result, expected);
+        Ok(())
+    }
+
+    #[test]
+    fn test_varbinview_byte_length() -> VortexResult<()> {
+        let array = VarBinViewArray::from_iter_str(["short", "a longer string here"]).into_array();
+        let result = array.apply(&byte_length(root()))?;
+        let expected = PrimitiveArray::from_iter(vec![5u64, 20]);
+        assert_arrays_eq!(result, expected);
+        Ok(())
+    }
+
+    #[test]
+    fn test_nullable_string_byte_length() -> VortexResult<()> {
+        let array = VarBinArray::from_nullable_strs(vec![Some("hello"), None, Some("Пуховички")])
+            .into_array();
+        let result = array.apply(&byte_length(root()))?;
+
+        let mut ctx = LEGACY_SESSION.create_execution_ctx();
+        assert!(result.is_valid(0, &mut ctx)?);
+        assert!(!result.is_valid(1, &mut ctx)?);
+        assert!(result.is_valid(2, &mut ctx)?);
+        assert_eq!(
+            result.execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx())?,
+            Scalar::primitive(5u64, Nullability::Nullable),
+        );
+        assert_eq!(
+            result.execute_scalar(2, &mut LEGACY_SESSION.create_execution_ctx())?,
+            Scalar::primitive(18u64, Nullability::Nullable),
+        );
+        Ok(())
+    }
+
+    #[test]
+    fn test_null_scalar_byte_length() -> VortexResult<()> {
+        let null_scalar = Scalar::null(DType::Utf8(Nullability::Nullable));
+        let array = ConstantArray::new(null_scalar, 2).into_array();
+        let result = array.apply(&byte_length(root()))?;
+        let mut ctx = LEGACY_SESSION.create_execution_ctx();
+        assert!(!result.is_valid(0, &mut ctx)?);
+        assert!(!result.is_valid(1, &mut ctx)?);
+        Ok(())
+    }
+
+    #[test]
+    fn test_display() {
+        let expr = byte_length(root());
+        assert_eq!(expr.to_string(), "vortex.byte_length($)");
+    }
+}

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

@@ -3,6 +3,7 @@
 
 pub mod between;
 pub mod binary;
+pub mod byte_length;
 pub mod case_when;
 pub mod cast;
 pub mod dynamic;