use scanbuilder::new

Signed-off-by: Mikhail Kot <mikhail@spiraldb.com>

Author: myrrc

vortex-bench/src/random_access/take.rs

@@ -24,7 +24,8 @@ use vortex::array::VortexSessionExecute;
 use vortex::array::stream::ArrayStreamExt;
 use vortex::buffer::Buffer;
 use vortex::file::OpenOptionsSessionExt;
-use vortex::file::VortexFile;
+use vortex::layout::LayoutReader;
+use vortex::layout::scan::scan_builder::ScanBuilder;
 use vortex::utils::aliases::hash_map::HashMap;
 
 use crate::Format;
@@ -38,7 +39,7 @@ use crate::random_access::RandomAccessorRet;
 pub struct VortexRandomAccessor {
     name: String,
     format: Format,
-    file: VortexFile,
+    reader: Arc<dyn LayoutReader>,
 }
 
 impl VortexRandomAccessor {
@@ -49,10 +50,11 @@ impl VortexRandomAccessor {
         format: Format,
     ) -> anyhow::Result<Self> {
         let file = SESSION.open_options().open_path(path.as_ref()).await?;
+        let reader = file.layout_reader()?;
         Ok(Self {
             name: name.into(),
             format,
-            file,
+            reader,
         })
     }
 }
@@ -69,9 +71,8 @@ impl RandomAccessor for VortexRandomAccessor {
 
     async fn take(&self, indices: &[u64]) -> anyhow::Result<RandomAccessorRet> {
         let indices_buf: Buffer<u64> = Buffer::from(indices.to_vec());
-        let array = self
-            .file
-            .scan()?
+        let reader = ScanBuilder::new(SESSION.clone(), self.reader.clone());
+        let array = reader
             .with_row_indices(indices_buf)
             .into_array_stream()?
             .read_all()

vortex-layout/src/layouts/row_idx/mod.rs

@@ -24,7 +24,6 @@ use vortex_array::dtype::FieldMask;
 use vortex_array::dtype::FieldName;
 use vortex_array::dtype::Nullability;
 use vortex_array::dtype::PType;
-use vortex_array::expr::ExactExpr;
 use vortex_array::expr::Expression;
 use vortex_array::expr::is_root;
 use vortex_array::expr::root;
@@ -50,7 +49,7 @@ pub struct RowIdxLayoutReader {
     name: Arc<str>,
     row_offset: u64,
     child: Arc<dyn LayoutReader>,
-    partition_cache: DashMap<ExactExpr, Arc<OnceLock<Partitioning>>>,
+    partition_cache: DashMap<Expression, Arc<OnceLock<Partitioning>>>,
     session: VortexSession,
 }
 
@@ -66,10 +65,7 @@ impl RowIdxLayoutReader {
     }
 
     fn partition_expr(&self, expr: &Expression) -> VortexResult<Partitioning> {
-        let key = ExactExpr(expr.clone());
-
-        // Check cache first with read-only lock.
-        if let Some(entry) = self.partition_cache.get(&key)
+        if let Some(entry) = self.partition_cache.get(&expr)
             && let Some(partitioning) = entry.value().get()
         {
             return Ok(partitioning.clone());
@@ -78,7 +74,7 @@ impl RowIdxLayoutReader {
         let result = self.compute_partitioning(expr)?;
 
         self.partition_cache
-            .entry(key)
+            .entry(expr.clone())
             .or_insert_with(|| Arc::new(OnceLock::new()))
             .get_or_init(|| result.clone());
 

vortex-layout/src/layouts/struct_/reader.rs

@@ -19,7 +19,6 @@ use vortex_array::dtype::FieldMask;
 use vortex_array::dtype::FieldName;
 use vortex_array::dtype::Nullability;
 use vortex_array::dtype::StructFields;
-use vortex_array::expr::ExactExpr;
 use vortex_array::expr::Expression;
 use vortex_array::expr::col;
 use vortex_array::expr::make_free_field_annotator;
@@ -59,7 +58,7 @@ pub struct StructReader {
     expanded_root_expr: Expression,
 
     field_lookup: Option<HashMap<FieldName, usize>>,
-    partitioned_expr_cache: DashMap<ExactExpr, Arc<OnceLock<Partitioned>>>,
+    partitioned_expr_cache: DashMap<Expression, Arc<OnceLock<Partitioned>>>,
 }
 
 impl StructReader {
@@ -154,18 +153,16 @@ impl StructReader {
 
     /// Utility for partitioning an expression over the fields of a struct.
     fn partition_expr(&self, expr: Expression) -> VortexResult<Partitioned> {
-        let key = ExactExpr(expr.clone());
-
-        if let Some(entry) = self.partitioned_expr_cache.get(&key)
+        if let Some(entry) = self.partitioned_expr_cache.get(&expr)
             && let Some(partitioning) = entry.value().get()
         {
             return Ok(partitioning.clone());
         }
 
-        let result = self.compute_partitioned_expr(expr)?;
+        let result = self.compute_partitioned_expr(expr.clone())?;
 
         self.partitioned_expr_cache
-            .entry(key)
+            .entry(expr)
             .or_insert_with(|| Arc::new(OnceLock::new()))
             .get_or_init(|| result.clone());