Optimize partial Flat random access with blocking I/O

Cargo.toml

@@ -169,6 +169,7 @@ geoarrow = "0.8.0"
 geoarrow-cast = "0.8.0"
 get_dir = "0.5.0"
 glob = "0.3.2"
+io-uring = "0.7.13"
 goldenfile = "1"
 half = { version = "2.7.1", features = ["std", "num-traits"] }
 hashbrown = "0.17.1"

benchmarks/datafusion-bench/src/main.rs

@@ -289,10 +289,14 @@ async fn register_v2_tables<B: Benchmark + ?Sized>(
             .runtime_env()
             .object_store(table_url.object_store())?;
 
-        let fs: FileSystemRef = Arc::new(ObjectStoreFileSystem::new(
-            Arc::clone(&store),
-            SESSION.handle(),
-        ));
+        let fs: FileSystemRef = if benchmark_base.scheme() == "file" {
+            Arc::new(ObjectStoreFileSystem::local(SESSION.handle()))
+        } else {
+            Arc::new(ObjectStoreFileSystem::new(
+                Arc::clone(&store),
+                SESSION.handle(),
+            ))
+        };
         let base_prefix = benchmark_base.path().trim_start_matches('/').to_string();
         let fs = fs.with_prefix(base_prefix);
 

encodings/alp/src/alp_rd/array.rs

@@ -71,6 +71,33 @@ pub struct ALPRDMetadata {
     patches: Option<PatchesMetadata>,
 }
 
+impl ALPRDMetadata {
+    pub fn right_bit_width(&self) -> VortexResult<u8> {
+        u8::try_from(self.right_bit_width).map_err(|_| {
+            vortex_err!(
+                "right bit width {} does not fit in u8",
+                self.right_bit_width
+            )
+        })
+    }
+
+    pub fn left_parts_dictionary(&self) -> VortexResult<Buffer<u16>> {
+        self.dict
+            .get(..usize::try_from(self.dict_len)?)
+            .ok_or_else(|| vortex_err!("ALPRD dictionary length is out of bounds"))?
+            .iter()
+            .map(|&value| {
+                u16::try_from(value)
+                    .map_err(|_| vortex_err!("ALPRD dictionary value {value} does not fit in u16"))
+            })
+            .collect()
+    }
+
+    pub fn patches(&self) -> Option<&PatchesMetadata> {
+        self.patches.as_ref()
+    }
+}
+
 impl ArrayHash for ALPRDData {
     fn array_hash<H: Hasher>(&self, state: &mut H, accuracy: EqMode) {
         self.left_parts_dictionary.array_hash(state, accuracy);

encodings/fastlanes/src/bitpacking/mod.rs

@@ -19,6 +19,7 @@ mod vtable;
 pub(crate) use plugin::BitPackedPatchedPlugin;
 pub use vtable::BitPacked;
 pub use vtable::BitPackedArray;
+pub use vtable::BitPackedMetadata;
 
 pub(crate) fn initialize(session: &vortex_session::VortexSession) {
     vtable::initialize(session);

encodings/fastlanes/src/bitpacking/vtable/mod.rs

@@ -72,6 +72,22 @@ pub struct BitPackedMetadata {
     pub(crate) patches: Option<PatchesMetadata>,
 }
 
+impl BitPackedMetadata {
+    pub fn bit_width(&self) -> VortexResult<u8> {
+        u8::try_from(self.bit_width)
+            .map_err(|_| vortex_err!("bit width {} does not fit in u8", self.bit_width))
+    }
+
+    pub fn offset(&self) -> VortexResult<u16> {
+        u16::try_from(self.offset)
+            .map_err(|_| vortex_err!("bit-packed offset {} does not fit in u16", self.offset))
+    }
+
+    pub fn patches(&self) -> Option<&PatchesMetadata> {
+        self.patches.as_ref()
+    }
+}
+
 impl ArrayHash for BitPackedData {
     fn array_hash<H: Hasher>(&self, state: &mut H, accuracy: EqMode) {
         self.offset.hash(state);

vortex-array/src/arrays/list/mod.rs

@@ -14,6 +14,7 @@ pub(crate) mod compute;
 
 mod vtable;
 pub use vtable::List;
+pub use vtable::ListMetadata;
 
 pub(crate) fn initialize(session: &vortex_session::VortexSession) {
     compute::initialize(session);

vortex-array/src/arrays/list/vtable/mod.rs

@@ -52,6 +52,12 @@ pub struct ListMetadata {
     offset_ptype: i32,
 }
 
+impl ListMetadata {
+    pub fn elements_len(&self) -> u64 {
+        self.elements_len
+    }
+}
+
 impl ArrayHash for ListData {
     fn array_hash<H: Hasher>(&self, _state: &mut H, _accuracy: EqMode) {}
 }

vortex-array/src/arrays/struct_/compute/rules.rs

@@ -77,6 +77,32 @@ fn reduce_struct_cast(
     ))
 }
 
+pub(crate) fn struct_get_item(
+    array: ArrayView<'_, Struct>,
+    field_name: &crate::dtype::FieldName,
+) -> VortexResult<ArrayRef> {
+    let field = array
+        .unmasked_field_by_name_opt(field_name)
+        .ok_or_else(|| {
+            vortex_err!(
+                "Field {} missing from struct array {}",
+                field_name,
+                array.struct_fields().names()
+            )
+        })?;
+
+    match array.validity()? {
+        Validity::NonNullable => Ok(field.clone()),
+        Validity::AllValid => field.clone().cast(field.dtype().as_nullable()),
+        Validity::AllInvalid => Ok(ConstantArray::new(
+            Scalar::null(field.dtype().as_nullable()),
+            field.len(),
+        )
+        .into_array()),
+        Validity::Array(mask) => field.clone().mask(mask),
+    }
+}
+
 /// Rule to flatten get_item from struct by field name
 #[derive(Debug)]
 pub(crate) struct StructGetItemRule;
@@ -90,38 +116,7 @@ impl ArrayParentReduceRule<Struct> for StructGetItemRule {
         parent: ScalarFnArrayView<'_, GetItem>,
         _child_idx: usize,
     ) -> VortexResult<Option<ArrayRef>> {
-        let field_name = parent.options;
-        let field = child
-            .unmasked_field_by_name_opt(field_name)
-            .ok_or_else(|| {
-                vortex_err!(
-                    "Field '{}' missing from struct array {}",
-                    field_name,
-                    child.struct_fields().names()
-                )
-            })?;
-
-        match child.validity()? {
-            Validity::NonNullable => {
-                // If the struct is non-nullable, the field's validity is unchanged
-                Ok(Some(field.clone()))
-            }
-            Validity::AllValid => {
-                // If struct is nullable, field must also be nullable
-                field.clone().cast(field.dtype().as_nullable()).map(Some)
-            }
-            Validity::AllInvalid => {
-                // If everything is invalid, the field is also all invalid
-                Ok(Some(
-                    ConstantArray::new(Scalar::null(field.dtype().as_nullable()), field.len())
-                        .into_array(),
-                ))
-            }
-            Validity::Array(mask) => {
-                // If the validity is an array, we need to combine it with the field's validity
-                field.clone().mask(mask).map(Some)
-            }
-        }
+        struct_get_item(child, parent.options).map(Some)
     }
 }
 

vortex-array/src/expression.rs

@@ -9,10 +9,16 @@ use crate::ArrayRef;
 use crate::IntoArray;
 use crate::arrays::ConstantArray;
 use crate::arrays::ScalarFnArray;
+use crate::arrays::Struct;
+use crate::arrays::StructArray;
+use crate::arrays::struct_::compute::rules::struct_get_item;
 use crate::expr::BoundExpression;
 use crate::expr::Expression;
 use crate::optimizer::ArrayOptimizer;
+use crate::scalar_fn::fns::get_item::GetItem;
 use crate::scalar_fn::fns::literal::Literal;
+use crate::scalar_fn::fns::pack::Pack;
+use crate::validity::Validity;
 
 impl ArrayRef {
     /// Apply a bound expression to this array, producing a new array in constant time.
@@ -35,6 +41,24 @@ impl ArrayRef {
             .map(|child| self.clone().apply_bound(child))
             .try_collect()?;
 
+        if let Some(field_name) = scalar_fn.as_opt::<GetItem>()
+            && let [child] = children.as_slice()
+            && let Some(array) = child.as_opt::<Struct>()
+        {
+            return struct_get_item(array, field_name);
+        }
+
+        if let Some(pack) = scalar_fn.as_opt::<Pack>() {
+            let validity = match pack.nullability {
+                crate::dtype::Nullability::NonNullable => Validity::NonNullable,
+                crate::dtype::Nullability::Nullable => Validity::AllValid,
+            };
+            return Ok(
+                StructArray::try_new(pack.names.clone(), children, self.len(), validity)?
+                    .into_array(),
+            );
+        }
+
         let array =
             ScalarFnArray::try_new_with_len(scalar_fn.clone(), children, self.len())?.into_array();
 

vortex-array/src/mask_future.rs

@@ -20,6 +20,9 @@ use vortex_mask::Mask;
 pub struct MaskFuture {
     inner: Shared<BoxFuture<'static, SharedVortexResult<Mask>>>,
     len: usize,
+    upper_bound: Option<Mask>,
+    upper_bound_is_exact: bool,
+    partial_reads_allowed: bool,
 }
 
 impl MaskFuture {
@@ -40,6 +43,9 @@ impl MaskFuture {
                 .boxed()
                 .shared(),
             len,
+            upper_bound: None,
+            upper_bound_is_exact: false,
+            partial_reads_allowed: false,
         }
     }
 
@@ -55,7 +61,12 @@ impl MaskFuture {
 
     /// Create a MaskFuture from a ready mask.
     pub fn ready(mask: Mask) -> Self {
-        Self::new(mask.len(), async move { Ok(mask) })
+        let upper_bound = mask.clone();
+        let mut future = Self::new(mask.len(), async move { Ok(mask) });
+        future.upper_bound = Some(upper_bound);
+        future.upper_bound_is_exact = true;
+        future.partial_reads_allowed = true;
+        future
     }
 
     /// Create a MaskFuture that resolves to a mask with all values set to true.
@@ -72,7 +83,57 @@ impl MaskFuture {
         }
 
         let inner = self.inner.clone();
-        Self::new(range.len(), async move { Ok(inner.await?.slice(range)) })
+        let upper_bound = self
+            .upper_bound
+            .as_ref()
+            .map(|upper_bound| upper_bound.slice(range.clone()));
+        let mut sliced = Self::new(range.len(), async move { Ok(inner.await?.slice(range)) });
+        sliced.upper_bound = upper_bound;
+        sliced.upper_bound_is_exact = self.upper_bound_is_exact;
+        sliced.partial_reads_allowed = self.partial_reads_allowed;
+        sliced
+    }
+
+    /// Attach a conservative upper bound for the mask resolved by this future.
+    ///
+    /// Readers can use this to register I/O eagerly without waiting for filter evaluation. The
+    /// resolved mask must not contain a true row that is false in `upper_bound`.
+    pub fn with_upper_bound(mut self, upper_bound: Mask) -> Self {
+        assert_eq!(
+            upper_bound.len(),
+            self.len,
+            "MaskFuture upper bound length mismatch"
+        );
+        self.upper_bound = Some(upper_bound);
+        self.upper_bound_is_exact = false;
+        self
+    }
+
+    /// Return the conservative upper bound for this future, when one is known.
+    pub fn upper_bound(&self) -> Option<&Mask> {
+        self.upper_bound.as_ref()
+    }
+
+    /// Return whether the upper bound is the exact mask returned by this future.
+    pub fn upper_bound_is_exact(&self) -> bool {
+        self.upper_bound_is_exact
+    }
+
+    /// Permit readers to satisfy this selection using partial segment reads.
+    pub fn with_partial_reads(mut self) -> Self {
+        self.partial_reads_allowed = true;
+        self
+    }
+
+    /// Prevent readers from turning this mask into partial segment reads.
+    pub fn without_partial_reads(mut self) -> Self {
+        self.partial_reads_allowed = false;
+        self
+    }
+
+    /// Return whether readers may satisfy this selection using partial segment reads.
+    pub fn partial_reads_allowed(&self) -> bool {
+        self.partial_reads_allowed
     }
 
     pub fn inspect(
@@ -84,6 +145,9 @@ impl MaskFuture {
         Self {
             inner: self.inner.inspect(f).boxed().shared(),
             len,
+            upper_bound: self.upper_bound,
+            upper_bound_is_exact: self.upper_bound_is_exact,
+            partial_reads_allowed: self.partial_reads_allowed,
         }
     }
 }
@@ -119,8 +183,26 @@ mod tests {
 
             let partial = fut.slice(0..mask.len() - 1);
             assert_eq!(partial.len(), mask.len() - 1);
+            assert_eq!(partial.upper_bound(), Some(&mask.slice(0..mask.len() - 1)));
             assert_eq!(partial.await?, mask.slice(0..mask.len() - 1));
             Ok(())
         })
     }
+
+    #[test]
+    fn new_future_has_no_upper_bound_until_attached() {
+        let future = MaskFuture::new(3, async { Ok(Mask::new_false(3)) });
+        assert!(future.upper_bound().is_none());
+
+        let upper_bound = Mask::from_indices(3, [0, 2]);
+        let future = future.with_upper_bound(upper_bound.clone());
+        assert_eq!(future.upper_bound(), Some(&upper_bound));
+        assert!(!future.upper_bound_is_exact());
+    }
+
+    #[test]
+    fn ready_future_has_an_exact_upper_bound() {
+        let future = MaskFuture::ready(Mask::from_indices(3, [0, 2]));
+        assert!(future.upper_bound_is_exact());
+    }
 }