Skip to main content

laminar_connectors/schema/
traits.rs

1//! Capability traits for the connector schema framework.
2#![allow(clippy::disallowed_types)] // cold path: schema management
3
4use std::collections::HashMap;
5
6use super::error::SchemaResult;
7use super::types::RawRecord;
8use arrow_array::RecordBatch;
9use arrow_schema::{DataType, SchemaRef};
10
11// ── FormatDecoder ──────────────────────────────────────────────────
12
13/// Decodes raw bytes into Arrow `RecordBatch`es.
14///
15/// Unlike [`RecordDeserializer`](crate::serde::RecordDeserializer) which
16/// takes `&[u8]` slices, `FormatDecoder` works with [`RawRecord`]s that
17/// carry metadata, headers, and timestamps alongside the payload.
18pub trait FormatDecoder: Send + Sync {
19    /// Returns the Arrow schema produced by this decoder.
20    fn output_schema(&self) -> SchemaRef;
21
22    /// Decodes a batch of raw records into an Arrow `RecordBatch`.
23    ///
24    /// # Errors
25    ///
26    /// Returns [`SchemaError::DecodeError`](super::error::SchemaError::DecodeError)
27    /// if the input cannot be parsed.
28    fn decode_batch(&self, records: &[RawRecord]) -> SchemaResult<RecordBatch>;
29
30    /// Decodes a single raw record into an Arrow `RecordBatch` with one row.
31    ///
32    /// Default implementation delegates to [`decode_batch`](Self::decode_batch)
33    /// with a single-element slice.
34    ///
35    /// # Errors
36    ///
37    /// Returns [`SchemaError::DecodeError`](super::error::SchemaError::DecodeError)
38    /// if the input cannot be parsed.
39    fn decode_one(&self, record: &RawRecord) -> SchemaResult<RecordBatch> {
40        self.decode_batch(std::slice::from_ref(record))
41    }
42
43    /// Returns the name of the format this decoder handles (e.g., `"json"`).
44    fn format_name(&self) -> &'static str;
45}
46
47// ── FormatEncoder ──────────────────────────────────────────────────
48
49/// Encodes Arrow `RecordBatch`es into raw bytes.
50pub trait FormatEncoder: Send + Sync {
51    /// Returns the expected input schema.
52    fn input_schema(&self) -> SchemaRef;
53
54    /// Encodes a `RecordBatch` into a vector of byte records.
55    ///
56    /// Each element in the returned vector represents one serialized record.
57    ///
58    /// # Errors
59    ///
60    /// Returns a schema error if encoding fails.
61    fn encode_batch(&self, batch: &RecordBatch) -> SchemaResult<Vec<Vec<u8>>>;
62
63    /// Returns the name of the format this encoder produces (e.g., `"json"`).
64    fn format_name(&self) -> &'static str;
65}
66
67// ── Inference types ────────────────────────────────────────────────
68
69/// Configuration for schema inference.
70#[derive(Debug, Clone)]
71pub struct InferenceConfig {
72    /// Data format to use for inference.
73    pub format: String,
74
75    /// How to handle number type inference.
76    pub number_inference: NumberInference,
77
78    /// How to handle array type inference.
79    pub array_inference: ArrayInference,
80
81    /// Maximum number of samples to use.
82    pub max_samples: usize,
83
84    /// Minimum confidence threshold (0.0–1.0) for accepting an inferred type.
85    pub min_confidence: f64,
86
87    /// Type hints for specific fields.
88    pub type_hints: HashMap<String, DataType>,
89
90    /// Whether to treat empty strings as nulls.
91    pub empty_as_null: bool,
92}
93
94impl Default for InferenceConfig {
95    fn default() -> Self {
96        Self {
97            format: "json".to_string(),
98            number_inference: NumberInference::PreferLarger,
99            array_inference: ArrayInference::Utf8,
100            max_samples: 1000,
101            min_confidence: 0.8,
102            type_hints: HashMap::new(),
103            empty_as_null: false,
104        }
105    }
106}
107
108impl InferenceConfig {
109    /// Creates a new inference config for the given format.
110    #[must_use]
111    pub fn new(format: impl Into<String>) -> Self {
112        Self {
113            format: format.into(),
114            ..Self::default()
115        }
116    }
117
118    /// Sets the minimum confidence threshold.
119    #[must_use]
120    pub fn with_min_confidence(mut self, confidence: f64) -> Self {
121        self.min_confidence = confidence;
122        self
123    }
124
125    /// Sets the maximum number of samples.
126    #[must_use]
127    pub fn with_max_samples(mut self, n: usize) -> Self {
128        self.max_samples = n;
129        self
130    }
131
132    /// Adds a type hint for a specific field.
133    #[must_use]
134    pub fn with_type_hint(mut self, field: impl Into<String>, data_type: DataType) -> Self {
135        self.type_hints.insert(field.into(), data_type);
136        self
137    }
138
139    /// Enables treating empty strings as nulls.
140    #[must_use]
141    pub fn with_empty_as_null(mut self) -> Self {
142        self.empty_as_null = true;
143        self
144    }
145}
146
147/// How to infer numeric types.
148#[derive(Debug, Clone, Copy, PartialEq, Eq)]
149pub enum NumberInference {
150    /// Prefer the smallest type that fits (i32 before i64).
151    PreferSmallest,
152    /// Prefer larger types (always i64, always f64).
153    PreferLarger,
154}
155
156/// How to infer array/object types in JSON.
157#[derive(Debug, Clone, Copy, PartialEq, Eq)]
158pub enum ArrayInference {
159    /// Store arrays/objects as JSON-encoded Utf8 strings.
160    Utf8,
161    /// Attempt to infer Arrow List / Struct types.
162    NativeArrow,
163}
164
165/// The result of schema inference.
166#[derive(Debug, Clone)]
167pub struct InferredSchema {
168    /// The inferred Arrow schema.
169    pub schema: SchemaRef,
170
171    /// Overall confidence score (0.0–1.0).
172    pub confidence: f64,
173
174    /// Number of samples that were analyzed.
175    pub sample_count: usize,
176
177    /// Per-field inference details.
178    pub field_details: Vec<FieldInferenceDetail>,
179
180    /// Warnings generated during inference.
181    pub warnings: Vec<InferenceWarning>,
182}
183
184/// Per-field detail from inference.
185#[derive(Debug, Clone)]
186pub struct FieldInferenceDetail {
187    /// The field name.
188    pub field_name: String,
189
190    /// The inferred Arrow data type.
191    pub inferred_type: DataType,
192
193    /// Confidence for this specific field (0.0–1.0).
194    pub confidence: f64,
195
196    /// Number of non-null samples seen for this field.
197    pub non_null_count: usize,
198
199    /// Total number of samples that included this field.
200    pub total_count: usize,
201
202    /// Whether a type hint was applied.
203    pub hint_applied: bool,
204}
205
206/// A warning generated during inference.
207#[derive(Debug, Clone)]
208pub struct InferenceWarning {
209    /// The field this warning relates to, if any.
210    pub field: Option<String>,
211
212    /// Warning message.
213    pub message: String,
214
215    /// Severity level.
216    pub severity: WarningSeverity,
217}
218
219pub use laminar_core::error_codes::WarningSeverity;
220
221/// Compatibility mode for schema evolution.
222#[derive(Debug, Clone, Copy, PartialEq, Eq)]
223pub enum CompatibilityMode {
224    /// No compatibility checks.
225    None,
226    /// New schema can read old data.
227    Backward,
228    /// Old schema can read new data.
229    Forward,
230    /// Both backward and forward compatible.
231    Full,
232    /// Backward compatible with all previous versions.
233    BackwardTransitive,
234    /// Forward compatible with all previous versions.
235    ForwardTransitive,
236    /// Fully compatible with all previous versions.
237    FullTransitive,
238}
239
240/// A single schema change detected by schema diffing.
241#[derive(Debug, Clone, PartialEq)]
242pub enum SchemaChange {
243    /// A new column was added.
244    ColumnAdded {
245        /// Column name.
246        name: String,
247        /// The new column's data type.
248        data_type: DataType,
249        /// Whether the column is nullable.
250        nullable: bool,
251    },
252
253    /// An existing column was removed.
254    ColumnRemoved {
255        /// Column name.
256        name: String,
257    },
258
259    /// A column's data type changed.
260    TypeChanged {
261        /// Column name.
262        name: String,
263        /// Previous data type.
264        old_type: DataType,
265        /// New data type.
266        new_type: DataType,
267    },
268
269    /// A column's nullability changed.
270    NullabilityChanged {
271        /// Column name.
272        name: String,
273        /// Previous nullable flag.
274        was_nullable: bool,
275        /// New nullable flag.
276        now_nullable: bool,
277    },
278
279    /// A column was renamed.
280    ColumnRenamed {
281        /// Previous name.
282        old_name: String,
283        /// New name.
284        new_name: String,
285    },
286}
287
288/// The result of evaluating a set of schema changes.
289#[derive(Debug, Clone, PartialEq, Eq)]
290pub enum EvolutionVerdict {
291    /// All changes are compatible — evolution can proceed.
292    Compatible,
293
294    /// Changes require data migration but are feasible.
295    RequiresMigration,
296
297    /// Changes are incompatible — evolution is rejected.
298    Incompatible(String),
299}
300
301/// Describes how to project columns from the old schema to the new schema.
302#[derive(Debug, Clone)]
303pub struct ColumnProjection {
304    /// For each column in the new schema, the index in the old schema
305    /// (or `None` if the column is newly added and should be filled
306    /// with the default/null).
307    pub mappings: Vec<Option<usize>>,
308
309    /// The resulting schema after projection.
310    pub target_schema: SchemaRef,
311}
312
313// ── Object-safety assertions ───────────────────────────────────────
314
315// Compile-time checks that the codec traits are object-safe.
316const _: () = {
317    fn _assert_format_decoder_object_safe(_: &dyn FormatDecoder) {}
318    fn _assert_format_encoder_object_safe(_: &dyn FormatEncoder) {}
319};
320
321#[cfg(test)]
322mod tests {
323    use super::*;
324    use std::sync::Arc;
325
326    use arrow_schema::{Field, Schema};
327
328    #[test]
329    fn test_inference_config_defaults() {
330        let cfg = InferenceConfig::default();
331        assert_eq!(cfg.format, "json");
332        assert_eq!(cfg.max_samples, 1000);
333        assert!((cfg.min_confidence - 0.8).abs() < f64::EPSILON);
334        assert!(cfg.type_hints.is_empty());
335    }
336
337    #[test]
338    fn test_inference_config_builder() {
339        let cfg = InferenceConfig::new("csv")
340            .with_min_confidence(0.9)
341            .with_max_samples(500)
342            .with_type_hint("id", DataType::Int32)
343            .with_empty_as_null();
344
345        assert_eq!(cfg.format, "csv");
346        assert!((cfg.min_confidence - 0.9).abs() < f64::EPSILON);
347        assert_eq!(cfg.max_samples, 500);
348        assert_eq!(cfg.type_hints.get("id"), Some(&DataType::Int32));
349        assert!(cfg.empty_as_null);
350    }
351
352    #[test]
353    fn test_inferred_schema() {
354        let schema = Arc::new(Schema::new(vec![
355            Field::new("id", DataType::Int64, false),
356            Field::new("name", DataType::Utf8, true),
357        ]));
358
359        let inferred = InferredSchema {
360            schema: schema.clone(),
361            confidence: 0.95,
362            sample_count: 100,
363            field_details: vec![
364                FieldInferenceDetail {
365                    field_name: "id".into(),
366                    inferred_type: DataType::Int64,
367                    confidence: 1.0,
368                    non_null_count: 100,
369                    total_count: 100,
370                    hint_applied: false,
371                },
372                FieldInferenceDetail {
373                    field_name: "name".into(),
374                    inferred_type: DataType::Utf8,
375                    confidence: 0.9,
376                    non_null_count: 90,
377                    total_count: 100,
378                    hint_applied: false,
379                },
380            ],
381            warnings: vec![],
382        };
383
384        assert_eq!(inferred.schema.fields().len(), 2);
385        assert!((inferred.confidence - 0.95).abs() < f64::EPSILON);
386        assert_eq!(inferred.field_details.len(), 2);
387    }
388
389    #[test]
390    fn test_schema_change_variants() {
391        let changes = [
392            SchemaChange::ColumnAdded {
393                name: "email".into(),
394                data_type: DataType::Utf8,
395                nullable: true,
396            },
397            SchemaChange::ColumnRemoved {
398                name: "legacy".into(),
399            },
400            SchemaChange::TypeChanged {
401                name: "age".into(),
402                old_type: DataType::Int32,
403                new_type: DataType::Int64,
404            },
405            SchemaChange::NullabilityChanged {
406                name: "name".into(),
407                was_nullable: false,
408                now_nullable: true,
409            },
410            SchemaChange::ColumnRenamed {
411                old_name: "fname".into(),
412                new_name: "first_name".into(),
413            },
414        ];
415        assert_eq!(changes.len(), 5);
416    }
417
418    #[test]
419    fn test_evolution_verdict() {
420        assert_eq!(EvolutionVerdict::Compatible, EvolutionVerdict::Compatible);
421        assert_ne!(
422            EvolutionVerdict::Compatible,
423            EvolutionVerdict::RequiresMigration
424        );
425    }
426
427    #[test]
428    fn test_column_projection() {
429        let schema = Arc::new(Schema::new(vec![
430            Field::new("a", DataType::Int64, false),
431            Field::new("b", DataType::Utf8, true),
432            Field::new("c", DataType::Float64, false),
433        ]));
434
435        let proj = ColumnProjection {
436            mappings: vec![Some(0), None, Some(1)],
437            target_schema: schema,
438        };
439
440        assert_eq!(proj.mappings.len(), 3);
441        assert_eq!(proj.mappings[0], Some(0));
442        assert_eq!(proj.mappings[1], None); // new column
443        assert_eq!(proj.mappings[2], Some(1));
444    }
445
446    #[test]
447    fn test_warning_severity() {
448        let w = InferenceWarning {
449            field: Some("price".into()),
450            message: "mixed int/float".into(),
451            severity: WarningSeverity::Warning,
452        };
453        assert_eq!(w.severity, WarningSeverity::Warning);
454        assert_eq!(w.field.as_deref(), Some("price"));
455    }
456}