From 11ee4957c15e2434b74a8502cc08447a4668a357 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Wed, 22 Jul 2026 04:03:53 -0500 Subject: [PATCH 1/2] benchmarks: count simulated object store GET requests per query Registers a CountingObjectStore (stacked under the latency wrapper) that models S3 request counting: one GET per get_opts call, one GET per coalesced range for get_ranges. Also re-chunks streamed local-file GET bodies at 2MiB via ChunkedStore to approximate network transfer granularity. Each query iteration now reports its GET request and byte counts. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01MGnT9oETcFMydc9cp95HLG --- benchmarks/src/clickbench.rs | 9 +- benchmarks/src/tpcds/run.rs | 9 +- benchmarks/src/tpch/run.rs | 9 +- benchmarks/src/util/counting_object_store.rs | 198 +++++++++++++++++++ benchmarks/src/util/mod.rs | 2 + benchmarks/src/util/options.rs | 28 ++- 6 files changed, 241 insertions(+), 14 deletions(-) create mode 100644 benchmarks/src/util/counting_object_store.rs diff --git a/benchmarks/src/clickbench.rs b/benchmarks/src/clickbench.rs index 70aaeb7d2d192..8096e0a84464e 100644 --- a/benchmarks/src/clickbench.rs +++ b/benchmarks/src/clickbench.rs @@ -19,7 +19,9 @@ use std::fs; use std::io::ErrorKind; use std::path::{Path, PathBuf}; -use crate::util::{BenchmarkRun, CommonOpt, QueryResult, print_memory_stats}; +use crate::util::{ + BenchmarkRun, CommonOpt, QueryResult, RequestCounts, print_memory_stats, +}; use clap::Args; use datafusion::logical_expr::{ExplainFormat, ExplainOption}; use datafusion::{ @@ -254,14 +256,17 @@ impl RunOpt { let mut millis = Vec::with_capacity(self.iterations()); let mut query_results = vec![]; for i in 0..self.iterations() { + let requests_before = RequestCounts::snapshot(); let start = Instant::now(); let results = ctx.sql(sql).await?.collect().await?; let elapsed = start.elapsed(); + let requests = RequestCounts::snapshot().since(&requests_before); let ms = elapsed.as_secs_f64() * 1000.0; millis.push(ms); let row_count: usize = results.iter().map(|b| b.num_rows()).sum(); println!( - "Query {query_id} iteration {i} took {ms:.1} ms and returned {row_count} rows" + "Query {query_id} iteration {i} took {ms:.1} ms and returned {row_count} rows ({} GET requests, {} bytes)", + requests.get_requests, requests.get_bytes ); query_results.push(QueryResult { elapsed, row_count }) } diff --git a/benchmarks/src/tpcds/run.rs b/benchmarks/src/tpcds/run.rs index 2e0274c935de3..9b71c4e37b697 100644 --- a/benchmarks/src/tpcds/run.rs +++ b/benchmarks/src/tpcds/run.rs @@ -19,7 +19,9 @@ use std::fs; use std::path::PathBuf; use std::sync::Arc; -use crate::util::{BenchmarkRun, CommonOpt, QueryResult, print_memory_stats}; +use crate::util::{ + BenchmarkRun, CommonOpt, QueryResult, RequestCounts, print_memory_stats, +}; use arrow::datatypes::Schema; use arrow::record_batch::RecordBatch; @@ -265,6 +267,7 @@ impl RunOpt { } for i in 0..self.iterations() { + let requests_before = RequestCounts::snapshot(); let start = Instant::now(); // query 15 is special, with 3 statements. the second statement is the one from which we @@ -276,12 +279,14 @@ impl RunOpt { } let elapsed = start.elapsed(); + let requests = RequestCounts::snapshot().since(&requests_before); let ms = elapsed.as_secs_f64() * 1000.0; millis.push(ms); info!("output:\n\n{}\n\n", pretty_format_batches(&result)?); let row_count = result.iter().map(|b| b.num_rows()).sum(); println!( - "Query {query_id} iteration {i} took {ms:.1} ms and returned {row_count} rows" + "Query {query_id} iteration {i} took {ms:.1} ms and returned {row_count} rows ({} GET requests, {} bytes)", + requests.get_requests, requests.get_bytes ); query_results.push(QueryResult { elapsed, row_count }); } diff --git a/benchmarks/src/tpch/run.rs b/benchmarks/src/tpch/run.rs index 422bcec9ea066..907982f3e0849 100644 --- a/benchmarks/src/tpch/run.rs +++ b/benchmarks/src/tpch/run.rs @@ -22,7 +22,9 @@ use super::{ TPCH_QUERY_END_ID, TPCH_QUERY_START_ID, TPCH_TABLES, get_query_sql_for_scale_factor, get_tbl_tpch_table_schema, get_tpch_table_schema, table_constraints, }; -use crate::util::{BenchmarkRun, CommonOpt, QueryResult, print_memory_stats}; +use crate::util::{ + BenchmarkRun, CommonOpt, QueryResult, RequestCounts, print_memory_stats, +}; use arrow::record_batch::RecordBatch; use arrow::util::pretty::{self, pretty_format_batches}; @@ -174,6 +176,7 @@ impl RunOpt { let sql = &get_query_sql_for_scale_factor(query_id, scale_factor)?; for i in 0..self.iterations() { + let requests_before = RequestCounts::snapshot(); let start = Instant::now(); // query 15 is special, with 3 statements. the second statement is the one from which we @@ -194,12 +197,14 @@ impl RunOpt { } let elapsed = start.elapsed(); + let requests = RequestCounts::snapshot().since(&requests_before); let ms = elapsed.as_secs_f64() * 1000.0; millis.push(ms); info!("output:\n\n{}\n\n", pretty_format_batches(&result)?); let row_count = result.iter().map(|b| b.num_rows()).sum(); println!( - "Query {query_id} iteration {i} took {ms:.1} ms and returned {row_count} rows" + "Query {query_id} iteration {i} took {ms:.1} ms and returned {row_count} rows ({} GET requests, {} bytes)", + requests.get_requests, requests.get_bytes ); query_results.push(QueryResult { elapsed, row_count }); } diff --git a/benchmarks/src/util/counting_object_store.rs b/benchmarks/src/util/counting_object_store.rs new file mode 100644 index 0000000000000..32d6ea61d8c8a --- /dev/null +++ b/benchmarks/src/util/counting_object_store.rs @@ -0,0 +1,198 @@ +// 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. + +//! An ObjectStore wrapper that counts simulated GET requests. +//! +//! For real object stores every `get`/`get_opts` call is one GET, and a +//! `get_ranges` call becomes one GET per *coalesced* range (ranges within +//! [`OBJECT_STORE_COALESCE_DEFAULT`] of each other are merged into one +//! request, following `object_store::coalesce_ranges`). This wrapper models +//! that: it counts `get_opts` as one request and `get_ranges` as its merged +//! range count, regardless of how the local filesystem actually serves them. +//! +//! Counters are process-global so benchmark runners can snapshot them around +//! each query. + +use std::fmt; +use std::ops::Range; +use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; + +use async_trait::async_trait; +use futures::stream::BoxStream; +use object_store::path::Path; +use object_store::{ + CopyOptions, GetOptions, GetResult, ListResult, MultipartUpload, ObjectMeta, + ObjectStore, PutMultipartOptions, PutOptions, PutPayload, PutResult, Result, +}; + +/// Ranges within this many bytes of one another count as a single coalesced +/// GET, matching `object_store::OBJECT_STORE_COALESCE_DEFAULT`. +const COALESCE_GAP: u64 = 1024 * 1024; + +/// Simulated GET requests issued so far (process-global) +static GET_REQUESTS: AtomicUsize = AtomicUsize::new(0); +/// Bytes requested by GETs so far (process-global) +static GET_BYTES: AtomicU64 = AtomicU64::new(0); + +/// A snapshot of the request counters +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct RequestCounts { + pub get_requests: usize, + pub get_bytes: u64, +} + +impl RequestCounts { + /// Take a snapshot of the current process-global counters + pub fn snapshot() -> Self { + Self { + get_requests: GET_REQUESTS.load(Ordering::Relaxed), + get_bytes: GET_BYTES.load(Ordering::Relaxed), + } + } + + /// Counts accumulated since an earlier snapshot + pub fn since(&self, earlier: &Self) -> Self { + Self { + get_requests: self.get_requests - earlier.get_requests, + get_bytes: self.get_bytes - earlier.get_bytes, + } + } +} + +/// Number of coalesced requests `object_store` would issue for `ranges` +fn merged_request_count(ranges: &[Range]) -> usize { + if ranges.is_empty() { + return 0; + } + let mut sorted = ranges.to_vec(); + sorted.sort_unstable_by_key(|range| range.start); + let mut count = 1; + let mut range_end = sorted[0].end; + for range in &sorted[1..] { + if range + .start + .checked_sub(range_end) + .map(|delta| delta <= COALESCE_GAP) + .unwrap_or(true) + { + range_end = range_end.max(range.end); + } else { + count += 1; + range_end = range.end; + } + } + count +} + +/// An ObjectStore wrapper counting simulated GET requests +#[derive(Debug)] +pub struct CountingObjectStore { + inner: T, +} + +impl CountingObjectStore { + pub fn new(inner: T) -> Self { + Self { inner } + } +} + +impl fmt::Display for CountingObjectStore { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "CountingObjectStore({})", self.inner) + } +} + +#[async_trait] +impl ObjectStore for CountingObjectStore { + async fn put_opts( + &self, + location: &Path, + payload: PutPayload, + opts: PutOptions, + ) -> Result { + self.inner.put_opts(location, payload, opts).await + } + + async fn put_multipart_opts( + &self, + location: &Path, + opts: PutMultipartOptions, + ) -> Result> { + self.inner.put_multipart_opts(location, opts).await + } + + async fn get_opts(&self, location: &Path, options: GetOptions) -> Result { + GET_REQUESTS.fetch_add(1, Ordering::Relaxed); + let result = self.inner.get_opts(location, options).await?; + GET_BYTES.fetch_add(result.range.end - result.range.start, Ordering::Relaxed); + Ok(result) + } + + async fn get_ranges( + &self, + location: &Path, + ranges: &[Range], + ) -> Result> { + GET_REQUESTS.fetch_add(merged_request_count(ranges), Ordering::Relaxed); + GET_BYTES.fetch_add( + ranges.iter().map(|r| r.end - r.start).sum::(), + Ordering::Relaxed, + ); + self.inner.get_ranges(location, ranges).await + } + + fn delete_stream( + &self, + locations: BoxStream<'static, Result>, + ) -> BoxStream<'static, Result> { + self.inner.delete_stream(locations) + } + + fn list(&self, prefix: Option<&Path>) -> BoxStream<'static, Result> { + self.inner.list(prefix) + } + + async fn list_with_delimiter(&self, prefix: Option<&Path>) -> Result { + self.inner.list_with_delimiter(prefix).await + } + + async fn copy_opts( + &self, + from: &Path, + to: &Path, + options: CopyOptions, + ) -> Result<()> { + self.inner.copy_opts(from, to, options).await + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_merged_request_count() { + assert_eq!(merged_request_count(&[]), 0); + assert_eq!(merged_request_count(&[0..10]), 1); + // within 1MiB gap: one request + assert_eq!(merged_request_count(&[0..10, 20..30]), 1); + // over 1MiB apart: two requests + assert_eq!(merged_request_count(&[0..10, 2_000_000..2_000_010]), 2); + // unsorted input + assert_eq!(merged_request_count(&[2_000_000..2_000_010, 0..10]), 2); + } +} diff --git a/benchmarks/src/util/mod.rs b/benchmarks/src/util/mod.rs index 6dc11c0f425bd..143e38245c323 100644 --- a/benchmarks/src/util/mod.rs +++ b/benchmarks/src/util/mod.rs @@ -16,11 +16,13 @@ // under the License. //! Shared benchmark utilities +pub mod counting_object_store; pub mod latency_object_store; mod memory; mod options; mod run; +pub use counting_object_store::RequestCounts; pub use memory::print_memory_stats; pub use options::CommonOpt; pub use run::{BenchQuery, BenchmarkRun, QueryResult}; diff --git a/benchmarks/src/util/options.rs b/benchmarks/src/util/options.rs index a3e6d2a4c5538..e3f7406e105da 100644 --- a/benchmarks/src/util/options.rs +++ b/benchmarks/src/util/options.rs @@ -28,8 +28,10 @@ use datafusion::{ prelude::SessionConfig, }; use datafusion_common::{DataFusionError, Result}; +use object_store::chunked::ChunkedStore; use object_store::local::LocalFileSystem; +use super::counting_object_store::CountingObjectStore; use super::latency_object_store::LatencyObjectStore; // Common benchmark options (don't use doc comments otherwise this doc @@ -132,19 +134,29 @@ impl CommonOpt { Ok(rt_builder) } - /// Build the runtime environment, optionally wrapping the local filesystem - /// with a throttled object store to simulate remote storage latency. + /// Build the runtime environment, wrapping the local filesystem with a + /// request-counting object store (and, with `--simulate-latency`, a + /// throttled one that mimics remote storage latency). pub fn build_runtime(&self) -> Result> { let rt = self.runtime_env_builder()?.build_arc()?; - if self.simulate_latency { - let store: Arc = - Arc::new(LatencyObjectStore::new(LocalFileSystem::new())); - let url = ObjectStoreUrl::parse("file:///")?; - rt.register_object_store(url.as_ref(), store); + // Streamed GET bodies from the local filesystem arrive in 8KiB chunks + // by default; re-chunk them at 2MiB to approximate network transfer + // granularity (irrelevant for fully-materialized reads) + const STREAM_CHUNK_SIZE: usize = 2 * 1024 * 1024; + let base = CountingObjectStore::new(ChunkedStore::new( + Arc::new(LocalFileSystem::new()), + STREAM_CHUNK_SIZE, + )); + let store: Arc = if self.simulate_latency { println!( "Simulating S3-like object store latency (get: 25-200ms, list: 40-400ms)" ); - } + Arc::new(LatencyObjectStore::new(base)) + } else { + Arc::new(base) + }; + let url = ObjectStoreUrl::parse("file:///")?; + rt.register_object_store(url.as_ref(), store); Ok(rt) } } From 6709d58e3dcaea567a997ab178cb9a57ae4e0946 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Wed, 22 Jul 2026 11:39:25 -0500 Subject: [PATCH 2/2] Fix doc link and clippy lint in counting_object_store Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01MGnT9oETcFMydc9cp95HLG --- benchmarks/src/util/counting_object_store.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/benchmarks/src/util/counting_object_store.rs b/benchmarks/src/util/counting_object_store.rs index 32d6ea61d8c8a..d7033a859e5fe 100644 --- a/benchmarks/src/util/counting_object_store.rs +++ b/benchmarks/src/util/counting_object_store.rs @@ -19,7 +19,7 @@ //! //! For real object stores every `get`/`get_opts` call is one GET, and a //! `get_ranges` call becomes one GET per *coalesced* range (ranges within -//! [`OBJECT_STORE_COALESCE_DEFAULT`] of each other are merged into one +//! `OBJECT_STORE_COALESCE_DEFAULT` of each other are merged into one //! request, following `object_store::coalesce_ranges`). This wrapper models //! that: it counts `get_opts` as one request and `get_ranges` as its merged //! range count, regardless of how the local filesystem actually serves them. @@ -185,6 +185,7 @@ mod tests { use super::*; #[test] + #[expect(clippy::single_range_in_vec_init)] fn test_merged_request_count() { assert_eq!(merged_request_count(&[]), 0); assert_eq!(merged_request_count(&[0..10]), 1);