From 5540889f40972f0c5e43d9b2fb4a1d049208a9c4 Mon Sep 17 00:00:00 2001 From: Amanda Li Date: Thu, 30 Jul 2026 01:39:09 +0000 Subject: [PATCH] feat(codecs): add wire_to_arrow batch encoder Adds a `wire_to_arrow` batch encoding codec that decodes protobuf wire bytes from each event's `message` field directly into an Apache Arrow `RecordBatch`, pairing proto fields to Arrow columns by name. This bypasses the generic ProtobufDeserializer -> Event -> ArrowStreamSerializer chain for sinks that already carry raw proto bytes. Configured with a proto descriptor (`desc_file` + `message_type`) for the incoming bytes; the sink injects the output Arrow schema. Supports scalars, nested messages, repeated fields (packed and unpacked), maps, oneofs, and int64 -> Timestamp(Microsecond) coercion. Malformed rows are dropped and counted via the `wire_to_arrow_rows_dropped` metric rather than failing the whole batch. Co-authored-by: Isaac --- .../wire_to_arrow_batch_codec.feature.md | 10 + lib/codecs/Cargo.toml | 1 + lib/codecs/src/encoding/encoder.rs | 18 +- lib/codecs/src/encoding/format/mod.rs | 6 + .../encoding/format/wire_to_arrow/append.rs | 658 +++++++ .../encoding/format/wire_to_arrow/builders.rs | 703 +++++++ .../encoding/format/wire_to_arrow/encoder.rs | 98 + .../encoding/format/wire_to_arrow/errors.rs | 321 +++ .../src/encoding/format/wire_to_arrow/mod.rs | 58 + .../src/encoding/format/wire_to_arrow/plan.rs | 674 +++++++ .../src/encoding/format/wire_to_arrow/scan.rs | 260 +++ .../format/wire_to_arrow/serializer.rs | 134 ++ .../encoding/format/wire_to_arrow/tests.rs | 1738 +++++++++++++++++ .../src/encoding/format/wire_to_arrow/wire.rs | 544 ++++++ lib/codecs/src/encoding/mod.rs | 1 + lib/codecs/src/encoding/serializer.rs | 22 +- 16 files changed, 5243 insertions(+), 3 deletions(-) create mode 100644 changelog.d/wire_to_arrow_batch_codec.feature.md create mode 100644 lib/codecs/src/encoding/format/wire_to_arrow/append.rs create mode 100644 lib/codecs/src/encoding/format/wire_to_arrow/builders.rs create mode 100644 lib/codecs/src/encoding/format/wire_to_arrow/encoder.rs create mode 100644 lib/codecs/src/encoding/format/wire_to_arrow/errors.rs create mode 100644 lib/codecs/src/encoding/format/wire_to_arrow/mod.rs create mode 100644 lib/codecs/src/encoding/format/wire_to_arrow/plan.rs create mode 100644 lib/codecs/src/encoding/format/wire_to_arrow/scan.rs create mode 100644 lib/codecs/src/encoding/format/wire_to_arrow/serializer.rs create mode 100644 lib/codecs/src/encoding/format/wire_to_arrow/tests.rs create mode 100644 lib/codecs/src/encoding/format/wire_to_arrow/wire.rs diff --git a/changelog.d/wire_to_arrow_batch_codec.feature.md b/changelog.d/wire_to_arrow_batch_codec.feature.md new file mode 100644 index 0000000000000..414c2aed74b2c --- /dev/null +++ b/changelog.d/wire_to_arrow_batch_codec.feature.md @@ -0,0 +1,10 @@ +Added a `wire_to_arrow` batch encoding codec that decodes protobuf wire bytes from each +event's `message` field directly into an Apache Arrow `RecordBatch`, pairing proto fields to +Arrow columns by name. This bypasses the generic `ProtobufDeserializer -> Event -> +ArrowStreamSerializer` chain for sinks that already carry raw proto bytes, avoiding the +intermediate `DynamicMessage` / `LogEvent` representations. The codec is configured with a +proto descriptor (`desc_file` + `message_type`) for the incoming bytes; the sink injects the +output Arrow schema. Malformed rows are isolated (dropped and counted via the +`wire_to_arrow_rows_dropped` metric) rather than failing the whole batch. + +authors: amandaLi7 diff --git a/lib/codecs/Cargo.toml b/lib/codecs/Cargo.toml index 0c24d81913846..5f9e6b057b933 100644 --- a/lib/codecs/Cargo.toml +++ b/lib/codecs/Cargo.toml @@ -68,6 +68,7 @@ toml = { version = "0.9.8", optional = true } criterion.workspace = true futures.workspace = true indoc.workspace = true +proptest.workspace = true tokio = { workspace = true, features = ["test-util"] } toml.workspace = true similar-asserts = "1.7.0" diff --git a/lib/codecs/src/encoding/encoder.rs b/lib/codecs/src/encoding/encoder.rs index 2bbfba6cdbb55..8a69ac96e26c6 100644 --- a/lib/codecs/src/encoding/encoder.rs +++ b/lib/codecs/src/encoding/encoder.rs @@ -4,7 +4,7 @@ use vector_common::internal_event::emit; use vector_core::event::Event; #[cfg(feature = "arrow")] -use crate::encoding::ArrowStreamSerializer; +use crate::encoding::{ArrowStreamSerializer, WireToArrowSerializer}; #[cfg(feature = "parquet")] use crate::encoding::ParquetSerializer; use crate::{ @@ -31,6 +31,9 @@ pub enum BatchOutput { pub enum BatchSerializer { /// Arrow IPC stream format serializer. Arrow(ArrowStreamSerializer), + /// Wire-to-Arrow serializer: decodes proto wire bytes straight into an + /// Arrow `RecordBatch`. + WireToArrow(WireToArrowSerializer), /// Parquet format serializer. #[cfg(feature = "parquet")] Parquet(Box), @@ -58,7 +61,9 @@ impl BatchEncoder { /// Get the HTTP content type. pub const fn content_type(&self) -> Option<&'static str> { match &self.serializer { - BatchSerializer::Arrow(_) => Some("application/vnd.apache.arrow.stream"), + BatchSerializer::Arrow(_) | BatchSerializer::WireToArrow(_) => { + Some("application/vnd.apache.arrow.stream") + } #[cfg(feature = "parquet")] BatchSerializer::Parquet(_) => Some("application/vnd.apache.parquet"), } @@ -79,6 +84,12 @@ impl BatchEncoder { })?; Ok(BatchOutput::Arrow(record_batch)) } + BatchSerializer::WireToArrow(serializer) => { + let record_batch = serializer + .encode_to_record_batch(events) + .map_err(|err| Error::SerializingError(Box::new(err)))?; + Ok(BatchOutput::Arrow(record_batch)) + } #[cfg(feature = "parquet")] BatchSerializer::Parquet(_) => Err(Error::SerializingError(Box::from( "Parquet serializer does not support encode_batch; use the tokio Encoder interface instead", @@ -104,6 +115,9 @@ impl tokio_util::codec::Encoder> for BatchEncoder { } }) } + BatchSerializer::WireToArrow(_) => Err(Error::SerializingError(Box::from( + "WireToArrow serializer does not support the streaming Encoder interface; use encode_batch() instead", + ))), #[cfg(feature = "parquet")] BatchSerializer::Parquet(serializer) => serializer .encode(events, buffer) diff --git a/lib/codecs/src/encoding/format/mod.rs b/lib/codecs/src/encoding/format/mod.rs index f760ea5c8dd83..91d7b4979092c 100644 --- a/lib/codecs/src/encoding/format/mod.rs +++ b/lib/codecs/src/encoding/format/mod.rs @@ -23,6 +23,8 @@ mod raw_message; #[cfg(feature = "syslog")] mod syslog; mod text; +#[cfg(feature = "arrow")] +mod wire_to_arrow; use std::fmt::Debug; @@ -35,6 +37,10 @@ pub use arrow::{ ArrowEncodingError, ArrowStreamSerializer, ArrowStreamSerializerConfig, SchemaProvider, find_null_non_nullable_fields, }; +#[cfg(feature = "arrow")] +pub use wire_to_arrow::{ + WireToArrowEncoder, WireToArrowError, WireToArrowSerializer, WireToArrowSerializerConfig, +}; pub use avro::{AvroSerializer, AvroSerializerConfig, AvroSerializerOptions}; pub use cef::{CefSerializer, CefSerializerConfig}; use dyn_clone::DynClone; diff --git a/lib/codecs/src/encoding/format/wire_to_arrow/append.rs b/lib/codecs/src/encoding/format/wire_to_arrow/append.rs new file mode 100644 index 0000000000000..d8e675749b0da --- /dev/null +++ b/lib/codecs/src/encoding/format/wire_to_arrow/append.rs @@ -0,0 +1,658 @@ +//! Wire-value → Arrow builder append dispatch. +//! +//! The outer scan in `mod.rs` walks proto bytes tag-by-tag and hands each +//! decoded `WireValue` to one of the `append_*` functions here. For packed +//! repeated scalars the inner blob is walked tagless: `try_parse_field` +//! expects tag-prefixed fields and can't traverse it, so varints are read +//! via [`super::wire::try_read_varint`] and the two fixed-width wire +//! types are read inline (3-line LE chunk reads — not worth a helper). + +use super::wire::{WireValue, decode_zigzag32, decode_zigzag64, try_read_varint}; + +use super::builders::TypedBuilder; +use super::errors::{Result, WireToArrowError}; +use super::plan::{ScalarKind, WT_I32, WT_I64, WT_LEN, WT_VARINT}; + +/// Extract the inner bytes from a length-delimited `WireValue`, or error. +#[inline(always)] +pub(super) fn expect_len<'a>(wv: &'a WireValue<'a>) -> Result<&'a [u8]> { + match wv { + WireValue::Len(b) => Ok(b), + other => Err(WireToArrowError::WireTypeMismatch { + expected: WT_LEN, + actual: wire_type_byte(other), + }), + } +} + +/// Proto wire type numeric code for a `WireValue`. Used for error reporting. +#[inline] +pub(super) fn wire_type_byte(wv: &WireValue) -> u8 { + match wv { + WireValue::Varint(_) => WT_VARINT, + WireValue::I64(_) => WT_I64, + WireValue::Len(_) => WT_LEN, + WireValue::I32(_) => WT_I32, + } +} + +/// Append one scalar `WireValue` into the matching typed Arrow builder. +/// The only runtime error is a wire-type disagreement; (kind, builder) +/// pairing is enforced when the plan is built. +#[inline(always)] +pub(super) fn append_scalar_from_wire( + kind: ScalarKind, + wv: &WireValue, + tb: &mut TypedBuilder, +) -> Result<()> { + match kind { + ScalarKind::Int32 => { + let v = expect_varint(kind, wv)?; + if let TypedBuilder::Int32(b) = tb { + b.append_value(v as i32); + return Ok(()); + } + } + ScalarKind::Int64 => { + let v = expect_varint(kind, wv)?; + match tb { + TypedBuilder::Int64(b) => { + b.append_value(v as i64); + return Ok(()); + } + TypedBuilder::TimestampMicros(b) => { + b.append_value(v as i64); + return Ok(()); + } + _ => {} + } + } + ScalarKind::UInt32 => { + let v = expect_varint(kind, wv)?; + if let TypedBuilder::UInt32(b) = tb { + b.append_value(v as u32); + return Ok(()); + } + } + ScalarKind::UInt64 => { + let v = expect_varint(kind, wv)?; + if let TypedBuilder::UInt64(b) = tb { + b.append_value(v); + return Ok(()); + } + } + ScalarKind::SInt32 => { + let v = expect_varint(kind, wv)?; + if let TypedBuilder::Int32(b) = tb { + b.append_value(decode_zigzag32(v as u32)); + return Ok(()); + } + } + ScalarKind::SInt64 => { + let v = expect_varint(kind, wv)?; + match tb { + TypedBuilder::Int64(b) => { + b.append_value(decode_zigzag64(v)); + return Ok(()); + } + TypedBuilder::TimestampMicros(b) => { + b.append_value(decode_zigzag64(v)); + return Ok(()); + } + _ => {} + } + } + ScalarKind::Fixed32 => { + let v = expect_i32(kind, wv)?; + if let TypedBuilder::UInt32(b) = tb { + b.append_value(v); + return Ok(()); + } + } + ScalarKind::SFixed32 => { + let v = expect_i32(kind, wv)?; + if let TypedBuilder::Int32(b) = tb { + b.append_value(v as i32); + return Ok(()); + } + } + ScalarKind::Float => { + let v = expect_i32(kind, wv)?; + if let TypedBuilder::Float32(b) = tb { + b.append_value(f32::from_bits(v)); + return Ok(()); + } + } + ScalarKind::Fixed64 => { + let v = expect_i64(kind, wv)?; + if let TypedBuilder::UInt64(b) = tb { + b.append_value(v); + return Ok(()); + } + } + ScalarKind::SFixed64 => { + let v = expect_i64(kind, wv)?; + match tb { + TypedBuilder::Int64(b) => { + b.append_value(v as i64); + return Ok(()); + } + TypedBuilder::TimestampMicros(b) => { + b.append_value(v as i64); + return Ok(()); + } + _ => {} + } + } + ScalarKind::Double => { + let v = expect_i64(kind, wv)?; + if let TypedBuilder::Float64(b) = tb { + b.append_value(f64::from_bits(v)); + return Ok(()); + } + } + ScalarKind::Bool => { + let v = expect_varint(kind, wv)?; + if let TypedBuilder::Boolean(b) = tb { + b.append_value(v != 0); + return Ok(()); + } + } + ScalarKind::String => { + let bytes = expect_len(wv)?; + if let TypedBuilder::LargeUtf8(b) = tb { + let s = std::str::from_utf8(bytes).map_err(|_| WireToArrowError::InvalidUtf8)?; + b.append_value(s); + return Ok(()); + } + } + ScalarKind::Bytes => { + let bytes = expect_len(wv)?; + if let TypedBuilder::LargeBinary(b) = tb { + b.append_value(bytes); + return Ok(()); + } + } + } + // Wire bytes don't match the schema, or (defensively) the plan builder + // paired a scalar kind with a builder it can't write to. + Err(WireToArrowError::WireTypeMismatch { + expected: kind.wire_type(), + actual: wire_type_byte(wv), + }) +} + +/// Render a proto enum varint into a `LargeUtf8` builder as its enum-value +/// *name* (e.g. `1` -> `"SUCCESS"`), for schemas that model an enum column as +/// a string rather than an `Int32`. Rules: +/// - out-of-range value -> `UNKNOWN_ENUM_VALUE__` placeholder, never +/// a row drop (so an enum value added after this binary was built is still +/// representable); +/// - a wrong wire type (enum must be a varint) -> `WireTypeMismatch`. +/// The absent/proto3-default case (elided zero value -> null) is handled by the +/// builder-tree's `finalize_row`, not here. +#[inline] +pub(super) fn append_enum_string_from_wire( + desc: &prost_reflect::EnumDescriptor, + wv: &WireValue, + tb: &mut TypedBuilder, +) -> Result<()> { + let raw = match wv { + WireValue::Varint(v) => *v, + other => { + return Err(WireToArrowError::WireTypeMismatch { + expected: WT_VARINT, + actual: wire_type_byte(other), + }); + } + }; + // Enum values are int32 on the wire; proto sign-extends negatives to a + // 64-bit varint, so `append_enum_name` narrows via `as i32` before the + // descriptor lookup. + append_enum_name(desc, raw, tb) +} + +/// Append the proto3 default for `kind` to `tb`. Used at `finalize_row` time +/// for absent scalar slots inside Map entry sub-plans, where the Arrow Map +/// type declares the key non-nullable but proto3 elides the key tag whenever +/// it carries its default value (`""`, `0`, `false`, `b""`). Writing a null +/// here would fail `StructArray::try_new` at finish; writing the proto3 +/// default matches the semantics every standards-compliant proto consumer +/// applies. +/// +/// Infallible: [`MessagePlan::build_at_depth`] rejects any (`ScalarKind`, +/// Arrow leaf) pairing this helper can't handle via +/// [`ScalarKind::matches_arrow_type`], so by the time a builder tree exists +/// every `BuilderNode::Scalar` paired with `inside_map_entry = true` is +/// guaranteed to land in one of the matched arms below. Keep the pairings +/// here in sync with that check and with [`append_scalar_from_wire`]. +/// +/// [`MessagePlan::build_at_depth`]: super::plan::MessagePlan +/// [`ScalarKind::matches_arrow_type`]: super::plan::ScalarKind::matches_arrow_type +#[inline] +pub(super) fn append_proto3_default(kind: ScalarKind, tb: &mut TypedBuilder) { + match kind { + ScalarKind::Int32 | ScalarKind::SInt32 | ScalarKind::SFixed32 => { + if let TypedBuilder::Int32(b) = tb { + b.append_value(0); + return; + } + } + ScalarKind::Int64 | ScalarKind::SInt64 | ScalarKind::SFixed64 => match tb { + TypedBuilder::Int64(b) => { + b.append_value(0); + return; + } + TypedBuilder::TimestampMicros(b) => { + b.append_value(0); + return; + } + _ => {} + }, + ScalarKind::UInt32 | ScalarKind::Fixed32 => { + if let TypedBuilder::UInt32(b) = tb { + b.append_value(0); + return; + } + } + ScalarKind::UInt64 | ScalarKind::Fixed64 => { + if let TypedBuilder::UInt64(b) = tb { + b.append_value(0); + return; + } + } + ScalarKind::Float => { + if let TypedBuilder::Float32(b) = tb { + b.append_value(0.0); + return; + } + } + ScalarKind::Double => { + if let TypedBuilder::Float64(b) = tb { + b.append_value(0.0); + return; + } + } + ScalarKind::Bool => { + if let TypedBuilder::Boolean(b) = tb { + b.append_value(false); + return; + } + } + ScalarKind::String => { + if let TypedBuilder::LargeUtf8(b) = tb { + b.append_value(""); + return; + } + } + ScalarKind::Bytes => { + if let TypedBuilder::LargeBinary(b) = tb { + b.append_value(b"" as &[u8]); + return; + } + } + } + unreachable!( + "plan-build invariant: (ScalarKind, TypedBuilder) pairing is enforced \ + by ScalarKind::matches_arrow_type at MessagePlan::build_at_depth; \ + reaching this arm means the plan and builder tree diverged" + ); +} + +#[inline(always)] +fn expect_varint(kind: ScalarKind, wv: &WireValue) -> Result { + if let WireValue::Varint(v) = wv { + Ok(*v) + } else { + Err(WireToArrowError::WireTypeMismatch { + expected: kind.wire_type(), + actual: wire_type_byte(wv), + }) + } +} + +#[inline(always)] +fn expect_i32(kind: ScalarKind, wv: &WireValue) -> Result { + if let WireValue::I32(v) = wv { + Ok(*v) + } else { + Err(WireToArrowError::WireTypeMismatch { + expected: kind.wire_type(), + actual: wire_type_byte(wv), + }) + } +} + +#[inline(always)] +fn expect_i64(kind: ScalarKind, wv: &WireValue) -> Result { + if let WireValue::I64(v) = wv { + Ok(*v) + } else { + Err(WireToArrowError::WireTypeMismatch { + expected: kind.wire_type(), + actual: wire_type_byte(wv), + }) + } +} + +/// Append a repeated-scalar occurrence (either a single unpacked value or a +/// full packed blob) into `values`. +/// +/// `current_offset` is the running cumulative element count for the parent +/// Arrow `List`: bumped by 1 per element appended here (1 for +/// unpacked, N for a packed blob). The owning `BuilderNode::RepeatedScalar` +/// later pushes it onto its `offsets` buffer at row finalization, which is +/// how list lengths are recorded in Arrow's offsets-buffer layout. +pub(super) fn append_repeated_scalar( + kind: ScalarKind, + wv: &WireValue, + values: &mut TypedBuilder, + current_offset: &mut i32, +) -> Result<()> { + // Unpacked form: the `WireValue` variant matches the scalar's native + // wire type. Single append, regardless of scalar kind. + if wire_type_byte(wv) == kind.wire_type() { + append_scalar_from_wire(kind, wv, values)?; + // Cumulative across the batch — see scan.rs for the rationale. + *current_offset = + current_offset + .checked_add(1) + .ok_or(WireToArrowError::OffsetOverflow { + site: "append_repeated_scalar:unpacked", + })?; + return Ok(()); + } + + // Packed form: a `Len` blob holding a run of raw scalar values. Only + // valid when the scalar's native wire type is 0/1/5 (packable). + let WireValue::Len(inner) = wv else { + return Err(WireToArrowError::WireTypeMismatch { + expected: kind.wire_type(), + actual: wire_type_byte(wv), + }); + }; + // Proto spec forbids packed encoding for length-delimited scalars + // (string/bytes) — there's no length-prefix per element inside a packed + // blob, so a `Len`-typed `string`/`bytes` must arrive as one unpacked + // occurrence per value. Reject the combo here. + if kind.wire_type() == WT_LEN { + return Err(WireToArrowError::WireTypeMismatch { + expected: kind.wire_type(), + actual: wire_type_byte(wv), + }); + } + let mut remaining: &[u8] = inner; + while !remaining.is_empty() { + let decoded; + (decoded, remaining) = read_packed_element(kind, remaining)?; + append_scalar_from_wire(kind, &decoded, values)?; + *current_offset = + current_offset + .checked_add(1) + .ok_or(WireToArrowError::OffsetOverflow { + site: "append_repeated_scalar:packed", + })?; + } + Ok(()) +} + +/// Mirror of [`append_scalar_from_wire`] that runs the wire-type / UTF-8 +/// checks without touching a builder. Used by the encoder's pre-validate +/// pass so a malformed row can be detected and dropped before any column +/// builder is mutated (Arrow's `*Builder` types expose no rollback API, so +/// rejecting the row up front is how we keep per-row isolation). +/// +/// Must stay in lock-step with [`append_scalar_from_wire`]: every (kind, wv) +/// combination that succeeds here must also succeed there, and vice versa. +pub(super) fn validate_scalar_from_wire(kind: ScalarKind, wv: &WireValue) -> Result<()> { + match (kind, wv) { + (ScalarKind::Int32, WireValue::Varint(_)) + | (ScalarKind::Int64, WireValue::Varint(_)) + | (ScalarKind::UInt32, WireValue::Varint(_)) + | (ScalarKind::UInt64, WireValue::Varint(_)) + | (ScalarKind::SInt32, WireValue::Varint(_)) + | (ScalarKind::SInt64, WireValue::Varint(_)) + | (ScalarKind::Bool, WireValue::Varint(_)) + | (ScalarKind::Fixed32, WireValue::I32(_)) + | (ScalarKind::SFixed32, WireValue::I32(_)) + | (ScalarKind::Float, WireValue::I32(_)) + | (ScalarKind::Fixed64, WireValue::I64(_)) + | (ScalarKind::SFixed64, WireValue::I64(_)) + | (ScalarKind::Double, WireValue::I64(_)) + | (ScalarKind::Bytes, WireValue::Len(_)) => Ok(()), + (ScalarKind::String, WireValue::Len(bytes)) => std::str::from_utf8(bytes) + .map(|_| ()) + .map_err(|_| WireToArrowError::InvalidUtf8), + (_, wv) => Err(WireToArrowError::WireTypeMismatch { + expected: kind.wire_type(), + actual: wire_type_byte(wv), + }), + } +} + +/// Mirror of [`append_enum_string_from_wire`] that runs the wire-type check +/// without touching a builder. Used by the pre-validate pass. Must stay in +/// lock-step with [`append_enum_string_from_wire`]: every (desc, wv) +/// combination that succeeds here must also succeed there. An unrecognized +/// enum value is not a failure — it renders a placeholder on append. +pub(super) fn validate_enum_string_from_wire( + desc: &prost_reflect::EnumDescriptor, + wv: &WireValue, +) -> Result<()> { + let raw = match wv { + WireValue::Varint(v) => *v, + other => { + return Err(WireToArrowError::WireTypeMismatch { + expected: WT_VARINT, + actual: wire_type_byte(other), + }); + } + }; + validate_enum_number(desc, raw) +} + +/// Mirror of [`append_repeated_scalar`] that walks the value (or packed +/// blob) without appending. Used by the pre-validate pass — the packed-blob +/// inner loop in [`append_repeated_scalar`] is the one site in the encoder +/// where a partial append is possible (an EOF on element N leaves N-1 +/// values already in the builder), so dropping the row up front here is how +/// we keep per-row isolation for repeated scalars. +/// +/// Also enforces a per-row guard against a single packed blob that would +/// push the Arrow list's running offset past `i32::MAX`. The scan-time +/// `checked_add` catches batch-cumulative overflow as a clean error, but +/// surfacing the single-row case here keeps it inside per-row isolation — +/// the offending row drops, the rest of the batch survives. The cumulative +/// across-rows case is the irreducible remainder. +pub(super) fn validate_repeated_scalar(kind: ScalarKind, wv: &WireValue) -> Result<()> { + if wire_type_byte(wv) == kind.wire_type() { + return validate_scalar_from_wire(kind, wv); + } + let WireValue::Len(inner) = wv else { + return Err(WireToArrowError::WireTypeMismatch { + expected: kind.wire_type(), + actual: wire_type_byte(wv), + }); + }; + if kind.wire_type() == WT_LEN { + return Err(WireToArrowError::WireTypeMismatch { + expected: kind.wire_type(), + actual: wire_type_byte(wv), + }); + } + let mut remaining: &[u8] = inner; + // u64 so we can compare against i32::MAX without overflowing the counter + // itself; a packed blob is bounded by the wire-bytes length, which fits. + let mut count: u64 = 0; + while !remaining.is_empty() { + // Packed scalars are always varint / fixed32 / fixed64; the decoded + // `WireValue` is always shape-compatible with `kind`, so no further + // per-element validation is needed. + (_, remaining) = read_packed_element(kind, remaining)?; + count += 1; + if count > i32::MAX as u64 { + return Err(WireToArrowError::OffsetOverflow { + site: "validate_repeated_scalar:packed_row_exceeds_i32", + }); + } + } + Ok(()) +} + +/// Repeated analogue of [`append_enum_string_from_wire`]: render each enum +/// element's value *name* into a `LargeUtf8` list-values builder. Enums are +/// varints, so this accepts both wire forms — a single unpacked varint or a +/// packed blob of varints — mirroring [`append_repeated_scalar`]. `current_offset` +/// is bumped once per element for the parent `List` offsets buffer. +pub(super) fn append_repeated_enum_string( + desc: &prost_reflect::EnumDescriptor, + wv: &WireValue, + values: &mut TypedBuilder, + current_offset: &mut i32, +) -> Result<()> { + // Unpacked form: a single varint occurrence. + if let WireValue::Varint(raw) = wv { + append_enum_name(desc, *raw, values)?; + *current_offset = + current_offset + .checked_add(1) + .ok_or(WireToArrowError::OffsetOverflow { + site: "append_repeated_enum_string:unpacked", + })?; + return Ok(()); + } + + // Packed form: a `Len` blob holding a run of raw varints. + let WireValue::Len(inner) = wv else { + return Err(WireToArrowError::WireTypeMismatch { + expected: WT_VARINT, + actual: wire_type_byte(wv), + }); + }; + let mut remaining: &[u8] = inner; + while !remaining.is_empty() { + let (v, rest) = try_read_varint(remaining)?; + remaining = rest; + append_enum_name(desc, v, values)?; + *current_offset = + current_offset + .checked_add(1) + .ok_or(WireToArrowError::OffsetOverflow { + site: "append_repeated_enum_string:packed", + })?; + } + Ok(()) +} + +/// Look up `raw`'s enum-value name and append it to a `LargeUtf8` builder. +/// A value with no matching descriptor entry renders a synthetic placeholder +/// (see [`unknown_enum_placeholder`]), so a proto enum value added after this +/// binary was built lands as a string instead of dropping the row. +/// Shared by the singular and repeated enum-string paths. +#[inline] +fn append_enum_name( + desc: &prost_reflect::EnumDescriptor, + raw: u64, + tb: &mut TypedBuilder, +) -> Result<()> { + let value = raw as i32; + let TypedBuilder::LargeUtf8(b) = tb else { + return Err(WireToArrowError::PlanBuilderMismatch { + site: "append_enum_name", + }); + }; + match desc.get_value(value) { + Some(v) => b.append_value(v.name()), + None => b.append_value(unknown_enum_placeholder(desc, value)), + } + Ok(()) +} + +/// The placeholder string produced for an unrecognized enum value: +/// `UNKNOWN_ENUM_VALUE__`, where `` is the enum's simple +/// (unqualified) name from `desc.name()`. +#[inline] +fn unknown_enum_placeholder(desc: &prost_reflect::EnumDescriptor, value: i32) -> String { + format!("UNKNOWN_ENUM_VALUE_{}_{}", desc.name(), value) +} + +/// Mirror of [`append_repeated_enum_string`] that walks the value (or packed +/// blob) without appending — used by the pre-validate pass. Runs the same +/// wire-type check (and the packed-blob element-count guard) so a malformed +/// row is dropped before any builder is mutated. An unrecognized enum value is +/// not a failure — it renders a placeholder on append. +pub(super) fn validate_repeated_enum_string( + desc: &prost_reflect::EnumDescriptor, + wv: &WireValue, +) -> Result<()> { + match wv { + WireValue::Varint(raw) => validate_enum_number(desc, *raw), + WireValue::Len(inner) => { + let mut remaining: &[u8] = inner; + let mut count: u64 = 0; + while !remaining.is_empty() { + let (v, rest) = try_read_varint(remaining)?; + remaining = rest; + validate_enum_number(desc, v)?; + count += 1; + if count > i32::MAX as u64 { + return Err(WireToArrowError::OffsetOverflow { + site: "validate_repeated_enum_string:packed_row_exceeds_i32", + }); + } + } + Ok(()) + } + other => Err(WireToArrowError::WireTypeMismatch { + expected: WT_VARINT, + actual: wire_type_byte(other), + }), + } +} + +#[inline] +fn validate_enum_number(_desc: &prost_reflect::EnumDescriptor, _raw: u64) -> Result<()> { + // An unrecognized enum value is no longer a row-drop condition: the append + // path renders it as a placeholder string (see [`append_enum_name`]), so + // every varint is valid here. Kept as a named no-op so the validate pass + // stays in lock-step with the append path structurally. + Ok(()) +} + +/// Read one raw scalar value from a packed blob and yield it as a +/// `WireValue` alongside the remaining bytes. The caller reuses +/// [`append_scalar_from_wire`] for the actual append. +/// +/// Tagless: the inner blob of a packed-repeated field has no per-element +/// tags, so we dispatch on the scalar's wire type directly into the +/// tagless readers. +#[inline] +pub(super) fn read_packed_element<'a>( + kind: ScalarKind, + bytes: &'a [u8], +) -> Result<(WireValue<'a>, &'a [u8])> { + match kind.wire_type() { + WT_VARINT => { + let (v, rest) = try_read_varint(bytes)?; + Ok((WireValue::Varint(v), rest)) + } + WT_I64 => { + let Some((b, rest)) = bytes.split_first_chunk::<8>() else { + return Err(WireToArrowError::UnexpectedEof); + }; + Ok((WireValue::I64(u64::from_le_bytes(*b)), rest)) + } + WT_I32 => { + let Some((b, rest)) = bytes.split_first_chunk::<4>() else { + return Err(WireToArrowError::UnexpectedEof); + }; + Ok((WireValue::I32(u32::from_le_bytes(*b)), rest)) + } + // `WT_LEN` would be string/bytes — unreachable per the caller's guard. + // Any other value indicates a plan build bug. + _ => Err(WireToArrowError::PlanBuilderMismatch { + site: "read_packed_element:non_packable_wire_type", + }), + } +} + diff --git a/lib/codecs/src/encoding/format/wire_to_arrow/builders.rs b/lib/codecs/src/encoding/format/wire_to_arrow/builders.rs new file mode 100644 index 0000000000000..bbd067d524a98 --- /dev/null +++ b/lib/codecs/src/encoding/format/wire_to_arrow/builders.rs @@ -0,0 +1,703 @@ +//! Column builders used by the wire-to-Arrow encoder. +//! +//! Leaves (`TypedBuilder`) wrap `arrow::array::*Builder` without trait-object +//! indirection. Branch nodes (`BuilderNode::Struct`, `BuilderNode::RepeatedMessage`) +//! own their children + per-row bookkeeping (validity, list offsets). + +use std::sync::Arc; + +use arrow::array::{ + ArrayRef, BooleanBuilder, Float32Builder, Float64Builder, Int32Builder, Int64Builder, + LargeBinaryBuilder, LargeStringBuilder, ListArray, MapArray, StructArray, + TimestampMicrosecondBuilder, UInt32Builder, UInt64Builder, +}; +use arrow::buffer::{NullBuffer, OffsetBuffer, ScalarBuffer}; +use arrow::datatypes::{DataType, Field, TimeUnit}; + +use super::append::append_proto3_default; +use super::errors::{Result, WireToArrowError}; +use super::plan::{MessagePlan, PlanSlot, ScalarKind}; + +/// Leaf builder: one Arrow primitive column. Type-specific, no dyn dispatch. +pub enum TypedBuilder { + Int32(Int32Builder), + Int64(Int64Builder), + UInt32(UInt32Builder), + UInt64(UInt64Builder), + Float32(Float32Builder), + Float64(Float64Builder), + Boolean(BooleanBuilder), + LargeUtf8(LargeStringBuilder), + LargeBinary(LargeBinaryBuilder), + /// `TimestampMicrosecondBuilder` for the `_event_time` coercion and any + /// other proto int64 field whose Arrow column is declared as + /// `Timestamp(Microsecond, ...)`. The underlying i64 is written as + /// microseconds since Unix epoch — we don't transform values, only the + /// Arrow column type. + TimestampMicros(TimestampMicrosecondBuilder), +} + +/// Rough per-value byte-length hint used to pre-size the data buffer for +/// `LargeStringBuilder` / `LargeBinaryBuilder`. The builder grows on overflow, +/// so this only avoids the first few reallocations — picked to be in the right +/// order of magnitude for typical log-field values (short ids, short strings) +/// without over-allocating for columns that turn out to be mostly empty. +const AVG_VARLEN_BYTES_PER_VALUE: usize = 16; + +impl TypedBuilder { + /// True iff `dt` is one of the Arrow leaf data types this encoder can + /// build. Used at plan-build to reject unsupported types up front + /// instead of panicking in [`TypedBuilder::new`] on the first batch. + pub fn supports(dt: &DataType) -> bool { + matches!( + dt, + DataType::Int32 + | DataType::Int64 + | DataType::UInt32 + | DataType::UInt64 + | DataType::Float32 + | DataType::Float64 + | DataType::Boolean + | DataType::LargeUtf8 + | DataType::LargeBinary + | DataType::Timestamp(TimeUnit::Microsecond, _) + ) + } + + /// Construct a typed builder matching the given Arrow DataType. + /// + /// # Panics + /// Panics on unsupported primitive types. Plan-build validates the + /// Arrow leaf types via [`supports`](Self::supports) before any + /// builder is constructed, so reaching the panic here indicates a + /// plan/builder mismatch (build bug, not user-controllable input). + pub fn new(dt: &DataType, capacity: usize) -> Self { + match dt { + DataType::Int32 => TypedBuilder::Int32(Int32Builder::with_capacity(capacity)), + DataType::Int64 => TypedBuilder::Int64(Int64Builder::with_capacity(capacity)), + DataType::UInt32 => TypedBuilder::UInt32(UInt32Builder::with_capacity(capacity)), + DataType::UInt64 => TypedBuilder::UInt64(UInt64Builder::with_capacity(capacity)), + DataType::Float32 => TypedBuilder::Float32(Float32Builder::with_capacity(capacity)), + DataType::Float64 => TypedBuilder::Float64(Float64Builder::with_capacity(capacity)), + DataType::Boolean => TypedBuilder::Boolean(BooleanBuilder::with_capacity(capacity)), + DataType::LargeUtf8 => TypedBuilder::LargeUtf8(LargeStringBuilder::with_capacity( + capacity, + capacity * AVG_VARLEN_BYTES_PER_VALUE, + )), + DataType::LargeBinary => TypedBuilder::LargeBinary(LargeBinaryBuilder::with_capacity( + capacity, + capacity * AVG_VARLEN_BYTES_PER_VALUE, + )), + DataType::Timestamp(TimeUnit::Microsecond, tz) => { + let mut builder = TimestampMicrosecondBuilder::with_capacity(capacity); + if let Some(tz) = tz { + builder = builder.with_timezone(tz.clone()); + } + TypedBuilder::TimestampMicros(builder) + } + other => panic!("unsupported leaf DataType {other:?}"), + } + } + + pub fn append_null(&mut self) { + match self { + TypedBuilder::Int32(b) => b.append_null(), + TypedBuilder::Int64(b) => b.append_null(), + TypedBuilder::UInt32(b) => b.append_null(), + TypedBuilder::UInt64(b) => b.append_null(), + TypedBuilder::Float32(b) => b.append_null(), + TypedBuilder::Float64(b) => b.append_null(), + TypedBuilder::Boolean(b) => b.append_null(), + TypedBuilder::LargeUtf8(b) => b.append_null(), + TypedBuilder::LargeBinary(b) => b.append_null(), + TypedBuilder::TimestampMicros(b) => b.append_null(), + } + } + + pub fn finish(&mut self) -> ArrayRef { + match self { + TypedBuilder::Int32(b) => Arc::new(b.finish()), + TypedBuilder::Int64(b) => Arc::new(b.finish()), + TypedBuilder::UInt32(b) => Arc::new(b.finish()), + TypedBuilder::UInt64(b) => Arc::new(b.finish()), + TypedBuilder::Float32(b) => Arc::new(b.finish()), + TypedBuilder::Float64(b) => Arc::new(b.finish()), + TypedBuilder::Boolean(b) => Arc::new(b.finish()), + TypedBuilder::LargeUtf8(b) => Arc::new(b.finish()), + TypedBuilder::LargeBinary(b) => Arc::new(b.finish()), + TypedBuilder::TimestampMicros(b) => Arc::new(b.finish()), + } + } +} + +/// A tree of builders mirroring a `MessagePlan`. +pub struct BuilderNodeList { + pub(crate) nodes: Vec, + /// Per-row scratch: `present[i]` is `true` if slot `i` was touched while + /// scanning the current message at this level. Owned alongside `nodes` so + /// each sub-plan reuses its own buffer instead of allocating a fresh + /// `Vec` per nested-struct / list / map occurrence on the hot path. + /// Reset between rows / sub-rows via [`BuilderNodeList::reset_present`]. + pub(crate) present: Vec, +} + +/// One Arrow column's builder plus the dispatch info `scan_message` needs: +/// scalar kind for primitive variants, sub-plan for nested ones. Carrying it +/// here lets the scan loop dispatch on the node alone without re-indexing +/// `plan.slots`. +pub enum BuilderNode { + Scalar { + kind: ScalarKind, + builder: TypedBuilder, + }, + /// Proto enum field rendered into a STRING column by name (parity with the + /// arrow_stream path). `builder` is always a `TypedBuilder::LargeUtf8`; + /// `desc` supplies the number->name lookup at scan time. + EnumString { + desc: prost_reflect::EnumDescriptor, + builder: TypedBuilder, + }, + /// Repeated proto enum rendered into an Arrow `List` by name. + /// Same offset+current_offset bookkeeping as `RepeatedScalar`; `values` is + /// always a `TypedBuilder::LargeUtf8` and `desc` supplies the lookup. + RepeatedEnumString { + desc: prost_reflect::EnumDescriptor, + values: TypedBuilder, + offsets: Vec, + current_offset: i32, + }, + /// Singular nested message. `validity[i]` tells whether row `i` had this + /// field present (true) or absent (false — child values are null-filled). + Struct { + sub_plan: Arc, + children: BuilderNodeList, + validity: Vec, + }, + /// Repeated nested message -> Arrow `List`. + /// `offsets[i]` = total element count after row `i`. `offsets[0] = 0`. + /// `current_offset` tracks the running count across scan. + RepeatedMessage { + sub_plan: Arc, + children: BuilderNodeList, + offsets: Vec, + current_offset: i32, + }, + /// Repeated scalar -> Arrow `List`. Same offset+current_offset + /// bookkeeping as `RepeatedMessage`, but the child is a single typed + /// primitive builder rather than a tree. + RepeatedScalar { + kind: ScalarKind, + values: TypedBuilder, + offsets: Vec, + current_offset: i32, + }, + /// Proto map -> Arrow `Map`. Wire-level handling is identical + /// to `RepeatedMessage` (proto maps are `repeated MapEntry`); the finish + /// step assembles a `MapArray` reusing the user-supplied `entry_field` + /// verbatim — its name, nullability, and metadata are all preserved. + /// Arrow's Map type doesn't mandate a specific entry name (Spark favors + /// "key_value", the Arrow spec uses "entries"); honoring the caller's + /// choice is what lets `RecordBatch::try_new` accept the assembled batch. + Map { + sub_plan: Arc, + children: BuilderNodeList, + offsets: Vec, + current_offset: i32, + entry_field: Arc, + }, +} + +impl BuilderNodeList { + /// Allocate a builder tree matching `plan`, with capacity for `capacity` rows. + /// + /// Called **once per batch** by [`WireToArrowEncoder::encode_batch`], not + /// once per sink — Arrow's `*Builder::finish()` consumes the internal + /// buffers to produce the output `ArrayRef`, so the tree is single-use. + /// The shared, immutable state (`Arc`, `Arc`) lives + /// on the encoder and is what costs once per sink. + /// + /// [`WireToArrowEncoder::encode_batch`]: super::WireToArrowEncoder::encode_batch + pub fn with_capacity(plan: &MessagePlan, capacity: usize) -> Result { + let mut nodes = Vec::with_capacity(plan.slots.len()); + for (slot, field) in plan.slots.iter().zip(plan.arrow_fields.iter()) { + let node = match slot { + PlanSlot::Scalar(kind) => BuilderNode::Scalar { + kind: *kind, + builder: TypedBuilder::new(field.data_type(), capacity), + }, + PlanSlot::EnumString(desc) => BuilderNode::EnumString { + desc: desc.clone(), + builder: TypedBuilder::new(field.data_type(), capacity), + }, + PlanSlot::Struct(sub_plan) => BuilderNode::Struct { + sub_plan: Arc::clone(sub_plan), + children: BuilderNodeList::with_capacity(sub_plan, capacity)?, + validity: Vec::with_capacity(capacity), + }, + PlanSlot::RepeatedMessage(sub_plan) => { + let mut offsets = Vec::with_capacity(capacity + 1); + offsets.push(0); + BuilderNode::RepeatedMessage { + sub_plan: Arc::clone(sub_plan), + // List lengths tend to be small; 2x rows is a rough guess. + children: BuilderNodeList::with_capacity(sub_plan, capacity * 2)?, + offsets, + current_offset: 0, + } + } + PlanSlot::RepeatedScalar(kind) => { + let element_type = match field.data_type() { + DataType::List(element_field) => element_field.data_type(), + _ => { + return Err(WireToArrowError::PlanBuilderMismatch { + site: "with_capacity:repeated_scalar_non_list", + }); + } + }; + let mut offsets = Vec::with_capacity(capacity + 1); + offsets.push(0); + BuilderNode::RepeatedScalar { + kind: *kind, + values: TypedBuilder::new(element_type, capacity * 2), + offsets, + current_offset: 0, + } + } + PlanSlot::RepeatedEnumString(desc) => { + let element_type = match field.data_type() { + DataType::List(element_field) => element_field.data_type(), + _ => { + return Err(WireToArrowError::PlanBuilderMismatch { + site: "with_capacity:repeated_enum_string_non_list", + }); + } + }; + let mut offsets = Vec::with_capacity(capacity + 1); + offsets.push(0); + BuilderNode::RepeatedEnumString { + desc: desc.clone(), + values: TypedBuilder::new(element_type, capacity * 2), + offsets, + current_offset: 0, + } + } + PlanSlot::Map(sub_plan) => { + let entry_field = match field.data_type() { + DataType::Map(entry_field, _) => Arc::clone(entry_field), + _ => { + return Err(WireToArrowError::PlanBuilderMismatch { + site: "with_capacity:map_non_map_arrow_type", + }); + } + }; + let mut offsets = Vec::with_capacity(capacity + 1); + offsets.push(0); + BuilderNode::Map { + sub_plan: Arc::clone(sub_plan), + children: BuilderNodeList::with_capacity(sub_plan, capacity * 2)?, + offsets, + current_offset: 0, + entry_field, + } + } + // No proto tag points here, so the slot is null-padded each + // row by `finalize_row`'s "tag wasn't seen" branch. + PlanSlot::Absent => build_absent_node(field, capacity)?, + }; + nodes.push(node); + } + let present = vec![false; plan.slots.len()]; + Ok(Self { nodes, present }) + } + + /// Zero `present` ahead of scanning a row / sub-row. `slice::fill(false)` + /// lowers to memset. + #[inline] + pub fn reset_present(&mut self) { + self.present.fill(false); + } + + /// After scanning one message, push per-row bookkeeping (struct validity, + /// list offsets) and pad scalars whose tag wasn't seen. + /// + /// Padding rule for absent singular Scalars: + /// - **Inside a Map entry sub-plan** (`plan.inside_map_entry == true`): + /// write the proto3 scalar default (`""`, `0`, `false`, `b""`). Arrow's + /// Map type declares the key non-nullable; proto3 wire format elides + /// default-valued singular fields *inside* MapEntry messages too, so a + /// null would fail `StructArray::try_new` at finish. + /// - **Elsewhere**: write null. The plan-build non-nullability check + /// rejects schemas that can't tolerate the null up front. + /// + /// Infallible: the (`ScalarKind`, `TypedBuilder`) pairing the proto3-default + /// helper relies on is enforced at plan-build time + /// (`ScalarKind::matches_arrow_type` inside `MessagePlan::build_at_depth`), + /// so there's no per-row failure mode here that would otherwise force the + /// caller to tear down the whole batch. + #[inline] + pub fn finalize_row(&mut self, plan: &MessagePlan) { + let Self { nodes, present } = self; + debug_assert_eq!(plan.slots.len(), nodes.len()); + debug_assert_eq!(plan.slots.len(), present.len()); + let inside_map_entry = plan.inside_map_entry; + for (node, &was_present) in nodes.iter_mut().zip(present.iter()) { + match node { + BuilderNode::Scalar { kind, builder } => { + if !was_present { + if inside_map_entry { + append_proto3_default(*kind, builder); + } else { + builder.append_null(); + } + } + } + // An absent enum field is elided by proto3 at its zero value; + // `proto_to_value` only walks present fields, so it yields null + // (not the name of value 0). Match that: always null on absent. + // Enum fields never appear inside a map entry, so there is no + // proto3-default branch here. + BuilderNode::EnumString { builder, .. } => { + if !was_present { + builder.append_null(); + } + } + BuilderNode::Struct { + children, validity, .. + } => { + validity.push(was_present); + if !was_present { + children.fill_null_row(); + } + } + // All list-flavored slots push an offsets marker per row. + // For proto repeated fields (including maps), the outer list + // itself is never null — absent just means empty list. + BuilderNode::RepeatedMessage { + offsets, + current_offset, + .. + } + | BuilderNode::RepeatedScalar { + offsets, + current_offset, + .. + } + | BuilderNode::RepeatedEnumString { + offsets, + current_offset, + .. + } + | BuilderNode::Map { + offsets, + current_offset, + .. + } => { + offsets.push(*current_offset); + } + } + } + } + + /// Recursively append nulls / empty lists to the entire subtree so row + /// counts line up when a parent struct is null. + pub fn fill_null_row(&mut self) { + for node in self.nodes.iter_mut() { + match node { + BuilderNode::Scalar { builder, .. } + | BuilderNode::EnumString { builder, .. } => builder.append_null(), + BuilderNode::Struct { + children, validity, .. + } => { + validity.push(false); + children.fill_null_row(); + } + BuilderNode::RepeatedMessage { + offsets, + current_offset, + .. + } + | BuilderNode::RepeatedScalar { + offsets, + current_offset, + .. + } + | BuilderNode::RepeatedEnumString { + offsets, + current_offset, + .. + } + | BuilderNode::Map { + offsets, + current_offset, + .. + } => { + offsets.push(*current_offset); + } + } + } + } + + /// Finalize this level and return the Arrow arrays in schema order. + /// `Absent` slots fall through to the same per-variant branches: the + /// builders are already null-filled by `finalize_row`. + pub fn finish(&mut self, plan: &MessagePlan) -> Result> { + let mut out: Vec = Vec::with_capacity(plan.slots.len()); + for (idx, node) in self.nodes.iter_mut().enumerate() { + let arrow_field = &plan.arrow_fields[idx]; + let arr: ArrayRef = match node { + BuilderNode::Scalar { builder, .. } + | BuilderNode::EnumString { builder, .. } => builder.finish(), + BuilderNode::Struct { + sub_plan, + children, + validity, + } => { + let child_arrays = children.finish(sub_plan)?; + let null_buf = NullBuffer::from(std::mem::take(validity)); + Arc::new( + StructArray::try_new( + sub_plan.arrow_fields.clone(), + child_arrays, + Some(null_buf), + ) + .map_err(|e| WireToArrowError::ArrayAssembly { + kind: "struct", + source: e, + })?, + ) + } + BuilderNode::RepeatedMessage { + sub_plan, + children, + offsets, + .. + } => { + let child_arrays = children.finish(sub_plan)?; + let struct_arr = + StructArray::try_new(sub_plan.arrow_fields.clone(), child_arrays, None) + .map_err(|e| WireToArrowError::ArrayAssembly { + kind: "list element struct", + source: e, + })?; + let offset_buffer = + OffsetBuffer::new(ScalarBuffer::from(std::mem::take(offsets))); + let element_field = Arc::new(Field::new( + "item", + DataType::Struct(sub_plan.arrow_fields.clone()), + true, + )); + Arc::new( + ListArray::try_new( + element_field, + offset_buffer, + Arc::new(struct_arr), + None, + ) + .map_err(|e| WireToArrowError::ArrayAssembly { + kind: "list", + source: e, + })?, + ) + } + BuilderNode::RepeatedScalar { values, offsets, .. } => { + let values_array = values.finish(); + let offset_buffer = + OffsetBuffer::new(ScalarBuffer::from(std::mem::take(offsets))); + // Preserve the element field the schema declared (name + // typically "item", but follow the caller's choice). + let element_field = match arrow_field.data_type() { + DataType::List(f) => Arc::clone(f), + other => { + return Err(WireToArrowError::UnsupportedCombination { + name: arrow_field.name().to_string(), + kind: "RepeatedScalar".to_string(), + arrow_type: format!("{other:?}"), + repeated: true, + }); + } + }; + Arc::new( + ListArray::try_new(element_field, offset_buffer, values_array, None) + .map_err(|e| WireToArrowError::ArrayAssembly { + kind: "list (scalar)", + source: e, + })?, + ) + } + BuilderNode::RepeatedEnumString { values, offsets, .. } => { + let values_array = values.finish(); + let offset_buffer = + OffsetBuffer::new(ScalarBuffer::from(std::mem::take(offsets))); + let element_field = match arrow_field.data_type() { + DataType::List(f) => Arc::clone(f), + other => { + return Err(WireToArrowError::UnsupportedCombination { + name: arrow_field.name().to_string(), + kind: "RepeatedEnumString".to_string(), + arrow_type: format!("{other:?}"), + repeated: true, + }); + } + }; + Arc::new( + ListArray::try_new(element_field, offset_buffer, values_array, None) + .map_err(|e| WireToArrowError::ArrayAssembly { + kind: "list (enum string)", + source: e, + })?, + ) + } + BuilderNode::Map { + sub_plan, + children, + offsets, + entry_field, + .. + } => { + let child_arrays = children.finish(sub_plan)?; + let struct_arr = + StructArray::try_new(sub_plan.arrow_fields.clone(), child_arrays, None) + .map_err(|e| WireToArrowError::ArrayAssembly { + kind: "map entry struct", + source: e, + })?; + let offset_buffer = + OffsetBuffer::new(ScalarBuffer::from(std::mem::take(offsets))); + // Reuse the user-supplied entry Field unchanged: Arrow's + // Map spec doesn't pin the entry name ("entries" is + // canonical; Spark/Delta use "key_value"), and + // a name mismatch makes `RecordBatch::try_new` reject the + // whole batch at finish. Honoring the caller's name + + // nullability + metadata avoids that footgun. The encoder + // never emits null entries (only empty maps), so a + // declared non-nullable entry is also safe. + Arc::new( + MapArray::try_new( + Arc::clone(entry_field), + offset_buffer, + struct_arr, + None, + false, + ) + .map_err(|e| WireToArrowError::ArrayAssembly { + kind: "map", + source: e, + })?, + ) + } + }; + out.push(arr); + } + Ok(out) + } +} + +/// Build a null-filled builder tree matching the Arrow `field`'s shape, used +/// for [`PlanSlot::Absent`] columns. The builder is the same shape as a normal +/// column of that Arrow type, but no wire tags ever dispatch to it so it stays +/// fully null-padded by `finalize_row` / `fill_null_row`. +fn build_absent_node(field: &Field, capacity: usize) -> Result { + // Inert `kind` for variants that carry one — the slot is never + // dispatched, so the value never matters. + const ABSENT_KIND: ScalarKind = ScalarKind::Int32; + Ok(match field.data_type() { + DataType::Struct(inner_fields) => { + let sub_plan = Arc::new(MessagePlan::all_absent(inner_fields)); + BuilderNode::Struct { + children: BuilderNodeList::with_capacity(&sub_plan, capacity)?, + sub_plan, + validity: Vec::with_capacity(capacity), + } + } + DataType::List(element_field) => { + let mut offsets = Vec::with_capacity(capacity + 1); + offsets.push(0); + match element_field.data_type() { + DataType::Struct(inner_fields) => { + let sub_plan = Arc::new(MessagePlan::all_absent(inner_fields)); + BuilderNode::RepeatedMessage { + children: BuilderNodeList::with_capacity(&sub_plan, capacity * 2)?, + sub_plan, + offsets, + current_offset: 0, + } + } + _ => BuilderNode::RepeatedScalar { + kind: ABSENT_KIND, + values: TypedBuilder::new(element_field.data_type(), capacity * 2), + offsets, + current_offset: 0, + }, + } + } + DataType::Map(entry_field, _) => { + let inner_fields = match entry_field.data_type() { + DataType::Struct(fs) => fs, + _ => { + return Err(WireToArrowError::PlanBuilderMismatch { + site: "build_absent_node:map_entry_non_struct", + }); + } + }; + let sub_plan = Arc::new(MessagePlan::all_absent(inner_fields)); + let mut offsets = Vec::with_capacity(capacity + 1); + offsets.push(0); + BuilderNode::Map { + children: BuilderNodeList::with_capacity(&sub_plan, capacity * 2)?, + sub_plan, + offsets, + current_offset: 0, + entry_field: Arc::clone(entry_field), + } + } + // Scalar Arrow types — build a primitive builder. `TypedBuilder::new` + // still has an internal `unsupported leaf DataType` panic, but it's + // a build-bug-only path: `validate_arrow_leaf_types` rejects + // unsupported leaves at plan-build, so this call can't see one. + _ => BuilderNode::Scalar { + kind: ABSENT_KIND, + builder: TypedBuilder::new(field.data_type(), capacity), + }, + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::{Array, AsArray}; + + #[test] + fn int32_builder_roundtrip() { + let mut tb = TypedBuilder::new(&DataType::Int32, 4); + if let TypedBuilder::Int32(b) = &mut tb { + b.append_value(1); + b.append_value(2); + } + tb.append_null(); + let arr = tb.finish(); + let i32arr = arr.as_primitive::(); + assert_eq!(i32arr.len(), 3); + assert_eq!(i32arr.value(0), 1); + assert_eq!(i32arr.value(1), 2); + assert!(i32arr.is_null(2)); + } + + #[test] + fn string_builder_roundtrip() { + let mut tb = TypedBuilder::new(&DataType::LargeUtf8, 4); + if let TypedBuilder::LargeUtf8(b) = &mut tb { + b.append_value("hello"); + b.append_value("world"); + } + tb.append_null(); + let arr = tb.finish(); + assert_eq!(arr.len(), 3); + } + + #[test] + #[should_panic(expected = "unsupported leaf DataType")] + fn unsupported_type_panics() { + // Date32 is not in our scope. + let _ = TypedBuilder::new(&DataType::Date32, 1); + } +} diff --git a/lib/codecs/src/encoding/format/wire_to_arrow/encoder.rs b/lib/codecs/src/encoding/format/wire_to_arrow/encoder.rs new file mode 100644 index 0000000000000..36c6af6229c26 --- /dev/null +++ b/lib/codecs/src/encoding/format/wire_to_arrow/encoder.rs @@ -0,0 +1,98 @@ +//! [`WireToArrowEncoder`]: builds a [`MessagePlan`] once per +//! (proto descriptor, Arrow schema) pair and encodes batches of proto wire +//! bytes into Arrow `RecordBatch`es. + +use std::sync::Arc; + +use arrow::datatypes::{Fields, Schema}; +use arrow::record_batch::RecordBatch; +use bytes::Bytes; +use prost_reflect::MessageDescriptor; + +use super::builders::BuilderNodeList; +use super::errors::{Result, WireToArrowError}; +use super::plan::MessagePlan; +use super::scan::{scan_message, validate_message}; + +/// Streaming wire-format encoder. Build once per (proto message type, +/// Arrow schema) pair, then call [`WireToArrowEncoder::encode_batch`] +/// repeatedly. +pub struct WireToArrowEncoder { + plan: Arc, + schema: Arc, +} + +impl std::fmt::Debug for WireToArrowEncoder { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("WireToArrowEncoder") + .field("schema_fields", &self.schema.fields().len()) + .finish() + } +} + +impl WireToArrowEncoder { + /// Compile a plan for the given proto descriptor + Arrow schema. + /// + /// Every field in `schema` must exist (by name) in `descriptor`. Proto + /// fields absent from `schema` are silently skipped at scan time. + pub fn new(descriptor: &MessageDescriptor, schema: Schema) -> Result { + let plan = MessagePlan::build(descriptor, &Fields::from(schema.fields().clone()))?; + Ok(Self { + plan: Arc::new(plan), + schema: Arc::new(schema), + }) + } + + /// Encode a batch of serialized proto messages into a single `RecordBatch`. + /// + /// Per-row isolation: each message is pre-validated via + /// [`validate_message`] before any builder is touched. Rows that fail + /// validation are dropped from the output batch and counted via the + /// `wire_to_arrow_rows_dropped` metric (plus a rate-limit-friendly + /// warn log carrying a sample error). Returning an empty `RecordBatch` + /// is acceptable when every row was malformed. + /// + /// Errors out of this method are reserved for batch-level failures + /// that aren't attributable to a single row: a code-bug surface + /// (scan-vs-validate divergence surfaced as `PlanBuilderMismatch`) or + /// a `RecordBatchAssembly` rejection from Arrow. Row-finalize is + /// infallible, so it doesn't appear in this list. + pub fn encode_batch(&self, messages: &[Bytes]) -> Result { + let capacity = messages.len(); + let mut builders = BuilderNodeList::with_capacity(&self.plan, capacity)?; + let mut dropped = 0u64; + let mut sample_err: Option = None; + + for msg_bytes in messages { + // Pre-validate so a malformed row drops without poisoning any + // builder. Arrow `*Builder` has no public rollback API, and + // nested-struct `finalize_row` calls inside `scan_message` are + // not reversible, so an upfront validation pass is how we + // isolate per-row decode failures. + if let Err(err) = validate_message(&self.plan, msg_bytes) { + dropped += 1; + if sample_err.is_none() { + sample_err = Some(err); + } + continue; + } + builders.reset_present(); + scan_message(&self.plan, msg_bytes, &mut builders)?; + builders.finalize_row(&self.plan); + } + + if dropped > 0 { + metrics::counter!("wire_to_arrow_rows_dropped").increment(dropped); + tracing::warn!( + message = "wire-to-Arrow dropped malformed rows from batch", + dropped, + batch_size = messages.len(), + sample_error = ?sample_err, + ); + } + + let arrays = builders.finish(&self.plan)?; + RecordBatch::try_new(Arc::clone(&self.schema), arrays) + .map_err(|source| WireToArrowError::RecordBatchAssembly { source }) + } +} diff --git a/lib/codecs/src/encoding/format/wire_to_arrow/errors.rs b/lib/codecs/src/encoding/format/wire_to_arrow/errors.rs new file mode 100644 index 0000000000000..478074bdd1227 --- /dev/null +++ b/lib/codecs/src/encoding/format/wire_to_arrow/errors.rs @@ -0,0 +1,321 @@ +//! Error types for the streaming wire-to-Arrow encoder. + +use snafu::Snafu; + +/// Errors that can occur when building an encoding plan or encoding a batch. +#[derive(Debug, Snafu)] +#[snafu(visibility(pub(crate)))] +pub enum WireToArrowError { + /// Required serializer-config field (`schema`) was not populated before + /// `build()`. Sinks inject this at config build time. + #[snafu(display("wire-to-Arrow serializer requires a {field}"))] + ConfigurationMissing { + /// Which config field was missing. + field: &'static str, + }, + + /// Failed to load the proto descriptor from `desc_file` / `message_type`. + #[snafu(display("failed to load proto descriptor: {message}"))] + DescriptorLoad { + /// The underlying error message from `get_message_descriptor`. + message: String, + }, + + /// The batch had no events to encode. + #[snafu(display("cannot encode an empty batch"))] + NoEvents, + + /// An event in the batch had no `message` field. + #[snafu(display("event is missing a `message` field"))] + MessageBytesMissing, + + /// An event's `message` field was not a `Value::Bytes`. + #[snafu(display("event `message` is not bytes-typed"))] + MessageBytesWrongType, + + /// Proto descriptor is missing a field named in the Arrow schema. + #[snafu(display("proto field '{name}' not found in descriptor"))] + MissingProtoField { + /// The Arrow field name that wasn't found in the proto descriptor. + name: String, + }, + + /// The proto field's Kind cannot be represented by any supported scalar. + #[snafu(display("unsupported proto kind for field '{name}': {kind}"))] + UnsupportedKind { + /// The proto field name. + name: String, + /// The proto kind that isn't supported. + kind: String, + }, + + /// The combination of proto kind, Arrow type, and cardinality isn't supported. + #[snafu(display( + "field '{name}': unsupported combination \ + proto kind {kind} / arrow type {arrow_type} / repeated {repeated}" + ))] + UnsupportedCombination { + /// The proto field name. + name: String, + /// The proto kind involved. + kind: String, + /// The Arrow data type involved. + arrow_type: String, + /// Whether the proto field was repeated. + repeated: bool, + }, + + /// A repeated message field's Arrow element type isn't Struct. + #[snafu(display( + "field '{name}': repeated message requires List, got List<{element}>" + ))] + RepeatedNonStructList { + /// The proto field name. + name: String, + /// The Arrow list-element type that was found (expected Struct). + element: String, + }, + + /// Ran out of wire bytes before finishing a tag/field. + #[snafu(display("unexpected end of wire input"))] + UnexpectedEof, + + /// Varint exceeded the max 10-byte encoding. + #[snafu(display("varint exceeds 10 bytes"))] + VarintOverflow, + + /// Unknown proto wire type (should be 0, 1, 2, or 5). + #[snafu(display("invalid proto wire type {wire_type}"))] + InvalidWireType { + /// The unrecognized wire-type value from the tag. + wire_type: u8, + }, + + /// Wire type for a field doesn't match the plan's expectation. + #[snafu(display( + "wire type mismatch: plan expected {expected}, wire bytes had {actual}" + ))] + WireTypeMismatch { + /// The wire type the plan expected for this field. + expected: u8, + /// The wire type actually present in the bytes. + actual: u8, + }, + + /// String field contained non-UTF-8 bytes. + #[snafu(display("invalid UTF-8 in proto string field"))] + InvalidUtf8, + + /// Plan and builder trees diverged during scan / finish, or a code + /// path the encoder considers structurally impossible was reached. + /// Always a code bug — never user input. `site` is a short label + /// identifying which emit site fired so a bug report points at the + /// right path without needing a backtrace. + #[snafu(display("internal: plan/builder tree mismatch at {site}"))] + PlanBuilderMismatch { + /// Short label naming the emit site (e.g. `"scan_message"`, + /// `"finish:absent_struct_non_struct_arrow"`). Free-form but + /// expected to be a `&'static str` literal at the call site. + site: &'static str, + }, + + /// `arrow::record_batch::RecordBatch::try_new` rejected the assembled arrays. + #[snafu(display("failed to assemble RecordBatch: {source}"))] + RecordBatchAssembly { + /// The underlying Arrow error from `RecordBatch::try_new`. + source: arrow::error::ArrowError, + }, + + /// `arrow::array::StructArray::try_new` / `ListArray::try_new` rejected + /// the assembled arrays. + #[snafu(display("failed to assemble {kind} array: {source}"))] + ArrayAssembly { + /// Which kind of array failed to assemble (e.g. "struct", "list"). + kind: &'static str, + /// The underlying Arrow error from the array constructor. + source: arrow::error::ArrowError, + }, + + /// A wire-parse error the encoder doesn't model as one of the + /// variants above. + #[snafu(display("wire parse error: {source}"))] + ProtoParser { + /// The underlying [`super::wire::ParseError`]. + source: super::wire::ParseError, + }, + + /// Plan-build recursion exceeded [`MAX_NESTING_DEPTH`]. Caps both the + /// build-time walk over the Arrow schema and the scan-time walk over + /// wire bytes (which can't recurse deeper than the plan). + /// + /// [`MAX_NESTING_DEPTH`]: super::plan::MAX_NESTING_DEPTH + #[snafu(display( + "wire-to-Arrow plan exceeds max nesting depth of {limit}" + ))] + SchemaTooDeep { + /// The configured maximum depth. + limit: usize, + }, + + /// The Arrow schema declares a primitive leaf type the encoder can't + /// build a column for (e.g. `Date32`, `Time64`, decimal). Caught at + /// plan-build so the failure surfaces at serializer init rather than + /// panicking inside `TypedBuilder::new` on the first batch. + #[snafu(display( + "Arrow field '{name}' has unsupported leaf data type {arrow_type}" + ))] + UnsupportedArrowLeafType { + /// The Arrow field name carrying the unsupported leaf type. + name: String, + /// The unsupported Arrow data type (debug form). + arrow_type: String, + }, + + /// A repeated-list slot's running offset would overflow `i32`. + /// + /// Arrow's `ListArray` / `MapArray` use `i32` offsets, so the cumulative + /// element count across all rows in a batch is capped at `i32::MAX` + /// (~2.1B). The encoder normally bumps the counter via `+= 1` inside + /// `scan_message` / `append_repeated_scalar`; without bounds checking, + /// release-mode wrap-around silently produces a non-monotonic offsets + /// buffer and `OffsetBuffer::new` asserts at batch finish, taking the + /// whole process down. + /// + /// Two places guard against this: + /// - `validate_message`'s packed-scalar count drops a single row whose + /// own delta would already exceed `i32::MAX` (per-row isolation). + /// - The runtime appenders use `checked_add` and surface this variant if + /// the cumulative count (across rows + sub-rows in the batch) would + /// wrap. That path fails the batch cleanly with a structured error + /// rather than panicking the process — adversarial wire bytes can no + /// longer crash the encoder regardless of how many rows they span. + #[snafu(display("repeated-list offset would overflow i32 at {site}"))] + OffsetOverflow { + /// Short label naming the site that detected the overflow (e.g. + /// `"scan_message:repeated_message"`, `"append_repeated_scalar:packed"`). + site: &'static str, + }, + + /// A singular (non-repeated) proto field appeared more than once in a + /// single message. Proto3 parsers must accept this (last-wins for + /// scalars, merge for sub-messages), but the encoder appends to + /// Arrow column builders on every occurrence, which would diverge + /// column lengths and fail batch assembly. Surfaced from + /// `validate_message` so the offending row is dropped via the normal + /// per-row isolation path instead of poisoning the whole batch. + #[snafu(display( + "duplicate singular proto field {field_number} in one message" + ))] + DuplicateSingularField { + /// The proto field number whose tag appeared more than once. + field_number: u32, + }, + + /// Arrow Map's entry struct declares a field name that the proto + /// MapEntry descriptor doesn't carry. + /// + /// Proto `map` is generated as a `MapEntry` message with fields + /// named `key` (1) and `value` (2). Arrow's Map type doesn't enforce + /// these names on its inner Struct, so a user-supplied schema can + /// declare `Map` and have it pass type checking — but + /// there is no proto field for the encoder to read into the slot, the + /// `MapEntry`'s key non-null contract still applies, and the + /// absent-padding path can't honor both. Reject the mismatch at + /// plan-build so it surfaces clearly at sink init rather than as a + /// runtime panic when the proto3-default helper meets a kind/builder + /// pair it can't pad. + #[snafu(display( + "Arrow Map entry has field '{name}' which is not declared in the proto MapEntry; \ + entry struct field names must match proto MapEntry's 'key' and 'value'" + ))] + MapEntryFieldNotInProto { + /// The Arrow Map entry field name that doesn't match proto MapEntry. + name: String, + }, + + /// The Arrow schema declares a singular column as non-nullable, but the + /// encoder cannot guarantee a value will be present on every row. + /// + /// Proto3 omits default-valued singular fields on the wire, so the + /// encoder writes a null whenever a tag is absent. A non-nullable + /// declaration would then trip a generic `RecordBatch::try_new` + /// "non-nullable contains nulls" failure deep in `encode_batch`, + /// dropping the entire batch with no row context. We reject the + /// mismatch at plan-build time so it surfaces clearly at serializer + /// init. + /// + /// `List<…>` and `Map<…>` outer columns are exempt: the encoder always + /// emits at least an empty list / empty map per row, so the outer + /// column never contains a null. + #[snafu(display( + "Arrow field '{name}' is declared non-nullable but {reason}; \ + declare the column nullable in the schema or change the proto" + ))] + NonNullableNotGuaranteed { + /// The Arrow field name. + name: String, + /// Short explanation of why the encoder can't guarantee non-null + /// (e.g. "proto3 singular fields are omitted at default value"). + reason: &'static str, + }, +} + +/// Result alias for wire-to-Arrow encoder operations. +pub type Result = std::result::Result; + +impl From for WireToArrowError { + /// Collapse most [`super::wire::ParseError`]s onto this crate's pre-existing + /// variants; the long tail falls through into [`WireToArrowError::ProtoParser`]. + fn from(err: super::wire::ParseError) -> Self { + use super::wire::ParseError; + match err { + ParseError::TruncatedVarint | ParseError::BufferTooShort { .. } => { + WireToArrowError::UnexpectedEof + } + ParseError::VarintTooLong => WireToArrowError::VarintOverflow, + ParseError::InvalidWireType(wt) => WireToArrowError::InvalidWireType { wire_type: wt }, + _ => WireToArrowError::ProtoParser { source: err }, + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use super::super::wire::ParseError; + + #[test] + fn error_display_is_informative() { + let e = WireToArrowError::MissingProtoField { + name: "foo".to_string(), + }; + let msg = format!("{e}"); + assert!(msg.contains("foo"), "display must contain field name: {msg}"); + } + + #[test] + fn wire_type_mismatch_displays_both_numbers() { + let e = WireToArrowError::WireTypeMismatch { + expected: 2, + actual: 0, + }; + let msg = format!("{e}"); + assert!(msg.contains("2") && msg.contains("0"), "got: {msg}"); + } + + #[test] + fn from_parse_error_truncated_varint() { + assert!(matches!( + WireToArrowError::from(ParseError::TruncatedVarint), + WireToArrowError::UnexpectedEof + )); + } + + #[test] + fn from_parse_error_invalid_wire_type() { + assert!(matches!( + WireToArrowError::from(ParseError::InvalidWireType(7)), + WireToArrowError::InvalidWireType { wire_type: 7 } + )); + } +} diff --git a/lib/codecs/src/encoding/format/wire_to_arrow/mod.rs b/lib/codecs/src/encoding/format/wire_to_arrow/mod.rs new file mode 100644 index 0000000000000..5f75e1930aad4 --- /dev/null +++ b/lib/codecs/src/encoding/format/wire_to_arrow/mod.rs @@ -0,0 +1,58 @@ +//! Streaming wire-format to Arrow encoder. +//! +//! Parses proto wire bytes in a single pass and appends values directly into +//! Arrow `RecordBatch` column builders, skipping the `DynamicMessage` / +//! `LogEvent` intermediate representations used by the generic +//! `ProtobufDeserializer` + `ArrowStreamSerializer` path. +//! +//! Used as a [`BatchSerializerConfig`] variant — the upstream source/transform +//! is expected to place the original proto wire bytes in the event's `message` +//! field (Vector convention). +//! +//! Failure semantics are split: +//! * Event-shape problems (missing `message` field, non-`Bytes` value) fail +//! the batch — the pipeline is misconfigured if any event reaches here in +//! the wrong shape. +//! * Wire-format decode errors are isolated to the offending row: the row +//! is dropped from the output `RecordBatch`, counted via the +//! `wire_to_arrow_rows_dropped` metric, and a sample error is logged. +//! One poison message can't poison the whole batch. +//! +//! ## Scope +//! +//! The encoder takes one `MessageDescriptor` and decodes the bytes in +//! `event.message` against it, emitting one `RecordBatch` row per event. It +//! is agnostic to how the caller produced those bytes and to what any +//! particular schema represents. If the incoming payload requires any +//! pre-processing — multi-frame unwrapping, decompression, merging bytes from +//! multiple sources, sink-time / build-time stamps — perform it upstream (in +//! VRL or a custom transform) so that `event.message` holds a single +//! self-contained byte stream that matches the configured descriptor. +//! +//! ## Supported today +//! +//! - Scalar proto fields (int32/int64/uint32/uint64/sint32/sint64/fixed*/float/double/bool/string/bytes/enum) +//! - Singular nested messages -> Arrow `Struct` +//! - Repeated nested messages -> Arrow `List` +//! - Repeated scalars (packed and unpacked) -> Arrow `List` +//! - Proto maps (`map`) -> Arrow `Map` +//! - Oneof variants +//! - `int64 -> Timestamp(Microsecond, tz)` coercion +//! +//! [`BatchSerializerConfig`]: crate::encoding::BatchSerializerConfig + +mod append; +mod builders; +mod encoder; +mod errors; +mod plan; +mod scan; +mod serializer; +mod wire; + +#[cfg(test)] +mod tests; + +pub use encoder::WireToArrowEncoder; +pub use errors::WireToArrowError; +pub use serializer::{WireToArrowSerializer, WireToArrowSerializerConfig}; diff --git a/lib/codecs/src/encoding/format/wire_to_arrow/plan.rs b/lib/codecs/src/encoding/format/wire_to_arrow/plan.rs new file mode 100644 index 0000000000000..84acfd3590bff --- /dev/null +++ b/lib/codecs/src/encoding/format/wire_to_arrow/plan.rs @@ -0,0 +1,674 @@ +//! Encoding plan: a tree describing how to decode a given proto message into +//! a set of Arrow column builders. +//! +//! Built once per (proto descriptor, Arrow schema) pair. Immutable after build. + +use std::sync::Arc; + +use arrow::datatypes::{DataType, Fields, TimeUnit}; +use prost_reflect::{Cardinality, EnumDescriptor, Kind, MessageDescriptor}; + +use super::builders::TypedBuilder; +use super::errors::{Result, WireToArrowError}; + +/// Maximum nesting depth permitted in a plan. Arrow schemas can in principle +/// nest arbitrarily deep, but the build walk recurses 1:1 with structural +/// depth and would blow the Rust call stack on pathological input. Real-world +/// schemas are well under this; the cap exists to keep DoS-shaped input +/// (deep schema at plan-build time, or deep wire-bytes nesting at scan time) +/// from running to stack overflow. +/// +/// `scan_message` recursion is bounded by the plan, so capping the plan caps +/// both paths. +pub const MAX_NESTING_DEPTH: usize = 64; + +/// Proto wire-type codes (the low 3 bits of a tag). +/// +/// The [`super::wire::WireType`] enum is private to that module, so we +/// redeclare the codes here for use across this module's public-API +/// surface (`ScalarKind::wire_type`, error fields, packed-scalar dispatch). +/// Keep these in sync with the proto spec: +pub(super) const WT_VARINT: u8 = 0; +pub(super) const WT_I64: u8 = 1; +pub(super) const WT_LEN: u8 = 2; +pub(super) const WT_I32: u8 = 5; + +/// "No slot maps to this proto field number" in [`MessagePlan::slot_by_proto_field`]. +pub(crate) const SLOT_UNKNOWN: u32 = u32::MAX; + +/// Proto scalar kinds that this encoder can read off the wire and append to +/// Arrow primitive builders. Proto enums map to `Int32` by default (or to a +/// `LargeUtf8` value name when the target column is a string). +#[derive(Clone, Copy, Debug)] +pub enum ScalarKind { + Int32, + Int64, + UInt32, + UInt64, + SInt32, + SInt64, + Bool, + Fixed32, + SFixed32, + Float, + Fixed64, + SFixed64, + Double, + String, + Bytes, +} + +impl ScalarKind { + /// Proto wire type expected for values of this kind. + pub fn wire_type(self) -> u8 { + match self { + ScalarKind::Int32 + | ScalarKind::Int64 + | ScalarKind::UInt32 + | ScalarKind::UInt64 + | ScalarKind::SInt32 + | ScalarKind::SInt64 + | ScalarKind::Bool => WT_VARINT, + ScalarKind::Fixed32 | ScalarKind::SFixed32 | ScalarKind::Float => WT_I32, + ScalarKind::Fixed64 | ScalarKind::SFixed64 | ScalarKind::Double => WT_I64, + ScalarKind::String | ScalarKind::Bytes => WT_LEN, + } + } + + /// True iff `dt` is the Arrow leaf type that pairs with this proto + /// scalar kind. Enforced at plan-build by + /// [`MessagePlan::build_at_depth`] so the runtime appenders + /// ([`append_scalar_from_wire`] and [`append_proto3_default`]) can + /// assume the pairing is well-formed and don't need a runtime fallthrough + /// for kind/builder mismatches. + /// + /// [`append_scalar_from_wire`]: super::append::append_scalar_from_wire + /// [`append_proto3_default`]: super::append::append_proto3_default + pub(super) fn matches_arrow_type(self, dt: &DataType) -> bool { + match (self, dt) { + (ScalarKind::Int32 | ScalarKind::SInt32 | ScalarKind::SFixed32, DataType::Int32) => { + true + } + ( + ScalarKind::Int64 | ScalarKind::SInt64 | ScalarKind::SFixed64, + DataType::Int64 | DataType::Timestamp(TimeUnit::Microsecond, _), + ) => true, + (ScalarKind::UInt32 | ScalarKind::Fixed32, DataType::UInt32) => true, + (ScalarKind::UInt64 | ScalarKind::Fixed64, DataType::UInt64) => true, + (ScalarKind::Float, DataType::Float32) => true, + (ScalarKind::Double, DataType::Float64) => true, + (ScalarKind::Bool, DataType::Boolean) => true, + (ScalarKind::String, DataType::LargeUtf8) => true, + (ScalarKind::Bytes, DataType::LargeBinary) => true, + _ => false, + } + } + + /// Map a `prost_reflect::Kind` to a `ScalarKind`. Returns `None` for Kinds + /// that aren't scalars (Message types are handled at the plan level). + pub fn from_proto_kind(kind: &Kind) -> Option { + Some(match kind { + Kind::Int32 => ScalarKind::Int32, + Kind::Int64 => ScalarKind::Int64, + Kind::Uint32 => ScalarKind::UInt32, + Kind::Uint64 => ScalarKind::UInt64, + Kind::Sint32 => ScalarKind::SInt32, + Kind::Sint64 => ScalarKind::SInt64, + Kind::Fixed32 => ScalarKind::Fixed32, + Kind::Fixed64 => ScalarKind::Fixed64, + Kind::Sfixed32 => ScalarKind::SFixed32, + Kind::Sfixed64 => ScalarKind::SFixed64, + Kind::Float => ScalarKind::Float, + Kind::Double => ScalarKind::Double, + Kind::Bool => ScalarKind::Bool, + Kind::String => ScalarKind::String, + Kind::Bytes => ScalarKind::Bytes, + // Proto enums carry over as int32 on the Arrow side. + Kind::Enum(_) => ScalarKind::Int32, + Kind::Message(_) => return None, + }) + } +} + +/// One entry per Arrow field at this message level: describes how to route +/// wire-bytes values into the corresponding Arrow column builder. +#[derive(Debug)] +pub enum PlanSlot { + Scalar(ScalarKind), + Struct(Arc), + RepeatedMessage(Arc), + /// Repeated scalar field (e.g. `repeated int32`) -> Arrow `List`. + /// Handles both packed and unpacked wire encodings at scan time. + RepeatedScalar(ScalarKind), + /// Proto enum field paired to a STRING Arrow column: render the varint as + /// its enum-value *name* (e.g. `1` -> `"SUCCESS"`), matching the + /// arrow_stream / `proto_to_value` path. The default enum mapping is + /// `ScalarKind::Int32` (see [`ScalarKind::from_proto_kind`]); this variant + /// is only chosen when the target Arrow leaf is `LargeUtf8`. Carries the + /// [`EnumDescriptor`] for the number->name lookup at scan time. Kept out of + /// [`ScalarKind`] so that enum stays `Copy` and the primitive hot path is + /// untouched. + EnumString(EnumDescriptor), + /// Repeated proto enum field paired to an Arrow `List` column: + /// render each enum varint as its value name (the repeated analogue of + /// [`PlanSlot::EnumString`]). Handles both packed and unpacked wire + /// encodings at scan time, like [`PlanSlot::RepeatedScalar`]. A repeated + /// enum paired with `List` still falls through to + /// [`PlanSlot::RepeatedScalar`] and stays numeric. + RepeatedEnumString(EnumDescriptor), + /// Proto `map` -> Arrow `Map`. On the wire, maps + /// are encoded as `repeated MapEntry` where `MapEntry` is a generated + /// message with field 1 = key and field 2 = value; we scan them the same + /// way as `RepeatedMessage` and assemble a `MapArray` at finish time. + Map(Arc), + /// Arrow column has no matching proto field — always emits null (or empty + /// list / all-null struct). Happens when the Arrow schema has more columns + /// than the producer's proto — typically because a field was deleted from + /// the proto schema but the Arrow schema hasn't been updated yet, or the + /// producer is running an older version. The scanner never dispatches to + /// these slots; `finalize_row` null-pads them for every row. + Absent, +} + +/// Plan for encoding one proto message type into a set of Arrow column builders. +#[derive(Debug)] +pub struct MessagePlan { + /// One entry per Arrow field at this level, in schema order. + pub(crate) slots: Vec, + /// Reverse index from proto field number to slot index, with + /// [`SLOT_UNKNOWN`] marking unknown / out-of-range fields. Dense vector, + /// no hashing on the hot path. + pub(crate) slot_by_proto_field: Vec, + /// Arrow `Fields` at this level, kept for assembly of `StructArray` / `ListArray`. + pub(crate) arrow_fields: Fields, + /// True if this plan describes the entry sub-message of a `map` + /// slot. Proto3 elides singular fields at their default value on the + /// wire — including *inside* `MapEntry` messages — but Arrow's Map type + /// declares the key field non-nullable. `finalize_row` consults this + /// flag and materializes the proto3 scalar default (e.g. `""` for + /// String, `0` for Int32) for absent scalar slots instead of writing a + /// null, which would fail `StructArray::try_new` at batch finish. + pub(crate) inside_map_entry: bool, +} + +impl MessagePlan { + /// Build a plan from a proto message descriptor and a matching Arrow `Fields`. + /// + /// Fields in the Arrow schema must exist (by name) in the proto descriptor. + /// Proto fields absent from the Arrow schema are treated as unknown and will + /// be skipped at scan time. + /// + /// # Self-referential proto types + /// + /// Proto schemas can reference themselves (e.g. `message Tree { Tree left + /// = 1; }`), but Arrow schemas cannot carry a recursive type. Recursion + /// in this builder terminates because we only descend into + /// `Kind::Message(_)` fields when the Arrow target at that path is also + /// a nested type (`Struct` / `List` / `Map`). Arrow schemas are + /// finite by construction (Arrow schemas don't produce cyclic types), so + /// each recursion step strictly reduces the + /// remaining Arrow depth. Proto self-reference past the depth declared + /// in the Arrow schema is treated as an unknown field and skipped at + /// scan time. + /// + /// A hard depth cap of [`MAX_NESTING_DEPTH`] guards against pathological + /// schemas that would otherwise overflow the Rust call stack at build + /// time. `scan_message`'s recursion is bounded by the plan, so this cap + /// also bounds the scan-time recursion driven by attacker-controlled + /// wire bytes. + /// + /// # Oneof + /// + /// Proto `oneof` is purely an annotation; on the wire each variant is a + /// normal singular field with its own tag, and the receiver takes + /// whichever variant appeared last in the bytes. No special handling is + /// needed at the plan level — each variant becomes its own `PlanSlot` + /// (Scalar / Struct / etc.) and the normal "absent slot => null" + /// machinery produces the correct Arrow output. + pub fn build(descriptor: &MessageDescriptor, fields: &Fields) -> Result { + Self::build_at_depth(descriptor, fields, 0, /* inside_map_entry */ false) + } + + /// Recursive helper for [`build`]; `depth` is the current nesting level + /// (0 at the top), `inside_map_entry` is true when called for a Map's + /// entry sub-plan. Map entries have Arrow-spec-mandated nullability + /// (key non-nullable, value typically nullable), so the singular-field + /// nullability check is suppressed inside that recursion to avoid + /// false positives. Returns [`WireToArrowError::SchemaTooDeep`] once + /// the level being built would exceed [`MAX_NESTING_DEPTH`]. + fn build_at_depth( + descriptor: &MessageDescriptor, + fields: &Fields, + depth: usize, + inside_map_entry: bool, + ) -> Result { + if depth >= MAX_NESTING_DEPTH { + return Err(WireToArrowError::SchemaTooDeep { + limit: MAX_NESTING_DEPTH, + }); + } + let mut slots = Vec::with_capacity(fields.len()); + let mut max_field_num = 0u32; + // `slot_proto_numbers[i] = Some(n)` means slot i maps to proto field n; + // `None` means slot i is `Absent` (no proto tag maps here) and is skipped + // by the reverse-index build below. + let mut slot_proto_numbers: Vec> = Vec::with_capacity(fields.len()); + + for arrow_field in fields.iter() { + let Some(proto_field) = descriptor.get_field_by_name(arrow_field.name()) else { + // Inside a Map entry sub-plan, a name that doesn't match the + // proto MapEntry's `key`/`value` is structurally broken: there's + // no proto field to read from, the Map type's non-null key + // contract still applies, and the absent-padding path would + // hand a kind/builder pair to `append_proto3_default` that it + // can't satisfy — which is now a panic (`unreachable!`) rather + // than a Result. Reject up front so the failure shows up at + // sink init with a clear message. + if inside_map_entry { + return Err(WireToArrowError::MapEntryFieldNotInProto { + name: arrow_field.name().to_string(), + }); + } + // Schema drift: the Arrow column exists but the proto doesn't + // carry it. We can only emit all-null for such a column, so + // a non-nullable declaration is a hard mismatch — error + // early before any data flows. + if !arrow_field.is_nullable() { + return Err(WireToArrowError::NonNullableNotGuaranteed { + name: arrow_field.name().to_string(), + reason: "the proto descriptor does not carry this field, \ + so the column would be all-null", + }); + } + // Absent slots get their builders constructed via + // `build_absent_node` -> `TypedBuilder::new` for every + // primitive leaf in the Arrow type. Validate them at + // plan-build so an unsupported leaf type surfaces here + // instead of panicking on the first batch. + validate_arrow_leaf_types(arrow_field.name(), arrow_field.data_type())?; + // Log + metric + keep going — the column becomes + // always-null. Typical cause: a field was removed from the + // proto before the target schema was updated. + tracing::warn!( + message = "proto descriptor is missing a field declared in the Arrow schema; \ + the column will be emitted as all-null", + field = %arrow_field.name(), + descriptor = %descriptor.full_name(), + ); + metrics::counter!( + "wire_to_arrow_missing_proto_field", + "field" => arrow_field.name().to_string(), + "descriptor" => descriptor.full_name().to_string(), + ) + .increment(1); + slots.push(PlanSlot::Absent); + slot_proto_numbers.push(None); + continue; + }; + max_field_num = max_field_num.max(proto_field.number()); + slot_proto_numbers.push(Some(proto_field.number())); + + let is_repeated = proto_field.cardinality() == Cardinality::Repeated; + let kind = proto_field.kind(); + + // Maps take precedence: proto map fields have `is_map() == true` and + // cardinality Repeated, but we dispatch differently from a bare + // repeated-message field. + let slot = if proto_field.is_map() { + let entry_desc = match &kind { + Kind::Message(m) => m, + _ => { + return Err(WireToArrowError::UnsupportedCombination { + name: arrow_field.name().to_string(), + kind: format!("{kind:?}"), + arrow_type: format!("{:?}", arrow_field.data_type()), + repeated: is_repeated, + }); + } + }; + let entry_fields = match arrow_field.data_type() { + DataType::Map(entry_field, _keys_sorted) => match entry_field.data_type() { + DataType::Struct(fs) => fs, + other => { + return Err(WireToArrowError::UnsupportedCombination { + name: arrow_field.name().to_string(), + kind: format!("{kind:?}"), + arrow_type: format!("Map(entry_type = {other:?})"), + repeated: is_repeated, + }); + } + }, + other => { + return Err(WireToArrowError::UnsupportedCombination { + name: arrow_field.name().to_string(), + kind: format!("{kind:?}"), + arrow_type: format!("{other:?}"), + repeated: is_repeated, + }); + } + }; + let sub = MessagePlan::build_at_depth( + entry_desc, + entry_fields, + depth + 1, + /* inside_map_entry */ true, + )?; + PlanSlot::Map(Arc::new(sub)) + } else { + match (&kind, arrow_field.data_type(), is_repeated) { + // Singular proto enum -> STRING column: render the enum + // value *name* rather than its number. Enum + an integer + // column falls through to the generic scalar arm below, + // where `from_proto_kind` maps it to `Int32` as before, so + // existing enum->int tables are unaffected. + (Kind::Enum(enum_desc), DataType::LargeUtf8, false) => { + PlanSlot::EnumString(enum_desc.clone()) + } + // Singular scalar. + (_, dt, false) if !matches!(dt, DataType::Struct(_) | DataType::List(_)) => { + validate_arrow_leaf_types(arrow_field.name(), dt)?; + let sk = ScalarKind::from_proto_kind(&kind).ok_or_else(|| { + WireToArrowError::UnsupportedKind { + name: arrow_field.name().to_string(), + kind: format!("{kind:?}"), + } + })?; + // Reject mismatched (proto scalar, Arrow leaf) pairings + // up front. The runtime appenders rely on this invariant + // to avoid a per-row fallthrough that would otherwise + // fail the whole batch rather than the offending row. + if !sk.matches_arrow_type(dt) { + return Err(WireToArrowError::UnsupportedCombination { + name: arrow_field.name().to_string(), + kind: format!("{kind:?}"), + arrow_type: format!("{dt:?}"), + repeated: false, + }); + } + PlanSlot::Scalar(sk) + } + // Singular nested message. + (Kind::Message(inner_desc), DataType::Struct(inner_fields), false) => { + let sub = MessagePlan::build_at_depth( + inner_desc, + inner_fields, + depth + 1, + /* inside_map_entry */ false, + )?; + PlanSlot::Struct(Arc::new(sub)) + } + // Repeated nested message -> Arrow List. + (Kind::Message(inner_desc), DataType::List(element_field), true) => { + let inner_fields = match element_field.data_type() { + DataType::Struct(fs) => fs, + other => { + return Err(WireToArrowError::RepeatedNonStructList { + name: arrow_field.name().to_string(), + element: format!("{other:?}"), + }); + } + }; + let sub = MessagePlan::build_at_depth( + inner_desc, + inner_fields, + depth + 1, + /* inside_map_entry */ false, + )?; + PlanSlot::RepeatedMessage(Arc::new(sub)) + } + // Repeated proto enum -> Arrow List: render each + // element's value *name*. A repeated enum paired with an + // integer element type falls through to the RepeatedScalar + // arm below and stays numeric (Int32), as before. + (Kind::Enum(enum_desc), DataType::List(item_field), true) + if matches!(item_field.data_type(), DataType::LargeUtf8) => + { + PlanSlot::RepeatedEnumString(enum_desc.clone()) + } + // Repeated scalar -> Arrow List. + (_, DataType::List(item_field), true) => { + validate_arrow_leaf_types(item_field.name(), item_field.data_type())?; + let sk = ScalarKind::from_proto_kind(&kind).ok_or_else(|| { + WireToArrowError::UnsupportedKind { + name: arrow_field.name().to_string(), + kind: format!("{kind:?}"), + } + })?; + if !sk.matches_arrow_type(item_field.data_type()) { + return Err(WireToArrowError::UnsupportedCombination { + name: arrow_field.name().to_string(), + kind: format!("{kind:?}"), + arrow_type: format!("List<{:?}>", item_field.data_type()), + repeated: true, + }); + } + PlanSlot::RepeatedScalar(sk) + } + (k, dt, r) => { + return Err(WireToArrowError::UnsupportedCombination { + name: arrow_field.name().to_string(), + kind: format!("{k:?}"), + arrow_type: format!("{dt:?}"), + repeated: r, + }); + } + } + }; + // Singular slots (Scalar, Struct) emit a null whenever the + // tag is absent from the wire — which proto3 does by default + // for default-valued fields. A non-nullable Arrow declaration + // would trip a generic `RecordBatch::try_new` failure deep in + // encode_batch and drop the whole batch with no row context. + // Reject the mismatch up front. List<…>/Map<…> outer columns + // are exempt: the encoder always emits at least an empty + // list / empty map per row, so the outer column never holds + // a null. Map entry sub-plans are also exempt: Arrow's Map + // type itself dictates the non-null key contract, so the + // check would be a false positive there. + if !arrow_field.is_nullable() && !inside_map_entry { + match slot { + // EnumString is a singular field too: absent -> null (parity + // with `proto_to_value`), so a non-nullable column can't be + // guaranteed and is rejected alongside Scalar/Struct. + PlanSlot::Scalar(_) | PlanSlot::Struct(_) | PlanSlot::EnumString(_) => { + return Err(WireToArrowError::NonNullableNotGuaranteed { + name: arrow_field.name().to_string(), + reason: "proto3 singular fields are omitted at default value, \ + so the column may contain nulls", + }); + } + PlanSlot::RepeatedMessage(_) + | PlanSlot::RepeatedScalar(_) + | PlanSlot::RepeatedEnumString(_) + | PlanSlot::Map(_) + | PlanSlot::Absent => {} + } + } + slots.push(slot); + } + + let mut slot_by_proto_field = vec![SLOT_UNKNOWN; (max_field_num as usize) + 1]; + for (slot_idx, pn) in slot_proto_numbers.iter().enumerate() { + if let Some(pn) = pn { + slot_by_proto_field[*pn as usize] = slot_idx as u32; + } + } + + // Opposite-direction drift: proto fields the Arrow schema doesn't + // carry. These would be silently skipped at scan time (matching + // proto's standard "ignore unknown fields" behavior), but if the + // descriptor reflects the current producer schema, it signals + // "producer emits this field but the target schema hasn't caught up." + // Log + count once at plan build so operators notice. + let arrow_field_names: std::collections::HashSet<&str> = + fields.iter().map(|f| f.name().as_str()).collect(); + for proto_field in descriptor.fields() { + if !arrow_field_names.contains(proto_field.name()) { + tracing::warn!( + message = "proto descriptor has a field not declared in the Arrow schema; \ + occurrences on the wire will be silently skipped", + field = %proto_field.name(), + descriptor = %descriptor.full_name(), + ); + metrics::counter!( + "wire_to_arrow_extra_proto_field", + "field" => proto_field.name().to_string(), + "descriptor" => descriptor.full_name().to_string(), + ) + .increment(1); + } + } + + Ok(MessagePlan { + slots, + slot_by_proto_field, + arrow_fields: fields.clone(), + inside_map_entry, + }) + } + +} + +/// Recursively walk an Arrow `DataType` tree and verify every primitive +/// leaf is supported by [`TypedBuilder`]. Plan-build calls this so an +/// unsupported leaf (e.g. `Date32`) surfaces as a clean +/// [`WireToArrowError::UnsupportedArrowLeafType`] at serializer init, +/// instead of panicking in `TypedBuilder::new` on the first batch. +/// +/// `Struct` / `List` / `Map` are structural; recurse through them. The +/// terminal case is a primitive leaf that either passes +/// [`TypedBuilder::supports`] or fails the check. +pub(super) fn validate_arrow_leaf_types(field_name: &str, dt: &DataType) -> Result<()> { + match dt { + DataType::Struct(inner_fields) => { + for f in inner_fields.iter() { + validate_arrow_leaf_types(f.name(), f.data_type())?; + } + Ok(()) + } + DataType::List(item_field) => { + validate_arrow_leaf_types(item_field.name(), item_field.data_type()) + } + DataType::Map(entry_field, _) => { + if let DataType::Struct(entry_fields) = entry_field.data_type() { + for f in entry_fields.iter() { + validate_arrow_leaf_types(f.name(), f.data_type())?; + } + } + Ok(()) + } + leaf if TypedBuilder::supports(leaf) => Ok(()), + leaf => Err(WireToArrowError::UnsupportedArrowLeafType { + name: field_name.to_string(), + arrow_type: format!("{leaf:?}"), + }), + } +} + +impl MessagePlan { + /// Build a plan whose slots are all `Absent`. Used by the builder layer to + /// shape a null-filled sub-tree when an outer Arrow Struct / List / Map + /// column is itself `Absent` (so every nested child has to null-pad per row). + pub(crate) fn all_absent(fields: &Fields) -> Self { + let slots = (0..fields.len()).map(|_| PlanSlot::Absent).collect(); + MessagePlan { + slots, + slot_by_proto_field: Vec::new(), + arrow_fields: fields.clone(), + inside_map_entry: false, + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::datatypes::{Field, Schema}; + use prost_reflect::DescriptorPool; + use std::path::PathBuf; + + fn load_person_descriptor() -> MessageDescriptor { + let desc_path = PathBuf::from(env!("CARGO_MANIFEST_DIR")) + .join("tests/data/protobuf/protos/test_protobuf.desc"); + let bytes = std::fs::read(&desc_path).expect("read desc"); + DescriptorPool::decode(bytes.as_slice()) + .expect("decode pool") + .get_message_by_name("test_protobuf.Person") + .expect("Person descriptor") + } + + #[test] + fn build_scalar_plan() { + let desc = load_person_descriptor(); + let schema = Schema::new(vec![ + Field::new("name", DataType::LargeUtf8, true), + Field::new("id", DataType::Int32, true), + Field::new("email", DataType::LargeUtf8, true), + ]); + let plan = MessagePlan::build(&desc, &Fields::from(schema.fields().clone())).unwrap(); + assert_eq!(plan.slots.len(), 3); + assert!(matches!(plan.slots[0], PlanSlot::Scalar(ScalarKind::String))); + assert!(matches!(plan.slots[1], PlanSlot::Scalar(ScalarKind::Int32))); + assert!(matches!(plan.slots[2], PlanSlot::Scalar(ScalarKind::String))); + } + + #[test] + fn missing_proto_field_yields_absent_slot() { + // Schema-drift tolerance: if the Arrow schema declares a column the + // proto descriptor doesn't carry, the plan builder logs + increments + // a metric and emits a `PlanSlot::Absent` so the column comes out as + // all-null rather than failing the batch. Typical cause: a field was + // removed from the proto but the target schema still has the column. + let desc = load_person_descriptor(); + let schema = Schema::new(vec![ + Field::new("name", DataType::LargeUtf8, true), + Field::new("deleted_in_proto", DataType::Int32, true), + Field::new("id", DataType::Int32, true), + ]); + let plan = MessagePlan::build(&desc, &Fields::from(schema.fields().clone())).unwrap(); + assert!(matches!(plan.slots[0], PlanSlot::Scalar(ScalarKind::String))); + assert!(matches!(plan.slots[1], PlanSlot::Absent)); + assert!(matches!(plan.slots[2], PlanSlot::Scalar(ScalarKind::Int32))); + // No proto tag for slot 1 — so the reverse index never points at it. + assert!(plan.slot_by_proto_field.iter().all(|entry| *entry != 1)); + } + + #[test] + fn unsupported_combination_flagged() { + // Person.id is a scalar int32. If we claim it's a Struct in Arrow, + // the plan builder should reject. + let desc = load_person_descriptor(); + let schema = Schema::new(vec![Field::new( + "id", + DataType::Struct(Fields::from(vec![Field::new("x", DataType::Int32, true)])), + true, + )]); + let err = MessagePlan::build(&desc, &Fields::from(schema.fields().clone())) + .expect_err("should fail"); + assert!(matches!(err, WireToArrowError::UnsupportedCombination { .. })); + } + + #[test] + fn scalar_kind_arrow_type_mismatch_flagged() { + // Person.name is a proto String; declaring its Arrow column as Int32 + // is a mis-paired schema. Plan-build must catch this so the runtime + // appenders can assume the (ScalarKind, TypedBuilder) pairing is + // well-formed and don't need a fallthrough that fails the whole batch. + let desc = load_person_descriptor(); + let schema = Schema::new(vec![Field::new("name", DataType::Int32, true)]); + let err = MessagePlan::build(&desc, &Fields::from(schema.fields().clone())) + .expect_err("plan-build must reject String/Int32 pairing"); + assert!( + matches!(err, WireToArrowError::UnsupportedCombination { .. }), + "expected UnsupportedCombination, got {err:?}" + ); + } + + #[test] + fn wire_type_for_scalars() { + assert_eq!(ScalarKind::Int32.wire_type(), WT_VARINT); + assert_eq!(ScalarKind::String.wire_type(), WT_LEN); + assert_eq!(ScalarKind::Double.wire_type(), WT_I64); + assert_eq!(ScalarKind::Float.wire_type(), WT_I32); + } +} diff --git a/lib/codecs/src/encoding/format/wire_to_arrow/scan.rs b/lib/codecs/src/encoding/format/wire_to_arrow/scan.rs new file mode 100644 index 0000000000000..ab7d77c2c9435 --- /dev/null +++ b/lib/codecs/src/encoding/format/wire_to_arrow/scan.rs @@ -0,0 +1,260 @@ +//! Per-message wire walks. +//! +//! [`scan_message`] is the hot path — it walks one proto message's wire +//! bytes and appends decoded values into the matching Arrow column +//! builders. [`validate_message`] is the side-effect-free mirror used by +//! `WireToArrowEncoder::encode_batch` to detect malformed rows before +//! any builder is mutated, so a single poison row can be dropped without +//! poisoning the whole batch. + +use super::wire::try_parse_field; + +use super::append::{ + append_enum_string_from_wire, append_repeated_enum_string, append_repeated_scalar, + append_scalar_from_wire, expect_len, validate_enum_string_from_wire, + validate_repeated_enum_string, validate_repeated_scalar, + validate_scalar_from_wire, +}; +use super::builders::{self, BuilderNodeList}; +use super::errors::{Result, WireToArrowError}; +use super::plan::{MessagePlan, PlanSlot, SLOT_UNKNOWN}; + +/// Scan one proto message's wire bytes, appending values into `builders`. +/// +/// Sets `builders.present[i] = true` for each slot `i` touched by any tag in +/// this message. The caller is responsible for resetting `present` (via +/// [`BuilderNodeList::reset_present`]) before invoking, and for calling +/// [`BuilderNodeList::finalize_row`] afterwards. Sub-messages reuse their own +/// level's `present` buffer, so no per-occurrence allocation happens on the +/// hot path. +pub(super) fn scan_message( + plan: &MessagePlan, + mut bytes: &[u8], + builders: &mut BuilderNodeList, +) -> Result<()> { + let dispatch_table = plan.slot_by_proto_field.as_slice(); + while !bytes.is_empty() { + let (field, rest) = try_parse_field(bytes)?; + bytes = rest; + let field_number = field.field_num as usize; + + // Out-of-range or unknown field number — skip; `try_parse_field` + // already consumed the value. + let Some(&slot_idx) = dispatch_table.get(field_number) else { + continue; + }; + if slot_idx == SLOT_UNKNOWN { + continue; + } + let slot_idx = slot_idx as usize; + + let BuilderNodeList { nodes, present } = &mut *builders; + match &mut nodes[slot_idx] { + builders::BuilderNode::Scalar { kind, builder } => { + append_scalar_from_wire(*kind, &field.value, builder)?; + present[slot_idx] = true; + } + builders::BuilderNode::EnumString { desc, builder } => { + append_enum_string_from_wire(desc, &field.value, builder)?; + present[slot_idx] = true; + } + builders::BuilderNode::Struct { + sub_plan, children, .. + } => { + let sub_bytes = expect_len(&field.value)?; + children.reset_present(); + scan_message(sub_plan, sub_bytes, children)?; + children.finalize_row(sub_plan); + present[slot_idx] = true; + } + builders::BuilderNode::RepeatedMessage { + sub_plan, + children, + current_offset, + .. + } + | builders::BuilderNode::Map { + sub_plan, + children, + current_offset, + .. + } => { + let sub_bytes = expect_len(&field.value)?; + children.reset_present(); + scan_message(sub_plan, sub_bytes, children)?; + children.finalize_row(sub_plan); + // Cumulative across the batch — checked_add converts what + // would otherwise be a release-mode wrap + OffsetBuffer + // assertion (process panic at finish) into a clean batch + // failure. Per-row validate doesn't catch this case because + // the overflow can be aggregate across many small rows. + *current_offset = + current_offset + .checked_add(1) + .ok_or(WireToArrowError::OffsetOverflow { + site: "scan_message:repeated_message_or_map", + })?; + present[slot_idx] = true; + } + builders::BuilderNode::RepeatedScalar { + kind, + values, + current_offset, + .. + } => { + append_repeated_scalar(*kind, &field.value, values, current_offset)?; + present[slot_idx] = true; + } + builders::BuilderNode::RepeatedEnumString { + desc, + values, + current_offset, + .. + } => { + append_repeated_enum_string(desc, &field.value, values, current_offset)?; + present[slot_idx] = true; + } + } + } + Ok(()) +} + +/// Walk one proto message's wire bytes without touching any builders, surfacing +/// every decode error that [`scan_message`] would produce for the same input. +/// [`WireToArrowEncoder::encode_batch`] runs this as a pre-pass per message so +/// rows that fail can be dropped from the batch cleanly — no half-appended +/// leaves, no finalized nested sub-rows — and replaced with a `dropped` +/// counter instead of failing the entire batch. +/// +/// The two-pass cost is acceptable because (a) the parse walk is small +/// relative to value appends + buffer growth on the real scan, and (b) Arrow +/// `*Builder` types expose no public rollback API, so an in-place +/// "snapshot + truncate on error" alternative isn't viable. +/// +/// Must stay in lock-step with [`scan_message`]: any wire byte sequence that +/// is accepted here must also be accepted there, and vice versa. If the two +/// diverge (validate accepts but scan errors), the real scan's `?` in +/// `encode_batch` will bubble it out as a batch-level failure — that's a +/// clear signal of a code bug rather than user input. +/// +/// [`WireToArrowEncoder::encode_batch`]: super::encoder::WireToArrowEncoder::encode_batch +pub(super) fn validate_message(plan: &MessagePlan, mut bytes: &[u8]) -> Result<()> { + // Track which singular slots (Scalar / Struct) have been seen in this + // message so a duplicate tag drops the row instead of corrupting + // column-length alignment downstream in scan_message. Proto3 parsers + // are required to accept duplicate singular tags (last-wins for + // scalars, merge for sub-messages), but Arrow `*Builder` types don't + // expose retraction, so implementing last-wins would require either + // per-row scratch buffers or a lookahead pass. Dropping the row + // preserves per-row isolation; the wire_to_arrow_rows_dropped metric + // gives operators a signal if real producers start tripping this. + // + // Stack-allocated for plans with <=128 slots per level (covers every + // realistic Arrow schema we encode); heap-allocated bitvec for wider + // plans. Common case is zero allocations on the hot path. + let mut seen_singular = SeenSingular::with_capacity(plan.slots.len()); + + while !bytes.is_empty() { + let (field, rest) = try_parse_field(bytes)?; + bytes = rest; + let field_number = field.field_num as usize; + + let Some(&slot_idx) = plan.slot_by_proto_field.get(field_number) else { + continue; + }; + if slot_idx == SLOT_UNKNOWN { + continue; + } + let slot_idx = slot_idx as usize; + let slot = &plan.slots[slot_idx]; + + match slot { + PlanSlot::Scalar(sk) => { + if seen_singular.test_and_set(slot_idx) { + return Err(WireToArrowError::DuplicateSingularField { + field_number: field.field_num as u32, + }); + } + validate_scalar_from_wire(*sk, &field.value)?; + } + PlanSlot::Struct(sub_plan) => { + if seen_singular.test_and_set(slot_idx) { + return Err(WireToArrowError::DuplicateSingularField { + field_number: field.field_num as u32, + }); + } + let sub_bytes = expect_len(&field.value)?; + validate_message(sub_plan, sub_bytes)?; + } + PlanSlot::RepeatedMessage(sub_plan) | PlanSlot::Map(sub_plan) => { + let sub_bytes = expect_len(&field.value)?; + validate_message(sub_plan, sub_bytes)?; + } + PlanSlot::RepeatedScalar(sk) => validate_repeated_scalar(*sk, &field.value)?, + PlanSlot::RepeatedEnumString(desc) => { + validate_repeated_enum_string(desc, &field.value)? + } + PlanSlot::EnumString(desc) => { + if seen_singular.test_and_set(slot_idx) { + return Err(WireToArrowError::DuplicateSingularField { + field_number: field.field_num as u32, + }); + } + validate_enum_string_from_wire(desc, &field.value)?; + } + // No proto field number ever points at an Absent slot (Absent + // slots are Arrow columns the proto descriptor lacks), so this + // arm is unreachable in practice. Mirror `scan_message`'s + // fall-through and surface it as a code-bug signal. + PlanSlot::Absent => { + return Err(WireToArrowError::PlanBuilderMismatch { + site: "validate_message:absent_slot_unreachable", + }); + } + } + } + Ok(()) +} + +/// Bitset for tracking which singular slots have already been seen in one +/// `validate_message` call. Inline `u128` covers plans with up to 128 +/// singular Scalar/Struct slots per level — every realistic Arrow schema +/// fits — so the common case is allocation-free on the per-row hot path. +/// Wider plans fall back to a heap `Vec`. +enum SeenSingular { + Small(u128), + Large(Vec), +} + +impl SeenSingular { + #[inline] + fn with_capacity(slot_count: usize) -> Self { + if slot_count <= 128 { + Self::Small(0) + } else { + Self::Large(vec![0u64; slot_count.div_ceil(64)]) + } + } + + /// Set the bit for `idx` and return whether it was already set. + /// Used by `validate_message` to detect duplicate singular tags + /// (`true` on the second occurrence of any Scalar/Struct slot). + #[inline] + fn test_and_set(&mut self, idx: usize) -> bool { + match self { + Self::Small(bits) => { + let mask = 1u128 << idx; + let already = (*bits & mask) != 0; + *bits |= mask; + already + } + Self::Large(words) => { + let word = &mut words[idx / 64]; + let mask = 1u64 << (idx % 64); + let already = (*word & mask) != 0; + *word |= mask; + already + } + } + } +} diff --git a/lib/codecs/src/encoding/format/wire_to_arrow/serializer.rs b/lib/codecs/src/encoding/format/wire_to_arrow/serializer.rs new file mode 100644 index 0000000000000..b2e8ee0a9bcf0 --- /dev/null +++ b/lib/codecs/src/encoding/format/wire_to_arrow/serializer.rs @@ -0,0 +1,134 @@ +//! Public-facing config + serializer types. +//! +//! [`WireToArrowSerializerConfig`] is the `BatchSerializerConfig` variant +//! callers wire into a sink; [`WireToArrowSerializer`] is the runtime +//! object built from a config that turns a batch of events into a +//! `RecordBatch` via an underlying [`WireToArrowEncoder`]. + +use std::path::PathBuf; +use std::sync::Arc; + +use arrow::datatypes::Schema; +use arrow::record_batch::RecordBatch; +use prost_reflect::MessageDescriptor; +use vector_config::configurable_component; +use vector_core::{ + config::DataType, + event::{Event, Value}, + schema, +}; +use vrl::protobuf::descriptor::get_message_descriptor; + +use super::encoder::WireToArrowEncoder; +use super::errors::{Result, WireToArrowError}; + +/// Configuration for the wire-to-Arrow batch serializer. +/// +/// `desc_file` + `message_type` identify the proto descriptor for the +/// *incoming* wire bytes — the user must supply them directly, mirroring +/// [`ProtobufSerializerOptions`]. The sink injects the output Arrow `schema` +/// at build time (typically derived from its own schema source). +/// +/// The descriptor must describe the exact bytes present in `event.message`; +/// decoding uses the descriptor's field numbers as-is. If the payload needs +/// any pre-processing before it matches the descriptor, do it upstream. +/// +/// [`ProtobufSerializerOptions`]: crate::encoding::format::ProtobufSerializerOptions +#[configurable_component] +#[derive(Clone, Default)] +pub struct WireToArrowSerializerConfig { + /// Path to the protobuf descriptor set file describing the incoming wire bytes. + /// + /// Must correspond to the exact proto type serialized in `event.message`. + /// Typically the output of `protoc -I -o `. + #[configurable(metadata(docs::examples = "/etc/vector/protobuf_descriptor_set.desc"))] + pub desc_file: PathBuf, + + /// The fully-qualified message type within the descriptor file. Must name + /// the type of the bytes in `event.message`. + #[configurable(metadata(docs::examples = "package.Message"))] + pub message_type: String, + + /// The Arrow schema of the output `RecordBatch`. Injected by the sink. + #[serde(skip)] + #[configurable(derived)] + pub schema: Option, +} + +impl std::fmt::Debug for WireToArrowSerializerConfig { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("WireToArrowSerializerConfig") + .field("desc_file", &self.desc_file) + .field("message_type", &self.message_type) + .field( + "schema", + &self + .schema + .as_ref() + .map(|s| format!("{} fields", s.fields().len())), + ) + .finish() + } +} + +impl WireToArrowSerializerConfig { + /// The data type of events accepted by this serializer. + pub fn input_type(&self) -> DataType { + DataType::Log + } + + /// The schema required by the serializer. + pub fn schema_requirement(&self) -> schema::Requirement { + schema::Requirement::empty() + } +} + +/// Batch serializer that decodes proto wire bytes directly into an Arrow +/// `RecordBatch`, bypassing the generic `ProtobufDeserializer` chain. +#[derive(Clone, Debug)] +pub struct WireToArrowSerializer { + encoder: Arc, +} + +impl WireToArrowSerializer { + /// Build a serializer from the given configuration. Loads the proto + /// descriptor from `desc_file` + `message_type`; the output Arrow schema + /// must have been injected (via `config.schema`) by the sink. + pub fn new(config: WireToArrowSerializerConfig) -> Result { + let descriptor = get_message_descriptor(&config.desc_file, &config.message_type) + .map_err(|message| WireToArrowError::DescriptorLoad { message })?; + let schema = config + .schema + .ok_or_else(|| WireToArrowError::ConfigurationMissing { field: "schema" })?; + Self::from_descriptor(descriptor, schema) + } + + /// Build a serializer from an already-resolved descriptor and schema. + /// Mostly useful for tests and for callers that have the descriptor in + /// memory already. + pub fn from_descriptor(descriptor: MessageDescriptor, schema: Schema) -> Result { + let encoder = WireToArrowEncoder::new(&descriptor, schema)?; + Ok(Self { + encoder: Arc::new(encoder), + }) + } + + /// Encode a batch of events into a single Arrow `RecordBatch`. + /// + /// Every event must carry a `Value::Bytes`-typed `message` field holding + /// the original proto wire bytes; any miss rejects the batch. + pub fn encode_to_record_batch(&self, events: &[Event]) -> Result { + if events.is_empty() { + return Err(WireToArrowError::NoEvents); + } + let mut wire_bytes = Vec::with_capacity(events.len()); + for event in events { + match event.as_log().get_message() { + Some(Value::Bytes(b)) => wire_bytes.push(b.clone()), + Some(_) => return Err(WireToArrowError::MessageBytesWrongType), + None => return Err(WireToArrowError::MessageBytesMissing), + } + } + self.encoder.encode_batch(&wire_bytes) + } +} diff --git a/lib/codecs/src/encoding/format/wire_to_arrow/tests.rs b/lib/codecs/src/encoding/format/wire_to_arrow/tests.rs new file mode 100644 index 0000000000000..f4a47c350afd2 --- /dev/null +++ b/lib/codecs/src/encoding/format/wire_to_arrow/tests.rs @@ -0,0 +1,1738 @@ +use super::append::read_packed_element; +use super::plan::ScalarKind; +use super::{ + WireToArrowEncoder, WireToArrowError, WireToArrowSerializer, WireToArrowSerializerConfig, +}; + +use proptest::prelude::*; + +use arrow::array::{Array, AsArray}; +use arrow::datatypes::{DataType, Field, Fields as ArrowFields, Schema}; +use bytes::Bytes; +use prost_reflect::MessageDescriptor; +use prost_reflect::prost::Message as _; +use prost_reflect::prost_types::field_descriptor_proto::{Label, Type as ProtoType}; +use prost_reflect::prost_types::{ + DescriptorProto, FieldDescriptorProto, FileDescriptorProto, FileDescriptorSet, + OneofDescriptorProto, +}; +use prost_reflect::{DescriptorPool, DynamicMessage, Value as ProtoValue}; +use std::path::PathBuf; +use std::sync::Arc; +use vector_core::event::Event; +use vrl::event_path; +use super::wire::WireValue; + +fn encode_varint_for_test(mut value: u64) -> Vec { + let mut out = Vec::new(); + while value >= 0x80 { + out.push((value as u8) | 0x80); + value >>= 7; + } + out.push(value as u8); + out +} + +#[test] +fn packed_element_varint_roundtrip() { + for v in [0u64, 1, 127, 128, 255, 16384, u32::MAX as u64, u64::MAX] { + let encoded = encode_varint_for_test(v); + let (decoded, rest) = read_packed_element(ScalarKind::Int64, &encoded).unwrap(); + assert!(matches!(decoded, WireValue::Varint(d) if d == v), "mismatch on {v}"); + assert!(rest.is_empty(), "buffer not fully consumed for {v}"); + } +} + +#[test] +fn packed_element_varint_eof_maps_to_unexpected_eof() { + assert!(matches!( + read_packed_element(ScalarKind::Int64, &[0x80u8]), + Err(WireToArrowError::UnexpectedEof) + )); +} + +#[test] +fn packed_element_varint_overflow_maps_to_varint_overflow() { + assert!(matches!( + read_packed_element(ScalarKind::Int64, &[0xffu8; 11]), + Err(WireToArrowError::VarintOverflow) + )); +} + +#[test] +fn packed_element_fixed32_roundtrip() { + let bytes = 0x12345678u32.to_le_bytes(); + let (decoded, rest) = read_packed_element(ScalarKind::Fixed32, &bytes).unwrap(); + assert!(matches!(decoded, WireValue::I32(v) if v == 0x12345678)); + assert!(rest.is_empty()); +} + +#[test] +fn packed_element_fixed64_roundtrip() { + let v: u64 = 0x0011_2233_4455_6677; + let bytes = v.to_le_bytes(); + let (decoded, rest) = read_packed_element(ScalarKind::Fixed64, &bytes).unwrap(); + assert!(matches!(decoded, WireValue::I64(d) if d == v)); + assert!(rest.is_empty()); +} + +fn descriptor_pool(file: &str) -> DescriptorPool { + let desc_path = PathBuf::from(env!("CARGO_MANIFEST_DIR")) + .join("tests/data/protobuf/protos") + .join(file); + let bytes = std::fs::read(&desc_path).unwrap(); + DescriptorPool::decode(bytes.as_slice()).unwrap() +} + +fn scalar_descriptor() -> MessageDescriptor { + descriptor_pool("test_protobuf.desc") + .get_message_by_name("test_protobuf.Person") + .unwrap() +} + +fn rich_descriptor() -> MessageDescriptor { + descriptor_pool("test_protobuf3.desc") + .get_message_by_name("test_protobuf3.Person") + .unwrap() +} + +/// Build an ad-hoc `message Bag { repeated int32 numbers = 1; }` descriptor +/// programmatically, since none of the checked-in test protos have a bare +/// repeated scalar field. +fn repeated_int32_descriptor() -> MessageDescriptor { + let fd = FileDescriptorProto { + name: Some("wire_to_arrow_test.proto".into()), + package: Some("wire_to_arrow_test".into()), + message_type: vec![DescriptorProto { + name: Some("Bag".into()), + field: vec![FieldDescriptorProto { + name: Some("numbers".into()), + number: Some(1), + label: Some(Label::Repeated as i32), + r#type: Some(ProtoType::Int32 as i32), + ..Default::default() + }], + ..Default::default() + }], + ..Default::default() + }; + let set = FileDescriptorSet { file: vec![fd] }; + let mut bytes = Vec::new(); + set.encode(&mut bytes).unwrap(); + DescriptorPool::decode(bytes.as_slice()) + .unwrap() + .get_message_by_name("wire_to_arrow_test.Bag") + .unwrap() +} + +#[test] +fn scalar_roundtrip() { + let desc = scalar_descriptor(); + let schema = Schema::new(vec![ + Field::new("name", DataType::LargeUtf8, true), + Field::new("id", DataType::Int32, true), + Field::new("email", DataType::LargeUtf8, true), + ]); + let enc = WireToArrowEncoder::new(&desc, schema).unwrap(); + + let mut msg = DynamicMessage::new(desc.clone()); + msg.set_field_by_name("name", ProtoValue::String("Alice".into())); + msg.set_field_by_name("id", ProtoValue::I32(42)); + msg.set_field_by_name("email", ProtoValue::String("alice@x.com".into())); + let mut buf = Vec::new(); + msg.encode(&mut buf).unwrap(); + + let batch = enc.encode_batch(&[Bytes::from(buf)]).unwrap(); + assert_eq!(batch.num_rows(), 1); + assert_eq!(batch.num_columns(), 3); + assert_eq!(batch.column(0).as_string::().value(0), "Alice"); + assert_eq!( + batch + .column(1) + .as_primitive::() + .value(0), + 42 + ); +} + +#[test] +fn absent_fields_are_null() { + let desc = scalar_descriptor(); + let schema = Schema::new(vec![ + Field::new("name", DataType::LargeUtf8, true), + Field::new("id", DataType::Int32, true), + Field::new("email", DataType::LargeUtf8, true), + ]); + let enc = WireToArrowEncoder::new(&desc, schema).unwrap(); + + // Populate only `name`; `id` and `email` should show up as null. + let mut msg = DynamicMessage::new(desc.clone()); + msg.set_field_by_name("name", ProtoValue::String("only name".into())); + let mut buf = Vec::new(); + msg.encode(&mut buf).unwrap(); + + let batch = enc.encode_batch(&[Bytes::from(buf)]).unwrap(); + assert_eq!(batch.num_rows(), 1); + assert!(!batch.column(0).is_null(0)); + assert!(batch.column(1).is_null(0), "id should be null"); + assert!(batch.column(2).is_null(0), "email should be null"); +} + +#[test] +fn unknown_wire_fields_are_skipped() { + // Person.phones (field 4 in proto2 test_protobuf.Person) is not in + // our Arrow schema but will appear in wire bytes when populated. The + // encoder should skip it cleanly. + let desc = scalar_descriptor(); + let schema = Schema::new(vec![ + Field::new("name", DataType::LargeUtf8, true), + Field::new("id", DataType::Int32, true), + Field::new("email", DataType::LargeUtf8, true), + ]); + let enc = WireToArrowEncoder::new(&desc, schema).unwrap(); + + let phone_desc = desc + .get_field_by_name("phones") + .unwrap() + .kind() + .as_message() + .unwrap() + .clone(); + let mut phone = DynamicMessage::new(phone_desc); + phone.set_field_by_name("number", ProtoValue::String("555-0000".into())); + + let mut msg = DynamicMessage::new(desc.clone()); + msg.set_field_by_name("name", ProtoValue::String("Alice".into())); + msg.set_field_by_name("id", ProtoValue::I32(1)); + msg.set_field_by_name("phones", ProtoValue::List(vec![ProtoValue::Message(phone)])); + let mut buf = Vec::new(); + msg.encode(&mut buf).unwrap(); + + // Should not error — unknown fields get skipped. + let batch = enc.encode_batch(&[Bytes::from(buf)]).unwrap(); + assert_eq!(batch.num_rows(), 1); + assert_eq!(batch.column(0).as_string::().value(0), "Alice"); +} + +#[test] +fn absent_proto_field_becomes_all_null_column() { + // Schema-drift path: the Arrow schema has a column `missing_col` that + // the proto descriptor doesn't carry (simulates "the target schema has + // the column but the producer's proto dropped it"). The encoder must emit + // an all-null Arrow column for `missing_col` and otherwise populate normally. + let desc = scalar_descriptor(); + let schema = Schema::new(vec![ + Field::new("name", DataType::LargeUtf8, true), + Field::new("missing_col", DataType::Int64, true), + Field::new("id", DataType::Int32, true), + ]); + let enc = WireToArrowEncoder::new(&desc, schema).unwrap(); + + let mut msg1 = DynamicMessage::new(desc.clone()); + msg1.set_field_by_name("name", ProtoValue::String("Alice".into())); + msg1.set_field_by_name("id", ProtoValue::I32(1)); + let mut buf1 = Vec::new(); + msg1.encode(&mut buf1).unwrap(); + + let mut msg2 = DynamicMessage::new(desc.clone()); + msg2.set_field_by_name("name", ProtoValue::String("Bob".into())); + msg2.set_field_by_name("id", ProtoValue::I32(2)); + let mut buf2 = Vec::new(); + msg2.encode(&mut buf2).unwrap(); + + let batch = enc + .encode_batch(&[Bytes::from(buf1), Bytes::from(buf2)]) + .unwrap(); + assert_eq!(batch.num_rows(), 2); + assert_eq!(batch.num_columns(), 3); + + assert_eq!(batch.column(0).as_string::().value(0), "Alice"); + assert_eq!(batch.column(0).as_string::().value(1), "Bob"); + + // The missing column must be declared null for every row. + let missing = batch.column(1); + assert_eq!(missing.len(), 2); + assert!(missing.is_null(0)); + assert!(missing.is_null(1)); + + let ids = batch.column(2).as_primitive::(); + assert_eq!(ids.value(0), 1); + assert_eq!(ids.value(1), 2); +} + +fn serializer_for(desc: &MessageDescriptor, schema: Schema) -> WireToArrowSerializer { + WireToArrowSerializer::from_descriptor(desc.clone(), schema).expect("serializer build") +} + +fn event_with_message_bytes(bytes: Bytes) -> Event { + let mut e = Event::from(vector_core::event::LogEvent::default()); + e.as_mut_log().insert(event_path!("message"), bytes); + e +} + +#[test] +fn serializer_loads_descriptor_from_config() { + // Happy path via the user-facing constructor: `desc_file` + `message_type` + // point at a real descriptor set, `schema` is injected by the sink. + let schema = Schema::new(vec![Field::new("id", DataType::Int32, true)]); + let config = WireToArrowSerializerConfig { + desc_file: PathBuf::from(env!("CARGO_MANIFEST_DIR")) + .join("tests/data/protobuf/protos/test_protobuf.desc"), + message_type: "test_protobuf.Person".to_string(), + schema: Some(schema), + }; + let serializer = WireToArrowSerializer::new(config).expect("build"); + assert!(matches!( + serializer.encode_to_record_batch(&[]), + Err(WireToArrowError::NoEvents) + )); +} + +#[test] +fn serializer_errors_on_bad_descriptor_path() { + let schema = Schema::new(vec![Field::new("id", DataType::Int32, true)]); + let config = WireToArrowSerializerConfig { + desc_file: PathBuf::from("/nonexistent/path/to/schema.desc"), + message_type: "some.Message".to_string(), + schema: Some(schema), + }; + assert!(matches!( + WireToArrowSerializer::new(config), + Err(WireToArrowError::DescriptorLoad { .. }) + )); +} + +#[test] +fn serializer_errors_on_missing_schema() { + let config = WireToArrowSerializerConfig { + desc_file: PathBuf::from(env!("CARGO_MANIFEST_DIR")) + .join("tests/data/protobuf/protos/test_protobuf.desc"), + message_type: "test_protobuf.Person".to_string(), + schema: None, + }; + assert!(matches!( + WireToArrowSerializer::new(config), + Err(WireToArrowError::ConfigurationMissing { field: "schema" }) + )); +} + +#[test] +fn serializer_empty_batch_errors() { + let desc = scalar_descriptor(); + let schema = Schema::new(vec![Field::new("id", DataType::Int32, true)]); + let serializer = serializer_for(&desc, schema); + assert!(matches!( + serializer.encode_to_record_batch(&[]), + Err(WireToArrowError::NoEvents) + )); +} + +#[test] +fn serializer_missing_message_field_errors() { + let desc = scalar_descriptor(); + let schema = Schema::new(vec![Field::new("id", DataType::Int32, true)]); + let serializer = serializer_for(&desc, schema); + + let e1 = event_with_message_bytes(Bytes::from_static(b"")); + let e2 = Event::from(vector_core::event::LogEvent::default()); // no message + + assert!(matches!( + serializer.encode_to_record_batch(&[e1, e2]), + Err(WireToArrowError::MessageBytesMissing) + )); +} + +#[test] +fn serializer_wrong_type_message_errors() { + let desc = scalar_descriptor(); + let schema = Schema::new(vec![Field::new("id", DataType::Int32, true)]); + let serializer = serializer_for(&desc, schema); + + // Plain strings are represented as `Value::Bytes`, so use an integer + // to get a non-bytes variant for this negative case. + let mut e = Event::from(vector_core::event::LogEvent::default()); + e.as_mut_log().insert(event_path!("message"), 42_i64); + assert!(matches!( + serializer.encode_to_record_batch(&[e]), + Err(WireToArrowError::MessageBytesWrongType) + )); +} + +#[test] +fn serializer_end_to_end_matches_direct_encode() { + // Build events with wire bytes on `message`, encode via the + // serializer, and compare against calling the lower-level encoder + // directly. + let desc = scalar_descriptor(); + let schema = Schema::new(vec![ + Field::new("name", DataType::LargeUtf8, true), + Field::new("id", DataType::Int32, true), + Field::new("email", DataType::LargeUtf8, true), + ]); + let serializer = serializer_for(&desc, schema.clone()); + let enc = WireToArrowEncoder::new(&desc, schema).unwrap(); + + let mut wire_bytes_list = Vec::new(); + let mut events = Vec::new(); + for i in 0..5_i32 { + let mut msg = DynamicMessage::new(desc.clone()); + msg.set_field_by_name("name", ProtoValue::String(format!("n-{i}"))); + msg.set_field_by_name("id", ProtoValue::I32(i)); + let mut buf = Vec::new(); + msg.encode(&mut buf).unwrap(); + let bytes = Bytes::from(buf); + wire_bytes_list.push(bytes.clone()); + events.push(event_with_message_bytes(bytes)); + } + + let via_events = serializer.encode_to_record_batch(&events).unwrap(); + let via_bytes = enc.encode_batch(&wire_bytes_list).unwrap(); + assert_eq!(via_events.num_rows(), via_bytes.num_rows()); + assert_eq!(via_events.num_columns(), via_bytes.num_columns()); + for i in 0..via_events.num_columns() { + assert_eq!(via_events.column(i).as_ref(), via_bytes.column(i).as_ref()); + } +} + +#[test] +fn repeated_scalar_unpacked_roundtrip() { + // Unpacked: emit each element with its own tag. For proto3, this is + // the default for non-packed repeated scalars when the writer chooses + // not to pack (which can happen for proto2 as well). + let desc = repeated_int32_descriptor(); + let schema = Schema::new(vec![Field::new( + "numbers", + DataType::List(Arc::new(Field::new("item", DataType::Int32, true))), + true, + )]); + let enc = WireToArrowEncoder::new(&desc, schema).unwrap(); + + // Hand-craft wire bytes for two rows: + // row 0: numbers = [1, 2, 3] (unpacked: 3 tag+value pairs) + // row 1: numbers = [42] (single tag+value) + let tag = (1u8 << 3) | 0; // field 1, wire type 0 (varint) + let mut row0 = Vec::new(); + for v in [1i32, 2, 3] { + row0.push(tag); + encode_varint_into(&mut row0, v as u64); + } + let mut row1 = Vec::new(); + row1.push(tag); + encode_varint_into(&mut row1, 42); + + let batch = enc + .encode_batch(&[Bytes::from(row0), Bytes::from(row1)]) + .unwrap(); + assert_eq!(batch.num_rows(), 2); + let list = batch.column(0).as_list::(); + assert_eq!(list.value_length(0), 3); + assert_eq!(list.value_length(1), 1); + let values = list.values().as_primitive::(); + assert_eq!(values.len(), 4); + assert_eq!(values.value(0), 1); + assert_eq!(values.value(1), 2); + assert_eq!(values.value(2), 3); + assert_eq!(values.value(3), 42); +} + +#[test] +fn repeated_scalar_packed_roundtrip() { + // Packed: one length-delimited blob with concatenated varints. proto3 + // repeated scalars default to this encoding. + let desc = repeated_int32_descriptor(); + let schema = Schema::new(vec![Field::new( + "numbers", + DataType::List(Arc::new(Field::new("item", DataType::Int32, true))), + true, + )]); + let enc = WireToArrowEncoder::new(&desc, schema).unwrap(); + + // Tag for field 1 with wire_type 2 (length-delimited). + let tag = (1u8 << 3) | 2; + let mut payload = Vec::new(); + for v in [10i32, 20, 30, 40] { + encode_varint_into(&mut payload, v as u64); + } + let mut row = Vec::new(); + row.push(tag); + encode_varint_into(&mut row, payload.len() as u64); + row.extend_from_slice(&payload); + + let batch = enc.encode_batch(&[Bytes::from(row)]).unwrap(); + assert_eq!(batch.num_rows(), 1); + let list = batch.column(0).as_list::(); + assert_eq!(list.value_length(0), 4); + let values = list.values().as_primitive::(); + assert_eq!( + (0..4).map(|i| values.value(i)).collect::>(), + vec![10, 20, 30, 40] + ); +} + +#[test] +fn repeated_scalar_empty_row_produces_empty_list() { + // A row with no tag occurrences produces an empty list, not null. + let desc = repeated_int32_descriptor(); + let schema = Schema::new(vec![Field::new( + "numbers", + DataType::List(Arc::new(Field::new("item", DataType::Int32, true))), + true, + )]); + let enc = WireToArrowEncoder::new(&desc, schema).unwrap(); + + let batch = enc.encode_batch(&[Bytes::new()]).unwrap(); + assert_eq!(batch.num_rows(), 1); + let list = batch.column(0).as_list::(); + assert_eq!(list.value_length(0), 0); + assert!(!list.is_null(0), "list column itself should never be null"); +} + +#[test] +fn map_roundtrip() { + // Proto: test_protobuf3.Person.data = map + // where PhoneType is an enum (int32-encoded on the wire). + // + // Arrow side: Map with entry + // field named "key_value" per the sink's existing convention. + let desc = rich_descriptor(); + let entry_fields = ArrowFields::from(vec![ + Field::new("key", DataType::LargeUtf8, false), + Field::new("value", DataType::Int32, true), + ]); + let entry_field = Arc::new(Field::new( + "key_value", + DataType::Struct(entry_fields), + false, + )); + let schema = Schema::new(vec![Field::new( + "data", + DataType::Map(entry_field, false), + true, + )]); + let enc = WireToArrowEncoder::new(&desc, schema).unwrap(); + + // Populate a row with 2 map entries. + let mut msg = DynamicMessage::new(desc.clone()); + let mut entries: std::collections::HashMap = + std::collections::HashMap::new(); + entries.insert( + prost_reflect::MapKey::String("alpha".into()), + ProtoValue::EnumNumber(1), + ); + entries.insert( + prost_reflect::MapKey::String("beta".into()), + ProtoValue::EnumNumber(2), + ); + msg.set_field_by_name("data", ProtoValue::Map(entries)); + let mut buf = Vec::new(); + msg.encode(&mut buf).unwrap(); + + let batch = enc.encode_batch(&[Bytes::from(buf)]).unwrap(); + assert_eq!(batch.num_rows(), 1); + let map = batch + .column(0) + .as_any() + .downcast_ref::() + .expect("column should be MapArray"); + assert_eq!(map.value_length(0), 2, "expected 2 map entries"); + let keys = map.keys().as_string::(); + let values = map.values().as_primitive::(); + // Map entry iteration order is not guaranteed — collect then compare. + let pairs: std::collections::HashMap = (0..2) + .map(|i| (keys.value(i).to_string(), values.value(i))) + .collect(); + assert_eq!(pairs.get("alpha").copied(), Some(1)); + assert_eq!(pairs.get("beta").copied(), Some(2)); +} + +#[test] +fn map_entry_field_name_mismatch_rejected_at_plan_build() { + // Arrow's Map type doesn't enforce that the inner Struct's fields are + // named `key` and `value`, but proto MapEntry does. Declaring a Map + // whose entry struct has a different name ("k" here) leaves no proto + // field for the encoder to route bytes from and would force the + // absent-padding path to hand `append_proto3_default` a kind/builder + // pair it can't satisfy. Plan-build must reject this up front so the + // failure surfaces at sink init rather than as a runtime panic. + let desc = rich_descriptor(); + let entry_fields = ArrowFields::from(vec![ + Field::new("k", DataType::LargeUtf8, false), + Field::new("value", DataType::Int32, true), + ]); + let entry_field = Arc::new(Field::new( + "key_value", + DataType::Struct(entry_fields), + false, + )); + let schema = Schema::new(vec![Field::new( + "data", + DataType::Map(entry_field, false), + true, + )]); + let err = WireToArrowEncoder::new(&desc, schema) + .expect_err("plan-build must reject Map entry field 'k' (not in proto MapEntry)"); + assert!( + matches!( + err, + WireToArrowError::MapEntryFieldNotInProto { ref name } if name == "k" + ), + "expected MapEntryFieldNotInProto, got {err:?}" + ); +} + +/// Build a `data` field carrying one map entry, with raw bytes for the +/// MapEntry message (so we can elide the key tag, the value tag, or both — +/// proto3 default elision applies inside MapEntry messages just like every +/// other singular field). Schema: `test_protobuf3.Person.data` is field 4, +/// `map`. Wire format: outer tag `(4 << 3) | 2 = 0x22`, +/// LEN-prefixed entry body. +fn person_with_raw_map_entry(entry_body: &[u8]) -> Vec { + let mut out = Vec::with_capacity(entry_body.len() + 2); + out.push(0x22); // tag 4, LEN + out.push(entry_body.len() as u8); // body length (small for tests) + out.extend_from_slice(entry_body); + out +} + +#[test] +fn map_entry_with_empty_string_key_encodes_as_empty_string_not_null() { + // Proto3 elides default-valued singular fields *inside* MapEntry too. + // Arrow's Map type declares the key field non-nullable; before the fix + // an absent key tag produced a null, which fails `StructArray::try_new` + // at batch finish — taking the whole batch down even though the wire + // bytes are perfectly valid proto3. Producers in prost / Python / + // protoc-gen-cpp emit exactly this shape when given a map with key "". + let desc = rich_descriptor(); + let entry_fields = ArrowFields::from(vec![ + Field::new("key", DataType::LargeUtf8, false), + Field::new("value", DataType::Int32, true), + ]); + let entry_field = Arc::new(Field::new( + "key_value", + DataType::Struct(entry_fields), + false, + )); + let schema = Schema::new(vec![Field::new( + "data", + DataType::Map(entry_field, false), + true, + )]); + let enc = WireToArrowEncoder::new(&desc, schema).unwrap(); + + // MapEntry with only the value tag (key omitted, taking proto3 default ""). + // 0x10 = (2 << 3) | 0 = value tag, varint + // 0x01 = enum value 1 (HOME) + let bytes = person_with_raw_map_entry(&[0x10, 0x01]); + + let batch = enc + .encode_batch(&[Bytes::from(bytes)]) + .expect("default-keyed map entry must not fail the batch"); + assert_eq!(batch.num_rows(), 1); + let map = batch + .column(0) + .as_any() + .downcast_ref::() + .expect("MapArray"); + assert_eq!(map.value_length(0), 1); + let keys = map.keys().as_string::(); + assert_eq!(keys.value(0), "", "empty-default key must materialize as \"\""); + assert!(!keys.is_null(0), "key column must contain no nulls"); + let values = map.values().as_primitive::(); + assert_eq!(values.value(0), 1); +} + +#[test] +fn map_entry_with_default_int_value_encodes_as_zero_not_null() { + // Mirror of the key case for the value side: proto3 elides value=0 + // (default int) inside MapEntry. Even with a nullable Arrow value field, + // the proto semantics say "absent == 0", not "absent == null" — and + // when the Arrow value is *non*-nullable, a null here would crash the + // batch. Set value non-nullable to exercise both behaviors at once. + let desc = rich_descriptor(); + let entry_fields = ArrowFields::from(vec![ + Field::new("key", DataType::LargeUtf8, false), + Field::new("value", DataType::Int32, false), + ]); + let entry_field = Arc::new(Field::new( + "key_value", + DataType::Struct(entry_fields), + false, + )); + let schema = Schema::new(vec![Field::new( + "data", + DataType::Map(entry_field, false), + true, + )]); + let enc = WireToArrowEncoder::new(&desc, schema).unwrap(); + + // MapEntry with only the key tag (value omitted, taking proto3 default 0). + // 0x0a = (1 << 3) | 2 = key tag, LEN + // 0x03 = length + // "foo" + let bytes = person_with_raw_map_entry(&[0x0a, 0x03, b'f', b'o', b'o']); + + let batch = enc + .encode_batch(&[Bytes::from(bytes)]) + .expect("default-valued map entry must not fail the batch"); + assert_eq!(batch.num_rows(), 1); + let map = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let keys = map.keys().as_string::(); + let values = map.values().as_primitive::(); + assert_eq!(keys.value(0), "foo"); + assert_eq!(values.value(0), 0, "default-int value must materialize as 0"); + assert!(!values.is_null(0)); +} + +#[test] +fn map_entry_with_all_defaults_encodes_as_empty_default_pair() { + // An entirely-empty MapEntry on the wire (`0x22 0x00`) is what prost + // emits when you serialize `HashMap::from([("".to_string(), 0)])`. + // Both key and value tags are elided. Must produce a valid Arrow row + // with `("", 0)`. + let desc = rich_descriptor(); + let entry_fields = ArrowFields::from(vec![ + Field::new("key", DataType::LargeUtf8, false), + Field::new("value", DataType::Int32, false), + ]); + let entry_field = Arc::new(Field::new( + "key_value", + DataType::Struct(entry_fields), + false, + )); + let schema = Schema::new(vec![Field::new( + "data", + DataType::Map(entry_field, false), + true, + )]); + let enc = WireToArrowEncoder::new(&desc, schema).unwrap(); + + let bytes = person_with_raw_map_entry(&[]); // empty MapEntry + + let batch = enc + .encode_batch(&[Bytes::from(bytes)]) + .expect("empty MapEntry must not fail the batch"); + assert_eq!(batch.num_rows(), 1); + let map = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let keys = map.keys().as_string::(); + let values = map.values().as_primitive::(); + assert_eq!(keys.value(0), ""); + assert_eq!(values.value(0), 0); +} + +fn encode_varint_into(buf: &mut Vec, mut value: u64) { + while value >= 0x80 { + buf.push((value as u8) | 0x80); + value >>= 7; + } + buf.push(value as u8); +} + +/// `message Ts { int64 event_time = 1; }` — for the timestamp coercion test. +fn timestamp_descriptor() -> MessageDescriptor { + let fd = FileDescriptorProto { + name: Some("wire_to_arrow_test_ts.proto".into()), + package: Some("wire_to_arrow_test".into()), + message_type: vec![DescriptorProto { + name: Some("Ts".into()), + field: vec![FieldDescriptorProto { + name: Some("event_time".into()), + number: Some(1), + label: Some(Label::Optional as i32), + r#type: Some(ProtoType::Int64 as i32), + ..Default::default() + }], + ..Default::default() + }], + ..Default::default() + }; + let set = FileDescriptorSet { file: vec![fd] }; + let mut bytes = Vec::new(); + set.encode(&mut bytes).unwrap(); + DescriptorPool::decode(bytes.as_slice()) + .unwrap() + .get_message_by_name("wire_to_arrow_test.Ts") + .unwrap() +} + +/// `message Choice { oneof x { int32 a = 1; string b = 2; } }` +fn oneof_descriptor() -> MessageDescriptor { + let fd = FileDescriptorProto { + name: Some("wire_to_arrow_test_oneof.proto".into()), + package: Some("wire_to_arrow_test".into()), + message_type: vec![DescriptorProto { + name: Some("Choice".into()), + field: vec![ + FieldDescriptorProto { + name: Some("a".into()), + number: Some(1), + label: Some(Label::Optional as i32), + r#type: Some(ProtoType::Int32 as i32), + oneof_index: Some(0), + ..Default::default() + }, + FieldDescriptorProto { + name: Some("b".into()), + number: Some(2), + label: Some(Label::Optional as i32), + r#type: Some(ProtoType::String as i32), + oneof_index: Some(0), + ..Default::default() + }, + ], + oneof_decl: vec![OneofDescriptorProto { + name: Some("x".into()), + ..Default::default() + }], + ..Default::default() + }], + ..Default::default() + }; + let set = FileDescriptorSet { file: vec![fd] }; + let mut bytes = Vec::new(); + set.encode(&mut bytes).unwrap(); + DescriptorPool::decode(bytes.as_slice()) + .unwrap() + .get_message_by_name("wire_to_arrow_test.Choice") + .unwrap() +} + +#[test] +fn int64_to_timestamp_micros_coercion() { + // proto int64 field with the Arrow column declared as Timestamp(Micro, UTC). + let desc = timestamp_descriptor(); + let schema = Schema::new(vec![Field::new( + "event_time", + DataType::Timestamp(arrow::datatypes::TimeUnit::Microsecond, Some("UTC".into())), + true, + )]); + let enc = WireToArrowEncoder::new(&desc, schema).unwrap(); + + // Two rows, plus one with the field absent (should produce null). + let mut row0 = Vec::new(); + row0.push((1u8 << 3) | 0); // tag 1, varint + encode_varint_into(&mut row0, 1_700_000_000_000_000_u64); + let mut row1 = Vec::new(); + row1.push((1u8 << 3) | 0); + encode_varint_into(&mut row1, 1_800_000_000_000_000_u64); + + let batch = enc + .encode_batch(&[ + Bytes::from(row0), + Bytes::from(row1), + Bytes::new(), // absent -> null + ]) + .unwrap(); + + assert_eq!(batch.num_rows(), 3); + let col = batch + .column(0) + .as_any() + .downcast_ref::() + .expect("TimestampMicrosecondArray"); + assert_eq!(col.value(0), 1_700_000_000_000_000); + assert_eq!(col.value(1), 1_800_000_000_000_000); + assert!(col.is_null(2)); +} + +#[test] +fn oneof_variants_map_to_separate_columns() { + // With the wire-format identity (oneof variants look like regular + // singular fields), the encoder should populate whichever variant + // appears in the bytes and leave the others null. + let desc = oneof_descriptor(); + let schema = Schema::new(vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::LargeUtf8, true), + ]); + let enc = WireToArrowEncoder::new(&desc, schema).unwrap(); + + // Row 0: only `a = 42`. + let mut row0 = Vec::new(); + row0.push((1u8 << 3) | 0); // field 1, varint + encode_varint_into(&mut row0, 42); + + // Row 1: only `b = "hello"`. + let mut row1 = Vec::new(); + row1.push((2u8 << 3) | 2); // field 2, length-delimited + encode_varint_into(&mut row1, 5); + row1.extend_from_slice(b"hello"); + + let batch = enc + .encode_batch(&[Bytes::from(row0), Bytes::from(row1)]) + .unwrap(); + assert_eq!(batch.num_rows(), 2); + + let a = batch.column(0).as_primitive::(); + assert_eq!(a.value(0), 42); + assert!(a.is_null(1)); + + let b = batch.column(1).as_string::(); + assert!(batch.column(1).is_null(0)); + assert_eq!(b.value(1), "hello"); +} + +/// `message Tree { Tree next = 1; int32 leaf = 2; }` — a proto that's +/// self-referential by construction, used to drive deep plan/scan recursion. +fn self_referential_descriptor() -> MessageDescriptor { + let fd = FileDescriptorProto { + name: Some("wire_to_arrow_test_tree.proto".into()), + package: Some("wire_to_arrow_test".into()), + message_type: vec![DescriptorProto { + name: Some("Tree".into()), + field: vec![ + FieldDescriptorProto { + name: Some("next".into()), + number: Some(1), + label: Some(Label::Optional as i32), + r#type: Some(ProtoType::Message as i32), + type_name: Some(".wire_to_arrow_test.Tree".into()), + ..Default::default() + }, + FieldDescriptorProto { + name: Some("leaf".into()), + number: Some(2), + label: Some(Label::Optional as i32), + r#type: Some(ProtoType::Int32 as i32), + ..Default::default() + }, + ], + ..Default::default() + }], + ..Default::default() + }; + let set = FileDescriptorSet { file: vec![fd] }; + let mut bytes = Vec::new(); + set.encode(&mut bytes).unwrap(); + DescriptorPool::decode(bytes.as_slice()) + .unwrap() + .get_message_by_name("wire_to_arrow_test.Tree") + .unwrap() +} + +/// Build an Arrow Struct nested `levels` deep along a single `next` field, +/// with an `Int32` `leaf` at every level. Used to drive `MessagePlan::build` +/// recursion to a known depth. +fn nested_tree_struct(levels: usize) -> DataType { + if levels == 0 { + // Innermost level: just the leaf scalar, no `next`. + return DataType::Struct(ArrowFields::from(vec![Field::new( + "leaf", + DataType::Int32, + true, + )])); + } + DataType::Struct(ArrowFields::from(vec![ + Field::new("next", nested_tree_struct(levels - 1), true), + Field::new("leaf", DataType::Int32, true), + ])) +} + +#[test] +fn plan_build_rejects_schema_deeper_than_cap() { + use super::plan::MAX_NESTING_DEPTH; + + let desc = self_referential_descriptor(); + // One level over the cap. The wrapper Schema counts as depth 0, so we + // need MAX_NESTING_DEPTH levels of nested struct to step over the limit. + let deep_struct = nested_tree_struct(MAX_NESTING_DEPTH); + let schema = Schema::new(vec![Field::new("next", deep_struct, true)]); + let err = WireToArrowEncoder::new(&desc, schema).expect_err("should reject"); + assert!( + matches!(err, WireToArrowError::SchemaTooDeep { limit } if limit == MAX_NESTING_DEPTH), + "expected SchemaTooDeep, got {err:?}" + ); +} + +#[test] +fn plan_build_accepts_moderately_deep_schema() { + // A reasonably deep but legal schema must build without error. Pick a + // depth far below the cap so any reasonable real-world nesting is fine. + let desc = self_referential_descriptor(); + let deep_struct = nested_tree_struct(8); + let schema = Schema::new(vec![Field::new("next", deep_struct, true)]); + let enc = WireToArrowEncoder::new(&desc, schema).expect("8-deep schema should build"); + // And it should be able to encode an empty payload (every level absent). + let batch = enc.encode_batch(&[Bytes::new()]).unwrap(); + assert_eq!(batch.num_rows(), 1); +} + +#[test] +fn plan_build_rejects_non_nullable_singular_scalar() { + // proto3 omits default-valued singular scalars on the wire, so a + // column declared non-nullable would fail RecordBatch::try_new with + // a generic Arrow error deep in encode_batch — dropping the whole + // batch. The plan builder should reject this at init. + let desc = scalar_descriptor(); + let schema = Schema::new(vec![ + Field::new("name", DataType::LargeUtf8, true), + Field::new("id", DataType::Int32, /* nullable */ false), + Field::new("email", DataType::LargeUtf8, true), + ]); + let err = WireToArrowEncoder::new(&desc, schema).expect_err("should reject"); + assert!( + matches!( + &err, + WireToArrowError::NonNullableNotGuaranteed { name, .. } if name == "id" + ), + "expected NonNullableNotGuaranteed for 'id', got {err:?}" + ); +} + +#[test] +fn plan_build_allows_non_nullable_outer_list_and_map() { + // Repeated and map outer columns are always-present (the encoder + // emits an empty list / empty map for an absent occurrence), so a + // non-nullable declaration on the outer column is safe and must + // build without error. + let desc = rich_descriptor(); + let phone_struct = DataType::Struct(ArrowFields::from(vec![Field::new( + "number", + DataType::LargeUtf8, + true, + )])); + let phones_field = Field::new("item", phone_struct, true); + let entry_fields = ArrowFields::from(vec![ + // Arrow Map keys are mandated non-nullable by the Map type + // contract; the carve-out for Map entry sub-plans must let this + // through. + Field::new("key", DataType::LargeUtf8, false), + Field::new("value", DataType::Int32, true), + ]); + let entry_field = Arc::new(Field::new( + "key_value", + DataType::Struct(entry_fields), + false, + )); + let schema = Schema::new(vec![ + Field::new("name", DataType::LargeUtf8, true), + // Outer list non-nullable: OK, the encoder writes empty-list, not null. + Field::new( + "phones", + DataType::List(Arc::new(phones_field)), + /* nullable */ false, + ), + // Outer map non-nullable: OK, same reason. + Field::new("data", DataType::Map(entry_field, false), /* nullable */ false), + ]); + WireToArrowEncoder::new(&desc, schema).expect("should build"); +} + +#[test] +fn plan_build_rejects_non_nullable_absent_column() { + // Schema-drift case: Arrow schema has a column the proto descriptor + // doesn't carry. The encoder fills it with all nulls; a non-nullable + // declaration is a hard mismatch that must be rejected at init. + let desc = scalar_descriptor(); + let schema = Schema::new(vec![ + Field::new("name", DataType::LargeUtf8, true), + Field::new("dropped_from_proto", DataType::Int64, /* nullable */ false), + ]); + let err = WireToArrowEncoder::new(&desc, schema).expect_err("should reject"); + assert!( + matches!( + &err, + WireToArrowError::NonNullableNotGuaranteed { name, .. } + if name == "dropped_from_proto" + ), + "expected NonNullableNotGuaranteed for absent column, got {err:?}" + ); +} + +#[test] +fn plan_build_rejects_unsupported_arrow_leaf_in_scalar_slot() { + // proto says `id: int32`, Arrow says `id: Date32`. Date32 isn't in + // `TypedBuilder::supports`, so plan-build must reject up front + // rather than letting the first batch panic inside `TypedBuilder::new`. + let desc = scalar_descriptor(); + let schema = Schema::new(vec![ + Field::new("name", DataType::LargeUtf8, true), + Field::new("id", DataType::Date32, true), + Field::new("email", DataType::LargeUtf8, true), + ]); + let err = WireToArrowEncoder::new(&desc, schema).expect_err("should reject Date32"); + assert!( + matches!( + &err, + WireToArrowError::UnsupportedArrowLeafType { name, .. } if name == "id" + ), + "expected UnsupportedArrowLeafType for 'id', got {err:?}" + ); +} + +#[test] +fn plan_build_rejects_unsupported_arrow_leaf_in_absent_slot() { + // `created_at` doesn't exist in test_protobuf.Person, so it becomes + // PlanSlot::Absent. The Absent path builds via `build_absent_node`, + // which also calls `TypedBuilder::new` for leaves. Plan-build must + // validate the Arrow leaf type on absent slots too. + let desc = scalar_descriptor(); + let schema = Schema::new(vec![ + Field::new("name", DataType::LargeUtf8, true), + Field::new("created_at", DataType::Date32, true), + ]); + let err = WireToArrowEncoder::new(&desc, schema) + .expect_err("should reject Date32 on absent slot"); + assert!( + matches!( + &err, + WireToArrowError::UnsupportedArrowLeafType { name, .. } if name == "created_at" + ), + "expected UnsupportedArrowLeafType for 'created_at', got {err:?}" + ); +} + +#[test] +fn multiple_rows_preserve_order() { + let desc = scalar_descriptor(); + let schema = Schema::new(vec![Field::new("id", DataType::Int32, true)]); + let enc = WireToArrowEncoder::new(&desc, schema).unwrap(); + + let mut messages = Vec::new(); + for i in 0..5 { + let mut msg = DynamicMessage::new(desc.clone()); + msg.set_field_by_name("id", ProtoValue::I32(i * 10)); + let mut buf = Vec::new(); + msg.encode(&mut buf).unwrap(); + messages.push(Bytes::from(buf)); + } + + let batch = enc.encode_batch(&messages).unwrap(); + let ids = batch.column(0).as_primitive::(); + assert_eq!(ids.len(), 5); + for i in 0..5 { + assert_eq!(ids.value(i), (i as i32) * 10); + } +} + +#[test] +fn wire_descriptor_tags_drive_decode_regardless_of_arrow_schema_source() { + // The two mappings the encoder relies on: + // (a) Arrow schema ↔ proto descriptor — by name, at plan-build time. + // (b) Proto descriptor ↔ wire bytes — by tag number, at scan time. + // They are independent. An Arrow schema derived from any other + // descriptor (e.g. one synthesized with position-based tags) must + // not leak its tag numbers into the scan. Tags on the wire come from + // whatever descriptor the *wire* side was encoded with, and that's the + // only descriptor the serializer is told about. + // + // Here: wire descriptor uses tags 1001/1002/1003. If the scanner ever + // fell back to a position-based or otherwise-synthesized tag space + // (1/2/3), every wire tag would miss and the batch columns would be + // all-null. Full population proves decode uses the wire descriptor's + // tag numbers exclusively. + let wire_fd = FileDescriptorProto { + name: Some("wire_tag_divergence_test.proto".into()), + package: Some("wire_tag_divergence_test".into()), + message_type: vec![DescriptorProto { + name: Some("Row".into()), + field: vec![ + FieldDescriptorProto { + name: Some("name".into()), + number: Some(1001), + label: Some(Label::Optional as i32), + r#type: Some(ProtoType::String as i32), + ..Default::default() + }, + FieldDescriptorProto { + name: Some("id".into()), + number: Some(1002), + label: Some(Label::Optional as i32), + r#type: Some(ProtoType::Int32 as i32), + ..Default::default() + }, + FieldDescriptorProto { + name: Some("email".into()), + number: Some(1003), + label: Some(Label::Optional as i32), + r#type: Some(ProtoType::String as i32), + ..Default::default() + }, + ], + ..Default::default() + }], + ..Default::default() + }; + let set = FileDescriptorSet { + file: vec![wire_fd], + }; + let mut set_bytes = Vec::new(); + set.encode(&mut set_bytes).unwrap(); + let wire_desc = DescriptorPool::decode(set_bytes.as_slice()) + .unwrap() + .get_message_by_name("wire_tag_divergence_test.Row") + .unwrap(); + + let arrow_schema = Schema::new(vec![ + Field::new("name", DataType::LargeUtf8, true), + Field::new("id", DataType::Int32, true), + Field::new("email", DataType::LargeUtf8, true), + ]); + let serializer = WireToArrowSerializer::from_descriptor(wire_desc.clone(), arrow_schema) + .expect("serializer build"); + + let mut msg = DynamicMessage::new(wire_desc); + msg.set_field_by_name("name", ProtoValue::String("alice".into())); + msg.set_field_by_name("id", ProtoValue::I32(42)); + msg.set_field_by_name("email", ProtoValue::String("alice@example.com".into())); + let mut buf = Vec::new(); + msg.encode(&mut buf).unwrap(); + + let batch = serializer + .encode_to_record_batch(&[event_with_message_bytes(Bytes::from(buf))]) + .expect("encode"); + assert_eq!(batch.num_rows(), 1); + assert_eq!(batch.num_columns(), 3); + assert_eq!(batch.column(0).as_string::().value(0), "alice"); + assert_eq!( + batch + .column(1) + .as_primitive::() + .value(0), + 42 + ); + assert_eq!( + batch.column(2).as_string::().value(0), + "alice@example.com" + ); +} + +// ------------------------------------------------------------------------- +// Per-row isolation: a single malformed message in a batch must not poison +// the whole batch. The encoder pre-validates each message and drops bad +// rows from the output `RecordBatch`, counting them via the +// `wire_to_arrow_rows_dropped` metric. +// ------------------------------------------------------------------------- + +#[test] +fn encode_batch_drops_malformed_row_in_mixed_batch() { + // A single malformed message in the middle of an otherwise-valid batch + // must not fail the batch. The bad row is dropped; the valid rows still + // appear in the output `RecordBatch` in original order. + let desc = scalar_descriptor(); + let schema = Schema::new(vec![ + Field::new("name", DataType::LargeUtf8, true), + Field::new("id", DataType::Int32, true), + ]); + let enc = WireToArrowEncoder::new(&desc, schema).unwrap(); + + let mut a = DynamicMessage::new(desc.clone()); + a.set_field_by_name("name", ProtoValue::String("alice".into())); + a.set_field_by_name("id", ProtoValue::I32(1)); + let mut buf_a = Vec::new(); + a.encode(&mut buf_a).unwrap(); + + let mut b = DynamicMessage::new(desc.clone()); + b.set_field_by_name("name", ProtoValue::String("bob".into())); + b.set_field_by_name("id", ProtoValue::I32(2)); + let mut buf_b = Vec::new(); + b.encode(&mut buf_b).unwrap(); + + // Tag = (1 << 3) | 2 = 0x0a (`name`, LEN); declared length 5, only 2 + // payload bytes follow → BufferTooShort inside `try_parse_field`, + // which the encoder maps to `UnexpectedEof`. + let bad = vec![0x0a, 0x05, 0x01, 0x02]; + + let batch = enc + .encode_batch(&[Bytes::from(buf_a), Bytes::from(bad), Bytes::from(buf_b)]) + .expect("malformed row must not fail the batch"); + assert_eq!(batch.num_rows(), 2); + let names = batch.column(0).as_string::(); + assert_eq!(names.value(0), "alice"); + assert_eq!(names.value(1), "bob"); + let ids = batch.column(1).as_primitive::(); + assert_eq!(ids.value(0), 1); + assert_eq!(ids.value(1), 2); +} + +#[test] +fn encode_batch_drops_invalid_utf8_row() { + // String field with non-UTF-8 payload — the pre-validate pass catches + // this via the `from_utf8` check in `validate_scalar_from_wire`. + let desc = scalar_descriptor(); + let schema = Schema::new(vec![Field::new("name", DataType::LargeUtf8, true)]); + let enc = WireToArrowEncoder::new(&desc, schema).unwrap(); + + // Tag = (1 << 3) | 2 = 0x0a (`name`, LEN); length = 1; payload = 0xff + // (lone continuation byte, not a valid UTF-8 sequence). + let bad = vec![0x0a, 0x01, 0xff]; + + let mut ok = DynamicMessage::new(desc.clone()); + ok.set_field_by_name("name", ProtoValue::String("ok".into())); + let mut ok_buf = Vec::new(); + ok.encode(&mut ok_buf).unwrap(); + + let batch = enc + .encode_batch(&[Bytes::from(bad), Bytes::from(ok_buf)]) + .expect("invalid-UTF-8 row must not fail the batch"); + assert_eq!(batch.num_rows(), 1); + assert_eq!(batch.column(0).as_string::().value(0), "ok"); +} + +#[test] +fn encode_batch_all_malformed_returns_empty_batch() { + // Every row malformed → empty `RecordBatch` (schema preserved, zero + // rows). Surfacing the situation is the metric's job; the serializer + // doesn't conflate "every row was bad" with "configuration broken". + let desc = scalar_descriptor(); + let schema = Schema::new(vec![Field::new("name", DataType::LargeUtf8, true)]); + let enc = WireToArrowEncoder::new(&desc, schema).unwrap(); + + let bad1 = vec![0x0a, 0x05, 0x00, 0x01]; // truncated LEN value + let bad2 = vec![0x0a, 0x01, 0xff]; // invalid UTF-8 + + let batch = enc + .encode_batch(&[Bytes::from(bad1), Bytes::from(bad2)]) + .expect("all-malformed batch must still return an empty RecordBatch"); + assert_eq!(batch.num_rows(), 0); + assert_eq!(batch.num_columns(), 1); +} + +#[test] +fn encode_batch_drops_packed_scalar_eof_row() { + // Packed repeated int32 with a truncated varint in the inner blob — + // the one site (`append_repeated_scalar`'s packed loop) where the + // real scan could otherwise leave half the elements committed before + // erroring. Pre-validate sees the EOF and drops the row before any + // append. + let desc = repeated_int32_descriptor(); + let schema = Schema::new(vec![Field::new( + "numbers", + DataType::List(Arc::new(Field::new("item", DataType::Int32, true))), + true, + )]); + let enc = WireToArrowEncoder::new(&desc, schema).unwrap(); + + // Tag = (1 << 3) | 2 = 0x0a (`numbers`, LEN — packed form); length 2; + // inner = [0x80, 0x80] — a varint with continuation bits and no + // terminator → EOF inside the packed walk. + let bad = vec![0x0a, 0x02, 0x80, 0x80]; + // Packed encoding of `numbers = [7]`: one varint (7) inside a LEN blob. + let good = vec![0x0a, 0x01, 0x07]; + + let batch = enc + .encode_batch(&[Bytes::from(bad), Bytes::from(good)]) + .expect("packed EOF row must not fail the batch"); + assert_eq!(batch.num_rows(), 1); + let list = batch.column(0).as_list::(); + assert_eq!(list.value_length(0), 1); +} + +#[test] +fn encode_batch_drops_row_with_duplicate_singular_scalar_tag() { + // Proto3 parsers must accept duplicate singular tags (last-wins for + // scalars), but the encoder appends to Arrow column builders on every + // occurrence, so a second tag would diverge column lengths and fail + // `RecordBatch::try_new`. `validate_message` detects the duplicate and + // drops the row before any builder is touched. + let desc = scalar_descriptor(); + let schema = Schema::new(vec![ + Field::new("name", DataType::LargeUtf8, true), + Field::new("id", DataType::Int32, true), + ]); + let enc = WireToArrowEncoder::new(&desc, schema).unwrap(); + + let mut a = DynamicMessage::new(desc.clone()); + a.set_field_by_name("name", ProtoValue::String("alice".into())); + a.set_field_by_name("id", ProtoValue::I32(1)); + let mut buf_a = Vec::new(); + a.encode(&mut buf_a).unwrap(); + + let mut b = DynamicMessage::new(desc.clone()); + b.set_field_by_name("name", ProtoValue::String("bob".into())); + b.set_field_by_name("id", ProtoValue::I32(2)); + let mut buf_b = Vec::new(); + b.encode(&mut buf_b).unwrap(); + + // Hand-rolled bytes: two occurrences of singular tag 1 (`name`, LEN). + // 0x0a 0x03 "bob" — first occurrence, value "bob" + // 0x0a 0x05 "alice" — second occurrence, value "alice" + // Spec says last-wins ("alice"); the encoder cannot honor that without + // builder retraction, so the row must be dropped via validate. + let dup = vec![ + 0x0a, 0x03, b'b', b'o', b'b', 0x0a, 0x05, b'a', b'l', b'i', b'c', b'e', + ]; + + let batch = enc + .encode_batch(&[Bytes::from(buf_a), Bytes::from(dup), Bytes::from(buf_b)]) + .expect("duplicate-tag row must not fail the batch"); + assert_eq!(batch.num_rows(), 2); + let names = batch.column(0).as_string::(); + assert_eq!(names.value(0), "alice"); + assert_eq!(names.value(1), "bob"); + let ids = batch.column(1).as_primitive::(); + assert_eq!(ids.value(0), 1); + assert_eq!(ids.value(1), 2); +} + +#[test] +fn append_repeated_scalar_unpacked_offset_overflow_is_clean_error() { + // The Arrow `ListArray` uses i32 offsets, so `current_offset` can grow + // at most to `i32::MAX` across a batch. Without bounds checking the + // `+= 1` wraps in release mode and `OffsetBuffer::new` later asserts at + // batch finish, panicking the process. The encoder uses `checked_add` + // here so an overflow surfaces as a structured `OffsetOverflow` error. + // This converts a process-panic surface (adversarial input → crash) into + // a clean batch-level failure. + use super::append::append_repeated_scalar; + use super::builders::TypedBuilder; + let mut values = TypedBuilder::new(&DataType::Int32, 1); + let mut current_offset: i32 = i32::MAX; + let wv = WireValue::Varint(42); + let err = append_repeated_scalar(ScalarKind::Int32, &wv, &mut values, &mut current_offset) + .expect_err("unpacked repeated-scalar increment at i32::MAX must error"); + assert!( + matches!(err, WireToArrowError::OffsetOverflow { .. }), + "expected OffsetOverflow, got {err:?}", + ); +} + +#[test] +fn append_repeated_scalar_packed_offset_overflow_is_clean_error() { + // Same property for the packed-blob inner loop. A two-element packed + // varint starting from `current_offset == i32::MAX - 1` overflows on + // the second element; the loop's `checked_add` must surface that as + // `OffsetOverflow` rather than wrapping silently. + use super::append::append_repeated_scalar; + use super::builders::TypedBuilder; + let mut values = TypedBuilder::new(&DataType::Int32, 4); + // Two varint elements (7, 8) packed inline. We synthesize at the + // WireValue layer, so the outer tag + length prefix have already been + // stripped — only the inner packed payload appears here. + let blob: Vec = vec![0x07, 0x08]; + let wv = WireValue::Len(&blob); + let mut current_offset: i32 = i32::MAX - 1; + let err = append_repeated_scalar(ScalarKind::Int32, &wv, &mut values, &mut current_offset) + .expect_err("packed-scalar increment crossing i32::MAX must error"); + assert!( + matches!(err, WireToArrowError::OffsetOverflow { .. }), + "expected OffsetOverflow, got {err:?}", + ); +} + +#[test] +fn map_preserves_user_entry_field_name_and_metadata() { + // Arrow's Map spec doesn't pin a specific entry name — "entries" is + // canonical in Arrow itself; "key_value" is the Spark/Delta convention. + // Before the fix the encoder hardcoded + // "key_value" at finish time, so any caller whose schema declared a + // different entry name (or attached metadata) got `RecordBatch::try_new` + // rejection on every batch. The fix preserves the user-supplied entry + // Field unchanged. + let desc = rich_descriptor(); + let entry_fields = ArrowFields::from(vec![ + Field::new("key", DataType::LargeUtf8, false), + Field::new("value", DataType::Int32, true), + ]); + let user_entry = Arc::new( + Field::new("entries", DataType::Struct(entry_fields), false).with_metadata( + [("source".to_string(), "unit_test".to_string())] + .into_iter() + .collect(), + ), + ); + let outer_field = Field::new("data", DataType::Map(Arc::clone(&user_entry), false), true); + let schema = Schema::new(vec![outer_field.clone()]); + let enc = WireToArrowEncoder::new(&desc, schema).unwrap(); + + let mut msg = DynamicMessage::new(desc.clone()); + let mut entries: std::collections::HashMap = + std::collections::HashMap::new(); + entries.insert( + prost_reflect::MapKey::String("k1".into()), + ProtoValue::EnumNumber(1), + ); + msg.set_field_by_name("data", ProtoValue::Map(entries)); + let mut buf = Vec::new(); + msg.encode(&mut buf).unwrap(); + + let batch = enc + .encode_batch(&[Bytes::from(buf)]) + .expect("batch must finish with the user-supplied entry name preserved"); + assert_eq!(batch.num_rows(), 1); + // The schema the batch reports must match what we declared, including + // the non-default entry name and the metadata we attached. + let col_field = batch.schema().field(0).clone(); + let actual_entry = match col_field.data_type() { + DataType::Map(f, _) => Arc::clone(f), + other => panic!("expected Map, got {other:?}"), + }; + assert_eq!(actual_entry.name(), "entries"); + assert_eq!( + actual_entry.metadata().get("source").map(String::as_str), + Some("unit_test"), + "user-attached entry metadata must round-trip", + ); +} + +// ------------------------------------------------------------------------- +// Enum -> STRING column (PlanSlot::EnumString): render the proto enum varint +// as its value name, matching the arrow_stream / `proto_to_value` path. +// ------------------------------------------------------------------------- + +/// Build `message Msg { Outcome outcome = 1; }` with +/// `enum Outcome { UNSPECIFIED = 0; SUCCESS = 1; FAILURE = 2; }`. +fn singular_enum_descriptor() -> MessageDescriptor { + use prost_reflect::prost_types::{EnumDescriptorProto, EnumValueDescriptorProto}; + let fd = FileDescriptorProto { + name: Some("wire_to_arrow_enum_test.proto".into()), + package: Some("wire_to_arrow_enum_test".into()), + syntax: Some("proto3".into()), + enum_type: vec![EnumDescriptorProto { + name: Some("Outcome".into()), + value: vec![ + EnumValueDescriptorProto { + name: Some("UNSPECIFIED".into()), + number: Some(0), + ..Default::default() + }, + EnumValueDescriptorProto { + name: Some("SUCCESS".into()), + number: Some(1), + ..Default::default() + }, + EnumValueDescriptorProto { + name: Some("FAILURE".into()), + number: Some(2), + ..Default::default() + }, + ], + ..Default::default() + }], + message_type: vec![DescriptorProto { + name: Some("Msg".into()), + field: vec![ + FieldDescriptorProto { + name: Some("outcome".into()), + number: Some(1), + label: Some(Label::Optional as i32), + r#type: Some(ProtoType::Enum as i32), + type_name: Some(".wire_to_arrow_enum_test.Outcome".into()), + ..Default::default() + }, + FieldDescriptorProto { + name: Some("outcomes".into()), + number: Some(2), + label: Some(Label::Repeated as i32), + r#type: Some(ProtoType::Enum as i32), + type_name: Some(".wire_to_arrow_enum_test.Outcome".into()), + ..Default::default() + }, + ], + ..Default::default() + }], + ..Default::default() + }; + let set = FileDescriptorSet { file: vec![fd] }; + let mut bytes = Vec::new(); + set.encode(&mut bytes).unwrap(); + DescriptorPool::decode(bytes.as_slice()) + .unwrap() + .get_message_by_name("wire_to_arrow_enum_test.Msg") + .unwrap() +} + +#[test] +fn enum_to_string_renders_value_name() { + let desc = singular_enum_descriptor(); + let schema = Schema::new(vec![Field::new("outcome", DataType::LargeUtf8, true)]); + let enc = WireToArrowEncoder::new(&desc, schema).unwrap(); + + let mut msg = DynamicMessage::new(desc.clone()); + msg.set_field_by_name("outcome", ProtoValue::EnumNumber(1)); // SUCCESS + let mut buf = Vec::new(); + msg.encode(&mut buf).unwrap(); + + let batch = enc.encode_batch(&[Bytes::from(buf)]).unwrap(); + assert_eq!(batch.column(0).as_string::().value(0), "SUCCESS"); +} + +#[test] +fn enum_to_string_absent_is_null() { + // proto3 elides an enum at its zero value; `proto_to_value` only walks + // present fields, so an absent enum -> null (NOT the name of value 0). + let desc = singular_enum_descriptor(); + let schema = Schema::new(vec![Field::new("outcome", DataType::LargeUtf8, true)]); + let enc = WireToArrowEncoder::new(&desc, schema).unwrap(); + + // Empty message: `outcome` absent. + let batch = enc.encode_batch(&[Bytes::new()]).unwrap(); + assert_eq!(batch.num_rows(), 1); + assert!(batch.column(0).is_null(0), "absent enum must be null"); +} + +#[test] +fn enum_to_string_unknown_value_renders_placeholder() { + // An out-of-range enum number (e.g. a value added after this binary was + // built) has no descriptor entry. The encoder renders an + // `UNKNOWN_ENUM_VALUE__` placeholder rather than dropping the row. + let desc = singular_enum_descriptor(); + let schema = Schema::new(vec![Field::new("outcome", DataType::LargeUtf8, true)]); + let enc = WireToArrowEncoder::new(&desc, schema).unwrap(); + + let mut good = DynamicMessage::new(desc.clone()); + good.set_field_by_name("outcome", ProtoValue::EnumNumber(2)); // FAILURE + let mut good_buf = Vec::new(); + good.encode(&mut good_buf).unwrap(); + + let mut unknown = DynamicMessage::new(desc.clone()); + unknown.set_field_by_name("outcome", ProtoValue::EnumNumber(99)); // undefined + let mut unknown_buf = Vec::new(); + unknown.encode(&mut unknown_buf).unwrap(); + + let batch = enc + .encode_batch(&[Bytes::from(good_buf), Bytes::from(unknown_buf)]) + .expect("unknown-enum row must not fail the batch"); + assert_eq!(batch.num_rows(), 2, "the unknown-enum row must be kept"); + let col = batch.column(0).as_string::(); + assert_eq!(col.value(0), "FAILURE"); + assert_eq!(col.value(1), "UNKNOWN_ENUM_VALUE_Outcome_99"); +} + +#[test] +fn repeated_enum_to_string_renders_value_names() { + // `repeated Outcome` -> List: each element rendered by name. + let desc = singular_enum_descriptor(); + let outcomes = Field::new("item", DataType::LargeUtf8, true); + let schema = Schema::new(vec![Field::new( + "outcomes", + DataType::List(Arc::new(outcomes)), + true, + )]); + let enc = WireToArrowEncoder::new(&desc, schema).unwrap(); + + let mut msg = DynamicMessage::new(desc.clone()); + msg.set_field_by_name( + "outcomes", + ProtoValue::List(vec![ + ProtoValue::EnumNumber(1), // SUCCESS + ProtoValue::EnumNumber(2), // FAILURE + ProtoValue::EnumNumber(0), // UNSPECIFIED + ]), + ); + let mut buf = Vec::new(); + msg.encode(&mut buf).unwrap(); + + let batch = enc.encode_batch(&[Bytes::from(buf)]).unwrap(); + let list = batch.column(0).as_list::(); + let vals = list.value(0); + let strs = vals.as_string::(); + assert_eq!(strs.len(), 3); + assert_eq!(strs.value(0), "SUCCESS"); + assert_eq!(strs.value(1), "FAILURE"); + assert_eq!(strs.value(2), "UNSPECIFIED"); +} + +#[test] +fn repeated_enum_to_string_empty_when_absent() { + // An absent repeated field is an empty list (never null), like every + // other repeated slot. + let desc = singular_enum_descriptor(); + let outcomes = Field::new("item", DataType::LargeUtf8, true); + let schema = Schema::new(vec![Field::new( + "outcomes", + DataType::List(Arc::new(outcomes)), + true, + )]); + let enc = WireToArrowEncoder::new(&desc, schema).unwrap(); + + let batch = enc.encode_batch(&[Bytes::new()]).unwrap(); + assert_eq!(batch.num_rows(), 1); + let list = batch.column(0).as_list::(); + assert!(!list.is_null(0), "absent repeated enum must be empty list, not null"); + assert_eq!(list.value(0).len(), 0); +} + +#[test] +fn repeated_enum_to_string_unknown_value_renders_placeholder() { + // An out-of-range element renders its placeholder in place (parity with the + // singular case); the row is kept and known elements are unaffected. + let desc = singular_enum_descriptor(); + let outcomes = Field::new("item", DataType::LargeUtf8, true); + let schema = Schema::new(vec![Field::new( + "outcomes", + DataType::List(Arc::new(outcomes)), + true, + )]); + let enc = WireToArrowEncoder::new(&desc, schema).unwrap(); + + let mut msg = DynamicMessage::new(desc.clone()); + msg.set_field_by_name( + "outcomes", + ProtoValue::List(vec![ProtoValue::EnumNumber(1), ProtoValue::EnumNumber(99)]), + ); + let mut buf = Vec::new(); + msg.encode(&mut buf).unwrap(); + + let batch = enc + .encode_batch(&[Bytes::from(buf)]) + .expect("unknown-enum element must not fail the batch"); + assert_eq!(batch.num_rows(), 1, "the row with an unknown element must be kept"); + let list = batch.column(0).as_list::(); + let strs = list.value(0); + let strs = strs.as_string::(); + assert_eq!(strs.len(), 2); + assert_eq!(strs.value(0), "SUCCESS"); + assert_eq!(strs.value(1), "UNKNOWN_ENUM_VALUE_Outcome_99"); +} + +// ------------------------------------------------------------------------- +// Fuzz: random wire bytes through `encode_batch` must not panic. Any +// `Result` outcome is acceptable — we only care that bad input is reported +// as a normal error and that the scan-time recursion (which the depth cap +// also bounds) doesn't overflow the stack. +// ------------------------------------------------------------------------- + +proptest! { + #![proptest_config(ProptestConfig { + cases: 256, + // Per-case timeout in ms; bounds CI cost if a future change makes + // the scan path much slower for some inputs. + timeout: 2_000, + ..ProptestConfig::default() + })] + + /// Encoder must never panic on adversarial wire bytes against a scalar + /// schema. Most random byte sequences will hit `UnexpectedEof`, + /// `InvalidWireType`, or `WireTypeMismatch`; a few will parse but + /// produce nonsense values. All paths are fine as long as no panic. + #[test] + fn encode_batch_does_not_panic_on_random_bytes_scalar( + bytes in proptest::collection::vec(any::(), 0..256), + ) { + let desc = scalar_descriptor(); + let schema = Schema::new(vec![ + Field::new("name", DataType::LargeUtf8, true), + Field::new("id", DataType::Int32, true), + Field::new("email", DataType::LargeUtf8, true), + ]); + let enc = WireToArrowEncoder::new(&desc, schema).unwrap(); + // Any Result is acceptable — the assertion is "no panic". + let _ = enc.encode_batch(&[Bytes::from(bytes)]); + } + + /// Same property against a deeply-nestable self-referential descriptor. + /// This is the wire-side stack-overflow surface Flavio called out: + /// attacker-controlled bytes try to drive `scan_message` recursion + /// down to the plan's maximum depth. With the depth cap in place, + /// scanning bounded by the plan stays within the safe limit. + #[test] + fn encode_batch_does_not_panic_on_random_bytes_nested( + bytes in proptest::collection::vec(any::(), 0..512), + ) { + let desc = self_referential_descriptor(); + // Modest nesting in the Arrow schema — well under the cap, but + // enough that adversarial bytes have a real `next` field to chase. + let deep_struct = nested_tree_struct(8); + let schema = Schema::new(vec![Field::new("next", deep_struct, true)]); + let enc = WireToArrowEncoder::new(&desc, schema).unwrap(); + let _ = enc.encode_batch(&[Bytes::from(bytes)]); + } +} diff --git a/lib/codecs/src/encoding/format/wire_to_arrow/wire.rs b/lib/codecs/src/encoding/format/wire_to_arrow/wire.rs new file mode 100644 index 0000000000000..f6728c03a9ae1 --- /dev/null +++ b/lib/codecs/src/encoding/format/wire_to_arrow/wire.rs @@ -0,0 +1,544 @@ +//! Low-level protobuf wire-format decoding used by the wire-to-Arrow encoder. +//! +//! This is a self-contained subset of a proto wire parser: just enough to +//! walk a serialized message's tag/value pairs in a single pass and read +//! scalar values out of them. The encoder pairs these raw wire values against +//! a proto `MessageDescriptor` in [`super::scan`] / [`super::append`], so no +//! schema knowledge lives here. + +use std::fmt; + +/// Error type for protobuf wire-format parsing failures. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum ParseError { + /// Not enough bytes to read a fixed-size or length-delimited field. + BufferTooShort { + /// Number of bytes the field required. + needed: usize, + /// Number of bytes actually available. + available: usize, + /// The field number being parsed. + field_num: i32, + }, + /// Field number out of valid range (1 to 536,870,911). + InvalidFieldNumber { + /// The out-of-range field number. + field_num: i32, + }, + /// Invalid wire type value (must be 0-5). + InvalidWireType(u8), + /// Not enough bytes to parse a varint. + TruncatedVarint, + /// Deprecated group wire types are not supported. + UnsupportedGroupWireType, + /// Varint encoding uses more than 10 bytes. + VarintTooLong, +} + +impl fmt::Display for ParseError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + ParseError::BufferTooShort { + needed, + available, + field_num, + } => write!( + f, + "Field #{field_num}: Input buffer too short: need {needed} bytes, have {available}" + ), + ParseError::InvalidFieldNumber { field_num } => write!( + f, + "Field number {field_num} is out of valid range (must be 1 to 536,870,911)" + ), + ParseError::InvalidWireType(wt) => write!(f, "Invalid wire type: {wt}"), + ParseError::TruncatedVarint => write!(f, "Truncated varint"), + ParseError::UnsupportedGroupWireType => write!(f, "Group wire types are not supported"), + ParseError::VarintTooLong => write!(f, "Varint exceeds 10 bytes"), + } + } +} + +impl std::error::Error for ParseError {} + +/// Result type for wire-format parsing operations. +pub type ParseResult = Result; + +/// Raw wire value before schema interpretation. +#[derive(Debug, PartialEq, Clone, Copy)] +pub enum WireValue<'a> { + /// A base-128 varint (bool, int32/64, uint32/64, sint32/64 (zigzag), enum). + Varint(u64), + /// Fixed 8-byte value (fixed64, sfixed64, double). + I64(u64), + /// Length-delimited value (string, bytes, embedded message, packed repeated). + Len(&'a [u8]), + /// Fixed 4-byte value (fixed32, sfixed32, float). + I32(u32), +} + +/// ZigZag decode a 32-bit value (used for sint32). +#[inline(always)] +pub fn decode_zigzag32(n: u32) -> i32 { + ((n >> 1) as i32) ^ -((n & 1) as i32) +} + +/// ZigZag decode a 64-bit value (used for sint64). +#[inline(always)] +pub fn decode_zigzag64(n: u64) -> i64 { + ((n >> 1) as i64) ^ -((n & 1) as i64) +} + +/// A wire type as seen on the wire. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +#[repr(u8)] +enum WireType { + /// The Varint WireType indicates the value is a single VARINT. + Varint = 0, + /// The I64 WireType indicates that the value is precisely 8 bytes in + /// little-endian order containing a 64-bit signed integer or double type. + I64 = 1, + /// The Len WireType indicates that the value is a length represented as a + /// VARINT followed by exactly that number of bytes. + Len = 2, + /// Deprecated protobuf groups (start). + StartGroup = 3, + /// Deprecated protobuf groups (end). + EndGroup = 4, + /// The I32 WireType indicates that the value is precisely 4 bytes in + /// little-endian order containing a 32-bit signed integer or float type. + I32 = 5, +} + +impl TryFrom for WireType { + type Error = ParseError; + + #[inline(always)] + fn try_from(value: u64) -> Result { + match value { + 0 => Ok(WireType::Varint), + 1 => Ok(WireType::I64), + 2 => Ok(WireType::Len), + 3 => Ok(WireType::StartGroup), + 4 => Ok(WireType::EndGroup), + 5 => Ok(WireType::I32), + _ => Err(ParseError::InvalidWireType(value as u8)), + } + } +} + +/// Parsed wire field with field number and raw value. +#[derive(Debug, PartialEq, Clone, Copy)] +pub struct WireField<'a> { + /// The field number (tag >> 3). + pub field_num: i32, + /// The raw wire value. + pub value: WireValue<'a>, +} + +/// Parse a VARINT, returning the parsed value and the remaining bytes. +/// A 64-bit varint can require up to 10 bytes (64 bits / 7 bits per byte). +/// +/// Optimized with fast paths for 1-5 byte varints (covers ~99.9% of cases). +/// +/// Returns `Err(ParseError::TruncatedVarint)` if buffer is too short. +/// Returns `Err(ParseError::VarintTooLong)` if varint exceeds 10 bytes. +#[inline(always)] +pub fn try_read_varint(data: &[u8]) -> ParseResult<(u64, &[u8])> { + match *data { + // Empty buffer. + [] => Err(ParseError::TruncatedVarint), + // Fast path: 1-byte varint (values 0-127, very common for field tags and small ints). + [b0, ref rest @ ..] if b0 < 0x80 => Ok((b0 as u64, rest)), + // Only 1 byte but continuation bit set. + [_] => Err(ParseError::TruncatedVarint), + // Fast path: 2-byte varint (values 128-16383). + [b0, b1, ref rest @ ..] if b1 < 0x80 => { + Ok((((b0 & 0x7f) as u64) | ((b1 as u64) << 7), rest)) + } + // Only 2 bytes but continuation bit set. + [_, _] => Err(ParseError::TruncatedVarint), + // Fast path: 3-byte varint (values 16384-2097151). + [b0, b1, b2, ref rest @ ..] if b2 < 0x80 => Ok(( + ((b0 & 0x7f) as u64) | (((b1 & 0x7f) as u64) << 7) | ((b2 as u64) << 14), + rest, + )), + // Only 3 bytes but continuation bit set. + [_, _, _] => Err(ParseError::TruncatedVarint), + // Fast path: 4-byte varint (values 2097152-268435455). + [b0, b1, b2, b3, ref rest @ ..] if b3 < 0x80 => Ok(( + ((b0 & 0x7f) as u64) + | (((b1 & 0x7f) as u64) << 7) + | (((b2 & 0x7f) as u64) << 14) + | ((b3 as u64) << 21), + rest, + )), + // Only 4 bytes but continuation bit set. + [_, _, _, _] => Err(ParseError::TruncatedVarint), + // Fast path: 5-byte varint (values 268435456-34359738367). + [b0, b1, b2, b3, b4, ref rest @ ..] if b4 < 0x80 => Ok(( + ((b0 & 0x7f) as u64) + | (((b1 & 0x7f) as u64) << 7) + | (((b2 & 0x7f) as u64) << 14) + | (((b3 & 0x7f) as u64) << 21) + | ((b4 as u64) << 28), + rest, + )), + // Slow path: 6+ byte varints (rare). + _ => parse_varint_slow(data), + } +} + +/// Slow path for varints with 6+ bytes. +#[inline(always)] +fn parse_varint_slow(data: &[u8]) -> ParseResult<(u64, &[u8])> { + let mut value = 0u64; + let mut shift = 0; + + // Process bytes 0-8 (each contributes 7 bits). + for i in 0..9 { + let Some(&b) = data.get(i) else { + return Err(ParseError::TruncatedVarint); + }; + value |= ((b & 0x7f) as u64) << shift; + if b < 0x80 { + return Ok((value, &data[i + 1..])); + } + shift += 7; + } + + // 10th byte (index 9): can only contribute bit 0 (9*7 + 1 = 64 bits total). + // Bits 1-6 would overflow u64, bit 7 (continuation) would require 11+ bytes. + let Some(&b) = data.get(9) else { + return Err(ParseError::TruncatedVarint); + }; + if b > 0x01 { + return Err(ParseError::VarintTooLong); + } + value |= (b as u64) << shift; + Ok((value, &data[10..])) +} + +/// Convert a tag into a field number and a WireType. +/// Returns error if wire type is invalid or field number is out of range. +#[inline(always)] +fn try_unpack_tag(tag: u64) -> ParseResult<(i32, WireType)> { + let field_num = (tag >> 3) as i32; + let wire_type = WireType::try_from(tag & 0x7)?; + + // Validate field number range per protobuf spec. + // Field numbers must be 1 to 536,870,911 (2^29 - 1). + if !(1..=536_870_911).contains(&field_num) { + return Err(ParseError::InvalidFieldNumber { field_num }); + } + + Ok((field_num, wire_type)) +} + +/// Parse a field, returning the field and remaining bytes. +/// +/// Returns error on malformed input (truncated buffer, invalid wire type, etc.) +#[inline(always)] +pub fn try_parse_field(data: &[u8]) -> ParseResult<(WireField<'_>, &[u8])> { + let (tag, remainder) = try_read_varint(data)?; + let (field_num, wire_type) = try_unpack_tag(tag)?; + let (fieldvalue, remainder) = match wire_type { + WireType::Varint => { + let (value, remainder) = try_read_varint(remainder)?; + (WireValue::Varint(value), remainder) + } + WireType::I64 => { + // Fixed 8 bytes in little-endian order. + let Some((bytes, rest)) = remainder.split_first_chunk::<8>() else { + return Err(ParseError::BufferTooShort { + needed: 8, + available: remainder.len(), + field_num, + }); + }; + let value = u64::from_le_bytes(*bytes); + (WireValue::I64(value), rest) + } + WireType::Len => { + let (len, remainder) = try_read_varint(remainder)?; + let len = len as usize; + if remainder.len() < len { + return Err(ParseError::BufferTooShort { + needed: len, + available: remainder.len(), + field_num, + }); + } + let (value, remainder) = remainder.split_at(len); + (WireValue::Len(value), remainder) + } + WireType::I32 => { + // Fixed 4 bytes in little-endian order. + let Some((bytes, rest)) = remainder.split_first_chunk::<4>() else { + return Err(ParseError::BufferTooShort { + needed: 4, + available: remainder.len(), + field_num, + }); + }; + let value = u32::from_le_bytes(*bytes); + (WireValue::I32(value), rest) + } + WireType::StartGroup | WireType::EndGroup => { + return Err(ParseError::UnsupportedGroupWireType); + } + }; + Ok(( + WireField { + field_num, + value: fieldvalue, + }, + remainder, + )) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn varint_parsing() { + // Test cases: (input bytes, expected value, expected remaining bytes). + let cases = [ + // 1-byte varints. + (&[0x00][..], 0, &[][..]), + (&[0x01][..], 1, &[][..]), + (&[0x7F][..], 127, &[][..]), + // 2-byte varints. + (&[0x80, 0x01][..], 128, &[][..]), + // 3-byte varints. + (&[0x80, 0x80, 0x01][..], 16384, &[][..]), + // 4-byte varints. + (&[0x80, 0x80, 0x80, 0x01][..], 2097152, &[][..]), + (&[0xFF, 0xFF, 0xFF, 0x7F][..], 268435455, &[][..]), + // 5-byte varints. + (&[0x80, 0x80, 0x80, 0x80, 0x01][..], 268435456, &[][..]), + (&[0xFF, 0xFF, 0xFF, 0xFF, 0x7F][..], 34359738367, &[][..]), + // With trailing bytes. + (&[0x01, 0x02, 0x03][..], 1, &[0x02, 0x03][..]), + (&[0x80, 0x80, 0x80, 0x01, 0xFF][..], 2097152, &[0xFF][..]), + ]; + for (data, expected_val, expected_rest) in cases { + let (value, rest) = try_read_varint(data).expect("Failed to parse varint"); + assert_eq!(value, expected_val, "data: {data:?}"); + assert_eq!(rest, expected_rest, "data: {data:?}"); + } + } + + #[test] + fn varint_errors() { + // Truncated varints. + assert_eq!(try_read_varint(&[0x80]), Err(ParseError::TruncatedVarint)); + assert_eq!( + try_read_varint(&[0x80, 0x80]), + Err(ParseError::TruncatedVarint) + ); + assert_eq!( + try_read_varint(&[0x80, 0x80, 0x80]), + Err(ParseError::TruncatedVarint) + ); + assert_eq!( + try_read_varint(&[0x80, 0x80, 0x80, 0x80]), + Err(ParseError::TruncatedVarint) + ); + assert_eq!( + try_read_varint(&[0x80, 0x80, 0x80, 0x80, 0x80]), + Err(ParseError::TruncatedVarint) + ); + + // Varint too long (11 bytes). + let too_long = &[ + 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x01, + ]; + assert_eq!(try_read_varint(too_long), Err(ParseError::VarintTooLong)); + } + + #[test] + fn varint_10th_byte_validation() { + // 10-byte varints: bytes 0-8 all have continuation bit set, byte 9 is the 10th byte. + // The 10th byte can only have bit 0 set (bits 1-6 would overflow u64, bit 7 would need 11+ bytes). + + // Valid: 10th byte = 0x00 (contributes 0 to the value). + let valid_zero = &[0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x00]; + let (value, rest) = try_read_varint(valid_zero).unwrap(); + assert_eq!(value, 0); + assert!(rest.is_empty()); + + // Valid: 10th byte = 0x01 (sets bit 63, gives 2^63). + let valid_one = &[0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x01]; + let (value, rest) = try_read_varint(valid_one).unwrap(); + assert_eq!(value, 1u64 << 63); + assert!(rest.is_empty()); + + // Valid: u64::MAX = all bits set. + let max_u64 = &[0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0x01]; + let (value, rest) = try_read_varint(max_u64).unwrap(); + assert_eq!(value, u64::MAX); + assert!(rest.is_empty()); + + // Invalid: 10th byte = 0x02 (bit 1 set, would overflow u64). + let overflow_bit1 = &[0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x02]; + assert_eq!( + try_read_varint(overflow_bit1), + Err(ParseError::VarintTooLong) + ); + + // Invalid: 10th byte = 0x7F (bits 1-6 all set, would overflow u64). + let overflow_bits = &[0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x7F]; + assert_eq!( + try_read_varint(overflow_bits), + Err(ParseError::VarintTooLong) + ); + + // Invalid: 10th byte = 0x80 (continuation bit set, would need 11+ bytes). + let continuation = &[0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80]; + assert_eq!( + try_read_varint(continuation), + Err(ParseError::VarintTooLong) + ); + } + + #[test] + fn zigzag_decoding() { + let cases32 = [ + (0, 0), + (1, -1), + (2, 1), + (3, -2), + (4, 2), + (99, -50), + (100, 50), + ]; + for (encoded, expected) in cases32 { + assert_eq!(decode_zigzag32(encoded), expected, "zigzag32({encoded})"); + } + + let cases64 = [(0, 0), (1, -1), (2, 1), (3, -2), (4, 2)]; + for (encoded, expected) in cases64 { + assert_eq!(decode_zigzag64(encoded), expected, "zigzag64({encoded})"); + } + } + + #[test] + fn wire_type_conversion() { + let valid = [ + (0, WireType::Varint), + (1, WireType::I64), + (2, WireType::Len), + (3, WireType::StartGroup), + (4, WireType::EndGroup), + (5, WireType::I32), + ]; + for (val, expected) in valid { + assert_eq!(WireType::try_from(val), Ok(expected)); + } + + assert_eq!( + WireType::try_from(6u64), + Err(ParseError::InvalidWireType(6)) + ); + assert_eq!( + WireType::try_from(7u64), + Err(ParseError::InvalidWireType(7)) + ); + } + + #[test] + fn field_parsing() { + // (data, expected_field_num, expected_value). + let cases = [ + // Varint: field 1, value 150. Tag = 8, 150 = 0x96 0x01. + (&[8, 0x96, 0x01][..], 1, WireValue::Varint(150)), + // I32: field 1, tag = 13, value 0x01020304 little-endian. + ( + &[13, 0x04, 0x03, 0x02, 0x01][..], + 1, + WireValue::I32(0x01020304), + ), + // I64: field 1, tag = 9. + ( + &[9, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08][..], + 1, + WireValue::I64(0x0807060504030201), + ), + // Len: field 1, tag = 10, length 3. + ( + &[10, 3, 0xAA, 0xBB, 0xCC][..], + 1, + WireValue::Len(&[0xAA, 0xBB, 0xCC]), + ), + ]; + for (data, expected_num, expected_val) in cases { + let (field, rest) = try_parse_field(data).unwrap(); + assert_eq!(field.field_num, expected_num, "data: {data:?}"); + assert_eq!(field.value, expected_val, "data: {data:?}"); + assert!(rest.is_empty(), "data: {data:?}"); + } + } + + #[test] + fn field_parsing_errors() { + // Invalid wire type 6: tag = (1 << 3) | 6 = 14. + assert_eq!(try_parse_field(&[14]), Err(ParseError::InvalidWireType(6))); + + // Group wire type: tag = (1 << 3) | 3 = 11. + assert_eq!( + try_parse_field(&[11]), + Err(ParseError::UnsupportedGroupWireType) + ); + + // Buffer too short for I32: tag 13, only 2 bytes. + assert_eq!( + try_parse_field(&[13, 0x01, 0x02]), + Err(ParseError::BufferTooShort { + needed: 4, + available: 2, + field_num: 1, + }) + ); + + // Buffer too short for I64: tag 9, only 4 bytes. + assert_eq!( + try_parse_field(&[9, 0x01, 0x02, 0x03, 0x04]), + Err(ParseError::BufferTooShort { + needed: 8, + available: 4, + field_num: 1, + }) + ); + + // Buffer too short for Len: tag 10, length 100, only 3 bytes. + assert_eq!( + try_parse_field(&[10, 100, 0x01, 0x02, 0x03]), + Err(ParseError::BufferTooShort { + needed: 100, + available: 3, + field_num: 1, + }) + ); + + // Invalid field number 0: tag = (0 << 3) | 0 = 0. + assert_eq!( + try_parse_field(&[0]), + Err(ParseError::InvalidFieldNumber { field_num: 0 }) + ); + + // Invalid field number 536_870_912 (2^29) - exceeds max valid field number 536_870_911. + // Tag = (536_870_912 << 3) | 0 = 4_294_967_296 -> varint [0x80, 0x80, 0x80, 0x80, 0x10]. + assert_eq!( + try_parse_field(&[0x80, 0x80, 0x80, 0x80, 0x10]), + Err(ParseError::InvalidFieldNumber { + field_num: 536_870_912 + }) + ); + + // Max valid field number 536_870_911 (2^29 - 1) is okay. + // Tag = (536_870_911 << 3) | 0 = 4_294_967_288 -> varint [0xF8, 0xFF, 0xFF, 0xFF, 0x0F]. + let (field, _) = try_parse_field(&[0xF8, 0xFF, 0xFF, 0xFF, 0x0F, 0x01]).unwrap(); + assert_eq!(field.field_num, 536_870_911); + } +} diff --git a/lib/codecs/src/encoding/mod.rs b/lib/codecs/src/encoding/mod.rs index 361a62a301533..82fb54ec6f99d 100644 --- a/lib/codecs/src/encoding/mod.rs +++ b/lib/codecs/src/encoding/mod.rs @@ -16,6 +16,7 @@ pub use encoder::{Encoder, EncoderKind}; #[cfg(feature = "arrow")] pub use format::{ ArrowEncodingError, ArrowStreamSerializer, ArrowStreamSerializerConfig, SchemaProvider, + WireToArrowEncoder, WireToArrowError, WireToArrowSerializer, WireToArrowSerializerConfig, find_null_non_nullable_fields, }; pub use format::{ diff --git a/lib/codecs/src/encoding/serializer.rs b/lib/codecs/src/encoding/serializer.rs index 431f0356c3aee..b9a81aa2b625c 100644 --- a/lib/codecs/src/encoding/serializer.rs +++ b/lib/codecs/src/encoding/serializer.rs @@ -5,7 +5,10 @@ use vector_config::configurable_component; use vector_core::{config::DataType, event::Event, schema}; #[cfg(feature = "arrow")] -use super::format::{ArrowStreamSerializer, ArrowStreamSerializerConfig}; +use super::format::{ + ArrowStreamSerializer, ArrowStreamSerializerConfig, WireToArrowSerializer, + WireToArrowSerializerConfig, +}; #[cfg(feature = "opentelemetry")] use super::format::{OtlpSerializer, OtlpSerializerConfig}; #[cfg(feature = "parquet")] @@ -165,6 +168,17 @@ pub enum BatchSerializerConfig { /// [apache_arrow]: https://arrow.apache.org/ #[serde(rename = "arrow_stream")] ArrowStream(ArrowStreamSerializerConfig), + /// Decodes proto wire bytes from each event's `message` field directly + /// into an [Apache Arrow][apache_arrow] `RecordBatch`, bypassing the + /// generic `ProtobufDeserializer -> Event -> ArrowStreamSerializer` + /// chain. + /// + /// Requires a proto descriptor (`desc_file` + `message_type`) for the + /// incoming bytes; the sink injects the output Arrow `Schema`. + /// + /// [apache_arrow]: https://arrow.apache.org/ + #[serde(rename = "wire_to_arrow")] + WireToArrow(WireToArrowSerializerConfig), /// Encodes events in [Apache Parquet][apache_parquet] columnar format. /// /// [apache_parquet]: https://parquet.apache.org/ @@ -184,6 +198,10 @@ impl BatchSerializerConfig { let serializer = ArrowStreamSerializer::new(arrow_config.clone())?; Ok(super::BatchSerializer::Arrow(serializer)) } + BatchSerializerConfig::WireToArrow(config) => { + let serializer = WireToArrowSerializer::new(config.clone())?; + Ok(super::BatchSerializer::WireToArrow(serializer)) + } #[cfg(feature = "parquet")] BatchSerializerConfig::Parquet(parquet_config) => { let serializer = ParquetSerializer::new(parquet_config.clone())?; @@ -196,6 +214,7 @@ impl BatchSerializerConfig { pub fn input_type(&self) -> DataType { match self { BatchSerializerConfig::ArrowStream(arrow_config) => arrow_config.input_type(), + BatchSerializerConfig::WireToArrow(config) => config.input_type(), #[cfg(feature = "parquet")] BatchSerializerConfig::Parquet(parquet_config) => parquet_config.input_type(), } @@ -205,6 +224,7 @@ impl BatchSerializerConfig { pub fn schema_requirement(&self) -> schema::Requirement { match self { BatchSerializerConfig::ArrowStream(arrow_config) => arrow_config.schema_requirement(), + BatchSerializerConfig::WireToArrow(config) => config.schema_requirement(), #[cfg(feature = "parquet")] BatchSerializerConfig::Parquet(parquet_config) => parquet_config.schema_requirement(), }