Improve arrow-avro decoding for one-record messages - #10713
Conversation
github-actions
Bot
added
arrow
jordepic
changed the title
|
FYI @jecsand838 -- could you help review this PR? |
Absolutely! I'll have time tonight / tomorrow morning to review this. |
|
@jordepic is there currently a way to convert multiple Avro Datums in a single .avro to a RecordBatches directly? A Raw Binary Encoded Avro |
|
@kinshuk-bb Yes, for raw binary Avro containing multiple back-to-back datums encoded with the same known writer schema, this PR supports decoding them directly into RecordBatches without adding framing. Configure a Decoder with that writer schema, repeatedly call decode_datum on the remaining byte slice, advance by the returned consumed-byte count, and call flush whenever batch_is_full (plus once at the end). Each call decodes exactly one datum and leaves subsequent datums untouched. If by I am adding explicit coverage for concatenated raw datums and batch boundaries to make this use case clear. |
|
@jecsand838 if you have any time to take a look in the next few days I'd greatly appreciate it! This is a really big win for StreamFusion and any other system trying to beat out avro processing from the JVM! |
|
@jordepic thank you! i was relying on adding framing just to decode to be sure i wouldn't face issues. will try this out now. An addition to the cargo docs that this is directly possible now would be great. |
|
@kinshuk-bb Addressed in 92a4210: the crate-level documentation now includes a runnable example decoding consecutive bare Avro datums directly with |
jecsand838
left a comment
jecsand838
left a comment
There was a problem hiding this comment.
@jordepic Thank you so much for getting this PR up!
LMGTM. I left a few comments. Let me know what you think.
| /// Decode exactly one unframed Avro datum with the active writer schema. | ||
| /// | ||
| /// This is intended for transports such as Kafka where the message boundary is external to | ||
| /// Avro. It returns the number of datum bytes consumed, allowing the caller to ignore transport | ||
| /// payload bytes after the first datum when its format contract requires that behavior. | ||
| /// Consecutive unframed datums can be decoded by repeatedly passing the unconsumed suffix. | ||
| /// If the current batch is full, this method returns `Ok(0)` until [`Self::flush`] is called. | ||
| /// | ||
| /// The decoder must already have the desired active fingerprint, and this method does not | ||
| /// inspect or switch framing fingerprints. | ||
| /// | ||
| /// # Errors | ||
| /// | ||
| /// Returns an error if the datum is incomplete, malformed, or incompatible with the active | ||
| /// writer schema. | ||
| pub fn decode_datum(&mut self, data: &[u8]) -> Result<usize, AvroError> { | ||
| if self.remaining_capacity == 0 { | ||
| return Ok(0); | ||
| } | ||
| let consumed = self.active_decoder.decode(data, 1)?; | ||
| self.remaining_capacity -= 1; | ||
| self.awaiting_body = false; | ||
| Ok(consumed) | ||
| } |
There was a problem hiding this comment.
I'd recommend we avoid adding a second public decoding method and instead select the input grammar when building the decoder.
A decoder should generally consume one stable wire format for its lifetime, and this follows existing arrow-rs patterns such as arrow_json::StructMode, IPC’s DictionaryHandling, and the Avro writer’s construction-time choice between AvroSoeFormat and AvroBinaryFormat.
I suggest replacing decode_datum with:
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub enum DecoderMode {
#[default]
Framed,
UnframedDatum,
}Then expose:
pub fn with_decoder_mode(mut self, mode: DecoderMode) -> Self {
self.decoder_mode = mode;
self
}Decoder::decode can dispatch internally:
pub fn decode(&mut self, data: &[u8]) -> Result<usize, AvroError> {
match self.mode {
DecoderMode::Framed => self.decode_framed(data),
DecoderMode::UnframedDatum => self.decode_unframed(data),
}
}There was a problem hiding this comment.
Implemented in 2a5f2d9. The public decode_datum method is gone; ReaderBuilder::with_decoder_mode now selects DecoderMode::Framed (default) or DecoderMode::UnframedDatum at construction, and Decoder::decode dispatches internally. Updated the crate/module/builder docs, runnable example, README, and existing unframed tests to use the single decoding entry point.
| if self.remaining_capacity == 0 { | ||
| return Ok(0); | ||
| } |
There was a problem hiding this comment.
The new unframed mode may make Ok(0) too ambiguous. An empty record or a record containing only null fields is a valid datum that consumes zero bytes but still appends one row. The same result currently means the batch was already full and no row was appended. A caller that interprets zero as flush and retry can therefore duplicate a successfully decoded row.
I'd recommending giving batch-full a distinct outcome such as AvroError::BatchFull or something along those lines:
if self.capacity == 0 {
return Err(AvroError::BatchFull);
}
// A successful zero-width datum remains unambiguously Ok(0)
let consumed = self.active_decoder.decode(data, 1)?;
self.capacity -= 1;
Ok(consumed)There was a problem hiding this comment.
Addressed in 2a5f2d9 with AvroError::BatchFull for unframed decoding when capacity is exhausted. A successful zero-width datum still returns Ok(0), and the new regression test covers both an empty record and a record containing only a null field, including flushing and decoding another row without duplication.
| } | ||
| Self::Union(u) => u.append_null()?, | ||
| Self::Nullable(_, null_buffer, inner) => { | ||
| Self::Nullable(_, null_buffer, _, pending_nulls) => { |
There was a problem hiding this comment.
This defers child placeholders, but each null still calls NullBufferBuilder::append(false), so long null runs retain one bitmap write per row. The non-null path also calls the large append_nulls dispatcher when pending_nulls == 0.
Could we defer both the validity suffix and child placeholders behind one helper?
#[inline]
fn materialize_pending(
validity: &mut NullBufferBuilder,
values: &mut Decoder,
pending: &mut usize,
) -> Result<(), AvroError> {
let count = *pending;
if count == 0 {
return Ok(());
}
values.append_nulls(count)?;
validity.append_n_nulls(count);
*pending = 0;
Ok(())
}Null branches would only increment pending_nulls. Before a non-null/default value or flush, call this helper; after successfully decoding a value, call validity.append_non_null(). This makes long and all-null runs bulk operations and avoids zero-count dispatch on dense nullable data.
There was a problem hiding this comment.
Implemented in 2a5f2d9. NullableDecoder::materialize_pending now bulk-materializes child placeholders and the validity suffix together, while null branches only increment pending_nulls. Non-null/default paths skip the bulk dispatcher when no nulls are pending and append validity only after the value succeeds. Added a regression test that verifies neither bitmap nor child values are touched during a 64-row null run.
| other => { | ||
| for _ in 0..count { | ||
| other.append_null()?; | ||
| } |
There was a problem hiding this comment.
The catch-all calls append_null() count times, re-running the full Decoder match for every placeholder. Consequently, nullable arrays, maps, fixed values, UUIDs, enums, decimals, custom primitives, and REE do not receive the intended bulk optimization.
Perhaps add direct bulk arms where possible? for example:
Self::Uuid(values) => {
values.resize(values.len() + 16 * count, 0);
}
Self::Fixed(width, values) => {
values.resize(values.len() + (*width as usize) * count, 0);
}
Self::Enum(values, _, _) => {
values.resize(values.len() + count, 0);
}
Self::Decimal128(_, _, _, builder) => {
builder.append_value_n(0, count);
}
Self::Decimal256(_, _, _, builder) => {
builder.append_value_n(i256::ZERO, count);
}
Self::RunEndEncoded(_, len, inner) => {
inner.append_nulls(count)?;
*len += count;
}For offset-backed types, we could reserve count offsets before pushing repeated zero lengths. I would also make the match exhaustive and keep any necessarily per-row union handling explicit, so future variants cannot silently fall back to a slower implementation. This could always be looked into as a follow-up.
There was a problem hiding this comment.
Addressed in 2a5f2d9. append_nulls now matches Decoder exhaustively and has direct bulk arms for arrays/maps/strings with reserved offsets, fixed values, UUIDs, enums, all decimal widths, duration builders, custom primitive/temporal types, nested records, and run-end encoding. Union handling remains explicitly per-row where required, and append_null simply delegates to append_nulls(1).
| /// Nullable value plus trailing null placeholders not yet materialized in the child decoder. | ||
| Nullable(NullablePlan, NullBufferBuilder, Box<Decoder>, usize), |
There was a problem hiding this comment.
Adding pending_nulls as a fourth positional field makes the nullable state transitions difficult to audit across append_null, append_nulls, append_default, decode, and flush.
Could this state be represented by a named structure?
struct NullableDecoder {
plan: NullablePlan,
validity: NullBufferBuilder,
values: Box<Decoder>,
pending_nulls: usize,
}A materialize_pending() method on this structure would centralize the invariant and ensure every transition updates the child, validity bitmap, and pending count consistently. If append_nulls is made exhaustive, append_null() could also delegate to append_nulls(1), removing the duplicated per-variant implementation.
There was a problem hiding this comment.
Implemented in 2a5f2d9 with a named NullableDecoder containing plan, validity, values, and pending_nulls. Its materialize_pending method owns the shared invariant, and append_null, bulk append, defaults, decoding, and flush all transition through that representation consistently.
|
@alamb would you mind rerunning this guy? |
|
Sorry @alamb silly clippy formatting issue, could we give it another go? |
4 participants
Which issue does this PR close?
Rationale for this change
Some messaging systems deliver records one at a time. For example, Kafka gives a consumer the complete byte array for one message, so the consumer already knows where that message begins and ends. If that message contains one Avro record, the decoder does not need an Avro header to identify the record boundary.
This is not limited to schema-registry-framed Avro. Apache Flink supports both
format = 'avro-confluent', where each Kafka message includes a Confluent header containing its schema ID, andformat = 'avro', where each message contains a bare Avro datum with no Confluent header. For the latter, Flink derives the writer schema from the table definition, so the consumer already knows the schema even though the payload carries no schema ID or framing. Supporting this format with the existing decoder currently requires manufacturing a synthetic prefix and copying every message.In Avro terminology, one encoded value is called a datum. A datum can be a primitive value or a complete record.
The existing
Decoder::decodemethod expects each record to include an Avro framing prefix. That prefix identifies the Avro format and the writer schema. A caller that already has one complete Kafka message and already knows its schema must therefore add a temporary prefix and copy the payload before decoding it, or separately inspect the schema to determine how many bytes belong to the record.Nullable nested records have a separate performance cost. Consider an event with three optional nested records where only one is populated on each row. For every absent record, the decoder currently walks through all of its child fields and immediately appends placeholder values. Long sequences of absent records repeat the same work row by row.
What changes are included in this PR?
Decoder::decode_datum, which decodes one complete Avro value directly from the supplied bytes using the writer schema already selected on the decoder. It returns the number of bytes used by that value, leaving any remaining bytes untouched.The second change only affects how the decoder builds arrays internally. The returned Arrow arrays are unchanged, and all child arrays are brought to the correct length before a batch is returned.
Are these changes tested?
Yes.
cargo test -p arrow-avro --all-features: 479 unit tests passed; 26 documentation tests passed; 1 documentation test ignored.cargo clippy -p arrow-avro --all-targets --all-features -- -D warningscargo fmt --all -- --checkThe benchmark below decodes all 10,000 rows into one batch. It was run on an Apple M1 Max against a saved
mainbaseline.Command:
cargo bench -p arrow-avro --bench decoder -- 'SparseNested\(Struct\)/10000' --baseline mainAre there any user-facing changes?
Yes.
Decodergains a new, non-breakingdecode_datummethod for callers that already have the complete bytes for one Avro value and have already selected its writer schema. Existing decoding methods and framed input behavior are unchanged.AI assistance disclosure: Codex was used to help port the implementation, draft tests and the benchmark, and prepare the issue and PR text. The resulting code and all reported outputs were reviewed before submission, and I take responsibility for the contribution.