From 4fce0bcf8e97aa3163172e236520f0c9d6d88c83 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Wed, 12 Aug 2026 00:21:21 +0100 Subject: [PATCH] Broaden the string-to-Arrow benchmarks Rework `vortex-arrow`'s string export benchmarks into a matrix over encoding (offset, view, FSST, OnPair, Zstd), structure (flat, dict, chunked), operator (identity, filter, take, slice, mask, zip), and nullability, exported to both Arrow offset and view layouts, plus the same cases appended directly into `VarBinBuilder` and `VarBinViewBuilder`. Add a `string_to_arrow` benchmark in `vortex` that measures the end-to-end file scan producing Arrow offset arrays, so the export is also covered with data laid out by the default compressor rather than by hand. Signed-off-by: Robert Kruszewski --- Cargo.lock | 1 + vortex-arrow/benches/to_arrow.rs | 478 ++++++++++++++++++++---------- vortex/Cargo.toml | 5 + vortex/benches/string_to_arrow.rs | 222 ++++++++++++++ 4 files changed, 552 insertions(+), 154 deletions(-) create mode 100644 vortex/benches/string_to_arrow.rs diff --git a/Cargo.lock b/Cargo.lock index cd2662d2372..df568a1bb7e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9524,6 +9524,7 @@ dependencies = [ "arrow-array 58.4.0", "codspeed-divan-compat", "fastlanes", + "futures", "mimalloc", "parquet 58.4.0", "rand 0.10.2", diff --git a/vortex-arrow/benches/to_arrow.rs b/vortex-arrow/benches/to_arrow.rs index e6016ec9705..eea6105e3ff 100644 --- a/vortex-arrow/benches/to_arrow.rs +++ b/vortex-arrow/benches/to_arrow.rs @@ -3,6 +3,8 @@ #![expect(clippy::unwrap_used)] +use std::fmt::Display; +use std::fmt::Formatter; use std::sync::Arc; use std::sync::LazyLock; @@ -10,26 +12,32 @@ use arrow_schema::DataType; use arrow_schema::Field; use divan::Bencher; use divan::counter::ItemsCount; +use itertools::iproduct; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::array_session; +use vortex_array::arrays::BoolArray; use vortex_array::arrays::ChunkedArray; +use vortex_array::arrays::ConstantArray; use vortex_array::arrays::DecimalArray; use vortex_array::arrays::DictArray; use vortex_array::arrays::FilterArray; use vortex_array::arrays::ListArray; use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::SliceArray; use vortex_array::arrays::StructArray; use vortex_array::arrays::VarBinViewArray; use vortex_array::builders::VarBinBuilder; use vortex_array::builders::VarBinViewBuilder; +use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; use vortex_array::dtype::DecimalDType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::dtype::StructFields; +use vortex_array::scalar::Scalar; use vortex_array::session::ArraySessionExt; #[expect( deprecated, @@ -121,205 +129,367 @@ fn ArrowExportVTable_to_arrow_field(bencher: Bencher) { .bench_values(|dtype| SESSION.arrow().to_arrow_field("", &dtype).unwrap()) } -#[derive(Clone, Copy, Debug)] +#[derive(Clone, Copy)] enum StringEncoding { + Offset, View, Fsst, OnPair, Zstd, +} + +impl Display for StringEncoding { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + f.write_str(match self { + Self::Offset => "offset", + Self::View => "view", + Self::Fsst => "fsst", + Self::OnPair => "onpair", + Self::Zstd => "zstd", + }) + } +} + +#[derive(Clone, Copy)] +enum StringStructure { + Flat, Dict, - DictFsst, - DictZstd, - FilterFsst, - FilterZstd, - FilterDictFsst, - ChunkedFsst, - /// Every third row null, so the export walks a partial validity mask rather than an all-valid - /// one and has to interleave nulls with the decoded values. - NullableFsst, - NullableZstd, - NullableDict, + Chunked, +} + +impl Display for StringStructure { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + f.write_str(match self { + Self::Flat => "flat", + Self::Dict => "dict", + Self::Chunked => "chunked", + }) + } +} + +#[derive(Clone, Copy)] +enum StringOperator { + Identity, + Filter, + Take, + Slice, + Mask, + Zip, +} + +impl Display for StringOperator { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + f.write_str(match self { + Self::Identity => "identity", + Self::Filter => "filter", + Self::Take => "take", + Self::Slice => "slice", + Self::Mask => "mask", + Self::Zip => "zip", + }) + } +} + +#[derive(Clone, Copy)] +enum StringValidity { + NonNullable, + Nullable, +} + +impl StringValidity { + fn nullability(self) -> Nullability { + match self { + Self::NonNullable => Nullability::NonNullable, + Self::Nullable => Nullability::Nullable, + } + } +} + +impl Display for StringValidity { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + f.write_str(match self { + Self::NonNullable => "nonnull", + Self::Nullable => "nullable", + }) + } +} + +#[derive(Clone, Copy)] +struct StringCase { + encoding: StringEncoding, + structure: StringStructure, + operator: StringOperator, + validity: StringValidity, +} + +impl Display for StringCase { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!( + f, + "{}/{}/{}/{}", + self.encoding, self.structure, self.operator, self.validity + ) + } +} + +#[derive(Clone, Copy)] +enum ArrowStringLayout { + Offset, + View, +} + +impl ArrowStringLayout { + fn data_type(self) -> DataType { + match self { + Self::Offset => DataType::Utf8, + Self::View => DataType::Utf8View, + } + } +} + +impl Display for ArrowStringLayout { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + f.write_str(match self { + Self::Offset => "offset", + Self::View => "view", + }) + } +} + +#[derive(Clone, Copy)] +struct StringExportCase { + array: StringCase, + layout: ArrowStringLayout, +} + +impl Display for StringExportCase { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "{}/{}", self.layout, self.array) + } } const STRING_ENCODINGS: &[StringEncoding] = &[ + StringEncoding::Offset, StringEncoding::View, StringEncoding::Fsst, StringEncoding::OnPair, StringEncoding::Zstd, - StringEncoding::Dict, - StringEncoding::DictFsst, - StringEncoding::DictZstd, - StringEncoding::FilterFsst, - StringEncoding::FilterZstd, - StringEncoding::FilterDictFsst, - StringEncoding::ChunkedFsst, - StringEncoding::NullableFsst, - StringEncoding::NullableZstd, - StringEncoding::NullableDict, ]; - -/// Encodings whose `append_to_builder` the builder benchmarks reach directly. -/// -/// The Arrow export cannot stand in for these: `execute_until` stops at the first canonical array, -/// so a bare FSST/OnPair/Zstd root is canonicalized to `VarBinView` before any builder sees it. -/// Only `Chunked`, `Constant` and `VarBin` roots reach an encoding's own `append_to_builder` that -/// way, whereas the scan machinery appends encoded arrays into a builder directly. -const BUILDER_STRING_ENCODINGS: &[StringEncoding] = &[ - StringEncoding::View, - StringEncoding::Fsst, - StringEncoding::OnPair, - StringEncoding::Zstd, - StringEncoding::Dict, - StringEncoding::ChunkedFsst, - StringEncoding::NullableFsst, - StringEncoding::NullableZstd, - StringEncoding::NullableDict, +const STRING_STRUCTURES: &[StringStructure] = &[ + StringStructure::Flat, + StringStructure::Dict, + StringStructure::Chunked, ]; +const STRING_OPERATORS: &[StringOperator] = &[ + StringOperator::Identity, + StringOperator::Filter, + StringOperator::Take, + StringOperator::Slice, + StringOperator::Mask, + StringOperator::Zip, +]; +const STRING_VALIDITIES: &[StringValidity] = + &[StringValidity::NonNullable, StringValidity::Nullable]; +const ARROW_STRING_LAYOUTS: &[ArrowStringLayout] = + &[ArrowStringLayout::Offset, ArrowStringLayout::View]; -const OFFSET_STRING_ROWS: usize = 100_000; -const OFFSET_STRING_CHUNKS: usize = 4; +const STRING_ROWS: usize = 100_000; +const STRING_CHUNKS: usize = 4; const DICTIONARY_SIZE: usize = 2_048; -fn structured_strings(len: usize) -> VarBinViewArray { - let values = (0..len) - .map(|index| format!("https://example.com/common/path/{index:06}/shared-suffix")) - .collect::>(); - VarBinViewArray::from_iter_str(values.iter().map(String::as_str)) +fn string_cases() -> Vec { + iproduct!( + STRING_ENCODINGS.iter().copied(), + STRING_STRUCTURES.iter().copied(), + STRING_OPERATORS.iter().copied(), + STRING_VALIDITIES.iter().copied() + ) + .map(|(encoding, structure, operator, validity)| StringCase { + encoding, + structure, + operator, + validity, + }) + .collect() } -fn nullable_structured_strings(len: usize) -> VarBinViewArray { - let values = (0..len) - .map(|index| { - (!index.is_multiple_of(3)) - .then(|| format!("https://example.com/common/path/{index:06}/shared-suffix")) - }) - .collect::>(); - VarBinViewArray::from_iter( - values.iter().map(|value| value.as_deref()), - DType::Utf8(Nullability::Nullable), - ) +fn string_export_cases() -> Vec { + iproduct!(string_cases(), ARROW_STRING_LAYOUTS.iter().copied()) + .map(|(array, layout)| StringExportCase { array, layout }) + .collect() } -fn dictionary_values() -> VarBinViewArray { - structured_strings(DICTIONARY_SIZE) +fn structured_strings(len: usize, validity: StringValidity) -> VarBinViewArray { + match validity { + StringValidity::NonNullable => { + let values = (0..len) + .map(|index| format!("https://example.com/common/path/{index:06}/shared-suffix")) + .collect::>(); + VarBinViewArray::from_iter_str(values.iter().map(String::as_str)) + } + StringValidity::Nullable => { + let values = (0..len) + .map(|index| { + (!index.is_multiple_of(3)).then(|| { + format!("https://example.com/common/path/{index:06}/shared-suffix") + }) + }) + .collect::>(); + VarBinViewArray::from_iter( + values.iter().map(|value| value.as_deref()), + DType::Utf8(Nullability::Nullable), + ) + } + } } fn dictionary_codes() -> ArrayRef { PrimitiveArray::from_iter( - (0..OFFSET_STRING_ROWS).map(|index| u16::try_from(index % DICTIONARY_SIZE).unwrap()), + (0..STRING_ROWS).map(|index| u16::try_from(index % DICTIONARY_SIZE).unwrap()), ) .into_array() } -fn half_rows_mask() -> Mask { - Mask::from_iter((0..OFFSET_STRING_ROWS).map(|index| index.is_multiple_of(2))) +fn filtered(array: ArrayRef) -> ArrayRef { + let mask = Mask::from_iter((0..array.len()).map(|index| index.is_multiple_of(2))); + // Create a lazy FilterArray. The benchmark executes Filter during export. + FilterArray::new(array, mask).into_array() } -fn filtered(array: ArrayRef) -> ArrayRef { - // Keep Filter as a lazy intermediate so benchmark setup cannot optimize it away. - FilterArray::new(array, half_rows_mask()).into_array() +fn offset(source: VarBinViewArray, ctx: &mut ExecutionCtx) -> ArrayRef { + let source = source.into_array(); + let mut builder = VarBinBuilder::::with_capacity(source.dtype().clone(), source.len()); + source.append_to_builder(&mut builder, ctx).unwrap(); + builder.finish_into_varbin().into_array() } -fn chunked_fsst(ctx: &mut ExecutionCtx) -> ArrayRef { - let source = structured_strings(OFFSET_STRING_ROWS).into_array(); +fn fsst(source: ArrayRef, ctx: &mut ExecutionCtx) -> ArrayRef { let compressor = fsst_train_compressor(&source, ctx).unwrap(); - let chunk_size = OFFSET_STRING_ROWS / OFFSET_STRING_CHUNKS; - let chunks = (0..OFFSET_STRING_CHUNKS).map(|chunk_index| { + fsst_compress(&source, &compressor, ctx) + .unwrap() + .into_array() +} + +fn encode_strings( + source: VarBinViewArray, + encoding: StringEncoding, + ctx: &mut ExecutionCtx, +) -> ArrayRef { + match encoding { + StringEncoding::Offset => offset(source, ctx), + StringEncoding::View => source.into_array(), + StringEncoding::Fsst => fsst(source.into_array(), ctx), + StringEncoding::OnPair => { + onpair_compress(&source.into_array(), DEFAULT_CONFIG, ctx).unwrap() + } + StringEncoding::Zstd => Zstd::from_var_bin_view_without_dict(&source, 3, 8_192, ctx) + .unwrap() + .into_array(), + } +} + +fn dictionary_strings( + encoding: StringEncoding, + validity: StringValidity, + ctx: &mut ExecutionCtx, +) -> ArrayRef { + let values = encode_strings(structured_strings(DICTIONARY_SIZE, validity), encoding, ctx); + DictArray::try_new(dictionary_codes(), values) + .unwrap() + .into_array() +} + +fn chunked_strings( + encoding: StringEncoding, + validity: StringValidity, + ctx: &mut ExecutionCtx, +) -> ArrayRef { + let chunk_size = STRING_ROWS / STRING_CHUNKS; + let chunks = (0..STRING_CHUNKS).map(|chunk_index| { let start = chunk_index * chunk_size; - let end = if chunk_index + 1 == OFFSET_STRING_CHUNKS { - OFFSET_STRING_ROWS + let end = if chunk_index + 1 == STRING_CHUNKS { + STRING_ROWS } else { start + chunk_size }; - let chunk = source.slice(start..end).unwrap(); - fsst_compress(&chunk, &compressor, ctx) - .unwrap() - .into_array() + encode_strings(structured_strings(end - start, validity), encoding, ctx) }); - ChunkedArray::try_new(chunks, source.dtype().clone()) + ChunkedArray::try_new(chunks, DType::Utf8(validity.nullability())) .unwrap() .into_array() } -fn fsst(source: ArrayRef, ctx: &mut ExecutionCtx) -> ArrayRef { - let compressor = fsst_train_compressor(&source, ctx).unwrap(); - fsst_compress(&source, &compressor, ctx) +fn take(array: ArrayRef) -> ArrayRef { + let indices = PrimitiveArray::from_iter( + (0..array.len()) + .step_by(2) + .map(|index| u64::try_from(index).unwrap()), + ); + // Create a lazy DictArray. The benchmark executes Take during export. + DictArray::try_new(indices.into_array(), array) .unwrap() .into_array() } -fn string_array(encoding: StringEncoding) -> ArrayRef { +fn sliced(array: ArrayRef) -> ArrayRef { + let start = array.len() / 4; + let end = array.len() * 3 / 4; + // Create a lazy SliceArray. The benchmark executes Slice during export. + SliceArray::new(array, start..end).into_array() +} + +fn mask_array(len: usize) -> ArrayRef { + BoolArray::from_iter((0..len).map(|index| !index.is_multiple_of(3))).into_array() +} + +fn masked(array: ArrayRef) -> ArrayRef { + let mask = mask_array(array.len()); + array.mask(mask).unwrap() +} + +fn zipped(array: ArrayRef) -> ArrayRef { + let replacement = ConstantArray::new( + Scalar::utf8("replacement", array.dtype().nullability()), + array.len(), + ) + .into_array(); + mask_array(array.len()).zip(array, replacement).unwrap() +} + +fn apply_operator(array: ArrayRef, operator: StringOperator) -> ArrayRef { + match operator { + StringOperator::Identity => array, + StringOperator::Filter => filtered(array), + StringOperator::Take => take(array), + StringOperator::Slice => sliced(array), + StringOperator::Mask => masked(array), + StringOperator::Zip => zipped(array), + } +} + +fn string_array(case: StringCase) -> ArrayRef { let mut ctx = SESSION.create_execution_ctx(); - match encoding { - StringEncoding::View => structured_strings(OFFSET_STRING_ROWS).into_array(), - StringEncoding::Fsst => fsst( - structured_strings(OFFSET_STRING_ROWS).into_array(), - &mut ctx, - ), - StringEncoding::OnPair => onpair_compress( - &structured_strings(OFFSET_STRING_ROWS).into_array(), - DEFAULT_CONFIG, - &mut ctx, - ) - .unwrap(), - StringEncoding::Zstd => { - let source = structured_strings(OFFSET_STRING_ROWS); - Zstd::from_var_bin_view_without_dict(&source, 3, 8_192, &mut ctx) - .unwrap() - .into_array() - } - StringEncoding::Dict => { - DictArray::try_new(dictionary_codes(), dictionary_values().into_array()) - .unwrap() - .into_array() - } - StringEncoding::DictFsst => { - let values = fsst(dictionary_values().into_array(), &mut ctx); - DictArray::try_new(dictionary_codes(), values) - .unwrap() - .into_array() - } - StringEncoding::DictZstd => { - let values = dictionary_values(); - let compressed_values = - Zstd::from_var_bin_view_without_dict(&values, 3, 8_192, &mut ctx) - .unwrap() - .into_array(); - DictArray::try_new(dictionary_codes(), compressed_values) - .unwrap() - .into_array() - } - StringEncoding::FilterFsst => filtered(string_array(StringEncoding::Fsst)), - StringEncoding::FilterZstd => filtered(string_array(StringEncoding::Zstd)), - StringEncoding::FilterDictFsst => filtered(string_array(StringEncoding::DictFsst)), - StringEncoding::ChunkedFsst => chunked_fsst(&mut ctx), - StringEncoding::NullableFsst => fsst( - nullable_structured_strings(OFFSET_STRING_ROWS).into_array(), + let array = match case.structure { + StringStructure::Flat => encode_strings( + structured_strings(STRING_ROWS, case.validity), + case.encoding, &mut ctx, ), - StringEncoding::NullableZstd => { - let source = nullable_structured_strings(OFFSET_STRING_ROWS); - Zstd::from_var_bin_view_without_dict(&source, 3, 8_192, &mut ctx) - .unwrap() - .into_array() - } - StringEncoding::NullableDict => { - // Nulls live in the dictionary rather than the codes, so the export has to combine - // the two validities. - let values = nullable_structured_strings(DICTIONARY_SIZE); - DictArray::try_new(dictionary_codes(), values.into_array()) - .unwrap() - .into_array() - } - } + StringStructure::Dict => dictionary_strings(case.encoding, case.validity, &mut ctx), + StringStructure::Chunked => chunked_strings(case.encoding, case.validity, &mut ctx), + }; + apply_operator(array, case.operator) } -/// End-to-end export to Arrow `Utf8`, which is served through a `VarBinBuilder`. -#[divan::bench(args = STRING_ENCODINGS)] -fn offset_string_export(bencher: Bencher, encoding: StringEncoding) { - let array = string_array(encoding); - let field = Field::new("value", DataType::Utf8, array.dtype().is_nullable()); - +/// Measures export to Arrow offset arrays and Arrow view arrays. +#[divan::bench(args = string_export_cases())] +fn string_export(bencher: Bencher, case: StringExportCase) { + let array = string_array(case.array); + let field = Field::new( + "value", + case.layout.data_type(), + array.dtype().is_nullable(), + ); bencher .with_inputs(|| (array.clone(), SESSION.create_execution_ctx())) .input_counter(|(array, _)| ItemsCount::new(array.len())) @@ -331,10 +501,10 @@ fn offset_string_export(bencher: Bencher, encoding: StringEncoding) { }); } -/// Appends an encoded array straight into an offset builder. -#[divan::bench(args = BUILDER_STRING_ENCODINGS)] -fn append_to_varbin_builder(bencher: Bencher, encoding: StringEncoding) { - let array = string_array(encoding); +/// Measures a direct append to an offset builder. +#[divan::bench(args = string_cases())] +fn append_to_varbin_builder(bencher: Bencher, case: StringCase) { + let array = string_array(case); bencher .with_inputs(|| (array.clone(), SESSION.create_execution_ctx())) @@ -347,10 +517,10 @@ fn append_to_varbin_builder(bencher: Bencher, encoding: StringEncoding) { }); } -/// Appends an encoded array straight into a view builder. -#[divan::bench(args = BUILDER_STRING_ENCODINGS)] -fn append_to_view_builder(bencher: Bencher, encoding: StringEncoding) { - let array = string_array(encoding); +/// Measures a direct append to a view builder. +#[divan::bench(args = string_cases())] +fn append_to_view_builder(bencher: Bencher, case: StringCase) { + let array = string_array(case); bencher .with_inputs(|| (array.clone(), SESSION.create_execution_ctx())) diff --git a/vortex/Cargo.toml b/vortex/Cargo.toml index 6a2a840a500..bff10f924d6 100644 --- a/vortex/Cargo.toml +++ b/vortex/Cargo.toml @@ -60,6 +60,7 @@ anyhow = { workspace = true } arrow-array = { workspace = true } divan = { workspace = true } fastlanes = { workspace = true } +futures = { workspace = true } mimalloc = { workspace = true } parquet = { workspace = true } rand = { workspace = true } @@ -116,3 +117,7 @@ test = false [[bench]] name = "pipeline" harness = false + +[[bench]] +name = "string_to_arrow" +harness = false diff --git a/vortex/benches/string_to_arrow.rs b/vortex/benches/string_to_arrow.rs new file mode 100644 index 00000000000..8dca6032af6 --- /dev/null +++ b/vortex/benches/string_to_arrow.rs @@ -0,0 +1,222 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Measures Vortex file scans that produce Arrow offset arrays. +//! The file writer uses the default compressor for strings. + +#![expect(clippy::unwrap_used)] + +use std::fmt::Display; +use std::fmt::Formatter; +use std::sync::LazyLock; + +use arrow_array::ArrayRef as ArrowArrayRef; +use arrow_array::BinaryArray; +use arrow_array::StringArray; +use arrow_array::types::BinaryType; +use arrow_array::types::ByteArrayType; +use arrow_array::types::Utf8Type; +use divan::Bencher; +use divan::counter::ItemsCount; +use futures::StreamExt; +use vortex_array::ArrayRef; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::array_session; +use vortex_array::arrays::ChunkedArray; +use vortex_array::arrays::StructArray; +use vortex_array::arrays::VarBinViewArray; +use vortex_array::arrays::struct_::StructArrayExt; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::session::ArraySessionExt; +use vortex_array::stream::ArrayStreamExt; +#[expect( + deprecated, + reason = "the benchmark requests an explicit offset layout" +)] +use vortex_arrow::ArrowArrayExecutor; +use vortex_buffer::ByteBufferMut; +use vortex_edition::Edition; +use vortex_edition::EditionId; +use vortex_edition::EditionInclusion; +use vortex_edition::EditionSessionExt; +use vortex_file::OpenOptionsSessionExt; +use vortex_file::VortexFile; +use vortex_file::WriteOptionsSessionExt; +use vortex_file::WriteStrategyBuilder; +use vortex_io::session::RuntimeSession; +use vortex_io::session::RuntimeSessionExt; +use vortex_layout::session::LayoutSession; +use vortex_session::VortexSession; + +fn main() { + LazyLock::force(&FILE); + divan::main(); +} + +const ROWS_PER_CHUNK: usize = 65_536; +const CHUNKS: usize = 16; +const ROWS: usize = ROWS_PER_CHUNK * CHUNKS; +const BENCH_EDITION: EditionId = EditionId::new("bench", 2026, 8, 0); + +static RUNTIME: LazyLock = LazyLock::new(|| { + tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap() +}); + +static SESSION: LazyLock = LazyLock::new(|| { + let _guard = RUNTIME.enter(); + let session = array_session() + .with::() + .with::() + .with_tokio(); + vortex_file::register_default_encodings(&session); + enable_all_registered_array_encodings(&session); + session +}); + +fn enable_all_registered_array_encodings(session: &VortexSession) { + let editions = session.editions(); + editions + .declare_edition(Edition { + id: BENCH_EDITION, + min_vortex_version: None, + }) + .unwrap(); + let ids = session + .arrays() + .registry() + .read(|map| map.keys().copied().collect::>()); + for id in ids { + editions + .declare_inclusion(EditionInclusion::new(&id, BENCH_EDITION)) + .unwrap(); + } + session.enable_edition(BENCH_EDITION).unwrap(); +} + +#[derive(Clone, Copy)] +enum ByteKind { + Utf8, + Binary, +} + +impl Display for ByteKind { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + f.write_str(match self { + Self::Utf8 => "string_array", + Self::Binary => "binary_array", + }) + } +} + +const BYTE_KINDS: &[ByteKind] = &[ByteKind::Utf8, ByteKind::Binary]; + +static FILE: LazyLock = LazyLock::new(make_file); + +fn string_chunk(chunk: usize) -> StructArray { + let start = chunk * ROWS_PER_CHUNK; + let values = (start..start + ROWS_PER_CHUNK) + .map(|index| format!("https://example.com/common/path/{index:08}/shared-suffix")) + .collect::>(); + let values = VarBinViewArray::from_iter( + values.iter().map(|value| Some(value.as_str())), + DType::Utf8(Nullability::NonNullable), + ); + StructArray::from_fields(&[("value", values.into_array())]).unwrap() +} + +fn make_file() -> VortexFile { + let chunks = (0..CHUNKS) + .map(|chunk| string_chunk(chunk).into_array()) + .collect::>(); + let array = ChunkedArray::from_iter(chunks).into_array(); + let strategy = WriteStrategyBuilder::default() + .with_row_block_size(ROWS_PER_CHUNK) + .with_data_block_target_bytes(None) + .build(); + let mut bytes = ByteBufferMut::empty(); + RUNTIME + .block_on( + SESSION + .write_options() + .with_strategy(strategy) + .write(&mut bytes, array.to_array_stream()), + ) + .unwrap(); + SESSION.open_options().open_buffer(bytes).unwrap() +} + +#[expect( + deprecated, + reason = "the benchmark requests an explicit offset layout" +)] +fn to_offset_array(array: ArrayRef, kind: ByteKind) -> ArrowArrayRef { + let mut ctx = SESSION.create_execution_ctx(); + let struct_array = array.execute::(&mut ctx).unwrap(); + let values = struct_array + .unmasked_field_by_name("value") + .unwrap() + .clone(); + let arrow = match kind { + ByteKind::Utf8 => values.execute_arrow(Some(&Utf8Type::DATA_TYPE), &mut ctx), + ByteKind::Binary => values.execute_arrow(Some(&BinaryType::DATA_TYPE), &mut ctx), + } + .unwrap(); + + match kind { + ByteKind::Utf8 => assert!(arrow.as_any().is::()), + ByteKind::Binary => assert!(arrow.as_any().is::()), + } + arrow +} + +fn read_to_offset_array(file: &VortexFile, kind: ByteKind) -> ArrowArrayRef { + RUNTIME.block_on(async { + let array = file + .scan() + .unwrap() + .into_array_stream() + .unwrap() + .read_all() + .await + .unwrap(); + to_offset_array(array, kind) + }) +} + +fn read_to_offset_batches(file: &VortexFile, kind: ByteKind) -> Vec { + RUNTIME.block_on(async { + let mut stream = file.scan().unwrap().into_array_stream().unwrap(); + let mut arrays = Vec::new(); + let mut rows = 0; + while let Some(array) = stream.next().await { + let array = array.unwrap(); + rows += array.len(); + arrays.push(to_offset_array(array, kind)); + } + assert_eq!(rows, ROWS); + arrays + }) +} + +#[divan::bench(args = BYTE_KINDS)] +fn file_to_offset_array(bencher: Bencher, kind: ByteKind) { + let file = &*FILE; + bencher + .with_inputs(|| file) + .input_counter(|_| ItemsCount::new(ROWS)) + .bench_values(|file| read_to_offset_array(file, kind)); +} + +#[divan::bench(args = BYTE_KINDS)] +fn file_to_offset_batches(bencher: Bencher, kind: ByteKind) { + let file = &*FILE; + bencher + .with_inputs(|| file) + .input_counter(|_| ItemsCount::new(ROWS)) + .bench_values(|file| read_to_offset_batches(file, kind)); +}