Ling0925
/
buster
mirror of https://github.com/buster-so/buster.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

refactor(snowflake_query): optimize Arrow data processing with explicit row collection (#92)

This commit is contained in:
dal 2025-02-04 15:57:35 -08:00 committed by GitHub
parent d0ff21e10d
commit 0fde90b848
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 315 additions and 313 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

20
api/src/utils/query_engine/data_source_query_routes/snowflake_query.rs
View File

@ -56,12 +56,13 @@ pub async fn snowflake_query(
let rows = match snowflake_client.exec(&limited_query).await {
Ok(result) => match result {
snowflake_api::QueryResult::Arrow(result) => {
result
.iter()
.flat_map(|batch| {
let mut all_rows = Vec::new();
// Process each batch in order
for batch in result.iter() {
let schema = batch.schema();
(0..batch.num_rows()).map(move |row_idx| {
schema
for row_idx in 0..batch.num_rows() {
let row = schema
.fields()
.iter()
.enumerate()
@ -389,10 +390,11 @@ pub async fn snowflake_query(
};
(field.name().clone(), data_type)
})
.collect::<IndexMap<String, DataType>>()
})
})
.collect()
.collect::<IndexMap<String, DataType>>();
all_rows.push(row);
}
}
all_rows
}
_ => Vec::new(),
},