fix(s3select): preserve CSV input as strings (#5030)

* fix(s3select): preserve CSV input as strings

* fix(s3select): address CSV schema review findings
This commit is contained in:
GatewayJ
2026-07-20 21:02:54 +08:00
committed by GitHub
parent 376b90f61f
commit 48b2f3d6e3
4 changed files with 141 additions and 30 deletions
+11 -11
View File
@@ -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()
};
@@ -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()
+14 -14
View File
@@ -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);
@@ -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::<StringArray>()
.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::<StringArray>()
.expect("positional CSV column should use UTF-8 strings");
(0..column.len()).map(|row| column.value(row).to_string()).collect::<Vec<_>>()
})
.collect::<Vec<_>>();
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::<StringArray>()
.expect("CSV name column should use UTF-8 strings");
let ages = batch
.column(1)
.as_any()
.downcast_ref::<StringArray>()
.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]