fix(s3select): configure session partitions (#4511)

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<SessionCtx> {
         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());
+    }
+}

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<D: QueryDispatcher> {
     // 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::<usize>().ok()).unwrap_or(0)
+}
+
 pub async fn make_rustfsms(input: Arc<SelectObjectContentInput>, is_test: bool) -> QueryResult<impl DatabaseManagerSystem> {
     // 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<SqlQueryExecutionFactory>,
     default_table_provider: Arc<BaseTableProvider>,
 ) -> QueryResult<impl DatabaseManagerSystem> {
-    // 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() {