From 60c12275773bfb25ef1de9ec51f1d75422b661a5 Mon Sep 17 00:00:00 2001 From: Adam Gutglick Date: Fri, 17 Jul 2026 12:08:47 +0100 Subject: [PATCH 001/498] Assortment of `VortexSource` improvements and fixes (#8718) ## Rationale for this change General improvements of the `FileSource`-based Apache DataFusion integration. ## What changes are included in this PR? 1. DF metrics are in use, so make better use of them both on reads and writes, which is a step towards having the same default observability that `ParquetSource` provides. 2. Fix a subtle issue where a dynamic expression only references partition literals, which we currently don't handle correctly. 3. Drop `batch_size` from the source, by default DataFusion splits the data coming out of `DataScanExec` nodes. 4. Make sure we keep schema and field metadata through all the projections paths. ## What APIs are changed? Are there any user-facing changes? More metrics will be reported in DF, and batch size needs to be done through the `DataSourceExec` and `FileScanConfig` level. --------- Signed-off-by: Adam Gutglick --- vortex-datafusion/src/convert/schema.rs | 31 +- vortex-datafusion/src/persistent/opener.rs | 324 +++++++++++++++++---- vortex-datafusion/src/persistent/sink.rs | 60 +++- vortex-datafusion/src/persistent/source.rs | 26 +- 4 files changed, 363 insertions(+), 78 deletions(-) diff --git a/vortex-datafusion/src/convert/schema.rs b/vortex-datafusion/src/convert/schema.rs index 2ec17758593..907fd5cbbe1 100644 --- a/vortex-datafusion/src/convert/schema.rs +++ b/vortex-datafusion/src/convert/schema.rs @@ -56,7 +56,10 @@ pub fn calculate_physical_schema( }) .collect::>>()?; - Ok(Schema::new(fields)) + Ok(Schema::new_with_metadata( + fields, + reference_logical_schema.metadata().clone(), + )) } /// Calculate the physical Arrow type for a field, preferring the logical type when the @@ -246,6 +249,32 @@ mod tests { ); } + #[test] + fn test_schema_metadata_preserved() -> DFResult<()> { + let logical_schema = Schema::new_with_metadata( + vec![Field::new("col", DataType::Int32, false)], + [("table".to_string(), "metadata".to_string())] + .into_iter() + .collect(), + ); + let dtype = DType::Struct( + StructFields::from_iter([( + "col", + DType::Primitive(PType::I32, Nullability::NonNullable), + )]), + Nullability::NonNullable, + ); + + let physical_schema = + calculate_physical_schema(&dtype, &logical_schema, &ArrowSession::default())?; + + assert_eq!( + physical_schema.metadata().get("table"), + Some(&"metadata".to_string()) + ); + Ok(()) + } + #[test] fn test_utf8_variants_preserved() { // Non-view string types become view types after roundtrip through DType, diff --git a/vortex-datafusion/src/persistent/opener.rs b/vortex-datafusion/src/persistent/opener.rs index b1c999656f6..d0ffb472ebe 100644 --- a/vortex-datafusion/src/persistent/opener.rs +++ b/vortex-datafusion/src/persistent/opener.rs @@ -5,11 +5,13 @@ use std::ops::Range; use std::sync::Arc; use std::sync::Weak; +use arrow_array::RecordBatchOptions; use arrow_schema::Field; use arrow_schema::Schema; use datafusion_common::DataFusionError; use datafusion_common::Result as DFResult; use datafusion_common::ScalarValue; +use datafusion_common::Statistics; use datafusion_common::arrow::array::AsArray; use datafusion_common::arrow::array::RecordBatch; use datafusion_common::exec_datafusion_err; @@ -23,11 +25,13 @@ use datafusion_physical_expr::PhysicalExprRef; use datafusion_physical_expr::projection::ProjectionExprs; use datafusion_physical_expr::simplifier::PhysicalExprSimplifier; use datafusion_physical_expr::split_conjunction; +use datafusion_physical_expr::utils::collect_columns; use datafusion_physical_expr::utils::reassign_expr_columns; use datafusion_physical_expr_adapter::PhysicalExprAdapterFactory; use datafusion_physical_expr_adapter::replace_columns_with_literals; -use datafusion_physical_expr_common::physical_expr::is_dynamic_physical_expr; -use datafusion_physical_plan::metrics::Count; +use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet; +use datafusion_physical_plan::metrics::MetricBuilder; +use datafusion_physical_plan::metrics::MetricCategory; use datafusion_pruning::FilePruner; use futures::FutureExt; use futures::StreamExt; @@ -83,12 +87,12 @@ pub(crate) struct VortexOpener { /// This is the table's schema without partition columns. It may contain fields which do /// not exist in the file, and are supplied by the `schema_adapter_factory`. pub table_schema: TableSchema, - /// A hint for the desired row count of record batches returned from the scan. - pub batch_size: usize, /// If provided, the scan will not return more than this many rows. pub limit: Option, /// A metrics object for tracking performance of the scan. pub metrics_registry: Arc, + /// DataFusion-native metrics exposed through `DataSourceExec`. + pub df_metrics: ExecutionPlanMetricsSet, /// A shared cache of file readers. /// /// To save on the overhead of reparsing FlatBuffers and rebuilding the layout tree, we cache @@ -108,6 +112,12 @@ pub(crate) struct VortexOpener { impl FileOpener for VortexOpener { fn open(&self, file: PartitionedFile) -> DFResult { + // Calculate the output schema before replacing partition columns with literals so it + // retains the table and partition-field metadata declared by the plan. + let output_schema = Arc::new( + self.projection + .project_schema(self.table_schema.table_schema())?, + ); let session = self.session.clone(); let metrics_registry = Arc::clone(&self.metrics_registry); let labels = vec![ @@ -123,12 +133,11 @@ impl FileOpener for VortexOpener { let reader = InstrumentedReadAt::new_with_labels(reader, metrics_registry.as_ref(), labels.clone()); - let file_pruning_predicate = self.file_pruning_predicate.clone(); + let mut file_pruning_predicate = self.file_pruning_predicate.clone(); let expr_adapter_factory = Arc::clone(&self.expr_adapter_factory); let file_metadata_cache = self.file_metadata_cache.clone(); let unified_file_schema = Arc::clone(self.table_schema.file_schema()); - let batch_size = self.batch_size; let limit = self.limit; let layout_readers = Arc::clone(&self.layout_readers); let natural_split_ranges = Arc::clone(&self.natural_split_ranges); @@ -138,6 +147,10 @@ impl FileOpener for VortexOpener { let expr_convertor = Arc::clone(&self.expression_convertor); let projection_pushdown = self.projection_pushdown; + let predicate_creation_errors = MetricBuilder::new(&self.df_metrics) + .with_category(MetricCategory::Rows) + .global_counter("num_predicate_creation_errors"); + // Replace column access for partition columns with literals #[expect(clippy::disallowed_types)] let literal_value_cols = self @@ -149,6 +162,13 @@ impl FileOpener for VortexOpener { .zip(file.partition_values.clone()) .collect::>(); + let predicate_uses_partition_columns = + file_pruning_predicate.as_ref().is_some_and(|predicate| { + collect_columns(predicate) + .iter() + .any(|column| literal_value_cols.contains_key(column.name())) + }); + if !literal_value_cols.is_empty() { projection = projection.try_map_exprs(|expr| { replace_columns_with_literals(Arc::clone(&expr), &literal_value_cols) @@ -156,24 +176,30 @@ impl FileOpener for VortexOpener { filter = filter .map(|p| replace_columns_with_literals(p, &literal_value_cols)) .transpose()?; + file_pruning_predicate = file_pruning_predicate + .map(|p| replace_columns_with_literals(p, &literal_value_cols)) + .transpose()?; } Ok(async move { - // Create FilePruner when we have a predicate and either dynamic expressions - // or file statistics available. The pruner can eliminate files without - // opening them based on File-level statistics (min/max values per column) + // FilePruner requires a statistics object even when the rewritten predicate + // only contains partition literals. Supply unknown file-column statistics in + // that case so static and dynamic partition predicates can still prune. + let synthetic_statistics = (!file.has_statistics() && predicate_uses_partition_columns) + .then(|| { + file.clone() + .with_statistics(Arc::new(Statistics::new_unknown(&unified_file_schema))) + }); + let pruning_file = synthetic_statistics.as_ref().unwrap_or(&file); + let mut file_pruner = file_pruning_predicate - .filter(|p| { - // Only create pruner if we have dynamic expressions or file statistics - // to work with. Static predicates without stats won't benefit from pruning. - is_dynamic_physical_expr(p) || file.has_statistics() - }) + .filter(|_| file.has_statistics() || predicate_uses_partition_columns) .and_then(|predicate| { FilePruner::try_new( Arc::clone(&predicate), &unified_file_schema, - &file, - Count::default(), + pruning_file, + predicate_creation_errors, ) }); @@ -240,8 +266,6 @@ impl FileOpener for VortexOpener { &session.arrow(), )?); - let projected_physical_schema = projection.project_schema(&unified_file_schema)?; - let expr_adapter = expr_adapter_factory.create( Arc::clone(&unified_file_schema), Arc::clone(&this_file_schema), @@ -268,7 +292,7 @@ impl FileOpener for VortexOpener { expr_convertor.split_projection( projection.clone(), &this_file_schema, - &projected_physical_schema, + output_schema.as_ref(), )? } else { // When projection pushdown is disabled, read only the required columns @@ -285,7 +309,7 @@ impl FileOpener for VortexOpener { // When projection pushdown is enabled, the scan outputs the projected columns. // When disabled, the scan outputs raw columns and the projection is applied after. let scan_reference_schema = if projection_pushdown { - projected_physical_schema + (*output_schema).clone() } else { // Build schema from the raw columns being read let column_indices = projection.column_indices(); @@ -293,7 +317,7 @@ impl FileOpener for VortexOpener { .into_iter() .map(|idx| this_file_schema.field(idx).clone()) .collect(); - Schema::new(fields) + Schema::new_with_metadata(fields, this_file_schema.metadata().clone()) }; let stream_schema = calculate_physical_schema(&scan_dtype, &scan_reference_schema, &session.arrow())?; @@ -424,40 +448,26 @@ impl FileOpener for VortexOpener { }) .into_stream() .map_err(|e| exec_datafusion_err!("Failed to create Vortex stream: {e}"))? - .map_ok(move |rb| { - // We try and slice the stream into respecting datafusion's configured batch size. - stream::iter( - (0..rb.num_rows().div_ceil(batch_size * 2)) - .flat_map(move |block_idx| { - let offset = block_idx * batch_size * 2; - - // If we have less than two batches worth of rows left, we keep them together as a single batch. - if rb.num_rows() - offset < 2 * batch_size { - let length = rb.num_rows() - offset; - [Some(rb.slice(offset, length)), None].into_iter() - } else { - let first = rb.slice(offset, batch_size); - let second = rb.slice(offset + batch_size, batch_size); - [Some(first), Some(second)].into_iter() - } - }) - .flatten() - .map(Ok), - ) - }) .map_err(move |e: VortexError| { DataFusionError::External(Box::new(e.with_context(format!( "Failed to read Vortex file: {}", file.object_meta.location )))) }) - .try_flatten() .map(move |batch| { - if projector.projection().as_ref().is_empty() { + let batch = if projector.projection().as_ref().is_empty() { batch } else { batch.and_then(|b| projector.project_batch(&b)) - } + }?; + + let (_, columns, row_count) = batch.into_parts(); + RecordBatch::try_new_with_options( + Arc::clone(&output_schema), + columns, + &RecordBatchOptions::new().with_row_count(Some(row_count)), + ) + .map_err(Into::into) }) .boxed(); @@ -564,6 +574,7 @@ fn split_midpoint_to_byte(split_range: &Range, row_count: u64, total_size: #[cfg(test)] mod tests { + use std::fmt; use std::sync::Arc; use std::sync::LazyLock; @@ -571,6 +582,7 @@ mod tests { use arrow_schema::Fields; use arrow_schema::SchemaRef; use datafusion::arrow::array::DictionaryArray; + use datafusion::arrow::array::Int32Array; use datafusion::arrow::array::RecordBatch; use datafusion::arrow::array::StringArray; use datafusion::arrow::array::StructArray; @@ -585,9 +597,12 @@ mod tests { use datafusion::physical_expr::planner::logical2physical; use datafusion::physical_expr_adapter::DefaultPhysicalExprAdapterFactory; use datafusion::scalar::ScalarValue; + use datafusion_common::stats::Precision; use datafusion_execution::cache::DefaultFilesMetadataCache; use datafusion_expr::Operator; + use datafusion_physical_expr::PhysicalExpr; use datafusion_physical_expr::expressions as df_expr; + use datafusion_physical_expr::expressions::DynamicFilterPhysicalExpr; use datafusion_physical_expr::projection::ProjectionExpr; use insta::assert_snapshot; use itertools::Itertools; @@ -612,6 +627,54 @@ mod tests { static SESSION: LazyLock = LazyLock::new(VortexSession::default); + /// Test-only expr used to test error reporting. + #[derive(Debug, Eq, Hash, PartialEq)] + struct SnapshotErrorExpr; + + impl fmt::Display for SnapshotErrorExpr { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "snapshot_error") + } + } + + impl PhysicalExpr for SnapshotErrorExpr { + fn data_type(&self, _input_schema: &Schema) -> DFResult { + Ok(DataType::Boolean) + } + + fn nullable(&self, _input_schema: &Schema) -> DFResult { + Ok(false) + } + + fn fmt_sql(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + fmt::Display::fmt(self, f) + } + + fn evaluate(&self, _batch: &RecordBatch) -> DFResult { + Err(DataFusionError::Internal( + "intentional snapshot error".to_owned(), + )) + } + + fn children(&self) -> Vec<&PhysicalExprRef> { + Vec::new() + } + + fn with_new_children( + self: Arc, + children: Vec, + ) -> DFResult { + assert!(children.is_empty()); + Ok(self) + } + + fn snapshot(&self) -> DFResult> { + Err(DataFusionError::Internal( + "intentional snapshot error".to_owned(), + )) + } + } + #[rstest] #[case(0..3, 10, vec![0..2, 2..5, 5..10], Some(0..2))] #[case(3..7, 10, vec![0..2, 2..5, 5..10], Some(2..5))] @@ -700,9 +763,9 @@ mod tests { file_pruning_predicate: None, expr_adapter_factory: Arc::new(DefaultPhysicalExprAdapterFactory), table_schema, - batch_size: 100, limit: None, metrics_registry: Arc::new(DefaultMetricsRegistry::default()), + df_metrics: ExecutionPlanMetricsSet::new(), layout_readers: Default::default(), natural_split_ranges: Default::default(), has_output_ordering: false, @@ -766,6 +829,165 @@ mod tests { Ok(()) } + #[tokio::test] + async fn test_open_preserves_declared_schema_metadata() -> anyhow::Result<()> { + let object_store = Arc::new(InMemory::new()) as Arc; + let file_path = "part=1/file.vortex"; + let batch = record_batch!(("a", Int32, vec![Some(1), Some(2), Some(3)]))?; + let data_size = + write_arrow_to_vortex(Arc::clone(&object_store), file_path, batch.clone()).await?; + + let file_schema = Arc::new( + batch.schema().as_ref().clone().with_metadata( + [("table".to_string(), "metadata".to_string())] + .into_iter() + .collect(), + ), + ); + let table_schema = TableSchema::new( + file_schema, + vec![Arc::new( + Field::new("part", DataType::Int32, false).with_metadata( + [("partition".to_string(), "metadata".to_string())] + .into_iter() + .collect(), + ), + )], + ); + let projection = ProjectionExprs::from_indices(&[0, 1], table_schema.table_schema()); + let expected_schema = Arc::new(projection.project_schema(table_schema.table_schema())?); + + assert_eq!( + expected_schema.metadata().get("table"), + Some(&"metadata".to_string()) + ); + assert_eq!( + expected_schema.field(1).metadata().get("partition"), + Some(&"metadata".to_string()) + ); + + for projection_pushdown in [false, true] { + let mut opener = make_opener(Arc::clone(&object_store), table_schema.clone(), None); + opener.projection = projection.clone(); + opener.projection_pushdown = projection_pushdown; + + let mut file = PartitionedFile::new(file_path.to_string(), data_size); + file.partition_values = vec![ScalarValue::Int32(Some(1))]; + let batches = opener.open(file)?.await?.try_collect::>().await?; + + assert!(!batches.is_empty()); + for batch in batches { + assert_eq!(batch.schema().as_ref(), expected_schema.as_ref()); + } + } + + Ok(()) + } + + #[tokio::test] + async fn test_open_all_valid_nullable_columns_with_nonnullable_table_schema() + -> anyhow::Result<()> { + let object_store = Arc::new(InMemory::new()) as Arc; + let file_path = "nullable/file.vortex"; + let batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)])), + vec![Arc::new(Int32Array::from(vec![Some(1), Some(2), Some(3)]))], + )?; + let data_size = write_arrow_to_vortex(Arc::clone(&object_store), file_path, batch).await?; + + let expected_schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let table_schema = TableSchema::from_file_schema(Arc::clone(&expected_schema)); + + for projection_pushdown in [false, true] { + let mut opener = make_opener(Arc::clone(&object_store), table_schema.clone(), None); + opener.projection_pushdown = projection_pushdown; + + let file = PartitionedFile::new(file_path.to_string(), data_size); + let batches = opener.open(file)?.await?.try_collect::>().await?; + + assert_eq!(batches.len(), 1); + assert_eq!(batches[0].schema().as_ref(), expected_schema.as_ref()); + } + + Ok(()) + } + + #[tokio::test] + async fn test_file_pruning_replaces_partition_columns_without_file_statistics() + -> anyhow::Result<()> { + let object_store = Arc::new(InMemory::new()) as Arc; + let file_schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let table_schema = TableSchema::new( + Arc::clone(&file_schema), + vec![Arc::new(Field::new("part", DataType::Int32, false))], + ); + + let partition_column = Arc::new(df_expr::Column::new("part", 1)) as PhysicalExprRef; + let predicate = Arc::new(df_expr::BinaryExpr::new( + Arc::clone(&partition_column), + Operator::Gt, + df_expr::lit(ScalarValue::Int32(Some(1))), + )) as PhysicalExprRef; + let dynamic_predicate = Arc::new(DynamicFilterPhysicalExpr::new( + vec![partition_column], + predicate, + )) as PhysicalExprRef; + + let mut opener = make_opener(object_store, table_schema, None); + opener.file_pruning_predicate = Some(dynamic_predicate); + let df_metrics = opener.df_metrics.clone(); + + // The file does not exist and has no statistics. Replacing `part` with 1 + // makes the predicate false, so pruning must happen before any file I/O. + let mut file = PartitionedFile::new("missing.vortex", 1); + file.partition_values = vec![ScalarValue::Int32(Some(1))]; + let batches = opener.open(file)?.await?.try_collect::>().await?; + + assert!(batches.is_empty()); + assert_eq!( + df_metrics + .clone_inner() + .sum_by_name("num_predicate_creation_errors") + .map(|metric| metric.as_usize()), + Some(0) + ); + + Ok(()) + } + + #[tokio::test] + async fn test_file_pruning_creation_errors_are_reported() -> anyhow::Result<()> { + let object_store = Arc::new(InMemory::new()) as Arc; + let file_path = "metrics/file.vortex"; + let batch = record_batch!(("a", Int32, vec![Some(1), Some(2), Some(3)])).unwrap(); + let data_size = + write_arrow_to_vortex(Arc::clone(&object_store), file_path, batch.clone()).await?; + let mut statistics = Statistics::new_unknown(batch.schema().as_ref()); + statistics.column_statistics[0].null_count = Precision::Exact(0); + let file = PartitionedFile::new(file_path, data_size).with_statistics(Arc::new(statistics)); + + let mut opener = make_opener( + object_store, + TableSchema::from_file_schema(batch.schema()), + None, + ); + opener.file_pruning_predicate = Some(Arc::new(SnapshotErrorExpr)); + let df_metrics = opener.df_metrics.clone(); + + let batches = opener.open(file)?.await?.try_collect::>().await?; + + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 3); + assert_eq!( + df_metrics + .clone_inner() + .sum_by_name("num_predicate_creation_errors") + .map(|metric| metric.as_usize()), + Some(1) + ); + + Ok(()) + } + #[tokio::test] async fn test_open_empty_file() -> anyhow::Result<()> { use futures::TryStreamExt; @@ -872,9 +1094,9 @@ mod tests { file_pruning_predicate: None, expr_adapter_factory: Arc::new(DefaultPhysicalExprAdapterFactory), table_schema: table_schema.clone(), - batch_size: 100, limit: None, metrics_registry: Arc::new(DefaultMetricsRegistry::default()), + df_metrics: ExecutionPlanMetricsSet::new(), layout_readers: Default::default(), natural_split_ranges: Default::default(), has_output_ordering: false, @@ -959,9 +1181,9 @@ mod tests { file_pruning_predicate: None, expr_adapter_factory: Arc::new(DefaultPhysicalExprAdapterFactory), table_schema: TableSchema::from_file_schema(Arc::clone(&table_schema)), - batch_size: 100, limit: None, metrics_registry: Arc::new(DefaultMetricsRegistry::default()), + df_metrics: ExecutionPlanMetricsSet::new(), layout_readers: Default::default(), natural_split_ranges: Default::default(), has_output_ordering: false, @@ -1116,9 +1338,9 @@ mod tests { file_pruning_predicate: None, expr_adapter_factory: Arc::new(DefaultPhysicalExprAdapterFactory), table_schema: table_schema.clone(), - batch_size: 100, limit: None, metrics_registry: Arc::new(DefaultMetricsRegistry::default()), + df_metrics: ExecutionPlanMetricsSet::new(), layout_readers: Default::default(), natural_split_ranges: Default::default(), has_output_ordering: false, @@ -1176,9 +1398,9 @@ mod tests { file_pruning_predicate: None, expr_adapter_factory: Arc::new(DefaultPhysicalExprAdapterFactory), table_schema: TableSchema::from_file_schema(schema), - batch_size: 100, limit: None, metrics_registry: Arc::new(DefaultMetricsRegistry::default()), + df_metrics: ExecutionPlanMetricsSet::new(), layout_readers: Default::default(), natural_split_ranges: Default::default(), has_output_ordering: false, @@ -1385,9 +1607,9 @@ mod tests { file_pruning_predicate: None, expr_adapter_factory: Arc::new(DefaultPhysicalExprAdapterFactory), table_schema, - batch_size: 100, limit: None, metrics_registry: Arc::new(DefaultMetricsRegistry::default()), + df_metrics: ExecutionPlanMetricsSet::new(), layout_readers: Default::default(), natural_split_ranges: Default::default(), has_output_ordering: false, diff --git a/vortex-datafusion/src/persistent/sink.rs b/vortex-datafusion/src/persistent/sink.rs index 32be0c6873b..ed655b7e5d0 100644 --- a/vortex-datafusion/src/persistent/sink.rs +++ b/vortex-datafusion/src/persistent/sink.rs @@ -20,6 +20,10 @@ use datafusion_execution::SendableRecordBatchStream; use datafusion_execution::TaskContext; use datafusion_physical_plan::DisplayAs; use datafusion_physical_plan::DisplayFormatType; +use datafusion_physical_plan::metrics::Count; +use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet; +use datafusion_physical_plan::metrics::MetricBuilder; +use datafusion_physical_plan::metrics::MetricCategory; use datafusion_physical_plan::metrics::MetricsSet; use futures::StreamExt; use object_store::ObjectStore; @@ -38,15 +42,29 @@ pub struct VortexSink { config: FileSinkConfig, schema: SchemaRef, session: VortexSession, + metrics: ExecutionPlanMetricsSet, + rows_written: Count, + bytes_written: Count, } impl VortexSink { /// Creates a new [`VortexSink`] instance. pub fn new(config: FileSinkConfig, schema: SchemaRef, session: VortexSession) -> Self { + let metrics = ExecutionPlanMetricsSet::new(); + let rows_written = MetricBuilder::new(&metrics) + .with_category(MetricCategory::Rows) + .global_counter("rows_written"); + let bytes_written = MetricBuilder::new(&metrics) + .with_category(MetricCategory::Bytes) + .global_counter("bytes_written"); + Self { config, schema, session, + metrics, + rows_written, + bytes_written, } } } @@ -75,7 +93,7 @@ impl DisplayAs for VortexSink { #[async_trait] impl DataSink for VortexSink { fn metrics(&self) -> Option { - None + Some(self.metrics.clone_inner()) } /// Returns the sink schema @@ -161,7 +179,12 @@ impl FileSink for VortexSink { Ok(r) => { let (path, summary) = r?; - row_count += summary.row_count(); + let rows = summary.row_count(); + row_count += rows; + self.rows_written + .add(usize::try_from(rows).unwrap_or(usize::MAX)); + self.bytes_written + .add(usize::try_from(summary.size()).unwrap_or(usize::MAX)); tracing::info!(path = %path, "Successfully written file"); } @@ -234,12 +257,29 @@ mod tests { ) .await?; - ctx.session + let insert = ctx + .session .sql("INSERT INTO my_tbl VALUES ('hello', 1), ('world', 2);") - .await? - .collect() + .await?; + let physical_plan = insert.create_physical_plan().await?; + datafusion_physical_plan::collect(Arc::clone(&physical_plan), ctx.session.task_ctx()) .await?; + let metrics = physical_plan + .metrics() + .ok_or_else(|| anyhow::anyhow!("Vortex insert did not expose sink metrics"))?; + assert_eq!( + metrics + .sum_by_name("rows_written") + .map(|metric| metric.as_usize()), + Some(2) + ); + assert!( + metrics + .sum_by_name("bytes_written") + .is_some_and(|metric| metric.as_usize() > 0) + ); + let batches = ctx .session .sql("SELECT * from my_tbl") @@ -482,10 +522,12 @@ mod tests { file_output_mode: FileOutputMode::SingleFile, }; - let get_sink = || VortexSink { - config: config.clone(), - schema: Arc::clone(table_schema.file_schema()), - session: session.clone(), + let get_sink = || { + VortexSink::new( + config.clone(), + Arc::clone(table_schema.file_schema()), + session.clone(), + ) }; insta::assert_snapshot!(DefaultDisplay(get_sink()).to_string(), @"VortexSink(file_groups=[])"); diff --git a/vortex-datafusion/src/persistent/source.rs b/vortex-datafusion/src/persistent/source.rs index 469759527f6..34746c0ea50 100644 --- a/vortex-datafusion/src/persistent/source.rs +++ b/vortex-datafusion/src/persistent/source.rs @@ -29,7 +29,6 @@ use datafusion_physical_plan::filter_pushdown::PushedDownPredicate; use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet; use object_store::ObjectStore; use object_store::path::Path; -use vortex::error::VortexExpect; use vortex::file::VORTEX_FILE_EXTENSION; use vortex::layout::LayoutReader; use vortex::metrics::DefaultMetricsRegistry; @@ -190,8 +189,8 @@ pub struct VortexSource { /// Subset of predicates that can be pushed down into Vortex scan operations. /// These are expressions that Vortex can efficiently evaluate during scanning. pub(crate) vortex_predicate: Option, - pub(crate) batch_size: Option, - _unused_df_metrics: ExecutionPlanMetricsSet, + /// DataFusion-native metrics exposed through `DataSourceExec`. + df_metrics: ExecutionPlanMetricsSet, /// Shared layout readers, the source only lives as long as one scan. /// /// Sharing the readers allows us to only read every layout once from the file, even across partitions. @@ -228,8 +227,7 @@ impl VortexSource { projection, full_predicate: None, vortex_predicate: None, - batch_size: None, - _unused_df_metrics: Default::default(), + df_metrics: Default::default(), layout_readers: Arc::new(DashMap::default()), natural_split_ranges: Arc::new(DashMap::default()), expression_convertor, @@ -335,10 +333,6 @@ impl VortexSource { base_config: &FileScanConfig, partition: usize, ) -> DFResult { - let batch_size = self - .batch_size - .vortex_expect("batch_size must be supplied to VortexSource"); - let expr_adapter_factory = base_config .expr_adapter_factory .clone() @@ -358,9 +352,9 @@ impl VortexSource { file_pruning_predicate: self.full_predicate.clone(), expr_adapter_factory, table_schema: self.table_schema.clone(), - batch_size, limit: base_config.limit.map(|l| l as u64), metrics_registry: Arc::clone(&self.vx_metrics_registry), + df_metrics: self.df_metrics.clone(), layout_readers: Arc::clone(&self.layout_readers), natural_split_ranges: Arc::clone(&self.natural_split_ranges), has_output_ordering: !base_config.output_ordering.is_empty() || self.ordered, @@ -388,10 +382,9 @@ impl FileSource for VortexSource { )?)) } - fn with_batch_size(&self, batch_size: usize) -> Arc { - let mut source = self.clone(); - source.batch_size = Some(batch_size); - Arc::new(source) + fn with_batch_size(&self, _batch_size: usize) -> Arc { + // DataSourceExec applies BatchSplitStream after the FileSource stream. + Arc::new(self.clone()) } fn filter(&self) -> Option> { @@ -399,7 +392,7 @@ impl FileSource for VortexSource { } fn metrics(&self) -> &ExecutionPlanMetricsSet { - &self._unused_df_metrics + &self.df_metrics } fn file_type(&self) -> &str { @@ -667,12 +660,11 @@ mod tests { inner: DefaultExpressionConvertor::default(), }) as Arc; - let mut source = VortexSource::new( + let source = VortexSource::new( TableSchema::from_file_schema(file_schema), VortexSession::default(), ) .with_expression_convertor(Arc::clone(&expression_convertor)); - source.batch_size = Some(100); let config = FileScanConfigBuilder::new( ObjectStoreUrl::local_filesystem(), From da76f5bdeae45ea48e0b07127cbf4701317cfa26 Mon Sep 17 00:00:00 2001 From: jackylee Date: Fri, 17 Jul 2026 20:53:14 +0800 Subject: [PATCH 002/498] fix(spark): reject MapType instead of emitting a childless Arrow List (#8811) `SparkToArrowSchema` mapped `MapType` to a childless `ArrowType.List`, which is an invalid Arrow list (a List must have exactly one child). Drop the branch so `MapType` falls through to the existing fail-fast `UnsupportedOperationException`, and add a test. Signed-off-by: jackylee --- .../dev/vortex/spark/write/SparkToArrowSchema.java | 4 ---- .../dev/vortex/spark/write/SparkToArrowSchemaTest.java | 10 ++++++++++ 2 files changed, 10 insertions(+), 4 deletions(-) diff --git a/java/vortex-spark/src/main/java/dev/vortex/spark/write/SparkToArrowSchema.java b/java/vortex-spark/src/main/java/dev/vortex/spark/write/SparkToArrowSchema.java index 3ed5a56e52f..c84410e45a3 100644 --- a/java/vortex-spark/src/main/java/dev/vortex/spark/write/SparkToArrowSchema.java +++ b/java/vortex-spark/src/main/java/dev/vortex/spark/write/SparkToArrowSchema.java @@ -23,7 +23,6 @@ import org.apache.spark.sql.types.FloatType; import org.apache.spark.sql.types.IntegerType; import org.apache.spark.sql.types.LongType; -import org.apache.spark.sql.types.MapType; import org.apache.spark.sql.types.ShortType; import org.apache.spark.sql.types.StringType; import org.apache.spark.sql.types.StructField; @@ -125,9 +124,6 @@ private static ArrowType convertType(DataType sparkType) { return new ArrowType.List(); } else if (sparkType instanceof StructType) { return new ArrowType.Struct(); - } else if (sparkType instanceof MapType) { - // Map is represented as List> in Arrow - return new ArrowType.List(); } else { throw new UnsupportedOperationException("Unsupported Spark type for Arrow conversion: " + sparkType.getClass().getName()); diff --git a/java/vortex-spark/src/test/java/dev/vortex/spark/write/SparkToArrowSchemaTest.java b/java/vortex-spark/src/test/java/dev/vortex/spark/write/SparkToArrowSchemaTest.java index cd9b3ca184c..2df9c5c4a50 100644 --- a/java/vortex-spark/src/test/java/dev/vortex/spark/write/SparkToArrowSchemaTest.java +++ b/java/vortex-spark/src/test/java/dev/vortex/spark/write/SparkToArrowSchemaTest.java @@ -173,4 +173,14 @@ void unsupportedTypeIsRejected() { assertThrows(UnsupportedOperationException.class, () -> SparkToArrowSchema.convert(schema)); assertTrue(e.getMessage().contains("CalendarIntervalType")); } + + @Test + @DisplayName("MapType is rejected rather than converted to a childless List") + void mapTypeIsRejected() { + StructType schema = + new StructType().add("m", DataTypes.createMapType(DataTypes.StringType, DataTypes.IntegerType)); + UnsupportedOperationException e = + assertThrows(UnsupportedOperationException.class, () -> SparkToArrowSchema.convert(schema)); + assertTrue(e.getMessage().contains("MapType")); + } } From 35133675b05f1d10abd9a3b70959bb27b7a90053 Mon Sep 17 00:00:00 2001 From: jackylee Date: Fri, 17 Jul 2026 20:58:47 +0800 Subject: [PATCH 003/498] build(java): bump Spark 3 from 3.5.8 to 3.5.9 (#8812) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Rationale for this change Spark 3.5.9 was released to Maven Central on 2026-07-12 as the latest maintenance release of the 3.5 line, containing bug fixes over 3.5.8. Keeping the Spark 3 compileOnly/test dependency current picks up those fixes for the `vortex-spark` connector's Spark 3 builds. ## What changes are included in this PR? A one-line version bump in `java/gradle/libs.versions.toml`: `spark3 = "3.5.8"` → `spark3 = "3.5.9"`. No code changes. Verified locally against 3.5.9 with `:vortex-spark_2.12:test`: 120 tests ran, all passed except `VortexDataSourceS3MockTest`, which failed only because the local Docker daemon was not running (Testcontainers infrastructure requirement, unrelated to the Spark version). ## What APIs are changed? Are there any user-facing changes? None. Spark 3.5.9 is a drop-in maintenance upgrade; no connector code or public API is touched. Signed-off-by: jackylee --- java/gradle/libs.versions.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/java/gradle/libs.versions.toml b/java/gradle/libs.versions.toml index 1135946d4ab..881e44eef9d 100644 --- a/java/gradle/libs.versions.toml +++ b/java/gradle/libs.versions.toml @@ -12,7 +12,7 @@ netty = "4.2.15.Final" nopen = "1.0.1" roaringbitmap = "1.6.14" slf4j = "2.0.18" -spark3 = "3.5.8" +spark3 = "3.5.9" spark4 = "4.1.1" s3mock = "5.0.0" testcontainers-jupiter = "1.21.4" From 94c4e96baa8ab966914fd38efe3a962c5c076d72 Mon Sep 17 00:00:00 2001 From: jackylee Date: Fri, 17 Jul 2026 21:10:12 +0800 Subject: [PATCH 004/498] build(java): bump Spark 4 from 4.1.1 to 4.1.2 (#8813) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Rationale for this change Spark 4.1.2 is the latest maintenance release of the 4.1 line on Maven Central, containing bug fixes over 4.1.1. Keeping the Spark 4 dependency current picks up those fixes for the `vortex-spark_2.13` builds. Companion to #8812, which bumps the Spark 3 line. ## What changes are included in this PR? A one-line version bump in `java/gradle/libs.versions.toml`: `spark4 = "4.1.1"` → `spark4 = "4.1.2"`. No code changes. Verified locally against 4.1.2 with `:vortex-spark_2.13:test`: 120 tests ran, all passed except `VortexDataSourceS3MockTest`, which failed only because the local Docker daemon was not running (Testcontainers infrastructure requirement, unrelated to the Spark version). ## What APIs are changed? Are there any user-facing changes? None. Spark 4.1.2 is a drop-in maintenance upgrade; no connector code or public API is touched. Signed-off-by: jackylee --- java/gradle/libs.versions.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/java/gradle/libs.versions.toml b/java/gradle/libs.versions.toml index 881e44eef9d..0b817553b96 100644 --- a/java/gradle/libs.versions.toml +++ b/java/gradle/libs.versions.toml @@ -13,7 +13,7 @@ nopen = "1.0.1" roaringbitmap = "1.6.14" slf4j = "2.0.18" spark3 = "3.5.9" -spark4 = "4.1.1" +spark4 = "4.1.2" s3mock = "5.0.0" testcontainers-jupiter = "1.21.4" From 10b81730e24a91f380306dab4dddf3c5f6cfc38b Mon Sep 17 00:00:00 2001 From: Adam Gutglick Date: Fri, 17 Jul 2026 14:18:57 +0100 Subject: [PATCH 005/498] Skip building binaries in when there are no tests (#8820) ## Rationale for this change We have many binaries in the main workspace, many of them don't contain any tests (in the final `main.rs` or equivalent file). ## What changes are included in this PR? Sets `test = false` for most binaries in the workspace, which means that running `cargo test` or `cargo nextest` won't build the final binary. I've also restructured `random-access` a bit (with claude's help) so it actually has a dedicated `lib.rs`. ## What APIs are changed? Are there any user-facing changes? None Signed-off-by: Adam Gutglick --- benchmarks/compress-bench/Cargo.toml | 7 + benchmarks/datafusion-bench/Cargo.toml | 7 + benchmarks/duckdb-bench/Cargo.toml | 7 + benchmarks/lance-bench/Cargo.toml | 7 + benchmarks/random-access-bench/Cargo.toml | 4 + benchmarks/random-access-bench/src/lib.rs | 488 +++++++++++++++++++ benchmarks/random-access-bench/src/main.rs | 518 ++------------------- vortex-bench/Cargo.toml | 4 + vortex-cuda/gpu-scan-cli/Cargo.toml | 4 + vortex-test/compat-gen/Cargo.toml | 1 + vortex-test/e2e-cuda/Cargo.toml | 4 + vortex-tui/Cargo.toml | 1 + xtask/Cargo.toml | 4 + 13 files changed, 565 insertions(+), 491 deletions(-) create mode 100644 benchmarks/random-access-bench/src/lib.rs diff --git a/benchmarks/compress-bench/Cargo.toml b/benchmarks/compress-bench/Cargo.toml index d214af23511..4e8fbc530b2 100644 --- a/benchmarks/compress-bench/Cargo.toml +++ b/benchmarks/compress-bench/Cargo.toml @@ -37,5 +37,12 @@ vortex-bench = { workspace = true } lance = ["dep:lance-bench"] unstable_encodings = ["vortex/unstable_encodings"] +[[bin]] +name = "compress-bench" +test = false + +[lib] +test = false + [lints] workspace = true diff --git a/benchmarks/datafusion-bench/Cargo.toml b/benchmarks/datafusion-bench/Cargo.toml index a195bbaa03f..e7e18e6376e 100644 --- a/benchmarks/datafusion-bench/Cargo.toml +++ b/benchmarks/datafusion-bench/Cargo.toml @@ -17,6 +17,13 @@ publish = false [package.metadata.cargo-shear] ignored = ["vortex-cuda"] +[[bin]] +name = "datafusion-bench" +test = false + +[lib] +test = false + [dependencies] anyhow = { workspace = true } clap = { workspace = true, features = ["derive"] } diff --git a/benchmarks/duckdb-bench/Cargo.toml b/benchmarks/duckdb-bench/Cargo.toml index 77c3d5280fc..609f2fe2150 100644 --- a/benchmarks/duckdb-bench/Cargo.toml +++ b/benchmarks/duckdb-bench/Cargo.toml @@ -17,6 +17,13 @@ publish = false [package.metadata.cargo-shear] ignored = ["vortex-cuda"] +[[bin]] +name = "duckdb-bench" +test = false + +[lib] +test = false + [dependencies] anyhow = { workspace = true } clap = { workspace = true, features = ["derive"] } diff --git a/benchmarks/lance-bench/Cargo.toml b/benchmarks/lance-bench/Cargo.toml index 7cab6af2119..44ae517c16a 100644 --- a/benchmarks/lance-bench/Cargo.toml +++ b/benchmarks/lance-bench/Cargo.toml @@ -32,5 +32,12 @@ vortex-bench = { workspace = true } [features] unstable_encodings = ["vortex-bench/unstable_encodings"] +[lib] +test = false + +[[bin]] +name = "lance-bench" +test = false + [lints] workspace = true diff --git a/benchmarks/random-access-bench/Cargo.toml b/benchmarks/random-access-bench/Cargo.toml index 500db9ec7c4..c45c1ac2bff 100644 --- a/benchmarks/random-access-bench/Cargo.toml +++ b/benchmarks/random-access-bench/Cargo.toml @@ -29,5 +29,9 @@ vortex-bench = { workspace = true } lance = ["dep:lance-bench"] unstable_encodings = ["vortex-bench/unstable_encodings"] +[[bin]] +name = "random-access-bench" +test = false + [lints] workspace = true diff --git a/benchmarks/random-access-bench/src/lib.rs b/benchmarks/random-access-bench/src/lib.rs new file mode 100644 index 00000000000..7d4b7d2c389 --- /dev/null +++ b/benchmarks/random-access-bench/src/lib.rs @@ -0,0 +1,488 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::path::PathBuf; +use std::time::Duration; +use std::time::Instant; + +use anyhow::Result; +use clap::ValueEnum; +use indicatif::ProgressBar; +use rand::RngExt; +use rand::SeedableRng; +use rand::rngs::StdRng; +use rand_distr::Distribution; +use rand_distr::Exp; +use vortex_bench::Engine; +use vortex_bench::Format; +use vortex_bench::Target; +use vortex_bench::create_output_writer; +use vortex_bench::display::DisplayFormat; +use vortex_bench::display::print_measurements_json; +use vortex_bench::measurements::TimingMeasurement; +use vortex_bench::random_access::BenchDataset; +use vortex_bench::random_access::ParquetRandomAccessor; +use vortex_bench::random_access::RandomAccessor; +use vortex_bench::random_access::VortexRandomAccessor; +use vortex_bench::utils::constants::STORAGE_NVME; +use vortex_bench::v3; + +use crate::render::RandomAccessRun; +use crate::render::render_random_access_table; + +mod render; + +// --------------------------------------------------------------------------- +// Access patterns +// --------------------------------------------------------------------------- + +/// Access pattern for random access benchmarks. +#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq, ValueEnum)] +pub enum AccessPattern { + /// Multiple clusters of sequential indices scattered across the dataset, + /// simulating workloads with spatial locality (e.g. scanning nearby records). + Correlated, + /// Indices generated by a Poisson process (exponential inter-arrival times) + /// spread uniformly across the dataset, simulating random lookups with no locality. + Uniform, +} + +impl AccessPattern { + fn name(&self) -> &'static str { + match self { + AccessPattern::Correlated => "correlated", + AccessPattern::Uniform => "uniform", + } + } +} + +/// Number of clusters for the correlated pattern. +const NUM_CLUSTERS: usize = 5; + +/// Number of consecutive indices per cluster. +const CLUSTER_SIZE: usize = 20; + +/// Expected number of indices for the Poisson (uniform) pattern. +const POISSON_EXPECTED_COUNT: usize = 100; + +/// Generate indices for the given dataset and access pattern. +fn generate_indices(dataset: &dyn BenchDataset, pattern: AccessPattern) -> Vec { + let row_count = dataset.row_count(); + let mut rng = StdRng::seed_from_u64(42); + + match pattern { + AccessPattern::Correlated => { + // Pick random cluster starts, then emit CLUSTER_SIZE consecutive indices from each. + let mut indices = Vec::with_capacity(NUM_CLUSTERS * CLUSTER_SIZE); + for _ in 0..NUM_CLUSTERS { + let start = rng.random_range(0..row_count.saturating_sub(CLUSTER_SIZE as u64)); + for offset in 0..CLUSTER_SIZE as u64 { + indices.push(start + offset); + } + } + indices.sort_unstable(); + indices + } + AccessPattern::Uniform => { + // Poisson process: exponential inter-arrival times with rate chosen to yield + // ~POISSON_EXPECTED_COUNT indices across the dataset. + let rate = POISSON_EXPECTED_COUNT as f64 / row_count as f64; + // SAFETY: rate is always positive (POISSON_EXPECTED_COUNT > 0, row_count > 0). + #[expect(clippy::unwrap_used)] + let exp = Exp::new(rate).unwrap(); + let mut indices = Vec::with_capacity(POISSON_EXPECTED_COUNT); + let mut pos = 0.0_f64; + loop { + let gap: f64 = exp.sample(&mut rng); + pos += gap; + #[expect(clippy::cast_possible_truncation)] + let idx = pos as u64; + if idx >= row_count { + break; + } + indices.push(idx); + } + indices + } + } +} + +// --------------------------------------------------------------------------- +// Benchmark configuration +// --------------------------------------------------------------------------- + +/// Controls whether the file handle is reused or reopened each iteration. +#[derive(ValueEnum, Clone, Copy, Debug, PartialEq, Eq)] +pub enum OpenMode { + /// Reuse the file handle across iterations (cached metadata). + #[clap(name = "cached")] + Cached, + /// Reopen the file each iteration (includes footer parsing). + #[clap(name = "reopen")] + Reopen, + /// Run both cached and reopen variants. + #[clap(name = "both")] + Both, +} + +// --------------------------------------------------------------------------- +// Benchmark core +// --------------------------------------------------------------------------- + +/// Run a random access benchmark. +/// +/// Runs the take operation repeatedly until the time limit is reached, +/// collecting timing for each run. When `reopen` is true, the accessor is +/// recreated from scratch before each iteration so that file metadata +/// parsing is included in the timing. +#[expect(clippy::too_many_arguments)] +async fn benchmark_random_access( + dataset: &dyn BenchDataset, + format: Format, + measurement_name: &str, + pattern: Option, + indices: &[u64], + time_limit_secs: u64, + storage: &str, + reopen: bool, +) -> Result { + let time_limit = Duration::from_secs(time_limit_secs); + let overall_start = Instant::now(); + let mut runs = Vec::new(); + let mut accessor = open_accessor(dataset, format).await?; + + loop { + let start = Instant::now(); + let arr = accessor.take(indices).await?; + runs.push(start.elapsed()); + drop(arr); + + if overall_start.elapsed() >= time_limit { + break; + } + + if reopen { + accessor = open_accessor(dataset, format).await?; + } + } + + let timing = TimingMeasurement { + name: measurement_name.to_string(), + storage: storage.to_string(), + target: Target::new(format_to_engine(format), format), + runs, + }; + Ok(RandomAccessRun { + display_name: display_name(dataset.name(), pattern), + dataset: dataset.name().to_string(), + pattern, + reopen, + timing, + }) +} + +/// Row label for the table view. Format is implied by the column header, so +/// it is omitted from the row label even though it stays in the +/// [`TimingMeasurement::name`] used for JSON back-compat. +fn display_name(dataset: &str, pattern: Option) -> String { + match pattern { + Some(p) => format!("random-access/{}/{}", dataset, p.name()), + None => format!("random-access/{}", dataset), + } +} + +/// Build a measurement name for a benchmark run. +/// +/// For taxi (legacy), the name is `random-access/{format}-tokio-local-disk` to preserve +/// historical continuity with existing benchmark data. +/// For other datasets, includes dataset and pattern: +/// `random-access/{dataset}/{pattern}/{format}-tokio-local-disk`. +fn measurement_name(dataset: &str, pattern: Option, format: Format) -> String { + let fmt = format.ext(); + match pattern { + Some(p) => format!( + "random-access/{}/{}/{}-tokio-local-disk", + dataset, + p.name(), + fmt + ), + None => format!("random-access/{}-tokio-local-disk", fmt), + } +} + +fn v3_random_access_dataset_name(dataset: &str, pattern: Option) -> String { + match pattern { + Some(pattern) => format!("{dataset}/{}", pattern.name()), + None => dataset.to_string(), + } +} + +fn push_v3_random_access_record(records: &mut Vec, run: &RandomAccessRun) { + if run.reopen { + return; + } + + let dataset = v3_random_access_dataset_name(&run.dataset, run.pattern); + records.push(v3::random_access_record(&run.timing, &dataset)); +} + +/// Map format to the appropriate engine for random access benchmarks. +fn format_to_engine(format: Format) -> Engine { + match format { + Format::OnDiskVortex | Format::VortexCompact => Engine::Vortex, + Format::Parquet => Engine::Arrow, + #[cfg(feature = "lance")] + Format::Lance => Engine::Arrow, // Is this right here? + _ => Engine::default(), + } +} + +/// Open a random accessor for any supported format. +/// +/// For Vortex and Parquet, the path comes from [`BenchDataset::path`]. +/// For Lance (behind the `lance` feature), the path is resolved from lance-bench helpers. +async fn open_accessor( + dataset: &dyn BenchDataset, + format: Format, +) -> Result> { + let name = format!( + "random-access/{}/{}-tokio-local-disk", + dataset.name(), + format.ext() + ); + match format { + Format::OnDiskVortex | Format::VortexCompact => { + let path = dataset.path(format).await?; + Ok(Box::new( + VortexRandomAccessor::open(path, name, format).await?, + )) + } + Format::Parquet => { + let path = dataset.path(format).await?; + Ok(Box::new(ParquetRandomAccessor::open(path, name).await?)) + } + #[cfg(feature = "lance")] + Format::Lance => { + use lance_bench::random_access; + let path = match dataset.name() { + "taxi" => random_access::taxi_data_lance().await?, + "feature-vectors" => random_access::feature_vectors_lance().await?, + "nested-lists" => random_access::nested_lists_lance().await?, + "nested-structs" => random_access::nested_structs_lance().await?, + other => anyhow::bail!("Unknown dataset for Lance: {other}"), + }; + Ok(Box::new( + random_access::LanceRandomAccessor::open(path, name).await?, + )) + } + other => unimplemented!("open_accessor not implemented for {other}"), + } +} + +/// The benchmark ID used for output path. +const BENCHMARK_ID: &str = "random-access"; + +/// Fixed indices used by the original taxi benchmark (preserved for historical continuity). +const FIXED_TAXI_INDICES: [u64; 6] = [10, 11, 12, 13, 100_000, 3_000_000]; + +/// Resolved configuration for a single random-access benchmark invocation. +pub struct RunConfig { + /// Datasets to benchmark. + pub datasets: Vec>, + /// File formats to measure for each dataset. + pub formats: Vec, + /// Index access patterns to measure. + pub patterns: Vec, + /// Maximum duration of each benchmark target in seconds. + pub time_limit: u64, + /// Whether each benchmark run reuses or reopens its file handle. + pub open_mode: OpenMode, + /// Format used to emit measurements. + pub display_format: DisplayFormat, + /// Optional path for the primary result output. + pub output_path: Option, + /// Optional path for v3 JSONL benchmark records. + pub gh_json_v3: Option, +} + +/// Run random-access benchmarks with `config`. +pub async fn run(config: RunConfig) -> Result<()> { + let RunConfig { + datasets, + formats, + patterns, + time_limit, + open_mode, + display_format, + output_path, + gh_json_v3, + } = config; + + let reopen_variants: &[bool] = match open_mode { + OpenMode::Cached => &[false], + OpenMode::Reopen => &[true], + OpenMode::Both => &[false, true], + }; + + let total_steps: usize = datasets + .iter() + .map(|d| { + let legacy_extra = if d.name() == "taxi" { formats.len() } else { 0 }; + (formats.len() * patterns.len() + legacy_extra) * reopen_variants.len() + }) + .sum(); + let progress = ProgressBar::new(total_steps as u64); + + let mut runs: Vec = Vec::new(); + let mut v3_records: Vec = Vec::new(); + + // Iteration order matters for the table renderer: row order is set by the + // first time each `(dataset, pattern)` pair is observed. + for dataset in &datasets { + for format in &formats { + if dataset.name() == "taxi" { + let name = measurement_name(dataset.name(), None, *format); + for &reopen in reopen_variants { + let bench_name = if reopen { + format!("{name}-footer") + } else { + name.clone() + }; + let run = benchmark_random_access( + dataset.as_ref(), + *format, + &bench_name, + None, + &FIXED_TAXI_INDICES, + time_limit, + STORAGE_NVME, + reopen, + ) + .await?; + + push_v3_random_access_record(&mut v3_records, &run); + runs.push(run); + progress.inc(1); + } + } + + for pattern in &patterns { + let indices = generate_indices(dataset.as_ref(), *pattern); + let name = measurement_name(dataset.name(), Some(*pattern), *format); + for &reopen in reopen_variants { + let bench_name = if reopen { + format!("{name}-footer") + } else { + name.clone() + }; + let run = benchmark_random_access( + dataset.as_ref(), + *format, + &bench_name, + Some(*pattern), + &indices, + time_limit, + STORAGE_NVME, + reopen, + ) + .await?; + + push_v3_random_access_record(&mut v3_records, &run); + runs.push(run); + progress.inc(1); + } + } + } + } + + progress.finish(); + + if let Some(path) = gh_json_v3 { + v3::write_jsonl_to_path(&path, &v3_records)?; + } + + let mut writer = create_output_writer(&display_format, output_path, BENCHMARK_ID)?; + + match display_format { + DisplayFormat::Table => { + render_random_access_table(&mut writer, &runs, &formats, reopen_variants)?; + } + DisplayFormat::GhJson => { + let timings: Vec = runs.into_iter().map(|r| r.timing).collect(); + print_measurements_json(&mut writer, timings)?; + } + } + + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn v3_random_access_dataset_names_match_schema_dims() { + assert_eq!(v3_random_access_dataset_name("taxi", None), "taxi"); + assert_eq!( + v3_random_access_dataset_name("taxi", Some(AccessPattern::Correlated)), + "taxi/correlated" + ); + assert_eq!( + v3_random_access_dataset_name("feature-vectors", Some(AccessPattern::Uniform)), + "feature-vectors/uniform" + ); + } + + fn fake_run(dataset: &str, pattern: Option, reopen: bool) -> RandomAccessRun { + RandomAccessRun { + timing: TimingMeasurement { + name: format!("random-access/{dataset}/parquet-tokio-local-disk"), + target: Target::new(Engine::Arrow, Format::Parquet), + storage: STORAGE_NVME.to_string(), + runs: vec![Duration::from_nanos(10)], + }, + dataset: dataset.to_string(), + pattern, + reopen, + display_name: display_name(dataset, pattern), + } + } + + #[test] + fn v3_random_access_records_skip_reopen_variants() { + let mut records = Vec::new(); + + push_v3_random_access_record(&mut records, &fake_run("taxi", None, false)); + push_v3_random_access_record( + &mut records, + &fake_run("taxi", Some(AccessPattern::Uniform), false), + ); + push_v3_random_access_record( + &mut records, + &fake_run("taxi", Some(AccessPattern::Correlated), true), + ); + + assert_eq!(records.len(), 2); + match &records[0] { + v3::V3Record::RandomAccessTime(record) => assert_eq!(record.dataset, "taxi"), + other => panic!("expected random-access record, got {other:?}"), + } + match &records[1] { + v3::V3Record::RandomAccessTime(record) => assert_eq!(record.dataset, "taxi/uniform"), + other => panic!("expected random-access record, got {other:?}"), + } + } + + #[test] + fn display_name_drops_format_extension() { + assert_eq!(display_name("taxi", None), "random-access/taxi"); + assert_eq!( + display_name("taxi", Some(AccessPattern::Uniform)), + "random-access/taxi/uniform" + ); + assert_eq!( + display_name("feature-vectors", Some(AccessPattern::Correlated)), + "random-access/feature-vectors/correlated" + ); + } +} diff --git a/benchmarks/random-access-bench/src/main.rs b/benchmarks/random-access-bench/src/main.rs index 6722265fde6..f3d46ef2ff6 100644 --- a/benchmarks/random-access-bench/src/main.rs +++ b/benchmarks/random-access-bench/src/main.rs @@ -2,134 +2,21 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use std::path::PathBuf; -use std::time::Duration; -use std::time::Instant; use anyhow::Result; use clap::Parser; use clap::ValueEnum; -use indicatif::ProgressBar; -use rand::RngExt; -use rand::SeedableRng; -use rand::rngs::StdRng; -use rand_distr::Distribution; -use rand_distr::Exp; -use vortex_bench::Engine; +use random_access_bench::AccessPattern; +use random_access_bench::OpenMode; +use random_access_bench::RunConfig; use vortex_bench::Format; -use vortex_bench::Target; -use vortex_bench::create_output_writer; use vortex_bench::datasets::feature_vectors::FeatureVectorsData; use vortex_bench::datasets::nested_lists::NestedListsData; use vortex_bench::datasets::nested_structs::NestedStructsData; use vortex_bench::datasets::taxi_data::TaxiData; use vortex_bench::display::DisplayFormat; -use vortex_bench::display::print_measurements_json; -use vortex_bench::measurements::TimingMeasurement; use vortex_bench::random_access::BenchDataset; -use vortex_bench::random_access::ParquetRandomAccessor; -use vortex_bench::random_access::RandomAccessor; -use vortex_bench::random_access::VortexRandomAccessor; use vortex_bench::setup_logging_and_tracing; -use vortex_bench::utils::constants::STORAGE_NVME; -use vortex_bench::v3; - -use crate::render::RandomAccessRun; -use crate::render::render_random_access_table; - -mod render; - -// --------------------------------------------------------------------------- -// Access patterns -// --------------------------------------------------------------------------- - -/// Access pattern for random access benchmarks. -#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq, ValueEnum)] -pub enum AccessPattern { - /// Multiple clusters of sequential indices scattered across the dataset, - /// simulating workloads with spatial locality (e.g. scanning nearby records). - Correlated, - /// Indices generated by a Poisson process (exponential inter-arrival times) - /// spread uniformly across the dataset, simulating random lookups with no locality. - Uniform, -} - -impl AccessPattern { - fn name(&self) -> &'static str { - match self { - AccessPattern::Correlated => "correlated", - AccessPattern::Uniform => "uniform", - } - } -} - -/// Number of clusters for the correlated pattern. -const NUM_CLUSTERS: usize = 5; - -/// Number of consecutive indices per cluster. -const CLUSTER_SIZE: usize = 20; - -/// Expected number of indices for the Poisson (uniform) pattern. -const POISSON_EXPECTED_COUNT: usize = 100; - -/// Generate indices for the given dataset and access pattern. -fn generate_indices(dataset: &dyn BenchDataset, pattern: AccessPattern) -> Vec { - let row_count = dataset.row_count(); - let mut rng = StdRng::seed_from_u64(42); - - match pattern { - AccessPattern::Correlated => { - // Pick random cluster starts, then emit CLUSTER_SIZE consecutive indices from each. - let mut indices = Vec::with_capacity(NUM_CLUSTERS * CLUSTER_SIZE); - for _ in 0..NUM_CLUSTERS { - let start = rng.random_range(0..row_count.saturating_sub(CLUSTER_SIZE as u64)); - for offset in 0..CLUSTER_SIZE as u64 { - indices.push(start + offset); - } - } - indices.sort_unstable(); - indices - } - AccessPattern::Uniform => { - // Poisson process: exponential inter-arrival times with rate chosen to yield - // ~POISSON_EXPECTED_COUNT indices across the dataset. - let rate = POISSON_EXPECTED_COUNT as f64 / row_count as f64; - // SAFETY: rate is always positive (POISSON_EXPECTED_COUNT > 0, row_count > 0). - #[expect(clippy::unwrap_used)] - let exp = Exp::new(rate).unwrap(); - let mut indices = Vec::with_capacity(POISSON_EXPECTED_COUNT); - let mut pos = 0.0_f64; - loop { - let gap: f64 = exp.sample(&mut rng); - pos += gap; - #[expect(clippy::cast_possible_truncation)] - let idx = pos as u64; - if idx >= row_count { - break; - } - indices.push(idx); - } - indices - } - } -} - -// --------------------------------------------------------------------------- -// CLI -// --------------------------------------------------------------------------- - -/// Controls whether the file handle is reused or reopened each iteration. -#[derive(ValueEnum, Clone, Copy, Debug, PartialEq, Eq)] -enum OpenMode { - /// Reuse the file handle across iterations (cached metadata). - #[clap(name = "cached")] - Cached, - /// Reopen the file each iteration (includes footer parsing). - #[clap(name = "reopen")] - Reopen, - /// Run both cached and reopen variants. - #[clap(name = "both")] - Both, -} /// Which synthetic dataset to benchmark. #[derive(ValueEnum, Clone, Copy, Debug)] @@ -147,10 +34,10 @@ enum DatasetArg { impl DatasetArg { fn into_dataset(self) -> Box { match self { - DatasetArg::Taxi => Box::new(TaxiData), - DatasetArg::FeatureVectors => Box::new(FeatureVectorsData), - DatasetArg::NestedLists => Box::new(NestedListsData), - DatasetArg::NestedStructs => Box::new(NestedStructsData), + Self::Taxi => Box::new(TaxiData), + Self::FeatureVectors => Box::new(FeatureVectorsData), + Self::NestedLists => Box::new(NestedListsData), + Self::NestedStructs => Box::new(NestedStructsData), } } } @@ -201,379 +88,28 @@ struct Args { open_mode: OpenMode, } -#[tokio::main] -async fn main() -> Result<()> { - let args = Args::parse(); - - setup_logging_and_tracing(args.verbose, args.tracing)?; - - let config = RunConfig { - datasets: args - .datasets - .into_iter() - .map(|d| d.into_dataset()) - .collect(), - formats: args.formats, - patterns: args.patterns, - time_limit: args.time_limit, - open_mode: args.open_mode, - display_format: args.display_format, - output_path: args.output_path, - gh_json_v3: args.gh_json_v3, - }; - - run_random_access(config).await -} - -// --------------------------------------------------------------------------- -// Benchmark core -// --------------------------------------------------------------------------- - -/// Run a random access benchmark. -/// -/// Runs the take operation repeatedly until the time limit is reached, -/// collecting timing for each run. When `reopen` is true, the accessor is -/// recreated from scratch before each iteration so that file metadata -/// parsing is included in the timing. -#[expect(clippy::too_many_arguments)] -async fn benchmark_random_access( - dataset: &dyn BenchDataset, - format: Format, - measurement_name: &str, - pattern: Option, - indices: &[u64], - time_limit_secs: u64, - storage: &str, - reopen: bool, -) -> Result { - let time_limit = Duration::from_secs(time_limit_secs); - let overall_start = Instant::now(); - let mut runs = Vec::new(); - let mut accessor = open_accessor(dataset, format).await?; - - loop { - let start = Instant::now(); - let arr = accessor.take(indices).await?; - runs.push(start.elapsed()); - drop(arr); - - if overall_start.elapsed() >= time_limit { - break; - } - - if reopen { - accessor = open_accessor(dataset, format).await?; +impl Args { + fn into_run_config(self) -> RunConfig { + RunConfig { + datasets: self + .datasets + .into_iter() + .map(DatasetArg::into_dataset) + .collect(), + formats: self.formats, + patterns: self.patterns, + time_limit: self.time_limit, + open_mode: self.open_mode, + display_format: self.display_format, + output_path: self.output_path, + gh_json_v3: self.gh_json_v3, } } - - let timing = TimingMeasurement { - name: measurement_name.to_string(), - storage: storage.to_string(), - target: Target::new(format_to_engine(format), format), - runs, - }; - Ok(RandomAccessRun { - display_name: display_name(dataset.name(), pattern), - dataset: dataset.name().to_string(), - pattern, - reopen, - timing, - }) } -/// Row label for the table view. Format is implied by the column header, so -/// it is omitted from the row label even though it stays in the -/// [`TimingMeasurement::name`] used for JSON back-compat. -fn display_name(dataset: &str, pattern: Option) -> String { - match pattern { - Some(p) => format!("random-access/{}/{}", dataset, p.name()), - None => format!("random-access/{}", dataset), - } -} - -/// Build a measurement name for a benchmark run. -/// -/// For taxi (legacy), the name is `random-access/{format}-tokio-local-disk` to preserve -/// historical continuity with existing benchmark data. -/// For other datasets, includes dataset and pattern: -/// `random-access/{dataset}/{pattern}/{format}-tokio-local-disk`. -fn measurement_name(dataset: &str, pattern: Option, format: Format) -> String { - let fmt = format.ext(); - match pattern { - Some(p) => format!( - "random-access/{}/{}/{}-tokio-local-disk", - dataset, - p.name(), - fmt - ), - None => format!("random-access/{}-tokio-local-disk", fmt), - } -} - -fn v3_random_access_dataset_name(dataset: &str, pattern: Option) -> String { - match pattern { - Some(pattern) => format!("{dataset}/{}", pattern.name()), - None => dataset.to_string(), - } -} - -fn push_v3_random_access_record(records: &mut Vec, run: &RandomAccessRun) { - if run.reopen { - return; - } - - let dataset = v3_random_access_dataset_name(&run.dataset, run.pattern); - records.push(v3::random_access_record(&run.timing, &dataset)); -} - -/// Map format to the appropriate engine for random access benchmarks. -fn format_to_engine(format: Format) -> Engine { - match format { - Format::OnDiskVortex | Format::VortexCompact => Engine::Vortex, - Format::Parquet => Engine::Arrow, - #[cfg(feature = "lance")] - Format::Lance => Engine::Arrow, // Is this right here? - _ => Engine::default(), - } -} - -/// Open a random accessor for any supported format. -/// -/// For Vortex and Parquet, the path comes from [`BenchDataset::path`]. -/// For Lance (behind the `lance` feature), the path is resolved from lance-bench helpers. -async fn open_accessor( - dataset: &dyn BenchDataset, - format: Format, -) -> Result> { - let name = format!( - "random-access/{}/{}-tokio-local-disk", - dataset.name(), - format.ext() - ); - match format { - Format::OnDiskVortex | Format::VortexCompact => { - let path = dataset.path(format).await?; - Ok(Box::new( - VortexRandomAccessor::open(path, name, format).await?, - )) - } - Format::Parquet => { - let path = dataset.path(format).await?; - Ok(Box::new(ParquetRandomAccessor::open(path, name).await?)) - } - #[cfg(feature = "lance")] - Format::Lance => { - use lance_bench::random_access; - let path = match dataset.name() { - "taxi" => random_access::taxi_data_lance().await?, - "feature-vectors" => random_access::feature_vectors_lance().await?, - "nested-lists" => random_access::nested_lists_lance().await?, - "nested-structs" => random_access::nested_structs_lance().await?, - other => anyhow::bail!("Unknown dataset for Lance: {other}"), - }; - Ok(Box::new( - random_access::LanceRandomAccessor::open(path, name).await?, - )) - } - other => unimplemented!("open_accessor not implemented for {other}"), - } -} - -/// The benchmark ID used for output path. -const BENCHMARK_ID: &str = "random-access"; - -/// Fixed indices used by the original taxi benchmark (preserved for historical continuity). -const FIXED_TAXI_INDICES: [u64; 6] = [10, 11, 12, 13, 100_000, 3_000_000]; - -/// Resolved configuration for a single random-access benchmark invocation. -struct RunConfig { - datasets: Vec>, - formats: Vec, - patterns: Vec, - time_limit: u64, - open_mode: OpenMode, - display_format: DisplayFormat, - output_path: Option, - gh_json_v3: Option, -} - -async fn run_random_access(config: RunConfig) -> Result<()> { - let RunConfig { - datasets, - formats, - patterns, - time_limit, - open_mode, - display_format, - output_path, - gh_json_v3, - } = config; - - let reopen_variants: &[bool] = match open_mode { - OpenMode::Cached => &[false], - OpenMode::Reopen => &[true], - OpenMode::Both => &[false, true], - }; - - let total_steps: usize = datasets - .iter() - .map(|d| { - let legacy_extra = if d.name() == "taxi" { formats.len() } else { 0 }; - (formats.len() * patterns.len() + legacy_extra) * reopen_variants.len() - }) - .sum(); - let progress = ProgressBar::new(total_steps as u64); - - let mut runs: Vec = Vec::new(); - let mut v3_records: Vec = Vec::new(); - - // Iteration order matters for the table renderer: row order is set by the - // first time each `(dataset, pattern)` pair is observed. - for dataset in &datasets { - for format in &formats { - if dataset.name() == "taxi" { - let name = measurement_name(dataset.name(), None, *format); - for &reopen in reopen_variants { - let bench_name = if reopen { - format!("{name}-footer") - } else { - name.clone() - }; - let run = benchmark_random_access( - dataset.as_ref(), - *format, - &bench_name, - None, - &FIXED_TAXI_INDICES, - time_limit, - STORAGE_NVME, - reopen, - ) - .await?; - - push_v3_random_access_record(&mut v3_records, &run); - runs.push(run); - progress.inc(1); - } - } - - for pattern in &patterns { - let indices = generate_indices(dataset.as_ref(), *pattern); - let name = measurement_name(dataset.name(), Some(*pattern), *format); - for &reopen in reopen_variants { - let bench_name = if reopen { - format!("{name}-footer") - } else { - name.clone() - }; - let run = benchmark_random_access( - dataset.as_ref(), - *format, - &bench_name, - Some(*pattern), - &indices, - time_limit, - STORAGE_NVME, - reopen, - ) - .await?; - - push_v3_random_access_record(&mut v3_records, &run); - runs.push(run); - progress.inc(1); - } - } - } - } - - progress.finish(); - - if let Some(path) = gh_json_v3 { - v3::write_jsonl_to_path(&path, &v3_records)?; - } - - let mut writer = create_output_writer(&display_format, output_path, BENCHMARK_ID)?; - - match display_format { - DisplayFormat::Table => { - render_random_access_table(&mut writer, &runs, &formats, reopen_variants)?; - } - DisplayFormat::GhJson => { - let timings: Vec = runs.into_iter().map(|r| r.timing).collect(); - print_measurements_json(&mut writer, timings)?; - } - } - - Ok(()) -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn v3_random_access_dataset_names_match_schema_dims() { - assert_eq!(v3_random_access_dataset_name("taxi", None), "taxi"); - assert_eq!( - v3_random_access_dataset_name("taxi", Some(AccessPattern::Correlated)), - "taxi/correlated" - ); - assert_eq!( - v3_random_access_dataset_name("feature-vectors", Some(AccessPattern::Uniform)), - "feature-vectors/uniform" - ); - } - - fn fake_run(dataset: &str, pattern: Option, reopen: bool) -> RandomAccessRun { - RandomAccessRun { - timing: TimingMeasurement { - name: format!("random-access/{dataset}/parquet-tokio-local-disk"), - target: Target::new(Engine::Arrow, Format::Parquet), - storage: STORAGE_NVME.to_string(), - runs: vec![Duration::from_nanos(10)], - }, - dataset: dataset.to_string(), - pattern, - reopen, - display_name: display_name(dataset, pattern), - } - } - - #[test] - fn v3_random_access_records_skip_reopen_variants() { - let mut records = Vec::new(); - - push_v3_random_access_record(&mut records, &fake_run("taxi", None, false)); - push_v3_random_access_record( - &mut records, - &fake_run("taxi", Some(AccessPattern::Uniform), false), - ); - push_v3_random_access_record( - &mut records, - &fake_run("taxi", Some(AccessPattern::Correlated), true), - ); - - assert_eq!(records.len(), 2); - match &records[0] { - v3::V3Record::RandomAccessTime(record) => assert_eq!(record.dataset, "taxi"), - other => panic!("expected random-access record, got {other:?}"), - } - match &records[1] { - v3::V3Record::RandomAccessTime(record) => assert_eq!(record.dataset, "taxi/uniform"), - other => panic!("expected random-access record, got {other:?}"), - } - } - - #[test] - fn display_name_drops_format_extension() { - assert_eq!(display_name("taxi", None), "random-access/taxi"); - assert_eq!( - display_name("taxi", Some(AccessPattern::Uniform)), - "random-access/taxi/uniform" - ); - assert_eq!( - display_name("feature-vectors", Some(AccessPattern::Correlated)), - "random-access/feature-vectors/correlated" - ); - } +#[tokio::main] +async fn main() -> Result<()> { + let args = Args::parse(); + setup_logging_and_tracing(args.verbose, args.tracing)?; + random_access_bench::run(args.into_run_config()).await } diff --git a/vortex-bench/Cargo.toml b/vortex-bench/Cargo.toml index f228e1571d3..02f043b8e68 100644 --- a/vortex-bench/Cargo.toml +++ b/vortex-bench/Cargo.toml @@ -13,6 +13,10 @@ repository = { workspace = true } rust-version = { workspace = true } version = { workspace = true } +[[bin]] +name = "data-gen" +test = false + [lints] workspace = true diff --git a/vortex-cuda/gpu-scan-cli/Cargo.toml b/vortex-cuda/gpu-scan-cli/Cargo.toml index 9234d303952..468158a35ba 100644 --- a/vortex-cuda/gpu-scan-cli/Cargo.toml +++ b/vortex-cuda/gpu-scan-cli/Cargo.toml @@ -12,6 +12,10 @@ repository = { workspace = true } rust-version = { workspace = true } version = { workspace = true } +[[bin]] +name = "gpu-scan-cli" +test = false + [lints] workspace = true diff --git a/vortex-test/compat-gen/Cargo.toml b/vortex-test/compat-gen/Cargo.toml index 4fbd90a6273..4a62aca3671 100644 --- a/vortex-test/compat-gen/Cargo.toml +++ b/vortex-test/compat-gen/Cargo.toml @@ -18,6 +18,7 @@ workspace = true [[bin]] name = "vortex-compat" path = "src/main.rs" +test = false [dependencies] # Vortex crates diff --git a/vortex-test/e2e-cuda/Cargo.toml b/vortex-test/e2e-cuda/Cargo.toml index 4ef30925b11..3985b9d11e5 100644 --- a/vortex-test/e2e-cuda/Cargo.toml +++ b/vortex-test/e2e-cuda/Cargo.toml @@ -15,6 +15,10 @@ version = { workspace = true } [lib] crate-type = ["cdylib"] +[[bin]] +name = "cudf_harness_runner" +test = false + [lints] workspace = true diff --git a/vortex-tui/Cargo.toml b/vortex-tui/Cargo.toml index 34b90ecfc05..09c929796d7 100644 --- a/vortex-tui/Cargo.toml +++ b/vortex-tui/Cargo.toml @@ -39,6 +39,7 @@ crate-type = ["cdylib", "rlib"] name = "vx" path = "src/main.rs" required-features = ["native"] +test = false [dependencies] # Shared dependencies diff --git a/xtask/Cargo.toml b/xtask/Cargo.toml index ebbc6987d46..eae2db43413 100644 --- a/xtask/Cargo.toml +++ b/xtask/Cargo.toml @@ -15,6 +15,10 @@ repository = { workspace = true } rust-version = { workspace = true } version = { workspace = true } +[[bin]] +name = "xtask" +test = false + [dependencies] anyhow = { workspace = true } clap = { workspace = true, features = ["derive"] } From 928738ed972defba69e56b114caa41f0bbb99e94 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Fri, 17 Jul 2026 15:16:00 +0100 Subject: [PATCH 006/498] LayoutEncodingAdapter::build doesn't need to panic, it can error (#8824) Those panics look like an omission --- vortex-layout/src/encoding.rs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/vortex-layout/src/encoding.rs b/vortex-layout/src/encoding.rs index 90d33eecb5e..b6e10b3b38a 100644 --- a/vortex-layout/src/encoding.rs +++ b/vortex-layout/src/encoding.rs @@ -11,7 +11,7 @@ use vortex_array::DeserializeMetadata; use vortex_array::dtype::DType; use vortex_error::VortexExpect; use vortex_error::VortexResult; -use vortex_error::vortex_panic; +use vortex_error::vortex_bail; use vortex_session::registry::Id; use crate::IntoLayout; @@ -88,14 +88,14 @@ impl LayoutEncoding for LayoutEncodingAdapter { // Validate that the builder function returned the expected values. if layout.row_count() != row_count { - vortex_panic!( + vortex_bail!( "Layout row count mismatch: {} != {}", layout.row_count(), row_count ); } if layout.dtype() != dtype { - vortex_panic!("Layout dtype mismatch: {} != {}", layout.dtype(), dtype); + vortex_bail!("Layout dtype mismatch: {} != {}", layout.dtype(), dtype); } Ok(layout.into_layout()) From cc6803a7e4e47de6a1389b9740f9da40ab640878 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Fri, 17 Jul 2026 10:58:09 -0400 Subject: [PATCH 007/498] Fix nested scalar hashing (#8804) Fixes https://github.com/vortex-data/vortex/issues/8744 `Scalar` equality ignores dtype nullability recursively, but hashing only removed top-level nullability. The typed scalar views had the same mismatch, so equal nested values could produce different hashes. - adds recursive nullability-insensitive dtype hashing matching `DType::eq_ignore_nullability` - uses it for `Scalar`, `ListScalar`, `StructScalar`, and `ExtScalar` - covers nested list, struct, Union dtype, and extension storage nullability regressions This is kind of ugly but I don't think there is a better way to do this --- vortex-array/src/dtype/dtype_impl.rs | 79 +++++++++++++++---- vortex-array/src/dtype/extension/erased.rs | 8 ++ vortex-array/src/dtype/struct_.rs | 26 +++++- vortex-array/src/dtype/union.rs | 26 +++++- vortex-array/src/scalar/scalar_impl.rs | 2 +- vortex-array/src/scalar/tests/primitives.rs | 26 ++++++ .../src/scalar/typed_view/extension/mod.rs | 2 +- .../src/scalar/typed_view/extension/tests.rs | 26 ++++++ vortex-array/src/scalar/typed_view/list.rs | 33 ++++++-- vortex-array/src/scalar/typed_view/struct_.rs | 42 +++++++++- 10 files changed, 245 insertions(+), 25 deletions(-) diff --git a/vortex-array/src/dtype/dtype_impl.rs b/vortex-array/src/dtype/dtype_impl.rs index 20ff3d6a182..52408db66b7 100644 --- a/vortex-array/src/dtype/dtype_impl.rs +++ b/vortex-array/src/dtype/dtype_impl.rs @@ -3,6 +3,8 @@ use std::fmt::Display; use std::fmt::Formatter; +use std::hash::Hash; +use std::hash::Hasher; use std::sync::Arc; use DType::*; @@ -120,21 +122,9 @@ impl DType { lhs_size == rhs_size && lhs_dtype.eq_ignore_nullability(rhs_dtype) } (Struct(lhs_dtype, _), Struct(rhs_dtype, _)) => { - (lhs_dtype.names() == rhs_dtype.names()) - && (lhs_dtype - .fields() - .zip_eq(rhs_dtype.fields()) - .all(|(l, r)| l.eq_ignore_nullability(&r))) - } - (Union(lhs, _), Union(rhs, _)) => { - // Equal `names` implies equal length by FieldNames equality. - lhs.names() == rhs.names() - && lhs.type_ids() == rhs.type_ids() - && lhs - .variants() - .zip_eq(rhs.variants()) - .all(|(l, r)| l.eq_ignore_nullability(&r)) + lhs_dtype.eq_ignore_nullability(rhs_dtype) } + (Union(lhs, _), Union(rhs, _)) => lhs.eq_ignore_nullability(rhs), (Variant(_), Variant(_)) => true, (Extension(lhs_extdtype), Extension(rhs_extdtype)) => { lhs_extdtype.eq_ignore_nullability(rhs_extdtype) @@ -143,6 +133,25 @@ impl DType { } } + /// Hash this dtype using the same equivalence relation as [`Self::eq_ignore_nullability`]. + pub(crate) fn hash_ignore_nullability(&self, state: &mut H) { + std::mem::discriminant(self).hash(state); + + match self { + Null | Bool(_) | Utf8(_) | Binary(_) | Variant(_) => {} + Primitive(ptype, _) => ptype.hash(state), + Decimal(decimal, _) => decimal.hash(state), + List(element, _) => element.hash_ignore_nullability(state), + FixedSizeList(element, size, _) => { + element.hash_ignore_nullability(state); + size.hash(state); + } + Struct(fields, _) => fields.hash_ignore_nullability(state), + Union(variants, _) => variants.hash_ignore_nullability(state), + Extension(ext) => ext.hash_ignore_nullability(state), + } + } + /// Returns `true` if `self` is a subset type of `other, otherwise `false`. /// /// If `self` is nullable, this means that the other `DType` must also be nullable (since a @@ -514,18 +523,29 @@ impl Display for DType { #[cfg(test)] mod tests { + use std::collections::hash_map::DefaultHasher; + use std::hash::Hasher; use std::sync::Arc; + use vortex_error::VortexResult; + use crate::dtype::DType; use crate::dtype::Nullability::NonNullable; use crate::dtype::Nullability::Nullable; use crate::dtype::PType; + use crate::dtype::UnionVariants; use crate::dtype::decimal::DecimalDType; use crate::extension::datetime::Date; use crate::extension::datetime::Time; use crate::extension::datetime::TimeUnit; use crate::extension::datetime::Timestamp; + fn hash_ignore_nullability(dtype: &DType) -> u64 { + let mut hasher = DefaultHasher::new(); + dtype.hash_ignore_nullability(&mut hasher); + hasher.finish() + } + #[test] fn test_ext_dtype_eq_ignore_nullability() { let d1 = DType::Extension(Time::new(TimeUnit::Seconds, Nullable).erased()); @@ -541,6 +561,37 @@ mod tests { assert!(!t1.eq_ignore_nullability(&t2)); } + #[test] + fn test_union_dtype_hash_ignores_variant_nullability() -> VortexResult<()> { + let lhs = DType::Union( + UnionVariants::try_new( + ["int", "string"].into(), + vec![ + DType::Primitive(PType::I32, Nullable), + DType::Utf8(NonNullable), + ], + vec![5, 9], + )?, + NonNullable, + ); + let rhs = DType::Union( + UnionVariants::try_new( + ["int", "string"].into(), + vec![ + DType::Primitive(PType::I32, NonNullable), + DType::Utf8(Nullable), + ], + vec![5, 9], + )?, + NonNullable, + ); + + assert!(lhs.eq_ignore_nullability(&rhs)); + assert_eq!(hash_ignore_nullability(&lhs), hash_ignore_nullability(&rhs)); + + Ok(()) + } + #[test] fn element_size_null() { assert_eq!(DType::Null.element_size(), Some(0)); diff --git a/vortex-array/src/dtype/extension/erased.rs b/vortex-array/src/dtype/extension/erased.rs index c84b198516c..a80202ba4e1 100644 --- a/vortex-array/src/dtype/extension/erased.rs +++ b/vortex-array/src/dtype/extension/erased.rs @@ -75,6 +75,14 @@ impl ExtDTypeRef { .eq_ignore_nullability(other.storage_dtype()) } + /// Hash this extension dtype using the same equivalence relation as + /// [`Self::eq_ignore_nullability`]. + pub(crate) fn hash_ignore_nullability(&self, state: &mut H) { + self.id().hash(state); + self.0.metadata_hash(state); + self.storage_dtype().hash_ignore_nullability(state); + } + // TODO(connor): We should add a different type that returns something that can be serialized. /// Serialize the metadata into a byte vector. pub fn serialize_metadata(&self) -> VortexResult> { diff --git a/vortex-array/src/dtype/struct_.rs b/vortex-array/src/dtype/struct_.rs index 4d276c35205..95f73295d52 100644 --- a/vortex-array/src/dtype/struct_.rs +++ b/vortex-array/src/dtype/struct_.rs @@ -4,6 +4,7 @@ use std::fmt::Display; use std::fmt::Formatter; use std::hash::Hash; +use std::hash::Hasher; use std::sync::Arc; use std::sync::OnceLock; @@ -83,7 +84,7 @@ impl PartialEq for FieldDTypeInner { impl Eq for FieldDTypeInner {} impl Hash for FieldDTypeInner { - fn hash(&self, state: &mut H) { + fn hash(&self, state: &mut H) { match self { FieldDTypeInner::Owned(owned) => { owned.hash(state); @@ -255,12 +256,33 @@ impl PartialEq for StructFieldsInner { impl Eq for StructFieldsInner {} impl Hash for StructFieldsInner { - fn hash(&self, state: &mut H) { + fn hash(&self, state: &mut H) { self.names.hash(state); self.dtypes.hash(state); } } +impl StructFields { + /// Check if these struct fields are equal, ignoring field dtype nullability recursively. + pub fn eq_ignore_nullability(&self, other: &Self) -> bool { + Arc::ptr_eq(&self.0, &other.0) + || (self.0.names == other.0.names + && self + .fields() + .zip_eq(other.fields()) + .all(|(lhs, rhs)| lhs.eq_ignore_nullability(&rhs))) + } + + /// Hash these struct fields using the same equivalence relation as + /// [`Self::eq_ignore_nullability`]. + pub(crate) fn hash_ignore_nullability(&self, state: &mut H) { + self.0.names.hash(state); + for field in self.fields() { + field.hash_ignore_nullability(state); + } + } +} + impl Default for StructFields { fn default() -> Self { Self::empty() diff --git a/vortex-array/src/dtype/union.rs b/vortex-array/src/dtype/union.rs index 7ce53155adc..6074e32e87a 100644 --- a/vortex-array/src/dtype/union.rs +++ b/vortex-array/src/dtype/union.rs @@ -3,6 +3,7 @@ use std::fmt; use std::hash::Hash; +use std::hash::Hasher; use std::sync::Arc; use itertools::Itertools; @@ -133,13 +134,36 @@ impl PartialEq for UnionVariantsInner { impl Eq for UnionVariantsInner {} impl Hash for UnionVariantsInner { - fn hash(&self, state: &mut H) { + fn hash(&self, state: &mut H) { self.names.hash(state); self.dtypes.hash(state); self.type_ids.hash(state); } } +impl UnionVariants { + /// Check if these union variants are equal, ignoring variant dtype nullability recursively. + pub fn eq_ignore_nullability(&self, other: &Self) -> bool { + Arc::ptr_eq(&self.0, &other.0) + || (self.0.names == other.0.names + && self + .variants() + .zip_eq(other.variants()) + .all(|(lhs, rhs)| lhs.eq_ignore_nullability(&rhs)) + && self.0.type_ids == other.0.type_ids) + } + + /// Hash these union variants using the same equivalence relation as + /// [`Self::eq_ignore_nullability`]. + pub(crate) fn hash_ignore_nullability(&self, state: &mut H) { + self.0.names.hash(state); + for variant in self.variants() { + variant.hash_ignore_nullability(state); + } + self.0.type_ids.hash(state); + } +} + impl Default for UnionVariants { fn default() -> Self { Self::empty() diff --git a/vortex-array/src/scalar/scalar_impl.rs b/vortex-array/src/scalar/scalar_impl.rs index 60a675e42b9..edbb1fac76e 100644 --- a/vortex-array/src/scalar/scalar_impl.rs +++ b/vortex-array/src/scalar/scalar_impl.rs @@ -282,7 +282,7 @@ impl Scalar { /// values have equal hashes. impl Hash for Scalar { fn hash(&self, state: &mut H) { - self.dtype.as_nonnullable().hash(state); + self.dtype.hash_ignore_nullability(state); self.value.hash(state); } } diff --git a/vortex-array/src/scalar/tests/primitives.rs b/vortex-array/src/scalar/tests/primitives.rs index d2e32ce6ef0..5559b505533 100644 --- a/vortex-array/src/scalar/tests/primitives.rs +++ b/vortex-array/src/scalar/tests/primitives.rs @@ -5,6 +5,9 @@ #[cfg(test)] mod tests { + use std::collections::hash_map::DefaultHasher; + use std::hash::Hash; + use std::hash::Hasher; use std::sync::Arc; use vortex_buffer::ByteBuffer; @@ -24,6 +27,12 @@ mod tests { use crate::scalar::Scalar; use crate::scalar::ScalarValue; + fn scalar_hash(scalar: &Scalar) -> u64 { + let mut hasher = DefaultHasher::new(); + scalar.hash(&mut hasher); + hasher.finish() + } + #[test] fn default_value_for_complex_dtype() { let struct_dtype = DType::struct_( @@ -412,6 +421,23 @@ mod tests { assert_eq!(set.len(), 5); } + #[test] + fn test_scalar_hash_ignores_nested_nullability() { + let nullable = Scalar::list( + DType::Primitive(PType::I32, Nullability::Nullable), + vec![Scalar::primitive(42_i32, Nullability::Nullable)], + Nullability::NonNullable, + ); + let non_nullable = Scalar::list( + DType::Primitive(PType::I32, Nullability::NonNullable), + vec![Scalar::primitive(42_i32, Nullability::NonNullable)], + Nullability::NonNullable, + ); + + assert_eq!(nullable, non_nullable); + assert_eq!(scalar_hash(&nullable), scalar_hash(&non_nullable)); + } + #[test] fn test_scalar_partial_ord_incompatible_types() { let int_scalar = Scalar::primitive(42i32, Nullability::NonNullable); diff --git a/vortex-array/src/scalar/typed_view/extension/mod.rs b/vortex-array/src/scalar/typed_view/extension/mod.rs index 3508883fc92..b9f524c7077 100644 --- a/vortex-array/src/scalar/typed_view/extension/mod.rs +++ b/vortex-array/src/scalar/typed_view/extension/mod.rs @@ -139,7 +139,7 @@ impl PartialOrd for ExtScalar<'_> { impl Hash for ExtScalar<'_> { fn hash(&self, state: &mut H) { - self.ext_dtype.hash(state); + self.ext_dtype.hash_ignore_nullability(state); self.to_storage_scalar().hash(state); } } diff --git a/vortex-array/src/scalar/typed_view/extension/tests.rs b/vortex-array/src/scalar/typed_view/extension/tests.rs index 2b08c530af6..d9411806db6 100644 --- a/vortex-array/src/scalar/typed_view/extension/tests.rs +++ b/vortex-array/src/scalar/typed_view/extension/tests.rs @@ -1,6 +1,10 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::collections::hash_map::DefaultHasher; +use std::hash::Hash; +use std::hash::Hasher; + use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -174,6 +178,28 @@ fn test_ext_scalar_hash() { assert_eq!(set.len(), 2); } +#[test] +fn test_ext_scalar_hash_ignores_storage_nullability() { + let nullable = Scalar::extension::( + EmptyMetadata, + Scalar::primitive(42_i32, Nullability::Nullable), + ); + let non_nullable = Scalar::extension::( + EmptyMetadata, + Scalar::primitive(42_i32, Nullability::NonNullable), + ); + let nullable = nullable.as_extension(); + let non_nullable = non_nullable.as_extension(); + + assert_eq!(nullable, non_nullable); + + let mut nullable_hasher = DefaultHasher::new(); + nullable.hash(&mut nullable_hasher); + let mut non_nullable_hasher = DefaultHasher::new(); + non_nullable.hash(&mut non_nullable_hasher); + assert_eq!(nullable_hasher.finish(), non_nullable_hasher.finish()); +} + #[test] fn test_ext_scalar_storage() { let storage_scalar = Scalar::primitive(42i32, Nullability::NonNullable); diff --git a/vortex-array/src/scalar/typed_view/list.rs b/vortex-array/src/scalar/typed_view/list.rs index 15de731089a..f97857c92ed 100644 --- a/vortex-array/src/scalar/typed_view/list.rs +++ b/vortex-array/src/scalar/typed_view/list.rs @@ -90,7 +90,7 @@ impl PartialOrd for ListScalar<'_> { impl Hash for ListScalar<'_> { fn hash(&self, state: &mut H) { - self.dtype.hash(state); + self.dtype.hash_ignore_nullability(state); self.elements().hash(state); } } @@ -225,6 +225,9 @@ impl<'a> ListScalar<'a> { #[cfg(test)] mod tests { + use std::collections::hash_map::DefaultHasher; + use std::hash::Hash; + use std::hash::Hasher; use std::sync::Arc; use super::*; @@ -422,10 +425,6 @@ mod tests { #[test] fn test_list_hash() { - use std::collections::hash_map::DefaultHasher; - use std::hash::Hash; - use std::hash::Hasher; - let element_dtype = Arc::new(DType::Primitive(PType::I32, Nullability::NonNullable)); let children = vec![ Scalar::primitive(1i32, Nullability::NonNullable), @@ -446,6 +445,30 @@ mod tests { assert_eq!(hash1, hash2); } + #[test] + fn test_list_hash_ignores_nested_nullability() { + let nullable = Scalar::list( + DType::Primitive(PType::I32, Nullability::Nullable), + vec![Scalar::primitive(42_i32, Nullability::Nullable)], + Nullability::NonNullable, + ); + let non_nullable = Scalar::list( + DType::Primitive(PType::I32, Nullability::NonNullable), + vec![Scalar::primitive(42_i32, Nullability::NonNullable)], + Nullability::NonNullable, + ); + let nullable = nullable.as_list(); + let non_nullable = non_nullable.as_list(); + + assert_eq!(nullable, non_nullable); + + let mut nullable_hasher = DefaultHasher::new(); + nullable.hash(&mut nullable_hasher); + let mut non_nullable_hasher = DefaultHasher::new(); + non_nullable.hash(&mut non_nullable_hasher); + assert_eq!(nullable_hasher.finish(), non_nullable_hasher.finish()); + } + #[test] fn test_vec_conversion() { let element_dtype = Arc::new(DType::Primitive(PType::I32, Nullability::NonNullable)); diff --git a/vortex-array/src/scalar/typed_view/struct_.rs b/vortex-array/src/scalar/typed_view/struct_.rs index d8069bab77c..ead6bf332f7 100644 --- a/vortex-array/src/scalar/typed_view/struct_.rs +++ b/vortex-array/src/scalar/typed_view/struct_.rs @@ -103,7 +103,7 @@ impl PartialOrd for StructScalar<'_> { impl Hash for StructScalar<'_> { fn hash(&self, state: &mut H) { - self.dtype.hash(state); + self.dtype.hash_ignore_nullability(state); if let Some(fields) = self.fields_iter() { for f in fields { f.hash(state); @@ -329,6 +329,10 @@ impl Scalar { #[cfg(test)] mod tests { + use std::collections::hash_map::DefaultHasher; + use std::hash::Hash; + use std::hash::Hasher; + use super::*; use crate::dtype::DType; use crate::dtype::Nullability; @@ -614,6 +618,42 @@ mod tests { assert_ne!(scalar1.as_struct(), scalar3.as_struct()); } + #[test] + fn test_struct_hash_ignores_nested_nullability() { + let nullable_dtype = DType::Struct( + StructFields::new( + ["value"].into(), + vec![DType::Primitive(I32, Nullability::Nullable)], + ), + Nullability::NonNullable, + ); + let non_nullable_dtype = DType::Struct( + StructFields::new( + ["value"].into(), + vec![DType::Primitive(I32, Nullability::NonNullable)], + ), + Nullability::NonNullable, + ); + let nullable = Scalar::struct_( + nullable_dtype, + [Scalar::primitive(42_i32, Nullability::Nullable)], + ); + let non_nullable = Scalar::struct_( + non_nullable_dtype, + [Scalar::primitive(42_i32, Nullability::NonNullable)], + ); + let nullable = nullable.as_struct(); + let non_nullable = non_nullable.as_struct(); + + assert_eq!(nullable, non_nullable); + + let mut nullable_hasher = DefaultHasher::new(); + nullable.hash(&mut nullable_hasher); + let mut non_nullable_hasher = DefaultHasher::new(); + non_nullable.hash(&mut non_nullable_hasher); + assert_eq!(nullable_hasher.finish(), non_nullable_hasher.finish()); + } + #[test] fn test_struct_partial_ord() { let (_, _, dtype) = setup_types(); From 7cf53cb924cb043eff9405bc7839351d471cd8bf Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Fri, 17 Jul 2026 16:32:05 +0100 Subject: [PATCH 008/498] Add documentation links to benchmark suite PR comments (#8815) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Rationale for this change It's hard to know what each CI benchmark suite actually measures or where it is defined. Every suite should ship a short explainer doc, and the recurring `# Benchmarks: ` PR comments should link to it. ## What changes are included in this PR? 1. **A markdown explainer per benchmark suite**, living next to the suite definition: `vortex-bench/sql/{tpch,tpcds,appian,vortex}/README.md`, `vortex-bench/sql/{clickbench,fineweb,statpopgen,polarsignals,spatialbench,gharchive,public-bi}.md`, and `benchmarks/{random-access,compress}-bench/README.md`. Each covers what the suite measures, the dataset, the CI variants, and how to run it locally. 2. **The Rust benchmark code is the source of truth for the link.** `Benchmark::doc_path()` is a required trait method; the runner stamps it as a `doc` field on every results.json row (the standalone `random-access-bench`/`compress-bench` binaries do the same via a `DOC_PATH` const). 3. **`scripts/compare-benchmark-jsons.py` renders the comment title** as `# Benchmarks: [📖]()` from that field, replacing the `echo` lines in `bench-pr.yml`/`sql-benchmarks.yml`. Per review feedback, the link pins the results' commit SHA rather than `develop`, so it resolves before the PR merges and stays valid afterwards. ## What APIs are changed? Are there any user-facing changes? `vortex_bench::Benchmark` gains a required `doc_path()` method (new suites must ship a doc), and `print_measurements_json`/`export_results` take the doc path. results.json rows carry a new `doc` field. Benchmark PR comments now start with a 📖 link to the suite's explainer. --------- Signed-off-by: Joe Isaacs --- .github/workflows/bench-pr.yml | 4 +-- .github/workflows/sql-benchmarks.yml | 4 +-- benchmarks/compress-bench/README.md | 16 ++++++++++ benchmarks/compress-bench/src/main.rs | 7 +++-- benchmarks/random-access-bench/README.md | 22 ++++++++++++++ benchmarks/random-access-bench/src/lib.rs | 5 +++- scripts/compare-benchmark-jsons.py | 25 ++++++++++++++++ vortex-bench/REUSE.toml | 6 ++++ vortex-bench/sql/appian/README.md | 22 ++++++++++++++ vortex-bench/sql/clickbench.md | 30 +++++++++++++++++++ vortex-bench/sql/fineweb.md | 21 +++++++++++++ vortex-bench/sql/gharchive.md | 14 +++++++++ vortex-bench/sql/polarsignals.md | 23 ++++++++++++++ vortex-bench/sql/public-bi.md | 15 ++++++++++ vortex-bench/sql/spatialbench.md | 15 ++++++++++ vortex-bench/sql/statpopgen.md | 25 ++++++++++++++++ vortex-bench/sql/tpcds/README.md | 20 +++++++++++++ vortex-bench/sql/tpch/README.md | 23 ++++++++++++++ vortex-bench/sql/vortex/README.md | 26 ++++++++++++++++ vortex-bench/src/appian/mod.rs | 4 +++ vortex-bench/src/benchmark.rs | 4 +++ vortex-bench/src/clickbench/benchmark.rs | 8 +++++ vortex-bench/src/display.rs | 7 ++++- vortex-bench/src/fineweb/mod.rs | 4 +++ vortex-bench/src/polarsignals/benchmark.rs | 4 +++ vortex-bench/src/public_bi.rs | 4 +++ vortex-bench/src/realnest/gharchive.rs | 4 +++ vortex-bench/src/runner.rs | 10 +++++-- vortex-bench/src/spatialbench/benchmark.rs | 4 +++ .../src/statpopgen/statpopgen_benchmark.rs | 4 +++ vortex-bench/src/tpcds/tpcds_benchmark.rs | 4 +++ vortex-bench/src/tpch/benchmark.rs | 4 +++ vortex-bench/src/vortex_queries.rs | 4 +++ 33 files changed, 380 insertions(+), 12 deletions(-) create mode 100644 benchmarks/compress-bench/README.md create mode 100644 benchmarks/random-access-bench/README.md create mode 100644 vortex-bench/sql/appian/README.md create mode 100644 vortex-bench/sql/clickbench.md create mode 100644 vortex-bench/sql/fineweb.md create mode 100644 vortex-bench/sql/gharchive.md create mode 100644 vortex-bench/sql/polarsignals.md create mode 100644 vortex-bench/sql/public-bi.md create mode 100644 vortex-bench/sql/spatialbench.md create mode 100644 vortex-bench/sql/statpopgen.md create mode 100644 vortex-bench/sql/tpcds/README.md create mode 100644 vortex-bench/sql/tpch/README.md create mode 100644 vortex-bench/sql/vortex/README.md diff --git a/.github/workflows/bench-pr.yml b/.github/workflows/bench-pr.yml index 12137726122..b5b011150a6 100644 --- a/.github/workflows/bench-pr.yml +++ b/.github/workflows/bench-pr.yml @@ -113,10 +113,8 @@ jobs: python3 scripts/s3-download.py s3://vortex-ci-benchmark-results/data.json.gz data.json.gz --no-sign-request gzip -d -c data.json.gz > base.json - echo '# Benchmarks: ${{ matrix.benchmark.name }}' > comment.md - echo '' >> comment.md uv run --no-project scripts/compare-benchmark-jsons.py base.json results.json "${{ matrix.benchmark.name }}" \ - >> comment.md + > comment.md cat comment.md >> $GITHUB_STEP_SUMMARY - name: Comment PR diff --git a/.github/workflows/sql-benchmarks.yml b/.github/workflows/sql-benchmarks.yml index 243db83aff8..59dd9571b36 100644 --- a/.github/workflows/sql-benchmarks.yml +++ b/.github/workflows/sql-benchmarks.yml @@ -678,10 +678,8 @@ jobs: python3 scripts/s3-download.py s3://vortex-ci-benchmark-results/data.json.gz data.json.gz --no-sign-request gzip -d -c data.json.gz > base.json - echo "# Benchmarks: ${{ matrix.name }}" > comment.md - echo '' >> comment.md uv run --no-project scripts/compare-benchmark-jsons.py base.json results.json "${{ matrix.name }}" \ - >> comment.md + > comment.md cat comment.md >> "$GITHUB_STEP_SUMMARY" - name: Comment PR diff --git a/benchmarks/compress-bench/README.md b/benchmarks/compress-bench/README.md new file mode 100644 index 00000000000..df1027e2bd0 --- /dev/null +++ b/benchmarks/compress-bench/README.md @@ -0,0 +1,16 @@ +# Compression benchmark + +Measures compression and decompression throughput, plus resulting file sizes, for Vortex +versus Parquet (and optionally Lance) across a range of datasets: NYC taxi data, several +[Public BI](https://github.com/cwida/public_bi_benchmark) tables (Arade, Bimbo, +CMSprovider, Euro2016, Food, HashTags), TPC-H `l_comment` variants, and synthetic nested +data. This is the workload behind the `Compression` PR comment. + +See [`src/main.rs`](./src/main.rs) for the dataset list and CLI flags (`--formats`, +`--datasets`, `--ops compress,decompress`). + +## Running locally + +```bash +cargo run -p compress-bench --profile release_debug +``` diff --git a/benchmarks/compress-bench/src/main.rs b/benchmarks/compress-bench/src/main.rs index 8d9fa73f915..1cbf3f47c55 100644 --- a/benchmarks/compress-bench/src/main.rs +++ b/benchmarks/compress-bench/src/main.rs @@ -113,6 +113,9 @@ fn get_compressor(format: Format) -> Box { /// The benchmark ID used for output path. const BENCHMARK_ID: &str = "compress"; +/// Repo-relative path of the suite explainer linked from CI benchmark PR comments. +const DOC_PATH: &str = "benchmarks/compress-bench/README.md"; + async fn run_compress( iterations: usize, datasets_filter: Option, @@ -212,8 +215,8 @@ async fn run_compress( ) } DisplayFormat::GhJson => { - print_measurements_json(&mut writer, measurements.timings)?; - print_measurements_json(&mut writer, measurements.ratios) + print_measurements_json(&mut writer, measurements.timings, DOC_PATH)?; + print_measurements_json(&mut writer, measurements.ratios, DOC_PATH) } } } diff --git a/benchmarks/random-access-bench/README.md b/benchmarks/random-access-bench/README.md new file mode 100644 index 00000000000..14949c85fcb --- /dev/null +++ b/benchmarks/random-access-bench/README.md @@ -0,0 +1,22 @@ +# Random Access benchmark + +Measures point-lookup latency: fetching individual rows by index from a file, rather than +scanning it. This is the workload behind the `Random Access` PR comment. + +Two access patterns are generated with a fixed seed (see [`src/main.rs`](./src/main.rs)): + +- **correlated**: several clusters of consecutive indices scattered across the dataset, + simulating lookups with spatial locality; +- **uniform**: indices drawn from a Poisson process spread uniformly across the dataset, + simulating lookups with no locality. + +Each pattern runs over four datasets (`taxi`, `feature-vectors`, `nested-lists`, +`nested-structs`) in Parquet, Lance, and Vortex, both with a cached open file handle and +reopening the file per lookup. CI drives the full matrix via +[`scripts/random-access-split.py`](../../scripts/random-access-split.py). + +## Running locally + +```bash +cargo run -p random-access-bench --profile release_debug --features lance +``` diff --git a/benchmarks/random-access-bench/src/lib.rs b/benchmarks/random-access-bench/src/lib.rs index 7d4b7d2c389..27be40d0342 100644 --- a/benchmarks/random-access-bench/src/lib.rs +++ b/benchmarks/random-access-bench/src/lib.rs @@ -282,6 +282,9 @@ async fn open_accessor( /// The benchmark ID used for output path. const BENCHMARK_ID: &str = "random-access"; +/// Repo-relative path of the suite explainer linked from CI benchmark PR comments. +const DOC_PATH: &str = "benchmarks/random-access-bench/README.md"; + /// Fixed indices used by the original taxi benchmark (preserved for historical continuity). const FIXED_TAXI_INDICES: [u64; 6] = [10, 11, 12, 13, 100_000, 3_000_000]; @@ -409,7 +412,7 @@ pub async fn run(config: RunConfig) -> Result<()> { } DisplayFormat::GhJson => { let timings: Vec = runs.into_iter().map(|r| r.timing).collect(); - print_measurements_json(&mut writer, timings)?; + print_measurements_json(&mut writer, timings, DOC_PATH)?; } } diff --git a/scripts/compare-benchmark-jsons.py b/scripts/compare-benchmark-jsons.py index 7d560c144a3..d7cd34793db 100644 --- a/scripts/compare-benchmark-jsons.py +++ b/scripts/compare-benchmark-jsons.py @@ -12,6 +12,7 @@ # SPDX-FileCopyrightText: Copyright the Vortex contributors import math +import os import re import sys from dataclasses import dataclass @@ -763,6 +764,27 @@ def format_within_engine_summary(analyses: dict[str, dict[str, Any]]) -> str | N return " · ".join(summaries) +def format_title(benchmark_name: str, pr: pd.DataFrame) -> str: + """Render the comment title, linking the suite explainer doc emitted by the benchmark binary. + + The doc path is a repo-relative markdown path carried on the PR result rows (the `doc` + field, populated from `Benchmark::doc_path` in Rust), so the benchmark code is the single + source of truth for where each suite is documented. The link pins the PR's own commit so + it resolves before the PR merges and stays valid afterwards. + """ + + title = f"# Benchmarks: {benchmark_name}" if benchmark_name else "# Benchmarks" + if "doc" in pr.columns: + docs = pr["doc"].dropna().unique() + if len(docs) > 0: + server_url = os.environ.get("GITHUB_SERVER_URL", "https://github.com") + repository = os.environ.get("GITHUB_REPOSITORY", "vortex-data/vortex") + commits = pr["commit_id"].dropna().unique() if "commit_id" in pr.columns else [] + ref = commits[0] if len(commits) > 0 else "develop" + title += f" [\N{OPEN BOOK}]({server_url}/{repository}/blob/{ref}/{docs[0]})" + return title + + def format_report_help() -> str: """Render explanatory markdown for the benchmark report headline fields.""" @@ -823,6 +845,7 @@ def main() -> None: benchmark_name = sys.argv[3] if len(sys.argv) > 3 else "" pr = pd.read_json(sys.argv[2], lines=True) + title = format_title(benchmark_name, pr) base = read_latest_baseline_rows(sys.argv[1], pr) base_commit_id = set(base["commit_id"].unique()) @@ -896,6 +919,8 @@ def main() -> None: shifts += f" · Median polish {format_ratio_change(float(np.exp(polish.overall)))}" summary_fields.append(f"**Shifts**: {shifts}") + print(title) + print("") print("
".join(summary_fields)) print("") print(format_report_help()) diff --git a/vortex-bench/REUSE.toml b/vortex-bench/REUSE.toml index 5fb9bd33466..a39e10c12b5 100644 --- a/vortex-bench/REUSE.toml +++ b/vortex-bench/REUSE.toml @@ -10,6 +10,12 @@ path = "**/*.csv" SPDX-FileCopyrightText = "Copyright the Vortex contributors" SPDX-License-Identifier = "Apache-2.0" +# Benchmark suite explainer docs linked from CI PR comments. +[[annotations]] +path = "sql/**/*.md" +SPDX-FileCopyrightText = "Copyright the Vortex contributors" +SPDX-License-Identifier = "CC-BY-4.0" + # `insta` snapshot files do not allow leading comment lines. [[annotations]] path = "src/snapshots/**.snap" diff --git a/vortex-bench/sql/appian/README.md b/vortex-bench/sql/appian/README.md new file mode 100644 index 00000000000..9304bcee2a0 --- /dev/null +++ b/vortex-bench/sql/appian/README.md @@ -0,0 +1,22 @@ +# Appian benchmark + +Mirrors the queries from DuckDB's in-tree +[`appian_benchmarks`](https://github.com/duckdb/duckdb/tree/main/benchmark/appian_benchmarks) +suite — a real-world business-application workload (customers, orders, addresses) with +selective filters and joins over camelCase-named columns. + +The eight queries live in this directory (`q1.sql` ... `q8.sql`). Upstream ships the data +as a single `.duckdb` blob (~593 MB); the harness downloads it once and projects each +table into Parquet with lowercased column names so that DataFusion and DuckDB resolve the +verbatim queries identically — see [`src/appian`](../../src/appian) for the details. + +## CI variant + +CI runs this suite from local NVMe as the `Appian on NVME` PR comment, comparing +DataFusion and DuckDB over Parquet and Vortex files (plus a native DuckDB baseline). + +## Running locally + +```bash +vx-bench run appian --engine datafusion,duckdb --format parquet,vortex +``` diff --git a/vortex-bench/sql/clickbench.md b/vortex-bench/sql/clickbench.md new file mode 100644 index 00000000000..eb3db810649 --- /dev/null +++ b/vortex-bench/sql/clickbench.md @@ -0,0 +1,30 @@ +# ClickBench benchmark + +[ClickBench](https://github.com/ClickHouse/ClickBench) is ClickHouse's web-analytics +benchmark: 43 queries over a single wide `hits` table (~100M rows of real-ish traffic +data). It is heavy on aggregations, `GROUP BY`s over high-cardinality columns, and +selective string filters, and is the main "wide table scan" workload in CI. + +The queries live in [`clickbench_queries.sql`](./clickbench_queries.sql) (one query per +line, numbered from Q0 in file order). The harness lives in +[`src/clickbench`](../src/clickbench). + +## CI variant + +CI runs this suite from local NVMe as the `Clickbench on NVME` PR comment, comparing +DataFusion and DuckDB over Parquet, Vortex, and vortex-compact files (plus a native +DuckDB baseline). + +## Sorted variant + +`Clickbench Sorted on NVME` runs the same table sorted by event time, split into 100 +shards whose filenames are shuffled so engines cannot rely on file order, exercising sort +pushdown and zone-map-style pruning on a subset of the queries. See +[`ClickBenchSortedBenchmark`](../src/clickbench/benchmark.rs). + +## Running locally + +```bash +vx-bench run clickbench --engine datafusion,duckdb --format parquet,vortex +vx-bench run clickbench-sorted --engine datafusion,duckdb --format parquet,vortex +``` diff --git a/vortex-bench/sql/fineweb.md b/vortex-bench/sql/fineweb.md new file mode 100644 index 00000000000..13421cc5456 --- /dev/null +++ b/vortex-bench/sql/fineweb.md @@ -0,0 +1,21 @@ +# FineWeb benchmark + +A string-heavy benchmark over a sample shard of the HuggingFace +[FineWeb](https://huggingface.co/datasets/HuggingFaceFW/fineweb) web-crawl dataset +(URLs, dates, and full page text). The dataset exercises dictionary and FSST string +compression heavily, and the hand-crafted queries in [`fineweb.sql`](./fineweb.sql) +(numbered from Q0 in file order) focus on string predicates: equality filters, `LIKE` +prefix and containment patterns, and aggregations over string columns. + +The harness lives in [`src/fineweb`](../src/fineweb). + +## CI variants + +CI runs this suite from local NVMe (`FineWeb NVMe`) and from S3 (`FineWeb S3`), comparing +DataFusion and DuckDB over Parquet, Vortex, and vortex-compact files. + +## Running locally + +```bash +vx-bench run fineweb --engine datafusion,duckdb --format parquet,vortex +``` diff --git a/vortex-bench/sql/gharchive.md b/vortex-bench/sql/gharchive.md new file mode 100644 index 00000000000..16cebd2d1a9 --- /dev/null +++ b/vortex-bench/sql/gharchive.md @@ -0,0 +1,14 @@ +# GitHub Archive benchmark + +A deeply nested real-world dataset from the [GitHub Archive](https://www.gharchive.org/): +GitHub event records whose `payload`, `repo`, `actor`, and `org` fields are nested structs. +The queries in [`gharchive.sql`](./gharchive.sql) (numbered from Q0 in file order) filter +and aggregate on those nested fields, making this the main struct-field-pushdown workload. + +The harness lives in [`src/realnest/gharchive.rs`](../src/realnest/gharchive.rs). + +## Running locally + +```bash +cargo run --bin datafusion-bench --profile release_debug -- --benchmark gharchive +``` diff --git a/vortex-bench/sql/polarsignals.md b/vortex-bench/sql/polarsignals.md new file mode 100644 index 00000000000..499466593e7 --- /dev/null +++ b/vortex-bench/sql/polarsignals.md @@ -0,0 +1,23 @@ +# PolarSignals Profiling benchmark + +A benchmark over continuous-profiling stacktrace data from +[Polar Signals](https://www.polarsignals.com/). The schema features a sparse struct (ten +nullable label fields covering five fill-rate tiers), deeply nested locations +(`List>>>`), and several low-cardinality string columns — +making it the main deeply-nested / sparse-data workload in CI. + +The queries in [`polarsignals.sql`](./polarsignals.sql) select stacktrace samples over +increasingly wide time ranges with equality filters on the low-cardinality metadata +columns. The harness lives in [`src/polarsignals`](../src/polarsignals). + +## CI variant + +CI runs this suite from local NVMe as the `PolarSignals Profiling` PR comment. It runs +DataFusion over Vortex only (there is no Parquet control for this suite, so its geomean is +reported without a drift-corrected verdict). + +## Running locally + +```bash +vx-bench run polarsignals --engine datafusion --format vortex +``` diff --git a/vortex-bench/sql/public-bi.md b/vortex-bench/sql/public-bi.md new file mode 100644 index 00000000000..60859677741 --- /dev/null +++ b/vortex-bench/sql/public-bi.md @@ -0,0 +1,15 @@ +# Public BI benchmark + +Real-world BI datasets from the +[Public BI benchmark](https://github.com/cwida/public_bi_benchmark) — anonymized Tableau +workbooks with messy, wide, string-heavy tables. A subset of the datasets (Arade, Bimbo, +CMSprovider, Euro2016, Food, HashTags) also feeds the +[Compression benchmark](../../benchmarks/compress-bench/README.md). + +The datasets, schemas, and queries are defined in [`src/public_bi.rs`](../src/public_bi.rs). + +## Running locally + +```bash +vx-bench run public-bi +``` diff --git a/vortex-bench/sql/spatialbench.md b/vortex-bench/sql/spatialbench.md new file mode 100644 index 00000000000..f959e006887 --- /dev/null +++ b/vortex-bench/sql/spatialbench.md @@ -0,0 +1,15 @@ +# SpatialBench benchmark + +The [Apache Sedona SpatialBench](https://sedona.apache.org/spatialbench/) geospatial +analytics benchmark: twelve queries (Q1 ... Q12 in [`spatialbench.sql`](./spatialbench.sql), +DuckDB dialect) over a trips/zones schema, exercising spatial predicates and functions such +as `ST_DWithin`, `ST_Intersects`, and `ST_Distance`. The query logic matches upstream +`sedona-spatialbench`; only formatting differs. + +The harness lives in [`src/spatialbench`](../src/spatialbench). + +## Running locally + +```bash +vx-bench run spatialbench +``` diff --git a/vortex-bench/sql/statpopgen.md b/vortex-bench/sql/statpopgen.md new file mode 100644 index 00000000000..c0b90a54f54 --- /dev/null +++ b/vortex-bench/sql/statpopgen.md @@ -0,0 +1,25 @@ +# Statistical and Population Genetics benchmark + +A genomics analytics benchmark using the gnomAD v3.1.2 release of the jointly called One +Thousand Genomes + Human Genome Diversity Project (1kG+HGDP) dataset — a prefix of +chromosome 21. The data source is . + +The distinguishing feature is the `GT` genotype column: a large variable-length list of +small nullable integers per row, which stresses list encodings and list-aware compute. +The queries in [`statpopgen.sql`](./statpopgen.sql) cover common genomics patterns: +allele-frequency calculations, Hardy-Weinberg equilibrium statistics, variant filtering by +frequency and by locus interval, and random access to a specific variant. + +The harness lives in [`src/statpopgen`](../src/statpopgen). + +## CI variant + +CI runs this suite from local NVMe as the `Statistical and Population Genetics` PR +comment. Only DuckDB is exercised (over Parquet, Vortex, and vortex-compact files), +because the queries rely on DuckDB list lambdas. + +## Running locally + +```bash +vx-bench run statpopgen --engine duckdb --format parquet,vortex +``` diff --git a/vortex-bench/sql/tpcds/README.md b/vortex-bench/sql/tpcds/README.md new file mode 100644 index 00000000000..3e4258ce251 --- /dev/null +++ b/vortex-bench/sql/tpcds/README.md @@ -0,0 +1,20 @@ +# TPC-DS benchmark + +The industry-standard [TPC-DS](https://www.tpc.org/tpcds/) decision-support benchmark: 99 +analytical queries (`01.sql` ... `99.sql`) over a retail snowflake schema. Compared to +TPC-H it has many more queries, a wider schema, and heavier use of joins, subqueries, and +window functions, making it a broad regression net for scan and pushdown behavior. + +The benchmark harness lives in [`src/tpcds`](../../src/tpcds). + +## CI variant + +CI runs this suite at scale factor 1 from local NVMe as the `TPC-DS SF=1 on NVME` PR +comment, comparing DataFusion and DuckDB over Parquet, Vortex, and vortex-compact files +(plus a native DuckDB baseline). + +## Running locally + +```bash +vx-bench run tpcds --engine datafusion,duckdb --format parquet,vortex --opt scale-factor=1.0 +``` diff --git a/vortex-bench/sql/tpch/README.md b/vortex-bench/sql/tpch/README.md new file mode 100644 index 00000000000..aa5a16fbab2 --- /dev/null +++ b/vortex-bench/sql/tpch/README.md @@ -0,0 +1,23 @@ +# TPC-H benchmark + +The industry-standard [TPC-H](https://www.tpc.org/tpch/) decision-support benchmark: 22 +analytical queries over an 8-table warehouse schema (`lineitem`, `orders`, `customer`, ...). +It stresses scan throughput, filter pushdown, joins, and aggregations over mostly numeric +and date columns. + +The queries in this directory (`q1.sql` ... `q22.sql`) are the standard TPC-H queries. +Data is generated deterministically by [`tpchgen`](../../src/tpch/tpchgen.rs); the benchmark +harness lives in [`src/tpch`](../../src/tpch). + +## CI variants + +CI runs this suite at scale factor 1 and 10, reading either from local NVMe or from S3 +(the data is uploaded to a bucket first), as the `TPC-H SF={1,10} on {NVME,S3}` PR +comments. Each variant compares DataFusion and DuckDB over Parquet, Vortex, and +vortex-compact files (plus in-memory Arrow and native DuckDB baselines on NVMe). + +## Running locally + +```bash +vx-bench run tpch --engine datafusion,duckdb --format parquet,vortex --opt scale-factor=1.0 +``` diff --git a/vortex-bench/sql/vortex/README.md b/vortex-bench/sql/vortex/README.md new file mode 100644 index 00000000000..793f7fb0f67 --- /dev/null +++ b/vortex-bench/sql/vortex/README.md @@ -0,0 +1,26 @@ +# Vortex queries benchmark + +A small suite of microbenchmark queries targeting Vortex-specific scan paths, run on +**every PR commit** (unlike the label-gated suites) with a high iteration count, so it is +the most sensitive — and most frequently seen — benchmark comment. + +[`init.sql`](./init.sql) generates a 25M-row two-column table; the numbered queries each +pin down one scan behavior: + +- [`0_sum-with-filter.sql`](./0_sum-with-filter.sql): a filtered aggregation that forces a + linear scan today; once statistics are propagated to arrays it should use zone maps + instead of decoding every row. +- [`1_sum.sql`](./1_sum.sql): an unfiltered aggregation that should eventually be answered + from footer statistics without decoding data. + +## CI variant + +CI runs this as the `Vortex queries` PR comment (see +[`.github/workflows/sql-vortex-pr.yml`](../../../.github/workflows/sql-vortex-pr.yml)), +comparing DataFusion and DuckDB over Parquet and Vortex files with 100 iterations. + +## Running locally + +```bash +vx-bench run vortex --engine datafusion,duckdb --format parquet,vortex +``` diff --git a/vortex-bench/src/appian/mod.rs b/vortex-bench/src/appian/mod.rs index cdc82ac200a..20dabb5676a 100644 --- a/vortex-bench/src/appian/mod.rs +++ b/vortex-bench/src/appian/mod.rs @@ -121,6 +121,10 @@ impl AppianBenchmark { #[async_trait::async_trait] impl Benchmark for AppianBenchmark { + fn doc_path(&self) -> &'static str { + "vortex-bench/sql/appian/README.md" + } + fn queries(&self) -> anyhow::Result> { Ok(appian_queries().collect()) } diff --git a/vortex-bench/src/benchmark.rs b/vortex-bench/src/benchmark.rs index d59708886a6..b7015d242a7 100644 --- a/vortex-bench/src/benchmark.rs +++ b/vortex-bench/src/benchmark.rs @@ -64,6 +64,10 @@ pub trait Benchmark: Send + Sync { fn dataset(&self) -> BenchmarkDataset; + /// Repo-relative path of the markdown explainer for this benchmark suite, linked from the + /// title of CI benchmark PR comments. Required: every suite must ship a doc. + fn doc_path(&self) -> &'static str; + /// Get the name of the benchmark dataset fn dataset_name(&self) -> &str; diff --git a/vortex-bench/src/clickbench/benchmark.rs b/vortex-bench/src/clickbench/benchmark.rs index 346496dbda0..fcba1b9f9c1 100644 --- a/vortex-bench/src/clickbench/benchmark.rs +++ b/vortex-bench/src/clickbench/benchmark.rs @@ -76,6 +76,10 @@ fn read_clickbench_queries(queries_file: Option<&str>) -> Result &'static str { + "vortex-bench/sql/clickbench.md" + } + fn queries(&self) -> Result> { read_clickbench_queries(self.queries_file.as_deref()) } @@ -120,6 +124,10 @@ impl Benchmark for ClickBenchBenchmark { #[async_trait::async_trait] impl Benchmark for ClickBenchSortedBenchmark { + fn doc_path(&self) -> &'static str { + "vortex-bench/sql/clickbench.md#sorted-variant" + } + fn queries(&self) -> Result> { Ok(read_clickbench_queries(self.queries_file.as_deref())? .into_iter() diff --git a/vortex-bench/src/display.rs b/vortex-bench/src/display.rs index a8e903df25a..343c1a06612 100644 --- a/vortex-bench/src/display.rs +++ b/vortex-bench/src/display.rs @@ -105,9 +105,14 @@ pub fn render_table( pub fn print_measurements_json( writer: &mut dyn Write, all_measurements: Vec, + doc: &str, ) -> anyhow::Result<()> { for measurement in all_measurements { - writeln!(writer, "{}", measurement.to_json())?; + let mut json = measurement.to_json(); + if let Some(obj) = json.as_object_mut() { + obj.insert("doc".to_string(), doc.into()); + } + writeln!(writer, "{json}")?; } Ok(()) diff --git a/vortex-bench/src/fineweb/mod.rs b/vortex-bench/src/fineweb/mod.rs index b8c86925905..aefa796af3a 100644 --- a/vortex-bench/src/fineweb/mod.rs +++ b/vortex-bench/src/fineweb/mod.rs @@ -53,6 +53,10 @@ impl FinewebBenchmark { #[async_trait::async_trait] impl Benchmark for FinewebBenchmark { + fn doc_path(&self) -> &'static str { + "vortex-bench/sql/fineweb.md" + } + /// Some basic string-focused queries, numbered from Q0 in `sql/fineweb.sql` file order. fn queries(&self) -> anyhow::Result> { // `;`-separated; a `;` must not appear in a comment, or it would split a statement in two. diff --git a/vortex-bench/src/polarsignals/benchmark.rs b/vortex-bench/src/polarsignals/benchmark.rs index fce692ceb68..d1097e83bd1 100644 --- a/vortex-bench/src/polarsignals/benchmark.rs +++ b/vortex-bench/src/polarsignals/benchmark.rs @@ -59,6 +59,10 @@ impl PolarSignalsBenchmark { #[async_trait::async_trait] impl Benchmark for PolarSignalsBenchmark { + fn doc_path(&self) -> &'static str { + "vortex-bench/sql/polarsignals.md" + } + fn queries(&self) -> Result> { let queries_file = workspace_root() .join("vortex-bench") diff --git a/vortex-bench/src/public_bi.rs b/vortex-bench/src/public_bi.rs index f5c22bdea8f..3e651df5bb7 100644 --- a/vortex-bench/src/public_bi.rs +++ b/vortex-bench/src/public_bi.rs @@ -530,6 +530,10 @@ impl PublicBiBenchmark { #[async_trait] impl Benchmark for PublicBiBenchmark { + fn doc_path(&self) -> &'static str { + "vortex-bench/sql/public-bi.md" + } + fn queries(&self) -> anyhow::Result> { self.pbi_benchmark().queries() } diff --git a/vortex-bench/src/realnest/gharchive.rs b/vortex-bench/src/realnest/gharchive.rs index a84395f4a10..6596598d5ba 100644 --- a/vortex-bench/src/realnest/gharchive.rs +++ b/vortex-bench/src/realnest/gharchive.rs @@ -66,6 +66,10 @@ impl GithubArchiveBenchmark { #[async_trait::async_trait] impl Benchmark for GithubArchiveBenchmark { + fn doc_path(&self) -> &'static str { + "vortex-bench/sql/gharchive.md" + } + /// GitHub Archive queries, numbered from Q0 in `sql/gharchive.sql` file order. fn queries(&self) -> anyhow::Result> { // `;`-separated; a `;` must not appear in a comment, or it would split a statement in two. diff --git a/vortex-bench/src/runner.rs b/vortex-bench/src/runner.rs index 1983c63ba05..9e5c4db3d08 100644 --- a/vortex-bench/src/runner.rs +++ b/vortex-bench/src/runner.rs @@ -75,6 +75,7 @@ pub struct SqlBenchmarkRunner { formats: Vec, memory_tracker: Option, hide_progress_bar: bool, + doc: &'static str, query_measurements: Vec, memory_measurements: Vec, } @@ -106,6 +107,7 @@ impl SqlBenchmarkRunner { formats, memory_tracker, hide_progress_bar, + doc: benchmark.doc_path(), query_measurements: Vec::new(), memory_measurements: Vec::new(), }) @@ -223,6 +225,7 @@ impl SqlBenchmarkRunner { display_format, self.engine, &self.formats, + self.doc, f, ) } @@ -232,6 +235,7 @@ impl SqlBenchmarkRunner { display_format, self.engine, &self.formats, + self.doc, std::io::stdout().lock(), ), } @@ -249,6 +253,7 @@ impl SqlBenchmarkRunner { display_format, self.engine, &self.formats, + self.doc, output, ) } @@ -457,6 +462,7 @@ pub fn export_results( display_format: &DisplayFormat, engine: Engine, formats: &[Format], + doc: &str, mut output: W, ) -> anyhow::Result<()> { let targets = formats @@ -467,13 +473,13 @@ pub fn export_results( if !memory.is_empty() { match display_format { DisplayFormat::Table => render_table(&mut output, memory, &targets)?, - DisplayFormat::GhJson => print_measurements_json(&mut output, memory)?, + DisplayFormat::GhJson => print_measurements_json(&mut output, memory, doc)?, }; } match display_format { DisplayFormat::Table => render_table(&mut output, queries, &targets)?, - DisplayFormat::GhJson => print_measurements_json(&mut output, queries)?, + DisplayFormat::GhJson => print_measurements_json(&mut output, queries, doc)?, }; Ok(()) diff --git a/vortex-bench/src/spatialbench/benchmark.rs b/vortex-bench/src/spatialbench/benchmark.rs index 48fa84f8ec1..5df7671bd1f 100644 --- a/vortex-bench/src/spatialbench/benchmark.rs +++ b/vortex-bench/src/spatialbench/benchmark.rs @@ -58,6 +58,10 @@ impl SpatialBenchBenchmark { #[async_trait::async_trait] impl Benchmark for SpatialBenchBenchmark { + fn doc_path(&self) -> &'static str { + "vortex-bench/sql/spatialbench.md" + } + /// All SpatialBench queries, numbered started at Q1 in `spatialbench.sql` file order. fn queries(&self) -> anyhow::Result> { // `;`-separated; a `;` must not appear in a comment, or it would split a statement in two. diff --git a/vortex-bench/src/statpopgen/statpopgen_benchmark.rs b/vortex-bench/src/statpopgen/statpopgen_benchmark.rs index 0861f7bfc88..2c677fb79a3 100644 --- a/vortex-bench/src/statpopgen/statpopgen_benchmark.rs +++ b/vortex-bench/src/statpopgen/statpopgen_benchmark.rs @@ -103,6 +103,10 @@ impl StatPopGenBenchmark { #[async_trait::async_trait] impl Benchmark for StatPopGenBenchmark { + fn doc_path(&self) -> &'static str { + "vortex-bench/sql/statpopgen.md" + } + fn queries(&self) -> Result> { let queries_file = workspace_root() .join("vortex-bench") diff --git a/vortex-bench/src/tpcds/tpcds_benchmark.rs b/vortex-bench/src/tpcds/tpcds_benchmark.rs index 7f8f2d5281e..5bacae73516 100644 --- a/vortex-bench/src/tpcds/tpcds_benchmark.rs +++ b/vortex-bench/src/tpcds/tpcds_benchmark.rs @@ -56,6 +56,10 @@ impl TpcDsBenchmark { #[async_trait::async_trait] impl Benchmark for TpcDsBenchmark { + fn doc_path(&self) -> &'static str { + "vortex-bench/sql/tpcds/README.md" + } + fn queries(&self) -> Result> { Ok(tpcds_queries().collect()) } diff --git a/vortex-bench/src/tpch/benchmark.rs b/vortex-bench/src/tpch/benchmark.rs index 6851e5f5f06..4dfdaf1207d 100644 --- a/vortex-bench/src/tpch/benchmark.rs +++ b/vortex-bench/src/tpch/benchmark.rs @@ -75,6 +75,10 @@ impl TpcHBenchmark { #[async_trait::async_trait] impl Benchmark for TpcHBenchmark { + fn doc_path(&self) -> &'static str { + "vortex-bench/sql/tpch/README.md" + } + fn queries(&self) -> anyhow::Result> { Ok(tpch_queries().collect()) } diff --git a/vortex-bench/src/vortex_queries.rs b/vortex-bench/src/vortex_queries.rs index f611238bb6b..63e9a658068 100644 --- a/vortex-bench/src/vortex_queries.rs +++ b/vortex-bench/src/vortex_queries.rs @@ -89,6 +89,10 @@ impl VortexBenchmark { #[async_trait::async_trait] impl Benchmark for VortexBenchmark { + fn doc_path(&self) -> &'static str { + "vortex-bench/sql/vortex/README.md" + } + fn queries(&self) -> Result> { self.query_files()? .iter() From ef056d56627702a7ecd0eb7eceebeb003441fa32 Mon Sep 17 00:00:00 2001 From: Francesco Gargiulo Date: Fri, 17 Jul 2026 17:51:03 +0100 Subject: [PATCH 009/498] feat(vortex-onpair): update to the onpair 0.1 API (#8675) ## Summary Updates the experimental `vortex-onpair` encoding to the public API introduced in `onpair` 0.1. - Migrates compression and decoding to the new configuration, raw column-parts, and validated compact-dictionary APIs. - Updates compressed-domain equality comparisons for the new dictionary representation. - Supports adaptive `u32`/`u64` code offsets. - Validates dictionaries loaded from serialized data, including rejecting dictionaries that exceed the `u16` token address space. - Adds regression coverage for empty and all-null inputs, large dictionaries, narrowed children, sliced arrays, and `u64` code offsets. ## Scope This PR intentionally does not include the new `LIKE`, `list_contains`, or `take` pushdowns. Those will be submitted in a follow-up PR. A separate follow-up will enable OnPair on the default read path without enabling it in the default compressor. ## Dependency note The dependency is temporarily pinned to an upstream OnPair commit containing the dictionary-validation fix. Before merging, that fix will be published in a new crates.io release and the Git dependency will be replaced with the released version. ## Testing - `cargo +nightly fmt --all -- --check` - `cargo test -p vortex-onpair` - `cargo test -p vortex-btrblocks --features unstable_encodings` - `cargo clippy -p vortex-onpair -p vortex-btrblocks --all-targets --all-features -- -D warnings` --------- Signed-off-by: Francesco Gargiulo Signed-off-by: Francesco Gargiulo Co-authored-by: Claude Fable 5 --- Cargo.lock | 4 +- Cargo.toml | 2 +- encodings/experimental/onpair/README.md | 25 ++- .../experimental/onpair/benches/decode.rs | 55 ++--- .../onpair/goldenfiles/onpair.metadata | 2 +- encodings/experimental/onpair/src/array.rs | 192 +++++++++++------- .../experimental/onpair/src/canonical.rs | 32 +-- encodings/experimental/onpair/src/compress.rs | 165 ++++++--------- .../experimental/onpair/src/compute/cast.rs | 105 ++++++++-- .../onpair/src/compute/compare.rs | 176 ++++++++++++++-- .../experimental/onpair/src/compute/filter.rs | 3 +- .../experimental/onpair/src/compute/slice.rs | 3 +- encodings/experimental/onpair/src/decode.rs | 72 ++++++- encodings/experimental/onpair/src/kernel.rs | 4 +- encodings/experimental/onpair/src/lib.rs | 9 +- encodings/experimental/onpair/src/ops.rs | 28 +-- encodings/experimental/onpair/src/tests.rs | 167 ++++++++++++++- .../experimental/onpair/tests/big_data.rs | 5 +- vortex-btrblocks/src/schemes/string/onpair.rs | 8 +- 19 files changed, 758 insertions(+), 299 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index e23ac0d40f8..e79211b5436 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6451,9 +6451,9 @@ checksum = "cfe21416a02c693fb9f980befcb230ecc70b0b3d1cc4abf88b9675c4c1457f0c" [[package]] name = "onpair" -version = "0.0.4" +version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e7f54ca0079320b4b43a7502822b8a2e8a2c713adb7fc65ae4e65f940921a7eb" +checksum = "daf79077d8decdf5714242f1813f844f23c80a0cf009aba52881effb4c36ecbe" dependencies = [ "hashbrown 0.16.1", "rand 0.9.4", diff --git a/Cargo.toml b/Cargo.toml index 12e91c7ef20..68be0c50653 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -192,7 +192,7 @@ num_enum = { version = "0.7.3", default-features = false } object_store = { version = "0.13.2", default-features = false } once_cell = "1.21" oneshot = { version = "0.2.0", features = ["async"] } -onpair = { version = "0.0.4" } +onpair = "0.1.1" opentelemetry = "0.32.0" opentelemetry-otlp = "0.32.0" opentelemetry_sdk = "0.32.0" diff --git a/encodings/experimental/onpair/README.md b/encodings/experimental/onpair/README.md index 9628c006201..f19697a412d 100644 --- a/encodings/experimental/onpair/README.md +++ b/encodings/experimental/onpair/README.md @@ -10,26 +10,25 @@ cascading-compressor support on every integer child. ## Compute -Like the FSST encoding, this crate provides `cast` and `filter` -pushdown. Other operators fall back to ordinary decompression. +Like the FSST encoding, this crate pushes down common operations over the +encoded representation. It supports `cast`, `filter`, byte length, and +constant equality / inequality. Unsupported operators fall back to ordinary +decompression. ## Default Configuration -The default training preset is **dict-12**: 12 bits per token, -dictionary capped at 4 096 entries. Token codes are stored as a -`PrimitiveArray`; downstream `FastLanes::BitPacking` losslessly -narrows the child to exactly `bits`-bit codes on disk. +The default training configuration uses OnPair's default dictionary budget and +a fixed seed. Vortex stores token codes as an integer child array; downstream +integer compression may narrow or bit-pack that child independently. ## Layout - Buffer 0 — `dict_bytes`: dictionary blob built by the OnPair trainer, - padded with `MAX_TOKEN_SIZE` trailing zero bytes so the over-copy - decoder can read 16 bytes past the last token. -- Slot 0 — `dict_offsets`: `PrimitiveArray`, len `dict_size + 1`. -- Slot 1 — `codes`: `PrimitiveArray`, length `total_tokens`. -- Slot 2 — `codes_offsets`: `PrimitiveArray`, length `num_rows + 1`. -- Slot 3 — `uncompressed_lengths`: integer `PrimitiveArray`, length - `num_rows`. + including the read padding required by the decoder. +- Slot 0 — `dict_offsets`: integer child, len `dict_size + 1`. +- Slot 1 — `codes`: integer child, length `total_tokens`. +- Slot 2 — `codes_offsets`: integer child, length `num_rows + 1`. +- Slot 3 — `uncompressed_lengths`: integer child, length `num_rows`. - Slot 4 — optional validity child. All four integer slot children flow through the standard cascading diff --git a/encodings/experimental/onpair/benches/decode.rs b/encodings/experimental/onpair/benches/decode.rs index ff4158cda32..8a0f54bc2d8 100644 --- a/encodings/experimental/onpair/benches/decode.rs +++ b/encodings/experimental/onpair/benches/decode.rs @@ -3,7 +3,7 @@ // //! Decode-path microbenchmarks for the OnPair Vortex array. //! -//! * `decompress_into` — the upstream `onpair::decompress_into` decoder hot +//! * `decode_into` — the upstream `onpair::decode_into` decoder hot //! loop, fed by pre-materialised [`DecodeInputs`]. Measures the inner loop //! only (no child `execute`, no allocation). //! * `canonicalize_to_varbinview` — the full Vortex @@ -28,7 +28,7 @@ use std::mem::MaybeUninit; use std::sync::LazyLock; use divan::Bencher; -use onpair::Parts; +use onpair::CompactDictionaryView; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; @@ -38,12 +38,12 @@ use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::VarBinArray; use vortex_array::arrays::VarBinViewArray; use vortex_array::arrays::filter::FilterKernel; +use vortex_array::buffer::BufferHandle; use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; use vortex_array::dtype::NativePType; use vortex_array::dtype::Nullability; use vortex_buffer::Buffer; -use vortex_buffer::ByteBuffer; use vortex_mask::Mask; use vortex_onpair::DEFAULT_DICT12_CONFIG; use vortex_onpair::OnPair; @@ -51,30 +51,32 @@ use vortex_onpair::OnPairArray; use vortex_onpair::OnPairArraySlotsExt; /// Host-resident decode inputs, materialised once so the decode-loop benchmark -/// measures only `onpair::decompress_into` (not child `execute`/allocation). +/// measures only `onpair::decode_into` (not child `execute`/allocation). struct DecodeInputs { - dict_bytes: ByteBuffer, + dict_bytes: BufferHandle, dict_offsets: Buffer, codes: Buffer, - bits: u32, } impl DecodeInputs { - fn as_parts(&self) -> Parts<'_> { - Parts { - dict_bytes: self.dict_bytes.as_slice(), - dict_offsets: self.dict_offsets.as_slice(), - bits: self.bits, - codes: self.codes.as_slice(), + fn dict(&self) -> CompactDictionaryView<'_> { + // SAFETY: `materialise` validates this borrowed dictionary once after + // widening offsets. The benchmark then keeps both buffers immutable. + unsafe { + CompactDictionaryView::new_unchecked( + self.dict_bytes.as_host().as_slice(), + self.dict_offsets.as_slice(), + ) } } - fn decompressed_len(&self) -> usize { - onpair::decompressed_len(self.as_parts()) + fn decoded_len(&self) -> usize { + onpair::decoded_len(self.codes.as_slice(), self.dict()) } - fn decompress_into(&self, out: &mut [MaybeUninit]) -> usize { - onpair::decompress_into(self.as_parts(), out) + fn decode_into(&self, out: &mut [MaybeUninit]) -> usize { + // SAFETY: callers allocate `decoded_len + DECODE_PADDING` bytes. + unsafe { onpair::decode_into(self.codes.as_slice(), self.dict(), out) } } } use vortex_onpair::onpair_compress; @@ -162,6 +164,8 @@ fn compress(n: usize, shape: Shape, ctx: &mut ExecutionCtx) -> OnPairArray { ); onpair_compress(varbin.as_array(), DEFAULT_DICT12_CONFIG, ctx) .unwrap_or_else(|e| panic!("onpair_compress failed: {e}")) + .try_downcast::() + .unwrap_or_else(|array| panic!("expected OnPair array, got {}", array.encoding_id())) } /// Canonicalise a slot child to the decoder's native primitive width. @@ -175,13 +179,16 @@ fn widen(arr: &ArrayRef, ctx: &mut ExecutionCtx) -> Buffer { fn materialise(arr: &OnPairArray, ctx: &mut ExecutionCtx) -> (DecodeInputs, usize) { let view = arr.as_view(); + let dict_offsets = widen::(view.dict_offsets(), ctx); + let dict_bytes = view.dict_bytes_handle().clone(); + CompactDictionaryView::validate(dict_bytes.as_host().as_slice(), dict_offsets.as_slice()) + .expect("valid OnPair dictionary"); let inputs = DecodeInputs { - dict_bytes: view.dict_bytes().clone(), - dict_offsets: widen::(view.dict_offsets(), ctx), + dict_bytes, + dict_offsets, codes: widen::(view.codes(), ctx), - bits: view.bits(), }; - let total = inputs.decompressed_len(); + let total = inputs.decoded_len(); (inputs, total) } @@ -194,16 +201,16 @@ const CASES: &[(Shape, usize)] = &[ ]; /// Raw decode loop time, excluding child `execute` and the output allocation. -/// Hits `onpair::decompress_into` directly. +/// Hits `onpair::decode_into` directly. #[divan::bench(args = CASES)] -fn decompress_into_bench(bencher: Bencher, case: (Shape, usize)) { +fn decode_into_bench(bencher: Bencher, case: (Shape, usize)) { let mut ctx = SESSION.create_execution_ctx(); let (shape, n) = case; let arr = compress(n, shape, &mut ctx); let (inputs, total) = materialise(&arr, &mut ctx); bencher.bench_local(|| { - let mut out: Vec = Vec::with_capacity(total); - let written = inputs.decompress_into(out.spare_capacity_mut()); + let mut out: Vec = Vec::with_capacity(total + onpair::DECODE_PADDING); + let written = inputs.decode_into(out.spare_capacity_mut()); unsafe { out.set_len(written) }; divan::black_box(out); }); diff --git a/encodings/experimental/onpair/goldenfiles/onpair.metadata b/encodings/experimental/onpair/goldenfiles/onpair.metadata index e96baf1a0ab..80b7fcdef71 100644 --- a/encodings/experimental/onpair/goldenfiles/onpair.metadata +++ b/encodings/experimental/onpair/goldenfiles/onpair.metadata @@ -1 +1 @@ - € €è(08 \ No newline at end of file +€ €è(08 \ No newline at end of file diff --git a/encodings/experimental/onpair/src/array.rs b/encodings/experimental/onpair/src/array.rs index 519243b100e..111269ac484 100644 --- a/encodings/experimental/onpair/src/array.rs +++ b/encodings/experimental/onpair/src/array.rs @@ -5,7 +5,10 @@ use std::fmt::Debug; use std::fmt::Display; use std::fmt::Formatter; use std::hash::Hasher; +use std::sync::Arc; +use std::sync::OnceLock; +use onpair::CompactDictionaryView; use prost::Message as _; use vortex_array::Array; use vortex_array::ArrayEq; @@ -32,6 +35,7 @@ use vortex_array::vtable::VTable; use vortex_array::vtable::ValidityVTable; use vortex_array::vtable::child_to_validity; use vortex_array::vtable::validity_to_child; +use vortex_buffer::Buffer; use vortex_buffer::ByteBuffer; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -43,6 +47,7 @@ use vortex_session::registry::CachedId; use crate::canonical::canonicalize_onpair; use crate::canonical::onpair_decode_views; +use crate::decode::collect_widened; use crate::rules::RULES; /// An [`OnPair`]-encoded Vortex array. @@ -52,9 +57,8 @@ pub type OnPairArray = Array; /// /// On disk the layout is FSST-shape: /// -/// * Buffer 0 — `dict_bytes`: the dictionary blob built by the OnPair trainer, -/// padded with `onpair::MAX_TOKEN_SIZE` trailing zero -/// bytes so the over-copy decoder can read 16 bytes past the last token. +/// * Buffer 0 — `dict_bytes`: the read-padded dictionary blob built by the +/// OnPair trainer. /// * Slots — see [`OnPairSlots`]. /// /// The four integer slot children flow through the standard `compress_child` @@ -66,24 +70,18 @@ pub struct OnPairMetadata { /// Width of the per-row primitive `uncompressed_lengths` child. #[prost(enumeration = "PType", tag = "1")] pub uncompressed_lengths_ptype: i32, - /// Bits-per-token the column was compressed with (9..=16). Every value - /// in the `codes` child only uses its low `bits` bits. - #[prost(uint32, tag = "2")] - pub bits: u32, /// Number of dictionary tokens. `dict_offsets` has length `dict_size + 1`. - /// Bounded by `2^bits ≤ 2^16 = 65_536`, so `u32` is comfortably wide. #[prost(uint32, tag = "3")] pub dict_size: u32, - /// Total number of tokens across all rows. `codes` has this length; - /// `codes_offsets.last() == total_tokens`. + /// Length of the `codes` slot child. A sliced array may retain codes that + /// fall outside its visible row range. #[prost(uint64, tag = "4")] - pub total_tokens: u64, + pub codes_len: u64, /// PType of the `dict_offsets` slot child (defaults to U32, may be /// narrowed to U16/U8 by the cascading compressor when values fit). #[prost(enumeration = "PType", tag = "5")] pub dict_offsets_ptype: i32, - /// PType of the `codes` slot child (typically U16, may be narrowed to U8 - /// when `bits <= 8`). + /// PType of the `codes` slot child. #[prost(enumeration = "PType", tag = "6")] pub codes_ptype: i32, /// PType of the `codes_offsets` slot child. @@ -92,6 +90,7 @@ pub struct OnPairMetadata { } impl OnPairMetadata { + /// Decode the recorded [`PType`] of the `uncompressed_lengths` slot child. pub fn get_uncompressed_lengths_ptype(&self) -> VortexResult { PType::try_from(self.uncompressed_lengths_ptype) .map_err(|_| vortex_err!("Invalid PType {}", self.uncompressed_lengths_ptype)) @@ -100,18 +99,17 @@ impl OnPairMetadata { #[array_slots(OnPair)] pub struct OnPairSlots { - /// `PrimitiveArray`, length `dict_size + 1`. Cascading compressor may - /// narrow the ptype to U16/U8. + /// Primitive integer dictionary offsets, length `dict_size + 1`. The + /// cascading compressor may re-encode this child independently. pub dict_offsets: ArrayRef, - /// `PrimitiveArray`. Each value only uses its low `bits` bits; - /// downstream `FastLanes::BitPacking` losslessly shrinks the child to - /// exactly `bits`-bit codes on disk. + /// Primitive integer token codes. Downstream integer compression may + /// narrow or bit-pack this child independently of the OnPair metadata. pub codes: ArrayRef, - /// `PrimitiveArray`, length `num_rows + 1`. FoR / RunEnd / etc. apply - /// naturally via the cascading compressor. + /// Primitive integer row offsets into `codes`, length `num_rows + 1`. The + /// cascading compressor may re-encode this child independently. pub codes_offsets: ArrayRef, - /// Integer `PrimitiveArray`, length `num_rows`. Used to size the canonical - /// output buffer. + /// Integer decoded-length child, length `num_rows`. Used to size the + /// canonical output buffer. pub uncompressed_lengths: ArrayRef, /// Optional validity child for the outer string column. pub validity: Option, @@ -127,58 +125,114 @@ pub struct OnPairSlots { pub struct OnPairData { /// The dictionary blob (buffer 0). /// - /// INVARIANT: this buffer must be over-padded past its logical end - /// (`dict_offsets.last()`) by the decoder's fixed token read width, - /// `onpair::MAX_TOKEN_SIZE`. The over-copy decoder reads - /// every dictionary entry with one fixed-width load and then advances the - /// cursor by the token's true length, so the load for the final, shortest - /// token over-reads past the logical end of the dictionary. This is the - /// same over-read the decoder accounts for on the final few codes; the - /// trailing padding absorbs it so that any entry can be read in bounds. - /// `onpair_compress` establishes this padding (see `parts_to_children`); - /// the over-copy decoder lives in the `onpair` crate. + /// INVARIANT: this buffer is an OnPair compact dictionary byte buffer, + /// including its trailing read padding. dict_bytes: BufferHandle, - bits: u32, len: usize, + /// The `dict_offsets` child widened to `u32`, memoized on first use so the + /// child is decompressed and the dictionary content is validated at most + /// once per dictionary — never per operation. (`dict_bytes` needs no such + /// cache: it is a raw buffer, so access is already zero-cost.) + /// + /// INVARIANT: once populated, the offsets passed + /// [`CompactDictionaryView::validate`] against `dict_bytes` (or came from + /// the trainer via [`init_dict_offsets`](Self::init_dict_offsets)), and + /// they are the widened values of the array's `dict_offsets` child. The + /// `Arc` cell is shared only between arrays with identical `dict_bytes` + /// and logically identical `dict_offsets` (slice / filter / cast keep + /// both). + dict_offsets: Arc>>, } impl OnPairData { - pub fn new(dict_bytes: BufferHandle, bits: u32, len: usize) -> Self { + /// Build [`OnPairData`] from the dictionary blob and the number of rows. + pub fn new(dict_bytes: BufferHandle, len: usize) -> Self { Self { dict_bytes, - bits, len, + dict_offsets: Arc::new(OnceLock::new()), } } + /// Seed the widened-offsets cell with dictionary offsets that are already + /// known to be conformant, so the first operation skips validation. + /// + /// This is the crate-internal trust mint (the moral equivalent of + /// [`onpair::CompactDictionary::new_unchecked`]): compression seeds it + /// with the trainer's offsets, which are conformant by construction. + /// + /// # Safety + /// `(self.dict_bytes, offsets)` must satisfy every [`onpair`] compact + /// dictionary invariant (i.e. [`CompactDictionaryView::validate`] would + /// succeed on them), and `offsets` must be the widened values of the + /// array's `dict_offsets` child. [`dict_view`] relies on this to build + /// unchecked dictionary views. + pub(crate) unsafe fn init_dict_offsets(&self, offsets: Buffer) { + debug_assert!( + CompactDictionaryView::validate(self.dict_bytes().as_slice(), offsets.as_slice()) + .is_ok(), + "init_dict_offsets called with a non-conformant dictionary" + ); + // A benign race can only ever install another conformant value. + drop(self.dict_offsets.set(offsets)); + } + + /// Number of rows in the array. pub fn len(&self) -> usize { self.len } + /// Whether the array has zero rows. pub fn is_empty(&self) -> bool { self.len == 0 } - pub fn bits(&self) -> u32 { - self.bits - } - + /// The dictionary blob as a host byte buffer. pub fn dict_bytes(&self) -> &ByteBuffer { self.dict_bytes.as_host() } + /// The [`BufferHandle`] holding the dictionary blob (buffer 0). pub fn dict_bytes_handle(&self) -> &BufferHandle { &self.dict_bytes } } +/// A conformant [`CompactDictionaryView`] over `array`'s dictionary. +/// +/// The first call per dictionary widens the `dict_offsets` child to `u32` and +/// validates the dictionary content; both results are memoized in +/// [`OnPairData`], so subsequent calls — including on arrays derived by +/// slice / filter / cast, which share the cell — pay neither cost again. +pub(crate) fn dict_view<'a>( + array: ArrayView<'a, OnPair>, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + let data = array.data(); + let offsets = match data.dict_offsets.get() { + Some(offsets) => offsets, + None => { + let widened = collect_widened::(array.dict_offsets(), ctx)?; + CompactDictionaryView::validate(data.dict_bytes().as_slice(), widened.as_slice()) + .map_err(|e| vortex_err!(InvalidArgument: "Invalid OnPair dictionary: {e}"))?; + data.dict_offsets.get_or_init(|| widened) + } + }; + // SAFETY: the cell only ever holds offsets that satisfy the compact + // dictionary invariants against this `dict_bytes` (validated above, or + // guaranteed by `init_dict_offsets`'s contract), and both buffers are + // immutable. + Ok(unsafe { + CompactDictionaryView::new_unchecked(data.dict_bytes().as_slice(), offsets.as_slice()) + }) +} + impl Display for OnPairData { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { write!( f, - "len: {}, bits: {}, dict_bytes_len: {}", + "len: {}, dict_bytes_len: {}", self.len, - self.bits, self.dict_bytes.len() ) } @@ -188,7 +242,6 @@ impl Debug for OnPairData { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { f.debug_struct("OnPairData") .field("len", &self.len) - .field("bits", &self.bits) .field("dict_bytes_len", &self.dict_bytes.len()) .finish() } @@ -197,17 +250,14 @@ impl Debug for OnPairData { impl ArrayHash for OnPairData { fn array_hash(&self, state: &mut H, accuracy: EqMode) { self.dict_bytes.as_host().array_hash(state, accuracy); - state.write_u32(self.bits); } } impl ArrayEq for OnPairData { fn array_eq(&self, other: &Self, accuracy: EqMode) -> bool { - self.bits == other.bits - && self - .dict_bytes - .as_host() - .array_eq(other.dict_bytes.as_host(), accuracy) + self.dict_bytes + .as_host() + .array_eq(other.dict_bytes.as_host(), accuracy) } } @@ -217,7 +267,6 @@ pub struct OnPair; impl OnPair { /// Build an [`OnPairArray`] from already-materialised parts. - #[expect(clippy::too_many_arguments, reason = "every child is a real input")] pub fn try_new( dtype: DType, dict_bytes: BufferHandle, @@ -226,7 +275,6 @@ impl OnPair { codes_offsets: ArrayRef, uncompressed_lengths: ArrayRef, validity: Validity, - bits: u32, ) -> VortexResult { validate_parts( &dtype, @@ -234,10 +282,9 @@ impl OnPair { &codes, &codes_offsets, &uncompressed_lengths, - bits, )?; let len = uncompressed_lengths.len(); - let data = OnPairData::new(dict_bytes, bits, len); + let data = OnPairData::new(dict_bytes, len); let slots = OnPairSlots { dict_offsets, codes, @@ -251,19 +298,25 @@ impl OnPair { }) } - #[expect(clippy::too_many_arguments, reason = "every child is a real input")] + /// Build an [`OnPairArray`] without validation, carrying `data` — and with + /// it the memoized widened `dict_offsets` — from an existing array. + /// + /// # Safety + /// The parts must satisfy the same invariants [`try_new`](Self::try_new) + /// checks. If `data`'s widened-offsets cell is populated (or shared with a + /// live array), `dict_offsets` must hold the same logical offsets the cell + /// was built from. pub(crate) unsafe fn new_unchecked( dtype: DType, - dict_bytes: BufferHandle, + mut data: OnPairData, dict_offsets: ArrayRef, codes: ArrayRef, codes_offsets: ArrayRef, uncompressed_lengths: ArrayRef, validity: Validity, - bits: u32, ) -> OnPairArray { let len = uncompressed_lengths.len(); - let data = OnPairData::new(dict_bytes, bits, len); + data.len = len; let slots = OnPairSlots { dict_offsets, codes, @@ -284,13 +337,11 @@ fn validate_parts( codes: &ArrayRef, codes_offsets: &ArrayRef, uncompressed_lengths: &ArrayRef, - bits: u32, ) -> VortexResult<()> { vortex_ensure!( matches!(dtype, DType::Binary(_) | DType::Utf8(_)), "OnPair arrays must be Binary or Utf8, found {dtype}" ); - vortex_ensure!((9..=16).contains(&bits), "bits {bits} out of range [9, 16]"); if !dict_offsets.dtype().is_int() || dict_offsets.dtype().is_nullable() { vortex_bail!(InvalidArgument: "dict_offsets must be non-nullable integer"); @@ -338,7 +389,6 @@ impl VTable for OnPair { s.codes, s.codes_offsets, s.uncompressed_lengths, - data.bits, )?; if s.uncompressed_lengths.len() != len { vortex_bail!(InvalidArgument: "uncompressed_lengths must have same len as outer array"); @@ -379,6 +429,10 @@ impl VTable for OnPair { ); let mut data = array.data().clone(); data.dict_bytes = buffers[0].clone(); + // The replacement blob may differ from the one the memoized offsets + // were validated against, so drop the (shared) cell rather than + // carry a claim we can no longer prove. + data.dict_offsets = Arc::new(OnceLock::new()); Ok( ArrayParts::new(self.clone(), array.dtype().clone(), array.len(), data) .with_slots(array.slots().iter().cloned().collect()), @@ -391,13 +445,12 @@ impl VTable for OnPair { ) -> VortexResult>> { let dict_size = u32::try_from(array.dict_offsets().len().saturating_sub(1)) .map_err(|_| vortex_err!("OnPair dict_size exceeds u32"))?; - let total_tokens = array.codes().len() as u64; + let codes_len = array.codes().len() as u64; Ok(Some( OnPairMetadata { uncompressed_lengths_ptype: array.uncompressed_lengths().dtype().as_ptype().into(), - bits: array.bits(), dict_size, - total_tokens, + codes_len, dict_offsets_ptype: array.dict_offsets().dtype().as_ptype().into(), codes_ptype: array.codes().dtype().as_ptype().into(), codes_offsets_ptype: array.codes_offsets().dtype().as_ptype().into(), @@ -421,12 +474,11 @@ impl VTable for OnPair { let metadata = OnPairMetadata::decode(metadata)?; let uncompressed_ptype = metadata.get_uncompressed_lengths_ptype()?; - // Slot children. We pass `usize::MAX` for slots whose length we - // don't know up front (`dict_offsets` and `codes`). `codes_offsets` - // has known length `len + 1`. + // Slot children do not persist their own lengths, so metadata records + // the dictionary and code-stream sizes needed to deserialize them. let dict_offsets_len = metadata.dict_size as usize + 1; - let total_tokens = usize::try_from(metadata.total_tokens) - .map_err(|_| vortex_err!("total_tokens {} overflows usize", metadata.total_tokens))?; + let codes_len = usize::try_from(metadata.codes_len) + .map_err(|_| vortex_err!("codes_len {} overflows usize", metadata.codes_len))?; // The cascading compressor may have narrowed any of these integer // children to a tighter ptype; the recorded ptype tells the framework // exactly which dtype to materialise as. @@ -449,7 +501,7 @@ impl VTable for OnPair { let codes = children.get( 1, &DType::Primitive(codes_ptype, Nullability::NonNullable), - total_tokens, + codes_len, )?; let codes_offsets = children.get( 2, @@ -467,7 +519,7 @@ impl VTable for OnPair { other => vortex_bail!(InvalidArgument: "Expected 4 or 5 children, got {other}"), }; - let data = OnPairData::new(buffers[0].clone(), metadata.bits, len); + let data = OnPairData::new(buffers[0].clone(), len); let slots = OnPairSlots { dict_offsets, codes, @@ -534,6 +586,8 @@ impl ValidityVTable for OnPair { /// Convenience methods on top of the macro-generated [`OnPairArraySlotsExt`]. pub trait OnPairArrayExt: OnPairArraySlotsExt { + /// The array's [`Validity`], derived from the optional validity child and + /// the outer dtype's nullability. fn array_validity(&self) -> Validity { child_to_validity( self.as_ref().slots()[OnPairSlots::VALIDITY].as_ref(), diff --git a/encodings/experimental/onpair/src/canonical.rs b/encodings/experimental/onpair/src/canonical.rs index 20a058e7f90..f5e55467eff 100644 --- a/encodings/experimental/onpair/src/canonical.rs +++ b/encodings/experimental/onpair/src/canonical.rs @@ -2,14 +2,13 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors // //! Convert an [`OnPairArray`] to its canonical `VarBinViewArray` by handing -//! the materialised parts to `onpair::decompress_into`. +//! the materialised parts to `onpair::try_decode_into`. //! //! [`OnPairArray`]: crate::OnPairArray use std::sync::Arc; use num_traits::AsPrimitive; -use onpair::Parts; use vortex_array::ArrayRef; use vortex_array::ArrayView; use vortex_array::ExecutionCtx; @@ -25,9 +24,11 @@ use vortex_buffer::ByteBuffer; use vortex_buffer::ByteBufferMut; use vortex_error::VortexResult; use vortex_error::vortex_ensure; +use vortex_error::vortex_panic; use crate::OnPair; use crate::OnPairArraySlotsExt; +use crate::array::dict_view; use crate::decode::code_boundary_at; use crate::decode::collect_widened; @@ -89,21 +90,22 @@ pub(crate) fn onpair_decode_views( // contiguous decoder walks `codes` in order and never reads the per-row // boundaries, so an empty boundary slice is sound. let codes = collect_widened::(&array.codes().slice(code_start..code_end)?, ctx)?; - let dict_offsets = collect_widened::(array.dict_offsets(), ctx)?; + let dict = dict_view(array, ctx)?; let mut out_bytes = ByteBufferMut::with_capacity(total_size); - let written = onpair::decompress_into( - Parts { - dict_bytes: array.dict_bytes().as_slice(), - dict_offsets: dict_offsets.as_slice(), - bits: array.bits(), - codes: codes.as_slice(), - }, - out_bytes.spare_capacity_mut(), - ); - debug_assert_eq!(written, total_size); - // SAFETY: `decompress_into` initialised exactly `written` bytes of the - // spare capacity reserved above. + let written = + match onpair::try_decode_into(codes.as_slice(), dict, out_bytes.spare_capacity_mut()) { + Ok(written) => written, + Err(_) => { + vortex_panic!("OnPair codes decode to more bytes than uncompressed_lengths records") + } + }; + if written != total_size { + vortex_panic!( + "OnPair codes decoded to {written} bytes but uncompressed_lengths records {total_size}" + ); + } + // SAFETY: `try_decode_into` initialised exactly `written` bytes. unsafe { out_bytes.set_len(written) }; match_each_integer_ptype!(lengths.ptype(), |P| { diff --git a/encodings/experimental/onpair/src/compress.rs b/encodings/experimental/onpair/src/compress.rs index 3616ec7e717..a2e63f2c8a2 100644 --- a/encodings/experimental/onpair/src/compress.rs +++ b/encodings/experimental/onpair/src/compress.rs @@ -4,7 +4,6 @@ //! Train + compress entry points for the OnPair encoding. use onpair::Config; -use onpair::Offset; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; @@ -12,7 +11,7 @@ use vortex_array::arrays::ConstantArray; use vortex_array::arrays::VarBinViewArray; use vortex_array::arrays::varbinview::BinaryView; use vortex_array::buffer::BufferHandle; -use vortex_array::validity::Validity; +use vortex_array::scalar::Scalar; use vortex_buffer::Alignment; use vortex_buffer::Buffer; use vortex_buffer::BufferMut; @@ -24,7 +23,6 @@ use vortex_error::vortex_err; use vortex_mask::AllOr; use crate::OnPair; -use crate::OnPairArray; /// Default OnPair training configuration: 12-bit codes ("dict-12"). pub const DEFAULT_DICT12_CONFIG: Config = Config { @@ -32,34 +30,32 @@ pub const DEFAULT_DICT12_CONFIG: Config = Config { ..onpair::DEFAULT_CONFIG }; -fn onpair_compress_varbinview( - array: VarBinViewArray, +/// Compress any [`ArrayRef`] whose canonical form is a string array. +/// +/// All-null inputs are returned as a [`ConstantArray`]. +pub fn onpair_compress( + array: &ArrayRef, config: Config, ctx: &mut ExecutionCtx, -) -> VortexResult -where - O: Offset, -{ +) -> VortexResult { + let array = array.clone().execute::(ctx)?; let len = array.len(); - let mask = array.validity()?.execute_mask(len, ctx)?; - if mask.all_false() { - return OnPair::try_new( - array.dtype().clone(), - BufferHandle::new_host(ByteBuffer::empty()), - ConstantArray::new(0, len).into_array(), - ConstantArray::new(0u16, len).into_array(), - ConstantArray::new(0u32, len + 1).into_array(), - ConstantArray::new(0i32, len).into_array(), - Validity::AllInvalid, - 9, - ); + let validity = array.validity()?; + let mask = validity.execute_mask(len, ctx)?; + if matches!(mask.bit_buffer(), AllOr::None) { + // CascadingCompressor handles this earlier, but direct callers can reach it. + return Ok(ConstantArray::new(Scalar::null(array.dtype().clone()), len).into_array()); } - let mut flat: Vec = Vec::with_capacity(len * 16); - let mut offsets: Vec = Vec::with_capacity(len + 1); - let mut uncompressed_lengths: BufferMut = BufferMut::with_capacity(len); - offsets.push(O::from_usize(0)); let views = array.views(); + let flat_bytes: usize = views.iter().map(|v| v.len() as usize).sum(); + + // TODO(francesco): we flatten because onpair training needs a contiguous `(bytes, offsets)` + // pair. Allowing onpair to train on a slice-of-slices would let us skip this copy. + let mut flat: Vec = Vec::with_capacity(flat_bytes); + let mut offsets: Vec = Vec::with_capacity(len + 1); + let mut uncompressed_lengths: BufferMut = BufferMut::with_capacity(len); + offsets.push(0); let buffers = array .data_buffers() .as_ref() @@ -72,24 +68,22 @@ where for view in views { let bytes = view_bytes(view, &buffers); flat.extend_from_slice(bytes); - offsets.push(O::from_usize(flat.len())); + offsets.push(u64::try_from(flat.len()).vortex_expect("offset must fit in u64")); uncompressed_lengths .push(i32::try_from(view.len()).vortex_expect("must fit in i32")); } } - AllOr::None => { - unreachable!("all_false() should have been caught earlier"); - } + AllOr::None => unreachable!("all-null input handled above"), AllOr::Some(validity) => { for (view, valid) in views.iter().zip(validity.iter()) { if valid { let bytes = view_bytes(view, &buffers); flat.extend_from_slice(bytes); - offsets.push(O::from_usize(flat.len())); + offsets.push(u64::try_from(flat.len()).vortex_expect("offset must fit in u64")); uncompressed_lengths .push(i32::try_from(view.len()).vortex_expect("must fit in i32")); } else { - offsets.push(O::from_usize(flat.len())); + offsets.push(u64::try_from(flat.len()).vortex_expect("offset must fit in u64")); uncompressed_lengths.push(0); } } @@ -98,25 +92,29 @@ where let column = onpair::compress(&flat, &offsets, config) .map_err(|e| vortex_err!("OnPair compress failed: {e}"))?; - let bits = column.bits; - let dict_bytes = dict_bytes_to_buffer(column.dict_bytes); - let codes_offsets = - build_codes_offsets(&column.codes, &column.dict_offsets, &offsets)?.into_array(); - let codes = Buffer::from(column.codes).into_array(); - let dict_offsets = Buffer::from(column.dict_offsets).into_array(); + let (dict, codes, row_offsets) = column.into_raw(); + let (dict_bytes, dict_offsets) = dict.into_raw(); + let codes_offsets = codes_offsets_array(&row_offsets); + let codes = Buffer::from(codes).into_array(); + // The `dict_offsets` child and the memoized widened-offsets cell share + // this buffer, so seeding below costs no copy. + let dict_offsets = Buffer::from(dict_offsets); let uncompressed_lengths = uncompressed_lengths.into_array(); - OnPair::try_new( + let encoded = OnPair::try_new( array.dtype().clone(), - dict_bytes, - dict_offsets, + dict_bytes_to_buffer(dict_bytes), + dict_offsets.clone().into_array(), codes, codes_offsets, uncompressed_lengths, - array.validity()?, - bits, - ) + validity, + )?; + // SAFETY: the trainer's dictionary is conformant by construction, and + // `dict_offsets` is exactly the u32 child attached above. + unsafe { encoded.init_dict_offsets(dict_offsets) }; + Ok(encoded.into_array()) } fn view_bytes<'a>(view: &'a BinaryView, buffers: &'a [&ByteBuffer]) -> &'a [u8] { @@ -128,66 +126,33 @@ fn view_bytes<'a>(view: &'a BinaryView, buffers: &'a [&ByteBuffer]) -> &'a [u8] } } -/// Lift compressed dictionary bytes into the Vortex buffer slot. fn dict_bytes_to_buffer(dict_bytes: Vec) -> BufferHandle { - // Pad the dictionary blob with MAX_TOKEN_SIZE zero bytes so the - // over-copy decoder can issue a fixed 16-byte load for every token - // without risking an OOB read on the last entry. - // // Align dict_bytes to 8 bytes so the segment that ultimately holds the - // OnPair tree starts at an 8-aligned in-memory address. Without this - // anchor, the per-buffer padding the serializer inserts is only - // *relative* to the segment start; if the segment lands at a u8-aligned - // heap address, downstream `PrimitiveArray::deserialize` panics - // with `Misaligned buffer cannot be used to build PrimitiveArray of u32`. - let mut padded = ByteBufferMut::with_capacity_aligned( - dict_bytes.len() + onpair::MAX_TOKEN_SIZE, - Alignment::new(8), - ); - padded.extend_from_slice(&dict_bytes); - unsafe { padded.push_n_unchecked(0, dict_bytes.len() + onpair::MAX_TOKEN_SIZE - padded.len()) }; - BufferHandle::new_host(padded.freeze()) + // OnPair tree starts at an 8-aligned in-memory address. Without this anchor, + // downstream primitive children may deserialize from a misaligned segment. + let mut aligned = ByteBufferMut::with_capacity_aligned(dict_bytes.len(), Alignment::new(8)); + aligned.extend_from_slice(&dict_bytes); + BufferHandle::new_host(aligned.freeze()) } -/// Reconstruct the per-row `codes_offsets` from the flat `codes`, the -/// dictionary `dict_offsets` (token byte lengths) and the per-row decoded byte -/// boundaries. Returns `nrows + 1` cumulative code counts (`u32`). -// TODO(joe): can we compute this while compressing the array, yes but a worse API. -fn build_codes_offsets( - codes: &[u16], - dict_offsets: &[u32], - row_byte_offsets: &[O], -) -> VortexResult> { - let nrows = row_byte_offsets.len() - 1; - let mut codes_offsets = BufferMut::with_capacity(nrows + 1); - codes_offsets.push(0u32); - let mut decoded_bytes: u64 = 0; - let mut code_idx: usize = 0; - for r in 0..nrows { - let target = row_byte_offsets[r + 1] - .to_usize() - .ok_or_else(|| vortex_err!("OnPair row byte offset does not fit usize"))? - as u64; - while decoded_bytes < target { - let code = codes[code_idx] as usize; - decoded_bytes += u64::from(dict_offsets[code + 1] - dict_offsets[code]); - code_idx += 1; - } - codes_offsets.push( - u32::try_from(code_idx) - .map_err(|_| vortex_err!("OnPair: code boundary {code_idx} does not fit u32"))?, - ); +/// Build the `codes_offsets` child from the library's per-row code boundaries, +/// storing the narrowest of `u32`/`u64` that holds the largest boundary. +/// `row_offsets` is non-decreasing, so its last entry is that maximum and one +/// bound check picks the width. `u32` covers the common case (the cascading +/// compressor narrows it further to `u16`/`u8`); `u64` engages only when a +/// single chunk carries more than `u32::MAX` tokens, matching the `u64` byte +/// offsets accepted at compression. +fn codes_offsets_array(row_offsets: &[u64]) -> ArrayRef { + let total_tokens = row_offsets.last().copied().unwrap_or(0); + if u32::try_from(total_tokens).is_ok() { + Buffer::from( + row_offsets + .iter() + .map(|&o| u32::try_from(o).vortex_expect("code boundary fits u32")) + .collect::>(), + ) + .into_array() + } else { + Buffer::from(row_offsets.to_vec()).into_array() } - Ok(codes_offsets.freeze()) -} - -/// Compress any [`ArrayRef`] whose canonical form is a string array, by first -/// canonicalising to `VarBinViewArray`. -pub fn onpair_compress( - array: &ArrayRef, - config: Config, - ctx: &mut ExecutionCtx, -) -> VortexResult { - let view = array.clone().execute::(ctx)?; - onpair_compress_varbinview::(view, config, ctx) } diff --git a/encodings/experimental/onpair/src/compute/cast.rs b/encodings/experimental/onpair/src/compute/cast.rs index 93e1fdd8f8a..c31893f82fd 100644 --- a/encodings/experimental/onpair/src/compute/cast.rs +++ b/encodings/experimental/onpair/src/compute/cast.rs @@ -3,16 +3,38 @@ use vortex_array::ArrayRef; use vortex_array::ArrayView; +use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::dtype::DType; +use vortex_array::scalar_fn::fns::cast::CastKernel; use vortex_array::scalar_fn::fns::cast::CastReduce; +use vortex_array::validity::Validity; use vortex_error::VortexResult; use crate::OnPair; use crate::OnPairArraySlotsExt; -/// Cast between `Utf8` and `Binary` (or adjust nullability) without touching -/// any of the encoded payload — we only rewrap into a new outer DType. +/// Adjust nullability without touching any encoded payload — we only rewrap +/// into a new outer DType. +fn build_with_validity( + array: ArrayView<'_, OnPair>, + dtype: &DType, + validity: Validity, +) -> ArrayRef { + unsafe { + OnPair::new_unchecked( + dtype.clone(), + array.data().clone(), + array.dict_offsets().clone(), + array.codes().clone(), + array.codes_offsets().clone(), + array.uncompressed_lengths().clone(), + validity, + ) + } + .into_array() +} + impl CastReduce for OnPair { fn cast(array: ArrayView<'_, Self>, dtype: &DType) -> VortexResult> { if !array.dtype().eq_ignore_nullability(dtype) { @@ -24,20 +46,69 @@ impl CastReduce for OnPair { else { return Ok(None); }; - Ok(Some( - unsafe { - OnPair::new_unchecked( - dtype.clone(), - array.dict_bytes_handle().clone(), - array.dict_offsets().clone(), - array.codes().clone(), - array.codes_offsets().clone(), - array.uncompressed_lengths().clone(), - new_validity, - array.bits(), - ) - } - .into_array(), - )) + Ok(Some(build_with_validity(array, dtype, new_validity))) + } +} + +impl CastKernel for OnPair { + fn cast( + array: ArrayView<'_, Self>, + dtype: &DType, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + if !array.dtype().eq_ignore_nullability(dtype) { + return Ok(None); + } + let new_validity = array.array().validity()?.cast_nullability( + dtype.nullability(), + array.array().len(), + ctx, + )?; + Ok(Some(build_with_validity(array, dtype, new_validity))) + } +} + +#[cfg(test)] +mod tests { + use std::sync::LazyLock; + + use rstest::rstest; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::arrays::VarBinArray; + use vortex_array::compute::conformance::cast::test_cast_conformance; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_error::VortexResult; + use vortex_session::VortexSession; + + use crate::compress::DEFAULT_DICT12_CONFIG; + use crate::compress::onpair_compress; + + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + + #[rstest] + #[case(VarBinArray::from_iter( + vec![Some("hello"), Some("world"), Some("hello world")], + DType::Utf8(Nullability::NonNullable) + ))] + #[case(VarBinArray::from_iter( + vec![Some("foo"), None, Some("bar"), Some("foobar")], + DType::Utf8(Nullability::Nullable) + ))] + #[case(VarBinArray::from_iter( + vec![Some("test")], + DType::Utf8(Nullability::NonNullable) + ))] + fn test_cast_onpair_conformance(#[case] array: VarBinArray) -> VortexResult<()> { + let array = array.into_array(); + let mut ctx = SESSION.create_execution_ctx(); + let onpair = onpair_compress(&array, DEFAULT_DICT12_CONFIG, &mut ctx)?; + test_cast_conformance(&onpair.into_array(), &mut ctx); + Ok(()) } } diff --git a/encodings/experimental/onpair/src/compute/compare.rs b/encodings/experimental/onpair/src/compute/compare.rs index 939a1a2fb42..bb10fcfd6aa 100644 --- a/encodings/experimental/onpair/src/compute/compare.rs +++ b/encodings/experimental/onpair/src/compute/compare.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use onpair::search; use vortex_array::ArrayRef; use vortex_array::ArrayView; use vortex_array::ExecutionCtx; @@ -17,6 +18,8 @@ use vortex_error::VortexResult; use crate::OnPair; use crate::OnPairArraySlotsExt; +use crate::array::dict_view; +use crate::decode::collect_codes_window; impl CompareKernel for OnPair { fn compare( @@ -28,29 +31,51 @@ impl CompareKernel for OnPair { let Some(constant) = rhs.as_constant() else { return Ok(None); }; - let is_empty = match constant.dtype() { - DType::Utf8(_) => constant.as_utf8().is_empty(), - DType::Binary(_) => constant.as_binary().is_empty(), + let needle = match constant.dtype() { + DType::Utf8(_) => constant + .as_utf8() + .value() + .map(|value| value.as_bytes().to_vec()), + DType::Binary(_) => constant + .as_binary() + .value() + .map(|value| value.as_slice().to_vec()), _ => return Ok(None), }; - if is_empty != Some(true) { + let Some(needle) = needle else { return Ok(None); - } - - let lengths = lhs.uncompressed_lengths(); - let buffer = match operator { - // every value is greater than an empty string - CompareOperator::Gte => BitBuffer::new_set(lhs.len()), - // no value is less than an empty string - CompareOperator::Lt => BitBuffer::new_unset(lhs.len()), - _ => lengths - .binary( - ConstantArray::new(Scalar::zero_value(lengths.dtype()), lengths.len()) - .into_array(), - operator.into(), - )? - .execute(ctx)?, }; + + let buffer = if needle.is_empty() { + let lengths = lhs.uncompressed_lengths(); + match operator { + // every value is greater than an empty string + CompareOperator::Gte => BitBuffer::new_set(lhs.len()), + // no value is less than an empty string + CompareOperator::Lt => BitBuffer::new_unset(lhs.len()), + _ => lengths + .binary( + ConstantArray::new(Scalar::zero_value(lengths.dtype()), lengths.len()) + .into_array(), + operator.into(), + )? + .execute(ctx)?, + } + } else { + if !matches!(operator, CompareOperator::Eq | CompareOperator::NotEq) { + return Ok(None); + } + + let dict = dict_view(lhs, ctx)?; + let query = search::tokenize(&needle, dict); + let window = collect_codes_window(lhs, ctx)?; + + let negated = operator == CompareOperator::NotEq; + BitBuffer::collect_bool(lhs.len(), |i| { + search::equals(window.row(i), &query) != negated + }) + }; + Ok(Some( BoolArray::new( buffer, @@ -69,7 +94,6 @@ mod tests { use rstest::rstest; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; - use vortex_array::array_session; use vortex_array::arrays::BoolArray; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::VarBinArray; @@ -78,14 +102,22 @@ mod tests { use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::scalar::Scalar; + use vortex_array::scalar_fn::fns::binary::CompareKernel; + use vortex_array::scalar_fn::fns::operators::CompareOperator; use vortex_array::scalar_fn::fns::operators::Operator; use vortex_error::VortexResult; + use vortex_error::vortex_err; use vortex_session::VortexSession; + use crate::OnPair; use crate::compress::DEFAULT_DICT12_CONFIG; use crate::compress::onpair_compress; - static SESSION: LazyLock = LazyLock::new(array_session); + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); #[cfg_attr(miri, ignore)] #[rstest] @@ -142,4 +174,106 @@ mod tests { ); Ok(()) } + + #[cfg_attr(miri, ignore)] + #[test] + fn compare_nonempty_string_nullable() -> VortexResult<()> { + let input = VarBinArray::from_iter( + [Some("hello"), None, Some("world"), Some("hello")], + DType::Utf8(Nullability::Nullable), + ); + let mut ctx = SESSION.create_execution_ctx(); + let arr = onpair_compress(input.as_array(), DEFAULT_DICT12_CONFIG, &mut ctx)?.into_array(); + let rhs = ConstantArray::new("hello", arr.len()).into_array(); + + let eq = arr + .binary(rhs.clone(), Operator::Eq)? + .execute::(&mut ctx)?; + assert_arrays_eq!( + &eq, + &BoolArray::from_iter([Some(true), None, Some(false), Some(true)]), + &mut ctx + ); + + let neq = arr + .binary(rhs, Operator::NotEq)? + .execute::(&mut ctx)?; + assert_arrays_eq!( + &neq, + &BoolArray::from_iter([Some(false), None, Some(true), Some(false)]), + &mut ctx + ); + Ok(()) + } + + /// Call `CompareKernel::compare` directly and verify it returns `Some` + /// (i.e. the kernel handles the constant needle rather than silently + /// falling back to canonical decompression), and that ordering against a + /// non-empty needle declines — code sequences are not order-preserving. + #[cfg_attr(miri, ignore)] + #[test] + fn compare_kernel_handles_constant_needle() -> VortexResult<()> { + let input = VarBinArray::from_iter( + [Some("hello"), Some("world"), Some("hello"), Some("he")], + DType::Utf8(Nullability::NonNullable), + ); + let mut ctx = SESSION.create_execution_ctx(); + let arr = onpair_compress(input.as_array(), DEFAULT_DICT12_CONFIG, &mut ctx)? + .try_downcast::() + .map_err(|array| vortex_err!("expected OnPair array, got {}", array.encoding_id()))?; + let rhs = ConstantArray::new("hello", arr.len()).into_array(); + + let eq = + ::compare(arr.as_view(), &rhs, CompareOperator::Eq, &mut ctx)? + .expect("OnPair CompareKernel should handle a constant needle"); + assert_arrays_eq!( + eq, + BoolArray::from_iter([true, false, true, false]), + &mut ctx + ); + + let lt = + ::compare(arr.as_view(), &rhs, CompareOperator::Lt, &mut ctx)?; + assert!( + lt.is_none(), + "ordering compare against a non-empty needle must fall back to canonical" + ); + Ok(()) + } + + /// The equality path resolves row windows relative to `codes_offsets[0]`, + /// which is nonzero for a sliced array. Exercise the kernel directly on a + /// slice so the windowed row arithmetic is covered. + #[cfg_attr(miri, ignore)] + #[test] + fn compare_kernel_on_sliced_array() -> VortexResult<()> { + let input = VarBinArray::from_iter( + [ + Some("aardvark"), + Some("hello"), + Some("world"), + Some("hello"), + Some("zebra"), + ], + DType::Utf8(Nullability::NonNullable), + ); + let mut ctx = SESSION.create_execution_ctx(); + let arr = onpair_compress(input.as_array(), DEFAULT_DICT12_CONFIG, &mut ctx)?.into_array(); + let sliced = arr.slice(1..4)?; + assert!(sliced.is::(), "slice dropped OnPair encoding"); + let sliced = sliced + .try_downcast::() + .map_err(|_| vortex_err!("sliced array was not OnPair"))?; + + let rhs = ConstantArray::new("hello", sliced.len()).into_array(); + let eq = ::compare( + sliced.as_view(), + &rhs, + CompareOperator::Eq, + &mut ctx, + )? + .expect("OnPair CompareKernel should handle a constant needle"); + assert_arrays_eq!(eq, BoolArray::from_iter([true, false, true]), &mut ctx); + Ok(()) + } } diff --git a/encodings/experimental/onpair/src/compute/filter.rs b/encodings/experimental/onpair/src/compute/filter.rs index c26f3eeeacd..1543bff136a 100644 --- a/encodings/experimental/onpair/src/compute/filter.rs +++ b/encodings/experimental/onpair/src/compute/filter.rs @@ -63,13 +63,12 @@ impl FilterKernel for OnPair { unsafe { OnPair::new_unchecked( array.dtype().clone(), - array.dict_bytes_handle().clone(), + array.data().clone(), array.dict_offsets().clone(), filtered_codes.elements().clone(), filtered_codes.offsets().clone(), uncompressed_lengths, validity, - array.bits(), ) } .into_array(), diff --git a/encodings/experimental/onpair/src/compute/slice.rs b/encodings/experimental/onpair/src/compute/slice.rs index fcfebf413bf..49d840f2029 100644 --- a/encodings/experimental/onpair/src/compute/slice.rs +++ b/encodings/experimental/onpair/src/compute/slice.rs @@ -28,13 +28,12 @@ impl SliceReduce for OnPair { unsafe { OnPair::new_unchecked( array.dtype().clone(), - array.dict_bytes_handle().clone(), + array.data().clone(), array.dict_offsets().clone(), array.codes().clone(), codes_offsets, uncompressed_lengths, validity, - array.bits(), ) } .into_array(), diff --git a/encodings/experimental/onpair/src/decode.rs b/encodings/experimental/onpair/src/decode.rs index fd0b02b0284..96add100e36 100644 --- a/encodings/experimental/onpair/src/decode.rs +++ b/encodings/experimental/onpair/src/decode.rs @@ -1,19 +1,25 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors // -//! Helpers for turning [`OnPair`](super::OnPair) slot children into the inputs the upstream +//! Helpers for turning [`OnPair`] slot children into the inputs the upstream //! `onpair` decoder consumes. use vortex_array::ArrayRef; +use vortex_array::ArrayView; use vortex_array::ExecutionCtx; use vortex_array::arrays::PrimitiveArray; use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; use vortex_array::dtype::NativePType; use vortex_buffer::Buffer; +use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_error::vortex_ensure; use vortex_error::vortex_err; +use crate::OnPair; +use crate::OnPairArraySlotsExt; + /// Canonicalise a slot child to the decoder's native primitive width. pub(crate) fn collect_widened( arr: &ArrayRef, @@ -41,3 +47,67 @@ pub(crate) fn code_boundary_at( .as_::() .ok_or_else(|| vortex_err!("OnPair codes_offsets[{index}] is null")) } + +/// A validated, materialised window over an array's `codes`: the widened +/// per-row `codes_offsets` boundaries plus the codes they bound. +/// +/// `slice` keeps the full `codes` child and only narrows `codes_offsets`, so +/// for a sliced array the window starts at `offsets[0] > 0`; [`row`] resolves +/// row indices relative to that start. Built once per query by the +/// compressed-domain compare kernel. +/// +/// [`row`]: CodesWindow::row +pub(crate) struct CodesWindow { + offsets: Buffer, + codes: Buffer, + code_start: usize, +} + +impl CodesWindow { + /// The codes for row `i`. + pub(crate) fn row(&self, i: usize) -> &[u16] { + &self.codes[self.local(i)..self.local(i + 1)] + } + + /// Offset `i` rebased into the window's local `codes` slice. Offsets are + /// bounded by `codes.len()` (a `usize`), so the conversion never truncates. + fn local(&self, i: usize) -> usize { + usize::try_from(self.offsets[i]).vortex_expect("code offset fits usize") - self.code_start + } +} + +/// Materialise the [`CodesWindow`] for every row of `array`: offsets must be +/// nondecreasing and end within the `codes` child. Codes themselves are +/// trusted to the upstream decoder/search primitives, which bounds-check them +/// in-loop and panic on a malformed value. +pub(crate) fn collect_codes_window( + array: ArrayView<'_, OnPair>, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let len = array.len(); + let offsets = collect_widened::(array.codes_offsets(), ctx)?; + vortex_ensure!( + offsets.len() == len + 1, + "OnPair codes_offsets has {} entries, expected len + 1 = {}", + offsets.len(), + len + 1 + ); + vortex_ensure!( + offsets.is_sorted(), + "OnPair codes_offsets must be nondecreasing" + ); + let code_start = usize::try_from(offsets[0]).vortex_expect("code offset fits usize"); + let code_end = usize::try_from(offsets[len]).vortex_expect("code offset fits usize"); + vortex_ensure!( + code_end <= array.codes().len(), + "OnPair codes_offsets end {} exceeds codes len {}", + code_end, + array.codes().len() + ); + let codes = collect_widened::(&array.codes().slice(code_start..code_end)?, ctx)?; + Ok(CodesWindow { + offsets, + codes, + code_start, + }) +} diff --git a/encodings/experimental/onpair/src/kernel.rs b/encodings/experimental/onpair/src/kernel.rs index 8863d750a72..e6ae93248b5 100644 --- a/encodings/experimental/onpair/src/kernel.rs +++ b/encodings/experimental/onpair/src/kernel.rs @@ -10,13 +10,15 @@ use vortex_array::scalar_fn::fns::binary::Binary; use vortex_array::scalar_fn::fns::binary::CompareExecuteAdaptor; use vortex_array::scalar_fn::fns::byte_length::ByteLength; use vortex_array::scalar_fn::fns::byte_length::ByteLengthExecuteAdaptor; +use vortex_array::scalar_fn::fns::cast::Cast; +use vortex_array::scalar_fn::fns::cast::CastExecuteAdaptor; use vortex_session::VortexSession; use crate::OnPair; -// TODO: implement ListExecute & TakeExecute for OnPair pub(super) fn initialize(session: &VortexSession) { let kernels = session.kernels(); + kernels.register_execute_parent_kernel(Cast.id(), OnPair, CastExecuteAdaptor(OnPair)); kernels.register_execute_parent_kernel(Filter.id(), OnPair, FilterExecuteAdaptor(OnPair)); kernels.register_execute_parent_kernel(Binary.id(), OnPair, CompareExecuteAdaptor(OnPair)); kernels.register_execute_parent_kernel( diff --git a/encodings/experimental/onpair/src/lib.rs b/encodings/experimental/onpair/src/lib.rs index 11b22f63bc1..4c2ae4529cc 100644 --- a/encodings/experimental/onpair/src/lib.rs +++ b/encodings/experimental/onpair/src/lib.rs @@ -2,11 +2,10 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors //! Vortex string array backed by the [OnPair][onpair] short-string -//! compression library, with `cast` and `filter` pushdown. +//! compression library, with pushdown for common string operations. //! -//! The default training preset is `dict-12` (12 bits per token, dictionary -//! capped at 4 096 entries). See [`onpair_compress`] for the entry point and -//! [`OnPairArray`] for the resulting array type. +//! See [`onpair_compress`] for the compression entry point and [`OnPairArray`] +//! for the encoded representation of non-null inputs. //! //! [onpair]: https://arxiv.org/abs/2508.02280 @@ -23,9 +22,9 @@ mod tests; pub use array::*; pub use compress::*; -pub use onpair::Bits; pub use onpair::Config; pub use onpair::Error as OnPairError; +pub use onpair::MaxDictBits; pub use onpair::Threshold; use vortex_array::session::ArraySessionExt; use vortex_session::VortexSession; diff --git a/encodings/experimental/onpair/src/ops.rs b/encodings/experimental/onpair/src/ops.rs index a6e097bbfd4..728e5a0e6f2 100644 --- a/encodings/experimental/onpair/src/ops.rs +++ b/encodings/experimental/onpair/src/ops.rs @@ -1,7 +1,6 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use onpair::Parts; use vortex_array::ArrayView; use vortex_array::ExecutionCtx; use vortex_array::arrays::varbin::varbin_scalar; @@ -10,9 +9,11 @@ use vortex_array::vtable::OperationsVTable; use vortex_buffer::ByteBuffer; use vortex_error::VortexResult; use vortex_error::vortex_err; +use vortex_error::vortex_panic; use crate::OnPair; use crate::OnPairArraySlotsExt; +use crate::array::dict_view; use crate::decode::code_boundary_at; use crate::decode::collect_widened; @@ -32,13 +33,7 @@ impl OperationsVTable for OnPair { let row_end = code_boundary_at(codes_offsets, index + 1, ctx)?; let codes = collect_widened::(&array.codes().slice(row_start..row_end)?, ctx)?; - let dict_offsets = collect_widened::(array.dict_offsets(), ctx)?; - let parts = Parts { - dict_bytes: array.dict_bytes().as_slice(), - dict_offsets: dict_offsets.as_slice(), - bits: array.bits(), - codes: codes.as_slice(), - }; + let dict = dict_view(array, ctx)?; // The per-row decoded length is recorded in the `uncompressed_lengths` // child, so read it directly instead of asking the decoder to compute it. @@ -49,10 +44,19 @@ impl OperationsVTable for OnPair { .as_::() .ok_or_else(|| vortex_err!("OnPair uncompressed_lengths[{index}] is null"))?; let mut buf: Vec = Vec::with_capacity(len); - let written = onpair::decompress_into(parts, buf.spare_capacity_mut()); - debug_assert_eq!(written, len); - // SAFETY: `decompress_into` initialised `written` bytes of the spare - // capacity reserved above. + let written = + match onpair::try_decode_into(codes.as_slice(), dict, buf.spare_capacity_mut()) { + Ok(written) => written, + Err(_) => vortex_panic!( + "OnPair row {index} decodes to more bytes than uncompressed_lengths records" + ), + }; + if written != len { + vortex_panic!( + "OnPair row {index} decoded to {written} bytes but uncompressed_lengths records {len}" + ); + } + // SAFETY: `try_decode_into` initialised exactly `written` bytes. unsafe { buf.set_len(written) }; Ok(varbin_scalar(ByteBuffer::from(buf), array.dtype())) } diff --git a/encodings/experimental/onpair/src/tests.rs b/encodings/experimental/onpair/src/tests.rs index 4f2124025f2..20aa7b4e061 100644 --- a/encodings/experimental/onpair/src/tests.rs +++ b/encodings/experimental/onpair/src/tests.rs @@ -3,17 +3,25 @@ use std::sync::LazyLock; +use onpair::CompactDictionaryView; use prost::Message; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; +use vortex_array::arrays::BoolArray; +use vortex_array::arrays::Constant; +use vortex_array::arrays::ConstantArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::VarBinArray; use vortex_array::arrays::VarBinViewArray; use vortex_array::arrays::filter::FilterKernel; +use vortex_array::assert_arrays_eq; +use vortex_array::buffer::BufferHandle; +use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::match_each_integer_ptype; +use vortex_array::scalar_fn::fns::operators::Operator; use vortex_array::test_harness::check_metadata; use vortex_array::validity::Validity; use vortex_buffer::BufferMut; @@ -28,6 +36,17 @@ use crate::compress::onpair_compress; static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +fn compress_onpair( + array: &vortex_array::ArrayRef, + ctx: &mut vortex_array::ExecutionCtx, +) -> vortex_error::VortexResult { + onpair_compress(array, DEFAULT_DICT12_CONFIG, ctx)? + .try_downcast::() + .map_err(|array| { + vortex_error::vortex_err!("expected OnPair array, got {}", array.encoding_id()) + }) +} + fn sample_input() -> VarBinArray { VarBinArray::from_iter( [ @@ -41,6 +60,68 @@ fn sample_input() -> VarBinArray { ) } +#[test] +fn test_onpair_rejects_100k_token_dictionary() -> vortex_error::VortexResult<()> { + let num_tokens = 100_000usize; + let mut tokens = Vec::with_capacity(num_tokens); + tokens.extend((u8::MIN..=u8::MAX).map(|byte| vec![byte])); + let additional_tokens = u32::try_from(num_tokens - tokens.len())?; + tokens.extend((0..additional_tokens).map(|value| { + let bytes = value.to_be_bytes(); + bytes[1..].to_vec() + })); + tokens.sort_unstable(); + + let mut dict_bytes = Vec::new(); + let mut dict_offsets = Vec::with_capacity(num_tokens + 1); + dict_offsets.push(0u32); + for token in tokens { + dict_bytes.extend_from_slice(&token); + dict_offsets.push(u32::try_from(dict_bytes.len())?); + } + dict_bytes.resize(dict_bytes.len() + onpair::MAX_TOKEN_SIZE, 0); + + assert!(CompactDictionaryView::validate(&dict_bytes, &dict_offsets).is_err()); + Ok(()) +} + +/// Dictionary content is validated lazily — construction stays lightweight, +/// and a corrupt dictionary is rejected by the first operation that decodes +/// or searches through it, including on derived (sliced) arrays. +#[cfg_attr(miri, ignore)] +#[test] +fn test_corrupt_dictionary_rejected_on_first_use() -> vortex_error::VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let arr = compress_onpair(&sample_input().into_array(), &mut ctx)?; + let view = arr.as_view(); + + // Chop one byte off the blob so the trailing read-padding invariant fails. + let truncated = view.dict_bytes().slice(0..view.dict_bytes().len() - 1); + let corrupt = OnPair::try_new( + view.dtype().clone(), + BufferHandle::new_host(truncated), + view.dict_offsets().clone(), + view.codes().clone(), + view.codes_offsets().clone(), + view.uncompressed_lengths().clone(), + view.array_validity(), + )?; + + let sliced = corrupt.clone().into_array().slice(1..3)?; + // Canonical decode, random access, and ops on a derived slice all pass + // through the same validation door. + assert!( + corrupt + .clone() + .into_array() + .execute::(&mut ctx) + .is_err() + ); + assert!(corrupt.into_array().execute_scalar(0, &mut ctx).is_err()); + assert!(sliced.execute::(&mut ctx).is_err()); + Ok(()) +} + #[cfg_attr(miri, ignore)] #[test] fn test_onpair_metadata_golden() { @@ -48,9 +129,8 @@ fn test_onpair_metadata_golden() { "onpair.metadata", &OnPairMetadata { uncompressed_lengths_ptype: PType::I32 as i32, - bits: 12, dict_size: 4096, - total_tokens: 128_000, + codes_len: 128_000, dict_offsets_ptype: PType::U32 as i32, codes_ptype: PType::U16 as i32, codes_offsets_ptype: PType::U32 as i32, @@ -88,6 +168,63 @@ fn test_onpair_roundtrip() -> vortex_error::VortexResult<()> { Ok(()) } +/// The `u64` `codes_offsets` branch only engages past `u32::MAX` tokens (a +/// multi-GiB chunk), so exercise the read path by widening a small array's +/// `codes_offsets` child to `u64` by hand and asserting the decode paths +/// (canonical and the compressed-domain equality compare, which builds a +/// `CodesWindow`) treat it identically to the default `u32` width. +#[cfg_attr(miri, ignore)] +#[test] +fn test_onpair_u64_codes_offsets() -> vortex_error::VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let narrow = compress_onpair(&sample_input().into_array(), &mut ctx)?; + + // Rebuild with only codes_offsets widened to u64; every other child is the + // input's, so the two arrays differ solely in codes_offsets width. + let wide = { + let view = narrow.as_view(); + let wide_offsets = view + .codes_offsets() + .cast(DType::Primitive(PType::U64, Nullability::NonNullable))? + .execute::(&mut ctx)? + .into_array(); + OnPair::try_new( + view.dtype().clone(), + view.dict_bytes_handle().clone(), + view.dict_offsets().clone(), + view.codes().clone(), + wide_offsets, + view.uncompressed_lengths().clone(), + view.array_validity(), + )? + }; + assert_eq!( + wide.as_view().codes_offsets().dtype(), + &DType::Primitive(PType::U64, Nullability::NonNullable) + ); + + // Canonical decode is byte-identical across the two widths. + let narrow_decoded = narrow.into_array().execute::(&mut ctx)?; + let wide_decoded = wide + .clone() + .into_array() + .execute::(&mut ctx)?; + assert_arrays_eq!(&wide_decoded, &narrow_decoded, &mut ctx); + + // Equality compare drives CodesWindow over the u64 codes_offsets. + let needle = ConstantArray::new("https://www.example.com/page", wide.len()).into_array(); + let eq = wide + .into_array() + .binary(needle, Operator::Eq)? + .execute::(&mut ctx)?; + assert_arrays_eq!( + &eq, + &BoolArray::from_iter([true, false, false, false, true]), + &mut ctx + ); + Ok(()) +} + #[cfg_attr(miri, ignore)] #[test] fn test_onpair_nullable_canonicalize() -> vortex_error::VortexResult<()> { @@ -220,6 +357,23 @@ fn test_onpair_empty() -> vortex_error::VortexResult<()> { Ok(()) } +/// All-null input is already represented optimally by a constant null array. +#[cfg_attr(miri, ignore)] +#[test] +fn test_onpair_all_null() -> vortex_error::VortexResult<()> { + let input = VarBinArray::from_iter( + [None::<&str>, None, None], + DType::Utf8(Nullability::Nullable), + ) + .into_array(); + let mut ctx = SESSION.create_execution_ctx(); + let arr = onpair_compress(&input, DEFAULT_DICT12_CONFIG, &mut ctx)?; + + assert!(arr.is::()); + assert_arrays_eq!(arr, input, &mut ctx); + Ok(()) +} + /// Filter must share the dictionary — never recompress (this is the /// regression cause on TPC-H Q22 SF=10). Exercise both selectivities /// and check that the result is bit-exact and still an OnPairArray. @@ -235,7 +389,7 @@ fn test_onpair_filter_shares_dict() -> vortex_error::VortexResult<()> { DType::Utf8(Nullability::NonNullable), ); let mut ctx = SESSION.create_execution_ctx(); - let arr = onpair_compress(&varbin.into_array(), DEFAULT_DICT12_CONFIG, &mut ctx)?; + let arr = compress_onpair(&varbin.into_array(), &mut ctx)?; let dict_bytes_before = arr.dict_bytes().clone(); let dict_offsets_len_before = arr.dict_offsets().len(); @@ -309,13 +463,12 @@ fn narrow_codes_offsets(arr: &crate::OnPairArray, target: PType) -> crate::OnPai unsafe { OnPair::new_unchecked( view.dtype().clone(), - view.dict_bytes_handle().clone(), + view.data().clone(), view.dict_offsets().clone(), view.codes().clone(), narrowed_array, view.uncompressed_lengths().clone(), view.array_validity(), - view.bits(), ) } } @@ -339,7 +492,7 @@ fn test_onpair_filter_with_narrowed_codes_offsets_u16() -> vortex_error::VortexR DType::Utf8(Nullability::NonNullable), ); let mut ctx = SESSION.create_execution_ctx(); - let arr = onpair_compress(&varbin.into_array(), DEFAULT_DICT12_CONFIG, &mut ctx)?; + let arr = compress_onpair(&varbin.into_array(), &mut ctx)?; // Force `codes_offsets` to u16 so the panicking pre-fix // `as_slice::()` would fire. @@ -395,7 +548,7 @@ fn test_onpair_filter_with_narrowed_codes_offsets_u8() -> vortex_error::VortexRe DType::Utf8(Nullability::NonNullable), ); let mut ctx = SESSION.create_execution_ctx(); - let arr = onpair_compress(&varbin.into_array(), DEFAULT_DICT12_CONFIG, &mut ctx)?; + let arr = compress_onpair(&varbin.into_array(), &mut ctx)?; let arr = narrow_codes_offsets(&arr, PType::U8); assert_eq!(arr.as_view().codes_offsets().dtype().as_ptype(), PType::U8); diff --git a/encodings/experimental/onpair/tests/big_data.rs b/encodings/experimental/onpair/tests/big_data.rs index e7f74fdd782..15f50bb3322 100644 --- a/encodings/experimental/onpair/tests/big_data.rs +++ b/encodings/experimental/onpair/tests/big_data.rs @@ -76,10 +76,9 @@ fn smoke_100k_rows() -> vortex_error::VortexResult<()> { let t0 = Instant::now(); let arr = onpair_compress(&varbin.into_array(), DEFAULT_DICT12_CONFIG, &mut ctx)?; let compress_elapsed = t0.elapsed(); - let bits = arr.bits(); eprintln!( - "compressed {} rows ({} raw bytes) in {:?}, bits={}", - n, raw_bytes, compress_elapsed, bits + "compressed {} rows ({} raw bytes) in {:?}", + n, raw_bytes, compress_elapsed ); let arr_ref = arr.into_array(); diff --git a/vortex-btrblocks/src/schemes/string/onpair.rs b/vortex-btrblocks/src/schemes/string/onpair.rs index f0a8966089c..8c7b0561502 100644 --- a/vortex-btrblocks/src/schemes/string/onpair.rs +++ b/vortex-btrblocks/src/schemes/string/onpair.rs @@ -49,7 +49,7 @@ impl Scheme for OnPairScheme { /// 4 primitive slot children flow through the cascading compressor: /// `dict_offsets` (u32 → typically `FoR`/`BitPacked`), `codes` (u16 → - /// `FastLanes::BitPacked` to exactly `bits` = 12 by default), + /// usually `FastLanes::BitPacked` after scheme selection), /// `codes_offsets` (u32 → `FoR`), `uncompressed_lengths` (i32 → narrow /// + `FoR`). Validity stays untouched. fn num_children(&self) -> usize { @@ -73,7 +73,10 @@ impl Scheme for OnPairScheme { exec_ctx: &mut ExecutionCtx, ) -> VortexResult { let utf8 = data.array_as_varbinview().into_owned(); - let onpair_array = onpair_compress(utf8.as_array(), DEFAULT_DICT12_CONFIG, exec_ctx)?; + let encoded = onpair_compress(utf8.as_array(), DEFAULT_DICT12_CONFIG, exec_ctx)?; + let Some(onpair_array) = encoded.as_opt::() else { + return Ok(encoded); + }; let dict_offsets = compress_offsets_child( compressor, @@ -116,7 +119,6 @@ impl Scheme for OnPairScheme { codes_offsets, uncompressed_lengths, onpair_array.array_validity(), - onpair_array.bits(), )? .into_array()) } From e6eeaaef57ffa5545a1cfa3408cb7c3d8f2119a1 Mon Sep 17 00:00:00 2001 From: Nemo Yu <83347615+HarukiMoriarty@users.noreply.github.com> Date: Fri, 17 Jul 2026 13:57:37 -0400 Subject: [PATCH 010/498] feat(vortex-array): union_child_validities expr helper (#8829) ## What Adds `vortex_array::expr::union_child_validities(expr)`: the conjunction of an expression's child validities (result is null iff any operand is null). This is the common `ScalarFnVTable::validity()` for null-propagating kernels (comparisons, arithmetic, most geo/tensor ops), letting the planner derive the output null mask without executing the kernel. Migrates `vortex-tensor`'s four scalar functions (`InnerProduct`, `CosineSimilarity`, `L2Denorm`, `L2Norm`) onto it, replacing their inline `and(child validities)`. Signed-off-by: Nemo Yu --- vortex-array/src/expr/exprs.rs | 17 +++++++++++++++++ .../src/scalar_fns/cosine_similarity.rs | 7 ++----- vortex-tensor/src/scalar_fns/inner_product.rs | 7 ++----- vortex-tensor/src/scalar_fns/l2_denorm.rs | 7 ++----- vortex-tensor/src/scalar_fns/l2_norm.rs | 3 ++- 5 files changed, 25 insertions(+), 16 deletions(-) diff --git a/vortex-array/src/expr/exprs.rs b/vortex-array/src/expr/exprs.rs index da70427a0f7..9843277ede7 100644 --- a/vortex-array/src/expr/exprs.rs +++ b/vortex-array/src/expr/exprs.rs @@ -7,6 +7,7 @@ use std::sync::Arc; use std::sync::LazyLock; use vortex_error::VortexExpect; +use vortex_error::VortexResult; use vortex_error::vortex_panic; use vortex_utils::iter::ReduceBalancedIterExt; @@ -434,6 +435,22 @@ where iter.into_iter().reduce_balanced(and) } +/// The conjunction of an expression's child validities — i.e. the validity of a scalar function +/// whose result is null exactly when any operand is null. +/// +/// This is the `ScalarFnVTable::validity` for kernels that propagate nulls and never produce a +/// null from non-null inputs (comparisons, arithmetic, most geo and tensor ops). Returning it lets +/// the planner derive the output's null mask without executing the kernel. Yields `None` when the +/// expression has no children. +pub fn union_child_validities(expression: &Expression) -> VortexResult> { + let child_validities = expression + .children() + .iter() + .map(Expression::validity) + .collect::>>()?; + Ok(and_collect(child_validities)) +} + /// Create a new [`Binary`] using the [`Add`](Operator::Add) operator. /// /// ## Example usage diff --git a/vortex-tensor/src/scalar_fns/cosine_similarity.rs b/vortex-tensor/src/scalar_fns/cosine_similarity.rs index 82328a9b2d9..ffe198eb0fa 100644 --- a/vortex-tensor/src/scalar_fns/cosine_similarity.rs +++ b/vortex-tensor/src/scalar_fns/cosine_similarity.rs @@ -15,7 +15,7 @@ use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayVTable; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::expr::Expression; -use vortex_array::expr::and; +use vortex_array::expr::union_child_validities; use vortex_array::match_each_float_ptype; use vortex_array::scalar_fn::Arity; use vortex_array::scalar_fn::ChildName; @@ -180,10 +180,7 @@ impl ScalarFnVTable for CosineSimilarity { expression: &Expression, ) -> VortexResult> { // The result is null if either input tensor is null. - let lhs_validity = expression.child(0).validity()?; - let rhs_validity = expression.child(1).validity()?; - - Ok(Some(and(lhs_validity, rhs_validity))) + union_child_validities(expression) } fn is_null_sensitive(&self, _options: &Self::Options) -> bool { diff --git a/vortex-tensor/src/scalar_fns/inner_product.rs b/vortex-tensor/src/scalar_fns/inner_product.rs index cfea956f314..cb17d179161 100644 --- a/vortex-tensor/src/scalar_fns/inner_product.rs +++ b/vortex-tensor/src/scalar_fns/inner_product.rs @@ -18,7 +18,7 @@ use vortex_array::dtype::DType; use vortex_array::dtype::NativePType; use vortex_array::dtype::Nullability; use vortex_array::expr::Expression; -use vortex_array::expr::and; +use vortex_array::expr::union_child_validities; use vortex_array::match_each_float_ptype; use vortex_array::scalar_fn::Arity; use vortex_array::scalar_fn::ChildName; @@ -164,10 +164,7 @@ impl ScalarFnVTable for InnerProduct { expression: &Expression, ) -> VortexResult> { // The result is null if either input tensor is null. - let lhs_validity = expression.child(0).validity()?; - let rhs_validity = expression.child(1).validity()?; - - Ok(Some(and(lhs_validity, rhs_validity))) + union_child_validities(expression) } fn is_null_sensitive(&self, _options: &Self::Options) -> bool { diff --git a/vortex-tensor/src/scalar_fns/l2_denorm.rs b/vortex-tensor/src/scalar_fns/l2_denorm.rs index 4f4159b4a21..a4ac9e6c91e 100644 --- a/vortex-tensor/src/scalar_fns/l2_denorm.rs +++ b/vortex-tensor/src/scalar_fns/l2_denorm.rs @@ -31,7 +31,7 @@ use vortex_array::dtype::NativePType; use vortex_array::dtype::Nullability; use vortex_array::dtype::proto::dtype as pb; use vortex_array::expr::Expression; -use vortex_array::expr::and; +use vortex_array::expr::union_child_validities; use vortex_array::match_each_float_ptype; use vortex_array::scalar::Scalar; use vortex_array::scalar::ScalarValue; @@ -256,10 +256,7 @@ impl ScalarFnVTable for L2Denorm { _options: &Self::Options, expression: &Expression, ) -> VortexResult> { - let normalized_validity = expression.child(0).validity()?; - let norms_validity = expression.child(1).validity()?; - - Ok(Some(and(normalized_validity, norms_validity))) + union_child_validities(expression) } fn is_null_sensitive(&self, _options: &Self::Options) -> bool { diff --git a/vortex-tensor/src/scalar_fns/l2_norm.rs b/vortex-tensor/src/scalar_fns/l2_norm.rs index aa85ab4c78e..0eeb6777ab3 100644 --- a/vortex-tensor/src/scalar_fns/l2_norm.rs +++ b/vortex-tensor/src/scalar_fns/l2_norm.rs @@ -25,6 +25,7 @@ use vortex_array::dtype::NativePType; use vortex_array::dtype::Nullability; use vortex_array::dtype::proto::dtype as pb; use vortex_array::expr::Expression; +use vortex_array::expr::union_child_validities; use vortex_array::match_each_float_ptype; use vortex_array::scalar::Scalar; use vortex_array::scalar_fn::Arity; @@ -184,7 +185,7 @@ impl ScalarFnVTable for L2Norm { expression: &Expression, ) -> VortexResult> { // The result is null if the input tensor is null. - Ok(Some(expression.child(0).validity()?)) + union_child_validities(expression) } fn is_null_sensitive(&self, _options: &Self::Options) -> bool { From c5e075ab7f9babffd32f9592e8e35a1981c5cd49 Mon Sep 17 00:00:00 2001 From: Dan King Date: Fri, 17 Jul 2026 19:20:56 -0400 Subject: [PATCH 011/498] Add PiecewiseSequence index array (#8799) This array allows FixedSizeList, List, and ListView to implement take without materializing a dense array of indices of size `O(N_TAKE_INDICES * LIST_SIZE)`. Their elements child arrays can, in turn, use efficient memcpy calls instead of copying individual indices. ## Summary Adds `PiecewiseSequenceArray`, a serializable index encoding for piecewise arithmetic sequences: `starts[i] + j * multipliers[i]` for `j` in `0..lengths[i]`. This PR only introduces the index array and its execution/scalar behavior. It intentionally does not specialize `take`; dependent PRs consume this encoding for the constant-multiplier-1 contiguous case. ## Validation - `cargo fmt -p vortex-array` - `cargo test -p vortex-array piecewise_sequence` --------- Signed-off-by: Daniel King --- vortex-array/src/arrays/mod.rs | 4 + .../src/arrays/piecewise_sequence/array.rs | 70 +++++ .../src/arrays/piecewise_sequence/mod.rs | 125 ++++++++ .../src/arrays/piecewise_sequence/tests.rs | 154 ++++++++++ .../src/arrays/piecewise_sequence/vtable.rs | 275 ++++++++++++++++++ vortex-array/src/session/mod.rs | 2 + 6 files changed, 630 insertions(+) create mode 100644 vortex-array/src/arrays/piecewise_sequence/array.rs create mode 100644 vortex-array/src/arrays/piecewise_sequence/mod.rs create mode 100644 vortex-array/src/arrays/piecewise_sequence/tests.rs create mode 100644 vortex-array/src/arrays/piecewise_sequence/vtable.rs diff --git a/vortex-array/src/arrays/mod.rs b/vortex-array/src/arrays/mod.rs index 15b6b5aa6be..001e20dee65 100644 --- a/vortex-array/src/arrays/mod.rs +++ b/vortex-array/src/arrays/mod.rs @@ -88,6 +88,10 @@ pub mod patched; pub use patched::Patched; pub use patched::PatchedArray; +pub mod piecewise_sequence; +pub use piecewise_sequence::PiecewiseSequence; +pub use piecewise_sequence::PiecewiseSequenceArray; + pub mod primitive; pub use primitive::Primitive; pub use primitive::PrimitiveArray; diff --git a/vortex-array/src/arrays/piecewise_sequence/array.rs b/vortex-array/src/arrays/piecewise_sequence/array.rs new file mode 100644 index 00000000000..1559cf7d3ad --- /dev/null +++ b/vortex-array/src/arrays/piecewise_sequence/array.rs @@ -0,0 +1,70 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use smallvec::smallvec; +use vortex_error::VortexResult; + +use crate::ArrayRef; +use crate::array::Array; +use crate::array::ArrayParts; +use crate::array::EmptyArrayData; +use crate::array::TypedArrayRef; +use crate::array_slots; +use crate::arrays::PiecewiseSequence; +use crate::dtype::PType; + +#[array_slots(PiecewiseSequence)] +pub struct PiecewiseSequenceSlots { + /// The inclusive start index of each sequential piece. + pub starts: ArrayRef, + /// The length of each sequential piece. + pub lengths: ArrayRef, + /// The distance between consecutive indices in each piece. + pub multipliers: ArrayRef, +} + +/// Extension methods for [`Array`] values using the [`PiecewiseSequence`] encoding. +pub trait PiecewiseSequenceArrayExt: + TypedArrayRef + PiecewiseSequenceArraySlotsExt +{ +} +impl> PiecewiseSequenceArrayExt for T {} + +impl Array { + /// Constructs a new `PiecewiseSequenceArray` from start, length, and multiplier arrays. + /// + /// This validates only structural invariants: all children must be non-nullable unsigned + /// integer arrays with matching lengths, and the outer array length is the declared expanded + /// index length. Individual ranges are checked when the index array is executed or consumed by + /// a take implementation. + pub fn try_new( + starts: ArrayRef, + lengths: ArrayRef, + multipliers: ArrayRef, + len: usize, + ) -> VortexResult { + Array::try_from_parts( + ArrayParts::new(PiecewiseSequence, PType::U64.into(), len, EmptyArrayData) + .with_slots(smallvec![Some(starts), Some(lengths), Some(multipliers)]), + ) + } + + /// Constructs a new `PiecewiseSequenceArray` without validation. + /// + /// # Safety + /// + /// The caller must guarantee the same structural invariants as [`Self::try_new`]. + pub unsafe fn new_unchecked( + starts: ArrayRef, + lengths: ArrayRef, + multipliers: ArrayRef, + len: usize, + ) -> Self { + unsafe { + Array::from_parts_unchecked( + ArrayParts::new(PiecewiseSequence, PType::U64.into(), len, EmptyArrayData) + .with_slots(smallvec![Some(starts), Some(lengths), Some(multipliers)]), + ) + } + } +} diff --git a/vortex-array/src/arrays/piecewise_sequence/mod.rs b/vortex-array/src/arrays/piecewise_sequence/mod.rs new file mode 100644 index 00000000000..b914c5e9651 --- /dev/null +++ b/vortex-array/src/arrays/piecewise_sequence/mod.rs @@ -0,0 +1,125 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Index encoding for concatenated sequential ranges. +//! +//! A `PiecewiseSequenceArray` represents the expanded index sequence +//! `starts[i] + j * multipliers[i]` for `j` in `0..lengths[i]` for each piece `i`. It is +//! intended for take operations that can gather regular runs without materializing one index per +//! element. + +use itertools::Itertools; +use vortex_buffer::BufferMut; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; + +use crate::ArrayRef; +use crate::array::ArrayView; +use crate::arrays::PrimitiveArray; +use crate::arrays::piecewise_sequence::array::PiecewiseSequenceArraySlotsExt; +use crate::dtype::UnsignedPType; +use crate::executor::ExecutionCtx; + +pub mod array; +mod vtable; + +#[cfg(test)] +mod tests; + +pub use array::PiecewiseSequenceArrayExt; +pub use vtable::*; + +pub(crate) fn check_index_arrays( + starts: &ArrayRef, + lengths: &ArrayRef, + multipliers: &ArrayRef, +) -> VortexResult<()> { + check_index_array("starts", starts)?; + check_index_array("lengths", lengths)?; + check_index_array("multipliers", multipliers)?; + vortex_ensure!( + starts.len() == lengths.len(), + "PiecewiseSequenceArray starts length {} does not match lengths length {}", + starts.len(), + lengths.len() + ); + vortex_ensure!( + starts.len() == multipliers.len(), + "PiecewiseSequenceArray starts length {} does not match multipliers length {}", + starts.len(), + multipliers.len() + ); + Ok(()) +} + +pub(crate) fn execute_index_arrays( + array: ArrayView<'_, PiecewiseSequence>, + ctx: &mut ExecutionCtx, +) -> VortexResult<(PrimitiveArray, PrimitiveArray, PrimitiveArray)> { + let starts = array.starts().clone().execute::(ctx)?; + let lengths = array.lengths().clone().execute::(ctx)?; + let multipliers = array.multipliers().clone().execute::(ctx)?; + check_index_arrays(starts.as_ref(), lengths.as_ref(), multipliers.as_ref())?; + Ok((starts, lengths, multipliers)) +} + +fn check_index_array(name: &str, array: &ArrayRef) -> VortexResult<()> { + vortex_ensure!( + array.dtype().is_unsigned_int(), + "PiecewiseSequenceArray {name} must have unsigned integer dtype, got {}", + array.dtype() + ); + vortex_ensure!( + !array.dtype().is_nullable(), + "PiecewiseSequenceArray {name} must be non-nullable, got {}", + array.dtype() + ); + Ok(()) +} + +pub(crate) fn materialize_ranges( + starts: &PrimitiveArray, + lengths: &PrimitiveArray, + multipliers: &PrimitiveArray, + output_len: usize, +) -> VortexResult> +where + S: UnsignedPType, + L: UnsignedPType, + M: UnsignedPType, +{ + let starts = starts.as_slice::(); + let lengths = lengths.as_slice::(); + let multipliers = multipliers.as_slice::(); + let mut values = BufferMut::with_capacity(output_len); + let mut computed_len = 0usize; + + for ((&start, &length), &multiplier) in starts.iter().zip_eq(lengths).zip_eq(multipliers) { + let start: usize = start.as_(); + let length: usize = length.as_(); + let multiplier: usize = multiplier.as_(); + if length != 0 { + let last_offset = length - 1; + let last_delta = last_offset + .checked_mul(multiplier) + .ok_or_else(|| vortex_err!("PiecewiseSequenceArray range overflows usize"))?; + start + .checked_add(last_delta) + .ok_or_else(|| vortex_err!("PiecewiseSequenceArray range overflows usize"))?; + } + computed_len = computed_len + .checked_add(length) + .ok_or_else(|| vortex_err!("PiecewiseSequenceArray output length overflows usize"))?; + + values.extend((0..length).map(|offset| (start + offset * multiplier) as u64)); + } + + if computed_len != output_len { + vortex_bail!( + "PiecewiseSequenceArray expanded length {computed_len} does not match declared length {output_len}" + ); + } + Ok(values) +} diff --git a/vortex-array/src/arrays/piecewise_sequence/tests.rs b/vortex-array/src/arrays/piecewise_sequence/tests.rs new file mode 100644 index 00000000000..ab201943e3d --- /dev/null +++ b/vortex-array/src/arrays/piecewise_sequence/tests.rs @@ -0,0 +1,154 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_buffer::ByteBufferMut; +use vortex_buffer::buffer; +use vortex_error::VortexResult; +use vortex_session::registry::ReadContext; + +use crate::ArrayContext; +use crate::IntoArray; +use crate::VortexSessionExecute; +use crate::array_session; +use crate::arrays::ConstantArray; +use crate::arrays::PiecewiseSequence; +use crate::arrays::PiecewiseSequenceArray; +use crate::arrays::PrimitiveArray; +use crate::assert_arrays_eq; +use crate::serde::SerializeOptions; +use crate::serde::SerializedArray; + +#[test] +fn materializes_piecewise_indices() -> VortexResult<()> { + let starts = buffer![3u64, 15, 21].into_array(); + let lengths = buffer![3u64, 3, 3].into_array(); + let multipliers = ConstantArray::new(1u64, 3).into_array(); + let array = PiecewiseSequenceArray::try_new(starts, lengths, multipliers, 9)?.into_array(); + + let expected = PrimitiveArray::from_iter([3u64, 4, 5, 15, 16, 17, 21, 22, 23]).into_array(); + assert_arrays_eq!(array, expected, &mut array_session().create_execution_ctx()); + Ok(()) +} + +#[test] +fn materializes_repeated_and_empty_ranges() -> VortexResult<()> { + let starts = buffer![5u64, 2, 5].into_array(); + let lengths = buffer![2u64, 0, 2].into_array(); + let multipliers = ConstantArray::new(1u64, 3).into_array(); + let array = PiecewiseSequenceArray::try_new(starts, lengths, multipliers, 4)?.into_array(); + + let expected = PrimitiveArray::from_iter([5u64, 6, 5, 6]).into_array(); + assert_arrays_eq!(array, expected, &mut array_session().create_execution_ctx()); + Ok(()) +} + +#[test] +fn materializes_multiplied_ranges() -> VortexResult<()> { + let starts = buffer![3u64, 15].into_array(); + let lengths = buffer![3u64, 2].into_array(); + let multipliers = buffer![2u64, 4].into_array(); + let array = PiecewiseSequenceArray::try_new(starts, lengths, multipliers, 5)?.into_array(); + + let expected = PrimitiveArray::from_iter([3u64, 5, 7, 15, 19]).into_array(); + assert_arrays_eq!(array, expected, &mut array_session().create_execution_ctx()); + Ok(()) +} + +#[test] +fn supports_constant_lengths() -> VortexResult<()> { + let starts = buffer![0u64, 10, 20].into_array(); + let lengths = ConstantArray::new(2u64, 3).into_array(); + let multipliers = ConstantArray::new(1u64, 3).into_array(); + let array = PiecewiseSequenceArray::try_new(starts, lengths, multipliers, 6)?.into_array(); + + let expected = PrimitiveArray::from_iter([0u64, 1, 10, 11, 20, 21]).into_array(); + assert_arrays_eq!(array, expected, &mut array_session().create_execution_ctx()); + Ok(()) +} + +#[test] +fn scalar_at_maps_into_piece() -> VortexResult<()> { + let starts = buffer![3u64, 15, 21].into_array(); + let lengths = buffer![3u64, 3, 3].into_array(); + let multipliers = buffer![1u64, 1, 1].into_array(); + let array = PiecewiseSequenceArray::try_new(starts, lengths, multipliers, 9)?.into_array(); + let mut ctx = array_session().create_execution_ctx(); + + assert_eq!(array.execute_scalar(0, &mut ctx)?, 3u64.into()); + assert_eq!(array.execute_scalar(4, &mut ctx)?, 16u64.into()); + assert_eq!(array.execute_scalar(8, &mut ctx)?, 23u64.into()); + Ok(()) +} + +#[test] +fn constructor_defers_range_value_validation() -> VortexResult<()> { + let starts = buffer![u64::MAX].into_array(); + let lengths = buffer![2u64].into_array(); + let multipliers = buffer![1u64].into_array(); + let array = PiecewiseSequenceArray::try_new(starts, lengths, multipliers, 2)?.into_array(); + + let err = array + .execute::(&mut array_session().create_execution_ctx()) + .unwrap_err(); + assert!( + err.to_string() + .contains("PiecewiseSequenceArray range overflows usize"), + "{err}" + ); + Ok(()) +} + +#[test] +fn execution_checks_declared_length() -> VortexResult<()> { + let starts = buffer![0u64, 3].into_array(); + let lengths = buffer![2u64, 2].into_array(); + let multipliers = ConstantArray::new(1u64, 2).into_array(); + let array = PiecewiseSequenceArray::try_new(starts, lengths, multipliers, 3)?.into_array(); + + let err = array + .execute::(&mut array_session().create_execution_ctx()) + .unwrap_err(); + assert!( + err.to_string() + .contains("PiecewiseSequenceArray expanded length 4 does not match declared length 3"), + "{err}" + ); + Ok(()) +} + +#[test] +fn serde_roundtrip_preserves_piecewise_indices() -> VortexResult<()> { + let array = PiecewiseSequenceArray::try_new( + buffer![3u32, 15, 21].into_array(), + buffer![2u16, 0, 2].into_array(), + buffer![2u8, 1, 3].into_array(), + 4, + )? + .into_array(); + let dtype = array.dtype().clone(); + let len = array.len(); + + let array_ctx = ArrayContext::empty(); + let serialized = array.serialize(&array_ctx, &array_session(), &SerializeOptions::default())?; + + let mut concat = ByteBufferMut::empty(); + for buffer in serialized { + concat.extend_from_slice(buffer.as_ref()); + } + + let parts = SerializedArray::try_from(concat.freeze())?; + let decoded = parts.decode( + &dtype, + len, + &ReadContext::new(array_ctx.to_ids()), + &array_session(), + )?; + + assert!(decoded.is::()); + assert_arrays_eq!( + decoded, + PrimitiveArray::from_iter([3u64, 5, 21, 24]).into_array(), + &mut array_session().create_execution_ctx() + ); + Ok(()) +} diff --git a/vortex-array/src/arrays/piecewise_sequence/vtable.rs b/vortex-array/src/arrays/piecewise_sequence/vtable.rs new file mode 100644 index 00000000000..ac2eba61116 --- /dev/null +++ b/vortex-array/src/arrays/piecewise_sequence/vtable.rs @@ -0,0 +1,275 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use itertools::Itertools; +use prost::Message; +use smallvec::smallvec; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; +use vortex_error::vortex_panic; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +use crate::ArrayParts; +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::ExecutionResult; +use crate::IntoArray; +use crate::array::Array; +use crate::array::ArrayId; +use crate::array::ArrayView; +use crate::array::EmptyArrayData; +use crate::array::OperationsVTable; +use crate::array::VTable; +use crate::array::ValidityVTable; +use crate::array::with_empty_buffers; +use crate::arrays::PrimitiveArray; +use crate::arrays::piecewise_sequence::array::PiecewiseSequenceArraySlotsExt; +use crate::arrays::piecewise_sequence::array::PiecewiseSequenceSlots; +use crate::arrays::piecewise_sequence::check_index_arrays; +use crate::arrays::piecewise_sequence::execute_index_arrays; +use crate::arrays::piecewise_sequence::materialize_ranges; +use crate::arrays::primitive::PrimitiveArrayExt; +use crate::buffer::BufferHandle; +use crate::dtype::DType; +use crate::dtype::PType; +use crate::dtype::UnsignedPType; +use crate::match_each_unsigned_integer_ptype; +use crate::scalar::Scalar; +use crate::serde::ArrayChildren; +use crate::validity::Validity; + +/// A [`PiecewiseSequence`]-encoded Vortex index array. +pub type PiecewiseSequenceArray = Array; + +#[derive(Clone, prost::Message)] +struct PiecewiseSequenceMetadata { + #[prost(uint64, tag = "1")] + num_pieces: u64, + #[prost(enumeration = "PType", tag = "2")] + starts_ptype: i32, + #[prost(enumeration = "PType", tag = "3")] + lengths_ptype: i32, + #[prost(enumeration = "PType", tag = "4")] + multipliers_ptype: i32, +} + +#[derive(Clone, Debug)] +pub struct PiecewiseSequence; + +impl VTable for PiecewiseSequence { + type TypedArrayData = EmptyArrayData; + type OperationsVTable = Self; + type ValidityVTable = Self; + + fn id(&self) -> ArrayId { + static ID: CachedId = CachedId::new("vortex.piecewise-sequence"); + *ID + } + + fn validate( + &self, + _data: &Self::TypedArrayData, + dtype: &DType, + _len: usize, + slots: &[Option], + ) -> VortexResult<()> { + vortex_ensure!( + dtype == &DType::from(PType::U64), + "PiecewiseSequenceArray dtype must be u64, got {dtype}" + ); + vortex_ensure!( + slots.len() == PiecewiseSequenceSlots::NAMES.len(), + "PiecewiseSequenceArray requires {} slots, got {}", + PiecewiseSequenceSlots::NAMES.len(), + slots.len() + ); + let starts = slots[PiecewiseSequenceSlots::STARTS] + .as_ref() + .ok_or_else(|| vortex_err!("PiecewiseSequenceArray starts slot must be present"))?; + let lengths = slots[PiecewiseSequenceSlots::LENGTHS] + .as_ref() + .ok_or_else(|| vortex_err!("PiecewiseSequenceArray lengths slot must be present"))?; + let multipliers = slots[PiecewiseSequenceSlots::MULTIPLIERS] + .as_ref() + .ok_or_else(|| { + vortex_err!("PiecewiseSequenceArray multipliers slot must be present") + })?; + check_index_arrays(starts, lengths, multipliers) + } + + fn nbuffers(_array: ArrayView<'_, Self>) -> usize { + 0 + } + + fn buffer(_array: ArrayView<'_, Self>, _idx: usize) -> BufferHandle { + vortex_panic!("PiecewiseSequenceArray has no buffers") + } + + fn buffer_name(_array: ArrayView<'_, Self>, _idx: usize) -> Option { + None + } + + fn with_buffers( + &self, + array: ArrayView<'_, Self>, + buffers: &[BufferHandle], + ) -> VortexResult> { + with_empty_buffers(self, array, buffers) + } + + fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { + PiecewiseSequenceSlots::NAMES[idx].to_string() + } + + fn serialize( + array: ArrayView<'_, Self>, + _session: &VortexSession, + ) -> VortexResult>> { + Ok(Some( + PiecewiseSequenceMetadata { + num_pieces: u64::try_from(array.starts().len()).map_err(|_| { + vortex_err!( + "PiecewiseSequenceArray piece count {} overflowed u64", + array.starts().len() + ) + })?, + starts_ptype: PType::try_from(array.starts().dtype())? as i32, + lengths_ptype: PType::try_from(array.lengths().dtype())? as i32, + multipliers_ptype: PType::try_from(array.multipliers().dtype())? as i32, + } + .encode_to_vec(), + )) + } + + fn deserialize( + &self, + dtype: &DType, + len: usize, + metadata: &[u8], + buffers: &[BufferHandle], + children: &dyn ArrayChildren, + _session: &VortexSession, + ) -> VortexResult> { + let metadata = PiecewiseSequenceMetadata::decode(metadata)?; + vortex_ensure!( + dtype == &DType::from(PType::U64), + "PiecewiseSequenceArray dtype must be u64, got {dtype}" + ); + vortex_ensure!( + buffers.is_empty(), + "PiecewiseSequenceArray expects no buffers, got {}", + buffers.len() + ); + vortex_ensure!( + children.len() == PiecewiseSequenceSlots::NAMES.len(), + "PiecewiseSequenceArray expects {} children, got {}", + PiecewiseSequenceSlots::NAMES.len(), + children.len() + ); + + let num_pieces = usize::try_from(metadata.num_pieces).map_err(|_| { + vortex_err!( + "PiecewiseSequenceArray piece count {} does not fit in usize", + metadata.num_pieces + ) + })?; + let starts = children.get( + PiecewiseSequenceSlots::STARTS, + &metadata.starts_ptype().into(), + num_pieces, + )?; + let lengths = children.get( + PiecewiseSequenceSlots::LENGTHS, + &metadata.lengths_ptype().into(), + num_pieces, + )?; + let multipliers = children.get( + PiecewiseSequenceSlots::MULTIPLIERS, + &metadata.multipliers_ptype().into(), + num_pieces, + )?; + + Ok( + ArrayParts::new(self.clone(), dtype.clone(), len, EmptyArrayData) + .with_slots(smallvec![Some(starts), Some(lengths), Some(multipliers)]), + ) + } + + fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { + let (starts, lengths, multipliers) = execute_index_arrays(array.as_view(), ctx)?; + + let values = match_each_unsigned_integer_ptype!(starts.ptype(), |S| { + match_each_unsigned_integer_ptype!(lengths.ptype(), |L| { + match_each_unsigned_integer_ptype!(multipliers.ptype(), |M| { + materialize_ranges::(&starts, &lengths, &multipliers, array.len())? + }) + }) + }); + Ok(ExecutionResult::done( + PrimitiveArray::new(values.freeze(), Validity::NonNullable).into_array(), + )) + } +} + +impl OperationsVTable for PiecewiseSequence { + fn scalar_at( + array: ArrayView<'_, PiecewiseSequence>, + index: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let (starts, lengths, multipliers) = execute_index_arrays(array, ctx)?; + + let value = match_each_unsigned_integer_ptype!(starts.ptype(), |S| { + match_each_unsigned_integer_ptype!(lengths.ptype(), |L| { + match_each_unsigned_integer_ptype!(multipliers.ptype(), |M| { + scalar_at::(&starts, &lengths, &multipliers, index)? + }) + }) + }); + Ok(value.into()) + } +} + +impl ValidityVTable for PiecewiseSequence { + fn validity(_array: ArrayView<'_, PiecewiseSequence>) -> VortexResult { + Ok(Validity::NonNullable) + } +} + +fn scalar_at( + starts: &PrimitiveArray, + lengths: &PrimitiveArray, + multipliers: &PrimitiveArray, + index: usize, +) -> VortexResult +where + S: UnsignedPType, + L: UnsignedPType, + M: UnsignedPType, +{ + let mut remaining = index; + for ((&start, &length), &multiplier) in starts + .as_slice::() + .iter() + .zip_eq(lengths.as_slice::()) + .zip_eq(multipliers.as_slice::()) + { + let length: usize = length.as_(); + if remaining < length { + let start: usize = start.as_(); + let multiplier: usize = multiplier.as_(); + let offset = remaining + .checked_mul(multiplier) + .ok_or_else(|| vortex_err!("PiecewiseSequenceArray range overflows usize"))?; + let value = start + .checked_add(offset) + .ok_or_else(|| vortex_err!("PiecewiseSequenceArray range overflows usize"))?; + return Ok(value as u64); + } + remaining -= length; + } + vortex_bail!("PiecewiseSequenceArray index {index} out of bounds") +} diff --git a/vortex-array/src/session/mod.rs b/vortex-array/src/session/mod.rs index 2cb8414cbfd..98a769ceea9 100644 --- a/vortex-array/src/session/mod.rs +++ b/vortex-array/src/session/mod.rs @@ -25,6 +25,7 @@ use crate::arrays::List; use crate::arrays::ListView; use crate::arrays::Masked; use crate::arrays::Null; +use crate::arrays::PiecewiseSequence; use crate::arrays::Primitive; use crate::arrays::Struct; use crate::arrays::VarBin; @@ -81,6 +82,7 @@ impl Default for ArraySession { this.register(Dict); this.register(List); this.register(Masked); + this.register(PiecewiseSequence); this.register(VarBin); this From b7b01d373f03cb177376946bd33f7941ea9649a0 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Mon, 20 Jul 2026 14:31:28 +0100 Subject: [PATCH 012/498] Avoid redundant allocations in struct layout reader (#8851) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pre-size the dtype and name vectors to their final length and push the `validity` entry first, instead of collecting at field-count capacity and then `insert(0, ..)`. The previous code `collect`ed each `Vec` at exact field-count capacity, then `insert(0, ..)` forced a reallocation **and** an O(n) shift of every element. Both are now gone. ## Checks - `cargo build -p vortex-layout` ✅ 🤖 Generated with [Claude Code](https://claude.com/claude-code) Signed-off-by: Joe Isaacs Co-authored-by: Claude Opus 4.8 (1M context) --- vortex-layout/src/layouts/struct_/reader.rs | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/vortex-layout/src/layouts/struct_/reader.rs b/vortex-layout/src/layouts/struct_/reader.rs index ab03a4624fa..18a14a2c97b 100644 --- a/vortex-layout/src/layouts/struct_/reader.rs +++ b/vortex-layout/src/layouts/struct_/reader.rs @@ -82,17 +82,17 @@ impl StructReader { .collect() }); - let mut dtypes: Vec = struct_dt.fields().collect(); - let mut names: Vec> = struct_dt - .names() - .iter() - .map(|x| Arc::clone(x.inner())) - .collect(); - - if layout.dtype.is_nullable() { - dtypes.insert(0, DType::Bool(Nullability::NonNullable)); - names.insert(0, Arc::from("validity")); + let nullable = layout.dtype.is_nullable(); + let extra = nullable as usize; + + let mut dtypes: Vec = Vec::with_capacity(struct_dt.nfields() + extra); + let mut names: Vec> = Vec::with_capacity(struct_dt.nfields() + extra); + if nullable { + dtypes.push(DType::Bool(Nullability::NonNullable)); + names.push(Arc::from("validity")); } + dtypes.extend(struct_dt.fields()); + names.extend(struct_dt.names().iter().map(|x| Arc::clone(x.inner()))); let lazy_children = LazyReaderChildren::new( Arc::clone(&layout.children), From 00cb75d1eaa734232e23279be89ad3d48f6bf1ae Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Mon, 20 Jul 2026 16:49:15 +0100 Subject: [PATCH 013/498] Validate PostscriptFragment offsets when parsing the footer (#8818) Parser shouldn't crash, but error in these cases --- vortex-file/src/footer/deserializer.rs | 131 ++++++++++++++++++++++--- vortex-file/src/footer/mod.rs | 4 +- 2 files changed, 117 insertions(+), 18 deletions(-) diff --git a/vortex-file/src/footer/deserializer.rs b/vortex-file/src/footer/deserializer.rs index a0f12e85c8a..7c2d6cac3d7 100644 --- a/vortex-file/src/footer/deserializer.rs +++ b/vortex-file/src/footer/deserializer.rs @@ -229,9 +229,11 @@ impl FooterDeserializer { initial_read: &[u8], segment: &PostscriptSegment, ) -> VortexResult { - let offset = usize::try_from(segment.offset - initial_offset)?; - let sliced_buffer = - FlatBuffer::copy_from(&initial_read[offset..offset + (segment.length as usize)]); + let sliced_buffer = FlatBuffer::copy_from(checked_segment_slice( + initial_read, + initial_offset, + segment, + )?); DType::from_flatbuffer(sliced_buffer, &self.session) } @@ -244,11 +246,9 @@ impl FooterDeserializer { dtype: &DType, session: &VortexSession, ) -> VortexResult { - let offset = usize::try_from(segment.offset - initial_offset)?; - let sliced_buffer = - FlatBuffer::copy_from(&initial_read[offset..offset + (segment.length as usize)]); + let sliced_buffer = checked_segment_slice(initial_read, initial_offset, segment)?; - let fb = root::(&sliced_buffer)?; + let fb = root::(sliced_buffer)?; FileStatistics::from_flatbuffer(&fb, dtype, session) } @@ -262,20 +262,119 @@ impl FooterDeserializer { dtype: DType, file_stats: Option, ) -> VortexResult