github-actions[bot] commented on code in PR #4394: URL: https://github.com/apache/iggy/pull/4394#discussion_r4173903106
########## core/connectors/sinks/bigquery_sink/src/lib.rs: ########## @@ -0,0 +1,457 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Google BigQuery sink connector. +//! +//! Writes Iggy messages into a BigQuery table through the Storage Write API +//! `_default` stream, with rows encoded as Arrow record batches. +//! +//! Module map: +//! - [`schema`]: BigQuery table schema (from `tables.get`) to Arrow schema. Review Comment: warning: The crate docs link five private modules, so rustdoc warns and the doctest job fails on `-D warnings`. Remove the square brackets around `schema`, `encode`, `client`, `error` and `sink`. ########## core/connectors/sinks/bigquery_sink/src/error.rs: ########## @@ -0,0 +1,207 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Failure classification for BigQuery calls. +//! +//! Only transient failures are retried. Retrying a permanent one (bad +//! schema, missing permission) cannot succeed and only delays the error. + +use crate::schema::SchemaError; +use gcloud_gax::grpc::{Code, Status}; +use iggy_connector_sdk::Error; +use std::fmt; + +/// An `AppendRows` call that did not append. Row-level errors are not part +/// of this type: they are an expected outcome the caller handles by +/// dropping the reported rows. +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) enum AppendError { + /// The call failed, or the response carried an error status. + Rpc { code: Code, message: String }, + /// The response stream closed without a response. + NoResponse, +} + +/// A `tables.get` call that did not return a usable schema. +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) enum TableError { + Http { status: u16, body: String }, + Transport(String), + Token(String), + Schema(SchemaError), +} + +impl AppendError { + pub(crate) fn is_retryable(&self) -> bool { + match self { + AppendError::Rpc { code, .. } => is_retryable_code(*code), + // The append may or may not have landed. Retrying risks a + // duplicate, which the delivery guarantees allow. + AppendError::NoResponse => true, + } + } +} + +impl TableError { + pub(crate) fn is_retryable(&self) -> bool { + match self { + TableError::Http { status, .. } => *status == 429 || *status >= 500, + TableError::Transport(_) => true, + TableError::Token(_) | TableError::Schema(_) => false, Review Comment: warning: `TableError::Token` is not retryable, so `retry_async` aborts `open()` on the first attempt when a token refresh fails. Return `true` for `TableError::Token` so the connector retries transient token failures. ########## core/connectors/sinks/bigquery_sink/src/encode.rs: ########## @@ -0,0 +1,1128 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Turns a run of `ConsumedMessage`s into Arrow record batches that fit in +//! one `AppendRows` request each. +//! +//! A message that cannot become a row (wrong payload type, a value that does +//! not fit its column, a row larger than the request budget) is rejected on +//! its own and reported back with its offset. The rest of the batch is still +//! written. +//! +//! Mapped mode decodes JSON objects with `arrow-json`. The fast path decodes +//! the whole run at once. When that fails, the run is decoded again row by +//! row so that only the offending rows are dropped. +//! +//! The writer schema of a request only contains the columns that at least one +//! row in the run sets. Columns no row mentions are left out, so BigQuery +//! fills them according to `missing_value` (the column default, or NULL). + +use crate::schema::{BqType, Column, MetaColumn, Mode, RawKind, TableLayout, UTC}; +use arrow::array::{ + ArrayRef, BinaryArray, Int64Array, RecordBatch, StringArray, TimestampMicrosecondArray, +}; +use arrow::compute::{cast, concat_batches}; +use arrow::datatypes::{DataType, FieldRef, Schema, SchemaRef}; +use arrow::error::ArrowError; +use arrow::ipc::writer::{ + CompressionContext, DictionaryTracker, IpcDataGenerator, IpcWriteOptions, write_message, +}; +use arrow::json::ReaderBuilder; +use base64::Engine; +use base64::engine::general_purpose::STANDARD as BASE64; +use iggy_connector_sdk::{ConsumedMessage, MessagesMetadata, Payload, TopicMetadata}; +use simd_json::{OwnedValue, StaticNode}; +use std::fmt::Write as _; +use std::sync::Arc; + +const HEADER_ENCODING_BASE64: &str = "base64"; +const UTC_OFFSET: &str = "+00:00"; + +/// One `AppendRows` request worth of rows. +#[derive(Debug)] +pub(crate) struct Chunk { + pub batch: RecordBatch, + /// Iggy offset of each row, same order as `batch`. + pub offsets: Vec<u64>, + pub schema_bytes: Vec<u8>, + pub batch_bytes: Vec<u8>, +} + +/// A message that was not turned into a row. +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct Rejected { + pub offset: u64, + pub reason: String, +} + +#[derive(Debug, Default)] +pub(crate) struct Encoded { + pub chunks: Vec<Chunk>, + pub rejected: Vec<Rejected>, +} + +/// Per-run context shared by every row. +pub(crate) struct RunContext<'a> { + pub layout: &'a TableLayout, + pub topic: &'a TopicMetadata, + pub messages: &'a MessagesMetadata, + pub max_request_bytes: usize, +} + +/// Encode a run of messages. Fails only on an internal Arrow error, which +/// means the batch as a whole cannot be built. Per-row problems end up in +/// `Encoded::rejected`. +pub(crate) fn encode( + ctx: &RunContext<'_>, + messages: Vec<ConsumedMessage>, +) -> Result<Encoded, ArrowError> { + let mut rejected = Vec::new(); + let batch = match &ctx.layout.raw { + Some((kind, field)) => encode_raw(ctx, *kind, field, messages, &mut rejected)?, + None => encode_mapped(ctx, messages, &mut rejected)?, + }; + let mut chunks = Vec::new(); + if let Some((batch, offsets)) = batch { + split_into_chunks( + batch, + offsets, + ctx.max_request_bytes, + &mut chunks, + &mut rejected, + )?; + } + Ok(Encoded { chunks, rejected }) +} + +/// Build a new chunk holding only the rows of `chunk` whose index is not in +/// `drop`. Used after BigQuery reports row-level errors. +pub(crate) fn without_rows(chunk: &Chunk, drop: &[usize]) -> Result<Option<Chunk>, ArrowError> { + let keep: Vec<bool> = (0..chunk.batch.num_rows()) + .map(|row| !drop.contains(&row)) + .collect(); + let offsets: Vec<u64> = chunk + .offsets + .iter() + .zip(&keep) + .filter_map(|(offset, keep)| keep.then_some(*offset)) + .collect(); + if offsets.is_empty() { + return Ok(None); + } + let mask = arrow::array::BooleanArray::from(keep); + let batch = arrow::compute::filter_record_batch(&chunk.batch, &mask)?; + let (schema_bytes, batch_bytes) = encode_ipc(&batch)?; + Ok(Some(Chunk { + batch, + offsets, + schema_bytes, + batch_bytes, + })) +} + +// ─── Mapped mode ───────────────────────────────────────────────────────────── + +fn encode_mapped( + ctx: &RunContext<'_>, + messages: Vec<ConsumedMessage>, + rejected: &mut Vec<Rejected>, +) -> Result<Option<(RecordBatch, Vec<u64>)>, ArrowError> { + let layout = ctx.layout; + let mut rows = Vec::with_capacity(messages.len()); + let mut kept = Vec::with_capacity(messages.len()); + for mut message in messages { + let payload = std::mem::replace(&mut message.payload, Payload::Raw(Vec::new())); + match json_object(payload).and_then(|mut row| { + normalize_row(&mut row, layout)?; + Ok(row) + }) { + Ok(row) => { + rows.push(row); + kept.push(message); + } + Err(reason) => rejected.push(Rejected { + offset: message.offset, + reason, + }), + } + } + if rows.is_empty() { + return Ok(None); + } + + let present = present_columns(&rows, layout); + if present.is_empty() && layout.metadata.is_empty() { + rejected.extend(kept.iter().map(|m| Rejected { + offset: m.offset, + reason: "payload has no field matching a table column".to_owned(), + })); + return Ok(None); + } + let payload_fields: Vec<FieldRef> = present.iter().map(|&i| layout.fields[i].clone()).collect(); + + let (decoded, survivors) = decode_rows(&payload_fields, &rows, kept.len(), |row, reason| { + rejected.push(Rejected { + offset: kept[row].offset, + reason, + }); + })?; + if survivors.is_empty() { + return Ok(None); + } + let kept: Vec<&ConsumedMessage> = survivors.iter().map(|&row| &kept[row]).collect(); + let offsets = kept.iter().map(|m| m.offset).collect(); + + let mut fields = payload_fields; + let mut columns: Vec<ArrayRef> = decoded.map(|b| b.columns().to_vec()).unwrap_or_default(); + append_metadata(ctx, &kept, &mut fields, &mut columns)?; + let batch = RecordBatch::try_new(Arc::new(Schema::new(fields)), columns)?; + Ok(Some((batch, offsets))) +} + +/// Take the JSON object out of a payload. Text and proto text must parse as +/// a JSON object. +fn json_object(payload: Payload) -> Result<OwnedValue, String> { + let value = match payload { + Payload::Json(value) => value, + Payload::Text(text) | Payload::Proto(text) => { + let mut bytes = text.into_bytes(); + simd_json::to_owned_value(&mut bytes) + .map_err(|e| format!("payload text is not valid JSON: {e}"))? + } + other => { + return Err(format!( + "{} payload is not supported in mapped mode", + payload_kind(&other) + )); + } + }; + match value { + OwnedValue::Object(_) => Ok(value), + _ => Err("payload is not a JSON object".to_owned()), + } +} + +/// Prepare one row for `arrow-json`: +/// - a REQUIRED column without a default must be present and non-null, +/// - REPEATED columns that are absent or null become `[]`, +/// - JSON column values become their JSON text, +/// - BYTES column values are base64 in JSON and become hex, which is what +/// `arrow-json` decodes into binary. +fn normalize_row(row: &mut OwnedValue, layout: &TableLayout) -> Result<(), String> { + let OwnedValue::Object(object) = row else { + return Err("payload is not a JSON object".to_owned()); + }; + for column in &layout.columns { + let value = object.get_mut(column.name.as_str()); + match (value, column.mode) { + (None | Some(OwnedValue::Static(StaticNode::Null)), Mode::Repeated) => { + object.insert(column.name.clone(), empty_array()); + } + (None | Some(OwnedValue::Static(StaticNode::Null)), Mode::Required) + if !column.has_default => + { + return Err(format!( + "missing value for REQUIRED column '{}'", + column.name + )); + } + (Some(value), _) => normalize_value(value, column)?, + (None, _) => {} + } + } + Ok(()) +} + +fn normalize_value(value: &mut OwnedValue, column: &Column) -> Result<(), String> { + if column.mode == Mode::Repeated { + let OwnedValue::Array(items) = value else { + return Err(format!("column '{}' expects an array", column.name)); + }; + for item in items.iter_mut() { + normalize_scalar(item, column)?; + } + return Ok(()); + } + normalize_scalar(value, column) +} + +fn normalize_scalar(value: &mut OwnedValue, column: &Column) -> Result<(), String> { + if matches!(value, OwnedValue::Static(StaticNode::Null)) { + return Ok(()); + } + match &column.bq_type { + BqType::Json => { + let text = simd_json::to_string(value) + .map_err(|e| format!("column '{}': cannot serialize JSON: {e}", column.name))?; + *value = OwnedValue::String(text); + } + BqType::Bytes => { + let OwnedValue::String(text) = value else { + return Err(format!("column '{}' expects a base64 string", column.name)); + }; + let bytes = BASE64 + .decode(text.as_bytes()) + .map_err(|e| format!("column '{}': invalid base64: {e}", column.name))?; + *value = OwnedValue::String(to_hex(&bytes)); + } + BqType::Record(children) => { + let OwnedValue::Object(object) = value else { + return Ok(()); + }; + for child in children { + match (object.get_mut(child.name.as_str()), child.mode) { + (None | Some(OwnedValue::Static(StaticNode::Null)), Mode::Repeated) => { + object.insert(child.name.clone(), empty_array()); + } + (Some(child_value), _) => normalize_value(child_value, child)?, + (None, _) => {} + } + } + } + _ => {} + } + Ok(()) +} + +/// Indexes of the layout columns that at least one row sets, in table order. +fn present_columns(rows: &[OwnedValue], layout: &TableLayout) -> Vec<usize> { + let mut present = vec![false; layout.columns.len()]; + for row in rows { + if let OwnedValue::Object(object) = row { + for key in object.keys() { + if let Some(&index) = layout.index.get(key.as_str()) { + present[index] = true; + } + } + } + } + present + .iter() + .enumerate() + .filter_map(|(index, set)| set.then_some(index)) + .collect() +} + +/// Decode rows into one batch. Returns the batch (`None` when there are no +/// payload fields) and the indexes of the rows it holds. +fn decode_rows( + fields: &[FieldRef], + rows: &[OwnedValue], + row_count: usize, + mut on_reject: impl FnMut(usize, String), +) -> Result<(Option<RecordBatch>, Vec<usize>), ArrowError> { + if fields.is_empty() { + return Ok((None, (0..row_count).collect())); + } + let target: SchemaRef = Arc::new(Schema::new(fields.to_vec())); + let schema: SchemaRef = Arc::new(Schema::new( + fields + .iter() + .map(|f| { + f.as_ref() + .clone() + .with_data_type(decode_type(f.data_type())) + }) + .collect::<Vec<_>>(), + )); + if let Ok(batch) = decode(&schema, rows) { + return Ok((Some(to_target(batch, &target)?), (0..row_count).collect())); + } + + let mut batches = Vec::with_capacity(rows.len()); + let mut survivors = Vec::with_capacity(rows.len()); + for (index, row) in rows.iter().enumerate() { + match decode(&schema, std::slice::from_ref(row)) { + Ok(batch) => { + batches.push(batch); + survivors.push(index); + } + Err(e) => on_reject(index, format!("row does not match the table schema: {e}")), + } + } + if batches.is_empty() { + return Ok((None, survivors)); + } + let batch = concat_batches(&schema, &batches)?; + Ok((Some(to_target(batch, &target)?), survivors)) +} + +/// The type `arrow-json` decodes into. Without the `chrono-tz` feature it +/// only parses offset timezones, so `"UTC"` is decoded as `"+00:00"` and +/// relabelled afterwards. Both name the same instant, so the relabel is a +/// metadata change. +fn decode_type(data_type: &DataType) -> DataType { + match data_type { + DataType::Timestamp(unit, Some(_)) => DataType::Timestamp(*unit, Some(UTC_OFFSET.into())), + DataType::List(item) => DataType::List(Arc::new( + item.as_ref() + .clone() + .with_data_type(decode_type(item.data_type())), + )), + DataType::Struct(children) => DataType::Struct( + children + .iter() + .map(|f| { + f.as_ref() + .clone() + .with_data_type(decode_type(f.data_type())) + }) + .collect(), + ), + other => other.clone(), + } +} + +fn to_target(batch: RecordBatch, target: &SchemaRef) -> Result<RecordBatch, ArrowError> { + let columns = batch + .columns() + .iter() + .zip(target.fields()) + .map(|(column, field)| { + if column.data_type() == field.data_type() { + Ok(column.clone()) + } else { + cast(column, field.data_type()) + } + }) + .collect::<Result<Vec<_>, _>>()?; + RecordBatch::try_new(target.clone(), columns) +} + +fn decode(schema: &SchemaRef, rows: &[OwnedValue]) -> Result<RecordBatch, ArrowError> { + let mut decoder = ReaderBuilder::new(schema.clone()) + .with_batch_size(rows.len().max(1)) + .build_decoder()?; + decoder.serialize(rows)?; + decoder + .flush()? + .ok_or_else(|| ArrowError::JsonError("decoder produced no rows".to_owned())) +} + +// ─── Raw mode ──────────────────────────────────────────────────────────────── + +fn encode_raw( + ctx: &RunContext<'_>, + kind: RawKind, + payload_field: &FieldRef, + messages: Vec<ConsumedMessage>, + rejected: &mut Vec<Rejected>, +) -> Result<Option<(RecordBatch, Vec<u64>)>, ArrowError> { + let mut kept = Vec::with_capacity(messages.len()); + let payload_column: ArrayRef = match kind { + RawKind::Bytes => { + let mut values = Vec::with_capacity(messages.len()); + for message in messages { + match message.payload.try_to_bytes() { + Ok(bytes) => { + values.push(bytes); + kept.push(message); + } + Err(e) => rejected.push(Rejected { + offset: message.offset, + reason: e.to_string(), + }), + } + } + Arc::new(BinaryArray::from_iter_values(values)) + } + RawKind::Json | RawKind::String => { + let mut values = Vec::with_capacity(messages.len()); + for mut message in messages { + let payload = std::mem::replace(&mut message.payload, Payload::Raw(Vec::new())); + match raw_text(payload, kind) { + Ok(text) => { + values.push(text); + kept.push(message); + } + Err(reason) => rejected.push(Rejected { + offset: message.offset, + reason, + }), + } + } + Arc::new(StringArray::from_iter_values(values)) + } + }; + if kept.is_empty() { + return Ok(None); + } + let kept: Vec<&ConsumedMessage> = kept.iter().collect(); + let offsets = kept.iter().map(|m| m.offset).collect(); + let mut fields = vec![payload_field.clone()]; + let mut columns = vec![payload_column]; + append_metadata(ctx, &kept, &mut fields, &mut columns)?; + let batch = RecordBatch::try_new(Arc::new(Schema::new(fields)), columns)?; + Ok(Some((batch, offsets))) +} + +/// Payload as text for a JSON or STRING column. A JSON column only takes +/// text that parses as JSON. +fn raw_text(payload: Payload, kind: RawKind) -> Result<String, String> { + let text = match payload { + Payload::Json(value) => { + return simd_json::to_string(&value).map_err(|e| format!("cannot serialize JSON: {e}")); + } + Payload::Text(text) | Payload::Proto(text) => text, + Payload::Raw(bytes) => { + String::from_utf8(bytes).map_err(|_| "raw payload is not valid UTF-8".to_owned())? + } + other => { + return Err(format!( + "{} payload cannot be written to a text column", + payload_kind(&other) + )); + } + }; + if kind == RawKind::Json { + let mut bytes = text.as_bytes().to_vec(); + simd_json::to_owned_value(&mut bytes) + .map_err(|e| format!("payload is not valid JSON: {e}"))?; + } + Ok(text) +} + +// ─── Metadata ──────────────────────────────────────────────────────────────── + +fn append_metadata( + ctx: &RunContext<'_>, + messages: &[&ConsumedMessage], + fields: &mut Vec<FieldRef>, + columns: &mut Vec<ArrayRef>, +) -> Result<(), ArrowError> { + let rows = messages.len(); + for (meta, field) in &ctx.layout.metadata { + let column: ArrayRef = match meta { + MetaColumn::Stream => Arc::new(StringArray::from_iter_values(std::iter::repeat_n( + ctx.topic.stream.as_str(), + rows, + ))), + MetaColumn::Topic => Arc::new(StringArray::from_iter_values(std::iter::repeat_n( + ctx.topic.topic.as_str(), + rows, + ))), + MetaColumn::PartitionId => Arc::new(Int64Array::from_iter_values(std::iter::repeat_n( + i64::from(ctx.messages.partition_id), + rows, + ))), + MetaColumn::Offset => Arc::new(Int64Array::from_iter_values( + messages.iter().map(|m| saturating_i64(m.offset)), + )), + MetaColumn::Timestamp => Arc::new( + TimestampMicrosecondArray::from_iter_values( + messages.iter().map(|m| saturating_i64(m.timestamp)), + ) + .with_timezone(UTC), + ), + MetaColumn::Id => Arc::new(StringArray::from_iter_values( + messages.iter().map(|m| m.id.to_string()), + )), + MetaColumn::Headers => { + let required = !field.is_nullable(); + Arc::new(StringArray::from_iter( + messages.iter().map(|m| headers_json(m, required)), + )) + } + }; + fields.push(field.clone()); + columns.push(column); + } + Ok(()) +} + +/// Headers as a JSON object. Text values are written as strings, binary +/// values as `{"data": <base64>, "iggy_header_encoding": "base64"}`, the same +/// shape `http_sink` uses. +fn headers_json(message: &ConsumedMessage, required: bool) -> Option<String> { + let headers = message.headers.as_ref().filter(|h| !h.is_empty()); + let Some(headers) = headers else { + return required.then(|| "{}".to_owned()); + }; + let mut object = serde_json::Map::with_capacity(headers.len()); + for (key, value) in headers { + let encoded = match value.as_raw() { + Ok(raw) => serde_json::json!({ + "data": BASE64.encode(raw), + "iggy_header_encoding": HEADER_ENCODING_BASE64, + }), + Err(_) => serde_json::Value::String(value.to_string_value()), + }; + object.insert(key.to_string_value(), encoded); + } + Some(serde_json::Value::Object(object).to_string()) +} + +// ─── Request sizing ────────────────────────────────────────────────────────── + +/// Encode `batch` and halve it until every piece fits `max_bytes`. A single +/// row that does not fit is rejected. +fn split_into_chunks( + batch: RecordBatch, + offsets: Vec<u64>, + max_bytes: usize, + chunks: &mut Vec<Chunk>, + rejected: &mut Vec<Rejected>, +) -> Result<(), ArrowError> { + let (schema_bytes, batch_bytes) = encode_ipc(&batch)?; Review Comment: warning: `split_into_chunks` encodes the batch, discards those bytes when the batch exceeds the budget, and re-encodes both halves, so every split level encodes its part again. Split on an estimated row size and encode each chunk once. ########## core/connectors/sinks/bigquery_sink/src/encode.rs: ########## @@ -0,0 +1,1128 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Turns a run of `ConsumedMessage`s into Arrow record batches that fit in +//! one `AppendRows` request each. +//! +//! A message that cannot become a row (wrong payload type, a value that does +//! not fit its column, a row larger than the request budget) is rejected on +//! its own and reported back with its offset. The rest of the batch is still +//! written. +//! +//! Mapped mode decodes JSON objects with `arrow-json`. The fast path decodes +//! the whole run at once. When that fails, the run is decoded again row by +//! row so that only the offending rows are dropped. +//! +//! The writer schema of a request only contains the columns that at least one +//! row in the run sets. Columns no row mentions are left out, so BigQuery +//! fills them according to `missing_value` (the column default, or NULL). + +use crate::schema::{BqType, Column, MetaColumn, Mode, RawKind, TableLayout, UTC}; +use arrow::array::{ + ArrayRef, BinaryArray, Int64Array, RecordBatch, StringArray, TimestampMicrosecondArray, +}; +use arrow::compute::{cast, concat_batches}; +use arrow::datatypes::{DataType, FieldRef, Schema, SchemaRef}; +use arrow::error::ArrowError; +use arrow::ipc::writer::{ + CompressionContext, DictionaryTracker, IpcDataGenerator, IpcWriteOptions, write_message, +}; +use arrow::json::ReaderBuilder; +use base64::Engine; +use base64::engine::general_purpose::STANDARD as BASE64; +use iggy_connector_sdk::{ConsumedMessage, MessagesMetadata, Payload, TopicMetadata}; +use simd_json::{OwnedValue, StaticNode}; +use std::fmt::Write as _; +use std::sync::Arc; + +const HEADER_ENCODING_BASE64: &str = "base64"; +const UTC_OFFSET: &str = "+00:00"; + +/// One `AppendRows` request worth of rows. +#[derive(Debug)] +pub(crate) struct Chunk { + pub batch: RecordBatch, + /// Iggy offset of each row, same order as `batch`. + pub offsets: Vec<u64>, + pub schema_bytes: Vec<u8>, + pub batch_bytes: Vec<u8>, +} + +/// A message that was not turned into a row. +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct Rejected { + pub offset: u64, + pub reason: String, +} + +#[derive(Debug, Default)] +pub(crate) struct Encoded { + pub chunks: Vec<Chunk>, + pub rejected: Vec<Rejected>, +} + +/// Per-run context shared by every row. +pub(crate) struct RunContext<'a> { + pub layout: &'a TableLayout, + pub topic: &'a TopicMetadata, + pub messages: &'a MessagesMetadata, + pub max_request_bytes: usize, +} + +/// Encode a run of messages. Fails only on an internal Arrow error, which +/// means the batch as a whole cannot be built. Per-row problems end up in +/// `Encoded::rejected`. +pub(crate) fn encode( + ctx: &RunContext<'_>, + messages: Vec<ConsumedMessage>, +) -> Result<Encoded, ArrowError> { + let mut rejected = Vec::new(); + let batch = match &ctx.layout.raw { + Some((kind, field)) => encode_raw(ctx, *kind, field, messages, &mut rejected)?, + None => encode_mapped(ctx, messages, &mut rejected)?, + }; + let mut chunks = Vec::new(); + if let Some((batch, offsets)) = batch { + split_into_chunks( + batch, + offsets, + ctx.max_request_bytes, + &mut chunks, + &mut rejected, + )?; + } + Ok(Encoded { chunks, rejected }) +} + +/// Build a new chunk holding only the rows of `chunk` whose index is not in +/// `drop`. Used after BigQuery reports row-level errors. +pub(crate) fn without_rows(chunk: &Chunk, drop: &[usize]) -> Result<Option<Chunk>, ArrowError> { + let keep: Vec<bool> = (0..chunk.batch.num_rows()) + .map(|row| !drop.contains(&row)) + .collect(); + let offsets: Vec<u64> = chunk + .offsets + .iter() + .zip(&keep) + .filter_map(|(offset, keep)| keep.then_some(*offset)) + .collect(); + if offsets.is_empty() { + return Ok(None); + } + let mask = arrow::array::BooleanArray::from(keep); + let batch = arrow::compute::filter_record_batch(&chunk.batch, &mask)?; + let (schema_bytes, batch_bytes) = encode_ipc(&batch)?; + Ok(Some(Chunk { + batch, + offsets, + schema_bytes, + batch_bytes, + })) +} + +// ─── Mapped mode ───────────────────────────────────────────────────────────── + +fn encode_mapped( + ctx: &RunContext<'_>, + messages: Vec<ConsumedMessage>, + rejected: &mut Vec<Rejected>, +) -> Result<Option<(RecordBatch, Vec<u64>)>, ArrowError> { + let layout = ctx.layout; + let mut rows = Vec::with_capacity(messages.len()); + let mut kept = Vec::with_capacity(messages.len()); + for mut message in messages { + let payload = std::mem::replace(&mut message.payload, Payload::Raw(Vec::new())); + match json_object(payload).and_then(|mut row| { + normalize_row(&mut row, layout)?; + Ok(row) + }) { + Ok(row) => { + rows.push(row); + kept.push(message); + } + Err(reason) => rejected.push(Rejected { + offset: message.offset, + reason, + }), + } + } + if rows.is_empty() { + return Ok(None); + } + + let present = present_columns(&rows, layout); + if present.is_empty() && layout.metadata.is_empty() { + rejected.extend(kept.iter().map(|m| Rejected { + offset: m.offset, + reason: "payload has no field matching a table column".to_owned(), + })); + return Ok(None); + } + let payload_fields: Vec<FieldRef> = present.iter().map(|&i| layout.fields[i].clone()).collect(); + + let (decoded, survivors) = decode_rows(&payload_fields, &rows, kept.len(), |row, reason| { + rejected.push(Rejected { + offset: kept[row].offset, + reason, + }); + })?; + if survivors.is_empty() { + return Ok(None); + } + let kept: Vec<&ConsumedMessage> = survivors.iter().map(|&row| &kept[row]).collect(); + let offsets = kept.iter().map(|m| m.offset).collect(); + + let mut fields = payload_fields; + let mut columns: Vec<ArrayRef> = decoded.map(|b| b.columns().to_vec()).unwrap_or_default(); + append_metadata(ctx, &kept, &mut fields, &mut columns)?; + let batch = RecordBatch::try_new(Arc::new(Schema::new(fields)), columns)?; + Ok(Some((batch, offsets))) +} + +/// Take the JSON object out of a payload. Text and proto text must parse as +/// a JSON object. +fn json_object(payload: Payload) -> Result<OwnedValue, String> { + let value = match payload { + Payload::Json(value) => value, + Payload::Text(text) | Payload::Proto(text) => { + let mut bytes = text.into_bytes(); + simd_json::to_owned_value(&mut bytes) + .map_err(|e| format!("payload text is not valid JSON: {e}"))? + } + other => { + return Err(format!( + "{} payload is not supported in mapped mode", + payload_kind(&other) + )); + } + }; + match value { + OwnedValue::Object(_) => Ok(value), + _ => Err("payload is not a JSON object".to_owned()), + } +} + +/// Prepare one row for `arrow-json`: +/// - a REQUIRED column without a default must be present and non-null, +/// - REPEATED columns that are absent or null become `[]`, +/// - JSON column values become their JSON text, +/// - BYTES column values are base64 in JSON and become hex, which is what +/// `arrow-json` decodes into binary. +fn normalize_row(row: &mut OwnedValue, layout: &TableLayout) -> Result<(), String> { + let OwnedValue::Object(object) = row else { + return Err("payload is not a JSON object".to_owned()); + }; + for column in &layout.columns { + let value = object.get_mut(column.name.as_str()); + match (value, column.mode) { + (None | Some(OwnedValue::Static(StaticNode::Null)), Mode::Repeated) => { + object.insert(column.name.clone(), empty_array()); + } + (None | Some(OwnedValue::Static(StaticNode::Null)), Mode::Required) + if !column.has_default => + { + return Err(format!( + "missing value for REQUIRED column '{}'", + column.name + )); + } + (Some(value), _) => normalize_value(value, column)?, + (None, _) => {} + } + } + Ok(()) +} + +fn normalize_value(value: &mut OwnedValue, column: &Column) -> Result<(), String> { + if column.mode == Mode::Repeated { + let OwnedValue::Array(items) = value else { + return Err(format!("column '{}' expects an array", column.name)); + }; + for item in items.iter_mut() { + normalize_scalar(item, column)?; + } + return Ok(()); + } + normalize_scalar(value, column) +} + +fn normalize_scalar(value: &mut OwnedValue, column: &Column) -> Result<(), String> { + if matches!(value, OwnedValue::Static(StaticNode::Null)) { + return Ok(()); + } + match &column.bq_type { + BqType::Json => { + let text = simd_json::to_string(value) + .map_err(|e| format!("column '{}': cannot serialize JSON: {e}", column.name))?; + *value = OwnedValue::String(text); + } + BqType::Bytes => { + let OwnedValue::String(text) = value else { + return Err(format!("column '{}' expects a base64 string", column.name)); + }; + let bytes = BASE64 + .decode(text.as_bytes()) + .map_err(|e| format!("column '{}': invalid base64: {e}", column.name))?; + *value = OwnedValue::String(to_hex(&bytes)); + } + BqType::Record(children) => { + let OwnedValue::Object(object) = value else { + return Ok(()); + }; + for child in children { + match (object.get_mut(child.name.as_str()), child.mode) { + (None | Some(OwnedValue::Static(StaticNode::Null)), Mode::Repeated) => { + object.insert(child.name.clone(), empty_array()); + } + (Some(child_value), _) => normalize_value(child_value, child)?, + (None, _) => {} + } + } + } + _ => {} + } + Ok(()) +} + +/// Indexes of the layout columns that at least one row sets, in table order. +fn present_columns(rows: &[OwnedValue], layout: &TableLayout) -> Vec<usize> { + let mut present = vec![false; layout.columns.len()]; + for row in rows { + if let OwnedValue::Object(object) = row { + for key in object.keys() { + if let Some(&index) = layout.index.get(key.as_str()) { + present[index] = true; + } + } + } + } + present + .iter() + .enumerate() + .filter_map(|(index, set)| set.then_some(index)) + .collect() +} + +/// Decode rows into one batch. Returns the batch (`None` when there are no +/// payload fields) and the indexes of the rows it holds. +fn decode_rows( + fields: &[FieldRef], + rows: &[OwnedValue], + row_count: usize, + mut on_reject: impl FnMut(usize, String), +) -> Result<(Option<RecordBatch>, Vec<usize>), ArrowError> { + if fields.is_empty() { + return Ok((None, (0..row_count).collect())); + } + let target: SchemaRef = Arc::new(Schema::new(fields.to_vec())); + let schema: SchemaRef = Arc::new(Schema::new( + fields + .iter() + .map(|f| { + f.as_ref() + .clone() + .with_data_type(decode_type(f.data_type())) + }) + .collect::<Vec<_>>(), + )); + if let Ok(batch) = decode(&schema, rows) { + return Ok((Some(to_target(batch, &target)?), (0..row_count).collect())); + } + + let mut batches = Vec::with_capacity(rows.len()); + let mut survivors = Vec::with_capacity(rows.len()); + for (index, row) in rows.iter().enumerate() { + match decode(&schema, std::slice::from_ref(row)) { + Ok(batch) => { + batches.push(batch); + survivors.push(index); + } + Err(e) => on_reject(index, format!("row does not match the table schema: {e}")), + } + } + if batches.is_empty() { + return Ok((None, survivors)); + } + let batch = concat_batches(&schema, &batches)?; + Ok((Some(to_target(batch, &target)?), survivors)) +} + +/// The type `arrow-json` decodes into. Without the `chrono-tz` feature it +/// only parses offset timezones, so `"UTC"` is decoded as `"+00:00"` and +/// relabelled afterwards. Both name the same instant, so the relabel is a +/// metadata change. +fn decode_type(data_type: &DataType) -> DataType { + match data_type { + DataType::Timestamp(unit, Some(_)) => DataType::Timestamp(*unit, Some(UTC_OFFSET.into())), + DataType::List(item) => DataType::List(Arc::new( + item.as_ref() + .clone() + .with_data_type(decode_type(item.data_type())), + )), + DataType::Struct(children) => DataType::Struct( + children + .iter() + .map(|f| { + f.as_ref() + .clone() + .with_data_type(decode_type(f.data_type())) + }) + .collect(), + ), + other => other.clone(), + } +} + +fn to_target(batch: RecordBatch, target: &SchemaRef) -> Result<RecordBatch, ArrowError> { + let columns = batch + .columns() + .iter() + .zip(target.fields()) + .map(|(column, field)| { + if column.data_type() == field.data_type() { + Ok(column.clone()) + } else { + cast(column, field.data_type()) + } + }) + .collect::<Result<Vec<_>, _>>()?; + RecordBatch::try_new(target.clone(), columns) +} + +fn decode(schema: &SchemaRef, rows: &[OwnedValue]) -> Result<RecordBatch, ArrowError> { + let mut decoder = ReaderBuilder::new(schema.clone()) + .with_batch_size(rows.len().max(1)) + .build_decoder()?; + decoder.serialize(rows)?; + decoder + .flush()? + .ok_or_else(|| ArrowError::JsonError("decoder produced no rows".to_owned())) +} + +// ─── Raw mode ──────────────────────────────────────────────────────────────── + +fn encode_raw( + ctx: &RunContext<'_>, + kind: RawKind, + payload_field: &FieldRef, + messages: Vec<ConsumedMessage>, + rejected: &mut Vec<Rejected>, +) -> Result<Option<(RecordBatch, Vec<u64>)>, ArrowError> { + let mut kept = Vec::with_capacity(messages.len()); + let payload_column: ArrayRef = match kind { + RawKind::Bytes => { + let mut values = Vec::with_capacity(messages.len()); + for message in messages { + match message.payload.try_to_bytes() { + Ok(bytes) => { + values.push(bytes); + kept.push(message); + } + Err(e) => rejected.push(Rejected { + offset: message.offset, + reason: e.to_string(), + }), + } + } + Arc::new(BinaryArray::from_iter_values(values)) + } + RawKind::Json | RawKind::String => { + let mut values = Vec::with_capacity(messages.len()); + for mut message in messages { + let payload = std::mem::replace(&mut message.payload, Payload::Raw(Vec::new())); + match raw_text(payload, kind) { + Ok(text) => { + values.push(text); + kept.push(message); + } + Err(reason) => rejected.push(Rejected { + offset: message.offset, + reason, + }), + } + } + Arc::new(StringArray::from_iter_values(values)) + } + }; + if kept.is_empty() { + return Ok(None); + } + let kept: Vec<&ConsumedMessage> = kept.iter().collect(); + let offsets = kept.iter().map(|m| m.offset).collect(); + let mut fields = vec![payload_field.clone()]; + let mut columns = vec![payload_column]; + append_metadata(ctx, &kept, &mut fields, &mut columns)?; + let batch = RecordBatch::try_new(Arc::new(Schema::new(fields)), columns)?; + Ok(Some((batch, offsets))) +} + +/// Payload as text for a JSON or STRING column. A JSON column only takes +/// text that parses as JSON. +fn raw_text(payload: Payload, kind: RawKind) -> Result<String, String> { + let text = match payload { + Payload::Json(value) => { + return simd_json::to_string(&value).map_err(|e| format!("cannot serialize JSON: {e}")); + } + Payload::Text(text) | Payload::Proto(text) => text, + Payload::Raw(bytes) => { + String::from_utf8(bytes).map_err(|_| "raw payload is not valid UTF-8".to_owned())? + } + other => { + return Err(format!( + "{} payload cannot be written to a text column", + payload_kind(&other) + )); + } + }; + if kind == RawKind::Json { + let mut bytes = text.as_bytes().to_vec(); + simd_json::to_owned_value(&mut bytes) + .map_err(|e| format!("payload is not valid JSON: {e}"))?; + } + Ok(text) +} + +// ─── Metadata ──────────────────────────────────────────────────────────────── + +fn append_metadata( + ctx: &RunContext<'_>, + messages: &[&ConsumedMessage], + fields: &mut Vec<FieldRef>, + columns: &mut Vec<ArrayRef>, +) -> Result<(), ArrowError> { + let rows = messages.len(); + for (meta, field) in &ctx.layout.metadata { + let column: ArrayRef = match meta { + MetaColumn::Stream => Arc::new(StringArray::from_iter_values(std::iter::repeat_n( + ctx.topic.stream.as_str(), + rows, + ))), + MetaColumn::Topic => Arc::new(StringArray::from_iter_values(std::iter::repeat_n( + ctx.topic.topic.as_str(), + rows, + ))), + MetaColumn::PartitionId => Arc::new(Int64Array::from_iter_values(std::iter::repeat_n( + i64::from(ctx.messages.partition_id), + rows, + ))), + MetaColumn::Offset => Arc::new(Int64Array::from_iter_values( + messages.iter().map(|m| saturating_i64(m.offset)), + )), + MetaColumn::Timestamp => Arc::new( + TimestampMicrosecondArray::from_iter_values( + messages.iter().map(|m| saturating_i64(m.timestamp)), + ) + .with_timezone(UTC), + ), + MetaColumn::Id => Arc::new(StringArray::from_iter_values( + messages.iter().map(|m| m.id.to_string()), Review Comment: nit: `append_metadata` allocates one `String` per message id for the `iggy_id` column, and `include_metadata` defaults to true. Build that column with a `StringBuilder` instead. ########## core/connectors/sinks/bigquery_sink/src/lib.rs: ########## @@ -0,0 +1,457 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Google BigQuery sink connector. +//! +//! Writes Iggy messages into a BigQuery table through the Storage Write API +//! `_default` stream, with rows encoded as Arrow record batches. +//! +//! Module map: +//! - [`schema`]: BigQuery table schema (from `tables.get`) to Arrow schema. +//! - [`encode`]: `ConsumedMessage`s to Arrow record batches, split by size. +//! - [`client`]: credentials, `tables.get`, and `AppendRows` calls. +//! - [`error`]: gRPC status classification into retryable and permanent. +//! - [`sink`]: the `Sink` trait implementation. + +use humantime::Duration as HumanDuration; +use iggy_connector_sdk::{Error, sink_connector}; +use secrecy::SecretString; +use serde::Deserialize; +use std::str::FromStr; +use std::sync::atomic::AtomicU64; +use std::time::Duration; +use tracing::warn; + +mod client; +mod encode; +mod error; +mod schema; +mod sink; + +sink_connector!(BigQuerySink); + +const DEFAULT_PAYLOAD_COLUMN: &str = "payload"; +const DEFAULT_MAX_REQUEST_BYTES: usize = 8 * 1024 * 1024; +/// Hard ceiling below the 10 MB `AppendRows` limit, leaving room for the +/// request envelope (stream name, writer schema, protobuf framing). +const MAX_REQUEST_BYTES_CEILING: usize = 9 * 1024 * 1024; +/// Smallest accepted request budget. Anything lower makes most real rows +/// unwritable. +const MIN_REQUEST_BYTES: usize = 64 * 1024; +const DEFAULT_MAX_RETRIES: u32 = 3; +const DEFAULT_RETRY_DELAY: &str = "1s"; +const DEFAULT_MAX_RETRY_DELAY: &str = "30s"; +const DEFAULT_TIMEOUT: &str = "30s"; + +/// Plugin configuration, deserialized from `[plugin_config]`. +/// +/// `Deserialize` only. Nothing re-serializes a plugin config, and leaving +/// `Serialize` off is what keeps `credentials_json` unserializable. +#[derive(Debug, Clone, Deserialize)] +pub struct BigQuerySinkConfig { + pub project_id: String, + pub dataset: String, + pub table: String, + /// `mapped` (default) or `raw`. + pub mode: Option<WriteMode>, + /// Path to a service account key file. Mutually exclusive with + /// `credentials_json`. When neither is set, Application Default + /// Credentials are used. + pub credentials_path: Option<String>, + /// Inline service account key JSON. + pub credentials_json: Option<SecretString>, + /// Column holding the payload in `raw` mode (default `payload`). + pub payload_column: Option<String>, + /// Write `iggy_stream`, `iggy_topic`, `iggy_partition_id`, `iggy_offset`, + /// `iggy_timestamp` and `iggy_id` (default `true`). + pub include_metadata: Option<bool>, + /// Write message headers into an `iggy_headers` column (default `false`). + pub include_headers: Option<bool>, + /// How BigQuery fills columns absent from a request: `default` (column + /// default value, else NULL) or `null`. + pub missing_value: Option<MissingValue>, + /// Upper bound for one `AppendRows` request (default 8 MiB). + pub max_request_bytes: Option<usize>, + /// Total attempts per append, including the first (default 3). + pub max_retries: Option<u32>, + /// Base retry delay, e.g. `"1s"` (default `1s`). + pub retry_delay: Option<String>, + /// Upper bound for a single retry delay (default `30s`). + pub max_retry_delay: Option<String>, + /// Per-request timeout for `tables.get` and `AppendRows` (default `30s`). + pub timeout: Option<String>, + /// Log every batch at info level instead of debug (default `false`). + pub verbose_logging: Option<bool>, + /// REST endpoint override, e.g. `http://127.0.0.1:9050`. Must be set + /// together with `grpc_endpoint`; when both are set no credentials are + /// used. Intended for local fakes and emulators. + pub endpoint: Option<String>, + /// gRPC `host:port` override for the Storage Write API. + pub grpc_endpoint: Option<String>, +} + +/// How message payloads map onto table columns. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum WriteMode { + /// Top-level JSON fields map to table columns by name. + #[default] + Mapped, + /// The whole payload goes into one column. + Raw, +} + +/// Maps to `AppendRowsRequest.default_missing_value_interpretation`. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum MissingValue { + #[default] + Default, + Null, +} + +impl MissingValue { + /// Wire value of `AppendRowsRequest.MissingValueInterpretation`. + pub(crate) fn as_proto(self) -> i32 { + match self { + MissingValue::Null => 1, + MissingValue::Default => 2, + } + } +} + +/// Resolved, validated runtime settings. Built once in `new()`. +#[derive(Debug, Clone)] +pub(crate) struct Settings { + pub mode: WriteMode, + pub payload_column: String, + pub include_metadata: bool, + pub include_headers: bool, + pub missing_value: MissingValue, + pub max_request_bytes: usize, + pub max_retries: u32, + pub retry_delay: Duration, + pub max_retry_delay: Duration, + pub timeout: Duration, + pub verbose: bool, +} + +/// Counters updated on the hot path without taking a lock. +#[derive(Debug, Default)] +pub(crate) struct Counters { + pub rows_written: AtomicU64, + pub rows_rejected: AtomicU64, + pub rows_failed: AtomicU64, +} + +#[derive(Debug)] +pub struct BigQuerySink { + id: u32, + config: BigQuerySinkConfig, + settings: Settings, + /// `project.dataset.table`, safe to log. + target: String, + client: Option<client::BigQueryClient>, Review Comment: simplification: `BigQuerySink` stores `client` and `layout` in two `Option` fields, but `try_open` sets both and `close` clears both, so neither can exist alone. Store one `Option<(BigQueryClient, TableLayout)>` instead. ########## core/connectors/sinks/bigquery_sink/src/sink.rs: ########## @@ -0,0 +1,299 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::BigQuerySink; +use crate::client::{AppendOutcome, BigQueryClient}; +use crate::encode::{self, Chunk, RunContext}; +use crate::error::{AppendError, TableError}; +use crate::schema::{TableLayout, parse_table_schema}; +use async_trait::async_trait; +use iggy_connector_sdk::retry::{RetryPolicy, retry_async}; +use iggy_connector_sdk::{ConsumedMessage, Error, MessagesMetadata, Sink, TopicMetadata}; +use std::sync::atomic::Ordering; +use tracing::{debug, error, info, warn}; + +#[async_trait] +impl Sink for BigQuerySink { + async fn open(&mut self) -> Result<(), Error> { + info!( + "Opening BigQuery sink connector ID: {} for table {} (mode: {:?})", + self.id, self.target, self.settings.mode + ); + // The SDK turns a failed open into a bare status code, so this is the + // only place the reason reaches the runtime log. + self.try_open().await.inspect_err(|e| { + error!( + "Failed to open BigQuery sink connector ID: {} for table {}: {e}", + self.id, self.target + ); + }) + } + + async fn consume( + &self, + topic_metadata: &TopicMetadata, + messages_metadata: MessagesMetadata, + messages: Vec<ConsumedMessage>, + ) -> Result<(), Error> { + if messages.is_empty() { + return Ok(()); + } + let (Some(client), Some(layout)) = (&self.client, &self.layout) else { + return Err(Error::InitError("BigQuery sink is not open".into())); + }; + let received = messages.len(); + debug!( + "BigQuery sink ID: {} received {received} messages from {}/{} partition {} current_offset {}", + self.id, + topic_metadata.stream, + topic_metadata.topic, + messages_metadata.partition_id, + messages_metadata.current_offset + ); + + let ctx = RunContext { + layout, + topic: topic_metadata, + messages: &messages_metadata, + max_request_bytes: self.settings.max_request_bytes, + }; + let encoded = encode::encode(&ctx, messages).map_err(|e| { + self.counters + .rows_failed + .fetch_add(received as u64, Ordering::Relaxed); + error!( + "BigQuery sink ID: {} cannot encode {received} messages from {}/{} partition {}: {e}", + self.id, topic_metadata.stream, topic_metadata.topic, messages_metadata.partition_id + ); + Error::Serialization(e.to_string()) + })?; + + for rejected in &encoded.rejected { + warn!( + "BigQuery sink ID: {} dropped message stream={} topic={} partition={} offset={}: {}", + self.id, + topic_metadata.stream, + topic_metadata.topic, + messages_metadata.partition_id, + rejected.offset, + rejected.reason + ); + } + self.counters + .rows_rejected + .fetch_add(encoded.rejected.len() as u64, Ordering::Relaxed); + + let mut written = 0u64; + let mut last_error = None; + for chunk in &encoded.chunks { + match self + .write_chunk(client, chunk, topic_metadata, &messages_metadata) + .await + { + Ok(rows) => written += rows, + Err(e) => last_error = Some(e), + } + } + self.counters + .rows_written + .fetch_add(written, Ordering::Relaxed); + + let summary = format!( Review Comment: nit: `consume` builds the summary string before the log level test, so every batch allocates it even when both levels are off. Pass the fields to `info!` and `debug!` directly. ########## core/connectors/sinks/bigquery_sink/tests/bigquery_sink.rs: ########## @@ -0,0 +1,424 @@ +// Licensed to the Apache Software Foundation (ASF) under one Review Comment: warning: No test under `core/integration/tests/connectors/` runs this sink against a real backend. The repository requires an integration test for every sink that talks to an external service. Add a BigQuery fixture and test, or get maintainer sign-off for the gap. ########## core/connectors/sinks/bigquery_sink/src/encode.rs: ########## @@ -0,0 +1,1128 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Turns a run of `ConsumedMessage`s into Arrow record batches that fit in +//! one `AppendRows` request each. +//! +//! A message that cannot become a row (wrong payload type, a value that does +//! not fit its column, a row larger than the request budget) is rejected on +//! its own and reported back with its offset. The rest of the batch is still +//! written. +//! +//! Mapped mode decodes JSON objects with `arrow-json`. The fast path decodes +//! the whole run at once. When that fails, the run is decoded again row by +//! row so that only the offending rows are dropped. +//! +//! The writer schema of a request only contains the columns that at least one +//! row in the run sets. Columns no row mentions are left out, so BigQuery +//! fills them according to `missing_value` (the column default, or NULL). + +use crate::schema::{BqType, Column, MetaColumn, Mode, RawKind, TableLayout, UTC}; +use arrow::array::{ + ArrayRef, BinaryArray, Int64Array, RecordBatch, StringArray, TimestampMicrosecondArray, +}; +use arrow::compute::{cast, concat_batches}; +use arrow::datatypes::{DataType, FieldRef, Schema, SchemaRef}; +use arrow::error::ArrowError; +use arrow::ipc::writer::{ + CompressionContext, DictionaryTracker, IpcDataGenerator, IpcWriteOptions, write_message, +}; +use arrow::json::ReaderBuilder; +use base64::Engine; +use base64::engine::general_purpose::STANDARD as BASE64; +use iggy_connector_sdk::{ConsumedMessage, MessagesMetadata, Payload, TopicMetadata}; +use simd_json::{OwnedValue, StaticNode}; +use std::fmt::Write as _; +use std::sync::Arc; + +const HEADER_ENCODING_BASE64: &str = "base64"; +const UTC_OFFSET: &str = "+00:00"; + +/// One `AppendRows` request worth of rows. +#[derive(Debug)] +pub(crate) struct Chunk { + pub batch: RecordBatch, + /// Iggy offset of each row, same order as `batch`. + pub offsets: Vec<u64>, + pub schema_bytes: Vec<u8>, + pub batch_bytes: Vec<u8>, +} + +/// A message that was not turned into a row. +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct Rejected { + pub offset: u64, + pub reason: String, +} + +#[derive(Debug, Default)] +pub(crate) struct Encoded { + pub chunks: Vec<Chunk>, + pub rejected: Vec<Rejected>, +} + +/// Per-run context shared by every row. +pub(crate) struct RunContext<'a> { + pub layout: &'a TableLayout, + pub topic: &'a TopicMetadata, + pub messages: &'a MessagesMetadata, + pub max_request_bytes: usize, +} + +/// Encode a run of messages. Fails only on an internal Arrow error, which +/// means the batch as a whole cannot be built. Per-row problems end up in +/// `Encoded::rejected`. +pub(crate) fn encode( + ctx: &RunContext<'_>, + messages: Vec<ConsumedMessage>, +) -> Result<Encoded, ArrowError> { + let mut rejected = Vec::new(); + let batch = match &ctx.layout.raw { + Some((kind, field)) => encode_raw(ctx, *kind, field, messages, &mut rejected)?, + None => encode_mapped(ctx, messages, &mut rejected)?, + }; + let mut chunks = Vec::new(); + if let Some((batch, offsets)) = batch { + split_into_chunks( + batch, + offsets, + ctx.max_request_bytes, + &mut chunks, + &mut rejected, + )?; + } + Ok(Encoded { chunks, rejected }) +} + +/// Build a new chunk holding only the rows of `chunk` whose index is not in +/// `drop`. Used after BigQuery reports row-level errors. +pub(crate) fn without_rows(chunk: &Chunk, drop: &[usize]) -> Result<Option<Chunk>, ArrowError> { + let keep: Vec<bool> = (0..chunk.batch.num_rows()) + .map(|row| !drop.contains(&row)) + .collect(); + let offsets: Vec<u64> = chunk + .offsets + .iter() + .zip(&keep) + .filter_map(|(offset, keep)| keep.then_some(*offset)) + .collect(); + if offsets.is_empty() { + return Ok(None); + } + let mask = arrow::array::BooleanArray::from(keep); + let batch = arrow::compute::filter_record_batch(&chunk.batch, &mask)?; + let (schema_bytes, batch_bytes) = encode_ipc(&batch)?; + Ok(Some(Chunk { + batch, + offsets, + schema_bytes, + batch_bytes, + })) +} + +// ─── Mapped mode ───────────────────────────────────────────────────────────── + +fn encode_mapped( + ctx: &RunContext<'_>, + messages: Vec<ConsumedMessage>, + rejected: &mut Vec<Rejected>, +) -> Result<Option<(RecordBatch, Vec<u64>)>, ArrowError> { + let layout = ctx.layout; + let mut rows = Vec::with_capacity(messages.len()); + let mut kept = Vec::with_capacity(messages.len()); + for mut message in messages { + let payload = std::mem::replace(&mut message.payload, Payload::Raw(Vec::new())); + match json_object(payload).and_then(|mut row| { + normalize_row(&mut row, layout)?; + Ok(row) + }) { + Ok(row) => { + rows.push(row); + kept.push(message); + } + Err(reason) => rejected.push(Rejected { + offset: message.offset, + reason, + }), + } + } + if rows.is_empty() { + return Ok(None); + } + + let present = present_columns(&rows, layout); + if present.is_empty() && layout.metadata.is_empty() { + rejected.extend(kept.iter().map(|m| Rejected { + offset: m.offset, + reason: "payload has no field matching a table column".to_owned(), + })); + return Ok(None); + } + let payload_fields: Vec<FieldRef> = present.iter().map(|&i| layout.fields[i].clone()).collect(); + + let (decoded, survivors) = decode_rows(&payload_fields, &rows, kept.len(), |row, reason| { + rejected.push(Rejected { + offset: kept[row].offset, + reason, + }); + })?; + if survivors.is_empty() { + return Ok(None); + } + let kept: Vec<&ConsumedMessage> = survivors.iter().map(|&row| &kept[row]).collect(); + let offsets = kept.iter().map(|m| m.offset).collect(); + + let mut fields = payload_fields; + let mut columns: Vec<ArrayRef> = decoded.map(|b| b.columns().to_vec()).unwrap_or_default(); + append_metadata(ctx, &kept, &mut fields, &mut columns)?; + let batch = RecordBatch::try_new(Arc::new(Schema::new(fields)), columns)?; + Ok(Some((batch, offsets))) +} + +/// Take the JSON object out of a payload. Text and proto text must parse as +/// a JSON object. +fn json_object(payload: Payload) -> Result<OwnedValue, String> { + let value = match payload { + Payload::Json(value) => value, + Payload::Text(text) | Payload::Proto(text) => { + let mut bytes = text.into_bytes(); + simd_json::to_owned_value(&mut bytes) + .map_err(|e| format!("payload text is not valid JSON: {e}"))? + } + other => { + return Err(format!( + "{} payload is not supported in mapped mode", + payload_kind(&other) + )); + } + }; + match value { + OwnedValue::Object(_) => Ok(value), + _ => Err("payload is not a JSON object".to_owned()), + } +} + +/// Prepare one row for `arrow-json`: +/// - a REQUIRED column without a default must be present and non-null, +/// - REPEATED columns that are absent or null become `[]`, +/// - JSON column values become their JSON text, +/// - BYTES column values are base64 in JSON and become hex, which is what +/// `arrow-json` decodes into binary. +fn normalize_row(row: &mut OwnedValue, layout: &TableLayout) -> Result<(), String> { + let OwnedValue::Object(object) = row else { + return Err("payload is not a JSON object".to_owned()); + }; + for column in &layout.columns { + let value = object.get_mut(column.name.as_str()); + match (value, column.mode) { + (None | Some(OwnedValue::Static(StaticNode::Null)), Mode::Repeated) => { + object.insert(column.name.clone(), empty_array()); + } + (None | Some(OwnedValue::Static(StaticNode::Null)), Mode::Required) + if !column.has_default => + { + return Err(format!( + "missing value for REQUIRED column '{}'", + column.name + )); + } + (Some(value), _) => normalize_value(value, column)?, + (None, _) => {} + } + } + Ok(()) +} + +fn normalize_value(value: &mut OwnedValue, column: &Column) -> Result<(), String> { + if column.mode == Mode::Repeated { + let OwnedValue::Array(items) = value else { + return Err(format!("column '{}' expects an array", column.name)); + }; + for item in items.iter_mut() { + normalize_scalar(item, column)?; + } + return Ok(()); + } + normalize_scalar(value, column) +} + +fn normalize_scalar(value: &mut OwnedValue, column: &Column) -> Result<(), String> { + if matches!(value, OwnedValue::Static(StaticNode::Null)) { + return Ok(()); + } + match &column.bq_type { + BqType::Json => { + let text = simd_json::to_string(value) + .map_err(|e| format!("column '{}': cannot serialize JSON: {e}", column.name))?; + *value = OwnedValue::String(text); + } + BqType::Bytes => { + let OwnedValue::String(text) = value else { + return Err(format!("column '{}' expects a base64 string", column.name)); + }; + let bytes = BASE64 + .decode(text.as_bytes()) + .map_err(|e| format!("column '{}': invalid base64: {e}", column.name))?; + *value = OwnedValue::String(to_hex(&bytes)); + } + BqType::Record(children) => { + let OwnedValue::Object(object) = value else { + return Ok(()); + }; + for child in children { + match (object.get_mut(child.name.as_str()), child.mode) { + (None | Some(OwnedValue::Static(StaticNode::Null)), Mode::Repeated) => { + object.insert(child.name.clone(), empty_array()); + } + (Some(child_value), _) => normalize_value(child_value, child)?, + (None, _) => {} + } + } + } + _ => {} + } + Ok(()) +} + +/// Indexes of the layout columns that at least one row sets, in table order. +fn present_columns(rows: &[OwnedValue], layout: &TableLayout) -> Vec<usize> { + let mut present = vec![false; layout.columns.len()]; + for row in rows { + if let OwnedValue::Object(object) = row { + for key in object.keys() { + if let Some(&index) = layout.index.get(key.as_str()) { + present[index] = true; + } + } + } + } + present + .iter() + .enumerate() + .filter_map(|(index, set)| set.then_some(index)) + .collect() +} + +/// Decode rows into one batch. Returns the batch (`None` when there are no +/// payload fields) and the indexes of the rows it holds. +fn decode_rows( + fields: &[FieldRef], + rows: &[OwnedValue], + row_count: usize, + mut on_reject: impl FnMut(usize, String), +) -> Result<(Option<RecordBatch>, Vec<usize>), ArrowError> { + if fields.is_empty() { + return Ok((None, (0..row_count).collect())); + } + let target: SchemaRef = Arc::new(Schema::new(fields.to_vec())); + let schema: SchemaRef = Arc::new(Schema::new( + fields + .iter() + .map(|f| { + f.as_ref() + .clone() + .with_data_type(decode_type(f.data_type())) + }) + .collect::<Vec<_>>(), + )); + if let Ok(batch) = decode(&schema, rows) { + return Ok((Some(to_target(batch, &target)?), (0..row_count).collect())); + } + + let mut batches = Vec::with_capacity(rows.len()); + let mut survivors = Vec::with_capacity(rows.len()); + for (index, row) in rows.iter().enumerate() { + match decode(&schema, std::slice::from_ref(row)) { Review Comment: warning: When one row fails, `decode_rows` decodes every row on its own and `concat_batches` copies every survivor again. Bisect the failing run to find the bad rows, then decode only those. ########## core/connectors/sinks/bigquery_sink/src/encode.rs: ########## @@ -0,0 +1,1128 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Turns a run of `ConsumedMessage`s into Arrow record batches that fit in +//! one `AppendRows` request each. +//! +//! A message that cannot become a row (wrong payload type, a value that does +//! not fit its column, a row larger than the request budget) is rejected on +//! its own and reported back with its offset. The rest of the batch is still +//! written. +//! +//! Mapped mode decodes JSON objects with `arrow-json`. The fast path decodes +//! the whole run at once. When that fails, the run is decoded again row by +//! row so that only the offending rows are dropped. +//! +//! The writer schema of a request only contains the columns that at least one +//! row in the run sets. Columns no row mentions are left out, so BigQuery +//! fills them according to `missing_value` (the column default, or NULL). + +use crate::schema::{BqType, Column, MetaColumn, Mode, RawKind, TableLayout, UTC}; +use arrow::array::{ + ArrayRef, BinaryArray, Int64Array, RecordBatch, StringArray, TimestampMicrosecondArray, +}; +use arrow::compute::{cast, concat_batches}; +use arrow::datatypes::{DataType, FieldRef, Schema, SchemaRef}; +use arrow::error::ArrowError; +use arrow::ipc::writer::{ + CompressionContext, DictionaryTracker, IpcDataGenerator, IpcWriteOptions, write_message, +}; +use arrow::json::ReaderBuilder; +use base64::Engine; +use base64::engine::general_purpose::STANDARD as BASE64; +use iggy_connector_sdk::{ConsumedMessage, MessagesMetadata, Payload, TopicMetadata}; +use simd_json::{OwnedValue, StaticNode}; +use std::fmt::Write as _; +use std::sync::Arc; + +const HEADER_ENCODING_BASE64: &str = "base64"; +const UTC_OFFSET: &str = "+00:00"; + +/// One `AppendRows` request worth of rows. +#[derive(Debug)] +pub(crate) struct Chunk { + pub batch: RecordBatch, + /// Iggy offset of each row, same order as `batch`. + pub offsets: Vec<u64>, + pub schema_bytes: Vec<u8>, + pub batch_bytes: Vec<u8>, +} + +/// A message that was not turned into a row. +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct Rejected { + pub offset: u64, + pub reason: String, +} + +#[derive(Debug, Default)] +pub(crate) struct Encoded { + pub chunks: Vec<Chunk>, + pub rejected: Vec<Rejected>, +} + +/// Per-run context shared by every row. +pub(crate) struct RunContext<'a> { + pub layout: &'a TableLayout, + pub topic: &'a TopicMetadata, + pub messages: &'a MessagesMetadata, + pub max_request_bytes: usize, +} + +/// Encode a run of messages. Fails only on an internal Arrow error, which +/// means the batch as a whole cannot be built. Per-row problems end up in +/// `Encoded::rejected`. +pub(crate) fn encode( + ctx: &RunContext<'_>, + messages: Vec<ConsumedMessage>, +) -> Result<Encoded, ArrowError> { + let mut rejected = Vec::new(); + let batch = match &ctx.layout.raw { + Some((kind, field)) => encode_raw(ctx, *kind, field, messages, &mut rejected)?, + None => encode_mapped(ctx, messages, &mut rejected)?, + }; + let mut chunks = Vec::new(); + if let Some((batch, offsets)) = batch { + split_into_chunks( + batch, + offsets, + ctx.max_request_bytes, + &mut chunks, + &mut rejected, + )?; + } + Ok(Encoded { chunks, rejected }) +} + +/// Build a new chunk holding only the rows of `chunk` whose index is not in +/// `drop`. Used after BigQuery reports row-level errors. +pub(crate) fn without_rows(chunk: &Chunk, drop: &[usize]) -> Result<Option<Chunk>, ArrowError> { + let keep: Vec<bool> = (0..chunk.batch.num_rows()) + .map(|row| !drop.contains(&row)) + .collect(); + let offsets: Vec<u64> = chunk + .offsets + .iter() + .zip(&keep) + .filter_map(|(offset, keep)| keep.then_some(*offset)) + .collect(); + if offsets.is_empty() { + return Ok(None); + } + let mask = arrow::array::BooleanArray::from(keep); + let batch = arrow::compute::filter_record_batch(&chunk.batch, &mask)?; + let (schema_bytes, batch_bytes) = encode_ipc(&batch)?; + Ok(Some(Chunk { + batch, + offsets, + schema_bytes, + batch_bytes, + })) +} + +// ─── Mapped mode ───────────────────────────────────────────────────────────── + +fn encode_mapped( + ctx: &RunContext<'_>, + messages: Vec<ConsumedMessage>, + rejected: &mut Vec<Rejected>, +) -> Result<Option<(RecordBatch, Vec<u64>)>, ArrowError> { + let layout = ctx.layout; + let mut rows = Vec::with_capacity(messages.len()); + let mut kept = Vec::with_capacity(messages.len()); + for mut message in messages { + let payload = std::mem::replace(&mut message.payload, Payload::Raw(Vec::new())); + match json_object(payload).and_then(|mut row| { + normalize_row(&mut row, layout)?; + Ok(row) + }) { + Ok(row) => { + rows.push(row); + kept.push(message); + } + Err(reason) => rejected.push(Rejected { + offset: message.offset, + reason, + }), + } + } + if rows.is_empty() { + return Ok(None); + } + + let present = present_columns(&rows, layout); + if present.is_empty() && layout.metadata.is_empty() { + rejected.extend(kept.iter().map(|m| Rejected { + offset: m.offset, + reason: "payload has no field matching a table column".to_owned(), + })); + return Ok(None); + } + let payload_fields: Vec<FieldRef> = present.iter().map(|&i| layout.fields[i].clone()).collect(); + + let (decoded, survivors) = decode_rows(&payload_fields, &rows, kept.len(), |row, reason| { + rejected.push(Rejected { + offset: kept[row].offset, + reason, + }); + })?; + if survivors.is_empty() { + return Ok(None); + } + let kept: Vec<&ConsumedMessage> = survivors.iter().map(|&row| &kept[row]).collect(); + let offsets = kept.iter().map(|m| m.offset).collect(); + + let mut fields = payload_fields; + let mut columns: Vec<ArrayRef> = decoded.map(|b| b.columns().to_vec()).unwrap_or_default(); + append_metadata(ctx, &kept, &mut fields, &mut columns)?; + let batch = RecordBatch::try_new(Arc::new(Schema::new(fields)), columns)?; + Ok(Some((batch, offsets))) +} + +/// Take the JSON object out of a payload. Text and proto text must parse as +/// a JSON object. +fn json_object(payload: Payload) -> Result<OwnedValue, String> { + let value = match payload { + Payload::Json(value) => value, + Payload::Text(text) | Payload::Proto(text) => { + let mut bytes = text.into_bytes(); + simd_json::to_owned_value(&mut bytes) + .map_err(|e| format!("payload text is not valid JSON: {e}"))? + } + other => { + return Err(format!( + "{} payload is not supported in mapped mode", + payload_kind(&other) + )); + } + }; + match value { + OwnedValue::Object(_) => Ok(value), + _ => Err("payload is not a JSON object".to_owned()), + } +} + +/// Prepare one row for `arrow-json`: +/// - a REQUIRED column without a default must be present and non-null, +/// - REPEATED columns that are absent or null become `[]`, +/// - JSON column values become their JSON text, +/// - BYTES column values are base64 in JSON and become hex, which is what +/// `arrow-json` decodes into binary. +fn normalize_row(row: &mut OwnedValue, layout: &TableLayout) -> Result<(), String> { + let OwnedValue::Object(object) = row else { + return Err("payload is not a JSON object".to_owned()); + }; + for column in &layout.columns { + let value = object.get_mut(column.name.as_str()); + match (value, column.mode) { + (None | Some(OwnedValue::Static(StaticNode::Null)), Mode::Repeated) => { + object.insert(column.name.clone(), empty_array()); + } + (None | Some(OwnedValue::Static(StaticNode::Null)), Mode::Required) + if !column.has_default => + { + return Err(format!( + "missing value for REQUIRED column '{}'", + column.name + )); + } + (Some(value), _) => normalize_value(value, column)?, + (None, _) => {} + } + } + Ok(()) +} + +fn normalize_value(value: &mut OwnedValue, column: &Column) -> Result<(), String> { + if column.mode == Mode::Repeated { + let OwnedValue::Array(items) = value else { + return Err(format!("column '{}' expects an array", column.name)); + }; + for item in items.iter_mut() { + normalize_scalar(item, column)?; + } + return Ok(()); + } + normalize_scalar(value, column) +} + +fn normalize_scalar(value: &mut OwnedValue, column: &Column) -> Result<(), String> { + if matches!(value, OwnedValue::Static(StaticNode::Null)) { + return Ok(()); + } + match &column.bq_type { + BqType::Json => { + let text = simd_json::to_string(value) + .map_err(|e| format!("column '{}': cannot serialize JSON: {e}", column.name))?; + *value = OwnedValue::String(text); + } + BqType::Bytes => { + let OwnedValue::String(text) = value else { + return Err(format!("column '{}' expects a base64 string", column.name)); + }; + let bytes = BASE64 + .decode(text.as_bytes()) + .map_err(|e| format!("column '{}': invalid base64: {e}", column.name))?; + *value = OwnedValue::String(to_hex(&bytes)); + } + BqType::Record(children) => { + let OwnedValue::Object(object) = value else { + return Ok(()); + }; + for child in children { + match (object.get_mut(child.name.as_str()), child.mode) { + (None | Some(OwnedValue::Static(StaticNode::Null)), Mode::Repeated) => { + object.insert(child.name.clone(), empty_array()); + } + (Some(child_value), _) => normalize_value(child_value, child)?, + (None, _) => {} + } + } + } + _ => {} + } + Ok(()) +} + +/// Indexes of the layout columns that at least one row sets, in table order. +fn present_columns(rows: &[OwnedValue], layout: &TableLayout) -> Vec<usize> { + let mut present = vec![false; layout.columns.len()]; + for row in rows { + if let OwnedValue::Object(object) = row { + for key in object.keys() { + if let Some(&index) = layout.index.get(key.as_str()) { + present[index] = true; + } + } + } + } + present + .iter() + .enumerate() + .filter_map(|(index, set)| set.then_some(index)) + .collect() +} + +/// Decode rows into one batch. Returns the batch (`None` when there are no +/// payload fields) and the indexes of the rows it holds. +fn decode_rows( + fields: &[FieldRef], + rows: &[OwnedValue], + row_count: usize, + mut on_reject: impl FnMut(usize, String), +) -> Result<(Option<RecordBatch>, Vec<usize>), ArrowError> { + if fields.is_empty() { + return Ok((None, (0..row_count).collect())); + } + let target: SchemaRef = Arc::new(Schema::new(fields.to_vec())); + let schema: SchemaRef = Arc::new(Schema::new( + fields + .iter() + .map(|f| { + f.as_ref() + .clone() + .with_data_type(decode_type(f.data_type())) + }) + .collect::<Vec<_>>(), + )); + if let Ok(batch) = decode(&schema, rows) { + return Ok((Some(to_target(batch, &target)?), (0..row_count).collect())); + } + + let mut batches = Vec::with_capacity(rows.len()); + let mut survivors = Vec::with_capacity(rows.len()); + for (index, row) in rows.iter().enumerate() { + match decode(&schema, std::slice::from_ref(row)) { + Ok(batch) => { + batches.push(batch); + survivors.push(index); + } + Err(e) => on_reject(index, format!("row does not match the table schema: {e}")), + } + } + if batches.is_empty() { + return Ok((None, survivors)); + } + let batch = concat_batches(&schema, &batches)?; + Ok((Some(to_target(batch, &target)?), survivors)) +} + +/// The type `arrow-json` decodes into. Without the `chrono-tz` feature it +/// only parses offset timezones, so `"UTC"` is decoded as `"+00:00"` and +/// relabelled afterwards. Both name the same instant, so the relabel is a +/// metadata change. +fn decode_type(data_type: &DataType) -> DataType { + match data_type { + DataType::Timestamp(unit, Some(_)) => DataType::Timestamp(*unit, Some(UTC_OFFSET.into())), + DataType::List(item) => DataType::List(Arc::new( + item.as_ref() + .clone() + .with_data_type(decode_type(item.data_type())), + )), + DataType::Struct(children) => DataType::Struct( + children + .iter() + .map(|f| { + f.as_ref() + .clone() + .with_data_type(decode_type(f.data_type())) + }) + .collect(), + ), + other => other.clone(), + } +} + +fn to_target(batch: RecordBatch, target: &SchemaRef) -> Result<RecordBatch, ArrowError> { + let columns = batch + .columns() + .iter() + .zip(target.fields()) + .map(|(column, field)| { + if column.data_type() == field.data_type() { + Ok(column.clone()) + } else { + cast(column, field.data_type()) + } + }) + .collect::<Result<Vec<_>, _>>()?; + RecordBatch::try_new(target.clone(), columns) +} + +fn decode(schema: &SchemaRef, rows: &[OwnedValue]) -> Result<RecordBatch, ArrowError> { + let mut decoder = ReaderBuilder::new(schema.clone()) + .with_batch_size(rows.len().max(1)) + .build_decoder()?; + decoder.serialize(rows)?; + decoder + .flush()? + .ok_or_else(|| ArrowError::JsonError("decoder produced no rows".to_owned())) +} + +// ─── Raw mode ──────────────────────────────────────────────────────────────── + +fn encode_raw( + ctx: &RunContext<'_>, + kind: RawKind, + payload_field: &FieldRef, + messages: Vec<ConsumedMessage>, + rejected: &mut Vec<Rejected>, +) -> Result<Option<(RecordBatch, Vec<u64>)>, ArrowError> { + let mut kept = Vec::with_capacity(messages.len()); + let payload_column: ArrayRef = match kind { + RawKind::Bytes => { + let mut values = Vec::with_capacity(messages.len()); + for message in messages { + match message.payload.try_to_bytes() { + Ok(bytes) => { + values.push(bytes); + kept.push(message); + } + Err(e) => rejected.push(Rejected { + offset: message.offset, + reason: e.to_string(), + }), + } + } + Arc::new(BinaryArray::from_iter_values(values)) + } + RawKind::Json | RawKind::String => { + let mut values = Vec::with_capacity(messages.len()); + for mut message in messages { + let payload = std::mem::replace(&mut message.payload, Payload::Raw(Vec::new())); + match raw_text(payload, kind) { + Ok(text) => { + values.push(text); + kept.push(message); + } + Err(reason) => rejected.push(Rejected { + offset: message.offset, + reason, + }), + } + } + Arc::new(StringArray::from_iter_values(values)) + } + }; + if kept.is_empty() { + return Ok(None); + } + let kept: Vec<&ConsumedMessage> = kept.iter().collect(); + let offsets = kept.iter().map(|m| m.offset).collect(); + let mut fields = vec![payload_field.clone()]; + let mut columns = vec![payload_column]; + append_metadata(ctx, &kept, &mut fields, &mut columns)?; + let batch = RecordBatch::try_new(Arc::new(Schema::new(fields)), columns)?; + Ok(Some((batch, offsets))) +} + +/// Payload as text for a JSON or STRING column. A JSON column only takes +/// text that parses as JSON. +fn raw_text(payload: Payload, kind: RawKind) -> Result<String, String> { + let text = match payload { + Payload::Json(value) => { + return simd_json::to_string(&value).map_err(|e| format!("cannot serialize JSON: {e}")); + } + Payload::Text(text) | Payload::Proto(text) => text, + Payload::Raw(bytes) => { + String::from_utf8(bytes).map_err(|_| "raw payload is not valid UTF-8".to_owned())? + } + other => { + return Err(format!( + "{} payload cannot be written to a text column", + payload_kind(&other) + )); + } + }; + if kind == RawKind::Json { + let mut bytes = text.as_bytes().to_vec(); + simd_json::to_owned_value(&mut bytes) + .map_err(|e| format!("payload is not valid JSON: {e}"))?; + } + Ok(text) +} + +// ─── Metadata ──────────────────────────────────────────────────────────────── + +fn append_metadata( + ctx: &RunContext<'_>, + messages: &[&ConsumedMessage], + fields: &mut Vec<FieldRef>, + columns: &mut Vec<ArrayRef>, +) -> Result<(), ArrowError> { + let rows = messages.len(); + for (meta, field) in &ctx.layout.metadata { + let column: ArrayRef = match meta { + MetaColumn::Stream => Arc::new(StringArray::from_iter_values(std::iter::repeat_n( + ctx.topic.stream.as_str(), + rows, + ))), + MetaColumn::Topic => Arc::new(StringArray::from_iter_values(std::iter::repeat_n( + ctx.topic.topic.as_str(), + rows, + ))), + MetaColumn::PartitionId => Arc::new(Int64Array::from_iter_values(std::iter::repeat_n( + i64::from(ctx.messages.partition_id), + rows, + ))), + MetaColumn::Offset => Arc::new(Int64Array::from_iter_values( + messages.iter().map(|m| saturating_i64(m.offset)), + )), + MetaColumn::Timestamp => Arc::new( + TimestampMicrosecondArray::from_iter_values( + messages.iter().map(|m| saturating_i64(m.timestamp)), + ) + .with_timezone(UTC), + ), + MetaColumn::Id => Arc::new(StringArray::from_iter_values( + messages.iter().map(|m| m.id.to_string()), + )), + MetaColumn::Headers => { + let required = !field.is_nullable(); + Arc::new(StringArray::from_iter( + messages.iter().map(|m| headers_json(m, required)), + )) + } + }; + fields.push(field.clone()); + columns.push(column); + } + Ok(()) +} + +/// Headers as a JSON object. Text values are written as strings, binary +/// values as `{"data": <base64>, "iggy_header_encoding": "base64"}`, the same +/// shape `http_sink` uses. +fn headers_json(message: &ConsumedMessage, required: bool) -> Option<String> { + let headers = message.headers.as_ref().filter(|h| !h.is_empty()); + let Some(headers) = headers else { + return required.then(|| "{}".to_owned()); + }; + let mut object = serde_json::Map::with_capacity(headers.len()); + for (key, value) in headers { + let encoded = match value.as_raw() { + Ok(raw) => serde_json::json!({ + "data": BASE64.encode(raw), + "iggy_header_encoding": HEADER_ENCODING_BASE64, + }), + Err(_) => serde_json::Value::String(value.to_string_value()), + }; + object.insert(key.to_string_value(), encoded); + } + Some(serde_json::Value::Object(object).to_string()) +} + +// ─── Request sizing ────────────────────────────────────────────────────────── + +/// Encode `batch` and halve it until every piece fits `max_bytes`. A single +/// row that does not fit is rejected. +fn split_into_chunks( + batch: RecordBatch, + offsets: Vec<u64>, + max_bytes: usize, + chunks: &mut Vec<Chunk>, + rejected: &mut Vec<Rejected>, +) -> Result<(), ArrowError> { + let (schema_bytes, batch_bytes) = encode_ipc(&batch)?; + let size = schema_bytes.len() + batch_bytes.len(); + if size <= max_bytes { + chunks.push(Chunk { + batch, + offsets, + schema_bytes, + batch_bytes, + }); + return Ok(()); + } + if batch.num_rows() == 1 { + rejected.push(Rejected { + offset: offsets[0], + reason: format!("row is {size} bytes encoded, above max_request_bytes ({max_bytes})"), + }); + return Ok(()); + } + let half = batch.num_rows() / 2; + let (left_offsets, right_offsets) = offsets.split_at(half); + split_into_chunks( + batch.slice(0, half), + left_offsets.to_vec(), + max_bytes, + chunks, + rejected, + )?; + split_into_chunks( + batch.slice(half, batch.num_rows() - half), + right_offsets.to_vec(), + max_bytes, + chunks, + rejected, + ) +} + +/// Serialize the schema and the batch as Arrow IPC stream messages, the +/// format `AppendRowsRequest.arrow_rows` carries. +fn encode_ipc(batch: &RecordBatch) -> Result<(Vec<u8>, Vec<u8>), ArrowError> { + let options = IpcWriteOptions::default(); + let generator = IpcDataGenerator::default(); + let mut tracker = DictionaryTracker::new(true); + let mut compression = CompressionContext::default(); + + let schema = generator.schema_to_bytes_with_dictionary_tracker( + batch.schema_ref(), + &mut tracker, + &options, + ); + let mut schema_bytes = Vec::new(); + write_message(&mut schema_bytes, schema, &options)?; + + let (dictionaries, encoded) = + generator.encode(batch, &mut tracker, &options, &mut compression)?; + let mut batch_bytes = Vec::new(); + for dictionary in dictionaries { + write_message(&mut batch_bytes, dictionary, &options)?; + } + write_message(&mut batch_bytes, encoded, &options)?; + Ok((schema_bytes, batch_bytes)) +} + +// ─── Helpers ───────────────────────────────────────────────────────────────── + +fn payload_kind(payload: &Payload) -> &'static str { + match payload { + Payload::Json(_) => "JSON", + Payload::Raw(_) => "raw", + Payload::Text(_) => "text", + Payload::Proto(_) => "proto", + Payload::FlatBuffer(_) => "FlatBuffer", + Payload::Avro(_) => "Avro", + } +} + +fn empty_array() -> OwnedValue { + OwnedValue::Array(Box::default()) +} + +fn to_hex(bytes: &[u8]) -> String { Review Comment: simplification: `to_hex` hand-rolls what `hex::encode` does, and the workspace already depends on `hex`. Add `hex = { workspace = true }` to the crate, then delete `to_hex` and its `std::fmt::Write` import. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
