gsadmin/rustfs
1
0
Fork 0
mirror of https://github.com/rustfs/rustfs.git synced 2026-09-01 09:48:20 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(s3select): report uncompressed input byte metrics (#6865)

Browse Source
This commit is contained in:
GatewayJ
2026-08-30 12:07:20 +08:00
committed by GitHub
parent b428875bed
commit 96239fc034
10 changed files with 1151 additions and 115 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/s3select-api/src/lib.rs
+2
View File
@@ -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)]
crates/s3select-api/src/metrics.rs
+88
View File
@@ -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());
}
}
crates/s3select-api/src/object_store.rs
+386 -41
View File
@@ -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<String>,
input_metrics: Arc<SelectInputMetrics>,
memory_pool: Arc<dyn MemoryPool>,
query_tracker: Option<QueryExecutionTracker>,
store: Option<Arc<SelectStore>>,
@@ -172,21 +174,35 @@ pub struct InvalidScanRange;
impl EcObjectStore {
pub fn new(input: Arc<SelectObjectContentInput>) -> S3Result<Self> {
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<SelectObjectContentInput>, snapshot: Arc<SelectObjectSnapshot>) -> S3Result<Self> {
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<SelectObjectContentInput>,
memory_pool: Arc<dyn MemoryPool>,
input_metrics: Arc<SelectInputMetrics>,
snapshot: Option<Arc<SelectObjectSnapshot>>,
) -> std::result::Result<Self, EcObjectStoreBuildError> {
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<SelectObjectContentInput>,
memory_pool: Arc<dyn MemoryPool>,
query_tracker: QueryExecutionTracker,
input_metrics: Arc<SelectInputMetrics>,
snapshot: Option<Arc<SelectObjectSnapshot>>,
) -> std::result::Result<Self, EcObjectStoreBuildError> {
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<SelectObjectContentInput>,
memory_pool: Arc<dyn MemoryPool>,
query_tracker: Option<QueryExecutionTracker>,
input_metrics: Arc<SelectInputMetrics>,
) -> std::result::Result<Self, EcObjectStoreBuildError> {
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<SelectObjectContentInput>,
memory_pool: Arc<dyn MemoryPool>,
query_tracker: Option<QueryExecutionTracker>,
input_metrics: Arc<SelectInputMetrics>,
snapshot: Arc<SelectObjectSnapshot>,
) -> std::result::Result<Self, EcObjectStoreBuildError> {
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<SelectObjectContentInput>,
memory_pool: Arc<dyn MemoryPool>,
query_tracker: Option<QueryExecutionTracker>,
input_metrics: Arc<SelectInputMetrics>,
store: Option<Arc<SelectStore>>,
snapshot: Option<Arc<SelectObjectSnapshot>>,
) -> 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<dyn tokio::io::AsyncRead + Unpin + Send + Sync>,
original_size: u64,
json_sub_path: Option<String>,
input_metrics: Arc<SelectInputMetrics>,
memory_pool: Arc<dyn MemoryPool>,
query_tracker: Option<QueryExecutionTracker>,
) -> futures_core::stream::BoxStream<'static, Result<Bytes>> {
@@ -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<P>(
stream: Box<dyn tokio::io::AsyncRead + Unpin + Send + Sync>,
original_size: u64,
json_sub_path: Option<String>,
input_metrics: Arc<SelectInputMetrics>,
memory_pool: Arc<dyn MemoryPool>,
query_tracker: Option<QueryExecutionTracker>,
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<S, E>(
stream: S,
input_metrics: Arc<SelectInputMetrics>,
) -> impl Stream<Item = std::result::Result<Bytes, E>> + Send + 'static
where
S: Stream<Item = std::result::Result<Bytes, E>> + Send + 'static,
E: Send + 'static,
{
stream.inspect_ok(move |bytes| input_metrics.record_uncompressed(bytes.len()))
}
pub fn bytes_stream<S>(stream: S, content_length: usize) -> impl Stream<Item = Result<Bytes>> + Send + 'static
where
S: Stream<Item = Result<Bytes, std::io::Error>> + 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<SelectObjectContentInput> {
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::<Vec<_>>()
.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<Bytes> = 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::<Vec<_>>()
.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::<Vec<_>>()
.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::<Vec<_>>().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::<Vec<_>>()
.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<Bytes> = 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<dyn MemoryPool> = 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::<std::io::Error>().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 |_, _| {
crates/s3select-api/src/query/dispatcher.rs
+166
View File
@@ -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<DispatchedQuery> {
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<DispatchedQuery> {
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<QueryStateMachine>) -> QueryResult<Option<Plan>>;
async fn execute_logical_plan(&self, logical_plan: Plan, query_state_machine: Arc<QueryStateMachine>) -> QueryResult<Output>;
@@ -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<Vec<Arc<crate::SelectInputMetrics>>>,
}
#[async_trait]
impl QueryDispatcher for DefaultDispatchDispatcher {
async fn execute_query(&self, query: &Query) -> QueryResult<Output> {
self.executed_metrics.lock().push(Arc::clone(query.input_metrics()));
Ok(Output::Nil(()))
}
async fn build_logical_plan(&self, _query_state_machine: Arc<QueryStateMachine>) -> QueryResult<Option<Plan>> {
unreachable!("default dispatch test does not plan queries")
}
async fn execute_logical_plan(
&self,
_logical_plan: Plan,
_query_state_machine: Arc<QueryStateMachine>,
) -> QueryResult<Output> {
unreachable!("default dispatch test does not execute plans")
}
async fn build_query_state_machine(&self, _query: Query) -> QueryResult<Arc<QueryStateMachine>> {
unreachable!("default dispatch test does not build state machines")
}
}
#[derive(Default)]
struct DistinctAdmittedDispatcher {
plain_metrics: Mutex<Vec<Arc<crate::SelectInputMetrics>>>,
admitted_metrics: Mutex<Vec<Arc<crate::SelectInputMetrics>>>,
fail_plain: bool,
fail_admitted: bool,
}
#[async_trait]
impl QueryDispatcher for DistinctAdmittedDispatcher {
async fn execute_query(&self, query: &Query) -> QueryResult<Output> {
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<Output> {
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<QueryStateMachine>) -> QueryResult<Option<Plan>> {
unreachable!("dispatch routing test does not plan queries")
}
async fn execute_logical_plan(
&self,
_logical_plan: Plan,
_query_state_machine: Arc<QueryStateMachine>,
) -> QueryResult<Output> {
unreachable!("dispatch routing test does not execute plans")
}
async fn build_query_state_machine(&self, _query: Query) -> QueryResult<Arc<QueryStateMachine>> {
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);
}
}
crates/s3select-api/src/query/mod.rs
+46 -1
View File
@@ -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<Arc<SelectObjectSnapshot>>,
input_metrics: Arc<SelectInputMetrics>,
}
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<SelectObjectSnapshot>> {
self.snapshot.as_ref()
}
pub fn input_metrics(&self) -> &Arc<SelectInputMetrics> {
&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())
}
crates/s3select-api/src/query/session.rs
+65 -12
View File
@@ -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<SessionCtx> {
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<SessionCtx> {
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<SessionCtx> {
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<SessionCtx> {
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<Arc<SelectObjectSnapshot>>,
query_tracker: Option<QueryExecutionTracker>,
input_metrics: Option<Arc<SelectInputMetrics>>,
memory_limit_bytes: usize,
) -> QueryResult<SessionCtx> {
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<Arc<SelectObjectSnapshot>>,
query_tracker: Option<QueryExecutionTracker>,
input_metrics: Option<Arc<SelectInputMetrics>>,
memory_limit_bytes: usize,
) -> QueryResult<SessionContext> {
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");