Merge branch 'develop' into break/operators

Author: gatesn

vortex-array/src/arrays/varbinview/compute/take.rs

@@ -1,12 +1,15 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+use std::iter;
 use std::ops::Deref;
 
 use num_traits::AsPrimitive;
 use vortex_buffer::Buffer;
 use vortex_dtype::match_each_integer_ptype;
 use vortex_error::VortexResult;
+use vortex_mask::AllOr;
+use vortex_mask::Mask;
 use vortex_vector::binaryview::BinaryView;
 
 use crate::Array;
@@ -23,16 +26,16 @@ use crate::vtable::ValidityHelper;
 /// Take involves creating a new array that references the old array, just with the given set of views.
 impl TakeKernel for VarBinViewVTable {
     fn take(&self, array: &VarBinViewArray, indices: &dyn Array) -> VortexResult<ArrayRef> {
-        // Compute the new validity
-
-        // This is valid since all elements (of all arrays) even null values must be inside
-        // min-max valid range.
+        // Compute the new validity.
         let validity = array.validity().take(indices)?;
         let indices = indices.to_primitive();
 
         let views_buffer = match_each_integer_ptype!(indices.ptype(), |I| {
-            // This is valid since all elements even null values are inside the min-max valid range.
-            take_views(array.views(), indices.as_slice::<I>())
+            take_views(
+                array.views(),
+                indices.as_slice::<I>(),
+                &indices.validity_mask(),
+            )
         });
 
         // SAFETY: taking all components at same indices maintains invariants
@@ -55,15 +58,36 @@ register_kernel!(TakeKernelAdapter(VarBinViewVTable).lift());
 fn take_views<I: AsPrimitive<usize>>(
     views: &Buffer<BinaryView>,
     indices: &[I],
+    mask: &Mask,
 ) -> Buffer<BinaryView> {
     // NOTE(ngates): this deref is not actually trivial, so we run it once.
     let views_ref = views.deref();
-    Buffer::<BinaryView>::from_trusted_len_iter(indices.iter().map(|i| views_ref[i.as_()]))
+    // We do not use iter_bools directly, since the resulting dyn iterator cannot
+    // implement TrustedLen.
+    match mask.bit_buffer() {
+        AllOr::All => {
+            Buffer::<BinaryView>::from_trusted_len_iter(indices.iter().map(|i| views_ref[i.as_()]))
+        }
+        AllOr::None => Buffer::<BinaryView>::from_trusted_len_iter(iter::repeat_n(
+            BinaryView::default(),
+            indices.len(),
+        )),
+        AllOr::Some(buffer) => Buffer::<BinaryView>::from_trusted_len_iter(
+            buffer.iter().zip(indices.iter()).map(|(valid, idx)| {
+                if valid {
+                    views_ref[idx.as_()]
+                } else {
+                    BinaryView::default()
+                }
+            }),
+        ),
+    }
 }
 
 #[cfg(test)]
 mod tests {
     use rstest::rstest;
+    use vortex_buffer::BitBuffer;
     use vortex_buffer::buffer;
     use vortex_dtype::DType;
     use vortex_dtype::Nullability::NonNullable;
@@ -76,6 +100,7 @@ mod tests {
     use crate::canonical::ToCanonical;
     use crate::compute::conformance::take::test_take_conformance;
     use crate::compute::take;
+    use crate::validity::Validity;
 
     #[test]
     fn take_nullable() {
@@ -103,11 +128,13 @@ mod tests {
     fn take_nullable_indices() {
         let arr = VarBinViewArray::from_iter(["one", "two"].map(Some), DType::Utf8(NonNullable));
 
-        let taken = take(
-            arr.as_ref(),
-            PrimitiveArray::from_option_iter(vec![Some(1), None]).as_ref(),
-        )
-        .unwrap();
+        let indices = PrimitiveArray::new(
+            // Verify that garbage values at NULL indices are ignored.
+            buffer![1u64, 999],
+            Validity::from(BitBuffer::from(vec![true, false])),
+        );
+
+        let taken = take(arr.as_ref(), indices.as_ref()).unwrap();
 
         assert!(taken.dtype().is_nullable());
         assert_eq!(

vortex-buffer/src/trusted_len.rs

@@ -161,6 +161,8 @@ where
 {
 }
 
+unsafe impl<T: Clone> TrustedLen for std::iter::RepeatN<T> {}
+
 // Arrow bit iterators
 unsafe impl<'a> TrustedLen for crate::bit::BitIterator<'a> {}
 unsafe impl<'a> TrustedLen for crate::bit::BitChunkIterator<'a> {}

vortex-python/python/vortex/dataset.py

@@ -791,6 +791,22 @@ def __init__(
     def schema(self):
         return self._dataset.schema
 
+    @property
+    @override
+    def dataset_schema(self) -> pyarrow.Schema:
+        return self._dataset.schema
+
+    @property
+    @override
+    def projected_schema(self) -> pyarrow.Schema:
+        if self._columns:
+            fields: list[pa.Field[pa.DataType]] = [
+                self._dataset.schema.field(c)  # pyright: ignore[reportUnknownMemberType]
+                for c in self._columns
+            ]
+            return pyarrow.schema(fields)
+        return self._dataset.schema
+
     @override
     def count_rows(self):
         return self._dataset.count_rows(

vortex-python/test/test_scan.py

@@ -64,3 +64,13 @@ def test_scan_with_cast(vxfile: vx.VortexFile):
         [{"index": 1, "string": pa.scalar("1", pa.string_view()), "bool": False, "float": math.sqrt(1)}]
     )
     assert str(actual.to_arrow_array()) == str(expected)
+
+
+def test_scanner_property_projected(vxfile: vx.VortexFile):
+    assert vxfile.to_dataset().scanner(columns=["bool"]).projected_schema == pa.schema([("bool", pa.bool_())])
+
+
+def test_scanner_property_dataset_schema(vxfile: vx.VortexFile):
+    assert vxfile.to_dataset().scanner().dataset_schema == pa.schema(
+        [("bool", pa.bool_()), ("float", pa.float64()), ("index", pa.int64()), ("string", pa.string_view())]
+    )

vortex-tui/src/main.rs

@@ -23,6 +23,7 @@ use vortex::session::VortexSession;
 use crate::inspect::InspectArgs;
 
 #[derive(clap::Parser)]
+#[command(version)]
 struct Cli {
     #[clap(subcommand)]
     command: Commands,