feat(s3select): add JSON handling and flattening for EcObjectStore (#1930)

Signed-off-by: 0xdx2 <xuedamon2@gmail.com>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
Co-authored-by: houseme <housemecn@gmail.com>

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

@@ -218,7 +218,16 @@ impl SimpleQueryDispatcher {
                 (ListingOptions::new(Arc::new(file_format)).with_file_extension(".parquet"), false, false)
             } else if self.input.request.input_serialization.json.is_some() {
                 let file_format = JsonFormat::default();
-                (ListingOptions::new(Arc::new(file_format)).with_file_extension(".json"), false, false)
+                // 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)
             } else {
                 return Err(QueryError::NotImplemented {
                     err: "not support this file type".to_string(),

crates/s3select-query/src/test/integration_test.rs

@@ -20,8 +20,8 @@ mod integration_tests {
         query::{Context, Query},
     };
     use s3s::dto::{
-        CSVInput, CSVOutput, ExpressionType, FileHeaderInfo, InputSerialization, OutputSerialization, SelectObjectContentInput,
-        SelectObjectContentRequest,
+        CSVInput, CSVOutput, ExpressionType, FileHeaderInfo, InputSerialization, JSONInput, JSONOutput, JSONType,
+        OutputSerialization, SelectObjectContentInput, SelectObjectContentRequest,
     };
     use std::sync::Arc;
 
@@ -53,6 +53,36 @@ mod integration_tests {
         }
     }
 
+    /// Build a `SelectObjectContentInput` targeting a JSON DOCUMENT file.
+    /// Uses `JSONType::DOCUMENT` so the NDJSON-flattening path in
+    /// `EcObjectStore` is exercised.
+    fn create_test_json_input(sql: &str) -> SelectObjectContentInput {
+        SelectObjectContentInput {
+            bucket: "test-bucket".to_string(),
+            expected_bucket_owner: None,
+            key: "test.json".to_string(),
+            sse_customer_algorithm: None,
+            sse_customer_key: None,
+            sse_customer_key_md5: None,
+            request: SelectObjectContentRequest {
+                expression: sql.to_string(),
+                expression_type: ExpressionType::from_static("SQL"),
+                input_serialization: InputSerialization {
+                    json: Some(JSONInput {
+                        type_: Some(JSONType::from_static(JSONType::DOCUMENT)),
+                    }),
+                    ..Default::default()
+                },
+                output_serialization: OutputSerialization {
+                    json: Some(JSONOutput::default()),
+                    ..Default::default()
+                },
+                request_progress: None,
+                scan_range: None,
+            },
+        }
+    }
+
     #[tokio::test]
     async fn test_database_creation() {
         let input = create_test_input("SELECT * FROM S3Object");
@@ -225,4 +255,170 @@ mod integration_tests {
             assert!(result.is_ok());
         }
     }
+
+    // ──────────────────────────────────────────────
+    // JSON-input variants of all the above tests
+    // These exercise the JSONType::LINES (JSON lines) code path
+    // ──────────────────────────────────────────────
+
+    #[tokio::test]
+    async fn test_database_creation_json() {
+        let input = create_test_json_input("SELECT * FROM S3Object");
+        let result = make_rustfsms(Arc::new(input), true).await;
+        assert!(result.is_ok());
+    }
+
+    #[tokio::test]
+    async fn test_global_db_creation_json() {
+        let input = create_test_json_input("SELECT * FROM S3Object");
+        let result = get_global_db(input.clone(), true).await;
+        assert!(result.is_ok());
+    }
+
+    #[tokio::test]
+    async fn test_simple_select_query_json() {
+        let sql = "SELECT * FROM S3Object";
+        let input = create_test_json_input(sql);
+        let db = get_global_db(input.clone(), true).await.unwrap();
+        let query = Query::new(Context { input: Arc::new(input) }, sql.to_string());
+
+        let result = db.execute(&query).await;
+        assert!(result.is_ok());
+
+        let query_handle = result.unwrap();
+        let output = query_handle.result().chunk_result().await;
+        assert!(output.is_ok());
+    }
+
+    #[tokio::test]
+    async fn test_select_with_where_clause_json() {
+        let sql = "SELECT name, age FROM S3Object WHERE age > 30";
+        let input = create_test_json_input(sql);
+        let db = get_global_db(input.clone(), true).await.unwrap();
+        let query = Query::new(Context { input: Arc::new(input) }, sql.to_string());
+
+        let result = db.execute(&query).await;
+        assert!(result.is_ok());
+    }
+
+    #[tokio::test]
+    async fn test_select_with_aggregation_json() {
+        let sql = "SELECT department, COUNT(*) as count FROM S3Object GROUP BY department";
+        let input = create_test_json_input(sql);
+        let db = get_global_db(input.clone(), true).await.unwrap();
+        let query = Query::new(Context { input: Arc::new(input) }, sql.to_string());
+
+        let result = db.execute(&query).await;
+        // Aggregation queries may fail due to lack of actual data, which is acceptable
+        match result {
+            Ok(_) => {
+                // If successful, that's great
+            }
+            Err(_) => {
+                // Expected to fail due to no actual data source
+            }
+        }
+    }
+
+    #[tokio::test]
+    async fn test_invalid_sql_syntax_json() {
+        let sql = "INVALID SQL SYNTAX";
+        let input = create_test_json_input(sql);
+        let db = get_global_db(input.clone(), true).await.unwrap();
+        let query = Query::new(Context { input: Arc::new(input) }, sql.to_string());
+
+        let result = db.execute(&query).await;
+        assert!(result.is_err());
+    }
+
+    #[tokio::test]
+    async fn test_multi_statement_error_json() {
+        let sql = "SELECT * FROM S3Object; SELECT 1;";
+        let input = create_test_json_input(sql);
+        let db = get_global_db(input.clone(), true).await.unwrap();
+        let query = Query::new(Context { input: Arc::new(input) }, sql.to_string());
+
+        let result = db.execute(&query).await;
+        assert!(result.is_err());
+
+        if let Err(QueryError::MultiStatement { num, .. }) = result {
+            assert_eq!(num, 2);
+        } else {
+            panic!("Expected MultiStatement error");
+        }
+    }
+
+    #[tokio::test]
+    async fn test_query_state_machine_workflow_json() {
+        let sql = "SELECT * FROM S3Object";
+        let input = create_test_json_input(sql);
+        let db = get_global_db(input.clone(), true).await.unwrap();
+        let query = Query::new(Context { input: Arc::new(input) }, sql.to_string());
+
+        let state_machine = db.build_query_state_machine(query.clone()).await;
+        assert!(state_machine.is_ok());
+
+        let state_machine = state_machine.unwrap();
+
+        let logical_plan = db.build_logical_plan(state_machine.clone()).await;
+        assert!(logical_plan.is_ok());
+
+        if let Ok(Some(plan)) = logical_plan {
+            let execution_result = db.execute_logical_plan(plan, state_machine).await;
+            assert!(execution_result.is_ok());
+        }
+    }
+
+    #[tokio::test]
+    async fn test_query_with_limit_json() {
+        let sql = "SELECT * FROM S3Object LIMIT 5";
+        let input = create_test_json_input(sql);
+        let db = get_global_db(input.clone(), true).await.unwrap();
+        let query = Query::new(Context { input: Arc::new(input) }, sql.to_string());
+
+        let result = db.execute(&query).await;
+        assert!(result.is_ok());
+
+        let query_handle = result.unwrap();
+        let output = query_handle.result().chunk_result().await.unwrap();
+
+        let total_rows: usize = output.iter().map(|batch| batch.num_rows()).sum();
+        assert!(total_rows <= 5);
+    }
+
+    #[tokio::test]
+    async fn test_query_with_order_by_json() {
+        let sql = "SELECT name, age FROM S3Object ORDER BY age DESC";
+        let input = create_test_json_input(sql);
+        let db = get_global_db(input.clone(), true).await.unwrap();
+        let query = Query::new(Context { input: Arc::new(input) }, sql.to_string());
+
+        let result = db.execute(&query).await;
+        assert!(result.is_ok());
+    }
+
+    #[tokio::test]
+    async fn test_concurrent_queries_json() {
+        let sql = "SELECT * FROM S3Object";
+        let input = create_test_json_input(sql);
+        let db = get_global_db(input.clone(), true).await.unwrap();
+
+        let mut handles = vec![];
+        for i in 0..3 {
+            let query = Query::new(
+                Context {
+                    input: Arc::new(input.clone()),
+                },
+                format!("SELECT * FROM S3Object LIMIT {}", i + 1),
+            );
+            let db_clone = db.clone();
+            let handle = tokio::spawn(async move { db_clone.execute(&query).await });
+            handles.push(handle);
+        }
+
+        for handle in handles {
+            let result = handle.await.unwrap();
+            assert!(result.is_ok());
+        }
+    }
 }