diff --git a/Cargo.lock b/Cargo.lock index 3bb94e339d..77c8c2d436 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -88,6 +88,15 @@ dependencies = [ "alloc-no-stdlib", ] +[[package]] +name = "alloca" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5a7d05ea6aea7e9e64d25b9156ba2fee3fdd659e34e41063cd2fc7cd020d7f4" +dependencies = [ + "cc", +] + [[package]] name = "allocator-api2" version = "0.2.21" @@ -103,6 +112,12 @@ dependencies = [ "libc", ] +[[package]] +name = "anes" +version = "0.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4b46cbb362ab8752921c97e041f5e366ee6297bd428a31275b9fcf1e380f7299" + [[package]] name = "anstream" version = "1.0.0" @@ -1262,6 +1277,12 @@ dependencies = [ "either", ] +[[package]] +name = "cast" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "37b2a672a2cb129a2e41c10b1224bb368f9f37a2b16b612598138befd7b37eb5" + [[package]] name = "cbc" version = "0.1.2" @@ -1336,6 +1357,33 @@ dependencies = [ "phf", ] +[[package]] +name = "ciborium" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "42e69ffd6f0917f5c029256a24d0161db17cea3997d185db0d35926308770f0e" +dependencies = [ + "ciborium-io", + "ciborium-ll", + "serde", +] + +[[package]] +name = "ciborium-io" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "05afea1e0a06c9be33d539b876f1ce3692f4afea2cb41f740e7743225ed1c757" + +[[package]] +name = "ciborium-ll" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "57663b653d948a338bfb3eeba9bb2fd5fcfaecb9e199e87e1eda4d9e8b240fd9" +dependencies = [ + "ciborium-io", + "half", +] + [[package]] name = "cipher" version = "0.4.4" @@ -1413,7 +1461,7 @@ version = "3.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "faf9468729b8cbcea668e36183cb69d317348c2e08e994829fb56ebfdfbaac34" dependencies = [ - "windows-sys 0.48.0", + "windows-sys 0.61.2", ] [[package]] @@ -1566,6 +1614,39 @@ dependencies = [ "cfg-if 1.0.4", ] +[[package]] +name = "criterion" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "950046b2aa2492f9a536f5f4f9a3de7b9e2476e575e05bd6c333371add4d98f3" +dependencies = [ + "alloca", + "anes", + "cast", + "ciborium", + "clap", + "criterion-plot", + "itertools 0.13.0", + "num-traits", + "oorandom", + "page_size", + "regex", + "serde", + "serde_json", + "tinytemplate", + "walkdir", +] + +[[package]] +name = "criterion-plot" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d8d80a2f4f5b554395e47b5d8305bc3d27813bacb73493eb1001e8f76dae29ea" +dependencies = [ + "cast", + "itertools 0.13.0", +] + [[package]] name = "crossbeam-channel" version = "0.5.16" @@ -2720,6 +2801,7 @@ dependencies = [ "bytes", "chrono", "crc32fast", + "criterion", "derive_builder", "expect-test", "fastnum", @@ -3914,6 +3996,12 @@ version = "0.1.13" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "269bca4c2591a28585d6bf10d9ed0332b7d76900a1b02bec41bdc3a2cdcda107" +[[package]] +name = "oorandom" +version = "11.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6790f58c7ff633d8771f42965289203411a5e5c68388703c06e14f24770b41e" + [[package]] name = "opaque-debug" version = "0.3.1" @@ -4225,6 +4313,16 @@ version = "0.5.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1a80800c0488c3a21695ea981a54918fbb37abf04f4d0720c453632255e2ff0e" +[[package]] +name = "page_size" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "30d5b2194ed13191c1999ae0704b7839fb18384fa22e49b57eeaa97d79ce40da" +dependencies = [ + "libc", + "winapi", +] + [[package]] name = "parking" version = "2.2.1" @@ -6387,6 +6485,16 @@ dependencies = [ "zerovec", ] +[[package]] +name = "tinytemplate" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "be4d6b5f19ff7664e8c98d03e2139cb510db9b0a60b55f8e8709b689d939b6bc" +dependencies = [ + "serde", + "serde_json", +] + [[package]] name = "tinyvec" version = "1.11.0" @@ -7144,7 +7252,7 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys 0.48.0", + "windows-sys 0.61.2", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index bb9770275a..0b3a8a9ec3 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -80,6 +80,9 @@ bytes = "1.11" cfg-if = "1" chrono = "0.4.41" crc32fast = "1" +criterion = { version = "0.8", default-features = false, features = [ + "cargo_bench_support", +] } derive_builder = "0.20" expect-test = "1" fastnum = { version = "0.7", default-features = false, features = [ diff --git a/crates/iceberg/Cargo.toml b/crates/iceberg/Cargo.toml index f22ed59b9b..bda58449bf 100644 --- a/crates/iceberg/Cargo.toml +++ b/crates/iceberg/Cargo.toml @@ -86,6 +86,7 @@ zeroize = { workspace = true } zstd = { workspace = true } [dev-dependencies] +criterion = { workspace = true } expect-test = { workspace = true } iceberg_test_utils = { path = "../test_utils", features = ["tests"] } minijinja = { workspace = true } @@ -97,6 +98,10 @@ serde_arrow = { version = "0.14", features = ["arrow-59"] } tempfile = { workspace = true } tracing-subscriber = { workspace = true } +[[bench]] +harness = false +name = "manifest_parse" + [package.metadata.cargo-machete] # These dependencies are added to ensure minimal dependency version ignored = ["tap"] diff --git a/crates/iceberg/benches/manifest_parse.rs b/crates/iceberg/benches/manifest_parse.rs new file mode 100644 index 0000000000..9c7f04027d --- /dev/null +++ b/crates/iceberg/benches/manifest_parse.rs @@ -0,0 +1,219 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::hint::black_box; +use std::sync::Arc; + +use apache_avro::types::Value; +use apache_avro::{Codec, DeflateSettings}; +use criterion::{BenchmarkId, Criterion, Throughput, criterion_group, criterion_main}; +use iceberg::io::FileIO; +use iceberg::spec::{ + DataContentType, DataFileBuilder, DataFileFormat, Datum, Literal, Manifest, + ManifestWriterBuilder, NestedField, PartitionSpec, PrimitiveType, Schema, Struct, Transform, + Type, +}; + +const COLUMN_TYPES: [PrimitiveType; 12] = [ + PrimitiveType::Long, + PrimitiveType::String, + PrimitiveType::Int, + PrimitiveType::Long, + PrimitiveType::Float, + PrimitiveType::Double, + PrimitiveType::String, + PrimitiveType::Boolean, + PrimitiveType::Date, + PrimitiveType::Timestamp, + PrimitiveType::String, + PrimitiveType::Long, +]; + +fn columns(num_columns: i32) -> impl Iterator { + (1..=num_columns).map(|id| (id, &COLUMN_TYPES[(id as usize - 1) % COLUMN_TYPES.len()])) +} + +fn bound(ty: &PrimitiveType, v: i64) -> Datum { + match ty { + PrimitiveType::Boolean => Datum::bool(v % 2 == 1), + PrimitiveType::Int => Datum::int(v as i32), + PrimitiveType::Long => Datum::long(v), + PrimitiveType::Float => Datum::float(v as f32), + PrimitiveType::Double => Datum::double(v as f64), + PrimitiveType::Date => Datum::date(v as i32), + PrimitiveType::Timestamp => Datum::timestamp_micros(v), + PrimitiveType::String => Datum::string(format!("{v:016}")), + _ => unreachable!(), + } +} + +/// Returns a V2 data manifest with `num_entries` files, each with stats for +/// all `num_columns` columns and an identity partition on a string column. +fn write_manifest(num_entries: u64, num_columns: i32) -> Vec { + let columns = || columns(num_columns); + let schema = Arc::new( + Schema::builder() + .with_fields(columns().map(|(id, ty)| { + let name = if id == 2 { + "category".to_string() + } else { + format!("c{id}") + }; + NestedField::optional(id, name, Type::Primitive(ty.clone())).into() + })) + .build() + .unwrap(), + ); + let partition_spec = PartitionSpec::builder(schema.clone()) + .add_partition_field("category", "category", Transform::Identity) + .unwrap() + .build() + .unwrap(); + let file_io = FileIO::new_with_memory(); + let path = "memory://manifest.avro"; + let mut writer = ManifestWriterBuilder::new( + file_io.new_output(path).unwrap(), + Some(1), + schema, + partition_spec, + ) + .build_v2_data(); + + for i in 0..num_entries { + let n = i as i64; + let category = format!("c{}", i % 10); + let data_file = DataFileBuilder::default() + .content(DataContentType::Data) + .file_path(format!( + "s3://bucket/warehouse/db/table/data/category={category}/{i:08}-0-6f1e1a2b-3c4d-5e6f-7a8b-9c0d1e2f3a4b.parquet" + )) + .file_format(DataFileFormat::Parquet) + .partition(Struct::from_iter([Some(Literal::string(category))])) + .record_count(10_000 + i) + .file_size_in_bytes(1_000_000 + i) + .column_sizes(columns().map(|(id, _)| (id, 1000 + id as u64 + i)).collect()) + .value_counts(columns().map(|(id, _)| (id, 10_000 + i)).collect()) + .null_value_counts(columns().map(|(id, _)| (id, i % 7)).collect()) + .nan_value_counts( + columns() + .filter(|(_, ty)| matches!(ty, PrimitiveType::Float | PrimitiveType::Double)) + .map(|(id, _)| (id, 0)) + .collect(), + ) + .lower_bounds( + columns() + .map(|(id, ty)| (id, bound(ty, n * 13 + id as i64))) + .collect(), + ) + .upper_bounds( + columns() + .map(|(id, ty)| (id, bound(ty, n * 17 + id as i64 + 1000))) + .collect(), + ) + .split_offsets(Some(vec![4])) + .sort_order_id(0) + .build() + .unwrap(); + writer.add_file(data_file, 1).unwrap(); + } + + futures::executor::block_on(async { + writer.write_manifest_file().await.unwrap(); + file_io.new_input(path).unwrap().read().await.unwrap() + }) + .to_vec() +} + +/// Rewrites `bytes` the way Iceberg Java writes V2 manifests by default: +/// without the `data_file` fields it doesn't write, deflate, and 64 KB blocks. +fn rewrite_like_java(bytes: &[u8]) -> Vec { + const OMITTED: [&str; 3] = ["first_row_id", "content_offset", "content_size_in_bytes"]; + let retain = |fields: &mut Vec| { + fields.retain(|f| !OMITTED.contains(&f["name"].as_str().unwrap())) + }; + + let reader = apache_avro::Reader::new(bytes).unwrap(); + let mut schema = serde_json::to_value(reader.writer_schema()).unwrap(); + for field in schema["fields"].as_array_mut().unwrap() { + if field["name"] == "data_file" { + retain(field["type"]["fields"].as_array_mut().unwrap()); + } + } + let schema = apache_avro::Schema::parse(&schema).unwrap(); + let mut writer = apache_avro::Writer::builder() + .schema(&schema) + .writer(Vec::new()) + .codec(Codec::Deflate(DeflateSettings::default())) + .block_size(64_000) + .build() + .unwrap(); + for (key, value) in reader.user_metadata() { + writer.add_user_metadata(key.clone(), value).unwrap(); + } + for entry in reader { + let Value::Record(mut entry) = entry.unwrap() else { + unreachable!() + }; + for (name, data_file) in &mut entry { + if let ("data_file", Value::Record(fields)) = (name.as_str(), data_file) { + fields.retain(|(name, _)| !OMITTED.contains(&name.as_str())); + } + } + writer.append_value(Value::Record(entry)).unwrap(); + } + writer.into_inner().unwrap() +} + +fn bench_parse_avro(c: &mut Criterion) { + let mut group = c.benchmark_group("manifest_parse_avro"); + group.sample_size(10); + for num_entries in [1_000, 2_000, 5_000, 10_000, 20_000, 50_000, 100_000] { + let bytes = write_manifest(num_entries, 12); + let java_like = rewrite_like_java(&bytes); + assert_eq!( + Manifest::parse_avro(&java_like).unwrap(), + Manifest::parse_avro(&bytes).unwrap() + ); + group.throughput(Throughput::Elements(num_entries)); + group.bench_with_input( + BenchmarkId::from_parameter(num_entries), + &java_like, + |b, bytes| b.iter(|| Manifest::parse_avro(black_box(bytes)).unwrap()), + ); + } + group.finish(); +} + +/// Parses 10,000-entry manifests from about 1 MB to 100 MB, made larger by +/// adding columns with stats. +fn bench_parse_avro_by_size(c: &mut Criterion) { + let mut group = c.benchmark_group("manifest_parse_avro_by_columns"); + group.sample_size(10); + for num_columns in [9, 17, 44, 82, 145, 330, 610] { + let java_like = rewrite_like_java(&write_manifest(10_000, num_columns)); + group.throughput(Throughput::Bytes(java_like.len() as u64)); + group.bench_with_input( + BenchmarkId::from_parameter(num_columns), + &java_like, + |b, bytes| b.iter(|| Manifest::parse_avro(black_box(bytes)).unwrap()), + ); + } + group.finish(); +} + +criterion_group!(benches, bench_parse_avro, bench_parse_avro_by_size); +criterion_main!(benches);