kparisa commented on code in PR #4394: URL: https://github.com/apache/iggy/pull/4394#discussion_r4221513462
########## core/connectors/sinks/bigquery_sink/src/encode.rs: ########## @@ -0,0 +1,1559 @@ +// 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, bisection isolates the offending +//! rows before the survivors are decoded together. +//! +//! 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 std::fmt::Write as _; +use std::io; +use std::sync::Arc; + +use arrow::array::{ + Array, ArrayRef, BinaryBuilder, Int64Array, ListArray, RecordBatch, StringArray, StringBuilder, + StructArray, TimestampMicrosecondArray, +}; +use arrow::compute::cast; +use arrow::datatypes::{DataType, FieldRef, Fields, 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 serde::ser::{Serialize, SerializeMap, SerializeSeq, Serializer}; +use simd_json::{OwnedValue, StaticNode}; + +use crate::schema::{BqType, Column, MetaColumn, Mode, RawKind, TableLayout, UTC}; + +const HEADER_ENCODING_BASE64: &str = "base64"; +const UTC_OFFSET: &str = "+00:00"; +const IPC_BATCH_OVERHEAD: usize = 1024; +const IPC_COLUMN_OVERHEAD: usize = 128; + +/// 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, |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 remain base64 until they are decoded directly into +/// Arrow binary buffers. +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) Review Comment: intending to drop such.. 2026-10-08T13:54:06.544169Z WARN connector: connector_target="iggy_connector_bigquery_sink::sink" BigQuery sink ID: 1 dropped message stream=demo_stream topic=demo_topic partition=0 offset=14: row does not match the table schema: Json error: whilst decoding field 'user_id': failed to parse "bad" as Int64 -- 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]
