diff --git a/Cargo.lock b/Cargo.lock index c05aae6923e..2b74261977e 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/Cargo.toml b/vortex-arrow/Cargo.toml index 11f2308451c..11d468a6d75 100644 --- a/vortex-arrow/Cargo.toml +++ b/vortex-arrow/Cargo.toml @@ -47,4 +47,5 @@ vortex-zstd = { workspace = true } [[bench]] name = "to_arrow" +path = "benches/to_arrow/main.rs" harness = false diff --git a/vortex-arrow/benches/to_arrow.rs b/vortex-arrow/benches/to_arrow.rs deleted file mode 100644 index e6016ec9705..00000000000 --- a/vortex-arrow/benches/to_arrow.rs +++ /dev/null @@ -1,386 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -#![expect(clippy::unwrap_used)] - -use std::sync::Arc; -use std::sync::LazyLock; - -use arrow_schema::DataType; -use arrow_schema::Field; -use divan::Bencher; -use divan::counter::ItemsCount; -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::ChunkedArray; -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::StructArray; -use vortex_array::arrays::VarBinViewArray; -use vortex_array::builders::VarBinBuilder; -use vortex_array::builders::VarBinViewBuilder; -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::session::ArraySessionExt; -#[expect( - deprecated, - reason = "benchmark comparing deprecated method with new one" -)] -use vortex_arrow::ArrowArrayExecutor; -use vortex_arrow::ArrowSessionExt; -#[allow(deprecated)] -use vortex_arrow::dtype::ToArrowType as _; -use vortex_fsst::fsst_compress; -use vortex_fsst::fsst_train_compressor; -use vortex_mask::Mask; -use vortex_onpair::DEFAULT_CONFIG; -use vortex_onpair::onpair_compress; -use vortex_session::VortexSession; -use vortex_zstd::Zstd; - -fn main() { - LazyLock::force(&SESSION); - divan::main(); -} - -static SESSION: LazyLock = LazyLock::new(|| { - let session = array_session(); - vortex_fsst::initialize(&session); - vortex_onpair::initialize(&session); - session.arrays().register(Zstd); - session -}); - -fn schema() -> DType { - let fields = StructFields::from_iter([ - ( - "primitive", - DType::Primitive(PType::F32, Nullability::Nullable), - ), - ( - "list", - DType::List( - Arc::new(DType::Binary(Nullability::NonNullable)), - Nullability::Nullable, - ), - ), - ( - "decimal", - DType::Decimal(DecimalDType::new(19, 10), Nullability::Nullable), - ), - ]); - DType::Struct(fields, Nullability::NonNullable) -} - -fn array() -> ArrayRef { - StructArray::from_fields(&[ - ( - "primitive", - PrimitiveArray::from_iter(0i16..1024).into_array(), - ), - ( - "list", - ListArray::from_iter_slow::( - (0..1024).map(|_| vec!["a", "b", "c"]).collect::>(), - Arc::new(DType::Utf8(Nullability::NonNullable)), - ) - .unwrap() - .into_array(), - ), - ( - "decimal", - DecimalArray::from_iter(0i64..1024, DecimalDType::new(19, 2)).into_array(), - ), - ]) - .unwrap() - .into_array() -} - -#[divan::bench] -fn to_arrow_dtype(bencher: Bencher) { - bencher.with_inputs(schema).bench_values(|dtype| { - #[expect(deprecated, reason = "benchmarking deprecated code path")] - dtype.to_arrow_dtype().unwrap() - }); -} - -#[allow(non_snake_case)] -#[divan::bench] -fn ArrowExportVTable_to_arrow_field(bencher: Bencher) { - bencher - .with_inputs(schema) - .bench_values(|dtype| SESSION.arrow().to_arrow_field("", &dtype).unwrap()) -} - -#[derive(Clone, Copy, Debug)] -enum StringEncoding { - View, - Fsst, - OnPair, - Zstd, - 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, -} - -const STRING_ENCODINGS: &[StringEncoding] = &[ - 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 OFFSET_STRING_ROWS: usize = 100_000; -const OFFSET_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 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 dictionary_values() -> VarBinViewArray { - structured_strings(DICTIONARY_SIZE) -} - -fn dictionary_codes() -> ArrayRef { - PrimitiveArray::from_iter( - (0..OFFSET_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 { - // Keep Filter as a lazy intermediate so benchmark setup cannot optimize it away. - FilterArray::new(array, half_rows_mask()).into_array() -} - -fn chunked_fsst(ctx: &mut ExecutionCtx) -> ArrayRef { - let source = structured_strings(OFFSET_STRING_ROWS).into_array(); - 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| { - let start = chunk_index * chunk_size; - let end = if chunk_index + 1 == OFFSET_STRING_CHUNKS { - OFFSET_STRING_ROWS - } else { - start + chunk_size - }; - let chunk = source.slice(start..end).unwrap(); - fsst_compress(&chunk, &compressor, ctx) - .unwrap() - .into_array() - }); - ChunkedArray::try_new(chunks, source.dtype().clone()) - .unwrap() - .into_array() -} - -fn fsst(source: ArrayRef, ctx: &mut ExecutionCtx) -> ArrayRef { - let compressor = fsst_train_compressor(&source, ctx).unwrap(); - fsst_compress(&source, &compressor, ctx) - .unwrap() - .into_array() -} - -fn string_array(encoding: StringEncoding) -> 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(), - &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() - } - } -} - -/// 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()); - - bencher - .with_inputs(|| (array.clone(), SESSION.create_execution_ctx())) - .input_counter(|(array, _)| ItemsCount::new(array.len())) - .bench_values(|(array, mut ctx)| { - SESSION - .arrow() - .execute_arrow(array, Some(&field), &mut ctx) - .unwrap() - }); -} - -/// 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); - - bencher - .with_inputs(|| (array.clone(), SESSION.create_execution_ctx())) - .input_counter(|(array, _)| ItemsCount::new(array.len())) - .bench_values(|(array, mut ctx)| { - let mut builder = - VarBinBuilder::::with_capacity(array.dtype().clone(), array.len()); - array.append_to_builder(&mut builder, &mut ctx).unwrap(); - builder.finish_into_varbin() - }); -} - -/// 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); - - bencher - .with_inputs(|| (array.clone(), SESSION.create_execution_ctx())) - .input_counter(|(array, _)| ItemsCount::new(array.len())) - .bench_values(|(array, mut ctx)| { - let mut builder = VarBinViewBuilder::with_capacity(array.dtype().clone(), array.len()); - array.append_to_builder(&mut builder, &mut ctx).unwrap(); - builder.finish_into_varbinview() - }); -} - -#[divan::bench] -fn to_arrow_array(bencher: Bencher) { - bencher - .with_inputs(|| (array(), SESSION.create_execution_ctx())) - .bench_values(|(array, mut ctx)| { - #[expect(deprecated, reason = "benchmarking deprecated code path")] - array.execute_arrow(None, &mut ctx).unwrap() - }); -} - -#[allow(non_snake_case)] -#[divan::bench] -fn ArrowExportVTable_execute_arrow(bencher: Bencher) { - bencher - .with_inputs(|| (array(), SESSION.create_execution_ctx())) - .bench_values(|(array, mut ctx)| { - SESSION - .arrow() - .execute_arrow(array, None, &mut ctx) - .unwrap() - }) -} diff --git a/vortex-arrow/benches/to_arrow/main.rs b/vortex-arrow/benches/to_arrow/main.rs new file mode 100644 index 00000000000..51761d8c36d --- /dev/null +++ b/vortex-arrow/benches/to_arrow/main.rs @@ -0,0 +1,132 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#![expect(clippy::unwrap_used)] + +mod string; + +use std::sync::Arc; +use std::sync::LazyLock; + +use divan::Bencher; +use vortex_array::ArrayRef; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::array_session; +use vortex_array::arrays::DecimalArray; +use vortex_array::arrays::ListArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::StructArray; +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::session::ArraySessionExt; +#[expect( + deprecated, + reason = "benchmark comparing deprecated method with new one" +)] +use vortex_arrow::ArrowArrayExecutor; +use vortex_arrow::ArrowSessionExt; +#[allow(deprecated)] +use vortex_arrow::dtype::ToArrowType as _; +use vortex_session::VortexSession; +use vortex_zstd::Zstd; + +fn main() { + LazyLock::force(&SESSION); + divan::main(); +} + +static SESSION: LazyLock = LazyLock::new(|| { + let session = array_session(); + vortex_fsst::initialize(&session); + vortex_onpair::initialize(&session); + session.arrays().register(Zstd); + session +}); + +fn schema() -> DType { + let fields = StructFields::from_iter([ + ( + "primitive", + DType::Primitive(PType::F32, Nullability::Nullable), + ), + ( + "list", + DType::List( + Arc::new(DType::Binary(Nullability::NonNullable)), + Nullability::Nullable, + ), + ), + ( + "decimal", + DType::Decimal(DecimalDType::new(19, 10), Nullability::Nullable), + ), + ]); + DType::Struct(fields, Nullability::NonNullable) +} + +fn array() -> ArrayRef { + StructArray::from_fields(&[ + ( + "primitive", + PrimitiveArray::from_iter(0i16..1024).into_array(), + ), + ( + "list", + ListArray::from_iter_slow::( + (0..1024).map(|_| vec!["a", "b", "c"]).collect::>(), + Arc::new(DType::Utf8(Nullability::NonNullable)), + ) + .unwrap() + .into_array(), + ), + ( + "decimal", + DecimalArray::from_iter(0i64..1024, DecimalDType::new(19, 2)).into_array(), + ), + ]) + .unwrap() + .into_array() +} + +#[divan::bench] +fn to_arrow_dtype(bencher: Bencher) { + bencher.with_inputs(schema).bench_values(|dtype| { + #[expect(deprecated, reason = "benchmarking deprecated code path")] + dtype.to_arrow_dtype().unwrap() + }); +} + +#[allow(non_snake_case)] +#[divan::bench] +fn ArrowExportVTable_to_arrow_field(bencher: Bencher) { + bencher + .with_inputs(schema) + .bench_values(|dtype| SESSION.arrow().to_arrow_field("", &dtype).unwrap()) +} + +#[divan::bench] +fn to_arrow_array(bencher: Bencher) { + bencher + .with_inputs(|| (array(), SESSION.create_execution_ctx())) + .bench_values(|(array, mut ctx)| { + #[expect(deprecated, reason = "benchmarking deprecated code path")] + array.execute_arrow(None, &mut ctx).unwrap() + }); +} + +#[allow(non_snake_case)] +#[divan::bench] +fn ArrowExportVTable_execute_arrow(bencher: Bencher) { + bencher + .with_inputs(|| (array(), SESSION.create_execution_ctx())) + .bench_values(|(array, mut ctx)| { + SESSION + .arrow() + .execute_arrow(array, None, &mut ctx) + .unwrap() + }) +} diff --git a/vortex-arrow/benches/to_arrow/string.rs b/vortex-arrow/benches/to_arrow/string.rs new file mode 100644 index 00000000000..3c2d192dd92 --- /dev/null +++ b/vortex-arrow/benches/to_arrow/string.rs @@ -0,0 +1,443 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +mod nested; + +use std::fmt::Display; +use std::fmt::Formatter; + +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::arrays::BoolArray; +use vortex_array::arrays::ChunkedArray; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::DictArray; +use vortex_array::arrays::FilterArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::SliceArray; +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::Nullability; +use vortex_array::scalar::Scalar; +use vortex_arrow::ArrowSessionExt; +use vortex_fsst::fsst_compress; +use vortex_fsst::fsst_train_compressor; +use vortex_mask::Mask; +use vortex_onpair::DEFAULT_CONFIG; +use vortex_onpair::onpair_compress; +use vortex_zstd::Zstd; + +use crate::SESSION; + +#[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, + 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, +]; +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 STRING_ROWS: usize = 100_000; +const STRING_CHUNKS: usize = 4; +const DICTIONARY_SIZE: usize = 2_048; + +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 string_export_cases() -> Vec { + iproduct!(string_cases(), ARROW_STRING_LAYOUTS.iter().copied()) + .map(|(array, layout)| StringExportCase { array, layout }) + .collect() +} + +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..STRING_ROWS).map(|index| u16::try_from(index % DICTIONARY_SIZE).unwrap()), + ) + .into_array() +} + +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 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 fsst(source: ArrayRef, ctx: &mut ExecutionCtx) -> ArrayRef { + let compressor = fsst_train_compressor(&source, ctx).unwrap(); + 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 == STRING_CHUNKS { + STRING_ROWS + } else { + start + chunk_size + }; + encode_strings(structured_strings(end - start, validity), encoding, ctx) + }); + ChunkedArray::try_new(chunks, DType::Utf8(validity.nullability())) + .unwrap() + .into_array() +} + +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 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(); + let array = match case.structure { + StringStructure::Flat => encode_strings( + structured_strings(STRING_ROWS, case.validity), + case.encoding, + &mut ctx, + ), + 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) +} + +/// 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())) + .bench_values(|(array, mut ctx)| { + SESSION + .arrow() + .execute_arrow(array, Some(&field), &mut ctx) + .unwrap() + }); +} + +/// 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())) + .input_counter(|(array, _)| ItemsCount::new(array.len())) + .bench_values(|(array, mut ctx)| { + let mut builder = + VarBinBuilder::::with_capacity(array.dtype().clone(), array.len()); + array.append_to_builder(&mut builder, &mut ctx).unwrap(); + builder.finish_into_varbin() + }); +} + +/// 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())) + .input_counter(|(array, _)| ItemsCount::new(array.len())) + .bench_values(|(array, mut ctx)| { + let mut builder = VarBinViewBuilder::with_capacity(array.dtype().clone(), array.len()); + array.append_to_builder(&mut builder, &mut ctx).unwrap(); + builder.finish_into_varbinview() + }); +} diff --git a/vortex-arrow/benches/to_arrow/string/nested.rs b/vortex-arrow/benches/to_arrow/string/nested.rs new file mode 100644 index 00000000000..da095035d7b --- /dev/null +++ b/vortex-arrow/benches/to_arrow/string/nested.rs @@ -0,0 +1,119 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::fmt::Display; +use std::fmt::Formatter; + +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::VortexSessionExecute; +use vortex_arrow::ArrowSessionExt; + +use super::STRING_ROWS; +use super::StringEncoding; +use super::StringValidity; +use super::encode_strings; +use super::filtered; +use super::sliced; +use super::structured_strings; +use super::take; +use crate::SESSION; + +#[derive(Clone, Copy)] +enum NestedOperator { + TakeFilter, + FilterTake, + SliceFilter, + FilterSlice, +} + +impl Display for NestedOperator { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + f.write_str(match self { + Self::TakeFilter => "take_filter", + Self::FilterTake => "filter_take", + Self::SliceFilter => "slice_filter", + Self::FilterSlice => "filter_slice", + }) + } +} + +#[derive(Clone, Copy)] +struct NestedStringCase { + encoding: StringEncoding, + operators: NestedOperator, + validity: StringValidity, +} + +impl Display for NestedStringCase { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "{}/{}/{}", self.encoding, self.operators, self.validity) + } +} + +const ENCODINGS: &[StringEncoding] = &[ + StringEncoding::Fsst, + StringEncoding::OnPair, + StringEncoding::Zstd, +]; +const OPERATORS: &[NestedOperator] = &[ + NestedOperator::TakeFilter, + NestedOperator::FilterTake, + NestedOperator::SliceFilter, + NestedOperator::FilterSlice, +]; +const VALIDITIES: &[StringValidity] = &[StringValidity::NonNullable, StringValidity::Nullable]; + +fn nested_string_cases() -> Vec { + iproduct!( + ENCODINGS.iter().copied(), + OPERATORS.iter().copied(), + VALIDITIES.iter().copied() + ) + .map(|(encoding, operators, validity)| NestedStringCase { + encoding, + operators, + validity, + }) + .collect() +} + +fn apply_nested_operators(array: ArrayRef, operators: NestedOperator) -> ArrayRef { + match operators { + NestedOperator::TakeFilter => take(filtered(array)), + NestedOperator::FilterTake => filtered(take(array)), + NestedOperator::SliceFilter => sliced(filtered(array)), + NestedOperator::FilterSlice => filtered(sliced(array)), + } +} + +fn nested_string_array(case: NestedStringCase) -> ArrayRef { + let mut ctx = SESSION.create_execution_ctx(); + let array = encode_strings( + structured_strings(STRING_ROWS, case.validity), + case.encoding, + &mut ctx, + ); + apply_nested_operators(array, case.operators) +} + +/// Measures offset array export through two lazy operators. +#[divan::bench(args = nested_string_cases())] +fn nested_string_export(bencher: Bencher, case: NestedStringCase) { + let array = nested_string_array(case); + let field = Field::new("value", DataType::Utf8, array.dtype().is_nullable()); + + bencher + .with_inputs(|| (array.clone(), SESSION.create_execution_ctx())) + .input_counter(|(array, _)| ItemsCount::new(array.len())) + .bench_values(|(array, mut ctx)| { + SESSION + .arrow() + .execute_arrow(array, Some(&field), &mut ctx) + .unwrap() + }); +} diff --git a/vortex-arrow/src/executor/byte.rs b/vortex-arrow/src/executor/byte.rs index 434269b04e7..6760cfd7f6b 100644 --- a/vortex-arrow/src/executor/byte.rs +++ b/vortex-arrow/src/executor/byte.rs @@ -16,8 +16,10 @@ use vortex_array::ArrayRef; use vortex_array::ArrayView; use vortex_array::Canonical; use vortex_array::ExecutionCtx; -use vortex_array::arrays::Chunked; -use vortex_array::arrays::Constant; +use vortex_array::arrays::Dict; +use vortex_array::arrays::Filter; +use vortex_array::arrays::ScalarFn; +use vortex_array::arrays::Slice; use vortex_array::arrays::VarBin; use vortex_array::arrays::varbin::VarBinArraySlotsExt; use vortex_array::builders::VarBinBuilder; @@ -34,18 +36,24 @@ use vortex_error::vortex_err; use crate::executor::validity::to_arrow_null_buffer; -/// Matches the encodings [`to_arrow_byte_array`] requires for export. +/// Matches byte arrays that should directly append to a `VarBinBuilder`. /// -/// `Chunked` and `Constant` are matched to stop execution before it destroys them: they have -/// specialized `append_to_builder` impls (chunk-wise append, scalar repeat) that the builder -/// fallback exploits. -struct ArrowByteExportable; - -impl Matcher for ArrowByteExportable { +/// Lazy operators must run before export. +/// `execute_until` removes them before this matcher accepts the array. +/// The exporter then calls `append_to_builder`. +/// A specialized implementation decodes directly into the builder. +/// This avoids an intermediate canonical `VarBinView`. +struct ShouldDirectlyAppend; + +impl Matcher for ShouldDirectlyAppend { type Match<'a> = &'a ArrayRef; fn try_match(array: &ArrayRef) -> Option> { - (array.is::() || array.is::() || array.is::()).then_some(array) + (!array.is::() + && !array.is::() + && !array.is::() + && !array.is::()) + .then_some(array) } } @@ -72,7 +80,7 @@ where let target_is_utf8 = matches!(T::DATA_TYPE, DataType::Utf8 | DataType::LargeUtf8); let validate_utf8 = target_is_utf8 && !source_is_utf8; - let array = array.execute_until::(ctx)?; + let array = array.execute_until::(ctx)?; // If the Vortex array is in VarBin format, we can directly convert it. if let Some(array) = array.as_opt::() { 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)); +}