diff --git a/crates/s3select-api/src/object_store.rs b/crates/s3select-api/src/object_store.rs index 108e3e89b..2caccc484 100644 --- a/crates/s3select-api/src/object_store.rs +++ b/crates/s3select-api/src/object_store.rs @@ -1427,11 +1427,7 @@ mod test { #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[serial_test::serial] async fn test_get_opts_validates_raw_length_before_delimiter_conversion() { - let temp_root = tempfile::tempdir().expect("create s3select test temp root"); - let env = rustfs_test_utils::TestECStoreEnv::builder() - .base_dir(temp_root.path()) - .build() - .await; + let env = crate::storage_api::select_test_ecstore_env().await; let bucket = "s3select-multi-byte-delimiter"; let object = "input.csv"; let input_bytes = b"a&&1\n"; @@ -1475,7 +1471,7 @@ mod test { json_sub_path: None, memory_pool: Arc::new(GreedyMemoryPool::new(1024)), query_tracker: None, - store: env.ecstore, + store: Arc::clone(&env.ecstore), }; let result = store diff --git a/crates/s3select-api/src/query/session.rs b/crates/s3select-api/src/query/session.rs index 10c7e5f15..7c0208206 100644 --- a/crates/s3select-api/src/query/session.rs +++ b/crates/s3select-api/src/query/session.rs @@ -349,6 +349,13 @@ impl SessionCtxFactory { }; let rt = RuntimeEnvBuilder::new().with_memory_limit(memory_limit_bytes, 1.0).build()?; let config = SessionConfig::new().with_target_partitions(self.target_partitions); + 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 scan_range_requires_single_file_scan { + config.with_repartition_file_scans(false) + } else { + config + }; let memory_pool = Arc::clone(&rt.memory_pool); let df_session_state = SessionStateBuilder::new() .with_config(config) @@ -474,11 +481,20 @@ fn test_parquet_batch( #[cfg(test)] mod tests { use super::*; - use datafusion::execution::memory_pool::MemoryLimit; - use s3s::dto::{ - CSVInput, CSVOutput, ExpressionType, InputSerialization, OutputSerialization, SelectObjectContentInput, - SelectObjectContentRequest, + use crate::storage_api::SelectPutObjReader; + use crate::storage_api::object_store::ObjectIO as _; + use datafusion::{ + datasource::{ + file_format::csv::CsvFormat, + listing::{ListingOptions, ListingTable, ListingTableConfig, ListingTableUrl}, + }, + execution::memory_pool::MemoryLimit, }; + use s3s::dto::{ + CSVInput, CSVOutput, ExpressionType, InputSerialization, JSONInput, OutputSerialization, ParquetInput, ScanRange, + SelectObjectContentInput, SelectObjectContentRequest, + }; + use std::io::Write as _; fn test_context() -> Context { Context { @@ -527,6 +543,47 @@ mod tests { .expect("session should be created with configured target partitions"); assert_eq!(session.inner().config().target_partitions(), 3); + assert!(session.inner().config().options().optimizer.repartition_file_scans); + } + + #[tokio::test] + async fn parquet_scan_range_keeps_file_repartitioning() { + let mut context = test_context(); + let request = &mut Arc::make_mut(&mut context.input).request; + request.scan_range = Some(ScanRange { + start: Some(0), + end: Some(0), + }); + request.input_serialization.csv = None; + request.input_serialization.parquet = Some(ParquetInput {}); + + let session = SessionCtxFactory::new(true) + .with_target_partitions(2) + .create_session_ctx(&context) + .await + .expect("Parquet ScanRange session should be created"); + + assert!(session.inner().config().options().optimizer.repartition_file_scans); + } + + #[tokio::test] + async fn json_lines_scan_range_disables_file_repartitioning() { + let mut context = test_context(); + let request = &mut Arc::make_mut(&mut context.input).request; + request.scan_range = Some(ScanRange { + start: Some(0), + end: Some(0), + }); + request.input_serialization.csv = None; + request.input_serialization.json = Some(JSONInput::default()); + + let session = SessionCtxFactory::new(true) + .with_target_partitions(2) + .create_session_ctx(&context) + .await + .expect("JSON LINES ScanRange session should be created"); + + assert!(!session.inner().config().options().optimizer.repartition_file_scans); } #[tokio::test] @@ -557,12 +614,7 @@ mod tests { #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[serial_test::serial] async fn session_factory_propagates_query_guard_to_ec_store() { - let temp_root = tempfile::tempdir().expect("create session test temp root"); - let env = rustfs_test_utils::TestECStoreEnv::builder() - .base_dir(temp_root.path()) - .init_bucket_metadata(false) - .build() - .await; + let env = crate::storage_api::select_test_ecstore_env().await; let admission = Arc::new(tokio::sync::Semaphore::new(1)); let permit = Arc::clone(&admission) @@ -586,6 +638,94 @@ mod tests { assert_eq!(Arc::strong_count(&query_guard), 1); } + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + #[serial_test::serial] + async fn scan_range_is_preserved_across_large_csv_partition_boundary() { + const ROW_COUNT: usize = 200_000; + const ROW_WIDTH: usize = 7; + const SELECTED_ROW: usize = 100_000; + + let env = crate::storage_api::select_test_ecstore_env().await; + let mut data = Vec::with_capacity(ROW_COUNT * ROW_WIDTH); + for row in 0..ROW_COUNT { + writeln!(&mut data, "{row:06}").expect("write fixed-width CSV test row"); + } + assert!(data.len() > 1024 * 1024); + + let mut context = test_context(); + let selected_start = i64::try_from(SELECTED_ROW * ROW_WIDTH).expect("selected row offset should fit in i64"); + Arc::make_mut(&mut context.input).request.scan_range = Some(ScanRange { + start: Some(selected_start), + end: Some(selected_start), + }); + env.make_bucket(&context.input.bucket, false).await; + let mut reader = SelectPutObjReader::from_vec(data); + env.ecstore + .put_object(&context.input.bucket, &context.input.key, &mut reader, &Default::default()) + .await + .expect("put large ScanRange CSV fixture"); + + let admission = Arc::new(tokio::sync::Semaphore::new(1)); + let permit = Arc::clone(&admission) + .acquire_owned() + .await + .expect("query permit should be available"); + let query_tracker = QueryExecutionTracker::new( + &QueryExecutionOwner::new(), + Arc::new(permit), + Instant::now() + std::time::Duration::from_secs(300), + 300, + ); + let session = SessionCtxFactory::new(false) + .with_target_partitions(2) + .create_session_ctx_with_tracker_and_store(&context, query_tracker, Arc::clone(&env.ecstore)) + .await + .expect("create production ScanRange session"); + assert!(!session.inner().config().options().optimizer.repartition_file_scans); + + let table_path = ListingTableUrl::parse(format!("s3://{}/{}", context.input.bucket, context.input.key)) + .expect("parse ScanRange table URL"); + let listing_options = + ListingOptions::new(Arc::new(CsvFormat::default().with_schema_infer_max_rec(0).with_has_header(false))) + .with_file_extension(".csv"); + let schema = listing_options + .infer_schema(session.inner(), &table_path) + .await + .expect("infer ScanRange CSV schema"); + let table = ListingTable::try_new( + ListingTableConfig::new(table_path) + .with_listing_options(listing_options) + .with_schema(schema), + ) + .expect("build ScanRange listing table"); + let query_context = SessionContext::new_with_state(session.inner().clone()); + query_context + .register_table("scan_input", Arc::new(table)) + .expect("register ScanRange listing table"); + + let batches = query_context + .sql("SELECT * FROM scan_input") + .await + .expect("plan ScanRange CSV query") + .collect() + .await + .expect("execute ScanRange CSV query"); + let mut values = Vec::new(); + for batch in batches { + let column = batch + .column(0) + .as_any() + .downcast_ref::() + .expect("ScanRange CSV column should be Utf8"); + for row in 0..batch.num_rows() { + values.push(column.value(row).to_string()); + } + } + + assert_eq!(values.len(), 1); + assert_eq!(values, [format!("{SELECTED_ROW:06}")]); + } + #[tokio::test] async fn elapsed_deadline_does_not_yield_query_guard_before_timer_poll() { let admission = Arc::new(tokio::sync::Semaphore::new(1)); diff --git a/crates/s3select-api/src/storage_api.rs b/crates/s3select-api/src/storage_api.rs index eba20036b..bb9e9ed66 100644 --- a/crates/s3select-api/src/storage_api.rs +++ b/crates/s3select-api/src/storage_api.rs @@ -43,6 +43,13 @@ pub(crate) type SelectGetObjectReader = ::ObjectInfo; pub(crate) type SelectObjectOptions = ::ObjectOptions; +#[cfg(test)] +pub(crate) async fn select_test_ecstore_env() -> &'static rustfs_test_utils::TestECStoreEnv { + static ENV: tokio::sync::OnceCell = tokio::sync::OnceCell::const_new(); + ENV.get_or_init(|| async { rustfs_test_utils::TestECStoreEnv::builder().build().await }) + .await +} + pub(crate) fn resolve_select_object_store_handle() -> Option> { resolve_select_object_store_handle_from_backend() }