From 765e719aafad0d64706254695ab573754b6e78dc Mon Sep 17 00:00:00 2001 From: Zhengchao An Date: Thu, 9 Jul 2026 01:07:01 +0800 Subject: [PATCH] fix(s3select): configure session partitions (#4511) --- crates/s3select-api/src/query/session.rs | 75 +++++++++++++++++++++++- crates/s3select-query/src/instance.rs | 45 ++++++++++++-- 2 files changed, 114 insertions(+), 6 deletions(-) diff --git a/crates/s3select-api/src/query/session.rs b/crates/s3select-api/src/query/session.rs index 1a354d9f5..739a2ef23 100644 --- a/crates/s3select-api/src/query/session.rs +++ b/crates/s3select-api/src/query/session.rs @@ -20,7 +20,7 @@ use datafusion::{ datatypes::{DataType, Field, Schema}, record_batch::RecordBatch, }, - execution::{SessionStateBuilder, context::SessionState, runtime_env::RuntimeEnvBuilder}, + execution::{SessionStateBuilder, config::SessionConfig, context::SessionState, runtime_env::RuntimeEnvBuilder}, object_store::{ObjectStore, ObjectStoreExt, memory::InMemory, path::Path}, parquet::arrow::ArrowWriter, prelude::SessionContext, @@ -48,9 +48,22 @@ pub struct SessionCtxDesc { #[derive(Default)] pub struct SessionCtxFactory { pub is_test: bool, + pub target_partitions: usize, } impl SessionCtxFactory { + pub fn new(is_test: bool) -> Self { + Self { + is_test, + target_partitions: 0, + } + } + + pub fn with_target_partitions(mut self, target_partitions: usize) -> Self { + self.target_partitions = target_partitions; + self + } + pub async fn create_session_ctx(&self, context: &Context) -> QueryResult { let df_session_ctx = self.build_df_session_context(context).await?; @@ -64,7 +77,9 @@ impl SessionCtxFactory { let path = format!("s3://{}", context.input.bucket); let store_url = url::Url::parse(&path).unwrap(); let rt = RuntimeEnvBuilder::new().build()?; + let config = SessionConfig::new().with_target_partitions(self.target_partitions); let df_session_state = SessionStateBuilder::new() + .with_config(config) .with_runtime_env(Arc::new(rt)) .with_default_features(); @@ -178,3 +193,61 @@ fn test_parquet_batch( ) .map_err(|e| QueryError::StoreError { e: e.to_string() }) } + +#[cfg(test)] +mod tests { + use super::*; + use s3s::dto::{ + CSVInput, CSVOutput, ExpressionType, InputSerialization, OutputSerialization, SelectObjectContentInput, + SelectObjectContentRequest, + }; + + fn test_context() -> Context { + Context { + input: Arc::new(SelectObjectContentInput { + bucket: "test-bucket".to_string(), + expected_bucket_owner: None, + key: "test.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("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, + }, + }), + } + } + + #[tokio::test] + async fn session_factory_applies_target_partitions() { + let factory = SessionCtxFactory::new(true).with_target_partitions(3); + let session = factory + .create_session_ctx(&test_context()) + .await + .expect("session should be created with configured target partitions"); + + assert_eq!(session.inner().config().target_partitions(), 3); + } + + #[tokio::test] + async fn session_factory_zero_target_partitions_uses_datafusion_default() { + let factory = SessionCtxFactory::new(true); + let session = factory + .create_session_ctx(&test_context()) + .await + .expect("session should be created with default target partitions"); + + assert_eq!(session.inner().config().target_partitions(), SessionConfig::new().target_partitions()); + } +} diff --git a/crates/s3select-query/src/instance.rs b/crates/s3select-query/src/instance.rs index c041bca7a..950d31365 100644 --- a/crates/s3select-query/src/instance.rs +++ b/crates/s3select-query/src/instance.rs @@ -33,6 +33,8 @@ use crate::{ sql::{optimizer::CascadeOptimizerBuilder, parser::DefaultParser}, }; +const ENV_RUSTFS_S3SELECT_TARGET_PARTITIONS: &str = "RUSTFS_S3SELECT_TARGET_PARTITIONS"; + #[derive(Builder)] pub struct RustFSms { // query dispatcher & query execution @@ -77,14 +79,32 @@ where } } +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +struct S3SelectRuntimeConfig { + target_partitions: usize, +} + +impl S3SelectRuntimeConfig { + fn from_env() -> Self { + Self { + target_partitions: target_partitions_from_env_value( + std::env::var(ENV_RUSTFS_S3SELECT_TARGET_PARTITIONS).ok().as_deref(), + ), + } + } +} + +fn target_partitions_from_env_value(value: Option<&str>) -> usize { + value.and_then(|value| value.parse::().ok()).unwrap_or(0) +} + pub async fn make_rustfsms(input: Arc, is_test: bool) -> QueryResult { // init Function Manager, we can define some UDF if need let func_manager = SimpleFunctionMetadataManager::default(); - // TODO session config need load global system config - let session_factory = Arc::new(SessionCtxFactory { is_test }); + let runtime_config = S3SelectRuntimeConfig::from_env(); + let session_factory = Arc::new(SessionCtxFactory::new(is_test).with_target_partitions(runtime_config.target_partitions)); let parser = Arc::new(DefaultParser::default()); let optimizer = Arc::new(CascadeOptimizerBuilder::default().build()); - // TODO wrap, and num_threads configurable let scheduler = Arc::new(LocalScheduler {}); let query_execution_factory = Arc::new(SqlQueryExecutionFactory::new(optimizer, scheduler)); @@ -115,8 +135,8 @@ pub async fn make_rustfsms_with_components( query_execution_factory: Arc, default_table_provider: Arc, ) -> QueryResult { - // TODO session config need load global system config - let session_factory = Arc::new(SessionCtxFactory { is_test }); + let runtime_config = S3SelectRuntimeConfig::from_env(); + let session_factory = Arc::new(SessionCtxFactory::new(is_test).with_target_partitions(runtime_config.target_partitions)); let query_dispatcher = SimpleQueryDispatcherBuilder::default() .with_input(input) @@ -147,6 +167,21 @@ mod tests { use crate::get_global_db; + use super::{S3SelectRuntimeConfig, target_partitions_from_env_value}; + + #[test] + fn parses_target_partitions_from_env_value() { + assert_eq!(target_partitions_from_env_value(Some("4")), 4); + assert_eq!(target_partitions_from_env_value(Some("0")), 0); + assert_eq!(target_partitions_from_env_value(Some("not-a-number")), 0); + assert_eq!(target_partitions_from_env_value(None), 0); + } + + #[test] + fn default_runtime_config_uses_datafusion_default_partitions() { + assert_eq!(S3SelectRuntimeConfig::default().target_partitions, 0); + } + #[tokio::test] #[ignore] async fn test_simple_sql() {