laminar_connectors/kafka/avro/
mod.rs

1//! Avro deserialization using `arrow-avro` with Confluent Schema Registry.
2//!
3//! [`AvroDeserializer`] implements [`RecordDeserializer`] by wrapping the
4//! `arrow-avro` push-based [`Decoder`], which
5//! natively supports the Confluent wire format (`0x00` + 4-byte BE schema ID
6//! + Avro payload).
7
8use std::collections::HashSet;
9use std::sync::Arc;
10
11use arrow_array::RecordBatch;
12use arrow_avro::reader::{Decoder, ReaderBuilder};
13use arrow_avro::schema::{AvroSchema, Fingerprint, FingerprintAlgorithm, SchemaStore};
14use arrow_schema::SchemaRef;
15use parking_lot::Mutex;
16
17use crate::error::{ConnectorError, SerdeError};
18use crate::kafka::schema_registry::SchemaRegistryClient;
19use crate::serde::{Format, RecordDeserializer};
20
21const DECODER_BATCH_CAPACITY: usize = 8192;
22
23/// Confluent wire format magic byte.
24const CONFLUENT_MAGIC: u8 = 0x00;
25
26/// Size of the Confluent wire format header (1 magic + 4 schema ID).
27const CONFLUENT_HEADER_SIZE: usize = 5;
28
29/// Avro deserializer backed by `arrow-avro` with optional Schema Registry.
30///
31/// Supports both raw Avro and the Confluent wire format. When a Schema
32/// Registry client is provided, unknown schema IDs are fetched and
33/// registered automatically.
34pub struct AvroDeserializer {
35    /// Schema store shared with the Decoder.
36    schema_store: SchemaStore,
37    /// Optional Schema Registry client for resolving unknown IDs.
38    schema_registry: Option<Arc<SchemaRegistryClient>>,
39    /// Set of schema IDs already registered in the store.
40    known_ids: HashSet<i32>,
41    /// Reused across batches; rebuilt when `register_schema` runs.
42    decoder: Mutex<Option<Decoder>>,
43}
44
45impl AvroDeserializer {
46    /// Creates a new Avro deserializer without Schema Registry integration.
47    ///
48    /// The caller must register schemas manually via [`register_schema`](Self::register_schema).
49    #[must_use]
50    pub fn new() -> Self {
51        Self {
52            schema_store: SchemaStore::new_with_type(FingerprintAlgorithm::Id),
53            schema_registry: None,
54            known_ids: HashSet::new(),
55            decoder: Mutex::new(None),
56        }
57    }
58
59    /// Creates a new Avro deserializer with Schema Registry integration.
60    ///
61    /// Unknown schema IDs encountered in the Confluent wire format will
62    /// be fetched from the registry automatically.
63    #[must_use]
64    pub fn with_schema_registry(registry: Arc<SchemaRegistryClient>) -> Self {
65        Self {
66            schema_store: SchemaStore::new_with_type(FingerprintAlgorithm::Id),
67            schema_registry: Some(registry),
68            known_ids: HashSet::new(),
69            decoder: Mutex::new(None),
70        }
71    }
72
73    /// Registers an Avro schema with a Confluent schema ID.
74    ///
75    /// # Errors
76    ///
77    /// Returns `SerdeError` if the fingerprint cannot be set.
78    #[allow(clippy::cast_sign_loss)]
79    pub fn register_schema(
80        &mut self,
81        schema_id: i32,
82        avro_schema_json: &str,
83    ) -> Result<(), SerdeError> {
84        let avro_schema = AvroSchema::new(avro_schema_json.to_string());
85        // Use Fingerprint::Id directly — NOT load_fingerprint_id which
86        // applies from_be byte-swap meant for raw wire bytes.
87        let fp = Fingerprint::Id(schema_id as u32);
88        self.schema_store
89            .set(fp, avro_schema)
90            .map_err(|e| SerdeError::MalformedInput(format!("failed to register schema: {e}")))?;
91        self.known_ids.insert(schema_id);
92        // Schema store changed — drop the cached decoder so it rebuilds
93        // against the new store on the next deserialize_batch call.
94        *self.decoder.lock() = None;
95        Ok(())
96    }
97
98    /// Ensures a schema ID is registered, fetching from SR if needed.
99    ///
100    /// Called by the Kafka source connector when an unknown schema ID is
101    /// encountered in the Confluent wire format during poll.
102    ///
103    /// # Errors
104    ///
105    /// Preserves Schema Registry `ConnectorError` classification for remote
106    /// resolution failures and returns `ConnectorError::Serde` for local
107    /// decoder registration failures.
108    /// Returns `Ok(true)` if this was a newly registered schema ID,
109    /// `Ok(false)` if already known.
110    pub async fn ensure_schema_registered(
111        &mut self,
112        schema_id: i32,
113    ) -> Result<bool, ConnectorError> {
114        if self.known_ids.contains(&schema_id) {
115            return Ok(false);
116        }
117
118        let registry = self.schema_registry.as_ref().ok_or(ConnectorError::Serde(
119            SerdeError::SchemaNotFound { schema_id },
120        ))?;
121
122        let cached = registry.resolve_confluent_id(schema_id).await?;
123
124        self.register_schema(schema_id, &cached.schema_str)
125            .map_err(ConnectorError::Serde)?;
126        Ok(true)
127    }
128
129    /// Extracts the Confluent schema ID from a wire-format message.
130    ///
131    /// Returns `None` if the message is not in Confluent wire format.
132    #[must_use]
133    pub fn extract_confluent_id(data: &[u8]) -> Option<i32> {
134        if data.len() < CONFLUENT_HEADER_SIZE || data[0] != CONFLUENT_MAGIC {
135            return None;
136        }
137        let id = i32::from_be_bytes([data[1], data[2], data[3], data[4]]);
138        Some(id)
139    }
140}
141
142impl Default for AvroDeserializer {
143    fn default() -> Self {
144        Self::new()
145    }
146}
147
148impl RecordDeserializer for AvroDeserializer {
149    fn deserialize(&self, data: &[u8], schema: &SchemaRef) -> Result<RecordBatch, SerdeError> {
150        self.deserialize_batch(&[data], schema)
151    }
152
153    fn deserialize_batch(
154        &self,
155        records: &[&[u8]],
156        schema: &SchemaRef,
157    ) -> Result<RecordBatch, SerdeError> {
158        if records.is_empty() {
159            return Ok(RecordBatch::new_empty(schema.clone()));
160        }
161
162        let mut guard = self.decoder.lock();
163        let decoder = if let Some(d) = guard.as_mut() {
164            d
165        } else {
166            let d = ReaderBuilder::new()
167                .with_batch_size(DECODER_BATCH_CAPACITY)
168                .with_writer_schema_store(self.schema_store.clone())
169                .build_decoder()
170                .map_err(|e| SerdeError::MalformedInput(format!("failed to build decoder: {e}")))?;
171            guard.insert(d)
172        };
173
174        let mut partials: Vec<RecordBatch> = Vec::new();
175        for record in records {
176            let mut offset = 0;
177            while offset < record.len() {
178                let consumed = decoder
179                    .decode(&record[offset..])
180                    .map_err(|e| SerdeError::MalformedInput(format!("Avro decode error: {e}")))?;
181                if consumed == 0 {
182                    break;
183                }
184                offset += consumed;
185                if decoder.batch_is_full() {
186                    if let Some(b) = decoder
187                        .flush()
188                        .map_err(|e| SerdeError::MalformedInput(format!("Avro flush: {e}")))?
189                    {
190                        partials.push(b);
191                    }
192                }
193            }
194        }
195        if let Some(b) = decoder
196            .flush()
197            .map_err(|e| SerdeError::MalformedInput(format!("Avro flush: {e}")))?
198        {
199            partials.push(b);
200        }
201
202        match partials.len() {
203            0 => Err(SerdeError::MalformedInput("no records decoded".into())),
204            1 => Ok(partials.pop().unwrap()),
205            _ => arrow_select::concat::concat_batches(schema, &partials)
206                .map_err(|e| SerdeError::MalformedInput(format!("concat: {e}"))),
207        }
208    }
209
210    fn format(&self) -> Format {
211        Format::Avro
212    }
213
214    fn as_any_mut(&mut self) -> Option<&mut dyn std::any::Any> {
215        Some(self)
216    }
217}
218
219impl std::fmt::Debug for AvroDeserializer {
220    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
221        f.debug_struct("AvroDeserializer")
222            .field("known_ids", &self.known_ids)
223            .field("has_registry", &self.schema_registry.is_some())
224            .finish_non_exhaustive()
225    }
226}
227
228#[cfg(test)]
229mod tests;
laminar_connectors/kafka/avro/mod.rs

laminar_connectors/kafka/avro/
mod.rs
