feat(s3select): support compressed CSV and JSON input (#6915)

This commit is contained in:
GatewayJ
2026-08-31 13:35:59 +08:00
committed by GitHub
parent 1d606e1cf6
commit 589a954478
12 changed files with 3057 additions and 144 deletions
Generated
+6
View File
@@ -3926,6 +3926,7 @@ dependencies = [
"aws-sdk-s3",
"aws-sdk-sts",
"aws-smithy-http-client",
"aws-smithy-types",
"base64-simd",
"bytes",
"chrono",
@@ -10557,10 +10558,14 @@ dependencies = [
name = "rustfs-s3select-api"
version = "1.0.0-rc.4"
dependencies = [
"arc-swap",
"async-compression",
"async-trait",
"bytes",
"chrono",
"crc-fast",
"datafusion",
"flate2",
"futures",
"futures-core",
"hotpath",
@@ -10576,6 +10581,7 @@ dependencies = [
"serial_test",
"thiserror 2.0.20",
"tokio",
"tokio-stream",
"tokio-util",
"tracing",
"transform-stream",
+1
View File
@@ -100,6 +100,7 @@ aws-sdk-s3 = { workspace = true, default-features = false, features = ["sigv4a",
aws-sdk-sts = { workspace = true, default-features = false, features = ["default-https-client", "rt-tokio"] }
aws-config = { workspace = true }
aws-smithy-http-client = { workspace = true, default-features = false, features = ["rustls-aws-lc"] }
aws-smithy-types.workspace = true
async-compression = { workspace = true, features = ["tokio", "bzip2", "xz"] }
async-trait = { workspace = true }
flate2.workspace = true
+1
View File
@@ -21,5 +21,6 @@ mod head_tls_bodyless_test;
mod lifecycle;
mod lock;
mod node_interact_test;
mod s3_select_compression;
mod sql;
mod tiering;
@@ -0,0 +1,351 @@
#![cfg(test)]
// 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 crate::common::{RustFSTestEnvironment, init_logging};
use async_compression::tokio::write::BzEncoder;
use aws_sdk_s3::{
Client,
error::ProvideErrorMetadata,
operation::select_object_content::{SelectObjectContentOutput, builders::SelectObjectContentFluentBuilder},
types::{
CompressionType, CsvInput, CsvOutput, ExpressionType, FileHeaderInfo, InputSerialization, JsonInput, JsonOutput,
JsonType, OutputSerialization, SelectObjectContentEventStream,
},
};
use aws_smithy_types::event_stream::RawMessage;
use bytes::Bytes;
use flate2::{Compression, write::GzEncoder};
use std::{error::Error, io::Cursor, time::Duration};
use tokio::io::AsyncWriteExt;
const BUCKET: &str = "s3-select-compression";
const SELECT_RESPONSE_TIMEOUT: Duration = Duration::from_secs(30);
type TestResult<T> = Result<T, Box<dyn Error + Send + Sync>>;
async fn create_test_environment(extra_env: &[(&str, &str)]) -> TestResult<(RustFSTestEnvironment, Client)> {
init_logging();
let mut env = RustFSTestEnvironment::new().await?;
env.start_rustfs_server_with_env(vec![], extra_env).await?;
let client = env.create_s3_client();
client.create_bucket().bucket(BUCKET).send().await?;
Ok((env, client))
}
async fn put_object(client: &Client, key: &str, body: &[u8]) -> TestResult<()> {
client
.put_object()
.bucket(BUCKET)
.key(key)
.body(Bytes::copy_from_slice(body).into())
.send()
.await?;
Ok(())
}
fn gzip(input: &[u8]) -> TestResult<Vec<u8>> {
let mut encoder = GzEncoder::new(Vec::new(), Compression::default());
std::io::Write::write_all(&mut encoder, input)?;
Ok(encoder.finish()?)
}
async fn bzip2(input: &[u8]) -> TestResult<Vec<u8>> {
let mut encoder = BzEncoder::new(Cursor::new(Vec::new()));
encoder.write_all(input).await?;
encoder.shutdown().await?;
Ok(encoder.into_inner().into_inner())
}
fn csv_select_request(
client: &Client,
key: &str,
compression: CompressionType,
expression: &str,
) -> SelectObjectContentFluentBuilder {
client
.select_object_content()
.bucket(BUCKET)
.key(key)
.expression(expression)
.expression_type(ExpressionType::Sql)
.input_serialization(
InputSerialization::builder()
.compression_type(compression)
.csv(CsvInput::builder().file_header_info(FileHeaderInfo::Use).build())
.build(),
)
.output_serialization(OutputSerialization::builder().csv(CsvOutput::builder().build()).build())
}
fn json_select_request(
client: &Client,
key: &str,
compression: CompressionType,
json_type: JsonType,
) -> SelectObjectContentFluentBuilder {
client
.select_object_content()
.bucket(BUCKET)
.key(key)
.expression("SELECT name FROM S3Object")
.expression_type(ExpressionType::Sql)
.input_serialization(
InputSerialization::builder()
.compression_type(compression)
.json(JsonInput::builder().set_type(Some(json_type)).build())
.build(),
)
.output_serialization(OutputSerialization::builder().json(JsonOutput::builder().build()).build())
}
async fn collect_success(
mut response: SelectObjectContentOutput,
compressed_bytes: usize,
processed_bytes: usize,
) -> TestResult<Vec<u8>> {
tokio::time::timeout(SELECT_RESPONSE_TIMEOUT, async move {
let mut records = Vec::new();
let mut stats = None;
let mut saw_end = false;
while let Some(event) = response.payload.recv().await? {
assert!(!saw_end, "Select emitted an event after End");
match event {
SelectObjectContentEventStream::Records(event) => {
assert!(stats.is_none(), "Select emitted Records after Stats");
if let Some(payload) = event.payload {
records.extend_from_slice(payload.as_ref());
}
}
SelectObjectContentEventStream::Stats(event) => {
assert!(stats.is_none(), "Select emitted more than one Stats event");
stats = event.details;
}
SelectObjectContentEventStream::End(_) => {
assert!(stats.is_some(), "Select emitted End before Stats");
saw_end = true;
}
_ => assert!(stats.is_none(), "Select emitted a non-terminal event after Stats"),
}
}
let stats = stats.ok_or("Select response ended without a Stats event")?;
assert_eq!(stats.bytes_scanned(), Some(i64::try_from(compressed_bytes)?));
assert_eq!(stats.bytes_processed(), Some(i64::try_from(processed_bytes)?));
assert_eq!(stats.bytes_returned(), Some(i64::try_from(records.len())?));
assert!(saw_end, "Select response ended without an End event");
Ok::<_, Box<dyn Error + Send + Sync>>(records)
})
.await
.map_err(|_| -> Box<dyn Error + Send + Sync> { "Select response timed out".into() })?
}
async fn assert_truncated_stream_failure(mut response: SelectObjectContentOutput) -> TestResult<()> {
tokio::time::timeout(SELECT_RESPONSE_TIMEOUT, async move {
loop {
match response.payload.recv().await {
Err(error) => {
// S3 Select request-level errors use `error` frames, which this SDK version exposes as raw response errors.
if let Some(code) = error.code() {
assert_eq!(code, "TruncatedInput", "unexpected modeled event-stream error: {error:?}");
} else if let aws_sdk_s3::error::SdkError::ResponseError(context) = &error
&& let RawMessage::Decoded(message) = context.raw()
{
let header = |name: &str| {
message
.headers()
.iter()
.find(|header| header.name().as_str() == name)
.and_then(|header| header.value().as_string().ok())
.map(|value| value.as_str())
};
assert_eq!(header(":message-type"), Some("error"));
assert_eq!(header(":error-code"), Some("TruncatedInput"));
} else {
panic!("unexpected event-stream error: {error:?}");
}
return Ok(());
}
Ok(Some(SelectObjectContentEventStream::Stats(_))) | Ok(Some(SelectObjectContentEventStream::End(_))) => {
return Err("truncated compressed input reached a success terminal event".into());
}
Ok(Some(_)) => {}
Ok(None) => return Err("truncated compressed input ended without an error event".into()),
}
}
})
.await
.map_err(|_| -> Box<dyn Error + Send + Sync> { "truncated Select response timed out".into() })?
}
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn test_select_object_content_compressed_csv_and_json() -> TestResult<()> {
const CSV: &[u8] = b"name,age\nAlice,30\nBob,25\n";
const JSON_LINES: &[u8] = b"{\"name\":\"Alice\"}\n{\"name\":\"Bob\"}\n";
const JSON_DOCUMENT: &[u8] = br#"[{"name":"Alice"},{"name":"Bob"}]"#;
let (_env, client) = create_test_environment(&[]).await?;
let gzip_csv = gzip(CSV)?;
put_object(&client, "records.csv.gz", &gzip_csv).await?;
let gzip_csv_records = collect_success(
csv_select_request(&client, "records.csv.gz", CompressionType::Gzip, "SELECT * FROM S3Object")
.send()
.await?,
gzip_csv.len(),
CSV.len(),
)
.await?;
assert_eq!(gzip_csv_records, b"Alice,30\nBob,25\n");
let bzip_csv = bzip2(CSV).await?;
put_object(&client, "records.csv.bz2", &bzip_csv).await?;
let bzip_csv_records = collect_success(
csv_select_request(&client, "records.csv.bz2", CompressionType::Bzip2, "SELECT * FROM S3Object")
.send()
.await?,
bzip_csv.len(),
CSV.len(),
)
.await?;
assert_eq!(bzip_csv_records, gzip_csv_records);
let gzip_json_lines = gzip(JSON_LINES)?;
put_object(&client, "json-lines", &gzip_json_lines).await?;
let gzip_json_records = collect_success(
json_select_request(&client, "json-lines", CompressionType::Gzip, JsonType::Lines)
.send()
.await?,
gzip_json_lines.len(),
JSON_LINES.len(),
)
.await?;
assert_eq!(gzip_json_records, JSON_LINES);
let bzip_json_lines = bzip2(JSON_LINES).await?;
put_object(&client, "records.jsonl.bz2", &bzip_json_lines).await?;
let bzip_json_records = collect_success(
json_select_request(&client, "records.jsonl.bz2", CompressionType::Bzip2, JsonType::Lines)
.send()
.await?,
bzip_json_lines.len(),
JSON_LINES.len(),
)
.await?;
assert_eq!(bzip_json_records, gzip_json_records);
let gzip_json_document = gzip(JSON_DOCUMENT)?;
put_object(&client, "document.json.gz", &gzip_json_document).await?;
let document_records = collect_success(
json_select_request(&client, "document.json.gz", CompressionType::Gzip, JsonType::Document)
.send()
.await?,
gzip_json_document.len(),
JSON_DOCUMENT.len(),
)
.await?;
assert_eq!(document_records, JSON_LINES);
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn test_select_object_content_invalid_compressed_stream_fails() -> TestResult<()> {
const CSV: &[u8] = b"name\nAlice\n";
let (_env, client) = create_test_environment(&[]).await?;
put_object(&client, "invalid.csv.gz", CSV).await?;
let invalid = csv_select_request(&client, "invalid.csv.gz", CompressionType::Gzip, "SELECT * FROM S3Object")
.send()
.await
.expect_err("invalid GZIP header must fail before streaming");
assert_eq!(
invalid.as_service_error().and_then(ProvideErrorMetadata::code),
Some("InvalidCompressionFormat")
);
put_object(&client, "empty.csv.gz", b"").await?;
let empty = csv_select_request(&client, "empty.csv.gz", CompressionType::Gzip, "SELECT * FROM S3Object")
.send()
.await
.expect_err("empty GZIP input must fail as truncated");
assert_eq!(empty.as_service_error().and_then(ProvideErrorMetadata::code), Some("TruncatedInput"));
let mut truncated = bzip2(CSV).await?;
truncated.pop();
put_object(&client, "truncated.csv.bz2", &truncated).await?;
let truncated = csv_select_request(&client, "truncated.csv.bz2", CompressionType::Bzip2, "SELECT * FROM S3Object")
.send()
.await?;
assert_truncated_stream_failure(truncated).await?;
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn test_select_object_content_compressed_disconnect_releases_query() -> TestResult<()> {
const OBJECT: &str = "disconnect.csv.gz";
const ROWS: usize = 16 * 1024;
const RELEASE_ATTEMPTS: usize = 20;
const RELEASE_BACKOFF: Duration = Duration::from_millis(25);
let (_env, client) = create_test_environment(&[("RUSTFS_S3SELECT_MAX_CONCURRENT_QUERIES", "1")]).await?;
let row = format!("{}\n", "x".repeat(1023));
let mut body = Vec::with_capacity("value\n".len() + ROWS * row.len());
body.extend_from_slice(b"value\n");
for _ in 0..ROWS {
body.extend_from_slice(row.as_bytes());
}
let compressed = gzip(&body)?;
put_object(&client, OBJECT, &compressed).await?;
let first = csv_select_request(&client, OBJECT, CompressionType::Gzip, "SELECT * FROM S3Object")
.send()
.await?;
let saturated = csv_select_request(&client, OBJECT, CompressionType::Gzip, "SELECT * FROM S3Object")
.send()
.await
.expect_err("the unread compressed response should retain the only query permit");
assert_eq!(saturated.as_service_error().and_then(ProvideErrorMetadata::code), Some("SlowDown"));
drop(first);
let second = tokio::time::timeout(Duration::from_secs(5), async {
for attempt in 0..RELEASE_ATTEMPTS {
match csv_select_request(&client, OBJECT, CompressionType::Gzip, "SELECT * FROM S3Object")
.send()
.await
{
Ok(response) => return Ok::<_, Box<dyn Error + Send + Sync>>(response),
Err(error)
if error.as_service_error().and_then(ProvideErrorMetadata::code) == Some("SlowDown")
&& attempt + 1 < RELEASE_ATTEMPTS =>
{
tokio::time::sleep(RELEASE_BACKOFF).await;
}
Err(error) if error.as_service_error().and_then(ProvideErrorMetadata::code) == Some("SlowDown") => {
return Err("disconnected compressed Select retained its query permit".into());
}
Err(error) => return Err(format!("unexpected Select error after disconnect: {error}").into()),
}
}
Err("query permit release retry loop ended unexpectedly".into())
})
.await
.map_err(|_| -> Box<dyn Error + Send + Sync> { "compressed Select did not release its query permit".into() })??;
drop(second);
Ok(())
}
+5
View File
@@ -60,21 +60,26 @@ hotpath-cpu = [
[dependencies]
hotpath.workspace = true
metrics = { workspace = true }
async-compression = { workspace = true, features = ["tokio", "gzip", "bzip2"] }
async-trait.workspace = true
arc-swap.workspace = true
bytes = { workspace = true, features = ["serde"] }
chrono = { workspace = true, features = ["serde"] }
crc-fast.workspace = true
rustfs-common.workspace = true
datafusion = { workspace = true, default-features = false, features = ["parquet", "recursive_protection", "sql"] }
rustfs-ecstore.workspace = true
rustfs-storage-api.workspace = true
futures = { workspace = true }
futures-core = { workspace = true }
flate2.workspace = true
http.workspace = true
s3s = { workspace = true, features = ["minio"] }
serde_json = { workspace = true, features = ["raw_value"] }
thiserror = { workspace = true }
parking_lot.workspace = true
tokio = { workspace = true, features = ["fs", "rt-multi-thread"] }
tokio-stream.workspace = true
tokio-util = { workspace = true, features = ["io", "compat"] }
tracing.workspace = true
uuid.workspace = true
File diff suppressed because it is too large Load Diff
+10
View File
@@ -23,6 +23,7 @@ use datafusion::{
use std::{error::Error as StdError, fmt::Display};
use thiserror::Error;
mod input_stream;
mod metrics;
pub mod object_store;
pub mod query;
@@ -79,6 +80,9 @@ pub enum SelectError {
#[error("The file is not in a supported compression format. Only GZIP and BZIP2 are supported.")]
InvalidCompressionFormat,
#[error("{compression} is not applicable to the queried object. Please correct the request and try again.")]
InvalidCompressionFormatForObject { compression: &'static str },
#[error("The data source type is not valid. Only CSV, JSON, and Parquet are supported.")]
InvalidDataSource,
@@ -87,6 +91,9 @@ pub enum SelectError {
)]
TruncatedInput,
#[error("Scan range queries are not supported on this type of object.")]
UnsupportedScanRangeInput,
#[error("An error occurred while parsing the CSV file. Check the file and try again.")]
CsvParsingError,
@@ -96,6 +103,9 @@ pub enum SelectError {
#[error("An error occurred while parsing the Parquet file. Check the file and try again.")]
ParquetParsingError,
#[error("The length of a record in the input or result is greater than the maxCharsPerRecord limit of 1 MB.")]
OverMaxRecordSize,
#[error("{message}")]
ParseSelectFailure { message: String },
+90 -17
View File
@@ -12,7 +12,11 @@
// See the License for the specific language governing permissions and
// limitations under the License.
use std::sync::atomic::{AtomicU64, Ordering};
use arc_swap::ArcSwap;
use std::sync::{
Arc,
atomic::{AtomicU64, Ordering},
};
#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
pub struct SelectInputMetricsSnapshot {
@@ -20,33 +24,72 @@ pub struct SelectInputMetricsSnapshot {
pub bytes_processed: u64,
}
#[derive(Debug, Default)]
#[derive(Debug)]
pub struct SelectInputMetrics {
active: ArcSwap<SelectInputMetricBank>,
}
#[derive(Debug, Default)]
struct SelectInputMetricBank {
uncompressed_bytes: AtomicU64,
compressed_bytes_scanned: AtomicU64,
compressed_bytes_processed: AtomicU64,
}
#[derive(Clone, Debug)]
pub(crate) struct SelectInputMetricsRecorder {
bank: Arc<SelectInputMetricBank>,
}
impl Default for SelectInputMetrics {
fn default() -> Self {
Self {
active: ArcSwap::from_pointee(SelectInputMetricBank::default()),
}
}
}
impl SelectInputMetrics {
pub fn snapshot(&self) -> SelectInputMetricsSnapshot {
let uncompressed_bytes = self.uncompressed_bytes.load(Ordering::Relaxed);
let bank = self.active.load();
let uncompressed_bytes = bank.uncompressed_bytes.load(Ordering::Relaxed);
SelectInputMetricsSnapshot {
bytes_scanned: uncompressed_bytes,
bytes_processed: uncompressed_bytes,
bytes_scanned: uncompressed_bytes.saturating_add(bank.compressed_bytes_scanned.load(Ordering::Relaxed)),
bytes_processed: uncompressed_bytes.saturating_add(bank.compressed_bytes_processed.load(Ordering::Relaxed)),
}
}
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)));
pub(crate) fn recorder(&self) -> SelectInputMetricsRecorder {
SelectInputMetricsRecorder {
bank: self.active.load_full(),
}
}
/// Clears planner-only reads before query execution begins.
/// Publishes a fresh bank so late planner writes remain isolated.
pub fn reset(&self) {
self.uncompressed_bytes.store(0, Ordering::Relaxed);
self.active.store(Arc::new(SelectInputMetricBank::default()));
}
}
impl SelectInputMetricsRecorder {
pub(crate) fn record_uncompressed(&self, bytes: usize) {
saturating_add(&self.bank.uncompressed_bytes, bytes);
}
pub(crate) fn record_scanned(&self, bytes: usize) {
saturating_add(&self.bank.compressed_bytes_scanned, bytes);
}
pub(crate) fn record_processed(&self, bytes: usize) {
saturating_add(&self.bank.compressed_bytes_processed, bytes);
}
}
fn saturating_add(counter: &AtomicU64, bytes: usize) {
let increment = u64::try_from(bytes).unwrap_or(u64::MAX);
let _ = counter.fetch_update(Ordering::Relaxed, Ordering::Relaxed, |current| Some(current.saturating_add(increment)));
}
#[cfg(test)]
mod tests {
use super::*;
@@ -54,7 +97,7 @@ mod tests {
#[test]
fn records_uncompressed_input_at_both_boundaries() {
let metrics = SelectInputMetrics::default();
metrics.record_uncompressed(7);
metrics.recorder().record_uncompressed(7);
assert_eq!(
metrics.snapshot(),
@@ -68,21 +111,51 @@ mod tests {
#[test]
fn counters_saturate_instead_of_wrapping() {
let metrics = SelectInputMetrics::default();
metrics.uncompressed_bytes.store(u64::MAX - 1, Ordering::Relaxed);
metrics
.active
.load()
.uncompressed_bytes
.store(u64::MAX - 1, Ordering::Relaxed);
metrics.record_uncompressed(2);
metrics.recorder().record_uncompressed(2);
assert_eq!(metrics.snapshot().bytes_scanned, u64::MAX);
assert_eq!(metrics.snapshot().bytes_processed, u64::MAX);
}
#[test]
fn compressed_boundaries_are_counted_independently() {
let metrics = SelectInputMetrics::default();
let recorder = metrics.recorder();
recorder.record_scanned(39);
recorder.record_processed(19);
assert_eq!(
metrics.snapshot(),
SelectInputMetricsSnapshot {
bytes_scanned: 39,
bytes_processed: 19,
}
);
}
#[test]
fn reset_clears_schema_inference_bytes() {
let metrics = SelectInputMetrics::default();
metrics.record_uncompressed(9);
let planning = metrics.recorder();
planning.record_uncompressed(9);
metrics.reset();
planning.record_uncompressed(5);
let execution = metrics.recorder();
execution.record_uncompressed(3);
assert_eq!(metrics.snapshot(), SelectInputMetricsSnapshot::default());
assert_eq!(
metrics.snapshot(),
SelectInputMetricsSnapshot {
bytes_scanned: 3,
bytes_processed: 3,
}
);
}
}
File diff suppressed because it is too large Load Diff
+24
View File
@@ -30,6 +30,7 @@ use datafusion::{
prelude::SessionContext,
};
use parking_lot::Mutex;
use s3s::dto::CompressionType;
use std::sync::{
Arc, Weak,
atomic::{AtomicU8, Ordering},
@@ -446,11 +447,19 @@ impl SessionCtxFactory {
let scan_range_requires_single_file_scan =
context.input.request.scan_range.is_some() && context.input.request.input_serialization.parquet.is_none();
let json_document_requires_single_file_scan = is_json_document_input(&context.input);
let compressed_input_requires_single_file_scan = context
.input
.request
.input_serialization
.compression_type
.as_ref()
.is_some_and(|compression| compression.as_str() != CompressionType::NONE);
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
|| json_document_requires_single_file_scan
|| compressed_input_requires_single_file_scan
|| metered_input_requires_single_file_scan
{
config.with_repartition_file_scans(false)
@@ -847,6 +856,21 @@ mod tests {
assert!(session.inner().config().options().optimizer.repartition_file_scans);
}
#[tokio::test]
async fn compressed_input_disables_file_repartitioning_without_metrics() {
let mut context = test_context();
Arc::make_mut(&mut context.input).request.input_serialization.compression_type =
Some(CompressionType::from_static(CompressionType::GZIP));
let session = SessionCtxFactory::new(true)
.with_target_partitions(2)
.create_session_ctx(&context)
.await
.expect("compressed session should be created");
assert!(!session.inner().config().options().optimizer.repartition_file_scans);
}
#[tokio::test]
async fn json_document_disables_file_repartitioning() {
let mut context = test_context();
+28 -12
View File
@@ -53,7 +53,7 @@ use rustfs_s3select_api::{
},
},
};
use s3s::dto::{FileHeaderInfo, JSONType, SelectObjectContentInput};
use s3s::dto::{CompressionType, FileHeaderInfo, JSONType, SelectObjectContentInput};
use std::sync::LazyLock;
use tokio::{
sync::Semaphore,
@@ -72,6 +72,7 @@ use crate::{
static IGNORE: LazyLock<FileHeaderInfo> = LazyLock::new(|| FileHeaderInfo::from_static(FileHeaderInfo::IGNORE));
static NONE: LazyLock<FileHeaderInfo> = LazyLock::new(|| FileHeaderInfo::from_static(FileHeaderInfo::NONE));
static USE: LazyLock<FileHeaderInfo> = LazyLock::new(|| FileHeaderInfo::from_static(FileHeaderInfo::USE));
const EXACT_OBJECT_FILE_EXTENSION: &str = "";
#[derive(Clone)]
pub struct SimpleQueryDispatcher {
@@ -416,6 +417,13 @@ impl SimpleQueryDispatcher {
let path = format!("s3://{}/{}", self.input.bucket, self.input.key);
let table_path = ListingTableUrl::parse(path)?;
let compressed_input = self
.input
.request
.input_serialization
.compression_type
.as_ref()
.is_some_and(|compression| compression.as_str() != CompressionType::NONE);
let (listing_options, need_rename_volume_name, need_ignore_volume_name) =
if let Some(csv) = self.input.request.input_serialization.csv.as_ref() {
let mut need_rename_volume_name = false;
@@ -465,22 +473,30 @@ impl SimpleQueryDispatcher {
file_format = file_format.with_quote(quote.as_bytes().first().copied().unwrap_or_default());
}
(
ListingOptions::new(Arc::new(file_format)).with_file_extension(".csv"),
ListingOptions::new(Arc::new(file_format)).with_file_extension(if compressed_input {
EXACT_OBJECT_FILE_EXTENSION
} else {
".csv"
}),
need_rename_volume_name,
need_ignore_volume_name,
)
} else if self.input.request.input_serialization.json.is_some() {
let file_format = JsonFormat::default();
// Use the actual file extension from the object key so that files stored
// with a `.jsonl` suffix (newline-delimited JSON) are also matched by
// DataFusion's listing/schema-inference logic. Falling back to ".json"
// preserves behaviour for keys that have no extension.
let file_ext = std::path::Path::new(&self.input.key)
.extension()
.and_then(|e| e.to_str())
.map(|e| format!(".{e}"))
.unwrap_or_else(|| ".json".to_string());
(ListingOptions::new(Arc::new(file_format)).with_file_extension(file_ext), false, false)
let file_extension = if compressed_input {
EXACT_OBJECT_FILE_EXTENSION.to_string()
} else {
std::path::Path::new(&self.input.key)
.extension()
.and_then(|extension| extension.to_str())
.map(|extension| format!(".{extension}"))
.unwrap_or_else(|| ".json".to_string())
};
(
ListingOptions::new(Arc::new(file_format)).with_file_extension(file_extension),
false,
false,
)
} else {
return Err(SelectError::InvalidDataSource.into());
};
+138 -7
View File
@@ -22,7 +22,7 @@ use futures::StreamExt;
use http::{HeaderMap, StatusCode, header::RANGE};
use rustfs_s3select_api::{
QueryError, SelectError, SelectInputMetrics,
object_store::{INVALID_SCAN_RANGE_MESSAGE, validate_scan_range_bounds},
object_store::{INVALID_SCAN_RANGE_MESSAGE, is_noop_scan_range, validate_scan_range_bounds},
query::{Context, Query},
};
use rustfs_s3select_query::instance::s3_select_query_timeout;
@@ -44,6 +44,8 @@ const RECORDS_CHUNK_TARGET: usize = 128 * 1024;
const DATA_SOURCE_PATH_UNSUPPORTED_CODE: &str = "DataSourcePathUnsupported";
const INVALID_QUERY_CODE: &str = "InvalidQuery";
const PARSE_SELECT_FAILURE_CODE: &str = "ParseSelectFailure";
const INVALID_REQUEST_PARAMETER_MESSAGE: &str =
"The value of a parameter in the SelectRequest element is invalid. Check the service API documentation and try again.";
const BUSY_MESSAGE: &str = "The service is unavailable. Try again later.";
const EMPTY_SELECT_EXPRESSION_MESSAGE: &str = "empty SQL expression";
const SLOW_DOWN_MESSAGE: &str = "Reduce your request rate.";
@@ -103,7 +105,11 @@ pub async fn execute_select_object_content(
)
.await
.map_err(|_| select_query_timeout_error(query_timeout.as_secs()))??;
validate_scan_range_for_object_size(&input.request, snapshot.logical_size())?;
let object_size = snapshot.logical_size();
validate_scan_range_for_object_size(&input.request, object_size)?;
if object_size == 0 && is_compressed_input(&input.request.input_serialization) {
return Err(map_select_error_to_s3(&SelectError::TruncatedInput));
}
let snapshot = Arc::new(snapshot);
let query =
Query::new_with_snapshot(Context { input: input.clone() }, input.request.expression.clone(), Arc::clone(&snapshot));
@@ -262,7 +268,14 @@ fn validate_select_request(headers: &http::HeaderMap, input: &mut SelectObjectCo
}
normalize_input_serialization(&mut input.request.input_serialization)?;
let compressed_input = is_compressed_input(&input.request.input_serialization);
if compressed_input && input.request.scan_range.as_ref().is_some_and(is_noop_scan_range) {
input.request.scan_range = None;
}
validate_scan_range(&input.request)?;
if compressed_input && input.request.scan_range.is_some() {
return Err(map_select_error_to_s3(&SelectError::UnsupportedScanRangeInput));
}
let output_format = normalize_output_serialization(&mut input.request.output_serialization)?;
if input.request.expression.trim().is_empty() {
@@ -284,6 +297,13 @@ fn validate_select_request(headers: &http::HeaderMap, input: &mut SelectObjectCo
})
}
fn is_compressed_input(input: &InputSerialization) -> bool {
input
.compression_type
.as_ref()
.is_some_and(|compression| compression.as_str() != CompressionType::NONE)
}
fn normalize_input_serialization(input: &mut InputSerialization) -> S3Result<()> {
let format_count =
usize::from(input.csv.is_some()) + usize::from(input.json.is_some()) + usize::from(input.parquet.is_some());
@@ -298,15 +318,19 @@ fn normalize_input_serialization(input: &mut InputSerialization) -> S3Result<()>
match compression.as_str() {
CompressionType::NONE => {}
CompressionType::GZIP | CompressionType::BZIP2 => {
return Err(s3_error!(
NotImplemented,
"SelectObjectContent currently supports only uncompressed input"
));
if input.parquet.is_some() {
return Err(S3Error::with_message(
S3ErrorCode::InvalidRequestParameter,
INVALID_REQUEST_PARAMETER_MESSAGE,
));
}
}
_ => return Err(map_select_error_to_s3(&SelectError::InvalidCompressionFormat)),
}
}
input.compression_type = Some(CompressionType::from_static(CompressionType::NONE));
input
.compression_type
.get_or_insert_with(|| CompressionType::from_static(CompressionType::NONE));
if let Some(csv) = input.csv.as_mut() {
if csv.allow_quoted_record_delimiter.unwrap_or(false) {
@@ -667,11 +691,16 @@ fn map_query_error_to_s3(err: QueryError) -> S3Error {
fn map_select_error_to_s3(err: &SelectError) -> S3Error {
match err {
SelectError::InvalidCompressionFormat => S3Error::with_message(S3ErrorCode::InvalidCompressionFormat, err.to_string()),
SelectError::InvalidCompressionFormatForObject { .. } => {
S3Error::with_message(S3ErrorCode::InvalidCompressionFormat, err.to_string())
}
SelectError::InvalidDataSource => S3Error::with_message(S3ErrorCode::InvalidDataSource, err.to_string()),
SelectError::TruncatedInput => S3Error::with_message(S3ErrorCode::TruncatedInput, err.to_string()),
SelectError::UnsupportedScanRangeInput => S3Error::with_message(S3ErrorCode::UnsupportedScanRangeInput, err.to_string()),
SelectError::CsvParsingError => S3Error::with_message(S3ErrorCode::CSVParsingError, err.to_string()),
SelectError::JsonParsingError => S3Error::with_message(S3ErrorCode::JSONParsingError, err.to_string()),
SelectError::ParquetParsingError => S3Error::with_message(S3ErrorCode::ParquetParsingError, err.to_string()),
SelectError::OverMaxRecordSize => S3Error::with_message(S3ErrorCode::OverMaxRecordSize, err.to_string()),
SelectError::ParseSelectFailure { message } => custom_bad_request(PARSE_SELECT_FAILURE_CODE, message.clone()),
SelectError::InvalidQuery => custom_bad_request(INVALID_QUERY_CODE, err.to_string()),
SelectError::InvalidDataType => S3Error::with_message(S3ErrorCode::InvalidDataType, err.to_string()),
@@ -995,8 +1024,20 @@ mod tests {
S3ErrorCode::InvalidCompressionFormat,
StatusCode::BAD_REQUEST,
),
(
SelectError::InvalidCompressionFormatForObject {
compression: CompressionType::GZIP,
},
S3ErrorCode::InvalidCompressionFormat,
StatusCode::BAD_REQUEST,
),
(SelectError::InvalidDataSource, S3ErrorCode::InvalidDataSource, StatusCode::BAD_REQUEST),
(SelectError::TruncatedInput, S3ErrorCode::TruncatedInput, StatusCode::BAD_REQUEST),
(
SelectError::UnsupportedScanRangeInput,
S3ErrorCode::UnsupportedScanRangeInput,
StatusCode::BAD_REQUEST,
),
(SelectError::CsvParsingError, S3ErrorCode::CSVParsingError, StatusCode::BAD_REQUEST),
(SelectError::JsonParsingError, S3ErrorCode::JSONParsingError, StatusCode::BAD_REQUEST),
(
@@ -1004,6 +1045,7 @@ mod tests {
S3ErrorCode::ParquetParsingError,
StatusCode::BAD_REQUEST,
),
(SelectError::OverMaxRecordSize, S3ErrorCode::OverMaxRecordSize, StatusCode::BAD_REQUEST),
(
SelectError::ParseSelectFailure {
message: "invalid SELECT expression".to_string(),
@@ -1372,6 +1414,12 @@ mod tests {
assert_eq!(compression_status, StatusCode::BAD_REQUEST);
assert!(compression_body.contains("<Code>InvalidCompressionFormat</Code>"));
assert!(compression_body.contains("<Message>"));
let scan_range_error = map_select_error_to_s3(&SelectError::UnsupportedScanRangeInput);
let (scan_range_status, scan_range_body) = http_xml_error(scan_range_error).await;
assert_eq!(scan_range_status, StatusCode::BAD_REQUEST);
assert!(scan_range_body.contains("<Code>UnsupportedScanRangeInput</Code>"));
assert!(scan_range_body.contains("<Message>Scan range queries are not supported on this type of object.</Message>"));
}
#[tokio::test(start_paused = true)]
@@ -1625,6 +1673,89 @@ mod tests {
);
}
#[test]
fn validate_preserves_supported_compression_for_csv_and_json_lines() {
for compression in [CompressionType::GZIP, CompressionType::BZIP2] {
let mut csv_input = base_input();
csv_input.request.input_serialization.compression_type = Some(CompressionType::from_static(compression));
validate_select_request(&HeaderMap::new(), &mut csv_input).expect("compressed CSV should be accepted");
assert_eq!(
csv_input
.request
.input_serialization
.compression_type
.as_ref()
.map(|value| value.as_str()),
Some(compression)
);
let mut json_input = base_input();
json_input.request.input_serialization.csv = None;
json_input.request.input_serialization.json = Some(JSONInput {
type_: Some(JSONType::from_static(JSONType::LINES)),
});
json_input.request.input_serialization.compression_type = Some(CompressionType::from_static(compression));
validate_select_request(&HeaderMap::new(), &mut json_input).expect("compressed JSON LINES should be accepted");
assert_eq!(
json_input
.request
.input_serialization
.compression_type
.as_ref()
.map(|value| value.as_str()),
Some(compression)
);
}
}
#[test]
fn validate_rejects_parquet_compression_with_select_request_error() {
let mut input = base_input();
input.request.input_serialization.csv = None;
input.request.input_serialization.parquet = Some(ParquetInput {});
input.request.input_serialization.compression_type = Some(CompressionType::from_static(CompressionType::GZIP));
let error = validate_select_request(&HeaderMap::new(), &mut input).expect_err("compressed Parquet must fail");
assert_eq!(error.code(), &S3ErrorCode::InvalidRequestParameter);
assert_eq!(error.message(), Some(INVALID_REQUEST_PARAMETER_MESSAGE));
}
#[test]
fn validate_normalizes_noop_compressed_scan_range_and_rejects_real_ranges() {
let mut noop = base_input();
noop.request.input_serialization.compression_type = Some(CompressionType::from_static(CompressionType::GZIP));
noop.request.scan_range = Some(ScanRange {
start: Some(0),
end: None,
});
validate_select_request(&HeaderMap::new(), &mut noop).expect("zero-start full scan should be normalized");
assert!(noop.request.scan_range.is_none());
let mut ranged = base_input();
ranged.request.input_serialization.compression_type = Some(CompressionType::from_static(CompressionType::GZIP));
ranged.request.scan_range = Some(ScanRange {
start: Some(1),
end: None,
});
let error = validate_select_request(&HeaderMap::new(), &mut ranged)
.expect_err("compressed input with an effective ScanRange must fail before object I/O");
assert_eq!(error.code(), &S3ErrorCode::UnsupportedScanRangeInput);
assert_eq!(error.status_code(), Some(StatusCode::BAD_REQUEST));
assert_eq!(error.message(), Some("Scan range queries are not supported on this type of object."));
let mut malformed = base_input();
malformed.request.input_serialization.compression_type = Some(CompressionType::from_static(CompressionType::GZIP));
malformed.request.scan_range = Some(ScanRange {
start: Some(10),
end: Some(1),
});
let error = validate_select_request(&HeaderMap::new(), &mut malformed)
.expect_err("malformed ScanRange must fail before compression compatibility validation");
assert_eq!(error.code(), &S3ErrorCode::InvalidRequestParameter);
assert_eq!(error.message(), Some(INVALID_SCAN_RANGE_MESSAGE));
}
#[test]
fn validate_rejects_unknown_csv_header_mode_before_streaming() {
let mut input = base_input();