linliu-code commented on code in PR #668:
URL: https://github.com/apache/hudi-rs/pull/668#discussion_r3779923400
##########
crates/core/src/file_group/log_file/avro.rs:
##########
@@ -16,77 +16,148 @@
* specific language governing permissions and limitations
* under the License.
*/
-use apache_avro::AvroResult;
-use apache_avro::types::Value as AvroValue;
-use apache_avro::{Schema as AvroSchema, from_avro_datum};
-use std::io::Read;
+use crate::Result;
+use crate::error::CoreError;
+use arrow_array::{ArrayRef, RecordBatch};
+use arrow_avro::reader::{Decoder as ArrowAvroDecoder, ReaderBuilder};
+use arrow_avro::schema::{AvroSchema as ArrowAvroSchema, SINGLE_OBJECT_MAGIC,
SchemaStore};
+use arrow_cast::cast;
+use arrow_schema::{DataType, Field, Schema};
+use std::sync::Arc;
-pub struct AvroDataBlockContentReader<R: Read> {
- reader: R,
- writer_schema: AvroSchema,
- remaining_records: u32,
+/// Decodes the bare Avro record bodies in a Hudi log block straight into
Arrow.
+///
+/// A Hudi data block frames each datum with a four-byte length and no Avro
+/// framing of its own, while `arrow-avro`'s decoder expects each record to be
+/// introduced by a Single Object Encoding prefix — the two-byte marker and the
+/// writer schema's fingerprint. The two are reconciled by synthesizing that
+/// prefix: the schema is registered once, and the ten bytes it yields are
+/// written ahead of every body.
+///
+/// The alternative is `arrow-avro`'s `AvroBodyDecoder`, which takes bare
bodies
+/// directly but is not released yet.
+pub struct AvroBlockDecoder {
+ decoder: ArrowAvroDecoder,
+ /// Marker plus fingerprint, identical for every record in the block.
+ prefix: [u8; 10],
+ /// Reused across records so framing costs one copy, not one allocation.
+ framed: Vec<u8>,
}
-impl<R: Read> AvroDataBlockContentReader<R> {
- pub fn new(reader: R, writer_schema: &AvroSchema, num_records: u32) ->
Self {
- Self {
- reader,
- writer_schema: writer_schema.clone(),
- remaining_records: num_records,
- }
- }
-}
+impl AvroBlockDecoder {
+ /// Build a decoder for a block written with `writer_schema_json`.
+ pub fn try_new(writer_schema_json: &str, batch_size: usize) ->
Result<Self> {
+ let mut store = SchemaStore::new();
+ let fingerprint = store
+ .register(ArrowAvroSchema::new(writer_schema_json.to_string()))
+ .map_err(|e| {
+ CoreError::LogBlockError(format!("Failed to register block
writer schema: {e}"))
+ })?;
-impl<R: Read> Iterator for AvroDataBlockContentReader<R> {
- type Item = AvroResult<AvroValue>;
+ let arrow_avro::schema::Fingerprint::Rabin(rabin) = fingerprint else {
+ return Err(CoreError::LogBlockError(format!(
+ "Expected a Rabin fingerprint for the block writer schema, got
{fingerprint:?}"
+ )));
+ };
- fn next(&mut self) -> Option<Self::Item> {
- if self.remaining_records == 0 {
- return None;
- }
+ let decoder = ReaderBuilder::new()
+ .with_writer_schema_store(store)
+ .with_active_fingerprint(fingerprint)
+ .with_batch_size(batch_size)
+ .build_decoder()
+ .map_err(|e| {
+ CoreError::LogBlockError(format!("Failed to build the Avro
decoder: {e}"))
+ })?;
- self.remaining_records -= 1;
+ let mut prefix = [0u8; 10];
+ prefix[..2].copy_from_slice(&SINGLE_OBJECT_MAGIC);
+ prefix[2..].copy_from_slice(&rabin.to_le_bytes());
- let mut record_content_length = [0u8; 4];
- match self.reader.read_exact(&mut record_content_length) {
- Ok(_) => {}
- Err(e) => {
- return Some(Err(apache_avro::Error::new(
- apache_avro::error::Details::ReadBytes(e),
- )));
- }
- }
+ Ok(Self {
+ decoder,
+ prefix,
+ framed: Vec::new(),
+ })
+ }
- let record_content_length = u32::from_be_bytes(record_content_length);
+ /// Decode one record body, returning a batch once enough rows have
accrued.
+ pub fn decode(&mut self, body: &[u8]) -> Result<Option<RecordBatch>> {
+ self.framed.clear();
+ self.framed.reserve(self.prefix.len() + body.len());
+ self.framed.extend_from_slice(&self.prefix);
+ self.framed.extend_from_slice(body);
- let mut record_reader = (&mut self.reader).take(record_content_length
as u64);
+ let consumed = self
+ .decoder
+ .decode(&self.framed)
+ .map_err(|e| CoreError::LogBlockError(format!("Failed to decode a
log record: {e}")))?;
+ if consumed != self.framed.len() {
+ return Err(CoreError::LogBlockError(format!(
+ "Log record decoded partially: {consumed} of {} bytes",
+ self.framed.len()
+ )));
+ }
- let result = from_avro_datum(&self.writer_schema, &mut record_reader,
None);
+ if self.decoder.batch_is_full() {
+ return self.flush();
+ }
+ Ok(None)
+ }
- Some(result)
+ /// Drain whatever rows are held, if any.
+ pub fn flush(&mut self) -> Result<Option<RecordBatch>> {
+ let batch = self.decoder.flush().map_err(|e| {
+ CoreError::LogBlockError(format!("Failed to flush decoded records:
{e}"))
+ })?;
+ batch.map(normalize_utc_timestamps).transpose()
}
}
-#[cfg(test)]
-mod tests {
- use super::*;
- use std::io::Cursor;
-
- #[test]
- fn test_read_error_on_truncated_data() {
- // Create a simple Avro schema
- let schema = AvroSchema::parse_str(r#"{"type": "null"}"#).unwrap();
-
- // Create a reader with only 2 bytes when we need 4 bytes for record
length
- let truncated_data = vec![0u8, 1u8];
- let reader = Cursor::new(truncated_data);
+/// Spell a UTC timestamp's zone the way the parquet reader does.
+///
+/// `arrow-avro` writes the zone of an Avro `timestamp-*` as the offset
+/// `+00:00`; parquet writes `UTC`. They denote the same zone, but Arrow
+/// compares timezones as strings, so a log batch and the base batch it merges
+/// with would be judged to have different types and refuse to concatenate.
+///
+/// Only the field's type label changes — the values are already UTC instants,
+/// so this rebinds metadata rather than converting anything.
+fn normalize_utc_timestamps(batch: RecordBatch) -> Result<RecordBatch> {
+ fn is_utc_alias(tz: &str) -> bool {
+ matches!(tz, "+00:00" | "+0000" | "00:00" | "Z" | "z")
+ }
- // Create reader expecting 1 record but with insufficient data
- let mut avro_reader = AvroDataBlockContentReader::new(reader, &schema,
1);
+ let needs_fix = batch
+ .schema()
+ .fields()
+ .iter()
+ .any(|f| matches!(f.data_type(), DataType::Timestamp(_, Some(tz)) if
is_utc_alias(tz)));
Review Comment:
You're right that the normalizer only inspects top-level fields, and no,
`reconcile_batch_to_schema` does not normalize nested timezones downstream — so
a timestamp nested in a struct, list or map would keep arrow-avro's `+00:00`.
I checked whether it's reachable: I walked every parquet file across the
whole fixture corpus looking for a timestamp nested inside a struct/list/map,
and there are **zero**. So this is untested territory rather than a known
break, and writing recursive retyping for `StructArray`/`ListArray`/`MapArray`
would mean shipping code no fixture can exercise — which is the thing this port
has been trying not to do.
So I've documented the scope on the function instead of guessing at the
implementation: "Top-level fields only. A timestamp nested in a struct, list or
map keeps arrow-avro's spelling; no fixture in the corpus has one, so recursing
here would be untested code standing in for a case nothing can currently
exercise. Not yet implemented, deliberately." If a fixture with a nested
timestamp shows up, that's the trigger to implement it properly with a gold
comparison behind it.
##########
crates/core/src/file_group/log_file/avro.rs:
##########
@@ -16,77 +16,148 @@
* specific language governing permissions and limitations
* under the License.
*/
-use apache_avro::AvroResult;
-use apache_avro::types::Value as AvroValue;
-use apache_avro::{Schema as AvroSchema, from_avro_datum};
-use std::io::Read;
+use crate::Result;
+use crate::error::CoreError;
+use arrow_array::{ArrayRef, RecordBatch};
+use arrow_avro::reader::{Decoder as ArrowAvroDecoder, ReaderBuilder};
+use arrow_avro::schema::{AvroSchema as ArrowAvroSchema, SINGLE_OBJECT_MAGIC,
SchemaStore};
+use arrow_cast::cast;
+use arrow_schema::{DataType, Field, Schema};
+use std::sync::Arc;
-pub struct AvroDataBlockContentReader<R: Read> {
- reader: R,
- writer_schema: AvroSchema,
- remaining_records: u32,
+/// Decodes the bare Avro record bodies in a Hudi log block straight into
Arrow.
+///
+/// A Hudi data block frames each datum with a four-byte length and no Avro
+/// framing of its own, while `arrow-avro`'s decoder expects each record to be
+/// introduced by a Single Object Encoding prefix — the two-byte marker and the
+/// writer schema's fingerprint. The two are reconciled by synthesizing that
+/// prefix: the schema is registered once, and the ten bytes it yields are
+/// written ahead of every body.
+///
+/// The alternative is `arrow-avro`'s `AvroBodyDecoder`, which takes bare
bodies
+/// directly but is not released yet.
+pub struct AvroBlockDecoder {
+ decoder: ArrowAvroDecoder,
+ /// Marker plus fingerprint, identical for every record in the block.
+ prefix: [u8; 10],
+ /// Reused across records so framing costs one copy, not one allocation.
+ framed: Vec<u8>,
}
-impl<R: Read> AvroDataBlockContentReader<R> {
- pub fn new(reader: R, writer_schema: &AvroSchema, num_records: u32) ->
Self {
- Self {
- reader,
- writer_schema: writer_schema.clone(),
- remaining_records: num_records,
- }
- }
-}
+impl AvroBlockDecoder {
+ /// Build a decoder for a block written with `writer_schema_json`.
+ pub fn try_new(writer_schema_json: &str, batch_size: usize) ->
Result<Self> {
+ let mut store = SchemaStore::new();
+ let fingerprint = store
+ .register(ArrowAvroSchema::new(writer_schema_json.to_string()))
+ .map_err(|e| {
+ CoreError::LogBlockError(format!("Failed to register block
writer schema: {e}"))
+ })?;
-impl<R: Read> Iterator for AvroDataBlockContentReader<R> {
- type Item = AvroResult<AvroValue>;
+ let arrow_avro::schema::Fingerprint::Rabin(rabin) = fingerprint else {
+ return Err(CoreError::LogBlockError(format!(
+ "Expected a Rabin fingerprint for the block writer schema, got
{fingerprint:?}"
+ )));
+ };
- fn next(&mut self) -> Option<Self::Item> {
- if self.remaining_records == 0 {
- return None;
- }
+ let decoder = ReaderBuilder::new()
+ .with_writer_schema_store(store)
+ .with_active_fingerprint(fingerprint)
+ .with_batch_size(batch_size)
+ .build_decoder()
+ .map_err(|e| {
+ CoreError::LogBlockError(format!("Failed to build the Avro
decoder: {e}"))
+ })?;
- self.remaining_records -= 1;
+ let mut prefix = [0u8; 10];
+ prefix[..2].copy_from_slice(&SINGLE_OBJECT_MAGIC);
+ prefix[2..].copy_from_slice(&rabin.to_le_bytes());
- let mut record_content_length = [0u8; 4];
- match self.reader.read_exact(&mut record_content_length) {
- Ok(_) => {}
- Err(e) => {
- return Some(Err(apache_avro::Error::new(
- apache_avro::error::Details::ReadBytes(e),
- )));
- }
- }
+ Ok(Self {
+ decoder,
+ prefix,
+ framed: Vec::new(),
+ })
+ }
- let record_content_length = u32::from_be_bytes(record_content_length);
+ /// Decode one record body, returning a batch once enough rows have
accrued.
+ pub fn decode(&mut self, body: &[u8]) -> Result<Option<RecordBatch>> {
+ self.framed.clear();
+ self.framed.reserve(self.prefix.len() + body.len());
+ self.framed.extend_from_slice(&self.prefix);
+ self.framed.extend_from_slice(body);
- let mut record_reader = (&mut self.reader).take(record_content_length
as u64);
+ let consumed = self
+ .decoder
+ .decode(&self.framed)
+ .map_err(|e| CoreError::LogBlockError(format!("Failed to decode a
log record: {e}")))?;
+ if consumed != self.framed.len() {
+ return Err(CoreError::LogBlockError(format!(
+ "Log record decoded partially: {consumed} of {} bytes",
+ self.framed.len()
+ )));
+ }
- let result = from_avro_datum(&self.writer_schema, &mut record_reader,
None);
+ if self.decoder.batch_is_full() {
+ return self.flush();
+ }
+ Ok(None)
+ }
- Some(result)
+ /// Drain whatever rows are held, if any.
+ pub fn flush(&mut self) -> Result<Option<RecordBatch>> {
+ let batch = self.decoder.flush().map_err(|e| {
+ CoreError::LogBlockError(format!("Failed to flush decoded records:
{e}"))
+ })?;
+ batch.map(normalize_utc_timestamps).transpose()
}
}
-#[cfg(test)]
-mod tests {
- use super::*;
- use std::io::Cursor;
-
- #[test]
- fn test_read_error_on_truncated_data() {
- // Create a simple Avro schema
- let schema = AvroSchema::parse_str(r#"{"type": "null"}"#).unwrap();
-
- // Create a reader with only 2 bytes when we need 4 bytes for record
length
- let truncated_data = vec![0u8, 1u8];
- let reader = Cursor::new(truncated_data);
+/// Spell a UTC timestamp's zone the way the parquet reader does.
+///
+/// `arrow-avro` writes the zone of an Avro `timestamp-*` as the offset
+/// `+00:00`; parquet writes `UTC`. They denote the same zone, but Arrow
+/// compares timezones as strings, so a log batch and the base batch it merges
+/// with would be judged to have different types and refuse to concatenate.
+///
+/// Only the field's type label changes — the values are already UTC instants,
+/// so this rebinds metadata rather than converting anything.
+fn normalize_utc_timestamps(batch: RecordBatch) -> Result<RecordBatch> {
+ fn is_utc_alias(tz: &str) -> bool {
+ matches!(tz, "+00:00" | "+0000" | "00:00" | "Z" | "z")
+ }
- // Create reader expecting 1 record but with insufficient data
- let mut avro_reader = AvroDataBlockContentReader::new(reader, &schema,
1);
+ let needs_fix = batch
+ .schema()
+ .fields()
+ .iter()
+ .any(|f| matches!(f.data_type(), DataType::Timestamp(_, Some(tz)) if
is_utc_alias(tz)));
+ if !needs_fix {
+ return Ok(batch);
+ }
- // Should return an error because read_exact fails
- let result = avro_reader.next();
- assert!(result.is_some());
- assert!(result.unwrap().is_err());
+ let mut fields: Vec<Field> = Vec::with_capacity(batch.num_columns());
+ let mut columns: Vec<ArrayRef> = Vec::with_capacity(batch.num_columns());
+ for (field, column) in batch.schema().fields().iter().zip(batch.columns())
{
+ match field.data_type() {
+ DataType::Timestamp(unit, Some(tz)) if is_utc_alias(tz) => {
+ let retimed = cast(column, &DataType::Timestamp(*unit,
Some("UTC".into())))
Review Comment:
Fixed. Renamed to `relabeled`, with a comment at the binding: "Relabeled,
not converted: same instants, different spelling of the same zone."
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
