jecsand838 commented on code in PR #8006: URL: https://github.com/apache/arrow-rs/pull/8006#discussion_r2248553873
########## arrow-avro/src/reader/mod.rs: ########## @@ -182,21 +174,130 @@ impl Decoder { FingerprintAlgorithm::Rabin, SchemaStore::fingerprint_algorithm, ); + // The loop stops when the batch is full, a schema change is staged, + // or handle_prefix indicates we need more bytes (Some(0)). while total_consumed < data.len() && self.remaining_capacity > 0 { - if let Some(prefix_bytes) = self.handle_prefix(&data[total_consumed..], hash_type)? { - // A batch is complete when its `remaining_capacity` is 0. It may be completed early if - // a schema change is detected or there are insufficient bytes to read the next prefix. - // A schema change requires a new batch. - total_consumed += prefix_bytes; - break; + match self.handle_prefix(&data[total_consumed..], hash_type)? { + None => { + // No prefix: decode one row. + let n = self.active_decoder.decode(&data[total_consumed..], 1)?; + total_consumed += n; + self.remaining_capacity -= 1; + } + Some(0) => { + // Detected start of a prefix but need more bytes. + break; + } + Some(n) => { + // Consumed a complete prefix (n > 0). Stage flush and stop. + total_consumed += n; + break; + } } - let n = self.active_decoder.decode(&data[total_consumed..], 1)?; - total_consumed += n; - self.remaining_capacity -= 1; } Ok(total_consumed) } + // Attempt to handle a single‑object‑encoding prefix at the current position. + // + // * Ok(None) – buffer does not start with the prefix. + // * Ok(Some(0)) – prefix detected, but the buffer is too short; caller should await more bytes. + // * Ok(Some(n)) – consumed `n > 0` bytes of a complete prefix (magic and fingerprint). + fn handle_prefix( + &mut self, + buf: &[u8], + hash_type: FingerprintAlgorithm, + ) -> Result<Option<usize>, ArrowError> { + // If there is no schema store, prefixes are unrecognized. + if self.writer_schema_store.is_none() { + return Ok(None); // Continue to decode the next record + } + // Need at least the magic bytes to decide (2 bytes). + let Some(magic_bytes) = buf.get(..SINGLE_OBJECT_MAGIC.len()) else { + return Ok(Some(0)); // Get more bytes + }; + // Bail out early if the magic does not match. + if magic_bytes != SINGLE_OBJECT_MAGIC { + return Ok(None); // Continue to decode the next record + } + // Try to parse the fingerprint that follows the magic. + let fingerprint_size = match hash_type { + FingerprintAlgorithm::Rabin => self + .handle_fingerprint::<8>(&buf[SINGLE_OBJECT_MAGIC.len()..], |bytes| { Review Comment: Good call, ty for that. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: github-unsubscr...@arrow.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org