Merge branch 'develop' into ngates/stat-expression

Author: gatesn

.github/release-drafter.yml

@@ -15,6 +15,9 @@ categories:
     collapse-after: 8
     labels:
       - "fix"
+  - title: "📖 Documentation"
+    labels:
+      - "documentation"
   - title: "🧰 Maintenance"
     labels:
       - "chore"

.github/workflows/labels.yml

@@ -28,6 +28,7 @@ jobs:
         run: |
           REQUIRED_LABELS=(
             "chore"
+            "documentation"
             "feature"
             "fix"
             "performance"

encodings/zstd/src/array.rs

@@ -454,10 +454,32 @@ impl ZstdArray {
 
         let decompressed = decompressed.freeze();
         // Last, we slice the exact values requested out of the decompressed data.
-        let slice_validity = self
+        let mut slice_validity = self
             .unsliced_validity
             .slice(self.slice_start..self.slice_stop);
 
+        // NOTE: this block handles setting the output type when the validity and DType disagree.
+        //
+        // ZSTD is a compact block compressor, meaning that null values are not stored inline in
+        // the data frames. A ZSTD Array that was initialized must always hold onto its full
+        // validity bitmap, even if sliced to only include non-null values.
+        //
+        // We ensure that the validity of the decompressed array ALWAYS matches the validity
+        // implied by the DType.
+        if !self.dtype().is_nullable() && slice_validity != Validity::NonNullable {
+            assert!(
+                slice_validity.all_valid(slice_n_rows),
+                "ZSTD array expects to be non-nullable but there are nulls after decompression"
+            );
+
+            slice_validity = Validity::NonNullable;
+        } else if self.dtype.is_nullable() && slice_validity == Validity::NonNullable {
+            slice_validity = Validity::AllValid;
+        }
+        //
+        // END OF IMPORTANT BLOCK
+        //
+
         match &self.dtype {
             DType::Primitive(..) => {
                 let slice_values_buffer = decompressed.slice(
@@ -531,6 +553,21 @@ impl ZstdArray {
     }
 
     pub(crate) fn _slice(&self, start: usize, stop: usize) -> ZstdArray {
+        let new_start = self.slice_start + start;
+        let new_stop = self.slice_start + stop;
+
+        assert!(
+            new_start <= self.slice_stop,
+            "new slice start {new_start} exceeds end {}",
+            self.slice_stop
+        );
+
+        assert!(
+            new_stop <= self.slice_stop,
+            "new slice stop {new_stop} exceeds end {}",
+            self.slice_stop
+        );
+
         ZstdArray {
             slice_start: self.slice_start + start,
             slice_stop: self.slice_start + stop,

encodings/zstd/src/compute/cast.rs

@@ -2,46 +2,67 @@
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
 use vortex_array::compute::{CastKernel, CastKernelAdapter};
-use vortex_array::{ArrayRef, IntoArray, register_kernel};
-use vortex_dtype::DType;
+use vortex_array::{ArrayRef, register_kernel};
+use vortex_dtype::{DType, Nullability};
 use vortex_error::VortexResult;
 
 use crate::{ZstdArray, ZstdVTable};
 
 impl CastKernel for ZstdVTable {
     fn cast(&self, array: &ZstdArray, dtype: &DType) -> VortexResult<Option<ArrayRef>> {
-        if !dtype.is_nullable() || !array.all_valid() {
-            // We cannot cast to non-nullable since the validity containing nulls is used to decode
-            // the ZSTD array, this would require rewriting tables.
+        if !dtype.eq_ignore_nullability(array.dtype()) {
+            // Type changes can't be handled in ZSTD, need to decode and tweak.
+            // TODO(aduffy): handle trivial conversions like Binary -> UTF8, integer widening, etc.
             return Ok(None);
         }
-        // ZstdArray is a general-purpose compression encoding using Zstandard compression.
-        // It can handle nullability changes without decompression by updating the validity
-        // bitmap, but type changes require decompression since the compressed data is
-        // type-specific and Zstd operates on raw bytes.
-        if array.dtype().eq_ignore_nullability(dtype) {
-            // Create a new validity with the target nullability
-            let new_validity = array
-                .unsliced_validity
-                .clone()
-                .cast_nullability(dtype.nullability(), array.len())?;
-
-            return Ok(Some(
+
+        let src_nullability = array.dtype().nullability();
+        let target_nullability = dtype.nullability();
+
+        match (src_nullability, target_nullability) {
+            // Same type case. This should be handled in the layer above but for
+            // completeness of the match arms we also handle it here.
+            (Nullability::Nullable, Nullability::Nullable)
+            | (Nullability::NonNullable, Nullability::NonNullable) => Ok(Some(array.to_array())),
+            (Nullability::NonNullable, Nullability::Nullable) => Ok(Some(
+                // nonnull => null, trivial cast by altering the validity
                 ZstdArray::new(
                     array.dictionary.clone(),
                     array.frames.clone(),
                     dtype.clone(),
                     array.metadata.clone(),
                     array.unsliced_n_rows(),
-                    new_validity,
+                    array.unsliced_validity.clone(),
                 )
-                ._slice(array.slice_start(), array.slice_stop())
-                .into_array(),
-            ));
+                .slice(array.slice_start()..array.slice_stop()),
+            )),
+            (Nullability::Nullable, Nullability::NonNullable) => {
+                // null => non-null works if there are no nulls in the sliced range
+                let sliced_len = array.slice_stop() - array.slice_start();
+                let has_nulls = !array
+                    .unsliced_validity
+                    .slice(array.slice_start()..array.slice_stop())
+                    .all_valid(sliced_len);
+
+                // We don't attempt to handle casting when there are nulls.
+                if has_nulls {
+                    return Ok(None);
+                }
+
+                // If there are no nulls, the cast is trivial
+                Ok(Some(
+                    ZstdArray::new(
+                        array.dictionary.clone(),
+                        array.frames.clone(),
+                        dtype.clone(),
+                        array.metadata.clone(),
+                        array.unsliced_n_rows(),
+                        array.unsliced_validity.clone(),
+                    )
+                    .slice(array.slice_start()..array.slice_stop()),
+                ))
+            }
         }
-
-        // For other casts (e.g., type changes), decode to canonical and let the underlying array handle it
-        Ok(None)
     }
 }
 

vortex-buffer/src/bit/buf.rs

@@ -8,7 +8,7 @@ use crate::bit::{
     BitChunks, BitIndexIterator, BitIterator, BitSliceIterator, UnalignedBitChunk,
     get_bit_unchecked,
 };
-use crate::{Alignment, BitBufferMut, Buffer, BufferMut, ByteBuffer, buffer};
+use crate::{Alignment, BitBufferMut, Buffer, ByteBuffer, buffer};
 
 /// An immutable bitset stored as a packed byte buffer.
 #[derive(Debug, Clone, Eq)]
@@ -114,37 +114,9 @@ impl BitBuffer {
         }
     }
 
-    /// Invokes `f` with indexes `0..len` collecting the boolean results into a new `BitBuffer`
-    pub fn collect_bool<F: FnMut(usize) -> bool>(len: usize, mut f: F) -> Self {
-        let mut buffer = BufferMut::with_capacity(len.div_ceil(64) * 8);
-
-        let chunks = len / 64;
-        let remainder = len % 64;
-        for chunk in 0..chunks {
-            let mut packed = 0;
-            for bit_idx in 0..64 {
-                let i = bit_idx + chunk * 64;
-                packed |= (f(i) as u64) << bit_idx;
-            }
-
-            // SAFETY: Already allocated sufficient capacity
-            unsafe { buffer.push_unchecked(packed) }
-        }
-
-        if remainder != 0 {
-            let mut packed = 0;
-            for bit_idx in 0..remainder {
-                let i = bit_idx + chunks * 64;
-                packed |= (f(i) as u64) << bit_idx;
-            }
-
-            // SAFETY: Already allocated sufficient capacity
-            unsafe { buffer.push_unchecked(packed) }
-        }
-
-        buffer.truncate(len.div_ceil(8));
-
-        Self::new(buffer.freeze().into_byte_buffer(), len)
+    /// Invokes `f` with indexes `0..len` collecting the boolean results into a new [`BitBuffer`].
+    pub fn collect_bool<F: FnMut(usize) -> bool>(len: usize, f: F) -> Self {
+        BitBufferMut::collect_bool(len, f).freeze()
     }
 
     /// Get the logical length of this `BoolBuffer`.

vortex-buffer/src/bit/buf_mut.rs

@@ -117,6 +117,43 @@ impl BitBufferMut {
         }
     }
 
+    /// Invokes `f` with indexes `0..len` collecting the boolean results into a new `BitBufferMut`
+    pub fn collect_bool<F: FnMut(usize) -> bool>(len: usize, mut f: F) -> Self {
+        let mut buffer = BufferMut::with_capacity(len.div_ceil(64) * 8);
+
+        let chunks = len / 64;
+        let remainder = len % 64;
+        for chunk in 0..chunks {
+            let mut packed = 0;
+            for bit_idx in 0..64 {
+                let i = bit_idx + chunk * 64;
+                packed |= (f(i) as u64) << bit_idx;
+            }
+
+            // SAFETY: Already allocated sufficient capacity
+            unsafe { buffer.push_unchecked(packed) }
+        }
+
+        if remainder != 0 {
+            let mut packed = 0;
+            for bit_idx in 0..remainder {
+                let i = bit_idx + chunks * 64;
+                packed |= (f(i) as u64) << bit_idx;
+            }
+
+            // SAFETY: Already allocated sufficient capacity
+            unsafe { buffer.push_unchecked(packed) }
+        }
+
+        buffer.truncate(len.div_ceil(8));
+
+        Self {
+            buffer: buffer.into_byte_buffer(),
+            offset: 0,
+            len,
+        }
+    }
+
     /// Return the underlying byte buffer.
     pub fn inner(&self) -> &ByteBufferMut {
         &self.buffer

vortex-buffer/src/buffer_mut.rs

@@ -110,6 +110,16 @@ impl<T> BufferMut<T> {
         buffer
     }
 
+    /// Return the [`ByteBufferMut`] for this [`BufferMut<T>`].
+    pub fn into_byte_buffer(self) -> ByteBufferMut {
+        ByteBufferMut {
+            bytes: self.bytes,
+            length: self.length * size_of::<T>(),
+            alignment: self.alignment,
+            _marker: Default::default(),
+        }
+    }
+
     /// Get the alignment of the buffer.
     #[inline(always)]
     pub fn alignment(&self) -> Alignment {

vortex-compute/src/filter/bitbuffer.rs

@@ -2,14 +2,10 @@
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
 use vortex_buffer::{BitBuffer, BitBufferMut, get_bit};
-use vortex_mask::{Mask, MaskIter};
+use vortex_mask::Mask;
 
 use crate::filter::{Filter, MaskIndices};
 
-/// If the filter density is above 80%, we use slices to filter the array instead of indices.
-// TODO(ngates): we need more experimentation to determine the best threshold here.
-const FILTER_SLICES_DENSITY_THRESHOLD: f64 = 0.8;
-
 impl Filter<Mask> for &BitBuffer {
     type Output = BitBuffer;
 
@@ -23,12 +19,29 @@ impl Filter<Mask> for &BitBuffer {
         match selection_mask {
             Mask::AllTrue(_) => self.clone(),
             Mask::AllFalse(_) => BitBuffer::empty(),
-            Mask::Values(v) => match v.threshold_iter(FILTER_SLICES_DENSITY_THRESHOLD) {
-                MaskIter::Indices(indices) => filter_indices(self, indices),
-                MaskIter::Slices(slices) => {
-                    filter_slices(self, selection_mask.true_count(), slices)
-                }
-            },
+            Mask::Values(v) => {
+                filter_indices(self.inner().as_ref(), self.offset(), v.indices()).freeze()
+            }
+        }
+    }
+}
+
+impl Filter<Mask> for &mut BitBufferMut {
+    type Output = ();
+
+    fn filter(self, selection_mask: &Mask) {
+        assert_eq!(
+            selection_mask.len(),
+            self.len(),
+            "Selection mask length must equal the mask length"
+        );
+
+        match selection_mask {
+            Mask::AllTrue(_) => {}
+            Mask::AllFalse(_) => self.clear(),
+            Mask::Values(v) => {
+                *self = filter_indices(self.inner().as_slice(), self.offset(), v.indices())
+            }
         }
     }
 }
@@ -37,25 +50,24 @@ impl Filter<MaskIndices<'_>> for &BitBuffer {
     type Output = BitBuffer;
 
     fn filter(self, indices: &MaskIndices) -> BitBuffer {
-        filter_indices(self, indices)
+        filter_indices(self.inner().as_ref(), self.offset(), indices).freeze()
     }
 }
 
-fn filter_indices(bools: &BitBuffer, indices: &[usize]) -> BitBuffer {
-    let buffer = bools.inner().as_ref();
-    BitBuffer::collect_bool(indices.len(), |idx| {
-        let idx = *unsafe { indices.get_unchecked(idx) };
-        get_bit(buffer, bools.offset() + idx)
-    })
-}
+impl Filter<MaskIndices<'_>> for &mut BitBufferMut {
+    type Output = ();
 
-fn filter_slices(buffer: &BitBuffer, output_len: usize, slices: &[(usize, usize)]) -> BitBuffer {
-    let mut builder = BitBufferMut::with_capacity(output_len);
-    for (start, end) in slices {
-        // TODO(ngates): we probably want a borrowed slice for things like this.
-        builder.append_buffer(&buffer.slice(*start..*end));
+    fn filter(self, indices: &MaskIndices) {
+        *self = filter_indices(self.inner().as_ref(), self.offset(), indices)
     }
-    builder.freeze()
+}
+
+fn filter_indices(bools: &[u8], bit_offset: usize, indices: &[usize]) -> BitBufferMut {
+    // FIXME(ngates): this is slower than it could be!
+    BitBufferMut::collect_bool(indices.len(), |idx| {
+        let idx = *unsafe { indices.get_unchecked(idx) };
+        get_bit(bools, bit_offset + idx)
+    })
 }
 
 #[cfg(test)]
@@ -64,20 +76,10 @@ mod test {
 
     use super::*;
 
-    #[test]
-    fn filter_bool_by_slice_test() {
-        let bits = bitbuffer![1 1 0];
-
-        let filtered = filter_slices(&bits, 2, &[(0, 1), (2, 3)]);
-        assert_eq!(2, filtered.len());
-
-        assert_eq!(filtered, bitbuffer![1 0])
-    }
-
     #[test]
     fn filter_bool_by_index_test() {
         let buf = bitbuffer![1 1 0];
-        let filtered = filter_indices(&buf, &[0, 2]);
+        let filtered = filter_indices(buf.inner().as_ref(), 0, &[0, 2]).freeze();
         assert_eq!(2, filtered.len());
         assert_eq!(filtered, bitbuffer![1 0])
     }