From 48b2f3d6e355525e50096c1b7b68ce37bcb7998d Mon Sep 17 00:00:00 2001 From: GatewayJ <835269233@qq.com> Date: Mon, 20 Jul 2026 21:02:54 +0800 Subject: [PATCH] fix(s3select): preserve CSV input as strings (#5030) * fix(s3select): preserve CSV input as strings * fix(s3select): address CSV schema review findings --- crates/s3select-api/src/query/session.rs | 22 ++-- .../s3select-query/src/dispatcher/manager.rs | 1 + crates/s3select-query/src/instance.rs | 28 ++-- .../src/test/integration_test.rs | 120 +++++++++++++++++- 4 files changed, 141 insertions(+), 30 deletions(-) diff --git a/crates/s3select-api/src/query/session.rs b/crates/s3select-api/src/query/session.rs index 739a2ef23..3a5879e3a 100644 --- a/crates/s3select-api/src/query/session.rs +++ b/crates/s3select-api/src/query/session.rs @@ -105,17 +105,17 @@ impl SessionCtxFactory { {\"id\":10,\"name\":\"Jack\",\"age\":38,\"department\":\"Finance\",\"salary\":7500}\n" .to_vec() } else { - 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" + b"id,name,age,department,salary\n\ + 1,Alice,25,HR,05000\n\ + 2,Bob,30,IT,6000\n\ + 3,Charlie,35,Finance,7000\n\ + 4,Diana,22,Marketing,4500\n\ + 5,Eve,28,IT,5500\n\ + 6,Frank,40,Finance,8000\n\ + 7,Grace,26,HR,5200\n\ + 8,Henry,32,IT,6200\n\ + 9,Ivy,24,Marketing,4800\n\ + 10,Jack,38,Finance,7500" .to_vec() }; diff --git a/crates/s3select-query/src/dispatcher/manager.rs b/crates/s3select-query/src/dispatcher/manager.rs index 46a145204..8ee239687 100644 --- a/crates/s3select-query/src/dispatcher/manager.rs +++ b/crates/s3select-query/src/dispatcher/manager.rs @@ -180,6 +180,7 @@ impl SimpleQueryDispatcher { let mut need_rename_volume_name = false; let mut need_ignore_volume_name = false; let mut file_format = CsvFormat::default() + .with_schema_infer_max_rec(0) .with_comment( csv.comments .clone() diff --git a/crates/s3select-query/src/instance.rs b/crates/s3select-query/src/instance.rs index 950d31365..ce9d5edd3 100644 --- a/crates/s3select-query/src/instance.rs +++ b/crates/s3select-query/src/instance.rs @@ -224,20 +224,20 @@ mod tests { .to_vec(); let expected = [ - "+----------------+---------+-----+------------+--------+", - "| 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 |", - "+----------------+---------+-----+------------+--------+", + "+----+---------+-----+------------+--------+", + "| id | name | age | department | salary |", + "+----+---------+-----+------------+--------+", + "| 1 | Alice | 25 | HR | 05000 |", + "| 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 |", + "+----+---------+-----+------------+--------+", ]; assert_batches_eq!(expected, &results); diff --git a/crates/s3select-query/src/test/integration_test.rs b/crates/s3select-query/src/test/integration_test.rs index 4ff8c8c9a..eee801261 100644 --- a/crates/s3select-query/src/test/integration_test.rs +++ b/crates/s3select-query/src/test/integration_test.rs @@ -15,6 +15,7 @@ #[cfg(test)] mod integration_tests { use crate::{create_fresh_db, get_global_db, instance::make_rustfsms}; + use datafusion::arrow::array::{Array, StringArray}; use rustfs_s3select_api::{ QueryError, query::{Context, Query}, @@ -170,14 +171,123 @@ mod integration_tests { } #[tokio::test] - async fn test_select_with_where_clause() { - let sql = "SELECT name, age FROM S3Object WHERE age > 30"; + async fn test_csv_values_remain_strings() { + let sql = "SELECT salary FROM S3Object LIMIT 1"; let input = create_test_input(sql); - let db = get_global_db(input.clone(), true).await.unwrap(); + let db = get_global_db(input.clone(), true).await.expect("create CSV test database"); let query = Query::new(Context { input: Arc::new(input) }, sql.to_string()); - let result = db.execute(&query).await; - assert!(result.is_ok()); + let batches = db + .execute(&query) + .await + .expect("execute CSV query") + .result() + .chunk_result() + .await + .expect("collect CSV query output"); + let salaries = batches + .first() + .expect("CSV query should return one batch") + .column(0) + .as_any() + .downcast_ref::() + .expect("CSV column should use UTF-8 string values"); + + assert_eq!(salaries.value(0), "05000"); + } + + #[tokio::test] + async fn test_csv_header_modes_keep_positional_string_columns() { + for (header, expected) in [ + (FileHeaderInfo::IGNORE, ["05000", "6000"]), + (FileHeaderInfo::NONE, ["salary", "05000"]), + ] { + let sql = "SELECT _5 FROM S3Object LIMIT 2"; + let mut input = create_test_input(sql); + input + .request + .input_serialization + .csv + .as_mut() + .expect("CSV input should be configured") + .file_header_info = Some(FileHeaderInfo::from_static(header)); + let db = get_global_db(input.clone(), true).await.expect("create CSV test database"); + let query = Query::new(Context { input: Arc::new(input) }, sql.to_string()); + let batches = db + .execute(&query) + .await + .expect("execute positional CSV query") + .result() + .chunk_result() + .await + .expect("collect positional CSV query output"); + + let values = batches + .iter() + .flat_map(|batch| { + let column = batch + .column(0) + .as_any() + .downcast_ref::() + .expect("positional CSV column should use UTF-8 strings"); + (0..column.len()).map(|row| column.value(row).to_string()).collect::>() + }) + .collect::>(); + + assert_eq!( + values, + expected.map(|value| value.to_string()), + "unexpected values for FileHeaderInfo={header}" + ); + } + } + + #[tokio::test] + async fn test_csv_numeric_comparison_with_and_without_cast() { + for sql in [ + "SELECT name, age FROM S3Object WHERE age > 30", + "SELECT name, age FROM S3Object WHERE CAST(age AS INT) > 30", + ] { + let input = create_test_input(sql); + let db = get_global_db(input.clone(), true).await.expect("create CSV test database"); + let query = Query::new(Context { input: Arc::new(input) }, sql.to_string()); + let batches = db + .execute(&query) + .await + .expect("execute CSV numeric comparison") + .result() + .chunk_result() + .await + .expect("collect CSV numeric comparison output"); + + let mut rows = Vec::new(); + for batch in &batches { + let names = batch + .column(0) + .as_any() + .downcast_ref::() + .expect("CSV name column should use UTF-8 strings"); + let ages = batch + .column(1) + .as_any() + .downcast_ref::() + .expect("CSV age column should use UTF-8 strings"); + for row in 0..batch.num_rows() { + rows.push((names.value(row).to_string(), ages.value(row).to_string())); + } + } + + assert_eq!( + rows, + [ + ("Charlie".to_string(), "35".to_string()), + ("Frank".to_string(), "40".to_string()), + ("Henry".to_string(), "32".to_string()), + ("Jack".to_string(), "38".to_string()), + ], + "unexpected rows for query: {sql}" + ); + } } #[tokio::test]