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, VecDeque};
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/// Writer schemas kept registered at once. With a Schema Registry, the least recently used schema
24/// that the batch being decoded does not use is evicted and fetched again when it reappears.
25const WRITER_SCHEMA_CAPACITY: usize = 64;
26
27/// Confluent wire format magic byte.
28const CONFLUENT_MAGIC: u8 = 0x00;
29
30/// Size of the Confluent wire format header (1 magic + 4 schema ID).
31const CONFLUENT_HEADER_SIZE: usize = 5;
32
33/// Avro deserializer backed by `arrow-avro` with optional Schema Registry.
34///
35/// Supports both raw Avro and the Confluent wire format. When a Schema
36/// Registry client is provided, unknown schema IDs are fetched and
37/// registered automatically.
38pub struct AvroDeserializer {
39    /// Schema store shared with the Decoder.
40    schema_store: SchemaStore,
41    /// Optional Schema Registry client for resolving unknown IDs.
42    schema_registry: Option<Arc<SchemaRegistryClient>>,
43    /// Registered writer schemas, least recently used first.
44    writers: VecDeque<(i32, AvroSchema)>,
45    /// Writer schemas the batch being decoded uses; never evicted before it decodes.
46    batch_ids: HashSet<i32>,
47    /// Reused across batches; rebuilt when `register_schema` runs.
48    decoder: Mutex<Option<Decoder>>,
49    reader_schema: Option<AvroSchema>,
50    projection: Option<Vec<usize>>,
51    schema_metrics: Option<super::KafkaSourceMetrics>,
52}
53
54impl AvroDeserializer {
55    /// Creates a new Avro deserializer without Schema Registry integration.
56    ///
57    /// The caller must register schemas manually via [`register_schema`](Self::register_schema).
58    #[must_use]
59    pub fn new() -> Self {
60        Self {
61            schema_store: SchemaStore::new_with_type(FingerprintAlgorithm::Id),
62            schema_registry: None,
63            writers: VecDeque::new(),
64            batch_ids: HashSet::new(),
65            decoder: Mutex::new(None),
66            reader_schema: None,
67            projection: None,
68            schema_metrics: None,
69        }
70    }
71
72    /// Creates a new Avro deserializer with Schema Registry integration.
73    ///
74    /// Unknown schema IDs encountered in the Confluent wire format will
75    /// be fetched from the registry automatically.
76    #[must_use]
77    pub fn with_schema_registry(registry: Arc<SchemaRegistryClient>) -> Self {
78        Self {
79            schema_store: SchemaStore::new_with_type(FingerprintAlgorithm::Id),
80            schema_registry: Some(registry),
81            writers: VecDeque::new(),
82            batch_ids: HashSet::new(),
83            decoder: Mutex::new(None),
84            reader_schema: None,
85            projection: None,
86            schema_metrics: None,
87        }
88    }
89
90    /// Bind the immutable native reader and its SQL projection before consuming records.
91    ///
92    /// # Errors
93    /// Rejects missing native content, invalid IDs, and unsupported projections.
94    pub fn bind_reader(
95        &mut self,
96        binding: &crate::schema::resolution::SchemaBinding,
97        logical: &SchemaRef,
98    ) -> Result<(), ConnectorError> {
99        let native = binding.value.as_ref().ok_or_else(|| {
100            ConnectorError::SchemaMismatch("Avro reader has no native schema".into())
101        })?;
102        let resolved = native
103            .definition
104            .get("resolved")
105            .ok_or_else(|| {
106                ConnectorError::SchemaMismatch("Avro reader has no resolved native schema".into())
107            })?
108            .to_string();
109        let external = crate::kafka::schema_registry::avro_to_arrow_schema(&resolved)?;
110        let projection = logical
111            .fields()
112            .iter()
113            .map(|field| {
114                external.index_of(field.name()).map_err(|_| {
115                    ConnectorError::SchemaMismatch(format!(
116                        "reader field '{}' is absent from the native record",
117                        field.name()
118                    ))
119                })
120            })
121            .collect::<Result<Vec<_>, _>>()?;
122        let id = crate::kafka::schema_resolution::contract_schema_id(binding)?;
123        let id = i32::try_from(id).map_err(|_| {
124            ConnectorError::SchemaMismatch("Avro schema ID exceeds the supported wire range".into())
125        })?;
126        self.register_schema(id, &resolved)
127            .map_err(ConnectorError::Serde)?;
128        self.reader_schema = Some(AvroSchema::new(resolved));
129        self.projection = Some(projection);
130        Ok(())
131    }
132
133    pub(crate) fn set_schema_metrics(&mut self, metrics: super::KafkaSourceMetrics) {
134        self.schema_metrics = Some(metrics);
135    }
136
137    /// Registers an Avro schema with a Confluent schema ID.
138    ///
139    /// # Errors
140    ///
141    /// Returns `SerdeError` if the fingerprint cannot be set.
142    #[allow(clippy::cast_sign_loss)]
143    pub fn register_schema(
144        &mut self,
145        schema_id: i32,
146        avro_schema_json: &str,
147    ) -> Result<(), SerdeError> {
148        if schema_id <= 0 || avro_schema_json.len() > 1024 * 1024 {
149            return Err(SerdeError::MalformedInput(
150                "invalid schema ID or schema exceeds 1 MiB".into(),
151            ));
152        }
153        if self.touch(schema_id) {
154            return Ok(());
155        }
156        if self.writers.len() >= WRITER_SCHEMA_CAPACITY {
157            self.evict_one()?;
158        }
159        let avro_schema = AvroSchema::new(avro_schema_json.to_string());
160        // Use Fingerprint::Id directly — NOT load_fingerprint_id which
161        // applies from_be byte-swap meant for raw wire bytes.
162        let fp = Fingerprint::Id(schema_id as u32);
163        self.schema_store
164            .set(fp, avro_schema.clone())
165            .map_err(|e| SerdeError::MalformedInput(format!("failed to register schema: {e}")))?;
166        self.writers.push_back((schema_id, avro_schema));
167        // Schema store changed — drop the cached decoder so it rebuilds
168        // against the new store on the next deserialize_batch call.
169        *self.decoder.lock() = None;
170        Ok(())
171    }
172
173    /// Drop the least recently used writer schema the current batch does not use, rebuilding the
174    /// schema store without it. Only a registry-backed deserializer can fetch it again.
175    #[allow(clippy::cast_sign_loss)]
176    fn evict_one(&mut self) -> Result<(), SerdeError> {
177        if self.schema_registry.is_none() {
178            return Err(SerdeError::MalformedInput(format!(
179                "at most {WRITER_SCHEMA_CAPACITY} writer schemas can be registered without a \
180                 Schema Registry to fetch evicted ones"
181            )));
182        }
183        let position = self
184            .writers
185            .iter()
186            .position(|(id, _)| !self.batch_ids.contains(id))
187            .ok_or_else(|| {
188                SerdeError::MalformedInput(format!(
189                    "one poll references more than {WRITER_SCHEMA_CAPACITY} Avro writer \
190                     schemas; lower max.poll.records"
191                ))
192            })?;
193        self.writers.remove(position);
194        let mut store = SchemaStore::new_with_type(FingerprintAlgorithm::Id);
195        for (id, schema) in &self.writers {
196            store
197                .set(Fingerprint::Id(*id as u32), schema.clone())
198                .map_err(|e| SerdeError::MalformedInput(format!("rebuild schema store: {e}")))?;
199        }
200        self.schema_store = store;
201        Ok(())
202    }
203
204    /// Mark `schema_id` most recently used; `false` when it is not registered.
205    fn touch(&mut self, schema_id: i32) -> bool {
206        let Some(position) = self.writers.iter().position(|(id, _)| *id == schema_id) else {
207            return false;
208        };
209        if let Some(entry) = self.writers.remove(position) {
210            self.writers.push_back(entry);
211        }
212        true
213    }
214
215    /// Start resolving a new poll's writer schemas; earlier batches no longer pin any.
216    pub(crate) fn begin_batch(&mut self) {
217        self.batch_ids.clear();
218    }
219
220    /// Whether `schema_id` is currently registered.
221    #[must_use]
222    pub fn is_registered(&self, schema_id: i32) -> bool {
223        self.writers.iter().any(|(id, _)| *id == schema_id)
224    }
225
226    /// Ensures a schema ID is registered, fetching from SR if needed.
227    ///
228    /// Called by the Kafka source connector when an unknown schema ID is
229    /// encountered in the Confluent wire format during poll.
230    ///
231    /// # Errors
232    ///
233    /// Preserves Schema Registry `ConnectorError` classification for remote
234    /// resolution failures and returns `ConnectorError::Serde` for local
235    /// decoder registration failures.
236    /// Returns `Ok(true)` if this was a newly registered schema ID,
237    /// `Ok(false)` if already known.
238    pub async fn ensure_schema_registered(
239        &mut self,
240        schema_id: i32,
241    ) -> Result<bool, ConnectorError> {
242        self.batch_ids.insert(schema_id);
243        if self.touch(schema_id) {
244            return Ok(false);
245        }
246
247        let registry = self.schema_registry.as_ref().ok_or(ConnectorError::Serde(
248            SerdeError::SchemaNotFound { schema_id },
249        ))?;
250
251        let mut observation = WriterFetchObservation::new(self.schema_metrics.clone());
252        let cached = registry.resolve_confluent_id(schema_id).await?;
253        observation.succeeded = true;
254
255        self.register_schema(schema_id, &cached.resolved_schema_str)
256            .map_err(ConnectorError::Serde)?;
257        Ok(true)
258    }
259
260    /// Extracts the Confluent schema ID from a wire-format message.
261    ///
262    /// Returns `None` if the message is not in Confluent wire format.
263    #[must_use]
264    pub fn extract_confluent_id(data: &[u8]) -> Option<i32> {
265        if data.len() < CONFLUENT_HEADER_SIZE || data[0] != CONFLUENT_MAGIC {
266            return None;
267        }
268        let id = i32::from_be_bytes([data[1], data[2], data[3], data[4]]);
269        Some(id)
270    }
271}
272
273impl Default for AvroDeserializer {
274    fn default() -> Self {
275        Self::new()
276    }
277}
278
279impl RecordDeserializer for AvroDeserializer {
280    fn deserialize(&self, data: &[u8], schema: &SchemaRef) -> Result<RecordBatch, SerdeError> {
281        self.deserialize_batch(&[data], schema)
282    }
283
284    fn deserialize_batch(
285        &self,
286        records: &[&[u8]],
287        schema: &SchemaRef,
288    ) -> Result<RecordBatch, SerdeError> {
289        if records.is_empty() {
290            return Ok(RecordBatch::new_empty(schema.clone()));
291        }
292
293        let mut guard = self.decoder.lock();
294        let decoder = if let Some(d) = guard.as_mut() {
295            d
296        } else {
297            let mut builder = ReaderBuilder::new()
298                .with_batch_size(DECODER_BATCH_CAPACITY)
299                .with_writer_schema_store(self.schema_store.clone());
300            if let Some(reader) = &self.reader_schema {
301                builder = builder.with_reader_schema(reader.clone());
302            }
303            if let Some(projection) = &self.projection {
304                builder = builder.with_projection(projection.clone());
305            }
306            let d = builder
307                .build_decoder()
308                .map_err(|e| SerdeError::MalformedInput(format!("failed to build decoder: {e}")))?;
309            guard.insert(d)
310        };
311
312        let mut partials: Vec<RecordBatch> = Vec::new();
313        for record in records {
314            let mut offset = 0;
315            while offset < record.len() {
316                let consumed = decoder
317                    .decode(&record[offset..])
318                    .map_err(|e| SerdeError::MalformedInput(format!("Avro decode error: {e}")))?;
319                if consumed == 0 {
320                    break;
321                }
322                offset += consumed;
323                if decoder.batch_is_full() {
324                    if let Some(b) = decoder
325                        .flush()
326                        .map_err(|e| SerdeError::MalformedInput(format!("Avro flush: {e}")))?
327                    {
328                        partials.push(b);
329                    }
330                }
331            }
332        }
333        if let Some(b) = decoder
334            .flush()
335            .map_err(|e| SerdeError::MalformedInput(format!("Avro flush: {e}")))?
336        {
337            partials.push(b);
338        }
339
340        match partials.len() {
341            0 => Err(SerdeError::MalformedInput("no records decoded".into())),
342            1 => {
343                let batch = partials.remove(0);
344                if batch.schema().as_ref() == schema.as_ref() {
345                    return Ok(batch);
346                }
347                RecordBatch::try_new(Arc::clone(schema), batch.columns().to_vec()).map_err(
348                    |error| {
349                        SerdeError::MalformedInput(format!("committed reader mismatch: {error}"))
350                    },
351                )
352            }
353            _ => arrow_select::concat::concat_batches(schema, &partials)
354                .map_err(|e| SerdeError::MalformedInput(format!("concat: {e}"))),
355        }
356    }
357
358    fn format(&self) -> Format {
359        Format::Avro
360    }
361
362    fn as_any_mut(&mut self) -> Option<&mut dyn std::any::Any> {
363        Some(self)
364    }
365}
366
367impl std::fmt::Debug for AvroDeserializer {
368    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
369        f.debug_struct("AvroDeserializer")
370            .field(
371                "writer_ids",
372                &self.writers.iter().map(|(id, _)| *id).collect::<Vec<_>>(),
373            )
374            .field("has_registry", &self.schema_registry.is_some())
375            .finish_non_exhaustive()
376    }
377}
378
379#[cfg(test)]
380mod tests;
381
382struct WriterFetchObservation {
383    metrics: Option<super::KafkaSourceMetrics>,
384    started: std::time::Instant,
385    succeeded: bool,
386}
387
388impl WriterFetchObservation {
389    fn new(metrics: Option<super::KafkaSourceMetrics>) -> Self {
390        if let Some(metrics) = &metrics {
391            metrics.schema_cache_misses.inc();
392            metrics.schema_unresolved_records.inc();
393        }
394        Self {
395            metrics,
396            started: std::time::Instant::now(),
397            succeeded: false,
398        }
399    }
400}
401
402impl Drop for WriterFetchObservation {
403    fn drop(&mut self) {
404        if let Some(metrics) = &self.metrics {
405            metrics.schema_unresolved_records.dec();
406            metrics
407                .schema_fetch_microseconds
408                .inc_by(u64::try_from(self.started.elapsed().as_micros()).unwrap_or(u64::MAX));
409            if !self.succeeded {
410                metrics.schema_fetch_failures.inc();
411            }
412        }
413    }
414}
laminar_connectors/kafka/avro/mod.rs

laminar_connectors/kafka/avro/
mod.rs
