1#![allow(clippy::disallowed_types)] use std::collections::HashMap;
5
6use super::error::SchemaResult;
7use super::types::RawRecord;
8use arrow_array::RecordBatch;
9use arrow_schema::{DataType, SchemaRef};
10
11pub trait FormatDecoder: Send + Sync {
19 fn output_schema(&self) -> SchemaRef;
21
22 fn decode_batch(&self, records: &[RawRecord]) -> SchemaResult<RecordBatch>;
29
30 fn decode_one(&self, record: &RawRecord) -> SchemaResult<RecordBatch> {
40 self.decode_batch(std::slice::from_ref(record))
41 }
42
43 fn format_name(&self) -> &'static str;
45}
46
47pub trait FormatEncoder: Send + Sync {
51 fn input_schema(&self) -> SchemaRef;
53
54 fn encode_batch(&self, batch: &RecordBatch) -> SchemaResult<Vec<Vec<u8>>>;
62
63 fn format_name(&self) -> &'static str;
65}
66
67#[derive(Debug, Clone)]
71pub struct InferenceConfig {
72 pub format: String,
74
75 pub number_inference: NumberInference,
77
78 pub array_inference: ArrayInference,
80
81 pub max_samples: usize,
83
84 pub min_confidence: f64,
86
87 pub type_hints: HashMap<String, DataType>,
89
90 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 #[must_use]
111 pub fn new(format: impl Into<String>) -> Self {
112 Self {
113 format: format.into(),
114 ..Self::default()
115 }
116 }
117
118 #[must_use]
120 pub fn with_min_confidence(mut self, confidence: f64) -> Self {
121 self.min_confidence = confidence;
122 self
123 }
124
125 #[must_use]
127 pub fn with_max_samples(mut self, n: usize) -> Self {
128 self.max_samples = n;
129 self
130 }
131
132 #[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 #[must_use]
141 pub fn with_empty_as_null(mut self) -> Self {
142 self.empty_as_null = true;
143 self
144 }
145}
146
147#[derive(Debug, Clone, Copy, PartialEq, Eq)]
149pub enum NumberInference {
150 PreferSmallest,
152 PreferLarger,
154}
155
156#[derive(Debug, Clone, Copy, PartialEq, Eq)]
158pub enum ArrayInference {
159 Utf8,
161 NativeArrow,
163}
164
165#[derive(Debug, Clone)]
167pub struct InferredSchema {
168 pub schema: SchemaRef,
170
171 pub confidence: f64,
173
174 pub sample_count: usize,
176
177 pub field_details: Vec<FieldInferenceDetail>,
179
180 pub warnings: Vec<InferenceWarning>,
182}
183
184#[derive(Debug, Clone)]
186pub struct FieldInferenceDetail {
187 pub field_name: String,
189
190 pub inferred_type: DataType,
192
193 pub confidence: f64,
195
196 pub non_null_count: usize,
198
199 pub total_count: usize,
201
202 pub hint_applied: bool,
204}
205
206#[derive(Debug, Clone)]
208pub struct InferenceWarning {
209 pub field: Option<String>,
211
212 pub message: String,
214
215 pub severity: WarningSeverity,
217}
218
219pub use laminar_core::error_codes::WarningSeverity;
220
221#[derive(Debug, Clone, Copy, PartialEq, Eq)]
223pub enum CompatibilityMode {
224 None,
226 Backward,
228 Forward,
230 Full,
232 BackwardTransitive,
234 ForwardTransitive,
236 FullTransitive,
238}
239
240#[derive(Debug, Clone, PartialEq)]
242pub enum SchemaChange {
243 ColumnAdded {
245 name: String,
247 data_type: DataType,
249 nullable: bool,
251 },
252
253 ColumnRemoved {
255 name: String,
257 },
258
259 TypeChanged {
261 name: String,
263 old_type: DataType,
265 new_type: DataType,
267 },
268
269 NullabilityChanged {
271 name: String,
273 was_nullable: bool,
275 now_nullable: bool,
277 },
278
279 ColumnRenamed {
281 old_name: String,
283 new_name: String,
285 },
286}
287
288#[derive(Debug, Clone, PartialEq, Eq)]
290pub enum EvolutionVerdict {
291 Compatible,
293
294 RequiresMigration,
296
297 Incompatible(String),
299}
300
301#[derive(Debug, Clone)]
303pub struct ColumnProjection {
304 pub mappings: Vec<Option<usize>>,
308
309 pub target_schema: SchemaRef,
311}
312
313const _: () = {
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); 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}