impl(bigquery): use arrow format with jobs.query - #6472
Conversation
product-auto-label
Bot
added
the
api: bigquery
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request adds support for Arrow serialization in BigQuery query results, allowing the SDK to request and process Arrow record batches directly. Key changes include updating the query execution to request Arrow format with Zstd compression, implementing Arrow-to-Value conversion for various data types, and parsing Arrow schemas into BigQuery table schemas. Feedback on the changes suggests using unsigned_abs() to prevent overflow panics when formatting intervals, avoiding unnecessary clones when parsing range objects, and ensuring DataType::Interval is correctly mapped to "INTERVAL" in the schema conversion logic.
| let ym_sign = if v.months < 0 { "-" } else { "" }; | ||
| let years = v.months.abs() / 12; | ||
| let months = v.months.abs() % 12; | ||
|
|
||
| // Format Time H:MM:SS[.fffffffff] | ||
| let (time_sign, total_nanos) = if v.nanoseconds < 0 { | ||
| ("-", (-v.nanoseconds) as u64) | ||
| } else { | ||
| ("", v.nanoseconds as u64) | ||
| }; |
There was a problem hiding this comment.
Using .abs() on v.months and negation -v.nanoseconds can panic with overflow if they contain i32::MIN or i64::MIN respectively. Use unsigned_abs() to safely obtain the absolute value as an unsigned integer without any risk of overflow panics.
| let ym_sign = if v.months < 0 { "-" } else { "" }; | |
| let years = v.months.abs() / 12; | |
| let months = v.months.abs() % 12; | |
| // Format Time H:MM:SS[.fffffffff] | |
| let (time_sign, total_nanos) = if v.nanoseconds < 0 { | |
| ("-", (-v.nanoseconds) as u64) | |
| } else { | |
| ("", v.nanoseconds as u64) | |
| }; | |
| let ym_sign = if v.months < 0 { "-" } else { "" }; | |
| let months_abs = v.months.unsigned_abs(); | |
| let years = months_abs / 12; | |
| let months = months_abs % 12; | |
| // Format Time H:MM:SS[.fffffffff] | |
| let (time_sign, total_nanos) = if v.nanoseconds < 0 { | |
| ("-", v.nanoseconds.unsigned_abs()) | |
| } else { | |
| ("", v.nanoseconds as u64) | |
| }; |
| wkt::Value::Object(obj) => { | ||
| let start = match obj.get("start") { | ||
| Some(wkt::Value::Null) | None => None, | ||
| Some(val) => Some(T::from_sql(val.clone())?), | ||
| }; | ||
| let end = match obj.get("end") { | ||
| Some(wkt::Value::Null) | None => None, | ||
| Some(val) => Some(T::from_sql(val.clone())?), | ||
| }; | ||
| Ok(Range { start, end }) | ||
| } |
There was a problem hiding this comment.
Since value is passed by value to from_sql, the matched obj is an owned map. We can bind it mutably as mut obj and use remove instead of get and clone(). This avoids unnecessary cloning of the inner wkt::Values, which can be expensive.
| wkt::Value::Object(obj) => { | |
| let start = match obj.get("start") { | |
| Some(wkt::Value::Null) | None => None, | |
| Some(val) => Some(T::from_sql(val.clone())?), | |
| }; | |
| let end = match obj.get("end") { | |
| Some(wkt::Value::Null) | None => None, | |
| Some(val) => Some(T::from_sql(val.clone())?), | |
| }; | |
| Ok(Range { start, end }) | |
| } | |
| wkt::Value::Object(mut obj) => { | |
| let start = match obj.remove("start") { | |
| Some(wkt::Value::Null) | None => None, | |
| Some(val) => Some(T::from_sql(val)?), | |
| }; | |
| let end = match obj.remove("end") { | |
| Some(wkt::Value::Null) | None => None, | |
| Some(val) => Some(T::from_sql(val)?), | |
| }; | |
| Ok(Range { start, end }) | |
| } |
References
- Avoid unnecessary clones of owned values by taking ownership of the object and removing fields directly. (link)
| DataType::Time32(_) | DataType::Time64(_) => ("TIME", vec![]), | ||
| DataType::Timestamp(_, Some(_)) => ("TIMESTAMP", vec![]), | ||
| DataType::Timestamp(_, None) => ("DATETIME", vec![]), | ||
| DataType::Decimal128(_, _) | DataType::Decimal256(_, _) => ("NUMERIC", vec![]), |
There was a problem hiding this comment.
Map DataType::Interval to "INTERVAL" to ensure consistency with the data conversion logic in row.rs where DataType::Interval is explicitly handled.
DataType::Time32(_) | DataType::Time64(_) => ("TIME", vec![]),
DataType::Timestamp(_, Some(_)) => ("TIMESTAMP", vec![]),
DataType::Timestamp(_, None) => ("DATETIME", vec![]),
DataType::Decimal128(_, _) | DataType::Decimal256(_, _) => ("NUMERIC", vec![]),
DataType::Interval(_) => ("INTERVAL", vec![]),
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #6472 +/- ##
==========================================
- Coverage 96.36% 96.17% -0.19%
==========================================
Files 295 295
Lines 83697 84177 +480
==========================================
+ Hits 80655 80961 +306
- Misses 3042 3216 +174 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
1 participant
Trying out arrow support on jobs.query. This is gonna break on CI because support for it is behind an allowlist.