Faster pipelines

Signed-off-by: Nicholas Gates <[email protected]>

Author: gatesn

encodings/fastlanes/benches/pipeline_v2_bitpacking_basic.rs

@@ -26,10 +26,6 @@ const BENCH_PARAMS: &[(usize, f64)] = &[
     (10_000, 1.0),
     (100_000, 0.5),
     (100_000, 1.0),
-    (1_000_000, 0.5),
-    (1_000_000, 1.0),
-    (10_000_000, 0.5),
-    (10_000_000, 1.0),
 ];
 
 #[divan::bench(args = BENCH_PARAMS)]

encodings/fastlanes/src/bitpacking/array/bitpack_pipeline.rs

@@ -57,6 +57,8 @@ impl PipelinedNode for BitPackedArray {
             Ok(Box::new(AlignedBitPackedKernel::<T>::new(
                 packed_bit_width,
                 packed_buffer,
+                // FIXME(ngates): if we make sure the mask has offset zero, we know that split_off
+                //  inside the kernel is free.
                 self.validity.to_mask(self.len()).into_mut(),
             )) as Box<dyn Kernel>)
         })
@@ -180,7 +182,7 @@ impl<BP: PhysicalPType<Physical: BitPacking>> Kernel for AlignedBitPackedKernel<
             BitPacking::unchecked_unpack(
                 self.packed_bit_width,
                 packed_bytes,
-                transmute(elements.as_mut()),
+                transmute::<&mut [BP], &mut [BP::Physical]>(elements.as_mut()),
             );
         }
 

vortex-buffer/src/bit/buf.rs

@@ -568,7 +568,9 @@ mod tests {
     fn test_slice_offset_calculation() {
         let buf = BitBuffer::collect_bool(16, |_| true);
         let sliced = buf.slice(10..16);
-        assert_eq!(sliced.offset(), 10);
+        assert_eq!(sliced.len(), 6);
+        // Ensure the offset is modulo 8
+        assert_eq!(sliced.offset(), 2);
     }
 
     #[rstest]

vortex-buffer/src/bit/buf_mut.rs

@@ -437,35 +437,43 @@ impl BitBufferMut {
     ///
     /// Unlike bytes, if the split position is not on a byte-boundary this operation will copy
     /// data into the result type, and mutate self.
+    #[must_use = "consider BitBufferMut::truncate if you don't need the other half"]
     pub fn split_off(&mut self, at: usize) -> Self {
-        assert!(at <= self.len, "index {at} exceeds len {}", self.len);
+        assert!(
+            at <= self.capacity(),
+            "index {at} exceeds capacity {}",
+            self.capacity()
+        );
 
-        let new_offset = self.offset;
-        let new_len = self.len - at;
+        // The length of the tail is any bits after `at`
+        let tail_len = self.len.saturating_sub(at);
+        let byte_pos = (self.offset + at).div_ceil(8);
 
         // If we are splitting on a byte boundary, we can just slice the buffer
-        if (self.offset + at) % 8 == 0 {
-            let byte_pos = (self.offset + at) / 8;
-            let new_buffer = self.buffer.split_off(byte_pos);
-            self.len = at;
+        // Or if `at > self.len`, then the tail is empty anyway and we can just return as much
+        // of the existing capacity as possible.
+        if at > self.len() || ((self.offset + at) % 8 == 0) {
+            let tail_buffer = self.buffer.split_off(byte_pos);
+            self.len = self.len.min(at);
+
+            // Return the tail buffer
             return Self {
-                buffer: new_buffer,
-                offset: new_offset,
-                len: new_len,
+                buffer: tail_buffer,
+                offset: 0,
+                len: tail_len,
             };
         }
 
-        // Otherwise, we need to copy bits into a new buffer
-        let mut new_buffer = BitBufferMut::with_capacity(new_len);
-        for i in 0..new_len {
-            let value = self.value(at + i);
-            new_buffer.append(value);
-        }
+        // Otherwise, we truncate ourselves, and copy any bits into a new tail buffer.
+        // Note that in this case we do not preserve the capacity.
+        let u64_cap = tail_len.div_ceil(8);
+        let mut tail_buffer_u64 = BufferMut::<u64>::with_capacity(u64_cap);
+        tail_buffer_u64.extend(
+            BitChunks::new(self.buffer.as_slice(), self.offset + at, tail_len).iter_padded(),
+        );
 
-        // Truncate self to the split position
         self.truncate(at);
-
-        new_buffer
+        BitBufferMut::from_buffer(tail_buffer_u64.into_byte_buffer(), 0, tail_len)
     }
 
     /// Absorbs a mutable buffer that was previously split off.
@@ -1058,4 +1066,40 @@ mod tests {
             assert_eq!(bit_buf.value(i), i % 2 == 0);
         }
     }
+
+    #[test]
+    fn test_split_off() {
+        // Test splitting at various positions and across a byte boundary
+        for i in 0..10 {
+            let buf = bitbuffer![0 1 0 1 0 1 0 1 0 1];
+
+            let mut buf_mut = buf.clone().into_mut();
+            assert_eq!(buf_mut.len(), 10);
+
+            let tail = buf_mut.split_off(i);
+            assert_eq!(buf_mut.len(), i);
+            assert_eq!(buf_mut.freeze(), buf.slice(0..i));
+
+            assert_eq!(tail.len(), 10 - i);
+            assert_eq!(tail.freeze(), buf.slice(i..10));
+        }
+    }
+
+    #[test]
+    fn test_split_off_with_offset() {
+        // Test splitting at various positions and across a byte boundary
+        for i in 0..10 {
+            let buf = bitbuffer![0 1 0 1 0 1 0 1 0 1 0 1].slice(2..);
+
+            let mut buf_mut = buf.clone().into_mut();
+            assert_eq!(buf_mut.len(), 10);
+
+            let tail = buf_mut.split_off(i);
+            assert_eq!(buf_mut.len(), i);
+            assert_eq!(buf_mut.freeze(), buf.slice(0..i));
+
+            assert_eq!(tail.len(), 10 - i);
+            assert_eq!(tail.freeze(), buf.slice(i..10));
+        }
+    }
 }

vortex-buffer/src/buffer_mut.rs

@@ -391,6 +391,16 @@ impl<T> BufferMut<T> {
         self.length += other.length;
     }
 
+    /// Return the [`ByteBufferMut`] for this [`BufferMut`].
+    pub fn into_byte_buffer(self) -> ByteBufferMut {
+        ByteBufferMut {
+            bytes: self.bytes,
+            length: self.length * size_of::<T>(),
+            alignment: self.alignment,
+            _marker: Default::default(),
+        }
+    }
+
     /// Freeze the `BufferMut` into a `Buffer`.
     pub fn freeze(self) -> Buffer<T> {
         Buffer {