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.
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #6472 +/- ##
==========================================
- Coverage 96.52% 96.10% -0.42%
==========================================
Files 304 305 +1
Lines 87969 88690 +721
==========================================
+ Hits 84908 85234 +326
- Misses 3061 3456 +395 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request introduces Arrow results format acceleration for BigQuery query jobs (under the google_cloud_unstable_bigquery_arrow configuration flag), implementing direct deserialization from Arrow cells to Rust types and adding a performance benchmark. The review feedback focuses on enhancing robustness and performance: it suggests handling Arrow IPC stream parsing errors gracefully in RowIterator instead of panicking, extending support to LargeList and FixedSizeList array variants, and optimizing decimal string conversions to avoid unnecessary allocations.
| record_batches: VecDeque<Arc<RecordBatch>>, | ||
| row_index: usize, | ||
| rows: VecDeque<wkt::Struct>, | ||
| max_results: Option<u32>, | ||
| } |
There was a problem hiding this comment.
To handle Arrow IPC stream parsing errors gracefully without panicking, we should add an init_error field to the RowIterator struct. This allows us to defer returning the error to the next() call, adhering to the repository style guide which discourages unwrap() and expect() in production code.
| record_batches: VecDeque<Arc<RecordBatch>>, | |
| row_index: usize, | |
| rows: VecDeque<wkt::Struct>, | |
| max_results: Option<u32>, | |
| } | |
| record_batches: VecDeque<Arc<RecordBatch>>, | |
| row_index: usize, | |
| rows: VecDeque<wkt::Struct>, | |
| max_results: Option<u32>, | |
| init_error: Option<RowError>, | |
| } |
References
- Panics: unwrap() and expect() should typically be avoided in production code and examples (use ? or handle errors). (link)
| let (rows, record_batches) = match q.cached_data { | ||
| CachedData::Rows(rows) => (rows, VecDeque::new()), | ||
| CachedData::Arrow { | ||
| serialized_record_batch, | ||
| serialized_schema, | ||
| } => { | ||
| let reader = StreamReader::try_new( | ||
| Cursor::new(serialized_schema).chain(Cursor::new(serialized_record_batch)), | ||
| None, | ||
| ) | ||
| .expect("valid arrow IPC stream"); // TODO: convert error | ||
| let batches = reader | ||
| .map(|res| res.map(Arc::new)) | ||
| .collect::<std::result::Result<VecDeque<_>, _>>() | ||
| .expect("valid record batches"); // TODO: convert error | ||
| (VecDeque::new(), batches) | ||
| } | ||
| }; | ||
|
|
||
| Self { | ||
| job_service: q.job_service, | ||
| job_ref: q.job_ref, | ||
| schema: q.schema, | ||
| page_token: q.page_token, | ||
| rows: q.cached_rows, | ||
| record_batches, | ||
| row_index: 0, | ||
| rows, | ||
| max_results: q.max_results, | ||
| } | ||
| } |
There was a problem hiding this comment.
Instead of calling .expect() when parsing the Arrow IPC stream and record batches, we should capture any errors into the init_error field so they can be returned gracefully during iteration. This adheres to the repository style guide's rule against panics in production code.
let mut init_error = None;
let (rows, record_batches) = match q.cached_data {
CachedData::Rows(rows) => (rows, VecDeque::new()),
CachedData::Arrow {
serialized_record_batch,
serialized_schema,
} => match StreamReader::try_new(
Cursor::new(serialized_schema).chain(Cursor::new(serialized_record_batch)),
None,
) {
Ok(reader) => {
match reader
.map(|res| res.map(Arc::new))
.collect::<std::result::Result<VecDeque<_>, _>>()
{
Ok(batches) => (VecDeque::new(), batches),
Err(e) => {
init_error = Some(RowError::InvalidRowFormat(format!(
"failed to read record batches: {e}"
)));
(VecDeque::new(), VecDeque::new())
}
}
}
Err(e) => {
init_error = Some(RowError::InvalidRowFormat(format!(
"failed to parse arrow schema: {e}"
)));
(VecDeque::new(), VecDeque::new())
}
},
};
Self {
job_service: q.job_service,
job_ref: q.job_ref,
schema: q.schema,
page_token: q.page_token,
record_batches,
row_index: 0,
rows,
max_results: q.max_results,
init_error,
}
}References
- Panics: unwrap() and expect() should typically be avoided in production code and examples (use ? or handle errors). (link)
| pub async fn next(&mut self) -> Option<Result<Row>> { | ||
| loop { | ||
| while let Some(batch) = self.record_batches.front() { |
There was a problem hiding this comment.
Check and return the deferred init_error if any error occurred during the initialization of the Arrow IPC stream reader.
pub async fn next(&mut self) -> Option<Result<Row>> {
if let Some(err) = self.init_error.take() {
return Some(Err(err));
}
loop {
while let Some(batch) = self.record_batches.front() {| DataType::List(_) => { | ||
| let v = Vec::<wkt::Value>::from_arrow(cell)?; | ||
| Ok(wkt::Value::Array(v)) | ||
| } |
There was a problem hiding this comment.
To ensure robust handling of all list array variants that might be returned by the Arrow integration, we should support LargeList and FixedSizeList as well.
| DataType::List(_) => { | |
| let v = Vec::<wkt::Value>::from_arrow(cell)?; | |
| Ok(wkt::Value::Array(v)) | |
| } | |
| DataType::List(_) | DataType::LargeList(_) | DataType::FixedSizeList(_, _) => { | |
| let v = Vec::<wkt::Value>::from_arrow(cell)?; | |
| Ok(wkt::Value::Array(v)) | |
| } |
| fn from_arrow(cell: ArrowCell<'_>) -> Result<Self, ConvertError> { | ||
| if cell.is_null() { | ||
| return Err(ConvertError::NotNull); | ||
| } | ||
| if let Some(arr) = cell.downcast_ref::<arrow::array::ListArray>() { | ||
| let value_arr = arr.value(cell.row_idx); | ||
| let mut result = Vec::with_capacity(value_arr.len()); | ||
| for i in 0..value_arr.len() { | ||
| result.push(T::from_arrow(ArrowCell::new(value_arr.as_ref(), i))?); | ||
| } | ||
| return Ok(result); | ||
| } | ||
| Err(ConvertError::TypeMismatch { | ||
| expected: "list array", | ||
| got: wkt::Value::String(cell.data_type_str()), | ||
| }) | ||
| } | ||
| } |
There was a problem hiding this comment.
Support downcasting to LargeListArray and FixedSizeListArray in Vec<T>::from_arrow to prevent unexpected runtime type mismatch errors with other list array variants.
fn from_arrow(cell: ArrowCell<'_>) -> Result<Self, ConvertError> {
if cell.is_null() {
return Err(ConvertError::NotNull);
}
if let Some(arr) = cell.downcast_ref::<arrow::array::ListArray>() {
let value_arr = arr.value(cell.row_idx);
let mut result = Vec::with_capacity(value_arr.len());
for i in 0..value_arr.len() {
result.push(T::from_arrow(ArrowCell::new(value_arr.as_ref(), i))?);
}
return Ok(result);
}
if let Some(arr) = cell.downcast_ref::<arrow::array::LargeListArray>() {
let value_arr = arr.value(cell.row_idx);
let mut result = Vec::with_capacity(value_arr.len());
for i in 0..value_arr.len() {
result.push(T::from_arrow(ArrowCell::new(value_arr.as_ref(), i))?);
}
return Ok(result);
}
if let Some(arr) = cell.downcast_ref::<arrow::array::FixedSizeListArray>() {
let value_arr = arr.value(cell.row_idx);
let mut result = Vec::with_capacity(value_arr.len());
for i in 0..value_arr.len() {
result.push(T::from_arrow(ArrowCell::new(value_arr.as_ref(), i))?);
}
return Ok(result);
}
Err(ConvertError::TypeMismatch {
expected: "list array",
got: wkt::Value::String(cell.data_type_str()),
})
}| fn from_arrow(cell: ArrowCell<'_>) -> Result<Self, ConvertError> { | ||
| if cell.is_null() { | ||
| return Err(ConvertError::NotNull); | ||
| } | ||
| let row_idx = cell.row_idx; | ||
| if let Some(arr) = cell.downcast_ref::<arrow::array::Decimal128Array>() { | ||
| let val = arr.value(row_idx); | ||
| let scale = arr.scale() as u32; | ||
| return rust_decimal::Decimal::try_from_i128_with_scale(val, scale) | ||
| .map_err(|e| ConvertError::Convert(Box::new(e))); | ||
| } | ||
| if let Some(arr) = cell.downcast_ref::<arrow::array::Decimal256Array>() { | ||
| let s = arr.value_as_string(row_idx); | ||
| let trimmed = if let Some((int_part, frac_part)) = s.split_once('.') { | ||
| let frac_trimmed = frac_part.trim_end_matches('0'); | ||
| if frac_trimmed.is_empty() { | ||
| int_part.to_string() | ||
| } else { | ||
| format!("{int_part}.{frac_trimmed}") | ||
| } | ||
| } else { | ||
| s | ||
| }; | ||
| return trimmed | ||
| .parse::<rust_decimal::Decimal>() | ||
| .map_err(|e| ConvertError::Convert(Box::new(e))); | ||
| } | ||
| Err(ConvertError::TypeMismatch { | ||
| expected: "decimal", | ||
| got: wkt::Value::String(cell.data_type_str()), | ||
| }) | ||
| } | ||
| } |
There was a problem hiding this comment.
Instead of allocating new strings via int_part.to_string() and format!, we can modify the owned String returned by arr.value_as_string(row_idx) in-place using truncate. This completely avoids extra string allocations, adhering to the repository style guide's focus on eliminating unnecessary allocations. Since we are slicing the string directly, we should also document the safety guarantee with a comment.
fn from_arrow(cell: ArrowCell<'_>) -> Result<Self, ConvertError> {
if cell.is_null() {
return Err(ConvertError::NotNull);
}
let row_idx = cell.row_idx;
if let Some(arr) = cell.downcast_ref::<arrow::array::Decimal128Array>() {
let val = arr.value(row_idx);
let scale = arr.scale() as u32;
return rust_decimal::Decimal::try_from_i128_with_scale(val, scale)
.map_err(|e| ConvertError::Convert(Box::new(e)));
}
if let Some(arr) = cell.downcast_ref::<arrow::array::Decimal256Array>() {
let mut s = arr.value_as_string(row_idx);
if let Some(dot_idx) = s.find('.') {
// SAFETY: dot_idx is found via s.find('.'), so dot_idx + 1 is a valid char boundary and within bounds for a decimal string.
let trimmed = s[dot_idx + 1..].trim_end_matches('0');
if trimmed.is_empty() {
s.truncate(dot_idx);
} else {
s.truncate(dot_idx + 1 + trimmed.len());
}
}
return s
.parse::<rust_decimal::Decimal>()
.map_err(|e| ConvertError::Convert(Box::new(e)));
}
Err(ConvertError::TypeMismatch {
expected: "decimal",
got: wkt::Value::String(cell.data_type_str()),
})
}References
- Unnecessary clones: Scrutinize expensive uses of clone() (i.e. for Strings, not for Arcs). Look out for subtle copies, e.g. from chunks() or a to_() instead of into_(). Is it necessary to copy the data? Can we move the data instead? (link)
- When slicing a collection directly (e.g.,
buf[len..]) in Rust, if the operation is guaranteed to be safe due to specific invariants, document this safety guarantee with a comment rather than switching to safe but more verbose error-handling methods like.get().
| DataType::List(sub_field) | DataType::LargeList(sub_field) => { | ||
| let mut sub = arrow_field_to_table_field(sub_field); | ||
| sub.name = field.name().clone(); | ||
| sub.mode = "REPEATED".to_string(); | ||
| sub | ||
| } |
There was a problem hiding this comment.
Support FixedSizeList in arrow_field_to_table_field to ensure complete consistency with other list array variants.
| DataType::List(sub_field) | DataType::LargeList(sub_field) => { | |
| let mut sub = arrow_field_to_table_field(sub_field); | |
| sub.name = field.name().clone(); | |
| sub.mode = "REPEATED".to_string(); | |
| sub | |
| } | |
| DataType::List(sub_field) | DataType::LargeList(sub_field) | DataType::FixedSizeList(sub_field, _) => { | |
| let mut sub = arrow_field_to_table_field(sub_field); | |
| sub.name = field.name().clone(); | |
| sub.mode = "REPEATED".to_string(); | |
| sub | |
| } |
1 participant
Trying out arrow support on jobs.query. This is gonna break on CI because support for it is behind an allowlist.