some cleanup

Signed-off-by: Connor Tsui <[email protected]>

Author: connortsui20

vortex-compute/src/take/bit_buffer.rs

@@ -1,10 +1,19 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+//! Take operation on [`BitBuffer`].
+//!
+//! NB: We do NOT implement `impl<I: UnsignedPType> Take<PVector<I>> for &BitBuffer`, specifically
+//! because there is a very similar implementation on `Mask` that has special logic for working with
+//! null indices. That logic could also be implemented on `BitBuffer`, but since it is not
+//! immediately clear what should happen in the case of a null index when taking a `BitBuffer` (do
+//! you set it to true or false?), we do not implement this at all.
+
 use vortex_buffer::BitBuffer;
 use vortex_buffer::get_bit;
 use vortex_dtype::UnsignedPType;
 
+use crate::take::LINUX_PAGE_SIZE;
 use crate::take::Take;
 
 impl<I: UnsignedPType> Take<[I]> for &BitBuffer {
@@ -13,7 +22,7 @@ impl<I: UnsignedPType> Take<[I]> for &BitBuffer {
     fn take(self, indices: &[I]) -> BitBuffer {
         // For boolean arrays that roughly fit into a single page (at least, on Linux), it's worth
         // the overhead to convert to a `Vec<bool>`.
-        if self.len() <= 4096 {
+        if self.len() <= LINUX_PAGE_SIZE {
             let bools = self.iter().collect();
             take_byte_bool(bools, indices)
         } else {
@@ -22,12 +31,6 @@ impl<I: UnsignedPType> Take<[I]> for &BitBuffer {
     }
 }
 
-// NB: We do NOT implement `impl<I: UnsignedPType> Take<PVector<I>> for &BitBuffer`, specifically
-// because there is a very similar implementation on `Mask` that has special logic for working with
-// null indices. That logic could also be implemented on `BitBuffer`, but since it is not
-// immediately clear what should happen in the case of a null index when taking a `BitBuffer` (do
-// you set it to true or false?), we do not implement this at all.
-
 /// # Panics
 ///
 /// Panics if an index is out of bounds.

vortex-compute/src/take/mask.rs

@@ -8,6 +8,7 @@ use vortex_mask::Mask;
 use vortex_vector::VectorOps;
 use vortex_vector::primitive::PVector;
 
+use crate::take::LINUX_PAGE_SIZE;
 use crate::take::Take;
 
 impl<I: UnsignedPType> Take<[I]> for &Mask {
@@ -38,14 +39,10 @@ impl<I: UnsignedPType> Take<PVector<I>> for &Mask {
         let indices_validity = indices.validity();
         let indices_len = indices.len();
 
-        match indices_validity {
+        let indices_validity_values = match indices_validity {
             Mask::AllTrue(_) => return self.take(indices.elements().as_slice()),
             Mask::AllFalse(_) => return Mask::AllFalse(indices_len),
-            Mask::Values(_) => (),
-        };
-
-        let Mask::Values(indices_validity_values) = indices_validity else {
-            unreachable!("we just matched on the other cases above");
+            Mask::Values(indices_validity_values) => indices_validity_values,
         };
 
         match self {
@@ -56,7 +53,7 @@ impl<I: UnsignedPType> Take<PVector<I>> for &Mask {
             Mask::Values(mask_values) => {
                 // For boolean arrays that roughly fit into a single page (at least, on Linux), it's
                 // worth the overhead to convert to a `Vec<bool>`.
-                if self.len() <= 4096 {
+                if self.len() <= LINUX_PAGE_SIZE {
                     let bools = mask_values.bit_buffer().iter().collect();
                     Mask::from_buffer(take_byte_bool_nullable(bools, indices))
                 } else {

vortex-compute/src/take/mod.rs

@@ -9,6 +9,9 @@ mod mask;
 pub mod slice;
 mod vector;
 
+/// The size of a page in Linux.
+const LINUX_PAGE_SIZE: usize = 4096;
+
 /// Function for taking based on indices (which can have different representations).
 pub trait Take<Indices: ?Sized> {
     /// The result type after performing the operation.