Compressed Statistics (#2686)

Compress the stats tables

Author: gatesn

bench-vortex/src/tpch/mod.rs

@@ -25,8 +25,7 @@ use vortex::arrays::ChunkedArray;
 use vortex::arrow::FromArrowArray;
 use vortex::dtype::DType;
 use vortex::error::VortexExpect as _;
-use vortex::file::{DEFAULT_REGISTRY, VORTEX_FILE_EXTENSION, VortexWriteOptions};
-use vortex::layout::{LayoutRegistry, LayoutRegistryExt};
+use vortex::file::{VORTEX_FILE_EXTENSION, VortexWriteOptions};
 use vortex::{Array, ArrayRef, TryIntoArray};
 use vortex_datafusion::SessionContextExt;
 use vortex_datafusion::persistent::VortexFormat;
@@ -352,10 +351,7 @@ async fn register_vortex_file(
         .await?;
     }
 
-    let format = Arc::new(VortexFormat::new(
-        DEFAULT_REGISTRY.clone(),
-        Arc::new(LayoutRegistry::default()),
-    ));
+    let format = Arc::new(VortexFormat::default());
     let table_url = ListingTableUrl::parse(vtx_file.as_str())?;
     let config = ListingTableConfig::new(table_url)
         .with_listing_options(ListingOptions::new(format as _))

vortex-array/src/compute/is_sorted.rs

@@ -59,6 +59,11 @@ where
 }
 
 pub fn is_sorted(array: &dyn Array) -> VortexResult<bool> {
+    // We currently don't support sorting struct arrays.
+    if array.dtype().is_struct() {
+        return Ok(false);
+    }
+
     if let Some(Precision::Exact(value)) = array.statistics().get_as::<bool>(Stat::IsSorted) {
         return Ok(value);
     }
@@ -76,6 +81,11 @@ pub fn is_sorted(array: &dyn Array) -> VortexResult<bool> {
     Ok(is_sorted)
 }
 pub fn is_strict_sorted(array: &dyn Array) -> VortexResult<bool> {
+    // We currently don't support sorting struct arrays.
+    if array.dtype().is_struct() {
+        return Ok(false);
+    }
+
     if let Some(Precision::Exact(value)) = array.statistics().get_as::<bool>(Stat::IsStrictSorted) {
         return Ok(value);
     }

vortex-datafusion/src/persistent/format.rs

@@ -28,7 +28,7 @@ use vortex_dtype::arrow::FromArrowType;
 use vortex_error::{VortexExpect, VortexResult, vortex_err};
 use vortex_expr::datafusion::convert_expr_to_vortex;
 use vortex_expr::{VortexExpr, and};
-use vortex_file::{VORTEX_FILE_EXTENSION, VortexOpenOptions};
+use vortex_file::{DEFAULT_REGISTRY, VORTEX_FILE_EXTENSION, VortexOpenOptions};
 use vortex_io::ObjectStoreReadAt;
 use vortex_layout::{LayoutRegistry, LayoutRegistryExt};
 
@@ -69,7 +69,7 @@ impl VortexFormatFactory {
     // Because FileFormatFactory has a default method
     /// Create a new [`VortexFormatFactory`] with the default encoding context.
     pub fn default_config() -> Self {
-        Self::with_registry(Arc::new(ArrayRegistry::default()))
+        Self::with_registry(DEFAULT_REGISTRY.clone())
     }
 
     /// Create a new [`VortexFormatFactory`] that creates [`VortexFormat`] instances with the provided [`Context`](vortex_array::ArrayContext).

vortex-file/src/strategy.rs

@@ -2,6 +2,7 @@
 
 use std::sync::Arc;
 
+use vortex_array::arcref::ArcRef;
 use vortex_array::stats::{PRUNING_STATS, STATS_TO_WRITE};
 use vortex_array::{Array, ArrayContext, ArrayRef};
 use vortex_btrblocks::BtrBlocksCompressor;
@@ -34,7 +35,7 @@ impl LayoutStrategy for VortexLayoutStrategy {
                 ctx.clone(),
                 &DType::Null,
                 ChunkedLayoutOptions {
-                    chunk_strategy: Arc::new(FlatLayoutOptions::default()),
+                    chunk_strategy: ArcRef::new_arc(Arc::new(FlatLayoutOptions::default()) as _),
                 },
             )
             .boxed(),
@@ -59,7 +60,9 @@ impl LayoutStrategy for VortexLayoutStrategy {
                 ctx.clone(),
                 dtype,
                 writer,
-                Arc::new(FlatLayout),
+                ArcRef::new_arc(Arc::new(BtrBlocksCompressedStrategy {
+                    child: ArcRef::new_ref(&FlatLayout),
+                })),
                 StatsLayoutOptions {
                     block_size: 8192,
                     stats: PRUNING_STATS.into(),
@@ -79,6 +82,17 @@ impl LayoutStrategy for VortexLayoutStrategy {
     }
 }
 
+struct BtrBlocksCompressedStrategy {
+    child: ArcRef<dyn LayoutStrategy>,
+}
+
+impl LayoutStrategy for BtrBlocksCompressedStrategy {
+    fn new_writer(&self, ctx: &ArrayContext, dtype: &DType) -> VortexResult<Box<dyn LayoutWriter>> {
+        let child = self.child.new_writer(ctx, dtype)?;
+        Ok(BtrBlocksCompressedWriter { child }.boxed())
+    }
+}
+
 /// A layout writer that compresses chunks using a sampling compressor, and re-uses the previous
 /// compressed chunk as a hint for the next.
 struct BtrBlocksCompressedWriter {

vortex-layout/src/data.rs

@@ -260,7 +260,7 @@ impl Layout {
         &self,
         segment_reader: Arc<dyn AsyncSegmentReader>,
         ctx: ArrayContext,
-    ) -> VortexResult<Arc<dyn LayoutReader + 'static>> {
+    ) -> VortexResult<Arc<dyn LayoutReader>> {
         self.vtable().reader(self.clone(), ctx, segment_reader)
     }
 

vortex-layout/src/layouts/chunked/eval_expr.rs

@@ -51,7 +51,7 @@ impl ExprEvaluator for ChunkedReader {
         Ok(ChunkedArray::new_unchecked(chunks, dtype).into_array())
     }
 
-    async fn prune_mask(&self, row_mask: RowMask, _expr: ExprRef) -> VortexResult<RowMask> {
+    async fn refine_mask(&self, row_mask: RowMask, _expr: ExprRef) -> VortexResult<RowMask> {
         // TODO(ngates): we should push-down to each child
         Ok(row_mask)
     }

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

@@ -20,9 +20,6 @@ use crate::{CHUNKED_LAYOUT_ID, LayoutId};
 pub struct ChunkedLayout;
 
 /// In-memory representation of Chunked layout.
-///
-/// First child in the list is the metadata table
-/// Subsequent children are consecutive chunks of this layout
 impl LayoutVTable for ChunkedLayout {
     fn id(&self) -> LayoutId {
         CHUNKED_LAYOUT_ID
@@ -44,9 +41,8 @@ impl LayoutVTable for ChunkedLayout {
         row_offset: u64,
         splits: &mut BTreeSet<u64>,
     ) -> VortexResult<()> {
-        let nchunks = layout.nchildren() - (if layout.metadata().is_some() { 1 } else { 0 });
         let mut offset = row_offset;
-        for i in 0..nchunks {
+        for i in 0..layout.nchildren() {
             let child = layout.child(i, layout.dtype().clone(), format!("[{}]", i))?;
             child.register_splits(field_mask, offset, splits)?;
             offset += child.row_count();
@@ -63,9 +59,8 @@ impl LayoutVTable for ChunkedLayout {
         projection_field_mask: &[FieldMask],
         segments: &mut SegmentCollector,
     ) -> VortexResult<()> {
-        let nchunks = layout.nchildren() - (if layout.metadata().is_some() { 1 } else { 0 });
         let mut offset = row_offset;
-        for i in 0..nchunks {
+        for i in 0..layout.nchildren() {
             let child = layout.child(i, layout.dtype().clone(), format!("[{i}]"))?;
             child.required_segments(offset, filter_field_mask, projection_field_mask, segments)?;
             offset += child.row_count();

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

@@ -33,11 +33,7 @@ impl ChunkedReader {
         }
 
         // The number of chunks
-        let mut nchunks = layout.nchildren();
-        if layout.metadata().is_some() {
-            // The final child is the statistics table.
-            nchunks -= 1;
-        }
+        let nchunks = layout.nchildren();
 
         // Construct a lazy scan for each chunk of the layout.
         let chunk_readers = (0..nchunks).map(|_| OnceLock::new()).collect();

vortex-layout/src/layouts/chunked/writer.rs

@@ -1,5 +1,4 @@
-use std::sync::Arc;
-
+use vortex_array::arcref::ArcRef;
 use vortex_array::{ArrayContext, ArrayRef};
 use vortex_dtype::DType;
 use vortex_error::{VortexExpect, VortexResult};
@@ -14,13 +13,13 @@ use crate::writer::LayoutWriter;
 
 pub struct ChunkedLayoutOptions {
     /// The layout strategy for each chunk.
-    pub chunk_strategy: Arc<dyn LayoutStrategy>,
+    pub chunk_strategy: ArcRef<dyn LayoutStrategy>,
 }
 
 impl Default for ChunkedLayoutOptions {
     fn default() -> Self {
         Self {
-            chunk_strategy: Arc::new(FlatLayout),
+            chunk_strategy: ArcRef::new_ref(&FlatLayout),
         }
     }
 }

vortex-layout/src/layouts/flat/eval_expr.rs

@@ -40,7 +40,7 @@ impl ExprEvaluator for FlatReader {
         Ok(array)
     }
 
-    async fn prune_mask(&self, row_mask: RowMask, _expr: ExprRef) -> VortexResult<RowMask> {
+    async fn refine_mask(&self, row_mask: RowMask, _expr: ExprRef) -> VortexResult<RowMask> {
         // No cheap pruning for us to do without fetching data.
         Ok(row_mask)
     }