From 96239fc034a1bab56783b5ee36cc811846b8ccae Mon Sep 17 00:00:00 2001 From: GatewayJ <835269233@qq.com> Date: Sun, 30 Aug 2026 12:07:20 +0800 Subject: [PATCH] feat(s3select): report uncompressed input byte metrics (#6865) --- crates/e2e_test/src/reliant/sql.rs | 139 +++++- crates/s3select-api/src/lib.rs | 2 + crates/s3select-api/src/metrics.rs | 88 ++++ crates/s3select-api/src/object_store.rs | 427 ++++++++++++++++-- crates/s3select-api/src/query/dispatcher.rs | 166 +++++++ crates/s3select-api/src/query/mod.rs | 47 +- crates/s3select-api/src/query/session.rs | 77 +++- .../s3select-query/src/dispatcher/manager.rs | 100 +++- crates/s3select-query/src/instance.rs | 99 +++- rustfs/src/app/select_object.rs | 121 +++-- 10 files changed, 1151 insertions(+), 115 deletions(-) create mode 100644 crates/s3select-api/src/metrics.rs diff --git a/crates/e2e_test/src/reliant/sql.rs b/crates/e2e_test/src/reliant/sql.rs index c44682cf0..d9f821a15 100644 --- a/crates/e2e_test/src/reliant/sql.rs +++ b/crates/e2e_test/src/reliant/sql.rs @@ -17,7 +17,8 @@ use crate::common::{RustFSTestEnvironment, init_logging}; use aws_sdk_s3::Client; use aws_sdk_s3::error::ProvideErrorMetadata; use aws_sdk_s3::types::{ - CsvInput, CsvOutput, ExpressionType, FileHeaderInfo, InputSerialization, JsonInput, JsonOutput, JsonType, OutputSerialization, + CsvInput, CsvOutput, ExpressionType, FileHeaderInfo, InputSerialization, JsonInput, JsonOutput, JsonType, + OutputSerialization, RequestProgress, }; use bytes::Bytes; use std::error::Error; @@ -104,6 +105,142 @@ async fn process_select_response( .map_err(|_| -> Box { "Select response timed out".into() })? } +async fn assert_input_byte_stats( + client: &Client, + object: &str, + body: &[u8], + expression: &str, + input_serialization: InputSerialization, + output_serialization: OutputSerialization, + progress_enabled: bool, +) -> TestResult<()> { + client + .put_object() + .bucket(BUCKET) + .key(object) + .body(Bytes::copy_from_slice(body).into()) + .send() + .await?; + + let mut request = client + .select_object_content() + .bucket(BUCKET) + .key(object) + .expression(expression) + .expression_type(ExpressionType::Sql) + .input_serialization(input_serialization) + .output_serialization(output_serialization); + if progress_enabled { + request = request.request_progress(RequestProgress::builder().enabled(true).build()); + } + let response = request.send().await?; + + let mut payload = response.payload; + let mut records_len = 0_u64; + let mut last_progress: Option = None; + let mut stats = None; + let mut saw_end = false; + while let Some(event) = payload.recv().await? { + match event { + aws_sdk_s3::types::SelectObjectContentEventStream::Records(records) => { + if let Some(bytes) = records.payload { + records_len = records_len.saturating_add(u64::try_from(bytes.as_ref().len())?); + } + } + aws_sdk_s3::types::SelectObjectContentEventStream::Progress(event) => { + let details = event.details.ok_or("Progress event did not contain details")?; + if let Some(previous) = last_progress.as_ref() { + assert!(details.bytes_scanned() >= previous.bytes_scanned()); + assert!(details.bytes_processed() >= previous.bytes_processed()); + assert!(details.bytes_returned() >= previous.bytes_returned()); + } + last_progress = Some(details); + } + aws_sdk_s3::types::SelectObjectContentEventStream::Stats(event) => stats = event.details, + aws_sdk_s3::types::SelectObjectContentEventStream::End(_) => { + saw_end = true; + break; + } + _ => {} + } + } + + let stats = stats.ok_or("Select response ended without a Stats event")?; + let input_len = i64::try_from(body.len())?; + assert_eq!(stats.bytes_scanned(), Some(input_len)); + assert_eq!(stats.bytes_processed(), Some(input_len)); + assert_eq!(stats.bytes_returned(), Some(i64::try_from(records_len)?)); + if progress_enabled { + let progress = last_progress.ok_or("Select response ended without a Progress event")?; + assert_eq!(progress.bytes_scanned(), stats.bytes_scanned()); + assert_eq!(progress.bytes_processed(), stats.bytes_processed()); + assert_eq!(progress.bytes_returned(), stats.bytes_returned()); + } else { + assert!(last_progress.is_none(), "disabled request progress emitted a Progress event"); + } + assert!(saw_end, "Select response ended without an End event"); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn test_select_object_content_reports_input_byte_stats() -> TestResult<()> { + const CSV_BODY: &[u8] = b"name,age\nAlice,30\nBob,25\n"; + const JSON_LINES_BODY: &[u8] = b"{\"name\":\"Alice\"}\n{\"name\":\"Bob\"}\n"; + const JSON_DOCUMENT_BODY: &[u8] = b"[{\"name\":\"Alice\"},{\"name\":\"Bob\"}]"; + + let (_env, client) = create_test_environment().await?; + setup_test_bucket(&client).await?; + assert_input_byte_stats( + &client, + "input-metrics.csv", + CSV_BODY, + "SELECT name FROM S3Object", + InputSerialization::builder() + .csv(CsvInput::builder().file_header_info(FileHeaderInfo::Use).build()) + .build(), + OutputSerialization::builder().csv(CsvOutput::builder().build()).build(), + true, + ) + .await?; + assert_input_byte_stats( + &client, + "input-metrics.jsonl", + JSON_LINES_BODY, + "SELECT name FROM S3Object", + InputSerialization::builder() + .json(JsonInput::builder().set_type(Some(JsonType::Lines)).build()) + .build(), + OutputSerialization::builder().json(JsonOutput::builder().build()).build(), + true, + ) + .await?; + assert_input_byte_stats( + &client, + "input-metrics.json", + JSON_DOCUMENT_BODY, + "SELECT name FROM S3Object", + InputSerialization::builder() + .json(JsonInput::builder().set_type(Some(JsonType::Document)).build()) + .build(), + OutputSerialization::builder().json(JsonOutput::builder().build()).build(), + true, + ) + .await?; + assert_input_byte_stats( + &client, + "input-metrics-without-progress.csv", + CSV_BODY, + "SELECT name FROM S3Object", + InputSerialization::builder() + .csv(CsvInput::builder().file_header_info(FileHeaderInfo::Use).build()) + .build(), + OutputSerialization::builder().csv(CsvOutput::builder().build()).build(), + false, + ) + .await?; + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn test_select_object_content_csv_basic() -> TestResult<()> { let (_env, client) = create_test_environment().await?; diff --git a/crates/s3select-api/src/lib.rs b/crates/s3select-api/src/lib.rs index 1b5cd1539..7662b4a78 100644 --- a/crates/s3select-api/src/lib.rs +++ b/crates/s3select-api/src/lib.rs @@ -23,10 +23,12 @@ use datafusion::{ use std::{error::Error as StdError, fmt::Display}; use thiserror::Error; +mod metrics; pub mod object_store; pub mod query; pub mod server; mod storage_api; +pub use metrics::{SelectInputMetrics, SelectInputMetricsSnapshot}; pub use storage_api::SelectObjectSnapshot; #[cfg(test)] diff --git a/crates/s3select-api/src/metrics.rs b/crates/s3select-api/src/metrics.rs new file mode 100644 index 000000000..d38eefbc8 --- /dev/null +++ b/crates/s3select-api/src/metrics.rs @@ -0,0 +1,88 @@ +// Copyright 2024 RustFS Team +// +// Licensed 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::sync::atomic::{AtomicU64, Ordering}; + +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +pub struct SelectInputMetricsSnapshot { + pub bytes_scanned: u64, + pub bytes_processed: u64, +} + +#[derive(Debug, Default)] +pub struct SelectInputMetrics { + uncompressed_bytes: AtomicU64, +} + +impl SelectInputMetrics { + pub fn snapshot(&self) -> SelectInputMetricsSnapshot { + let uncompressed_bytes = self.uncompressed_bytes.load(Ordering::Relaxed); + SelectInputMetricsSnapshot { + bytes_scanned: uncompressed_bytes, + bytes_processed: uncompressed_bytes, + } + } + + pub(crate) fn record_uncompressed(&self, bytes: usize) { + let increment = u64::try_from(bytes).unwrap_or(u64::MAX); + let _ = self + .uncompressed_bytes + .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |current| Some(current.saturating_add(increment))); + } + + /// Clears planner-only reads before query execution begins. + pub fn reset(&self) { + self.uncompressed_bytes.store(0, Ordering::Relaxed); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn records_uncompressed_input_at_both_boundaries() { + let metrics = SelectInputMetrics::default(); + metrics.record_uncompressed(7); + + assert_eq!( + metrics.snapshot(), + SelectInputMetricsSnapshot { + bytes_scanned: 7, + bytes_processed: 7, + } + ); + } + + #[test] + fn counters_saturate_instead_of_wrapping() { + let metrics = SelectInputMetrics::default(); + metrics.uncompressed_bytes.store(u64::MAX - 1, Ordering::Relaxed); + + metrics.record_uncompressed(2); + + assert_eq!(metrics.snapshot().bytes_scanned, u64::MAX); + assert_eq!(metrics.snapshot().bytes_processed, u64::MAX); + } + + #[test] + fn reset_clears_schema_inference_bytes() { + let metrics = SelectInputMetrics::default(); + metrics.record_uncompressed(9); + + metrics.reset(); + + assert_eq!(metrics.snapshot(), SelectInputMetricsSnapshot::default()); + } +} diff --git a/crates/s3select-api/src/object_store.rs b/crates/s3select-api/src/object_store.rs index 476cbab5f..5134a43a7 100644 --- a/crates/s3select-api/src/object_store.rs +++ b/crates/s3select-api/src/object_store.rs @@ -13,8 +13,9 @@ // limitations under the License. use crate::{ - PrepareSelectObjectSnapshotError, SELECT_DEFAULT_READ_BUFFER_SIZE, SelectError, SelectGetObjectReader, SelectObjectOptions, - SelectObjectSnapshot, SelectObjectSnapshotReadError, SelectStorageError, SelectStore, SnapshotConsistencyError, + PrepareSelectObjectSnapshotError, SELECT_DEFAULT_READ_BUFFER_SIZE, SelectError, SelectGetObjectReader, SelectInputMetrics, + SelectObjectOptions, SelectObjectSnapshot, SelectObjectSnapshotReadError, SelectStorageError, SelectStore, + SnapshotConsistencyError, query::{ parser::RustFsDialect, session::{QueryExecutionGuard, QueryExecutionTracker}, @@ -38,7 +39,7 @@ use datafusion::{ }, }; use futures::pin_mut; -use futures::{Stream, StreamExt, future::ready, stream}; +use futures::{Stream, StreamExt, TryStreamExt, future::ready, stream}; use futures_core::stream::BoxStream; use http::{HeaderMap, HeaderValue, header::HeaderName}; use parking_lot::Mutex; @@ -99,6 +100,7 @@ pub struct EcObjectStore { /// expression. When set, `flatten_json_document_to_ndjson` navigates to /// this key in the root JSON object before flattening. json_sub_path: Option, + input_metrics: Arc, memory_pool: Arc, query_tracker: Option, store: Option>, @@ -172,21 +174,35 @@ pub struct InvalidScanRange; impl EcObjectStore { pub fn new(input: Arc) -> S3Result { - Self::build_lazy(input, Arc::new(UnboundedMemoryPool::default()), None).map_err(map_build_error_to_s3) + Self::build_lazy( + input, + Arc::new(UnboundedMemoryPool::default()), + None, + Arc::new(SelectInputMetrics::default()), + ) + .map_err(map_build_error_to_s3) } pub fn new_with_snapshot(input: Arc, snapshot: Arc) -> S3Result { - Self::build_with_snapshot(input, Arc::new(UnboundedMemoryPool::default()), None, snapshot).map_err(map_build_error_to_s3) + Self::build_with_snapshot( + input, + Arc::new(UnboundedMemoryPool::default()), + None, + Arc::new(SelectInputMetrics::default()), + snapshot, + ) + .map_err(map_build_error_to_s3) } pub(crate) fn new_with_memory_pool( input: Arc, memory_pool: Arc, + input_metrics: Arc, snapshot: Option>, ) -> std::result::Result { match snapshot { - Some(snapshot) => Self::build_with_snapshot(input, memory_pool, None, snapshot), - None => Self::build_lazy(input, memory_pool, None), + Some(snapshot) => Self::build_with_snapshot(input, memory_pool, None, input_metrics, snapshot), + None => Self::build_lazy(input, memory_pool, None, input_metrics), } } @@ -194,11 +210,12 @@ impl EcObjectStore { input: Arc, memory_pool: Arc, query_tracker: QueryExecutionTracker, + input_metrics: Arc, snapshot: Option>, ) -> std::result::Result { match snapshot { - Some(snapshot) => Self::build_with_snapshot(input, memory_pool, Some(query_tracker), snapshot), - None => Self::build_lazy(input, memory_pool, Some(query_tracker)), + Some(snapshot) => Self::build_with_snapshot(input, memory_pool, Some(query_tracker), input_metrics, snapshot), + None => Self::build_lazy(input, memory_pool, Some(query_tracker), input_metrics), } } @@ -206,27 +223,30 @@ impl EcObjectStore { input: Arc, memory_pool: Arc, query_tracker: Option, + input_metrics: Arc, ) -> std::result::Result { let store = resolve_select_object_store_handle().ok_or(EcObjectStoreBuildError::StoreUnavailable)?; - Ok(Self::build(input, memory_pool, query_tracker, Some(store), None)) + Ok(Self::build(input, memory_pool, query_tracker, input_metrics, Some(store), None)) } fn build_with_snapshot( input: Arc, memory_pool: Arc, query_tracker: Option, + input_metrics: Arc, snapshot: Arc, ) -> std::result::Result { if !snapshot.is_for(&input.bucket, &input.key) { return Err(EcObjectStoreBuildError::Snapshot(SnapshotConsistencyError::ObjectChanged)); } - Ok(Self::build(input, memory_pool, query_tracker, None, Some(snapshot))) + Ok(Self::build(input, memory_pool, query_tracker, input_metrics, None, Some(snapshot))) } fn build( input: Arc, memory_pool: Arc, query_tracker: Option, + input_metrics: Arc, store: Option>, snapshot: Option>, ) -> Self { @@ -269,6 +289,7 @@ impl EcObjectStore { delimiter, is_json_document, json_sub_path, + input_metrics, memory_pool, query_tracker, store, @@ -705,14 +726,18 @@ impl ObjectStore for EcObjectStore { self.object_reader(range).await? }; + let meter_input = self.input.request.input_serialization.parquet.is_none(); let payload = if options.range.is_some() { let size = usize::try_from(result_range.end - result_range.start).map_err(|err| o_Error::Generic { store: "EcObjectStore", source: Box::new(err), })?; - GetResultPayload::Stream( - bytes_stream(ReaderStream::with_capacity(reader.stream, SELECT_DEFAULT_READ_BUFFER_SIZE), size).boxed(), - ) + let stream = bytes_stream(ReaderStream::with_capacity(reader.stream, SELECT_DEFAULT_READ_BUFFER_SIZE), size); + if meter_input { + GetResultPayload::Stream(meter_uncompressed_input_stream(stream, Arc::clone(&self.input_metrics)).boxed()) + } else { + GetResultPayload::Stream(stream.boxed()) + } } else if self.is_json_document { // JSON DOCUMENT mode: gate on object size before doing any I/O. // @@ -731,6 +756,7 @@ impl ObjectStore for EcObjectStore { reader.stream, original_size, self.json_sub_path.clone(), + Arc::clone(&self.input_metrics), Arc::clone(&self.memory_pool), self.query_tracker.clone(), ); @@ -739,12 +765,17 @@ impl ObjectStore for EcObjectStore { let delimiter = self.record_delimiter(); let include_header = self.csv_has_header(); let header = if include_header && read_start > 0 { - Some(self.read_header_record(original_size, &delimiter).await?) + let header = self.read_header_record(original_size, &delimiter).await?; + self.input_metrics.record_uncompressed(header.len()); + Some(header) } else { None }; let stream = scan_range_stream( - ReaderStream::with_capacity(reader.stream, SELECT_DEFAULT_READ_BUFFER_SIZE), + meter_uncompressed_input_stream( + ReaderStream::with_capacity(reader.stream, SELECT_DEFAULT_READ_BUFFER_SIZE), + Arc::clone(&self.input_metrics), + ), delimiter, scan_range, include_header && header.is_none(), @@ -757,22 +788,23 @@ impl ObjectStore for EcObjectStore { } else { stream }; - GetResultPayload::Stream(convert_csv_delimiter_stream( - stream, - record_delimiter, - self.need_convert.then(|| self.delimiter.clone()), - )) + let stream = + convert_csv_delimiter_stream(stream, record_delimiter, self.need_convert.then(|| self.delimiter.clone())); + GetResultPayload::Stream(stream) } else { let stream_size = usize::try_from(original_size).map_err(|err| o_Error::Generic { store: "EcObjectStore", source: Box::new(err), })?; let stream = bytes_stream(ReaderStream::with_capacity(reader.stream, SELECT_DEFAULT_READ_BUFFER_SIZE), stream_size); - GetResultPayload::Stream(convert_csv_delimiter_stream( - stream, - record_delimiter, - self.need_convert.then(|| self.delimiter.clone()), - )) + if meter_input { + let stream = meter_uncompressed_input_stream(stream, Arc::clone(&self.input_metrics)); + let stream = + convert_csv_delimiter_stream(stream, record_delimiter, self.need_convert.then(|| self.delimiter.clone())); + GetResultPayload::Stream(stream) + } else { + GetResultPayload::Stream(stream.boxed()) + } }; Ok(GetResult { @@ -1110,6 +1142,7 @@ fn json_document_ndjson_stream( stream: Box, original_size: u64, json_sub_path: Option, + input_metrics: Arc, memory_pool: Arc, query_tracker: Option, ) -> futures_core::stream::BoxStream<'static, Result> { @@ -1117,6 +1150,7 @@ fn json_document_ndjson_stream( stream, original_size, json_sub_path, + input_metrics, memory_pool, query_tracker, |all_bytes, json_sub_path| parse_json_document_to_lines(&all_bytes, json_sub_path.as_deref()), @@ -1127,6 +1161,7 @@ fn json_document_ndjson_stream_with_parser

( stream: Box, original_size: u64, json_sub_path: Option, + input_metrics: Arc, memory_pool: Arc, query_tracker: Option, parser: P, @@ -1158,17 +1193,15 @@ where source: Box::new(err), })?; - pin_mut!(stream); // ── 1. Read phase (lazy: only runs when the stream is polled) ──── + pin_mut!(stream); let mut all_bytes = Vec::with_capacity(buffer_capacity); - stream - .take(original_size) - .read_to_end(&mut all_bytes) - .await - .map_err(|e| o_Error::Generic { - store: "EcObjectStore", - source: Box::new(e), - })?; + let read_result = stream.take(original_size).read_to_end(&mut all_bytes).await; + input_metrics.record_uncompressed(all_bytes.len()); + read_result.map_err(|e| o_Error::Generic { + store: "EcObjectStore", + source: Box::new(e), + })?; if all_bytes.len() != buffer_capacity { return Err(incomplete_object_stream_error(buffer_capacity - all_bytes.len())); } @@ -1322,6 +1355,17 @@ fn flatten_json_document_to_ndjson(bytes: &[u8], json_sub_path: Option<&str>) -> Ok(Bytes::from(output)) } +fn meter_uncompressed_input_stream( + stream: S, + input_metrics: Arc, +) -> impl Stream> + Send + 'static +where + S: Stream> + Send + 'static, + E: Send + 'static, +{ + stream.inspect_ok(move |bytes| input_metrics.record_uncompressed(bytes.len())) +} + pub fn bytes_stream(stream: S, content_length: usize) -> impl Stream> + Send + 'static where S: Stream> + Send + 'static, @@ -1382,13 +1426,13 @@ mod test { SELECT_DEFAULT_READ_BUFFER_SIZE, SelectObjectOptions, SelectObjectSnapshot, SelectScanRange, SnapshotConsistencyError, bytes_stream, convert_csv_delimiter_stream, convert_field_delimiter_stream, convert_record_delimiter_stream, extract_json_sub_path_from_expression, find_delimiter, flatten_json_document_to_ndjson, http_range_spec_from_get_range, - json_document_ndjson_stream, json_document_ndjson_stream_with_parser, map_storage_error, scan_range_from_bounds, - scan_range_stream, select_read_headers, snapshot_last_modified, validate_json_document_size, + json_document_ndjson_stream, json_document_ndjson_stream_with_parser, map_storage_error, meter_uncompressed_input_stream, + scan_range_from_bounds, scan_range_stream, select_read_headers, snapshot_last_modified, validate_json_document_size, }; use crate::query::session::{QueryExecutionGuard, QueryExecutionOwner, QueryExecutionTracker}; use crate::storage_api::SelectPutObjReader; use crate::storage_api::object_store::ObjectIO as _; - use crate::{QueryError, SelectError, SelectStorageError}; + use crate::{QueryError, SelectError, SelectInputMetrics, SelectStorageError}; use bytes::Bytes; use datafusion::{ common::DataFusionError, @@ -1403,8 +1447,8 @@ mod test { use rustfs_test_utils::PutObjectCommitBarrier; use s3s::S3ErrorCode; use s3s::dto::{ - CSVInput, CSVOutput, ExpressionType, FileHeaderInfo, InputSerialization, OutputSerialization, ScanRange, - SelectObjectContentInput, SelectObjectContentRequest, + CSVInput, CSVOutput, ExpressionType, FileHeaderInfo, InputSerialization, JSONInput, JSONOutput, JSONType, + OutputSerialization, ScanRange, SelectObjectContentInput, SelectObjectContentRequest, }; use s3s::header::{ X_AMZ_SERVER_SIDE_ENCRYPTION_CUSTOMER_ALGORITHM, X_AMZ_SERVER_SIDE_ENCRYPTION_CUSTOMER_KEY, @@ -1442,6 +1486,21 @@ mod test { }) } + fn json_input(bucket: &str, object: &str, json_type: &'static str) -> Arc { + let mut input = (*csv_input(bucket, object)).clone(); + input.request.input_serialization = InputSerialization { + json: Some(JSONInput { + type_: Some(JSONType::from_static(json_type)), + }), + ..Default::default() + }; + input.request.output_serialization = OutputSerialization { + json: Some(JSONOutput::default()), + ..Default::default() + }; + Arc::new(input) + } + #[test] fn lazy_snapshot_headers_preserve_ssec_context() { let mut input = (*csv_input("bucket", "object.csv")).clone(); @@ -2235,6 +2294,24 @@ mod test { assert_eq!(output, b"a,1\nb,2\n"); } + #[tokio::test] + async fn delimiter_conversion_keeps_uncompressed_metrics_equal() { + let input = Bytes::from_static(b"a&&1\nb&&2\n"); + let input_metrics = Arc::new(SelectInputMetrics::default()); + let stream = stream::iter([Ok::<_, object_store::Error>(input.clone())]); + let stream = meter_uncompressed_input_stream(stream, Arc::clone(&input_metrics)); + let output = convert_field_delimiter_stream(stream, "&&".to_string()) + .try_collect::>() + .await + .expect("delimiter conversion should succeed") + .concat(); + + assert_eq!(output, b"a,1\nb,2\n"); + let input_len = u64::try_from(input.len()).expect("fixture length should fit in u64"); + assert_eq!(input_metrics.snapshot().bytes_scanned, input_len); + assert_eq!(input_metrics.snapshot().bytes_processed, input_len); + } + #[tokio::test] async fn test_field_delimiter_stream_converts_delimiter_split_across_chunks() { let chunks = stream::iter(vec![ @@ -2310,6 +2387,7 @@ mod test { delimiter: String::new(), is_json_document: false, json_sub_path: None, + input_metrics: Arc::new(SelectInputMetrics::default()), memory_pool: Arc::new(GreedyMemoryPool::new(1024)), query_tracker: None, store: None, @@ -2496,6 +2574,7 @@ mod test { delimiter: String::new(), is_json_document: false, json_sub_path: None, + input_metrics: Arc::new(SelectInputMetrics::default()), memory_pool: Arc::new(GreedyMemoryPool::new(32 * 1024 * 1024)), query_tracker: None, store: None, @@ -2586,12 +2665,14 @@ mod test { }, }); let snapshot = prepare_test_snapshot(bucket, object).await; + let input_metrics = Arc::new(SelectInputMetrics::default()); let store = super::EcObjectStore { input, need_convert: true, delimiter: "\r\n".to_string(), is_json_document: false, json_sub_path: None, + input_metrics: Arc::clone(&input_metrics), memory_pool: Arc::new(GreedyMemoryPool::new(1024)), query_tracker: None, store: None, @@ -2609,6 +2690,9 @@ mod test { let chunks: Vec = stream.try_collect().await.expect("collect converted object stream"); assert_eq!(chunks.concat(), b"a,1\r\n"); + let input_len = u64::try_from(input_bytes.len()).expect("fixture length should fit in u64"); + assert_eq!(input_metrics.snapshot().bytes_scanned, input_len); + assert_eq!(input_metrics.snapshot().bytes_processed, input_len); let requested_range = 3..10; let ranges = store @@ -2641,6 +2725,33 @@ mod test { assert_eq!(poll_count.load(Ordering::SeqCst), 1); } + #[tokio::test] + async fn metered_stream_counts_only_polled_chunks() { + let poll_count = Arc::new(AtomicUsize::new(0)); + let source_poll_count = Arc::clone(&poll_count); + let source = stream::unfold(0, move |index| { + let source_poll_count = Arc::clone(&source_poll_count); + async move { + source_poll_count.fetch_add(1, Ordering::SeqCst); + let bytes = match index { + 0 => Bytes::from_static(b"abcd"), + 1 => Bytes::from_static(b"efgh"), + _ => return None, + }; + Some((Ok::<_, std::io::Error>(bytes), index + 1)) + } + }); + let input_metrics = Arc::new(SelectInputMetrics::default()); + let mut metered = Box::pin(meter_uncompressed_input_stream(source, Arc::clone(&input_metrics))); + + assert_eq!(metered.next().await.expect("first chunk").expect("valid chunk"), b"abcd"[..]); + drop(metered); + + assert_eq!(input_metrics.snapshot().bytes_scanned, 4); + assert_eq!(input_metrics.snapshot().bytes_processed, 4); + assert_eq!(poll_count.load(Ordering::SeqCst), 1); + } + #[tokio::test] async fn test_bytes_stream_rejects_early_eof() { let source = stream::iter(vec![Ok::<_, std::io::Error>(Bytes::from_static(b"ab"))]); @@ -2663,6 +2774,219 @@ mod test { assert!(output.next().await.is_none()); } + #[tokio::test] + async fn full_and_range_object_streams_record_input_metrics() { + const BUCKET: &str = "s3select-input-metrics"; + const OBJECT: &str = "input.csv"; + const DATA: &[u8] = b"id,name\n1,a\n"; + + let env = crate::storage_api::select_test_ecstore_env().await; + env.make_bucket(BUCKET, false).await; + let mut reader = SelectPutObjReader::from_vec(DATA.to_vec()); + env.ecstore + .put_object(BUCKET, OBJECT, &mut reader, &Default::default()) + .await + .expect("put input metrics fixture"); + let snapshot = prepare_test_snapshot(BUCKET, OBJECT).await; + let input_metrics = Arc::new(SelectInputMetrics::default()); + let store = EcObjectStore::build_with_snapshot( + csv_input(BUCKET, OBJECT), + Arc::new(GreedyMemoryPool::new(1024 * 1024)), + None, + Arc::clone(&input_metrics), + snapshot, + ) + .expect("build metrics-aware object store"); + + let result = store + .get_opts(&Path::from(OBJECT), GetOptions::default()) + .await + .expect("open full object stream"); + let GetResultPayload::Stream(stream) = result.payload else { + panic!("expected streaming object payload"); + }; + let body = stream + .try_collect::>() + .await + .expect("read full object stream") + .concat(); + assert_eq!(body, DATA); + let data_len = u64::try_from(DATA.len()).expect("fixture length should fit in u64"); + assert_eq!(input_metrics.snapshot().bytes_scanned, data_len); + assert_eq!(input_metrics.snapshot().bytes_processed, data_len); + + input_metrics.reset(); + + let result = store + .get_opts( + &Path::from(OBJECT), + GetOptions { + range: Some(GetRange::Bounded(0..2)), + ..Default::default() + }, + ) + .await + .expect("open schema-style range stream"); + let GetResultPayload::Stream(stream) = result.payload else { + panic!("expected range stream payload"); + }; + let range = stream + .try_collect::>() + .await + .expect("read schema-style range") + .concat(); + assert_eq!(range, b"id"[..]); + assert_eq!(input_metrics.snapshot().bytes_scanned, 2); + assert_eq!(input_metrics.snapshot().bytes_processed, 2); + } + + #[tokio::test] + async fn dropping_real_object_stream_counts_only_consumed_bytes() { + const BUCKET: &str = "s3select-partial-input-metrics"; + const OBJECT: &str = "large.csv"; + + let data = vec![b'x'; SELECT_DEFAULT_READ_BUFFER_SIZE * 3]; + let env = crate::storage_api::select_test_ecstore_env().await; + env.make_bucket(BUCKET, false).await; + let mut reader = SelectPutObjReader::from_vec(data.clone()); + env.ecstore + .put_object(BUCKET, OBJECT, &mut reader, &Default::default()) + .await + .expect("put partial input metrics fixture"); + let snapshot = prepare_test_snapshot(BUCKET, OBJECT).await; + let input_metrics = Arc::new(SelectInputMetrics::default()); + let store = EcObjectStore::build_with_snapshot( + csv_input(BUCKET, OBJECT), + Arc::new(GreedyMemoryPool::new(1024 * 1024)), + None, + Arc::clone(&input_metrics), + snapshot, + ) + .expect("build metrics-aware object store"); + + let result = store + .get_opts(&Path::from(OBJECT), GetOptions::default()) + .await + .expect("open partial object stream"); + let GetResultPayload::Stream(mut stream) = result.payload else { + panic!("expected streaming object payload"); + }; + let first = stream + .next() + .await + .expect("first object chunk") + .expect("first object chunk should be valid"); + drop(stream); + + assert!(first.len() < data.len(), "fixture must span multiple reader chunks"); + let consumed = u64::try_from(first.len()).expect("chunk length should fit in u64"); + assert_eq!(input_metrics.snapshot().bytes_scanned, consumed); + assert_eq!(input_metrics.snapshot().bytes_processed, consumed); + } + + #[tokio::test] + async fn json_object_streams_record_input_metrics() { + const BUCKET: &str = "s3select-json-input-metrics"; + const LINES_OBJECT: &str = "input.jsonl"; + const LINES_DATA: &[u8] = b"{\"id\":1}\n{\"id\":2}\n"; + const DOCUMENT_OBJECT: &str = "input.json"; + const DOCUMENT_DATA: &[u8] = b"[{\"id\":1},{\"id\":2}]"; + + let env = crate::storage_api::select_test_ecstore_env().await; + env.make_bucket(BUCKET, false).await; + for (object, data, json_type) in [ + (LINES_OBJECT, LINES_DATA, JSONType::LINES), + (DOCUMENT_OBJECT, DOCUMENT_DATA, JSONType::DOCUMENT), + ] { + let mut reader = SelectPutObjReader::from_vec(data.to_vec()); + env.ecstore + .put_object(BUCKET, object, &mut reader, &Default::default()) + .await + .expect("put JSON input metrics fixture"); + let snapshot = prepare_test_snapshot(BUCKET, object).await; + let input_metrics = Arc::new(SelectInputMetrics::default()); + let store = EcObjectStore::build_with_snapshot( + json_input(BUCKET, object, json_type), + Arc::new(GreedyMemoryPool::new(1024 * 1024)), + None, + Arc::clone(&input_metrics), + snapshot, + ) + .expect("build metrics-aware JSON object store"); + + let result = store + .get_opts(&Path::from(object), GetOptions::default()) + .await + .expect("open JSON object stream"); + let GetResultPayload::Stream(stream) = result.payload else { + panic!("expected streaming JSON payload"); + }; + stream.try_collect::>().await.expect("read JSON object stream"); + + let input_len = u64::try_from(data.len()).expect("fixture length should fit in u64"); + assert_eq!(input_metrics.snapshot().bytes_scanned, input_len, "JSON type {json_type}"); + assert_eq!(input_metrics.snapshot().bytes_processed, input_len, "JSON type {json_type}"); + } + } + + #[tokio::test] + async fn scan_range_metrics_include_header_and_raw_range_once() { + const BUCKET: &str = "s3select-scan-range-input-metrics"; + const OBJECT: &str = "input.csv"; + const DATA: &[u8] = b"h1,h2\nr1,a\nr2,b\n"; + const HEADER_LEN: usize = b"h1,h2\n".len(); + const RECORD_START: usize = b"h1,h2\nr1,a\n".len(); + const READ_START: usize = RECORD_START - 1; + + let env = crate::storage_api::select_test_ecstore_env().await; + env.make_bucket(BUCKET, false).await; + let mut reader = SelectPutObjReader::from_vec(DATA.to_vec()); + env.ecstore + .put_object(BUCKET, OBJECT, &mut reader, &Default::default()) + .await + .expect("put ScanRange input metrics fixture"); + let snapshot = prepare_test_snapshot(BUCKET, OBJECT).await; + let mut input = (*csv_input(BUCKET, OBJECT)).clone(); + input + .request + .input_serialization + .csv + .as_mut() + .expect("CSV input") + .file_header_info = Some(FileHeaderInfo::from_static(FileHeaderInfo::USE)); + input.request.scan_range = Some(ScanRange { + start: Some(i64::try_from(RECORD_START).expect("fixture offset should fit in i64")), + end: Some(i64::try_from(RECORD_START).expect("fixture offset should fit in i64")), + }); + let input_metrics = Arc::new(SelectInputMetrics::default()); + let store = EcObjectStore::build_with_snapshot( + Arc::new(input), + Arc::new(GreedyMemoryPool::new(1024 * 1024)), + None, + Arc::clone(&input_metrics), + snapshot, + ) + .expect("build ScanRange metrics-aware object store"); + + let result = store + .get_opts(&Path::from(OBJECT), GetOptions::default()) + .await + .expect("open ScanRange object stream"); + let GetResultPayload::Stream(stream) = result.payload else { + panic!("expected streaming ScanRange payload"); + }; + let body = stream + .try_collect::>() + .await + .expect("read ScanRange object stream") + .concat(); + + assert_eq!(body, b"h1,h2\nr2,b\n"); + let expected_input = u64::try_from(HEADER_LEN + DATA.len() - READ_START).expect("fixture length should fit in u64"); + assert_eq!(input_metrics.snapshot().bytes_scanned, expected_input); + assert_eq!(input_metrics.snapshot().bytes_processed, expected_input); + } + #[tokio::test] async fn test_json_document_stream_respects_query_memory_pool() { let input = b"{}".to_vec(); @@ -2672,6 +2996,7 @@ mod test { Box::new(std::io::Cursor::new(input.clone())), input.len() as u64, None, + Arc::new(SelectInputMetrics::default()), memory_pool, None, ); @@ -2695,16 +3020,21 @@ mod test { let input = b"[1,2]".to_vec(); let required = input.len() * JSON_DOCUMENT_MEMORY_RESERVATION_MULTIPLIER; let memory_pool = Arc::new(GreedyMemoryPool::new(required)); + let input_metrics = Arc::new(SelectInputMetrics::default()); let output: Vec = json_document_ndjson_stream( Box::new(std::io::Cursor::new(input.clone())), input.len() as u64, None, + Arc::clone(&input_metrics), memory_pool.clone(), None, ) .try_collect() .await .expect("JSON conversion should fit the pool"); + let input_len = u64::try_from(input.len()).expect("fixture length should fit in u64"); + assert_eq!(input_metrics.snapshot().bytes_scanned, input_len); + assert_eq!(input_metrics.snapshot().bytes_processed, input_len); assert_eq!(output, vec![Bytes::from_static(b"1\n"), Bytes::from_static(b"2\n")]); assert_eq!(memory_pool.reserved(), 0); @@ -2713,8 +3043,17 @@ mod test { #[tokio::test] async fn test_json_document_stream_rejects_early_eof() { let input = b"{}".to_vec(); + let input_len = u64::try_from(input.len()).expect("fixture length should fit in u64"); + let input_metrics = Arc::new(SelectInputMetrics::default()); let memory_pool: Arc = Arc::new(GreedyMemoryPool::new(4 * JSON_DOCUMENT_MEMORY_RESERVATION_MULTIPLIER)); - let mut output = json_document_ndjson_stream(Box::new(std::io::Cursor::new(input)), 4, None, memory_pool, None); + let mut output = json_document_ndjson_stream( + Box::new(std::io::Cursor::new(input)), + 4, + None, + Arc::clone(&input_metrics), + memory_pool, + None, + ); let err = output .next() @@ -2727,6 +3066,8 @@ mod test { let source = source.downcast_ref::().expect("I/O error source"); assert_eq!(source.kind(), std::io::ErrorKind::UnexpectedEof); assert!(source.to_string().contains("2 bytes remaining")); + assert_eq!(input_metrics.snapshot().bytes_scanned, input_len); + assert_eq!(input_metrics.snapshot().bytes_processed, input_len); assert!(output.next().await.is_none()); } @@ -2739,6 +3080,7 @@ mod test { Box::new(std::io::Cursor::new(input.clone())), input.len() as u64, None, + Arc::new(SelectInputMetrics::default()), memory_pool, None, ); @@ -2844,6 +3186,7 @@ mod test { Box::new(std::io::Cursor::new(input.clone())), input.len() as u64, None, + Arc::new(SelectInputMetrics::default()), memory_pool, Some(query_tracker), ); @@ -2906,6 +3249,7 @@ mod test { Box::new(std::io::Cursor::new(input.clone())), input.len() as u64, None, + Arc::new(SelectInputMetrics::default()), memory_pool, Some(query_tracker.clone()), move |_, _| { @@ -2969,6 +3313,7 @@ mod test { Box::new(std::io::Cursor::new(input.clone())), input.len() as u64, None, + Arc::new(SelectInputMetrics::default()), memory_pool, Some(query_tracker), move |_, _| { diff --git a/crates/s3select-api/src/query/dispatcher.rs b/crates/s3select-api/src/query/dispatcher.rs index 8a3e8a908..6646de880 100644 --- a/crates/s3select-api/src/query/dispatcher.rs +++ b/crates/s3select-api/src/query/dispatcher.rs @@ -25,6 +25,8 @@ use super::{ session::QueryAdmission, }; +pub type DispatchedQuery = (Query, Output); + #[async_trait] pub trait QueryDispatcher: Send + Sync { // fn create_query_id(&self) -> QueryId; @@ -41,6 +43,18 @@ pub trait QueryDispatcher: Send + Sync { self.execute_query(query).await } + async fn dispatch_query(&self, query: &Query) -> QueryResult { + let execution_query = query.for_execution(); + let output = self.execute_query(&execution_query).await?; + Ok((execution_query, output)) + } + + async fn dispatch_query_admitted(&self, query: &Query, admission: QueryAdmission) -> QueryResult { + let execution_query = query.for_execution(); + let output = self.execute_query_admitted(&execution_query, admission).await?; + Ok((execution_query, output)) + } + async fn build_logical_plan(&self, query_state_machine: Arc) -> QueryResult>; async fn execute_logical_plan(&self, logical_plan: Plan, query_state_machine: Arc) -> QueryResult; @@ -53,3 +67,155 @@ pub trait QueryDispatcher: Send + Sync { // fn cancel_query(&self, id: &QueryId); } + +#[cfg(test)] +mod tests { + use super::*; + use crate::query::test_query; + use parking_lot::Mutex; + + #[derive(Default)] + struct DefaultDispatchDispatcher { + executed_metrics: Mutex>>, + } + + #[async_trait] + impl QueryDispatcher for DefaultDispatchDispatcher { + async fn execute_query(&self, query: &Query) -> QueryResult { + self.executed_metrics.lock().push(Arc::clone(query.input_metrics())); + Ok(Output::Nil(())) + } + + async fn build_logical_plan(&self, _query_state_machine: Arc) -> QueryResult> { + unreachable!("default dispatch test does not plan queries") + } + + async fn execute_logical_plan( + &self, + _logical_plan: Plan, + _query_state_machine: Arc, + ) -> QueryResult { + unreachable!("default dispatch test does not execute plans") + } + + async fn build_query_state_machine(&self, _query: Query) -> QueryResult> { + unreachable!("default dispatch test does not build state machines") + } + } + + #[derive(Default)] + struct DistinctAdmittedDispatcher { + plain_metrics: Mutex>>, + admitted_metrics: Mutex>>, + fail_plain: bool, + fail_admitted: bool, + } + + #[async_trait] + impl QueryDispatcher for DistinctAdmittedDispatcher { + async fn execute_query(&self, query: &Query) -> QueryResult { + self.plain_metrics.lock().push(Arc::clone(query.input_metrics())); + if self.fail_plain { + Err(crate::QueryError::Cancel) + } else { + Ok(Output::Nil(())) + } + } + + async fn execute_query_admitted(&self, query: &Query, _admission: QueryAdmission) -> QueryResult { + self.admitted_metrics.lock().push(Arc::clone(query.input_metrics())); + if self.fail_admitted { + Err(crate::QueryError::Cancel) + } else { + Ok(Output::Nil(())) + } + } + + async fn build_logical_plan(&self, _query_state_machine: Arc) -> QueryResult> { + unreachable!("dispatch routing test does not plan queries") + } + + async fn execute_logical_plan( + &self, + _logical_plan: Plan, + _query_state_machine: Arc, + ) -> QueryResult { + unreachable!("dispatch routing test does not execute plans") + } + + async fn build_query_state_machine(&self, _query: Query) -> QueryResult> { + unreachable!("dispatch routing test does not build state machines") + } + } + + #[tokio::test] + async fn plain_dispatch_propagates_override_errors() { + let dispatcher = DistinctAdmittedDispatcher { + fail_plain: true, + ..Default::default() + }; + + let error = match dispatcher.dispatch_query(&test_query()).await { + Err(error) => error, + Ok(_) => panic!("plain override error should propagate"), + }; + assert!(matches!(error, crate::QueryError::Cancel)); + assert_eq!(dispatcher.plain_metrics.lock().len(), 1); + assert!(dispatcher.admitted_metrics.lock().is_empty()); + } + + #[tokio::test] + async fn default_dispatch_methods_use_distinct_execution_metrics() { + let dispatcher = DefaultDispatchDispatcher::default(); + let query = test_query(); + + let (first, _) = dispatcher + .dispatch_query(&query) + .await + .expect("first dispatch should execute"); + let (second, _) = dispatcher + .dispatch_query(&query) + .await + .expect("second dispatch should execute"); + let (admitted, _) = dispatcher + .dispatch_query_admitted(&query, QueryAdmission::unmanaged()) + .await + .expect("admitted dispatch should execute"); + let executed_metrics = dispatcher.executed_metrics.lock(); + + assert!(!Arc::ptr_eq(first.input_metrics(), second.input_metrics())); + assert!(!Arc::ptr_eq(first.input_metrics(), admitted.input_metrics())); + assert!(Arc::ptr_eq(first.input_metrics(), &executed_metrics[0])); + assert!(Arc::ptr_eq(second.input_metrics(), &executed_metrics[1])); + assert!(Arc::ptr_eq(admitted.input_metrics(), &executed_metrics[2])); + } + + #[tokio::test] + async fn admitted_dispatch_uses_the_admitted_override_and_propagates_errors() { + let dispatcher = DistinctAdmittedDispatcher::default(); + let query = test_query(); + + let (dispatched, _) = dispatcher + .dispatch_query_admitted(&query, QueryAdmission::unmanaged()) + .await + .expect("admitted dispatch should execute through its override"); + assert!(dispatcher.plain_metrics.lock().is_empty()); + { + let admitted_metrics = dispatcher.admitted_metrics.lock(); + assert_eq!(admitted_metrics.len(), 1); + assert!(Arc::ptr_eq(dispatched.input_metrics(), &admitted_metrics[0])); + } + + let failing = DistinctAdmittedDispatcher { + fail_admitted: true, + ..Default::default() + }; + let error = match failing.dispatch_query_admitted(&query, QueryAdmission::unmanaged()).await { + Err(error) => error, + Ok(_) => panic!("admitted override error should propagate"), + }; + assert!(matches!(error, crate::QueryError::Cancel)); + assert!(failing.plain_metrics.lock().is_empty()); + assert_eq!(failing.admitted_metrics.lock().len(), 1); + } +} diff --git a/crates/s3select-api/src/query/mod.rs b/crates/s3select-api/src/query/mod.rs index 11fd92ebf..c09640752 100644 --- a/crates/s3select-api/src/query/mod.rs +++ b/crates/s3select-api/src/query/mod.rs @@ -15,7 +15,7 @@ use s3s::dto::SelectObjectContentInput; use std::sync::Arc; -use crate::SelectObjectSnapshot; +use crate::{SelectInputMetrics, SelectObjectSnapshot}; pub mod analyzer; pub mod ast; @@ -40,6 +40,7 @@ pub struct Query { context: Context, content: String, snapshot: Option>, + input_metrics: Arc, } impl Query { @@ -49,6 +50,7 @@ impl Query { context, content, snapshot: None, + input_metrics: Arc::new(SelectInputMetrics::default()), } } @@ -58,6 +60,7 @@ impl Query { context, content, snapshot: Some(snapshot), + input_metrics: Arc::new(SelectInputMetrics::default()), } } @@ -72,4 +75,46 @@ impl Query { pub fn snapshot(&self) -> Option<&Arc> { self.snapshot.as_ref() } + + pub fn input_metrics(&self) -> &Arc { + &self.input_metrics + } + + pub fn for_execution(&self) -> Self { + Self { + context: self.context.clone(), + content: self.content.clone(), + snapshot: self.snapshot.clone(), + input_metrics: Arc::new(SelectInputMetrics::default()), + } + } +} + +#[cfg(test)] +fn test_query() -> Query { + use s3s::dto::{CSVInput, CSVOutput, ExpressionType, InputSerialization, OutputSerialization, SelectObjectContentRequest}; + + let input = SelectObjectContentInput { + bucket: "bucket".to_string(), + expected_bucket_owner: None, + key: "input.csv".to_string(), + sse_customer_algorithm: None, + sse_customer_key: None, + sse_customer_key_md5: None, + request: SelectObjectContentRequest { + expression: "SELECT * FROM S3Object".to_string(), + expression_type: ExpressionType::from_static(ExpressionType::SQL), + input_serialization: InputSerialization { + csv: Some(CSVInput::default()), + ..Default::default() + }, + output_serialization: OutputSerialization { + csv: Some(CSVOutput::default()), + ..Default::default() + }, + request_progress: None, + scan_range: None, + }, + }; + Query::new(Context { input: Arc::new(input) }, "SELECT * FROM S3Object".to_string()) } diff --git a/crates/s3select-api/src/query/session.rs b/crates/s3select-api/src/query/session.rs index 37c9e482e..dc5ca8d25 100644 --- a/crates/s3select-api/src/query/session.rs +++ b/crates/s3select-api/src/query/session.rs @@ -12,9 +12,9 @@ // See the License for the specific language governing permissions and // limitations under the License. -use crate::SelectObjectSnapshot; -use crate::query::Context; +use crate::query::{Context, Query}; use crate::{QueryError, QueryResult, object_store::EcObjectStore}; +use crate::{SelectInputMetrics, SelectObjectSnapshot}; use datafusion::{ arrow::{ array::{Int32Array, StringArray}, @@ -314,7 +314,7 @@ impl SessionCtxFactory { } pub async fn create_session_ctx(&self, context: &Context) -> QueryResult { - self.create_session_ctx_inner(context, None, None, DEFAULT_S3SELECT_MEMORY_LIMIT_BYTES) + self.create_session_ctx_inner(context, None, None, None, DEFAULT_S3SELECT_MEMORY_LIMIT_BYTES) .await } @@ -324,7 +324,7 @@ impl SessionCtxFactory { query_tracker: QueryExecutionTracker, memory_limit_bytes: usize, ) -> QueryResult { - self.create_session_ctx_inner(context, None, Some(query_tracker), memory_limit_bytes) + self.create_session_ctx_inner(context, None, Some(query_tracker), None, memory_limit_bytes) .await } @@ -335,19 +335,36 @@ impl SessionCtxFactory { query_tracker: QueryExecutionTracker, memory_limit_bytes: usize, ) -> QueryResult { - self.create_session_ctx_inner(context, Some(snapshot), Some(query_tracker), memory_limit_bytes) + self.create_session_ctx_inner(context, Some(snapshot), Some(query_tracker), None, memory_limit_bytes) .await } + pub async fn create_session_ctx_for_query_with_tracker_and_memory_limit( + &self, + query: &Query, + query_tracker: QueryExecutionTracker, + memory_limit_bytes: usize, + ) -> QueryResult { + self.create_session_ctx_inner( + query.context(), + query.snapshot().cloned(), + Some(query_tracker), + Some(Arc::clone(query.input_metrics())), + memory_limit_bytes, + ) + .await + } + async fn create_session_ctx_inner( &self, context: &Context, snapshot: Option>, query_tracker: Option, + input_metrics: Option>, memory_limit_bytes: usize, ) -> QueryResult { let df_session_ctx = self - .build_df_session_context(context, snapshot, query_tracker.clone(), memory_limit_bytes) + .build_df_session_context(context, snapshot, query_tracker.clone(), input_metrics, memory_limit_bytes) .await?; Ok(SessionCtx { @@ -362,6 +379,7 @@ impl SessionCtxFactory { context: &Context, snapshot: Option>, query_tracker: Option, + input_metrics: Option>, memory_limit_bytes: usize, ) -> QueryResult { let path = format!("s3://{}", context.input.bucket); @@ -383,7 +401,12 @@ impl SessionCtxFactory { .is_some_and(|delimiter| delimiter.len() == 2 && delimiter.as_bytes() != b"\r\n"); let scan_range_requires_single_file_scan = context.input.request.scan_range.is_some() && context.input.request.input_serialization.parquet.is_none(); - let config = if custom_two_byte_record_delimiter || scan_range_requires_single_file_scan { + let metered_input_requires_single_file_scan = + input_metrics.is_some() && context.input.request.input_serialization.parquet.is_none(); + let config = if custom_two_byte_record_delimiter + || scan_range_requires_single_file_scan + || metered_input_requires_single_file_scan + { config.with_repartition_file_scans(false) } else { config @@ -438,11 +461,16 @@ impl SessionCtxFactory { df_session_state.with_object_store(&store_url, store).build() } else { + let input_metrics = input_metrics.unwrap_or_else(|| Arc::new(SelectInputMetrics::default())); let store: EcObjectStore = match query_tracker { - Some(query_tracker) => { - EcObjectStore::new_with_query_tracker(context.input.clone(), memory_pool, query_tracker, snapshot) - } - None => EcObjectStore::new_with_memory_pool(context.input.clone(), memory_pool, snapshot), + Some(query_tracker) => EcObjectStore::new_with_query_tracker( + context.input.clone(), + memory_pool, + query_tracker, + input_metrics, + snapshot, + ), + None => EcObjectStore::new_with_memory_pool(context.input.clone(), memory_pool, input_metrics, snapshot), } .map_err(|err| QueryError::Datafusion { source: Box::new(DataFusionError::External(Box::new(err))), @@ -587,6 +615,31 @@ mod tests { assert!(session.inner().config().options().optimizer.repartition_file_scans); } + #[tokio::test] + async fn metered_csv_and_json_inputs_disable_file_repartitioning() { + let factory = SessionCtxFactory::new(true).with_target_partitions(3); + let csv_context = test_context(); + let mut json_context = test_context(); + let json_request = &mut Arc::make_mut(&mut json_context.input).request; + json_request.input_serialization.csv = None; + json_request.input_serialization.json = Some(JSONInput::default()); + + for context in [&csv_context, &json_context] { + let session = factory + .create_session_ctx_inner( + context, + None, + None, + Some(Arc::new(SelectInputMetrics::default())), + DEFAULT_S3SELECT_MEMORY_LIMIT_BYTES, + ) + .await + .expect("metered session should be created"); + + assert!(!session.inner().config().options().optimizer.repartition_file_scans); + } + } + #[tokio::test] async fn parquet_scan_range_keeps_file_repartitioning() { let mut context = test_context(); @@ -702,7 +755,7 @@ mod tests { async fn session_factory_applies_memory_limit() { let factory = SessionCtxFactory::new(true); let session = factory - .create_session_ctx_inner(&test_context(), None, None, 1024) + .create_session_ctx_inner(&test_context(), None, None, None, 1024) .await .expect("session should be created with a bounded memory pool"); diff --git a/crates/s3select-query/src/dispatcher/manager.rs b/crates/s3select-query/src/dispatcher/manager.rs index 993d77c01..77a989da3 100644 --- a/crates/s3select-query/src/dispatcher/manager.rs +++ b/crates/s3select-query/src/dispatcher/manager.rs @@ -42,7 +42,7 @@ use rustfs_s3select_api::{ query::{ Query, ast::ExtStatement, - dispatcher::QueryDispatcher, + dispatcher::{DispatchedQuery, QueryDispatcher}, execution::{Output, QueryStateMachine}, function::FuncMetaManagerRef, logical_planner::{LogicalPlanner, Plan}, @@ -120,7 +120,7 @@ impl Drop for QueryPhaseGuard<'_> { #[async_trait] impl QueryDispatcher for SimpleQueryDispatcher { async fn execute_query(&self, query: &Query) -> QueryResult { - self.execute_query_inner(query, None).await + self.execute_query_inner(query, None).await.map(|(_, output)| output) } fn try_reserve_query(&self) -> QueryResult { @@ -133,6 +133,16 @@ impl QueryDispatcher for SimpleQueryDispatcher { } async fn execute_query_admitted(&self, query: &Query, admission: QueryAdmission) -> QueryResult { + self.execute_query_inner(query, Some(admission)) + .await + .map(|(_, output)| output) + } + + async fn dispatch_query(&self, query: &Query) -> QueryResult { + self.execute_query_inner(query, None).await + } + + async fn dispatch_query_admitted(&self, query: &Query, admission: QueryAdmission) -> QueryResult { self.execute_query_inner(query, Some(admission)).await } @@ -171,11 +181,12 @@ impl QueryDispatcher for SimpleQueryDispatcher { }; let logical_plan = self - .statement_to_logical_plan(stmt, &logical_planner, query_state_machine) + .statement_to_logical_plan(stmt, &logical_planner, Arc::clone(&query_state_machine)) .await?; Ok(logical_plan) }) .await?; + query_state_machine.query.input_metrics().reset(); if !query_tracker.mark_planned(&self.query_execution_owner) { drop(logical_plan); return Err(self.query_tracker_error(&query_tracker)); @@ -212,19 +223,21 @@ impl QueryDispatcher for SimpleQueryDispatcher { } async fn build_query_state_machine(&self, query: Query) -> QueryResult> { - self.build_query_state_machine_inner(query, None).await + self.build_query_state_machine_inner(query.for_execution(), None).await } } impl SimpleQueryDispatcher { - async fn execute_query_inner(&self, query: &Query, admission: Option) -> QueryResult { - let query_state_machine = self.build_query_state_machine_inner(query.clone(), admission).await?; + async fn execute_query_inner(&self, query: &Query, admission: Option) -> QueryResult { + let query_state_machine = self.build_query_state_machine_inner(query.for_execution(), admission).await?; + let execution_query = query_state_machine.query.clone(); let logical_plan = self.build_logical_plan(Arc::clone(&query_state_machine)).await?; let Some(logical_plan) = logical_plan else { - return Ok(Output::Nil(())); + return Ok((execution_query, Output::Nil(()))); }; - self.execute_logical_plan(logical_plan, query_state_machine).await + let output = self.execute_logical_plan(logical_plan, query_state_machine).await?; + Ok((execution_query, output)) } async fn build_query_state_machine_inner( @@ -256,29 +269,17 @@ impl SimpleQueryDispatcher { self.query_timeout.as_secs(), ); let phase_guard = QueryPhaseGuard::new(&query_tracker, &self.query_execution_owner); - let session = if let Some(snapshot) = query.snapshot().cloned() { - self.run_with_query_deadline( + let session = self + .run_with_query_deadline( &query_tracker, self.session_factory - .create_session_ctx_with_snapshot_and_tracker_and_memory_limit( - query.context(), - snapshot, + .create_session_ctx_for_query_with_tracker_and_memory_limit( + &query, query_tracker.clone(), self.memory_limit_bytes, ), ) - .await? - } else { - self.run_with_query_deadline( - &query_tracker, - self.session_factory.create_session_ctx_with_tracker_and_memory_limit( - query.context(), - query_tracker.clone(), - self.memory_limit_bytes, - ), - ) - .await? - }; + .await?; if !query_tracker.mark_admitted(&self.query_execution_owner) { drop(session); return Err(self.query_tracker_error(&query_tracker)); @@ -1705,6 +1706,55 @@ mod tests { assert_eq!(admission.available_permits(), 1); } + #[tokio::test] + async fn reused_query_gets_execution_local_input_metrics() { + let admission = Arc::new(Semaphore::new(2)); + let (dispatcher, input) = test_dispatcher(Arc::clone(&admission), Duration::from_secs(300)); + let query = Query::new(QueryContext { input }, "SELECT * FROM S3Object".to_string()); + + let first = dispatcher + .build_query_state_machine(query.clone()) + .await + .expect("first execution state"); + let second = dispatcher + .build_query_state_machine(query) + .await + .expect("second execution state"); + + assert!(!Arc::ptr_eq(first.query.input_metrics(), second.query.input_metrics())); + drop((first, second)); + assert_eq!(admission.available_permits(), 2); + } + + #[tokio::test] + async fn dispatching_a_reused_query_returns_execution_local_input_metrics() { + let env = snapshot_test_env().await; + let mut input = test_input(); + input.bucket = "s3select-reused-query-metrics".to_string(); + input.key = "input.csv".to_string(); + let input = Arc::new(input); + env.make_bucket(&input.bucket, false).await; + env.put_object_bytes(&input.bucket, &input.key, b"name\nAlice\n".to_vec()) + .await; + let snapshot = env.prepare_select_object_snapshot(&input.bucket, &input.key).await; + let dispatcher = production_dispatcher(Arc::clone(&input)); + let query = Query::new_with_snapshot( + QueryContext { + input: Arc::clone(&input), + }, + input.request.expression.clone(), + snapshot, + ); + + let (first_query, first_output) = dispatcher.dispatch_query(&query).await.expect("first dispatch should start"); + let (second_query, second_output) = dispatcher.dispatch_query(&query).await.expect("second dispatch should start"); + + assert!(!Arc::ptr_eq(first_query.input_metrics(), second_query.input_metrics())); + assert!(!Arc::ptr_eq(query.input_metrics(), first_query.input_metrics())); + assert!(!Arc::ptr_eq(query.input_metrics(), second_query.input_metrics())); + drop((first_output, second_output)); + } + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn concurrent_planning_claim_has_single_winner() { let admission = Arc::new(Semaphore::new(1)); diff --git a/crates/s3select-query/src/instance.rs b/crates/s3select-query/src/instance.rs index 444b59725..b01dc1afb 100644 --- a/crates/s3select-query/src/instance.rs +++ b/crates/s3select-query/src/instance.rs @@ -69,15 +69,15 @@ where } async fn execute(&self, query: &Query) -> QueryResult { - let result = self.query_dispatcher.execute_query(query).await?; + let (query, result) = self.query_dispatcher.dispatch_query(query).await?; - Ok(QueryHandle::new(query.clone(), result)) + Ok(QueryHandle::new(query, result)) } async fn execute_admitted(&self, query: &Query, admission: QueryAdmission) -> QueryResult { - let result = self.query_dispatcher.execute_query_admitted(query, admission).await?; + let (query, result) = self.query_dispatcher.dispatch_query_admitted(query, admission).await?; - Ok(QueryHandle::new(query.clone(), result)) + Ok(QueryHandle::new(query, result)) } async fn build_query_state_machine(&self, query: Query) -> QueryResult { @@ -247,8 +247,19 @@ pub async fn make_rustfsms_with_components( mod tests { use std::sync::Arc; + use async_trait::async_trait; use datafusion::{arrow::util::pretty, assert_batches_eq}; - use rustfs_s3select_api::query::{Context, Query}; + use parking_lot::Mutex; + use rustfs_s3select_api::{ + QueryResult, SelectInputMetrics, + query::{ + Context, Query, + dispatcher::QueryDispatcher, + execution::{Output, QueryStateMachine}, + logical_planner::Plan, + }, + server::dbms::DatabaseManagerSystem, + }; use s3s::dto::{ CSVInput, CSVOutput, ExpressionType, FieldDelimiter, FileHeaderInfo, InputSerialization, OutputSerialization, RecordDelimiter, SelectObjectContentInput, SelectObjectContentRequest, @@ -257,10 +268,66 @@ mod tests { use crate::get_global_db; use super::{ - DEFAULT_MAX_CONCURRENT_QUERIES, DEFAULT_MEMORY_LIMIT_BYTES, DEFAULT_QUERY_TIMEOUT_SECS, MAX_QUERY_TIMEOUT_SECS, + DEFAULT_MAX_CONCURRENT_QUERIES, DEFAULT_MEMORY_LIMIT_BYTES, DEFAULT_QUERY_TIMEOUT_SECS, MAX_QUERY_TIMEOUT_SECS, RustFSms, S3SelectRuntimeConfig, bounded_u64_from_env_value, bounded_usize_from_env_value, target_partitions_from_env_value, }; + #[derive(Default)] + struct FreshMetricsDispatcher { + executed_metrics: Mutex>>, + } + + #[async_trait] + impl QueryDispatcher for FreshMetricsDispatcher { + async fn execute_query(&self, query: &Query) -> QueryResult { + self.executed_metrics.lock().push(Arc::clone(query.input_metrics())); + Ok(Output::Nil(())) + } + + async fn build_logical_plan(&self, _query_state_machine: Arc) -> QueryResult> { + unreachable!("fresh metrics test does not plan queries") + } + + async fn execute_logical_plan( + &self, + _logical_plan: Plan, + _query_state_machine: Arc, + ) -> QueryResult { + unreachable!("fresh metrics test does not execute plans") + } + + async fn build_query_state_machine(&self, _query: Query) -> QueryResult> { + unreachable!("fresh metrics test does not build state machines") + } + } + + fn metrics_test_query() -> Query { + let expression = "SELECT * FROM S3Object"; + let input = SelectObjectContentInput { + bucket: "bucket".to_string(), + expected_bucket_owner: None, + key: "input.csv".to_string(), + sse_customer_algorithm: None, + sse_customer_key: None, + sse_customer_key_md5: None, + request: SelectObjectContentRequest { + expression: expression.to_string(), + expression_type: ExpressionType::from_static(ExpressionType::SQL), + input_serialization: InputSerialization { + csv: Some(CSVInput::default()), + ..Default::default() + }, + output_serialization: OutputSerialization { + csv: Some(CSVOutput::default()), + ..Default::default() + }, + request_progress: None, + scan_range: None, + }, + }; + Query::new(Context { input: Arc::new(input) }, expression.to_string()) + } + #[test] fn parses_target_partitions_from_env_value() { assert_eq!(target_partitions_from_env_value(Some("4")), 4); @@ -291,6 +358,26 @@ mod tests { assert_eq!(bounded_u64_from_env_value(None, 300, MAX_QUERY_TIMEOUT_SECS), 300); } + #[tokio::test] + async fn repeated_execute_returns_the_fresh_dispatched_query_metrics() { + let dispatcher = Arc::new(FreshMetricsDispatcher::default()); + let db = RustFSms { + query_dispatcher: Arc::clone(&dispatcher), + }; + let query = metrics_test_query(); + + let first = db.execute(&query).await.expect("first execution should succeed"); + let second = db.execute(&query).await.expect("second execution should succeed"); + let executed_metrics = dispatcher.executed_metrics.lock(); + + assert_eq!(executed_metrics.len(), 2); + assert!(Arc::ptr_eq(first.query().input_metrics(), &executed_metrics[0])); + assert!(Arc::ptr_eq(second.query().input_metrics(), &executed_metrics[1])); + assert!(!Arc::ptr_eq(first.query().input_metrics(), second.query().input_metrics())); + assert!(!Arc::ptr_eq(first.query().input_metrics(), query.input_metrics())); + assert!(!Arc::ptr_eq(second.query().input_metrics(), query.input_metrics())); + } + #[tokio::test] #[ignore = "requires a live RustFS store with a pre-seeded test object (bucket 'dandan')"] async fn test_simple_sql() { diff --git a/rustfs/src/app/select_object.rs b/rustfs/src/app/select_object.rs index 600a10727..b4aa98408 100644 --- a/rustfs/src/app/select_object.rs +++ b/rustfs/src/app/select_object.rs @@ -21,7 +21,7 @@ use datafusion::physical_plan::SendableRecordBatchStream; use futures::StreamExt; use http::{HeaderMap, StatusCode, header::RANGE}; use rustfs_s3select_api::{ - QueryError, SelectError, + QueryError, SelectError, SelectInputMetrics, object_store::{INVALID_SCAN_RANGE_MESSAGE, validate_scan_range_bounds}, query::{Context, Query}, }; @@ -53,6 +53,7 @@ const UNSUPPORTED_SQL_STRUCTURE_MESSAGE: &str = "We encountered an unsupported S struct SelectValidation { output_format: SelectOutputFormat, progress_enabled: bool, + reports_input_metrics: bool, } #[derive(Clone, Debug)] @@ -66,6 +67,11 @@ enum SelectProducerOutcome { ReceiverClosed, } +struct SelectEventChannel { + tx: mpsc::Sender>, + terminal_permit: mpsc::OwnedPermit>, +} + trait SelectSnapshotFence { fn ensure_snapshot_valid(&self) -> S3Result<()>; } @@ -101,10 +107,12 @@ pub async fn execute_select_object_content( let snapshot = Arc::new(snapshot); let query = Query::new_with_snapshot(Context { input: input.clone() }, input.request.expression.clone(), Arc::clone(&snapshot)); - let output = timeout_at(query_deadline, db.execute_admitted(&query, admission)) + let query_handle = timeout_at(query_deadline, db.execute_admitted(&query, admission)) .await .map_err(|_| select_query_timeout_error(query_timeout.as_secs()))? - .map_err(map_query_error_to_s3)? + .map_err(map_query_error_to_s3)?; + let input_metrics = Arc::clone(query_handle.query().input_metrics()); + let output = query_handle .result() .into_record_batch_stream() .map_err(map_query_error_to_s3)?; @@ -121,9 +129,9 @@ pub async fn execute_select_object_content( spawn_traced(async move { send_select_events_until_deadline( output, - tx, - terminal_permit, + SelectEventChannel { tx, terminal_permit }, validation, + input_metrics, query_deadline, query_timeout.as_secs(), snapshot, @@ -136,14 +144,19 @@ pub async fn execute_select_object_content( async fn send_select_events_until_deadline( output: SendableRecordBatchStream, - tx: mpsc::Sender>, - terminal_permit: mpsc::OwnedPermit>, + event_channel: SelectEventChannel, validation: SelectValidation, + input_metrics: Arc, deadline: Instant, timeout_seconds: u64, snapshot_lease: L, ) { - let outcome = match timeout_at(deadline, send_select_events(output, &tx, validation, &snapshot_lease)).await { + let outcome = match timeout_at( + deadline, + send_select_events(output, &event_channel.tx, validation, input_metrics, &snapshot_lease), + ) + .await + { Ok(outcome) => outcome, Err(_) => SelectProducerOutcome::Terminal(Err(map_query_error_to_s3( SelectError::QueryTimeout { @@ -153,7 +166,7 @@ async fn send_select_events_until_deadline( ))), }; if let SelectProducerOutcome::Terminal(event) = outcome { - terminal_permit.send(event); + event_channel.terminal_permit.send(event); } drop(snapshot_lease); } @@ -162,10 +175,11 @@ async fn send_select_events( mut output: SendableRecordBatchStream, tx: &mpsc::Sender>, validation: SelectValidation, + input_metrics: Arc, snapshot_fence: &impl SelectSnapshotFence, ) -> SelectProducerOutcome { let mut encoder = SelectOutputEncoder::new(validation.output_format); - let mut progress = SelectProgress::default(); + let mut progress = SelectProgress::new(validation.reports_input_metrics.then_some(input_metrics)); if tx .send(Ok(SelectObjectContentEvent::Cont(ContinuationEvent::default()))) @@ -192,7 +206,7 @@ async fn send_select_events( match encoder.encode_batch(&batch) { Ok(payloads) => { for payload in payloads { - progress.add_returned(payload.len()); + let payload_len = payload.len(); if tx .send(Ok(SelectObjectContentEvent::Records(RecordsEvent { payload: Some(payload) }))) .await @@ -200,6 +214,7 @@ async fn send_select_events( { return SelectProducerOutcome::ReceiverClosed; } + progress.add_returned(payload_len); if validation.progress_enabled && tx .send(Ok(SelectObjectContentEvent::Progress(ProgressEvent { @@ -265,6 +280,7 @@ fn validate_select_request(headers: &http::HeaderMap, input: &mut SelectObjectCo Ok(SelectValidation { output_format, progress_enabled, + reports_input_metrics: input.request.input_serialization.parquet.is_none(), }) } @@ -602,29 +618,39 @@ fn split_records_payload(bytes: Vec) -> Vec { .collect() } -#[derive(Default)] struct SelectProgress { + input_metrics: Option>, bytes_returned: u64, } impl SelectProgress { + fn new(input_metrics: Option>) -> Self { + Self { + input_metrics, + bytes_returned: 0, + } + } + fn add_returned(&mut self, bytes: usize) { - self.bytes_returned = self.bytes_returned.saturating_add(bytes as u64); + let bytes = u64::try_from(bytes).unwrap_or(u64::MAX); + self.bytes_returned = self.bytes_returned.saturating_add(bytes); } fn to_progress(&self) -> Progress { + let input = self.input_metrics.as_ref().map(|metrics| metrics.snapshot()); Progress { - bytes_processed: None, + bytes_processed: input.map(|metrics| clamp_i64(metrics.bytes_processed)), bytes_returned: Some(clamp_i64(self.bytes_returned)), - bytes_scanned: None, + bytes_scanned: input.map(|metrics| clamp_i64(metrics.bytes_scanned)), } } fn to_stats(&self) -> Stats { + let input = self.input_metrics.as_ref().map(|metrics| metrics.snapshot()); Stats { - bytes_processed: None, + bytes_processed: input.map(|metrics| clamp_i64(metrics.bytes_processed)), bytes_returned: Some(clamp_i64(self.bytes_returned)), - bytes_scanned: None, + bytes_scanned: input.map(|metrics| clamp_i64(metrics.bytes_scanned)), } } } @@ -852,6 +878,7 @@ mod tests { SelectValidation { output_format: SelectOutputFormat::Csv(CSVOutput::default()), progress_enabled: false, + reports_input_metrics: true, } } @@ -871,9 +898,9 @@ mod tests { let (lease, lease_released) = lease_drop_signal(); let producer = tokio::spawn(send_select_events_until_deadline( output, - tx, - terminal_permit, + SelectEventChannel { tx, terminal_permit }, csv_validation(), + Arc::new(SelectInputMetrics::default()), Instant::now() + std::time::Duration::from_secs(1), 300, lease, @@ -1096,9 +1123,9 @@ mod tests { let (lease, lease_released) = lease_drop_signal(); let producer = tokio::spawn(send_select_events_until_deadline( output, - tx, - terminal_permit, + SelectEventChannel { tx, terminal_permit }, csv_validation(), + Arc::new(SelectInputMetrics::default()), Instant::now() + std::time::Duration::from_secs(1), 300, lease, @@ -1451,9 +1478,9 @@ mod tests { let (lease, lease_released) = lease_drop_signal(); let producer = send_select_events_until_deadline( output, - tx, - terminal_permit, + SelectEventChannel { tx, terminal_permit }, csv_validation(), + Arc::new(SelectInputMetrics::default()), Instant::now() + std::time::Duration::from_secs(1), 300, lease, @@ -1489,7 +1516,8 @@ mod tests { )); let (tx, mut rx) = mpsc::channel(2); let snapshot_fence = LeaseDropSignal(None); - let producer = send_select_events(output, &tx, csv_validation(), &snapshot_fence); + let producer = + send_select_events(output, &tx, csv_validation(), Arc::new(SelectInputMetrics::default()), &snapshot_fence); tokio::pin!(producer); assert!(futures::poll!(producer.as_mut()).is_pending()); @@ -1525,7 +1553,14 @@ mod tests { )); let (tx, mut rx) = mpsc::channel(4); - let outcome = send_select_events(output, &tx, csv_validation(), &FailingSnapshotFence).await; + let outcome = send_select_events( + output, + &tx, + csv_validation(), + Arc::new(SelectInputMetrics::default()), + &FailingSnapshotFence, + ) + .await; let SelectProducerOutcome::Terminal(Err(error)) = outcome else { panic!("failed final snapshot fence must produce a terminal error"); @@ -1548,7 +1583,8 @@ mod tests { .try_reserve_owned() .expect("test channel should reserve terminal capacity"); let snapshot_fence = FailsAfterFirstSnapshotFence(std::sync::atomic::AtomicUsize::new(0)); - let producer = send_select_events(output, &tx, csv_validation(), &snapshot_fence); + let producer = + send_select_events(output, &tx, csv_validation(), Arc::new(SelectInputMetrics::default()), &snapshot_fence); tokio::pin!(producer); assert!(futures::poll!(producer.as_mut()).is_pending()); @@ -1802,7 +1838,7 @@ mod tests { #[test] fn split_records_payload_uses_exact_returned_bytes() { let payloads = split_records_payload(vec![b'x'; RECORDS_CHUNK_TARGET + 7]); - let mut progress = SelectProgress::default(); + let mut progress = SelectProgress::new(Some(Arc::new(SelectInputMetrics::default()))); for payload in &payloads { progress.add_returned(payload.len()); } @@ -1922,13 +1958,40 @@ mod tests { } #[test] - fn progress_does_not_report_unknown_input_bytes_as_zero() { - let mut progress = SelectProgress::default(); + fn progress_reports_zero_for_an_empty_input() { + let mut progress = SelectProgress::new(Some(Arc::new(SelectInputMetrics::default()))); progress.add_returned(12); let stats = progress.to_stats(); assert_eq!(stats.bytes_returned, Some(12)); + assert_eq!(stats.bytes_scanned, Some(0)); + assert_eq!(stats.bytes_processed, Some(0)); + } + + #[test] + fn progress_dto_clamps_counters_to_signed_event_range() { + assert_eq!(clamp_i64(u64::MAX), i64::MAX); + } + + #[test] + fn parquet_progress_keeps_input_metrics_unspecified() { + let mut input = base_input(); + input.request.input_serialization = InputSerialization { + csv: None, + json: None, + parquet: Some(ParquetInput {}), + compression_type: None, + }; + let validation = validate_select_request(&HeaderMap::new(), &mut input).expect("Parquet request should validate"); + let progress = SelectProgress::new( + validation + .reports_input_metrics + .then(|| Arc::new(SelectInputMetrics::default())), + ); + + let stats = progress.to_stats(); assert_eq!(stats.bytes_scanned, None); assert_eq!(stats.bytes_processed, None); + assert_eq!(stats.bytes_returned, Some(0)); } #[test]