dataprof_core/
source.rs

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

dataprof_core/
source.rs
