refactor: Restructure project layout and clean up dependencies (#30)

This commit introduces a significant reorganization of the project structure to improve maintainability and clarity.

Key changes include:
- Adjusted the directory layout for a more logical module organization.
- Removed unused crate dependencies, reducing the overall project size and potentially speeding up build times.
- Updated import paths and configuration files to reflect the structural changes.

crates/s3select-api/src/query/analyzer.rs

@@ -0,0 +1,26 @@
+// 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::Arc;
+
+use datafusion::logical_expr::LogicalPlan;
+
+use super::session::SessionCtx;
+use crate::QueryResult;
+
+pub type AnalyzerRef = Arc<dyn Analyzer + Send + Sync>;
+
+pub trait Analyzer {
+    fn analyze(&self, plan: &LogicalPlan, session: &SessionCtx) -> QueryResult<LogicalPlan>;
+}

crates/s3select-api/src/query/ast.rs

@@ -0,0 +1,22 @@
+// 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 datafusion::sql::sqlparser::ast::Statement;
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub enum ExtStatement {
+    /// ANSI SQL AST node
+    SqlStatement(Box<Statement>),
+    // we can expand command
+}

crates/s3select-api/src/query/datasource/mod.rs

@@ -0,0 +1,13 @@
+// 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.

crates/s3select-api/src/query/dispatcher.rs

@@ -0,0 +1,46 @@
+// 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::Arc;
+
+use async_trait::async_trait;
+
+use crate::QueryResult;
+
+use super::{
+    Query,
+    execution::{Output, QueryStateMachine},
+    logical_planner::Plan,
+};
+
+#[async_trait]
+pub trait QueryDispatcher: Send + Sync {
+    // fn create_query_id(&self) -> QueryId;
+
+    // fn query_info(&self, id: &QueryId);
+
+    async fn execute_query(&self, query: &Query) -> QueryResult<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>;
+
+    async fn build_query_state_machine(&self, query: Query) -> QueryResult<Arc<QueryStateMachine>>;
+
+    // fn running_query_infos(&self) -> Vec<QueryInfo>;
+
+    // fn running_query_status(&self) -> Vec<QueryStatus>;
+
+    // fn cancel_query(&self, id: &QueryId);
+}

crates/s3select-api/src/query/execution.rs

@@ -0,0 +1,255 @@
+// 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::fmt::Display;
+use std::pin::Pin;
+use std::sync::Arc;
+use std::sync::atomic::{AtomicPtr, Ordering};
+use std::task::{Context, Poll};
+use std::time::{Duration, Instant};
+
+use async_trait::async_trait;
+use datafusion::arrow::datatypes::{Schema, SchemaRef};
+use datafusion::arrow::record_batch::RecordBatch;
+use datafusion::physical_plan::SendableRecordBatchStream;
+use futures::{Stream, StreamExt, TryStreamExt};
+
+use crate::{QueryError, QueryResult};
+
+use super::Query;
+use super::logical_planner::Plan;
+use super::session::SessionCtx;
+
+pub type QueryExecutionRef = Arc<dyn QueryExecution>;
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum QueryType {
+    Batch,
+    Stream,
+}
+
+impl Display for QueryType {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            Self::Batch => write!(f, "batch"),
+            Self::Stream => write!(f, "stream"),
+        }
+    }
+}
+
+#[async_trait]
+pub trait QueryExecution: Send + Sync {
+    fn query_type(&self) -> QueryType {
+        QueryType::Batch
+    }
+    // 开始
+    async fn start(&self) -> QueryResult<Output>;
+    // 停止
+    fn cancel(&self) -> QueryResult<()>;
+}
+
+pub enum Output {
+    StreamData(SendableRecordBatchStream),
+    Nil(()),
+}
+
+impl Output {
+    pub fn schema(&self) -> SchemaRef {
+        match self {
+            Self::StreamData(stream) => stream.schema(),
+            Self::Nil(_) => Arc::new(Schema::empty()),
+        }
+    }
+
+    pub async fn chunk_result(self) -> QueryResult<Vec<RecordBatch>> {
+        match self {
+            Self::Nil(_) => Ok(vec![]),
+            Self::StreamData(stream) => {
+                let schema = stream.schema();
+                let mut res: Vec<RecordBatch> = stream.try_collect::<Vec<RecordBatch>>().await?;
+                if res.is_empty() {
+                    res.push(RecordBatch::new_empty(schema));
+                }
+                Ok(res)
+            }
+        }
+    }
+
+    pub async fn num_rows(self) -> usize {
+        match self.chunk_result().await {
+            Ok(rb) => rb.iter().map(|e| e.num_rows()).sum(),
+            Err(_) => 0,
+        }
+    }
+
+    /// Returns the number of records affected by the query operation
+    ///
+    /// If it is a select statement, returns the number of rows in the result set
+    ///
+    /// -1 means unknown
+    ///
+    /// panic! when StreamData's number of records greater than i64::Max
+    pub async fn affected_rows(self) -> i64 {
+        self.num_rows().await as i64
+    }
+}
+
+impl Stream for Output {
+    type Item = Result<RecordBatch, QueryError>;
+
+    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        let this = self.get_mut();
+        match this {
+            Output::StreamData(stream) => stream.poll_next_unpin(cx).map_err(|e| e.into()),
+            Output::Nil(_) => Poll::Ready(None),
+        }
+    }
+}
+
+#[async_trait]
+pub trait QueryExecutionFactory {
+    async fn create_query_execution(
+        &self,
+        plan: Plan,
+        query_state_machine: QueryStateMachineRef,
+    ) -> QueryResult<QueryExecutionRef>;
+}
+
+pub type QueryStateMachineRef = Arc<QueryStateMachine>;
+
+pub struct QueryStateMachine {
+    pub session: SessionCtx,
+    pub query: Query,
+
+    state: AtomicPtr<QueryState>,
+    start: Instant,
+}
+
+impl QueryStateMachine {
+    pub fn begin(query: Query, session: SessionCtx) -> Self {
+        Self {
+            session,
+            query,
+            state: AtomicPtr::new(Box::into_raw(Box::new(QueryState::ACCEPTING))),
+            start: Instant::now(),
+        }
+    }
+
+    pub fn begin_analyze(&self) {
+        // TODO record time
+        self.translate_to(Box::new(QueryState::RUNNING(RUNNING::ANALYZING)));
+    }
+
+    pub fn end_analyze(&self) {
+        // TODO record time
+    }
+
+    pub fn begin_optimize(&self) {
+        // TODO record time
+        self.translate_to(Box::new(QueryState::RUNNING(RUNNING::OPTMIZING)));
+    }
+
+    pub fn end_optimize(&self) {
+        // TODO
+    }
+
+    pub fn begin_schedule(&self) {
+        // TODO
+        self.translate_to(Box::new(QueryState::RUNNING(RUNNING::SCHEDULING)));
+    }
+
+    pub fn end_schedule(&self) {
+        // TODO
+    }
+
+    pub fn finish(&self) {
+        // TODO
+        self.translate_to(Box::new(QueryState::DONE(DONE::FINISHED)));
+    }
+
+    pub fn cancel(&self) {
+        // TODO
+        self.translate_to(Box::new(QueryState::DONE(DONE::CANCELLED)));
+    }
+
+    pub fn fail(&self) {
+        // TODO
+        self.translate_to(Box::new(QueryState::DONE(DONE::FAILED)));
+    }
+
+    pub fn state(&self) -> &QueryState {
+        unsafe { &*self.state.load(Ordering::Relaxed) }
+    }
+
+    pub fn duration(&self) -> Duration {
+        self.start.elapsed()
+    }
+
+    fn translate_to(&self, state: Box<QueryState>) {
+        self.state.store(Box::into_raw(state), Ordering::Relaxed);
+    }
+}
+
+#[derive(Debug, Clone)]
+pub enum QueryState {
+    ACCEPTING,
+    RUNNING(RUNNING),
+    DONE(DONE),
+}
+
+impl AsRef<str> for QueryState {
+    fn as_ref(&self) -> &str {
+        match self {
+            QueryState::ACCEPTING => "ACCEPTING",
+            QueryState::RUNNING(e) => e.as_ref(),
+            QueryState::DONE(e) => e.as_ref(),
+        }
+    }
+}
+
+#[derive(Debug, Clone)]
+pub enum RUNNING {
+    DISPATCHING,
+    ANALYZING,
+    OPTMIZING,
+    SCHEDULING,
+}
+
+impl AsRef<str> for RUNNING {
+    fn as_ref(&self) -> &str {
+        match self {
+            Self::DISPATCHING => "DISPATCHING",
+            Self::ANALYZING => "ANALYZING",
+            Self::OPTMIZING => "OPTMIZING",
+            Self::SCHEDULING => "SCHEDULING",
+        }
+    }
+}
+
+#[derive(Debug, Clone)]
+pub enum DONE {
+    FINISHED,
+    FAILED,
+    CANCELLED,
+}
+
+impl AsRef<str> for DONE {
+    fn as_ref(&self) -> &str {
+        match self {
+            Self::FINISHED => "FINISHED",
+            Self::FAILED => "FAILED",
+            Self::CANCELLED => "CANCELLED",
+        }
+    }
+}

crates/s3select-api/src/query/function.rs

@@ -0,0 +1,38 @@
+// 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::Arc;
+
+use datafusion::logical_expr::{AggregateUDF, ScalarUDF, WindowUDF};
+
+use crate::QueryResult;
+
+pub type FuncMetaManagerRef = Arc<dyn FunctionMetadataManager + Send + Sync>;
+pub trait FunctionMetadataManager {
+    fn register_udf(&mut self, udf: Arc<ScalarUDF>) -> QueryResult<()>;
+
+    fn register_udaf(&mut self, udaf: Arc<AggregateUDF>) -> QueryResult<()>;
+
+    fn register_udwf(&mut self, udwf: Arc<WindowUDF>) -> QueryResult<()>;
+
+    fn udf(&self, name: &str) -> QueryResult<Arc<ScalarUDF>>;
+
+    fn udaf(&self, name: &str) -> QueryResult<Arc<AggregateUDF>>;
+
+    fn udwf(&self, name: &str) -> QueryResult<Arc<WindowUDF>>;
+
+    fn udfs(&self) -> Vec<String>;
+    fn udafs(&self) -> Vec<String>;
+    fn udwfs(&self) -> Vec<String>;
+}

crates/s3select-api/src/query/logical_planner.rs

@@ -0,0 +1,54 @@
+// 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 async_trait::async_trait;
+use datafusion::arrow::datatypes::SchemaRef;
+use datafusion::logical_expr::LogicalPlan as DFPlan;
+
+use crate::QueryResult;
+
+use super::ast::ExtStatement;
+use super::session::SessionCtx;
+
+#[derive(Clone)]
+pub enum Plan {
+    // only support query sql
+    /// Query plan
+    Query(QueryPlan),
+}
+
+impl Plan {
+    pub fn schema(&self) -> SchemaRef {
+        match self {
+            Self::Query(p) => SchemaRef::from(p.df_plan.schema().as_ref().to_owned()),
+        }
+    }
+}
+
+#[derive(Debug, Clone)]
+pub struct QueryPlan {
+    pub df_plan: DFPlan,
+    pub is_tag_scan: bool,
+}
+
+impl QueryPlan {
+    pub fn is_explain(&self) -> bool {
+        matches!(self.df_plan, DFPlan::Explain(_) | DFPlan::Analyze(_))
+    }
+}
+
+#[async_trait]
+pub trait LogicalPlanner {
+    async fn create_logical_plan(&self, statement: ExtStatement, session: &SessionCtx) -> QueryResult<Plan>;
+}

crates/s3select-api/src/query/mod.rs

@@ -0,0 +1,57 @@
+// 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::Arc;
+
+use s3s::dto::SelectObjectContentInput;
+
+pub mod analyzer;
+pub mod ast;
+pub mod datasource;
+pub mod dispatcher;
+pub mod execution;
+pub mod function;
+pub mod logical_planner;
+pub mod optimizer;
+pub mod parser;
+pub mod physical_planner;
+pub mod scheduler;
+pub mod session;
+
+#[derive(Clone)]
+pub struct Context {
+    // maybe we need transfer some info?
+    pub input: Arc<SelectObjectContentInput>,
+}
+
+#[derive(Clone)]
+pub struct Query {
+    context: Context,
+    content: String,
+}
+
+impl Query {
+    #[inline(always)]
+    pub fn new(context: Context, content: String) -> Self {
+        Self { context, content }
+    }
+
+    pub fn context(&self) -> &Context {
+        &self.context
+    }
+
+    pub fn content(&self) -> &str {
+        self.content.as_str()
+    }
+}

crates/s3select-api/src/query/optimizer.rs

@@ -0,0 +1,29 @@
+// 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::Arc;
+
+use async_trait::async_trait;
+use datafusion::physical_plan::ExecutionPlan;
+
+use super::logical_planner::QueryPlan;
+use super::session::SessionCtx;
+use crate::QueryResult;
+
+pub type OptimizerRef = Arc<dyn Optimizer + Send + Sync>;
+
+#[async_trait]
+pub trait Optimizer {
+    async fn optimize(&self, plan: &QueryPlan, session: &SessionCtx) -> QueryResult<Arc<dyn ExecutionPlan>>;
+}

crates/s3select-api/src/query/parser.rs

@@ -0,0 +1,22 @@
+// 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::collections::VecDeque;
+
+use super::ast::ExtStatement;
+use crate::QueryResult;
+
+pub trait Parser {
+    fn parse(&self, sql: &str) -> QueryResult<VecDeque<ExtStatement>>;
+}

crates/s3select-api/src/query/physical_planner.rs

@@ -0,0 +1,35 @@
+// 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::Arc;
+
+use async_trait::async_trait;
+use datafusion::logical_expr::LogicalPlan;
+use datafusion::physical_plan::ExecutionPlan;
+use datafusion::physical_planner::ExtensionPlanner;
+
+use super::session::SessionCtx;
+use crate::QueryResult;
+
+#[async_trait]
+pub trait PhysicalPlanner {
+    /// Given a `LogicalPlan`, create an `ExecutionPlan` suitable for execution
+    async fn create_physical_plan(
+        &self,
+        logical_plan: &LogicalPlan,
+        session_state: &SessionCtx,
+    ) -> QueryResult<Arc<dyn ExecutionPlan>>;
+
+    fn inject_physical_transform_rule(&mut self, rule: Arc<dyn ExtensionPlanner + Send + Sync>);
+}

crates/s3select-api/src/query/scheduler.rs

@@ -0,0 +1,46 @@
+// 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::Arc;
+
+use async_trait::async_trait;
+use datafusion::common::Result;
+use datafusion::execution::context::TaskContext;
+use datafusion::physical_plan::{ExecutionPlan, SendableRecordBatchStream};
+
+pub type SchedulerRef = Arc<dyn Scheduler + Send + Sync>;
+
+#[async_trait]
+pub trait Scheduler {
+    /// Schedule the provided [`ExecutionPlan`] on this [`Scheduler`].
+    ///
+    /// Returns a [`ExecutionResults`] that can be used to receive results as they are produced,
+    /// as a [`futures::Stream`] of [`RecordBatch`]
+    async fn schedule(&self, plan: Arc<dyn ExecutionPlan>, context: Arc<TaskContext>) -> Result<ExecutionResults>;
+}
+
+pub struct ExecutionResults {
+    stream: SendableRecordBatchStream,
+}
+
+impl ExecutionResults {
+    pub fn new(stream: SendableRecordBatchStream) -> Self {
+        Self { stream }
+    }
+
+    /// Returns a [`SendableRecordBatchStream`] of this execution
+    pub fn stream(self) -> SendableRecordBatchStream {
+        self.stream
+    }
+}

crates/s3select-api/src/query/session.rs

@@ -0,0 +1,112 @@
+// 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::Arc;
+
+use datafusion::{
+    execution::{SessionStateBuilder, context::SessionState, runtime_env::RuntimeEnvBuilder},
+    parquet::data_type::AsBytes,
+    prelude::SessionContext,
+};
+use object_store::{ObjectStore, memory::InMemory, path::Path};
+use tracing::error;
+
+use crate::{QueryError, QueryResult, object_store::EcObjectStore};
+
+use super::Context;
+
+#[derive(Clone)]
+pub struct SessionCtx {
+    _desc: Arc<SessionCtxDesc>,
+    inner: SessionState,
+}
+
+impl SessionCtx {
+    pub fn inner(&self) -> &SessionState {
+        &self.inner
+    }
+}
+
+#[derive(Clone)]
+pub struct SessionCtxDesc {
+    // maybe we need some info
+}
+
+#[derive(Default)]
+pub struct SessionCtxFactory {
+    pub is_test: bool,
+}
+
+impl SessionCtxFactory {
+    pub async fn create_session_ctx(&self, context: &Context) -> QueryResult<SessionCtx> {
+        let df_session_ctx = self.build_df_session_context(context).await?;
+
+        Ok(SessionCtx {
+            _desc: Arc::new(SessionCtxDesc {}),
+            inner: df_session_ctx.state(),
+        })
+    }
+
+    async fn build_df_session_context(&self, context: &Context) -> QueryResult<SessionContext> {
+        let path = format!("s3://{}", context.input.bucket);
+        let store_url = url::Url::parse(&path).unwrap();
+        let rt = RuntimeEnvBuilder::new().build()?;
+        let df_session_state = SessionStateBuilder::new()
+            .with_runtime_env(Arc::new(rt))
+            .with_default_features();
+
+        let df_session_state = if self.is_test {
+            let store: Arc<dyn ObjectStore> = Arc::new(InMemory::new());
+            let data = b"id,name,age,department,salary
+            1,Alice,25,HR,5000
+            2,Bob,30,IT,6000
+            3,Charlie,35,Finance,7000
+            4,Diana,22,Marketing,4500
+            5,Eve,28,IT,5500
+            6,Frank,40,Finance,8000
+            7,Grace,26,HR,5200
+            8,Henry,32,IT,6200
+            9,Ivy,24,Marketing,4800
+            10,Jack,38,Finance,7500";
+            let data_bytes = data.as_bytes();
+            // let data = r#""year"╦"gender"╦"ethnicity"╦"firstname"╦"count"╦"rank"
+            // "2011"╦"FEMALE"╦"ASIAN AND PACIFIC ISLANDER"╦"SOPHIA"╦"119"╦"1"
+            // "2011"╦"FEMALE"╦"ASIAN AND PACIFIC ISLANDER"╦"CHLOE"╦"106"╦"2"
+            // "2011"╦"FEMALE"╦"ASIAN AND PACIFIC ISLANDER"╦"EMILY"╦"93"╦"3"
+            // "2011"╦"FEMALE"╦"ASIAN AND PACIFIC ISLANDER"╦"OLIVIA"╦"89"╦"4"
+            // "2011"╦"FEMALE"╦"ASIAN AND PACIFIC ISLANDER"╦"EMMA"╦"75"╦"5"
+            // "2011"╦"FEMALE"╦"ASIAN AND PACIFIC ISLANDER"╦"ISABELLA"╦"67"╦"6"
+            // "2011"╦"FEMALE"╦"ASIAN AND PACIFIC ISLANDER"╦"TIFFANY"╦"54"╦"7"
+            // "2011"╦"FEMALE"╦"ASIAN AND PACIFIC ISLANDER"╦"ASHLEY"╦"52"╦"8"
+            // "2011"╦"FEMALE"╦"ASIAN AND PACIFIC ISLANDER"╦"FIONA"╦"48"╦"9"
+            // "2011"╦"FEMALE"╦"ASIAN AND PACIFIC ISLANDER"╦"ANGELA"╦"47"╦"10""#;
+            // let data_bytes = Bytes::from(data);
+            let path = Path::from(context.input.key.clone());
+            store.put(&path, data_bytes.into()).await.map_err(|e| {
+                error!("put data into memory failed: {}", e.to_string());
+                QueryError::StoreError { e: e.to_string() }
+            })?;
+
+            df_session_state.with_object_store(&store_url, Arc::new(store)).build()
+        } else {
+            let store =
+                EcObjectStore::new(context.input.clone()).map_err(|_| QueryError::NotImplemented { err: String::new() })?;
+            df_session_state.with_object_store(&store_url, Arc::new(store)).build()
+        };
+
+        let df_session_ctx = SessionContext::new_with_state(df_session_state);
+
+        Ok(df_session_ctx)
+    }
+}