dataprof_core/
source.rs

1/// Supported file formats for data profiling
2#[derive(
3    Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize, schemars::JsonSchema,
4)]
5#[serde(rename_all = "lowercase")]
6pub enum FileFormat {
7    Csv,
8    Json,
9    Jsonl,
10    Parquet,
11    #[serde(untagged)]
12    Unknown(String),
13}
14
15impl std::fmt::Display for FileFormat {
16    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
17        match self {
18            Self::Csv => write!(f, "csv"),
19            Self::Json => write!(f, "json"),
20            Self::Jsonl => write!(f, "jsonl"),
21            Self::Parquet => write!(f, "parquet"),
22            Self::Unknown(s) => write!(f, "{}", s),
23        }
24    }
25}
26
27/// How JSON/JSONL parsing reacts to a malformed record.
28///
29/// Shared across the file, in-memory, and async byte input paths so callers get
30/// the same recovery behavior regardless of transport.
31#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, serde::Serialize, serde::Deserialize)]
32#[serde(rename_all = "lowercase")]
33pub enum JsonErrorPolicy {
34    /// Skip the malformed record, count it, and keep scanning. The profile is
35    /// reported as partial via `execution.error_count`.
36    #[default]
37    Skip,
38    /// Abort on the first malformed record with a `JsonParsingError` carrying
39    /// line/column context (never the record contents).
40    Strict,
41}
42
43/// Supported query engines for SQL-based profiling
44#[derive(
45    Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize, schemars::JsonSchema,
46)]
47#[serde(rename_all = "lowercase")]
48pub enum QueryEngine {
49    Postgres,
50    MySql,
51    Sqlite,
52    Snowflake,
53    BigQuery,
54    #[serde(untagged)]
55    Custom(String),
56}
57
58impl std::fmt::Display for QueryEngine {
59    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
60        match self {
61            Self::Postgres => write!(f, "postgres"),
62            Self::MySql => write!(f, "mysql"),
63            Self::Sqlite => write!(f, "sqlite"),
64            Self::Snowflake => write!(f, "snowflake"),
65            Self::BigQuery => write!(f, "bigquery"),
66            Self::Custom(s) => write!(f, "{}", s),
67        }
68    }
69}
70
71/// Source library for in-memory DataFrame profiling
72#[derive(
73    Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize, schemars::JsonSchema,
74)]
75#[serde(rename_all = "lowercase")]
76pub enum DataFrameLibrary {
77    Pandas,
78    Polars,
79    PyArrow,
80    #[serde(untagged)]
81    Custom(String),
82}
83
84impl std::fmt::Display for DataFrameLibrary {
85    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
86        match self {
87            Self::Pandas => write!(f, "pandas"),
88            Self::Polars => write!(f, "polars"),
89            Self::PyArrow => write!(f, "pyarrow"),
90            Self::Custom(s) => write!(f, "{}", s),
91        }
92    }
93}
94
95/// Supported stream source systems
96#[derive(
97    Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize, schemars::JsonSchema,
98)]
99#[serde(rename_all = "lowercase")]
100pub enum StreamSourceSystem {
101    Kafka,
102    Kinesis,
103    Pulsar,
104    Http,
105    WebSocket,
106    #[serde(rename = "object_store")]
107    ObjectStore,
108    #[serde(rename = "message_queue")]
109    MessageQueue,
110    Database,
111    #[serde(untagged)]
112    Custom(String),
113}
114
115impl std::fmt::Display for StreamSourceSystem {
116    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
117        match self {
118            Self::Kafka => write!(f, "kafka"),
119            Self::Kinesis => write!(f, "kinesis"),
120            Self::Pulsar => write!(f, "pulsar"),
121            Self::Http => write!(f, "http"),
122            Self::WebSocket => write!(f, "websocket"),
123            Self::ObjectStore => write!(f, "object_store"),
124            Self::MessageQueue => write!(f, "message_queue"),
125            Self::Database => write!(f, "database"),
126            Self::Custom(s) => write!(f, "{}", s),
127        }
128    }
129}
130
131/// Metadata specific to Parquet files
132#[derive(Debug, Clone, serde::Serialize, serde::Deserialize, schemars::JsonSchema)]
133pub struct ParquetMetadata {
134    /// Number of row groups in the Parquet file
135    pub num_row_groups: usize,
136    /// Compression codec used (e.g., "SNAPPY", "GZIP", "ZSTD", "UNCOMPRESSED")
137    pub compression: String,
138    /// Parquet file version (e.g., "1.0", "2.0")
139    pub version: i32,
140    /// Arrow schema as string representation
141    pub schema_summary: String,
142    /// Total compressed size in bytes
143    pub compressed_size_bytes: u64,
144    /// Estimated uncompressed size if available
145    pub uncompressed_size_bytes: Option<u64>,
146}
147
148/// Source-agnostic data source metadata.
149#[derive(Debug, Clone, serde::Serialize, serde::Deserialize, schemars::JsonSchema)]
150#[serde(tag = "type", rename_all = "snake_case")]
151pub enum DataSource {
152    /// File-based data source (CSV, JSON, Parquet, etc.)
153    File {
154        /// Absolute or relative path to the file
155        path: String,
156        /// Detected or specified file format
157        format: FileFormat,
158        /// File size in bytes
159        size_bytes: u64,
160        /// Last modification timestamp (ISO 8601 / RFC 3339)
161        #[serde(skip_serializing_if = "Option::is_none")]
162        modified_at: Option<String>,
163        /// Parquet-specific metadata (only present for Parquet files)
164        #[serde(skip_serializing_if = "Option::is_none")]
165        parquet_metadata: Option<ParquetMetadata>,
166    },
167    /// SQL query-based data source
168    Query {
169        /// Database engine used for the query
170        engine: QueryEngine,
171        /// SQL statement executed
172        statement: String,
173        /// Target database name (if applicable)
174        #[serde(skip_serializing_if = "Option::is_none")]
175        database: Option<String>,
176        /// Unique execution identifier for tracing
177        #[serde(skip_serializing_if = "Option::is_none")]
178        execution_id: Option<String>,
179    },
180    /// In-memory DataFrame source (pandas/polars via PyCapsule)
181    #[serde(rename = "dataframe")]
182    DataFrame {
183        /// User-provided name for identification
184        name: String,
185        /// Source library (pandas, polars, pyarrow)
186        source_library: DataFrameLibrary,
187        /// Number of rows at profiling time
188        row_count: usize,
189        /// Number of columns
190        column_count: usize,
191        /// Memory usage in bytes (if available)
192        #[serde(skip_serializing_if = "Option::is_none")]
193        memory_bytes: Option<u64>,
194    },
195    /// In-memory byte buffer profiled without a path behind it. The buffer may
196    /// be text (CSV/JSON/JSONL) or binary (Parquet); `format` says which.
197    /// Distinguished from [`DataSource::DataFrame`]: the caller handed over
198    /// bytes, not a schema-bearing dataframe, so provenance must say so.
199    Bytes {
200        /// Human-readable name for identification
201        name: String,
202        /// Format the buffer was parsed as
203        format: FileFormat,
204        /// Length of the buffer handed over, not of the values decoded from it
205        size_bytes: u64,
206    },
207    /// Streaming data source
208    Stream {
209        /// Stream identifier (e.g., Kafka topic, Kinesis stream name)
210        topic: String,
211        /// Batch identifier for ordering and deduplication
212        batch_id: String,
213        /// Partition for parallel processing (optional)
214        #[serde(skip_serializing_if = "Option::is_none")]
215        partition: Option<u32>,
216        /// Consumer group for Kafka-style coordination (optional)
217        #[serde(skip_serializing_if = "Option::is_none")]
218        consumer_group: Option<String>,
219        /// Source system identifier (kafka, kinesis, pulsar, http, etc.)
220        source_system: StreamSourceSystem,
221        /// Session ID for multi-tenant scenarios
222        #[serde(skip_serializing_if = "Option::is_none")]
223        session_id: Option<String>,
224        /// Timestamp of first record in batch (ISO 8601)
225        #[serde(skip_serializing_if = "Option::is_none")]
226        first_record_at: Option<String>,
227        /// Timestamp of last record in batch (ISO 8601)
228        #[serde(skip_serializing_if = "Option::is_none")]
229        last_record_at: Option<String>,
230    },
231}
232
233impl DataSource {
234    /// Get a human-readable identifier for this data source.
235    pub fn identifier(&self) -> String {
236        match self {
237            Self::File { path, .. } => path.clone(),
238            Self::Query {
239                engine, statement, ..
240            } => {
241                let truncated = if statement.chars().count() > 50 {
242                    let mut prefix: String = statement.chars().take(47).collect();
243                    prefix.push_str("...");
244                    prefix
245                } else {
246                    statement.clone()
247                };
248                format!("{}: {}", engine, truncated)
249            }
250            Self::DataFrame {
251                name,
252                source_library,
253                ..
254            } => format!("{}[{}]", source_library, name),
255            Self::Bytes { name, .. } => format!("bytes[{name}]"),
256            Self::Stream {
257                source_system,
258                topic,
259                batch_id,
260                ..
261            } => format!("{}[{}]-batch:{}", source_system, topic, batch_id),
262        }
263    }
264
265    /// Get file size in megabytes if this is a file-based source or dataframe.
266    pub fn size_mb(&self) -> Option<f64> {
267        match self {
268            Self::File { size_bytes, .. } => Some(*size_bytes as f64 / 1_048_576.0),
269            Self::DataFrame { memory_bytes, .. } => memory_bytes.map(|b| b as f64 / 1_048_576.0),
270            Self::Bytes { size_bytes, .. } => Some(*size_bytes as f64 / 1_048_576.0),
271            Self::Query { .. } | Self::Stream { .. } => None,
272        }
273    }
274
275    /// Check if this is a file-based source.
276    pub fn is_file(&self) -> bool {
277        matches!(self, Self::File { .. })
278    }
279
280    /// Check if this is a query-based source.
281    pub fn is_query(&self) -> bool {
282        matches!(self, Self::Query { .. })
283    }
284
285    /// Check if this is a DataFrame-based source.
286    pub fn is_dataframe(&self) -> bool {
287        matches!(self, Self::DataFrame { .. })
288    }
289
290    /// Check if this is a Stream-based source.
291    pub fn is_stream(&self) -> bool {
292        matches!(self, Self::Stream { .. })
293    }
294
295    /// Check if this is an in-memory byte-buffer source.
296    pub fn is_bytes(&self) -> bool {
297        matches!(self, Self::Bytes { .. })
298    }
299
300    /// Get the file path if this is a file-based source.
301    pub fn file_path(&self) -> Option<&str> {
302        match self {
303            Self::File { path, .. } => Some(path),
304            _ => None,
305        }
306    }
307
308    /// Get the stream topic if this is a stream-based source.
309    pub fn stream_topic(&self) -> Option<&str> {
310        match self {
311            Self::Stream { topic, .. } => Some(topic),
312            _ => None,
313        }
314    }
315
316    /// Get the batch ID if this is a stream-based source.
317    pub fn batch_id(&self) -> Option<&str> {
318        match self {
319            Self::Stream { batch_id, .. } => Some(batch_id),
320            _ => None,
321        }
322    }
323}
324
325#[cfg(test)]
326mod tests {
327    use super::*;
328
329    #[test]
330    fn test_data_source_file_identifier() {
331        let ds = DataSource::File {
332            path: "/path/to/data.csv".to_string(),
333            format: FileFormat::Csv,
334            size_bytes: 0,
335            modified_at: None,
336            parquet_metadata: None,
337        };
338
339        assert_eq!(ds.identifier(), "/path/to/data.csv");
340        assert!(ds.is_file());
341        assert!(!ds.is_query());
342        assert!(!ds.is_dataframe());
343        assert!(!ds.is_stream());
344    }
345
346    #[test]
347    fn test_data_source_bytes_identifier_and_helpers() {
348        let ds = DataSource::Bytes {
349            name: "csv_bytes".to_string(),
350            format: FileFormat::Csv,
351            size_bytes: 42,
352        };
353
354        assert_eq!(ds.identifier(), "bytes[csv_bytes]");
355        assert!(ds.is_bytes());
356        assert!(!ds.is_file());
357        assert!(!ds.is_dataframe());
358        assert!(!ds.is_stream());
359        assert_eq!(ds.size_mb(), Some(42.0 / 1_048_576.0));
360    }
361
362    #[test]
363    fn test_data_source_stream_identifier_and_helpers() {
364        let ds = DataSource::Stream {
365            topic: "events".to_string(),
366            batch_id: "b1".to_string(),
367            partition: Some(0),
368            consumer_group: None,
369            source_system: StreamSourceSystem::Kafka,
370            session_id: None,
371            first_record_at: None,
372            last_record_at: None,
373        };
374
375        assert_eq!(ds.identifier(), "kafka[events]-batch:b1");
376        assert!(ds.is_stream());
377        assert_eq!(ds.stream_topic(), Some("events"));
378        assert_eq!(ds.batch_id(), Some("b1"));
379        assert!(!ds.is_file());
380        assert!(!ds.is_query());
381        assert!(ds.size_mb().is_none());
382    }
383
384    #[test]
385    fn test_stream_json_serialization() {
386        let ds = DataSource::Stream {
387            topic: "sensor-data".to_string(),
388            batch_id: "batch-789".to_string(),
389            partition: Some(2),
390            consumer_group: Some("processing-group".to_string()),
391            source_system: StreamSourceSystem::Kinesis,
392            session_id: Some("session-1".to_string()),
393            first_record_at: Some("2023-01-01T10:00:00Z".to_string()),
394            last_record_at: Some("2023-01-01T10:05:00Z".to_string()),
395        };
396
397        let json = serde_json::to_string(&ds).unwrap();
398        assert!(json.contains(r#""type":"stream""#));
399        assert!(json.contains(r#""source_system":"kinesis""#));
400        assert!(json.contains(r#""topic":"sensor-data""#));
401
402        let deserialized: DataSource = serde_json::from_str(&json).unwrap();
403        assert!(deserialized.is_stream());
404        assert_eq!(deserialized.stream_topic(), Some("sensor-data"));
405    }
406
407    #[test]
408    fn test_stream_source_system_serialization_names() {
409        let object_store = serde_json::to_string(&StreamSourceSystem::ObjectStore).unwrap();
410        let message_queue = serde_json::to_string(&StreamSourceSystem::MessageQueue).unwrap();
411        let database = serde_json::to_string(&StreamSourceSystem::Database).unwrap();
412
413        assert_eq!(object_store, r#""object_store""#);
414        assert_eq!(message_queue, r#""message_queue""#);
415        assert_eq!(database, r#""database""#);
416
417        let object_store: StreamSourceSystem = serde_json::from_str(r#""object_store""#).unwrap();
418        let message_queue: StreamSourceSystem = serde_json::from_str(r#""message_queue""#).unwrap();
419        let database: StreamSourceSystem = serde_json::from_str(r#""database""#).unwrap();
420
421        assert_eq!(object_store, StreamSourceSystem::ObjectStore);
422        assert_eq!(message_queue, StreamSourceSystem::MessageQueue);
423        assert_eq!(database, StreamSourceSystem::Database);
424    }
425}
dataprof_core/source.rs

dataprof_core/
source.rs
