Skip to main content

rbt/scan/
mod.rs

1//! `rbt::scan`: Multi-format bronze lake scanners producing Arrow `RecordBatch`es.
2//!
3//! Formats: JSONL (jshift), JSON, Parquet, CSV, Arrow IPC (file/stream), log, txt, TOML.
4
5use crate::core::frontmatter::{SourceFormat, StagingFrontmatter};
6use crate::core::paths::{resolve_project_path, validate_glob_patterns, PathGlobSet};
7use crate::core::project::{ScanConfig, DEFAULT_PROTOBUF_MAX_PAYLOAD_BYTES};
8use anyhow::{anyhow, bail, Context, Result};
9use arrow::array::{ArrayRef, BinaryBuilder, Int64Builder, StringBuilder};
10use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
11use arrow::record_batch::RecordBatch;
12use std::collections::HashMap;
13use std::fs::File;
14use std::io::{BufRead, BufReader, Read};
15use std::path::{Path, PathBuf};
16use std::sync::Arc;
17
18/// Scan request derived from staging frontmatter + project root.
19#[derive(Debug, Clone)]
20pub struct ScanRequest {
21    pub project_dir: PathBuf,
22    pub scan_path: String,
23    pub format: SourceFormat,
24    /// jshift / projection paths (JSONL selective extract).
25    pub paths: Vec<String>,
26    /// TOML array-of-tables key (optional).
27    pub toml_rows_key: Option<String>,
28    /// Hive-style partition keys to inject as Utf8 columns from the file path.
29    pub partition_by: Vec<String>,
30    /// Keep only files whose hive path matches these partition values.
31    pub require_partitions: std::collections::HashMap<String, String>,
32    /// Filename / relative globs under scan_path (OR). Empty = all format matches.
33    pub path_glob: Vec<String>,
34    /// Inject `_source_path` column (absolute path of the bronze file).
35    pub inject_source_path: bool,
36    /// Named roots for `$name` expansion in `scan_path` (from project config).
37    pub roots: HashMap<String, String>,
38    /// Max bytes for one opaque protobuf file (from `scan.protobuf_max_payload_bytes`).
39    pub protobuf_max_payload_bytes: u64,
40}
41
42impl ScanRequest {
43    pub fn from_frontmatter(
44        project_dir: impl AsRef<Path>,
45        fm: &StagingFrontmatter,
46    ) -> Result<Self> {
47        Self::from_frontmatter_with_roots(project_dir, fm, HashMap::new())
48    }
49
50    pub fn from_frontmatter_with_roots(
51        project_dir: impl AsRef<Path>,
52        fm: &StagingFrontmatter,
53        roots: HashMap<String, String>,
54    ) -> Result<Self> {
55        Self::from_frontmatter_with_config(project_dir, fm, roots, &ScanConfig::default())
56    }
57
58    pub fn from_frontmatter_with_config(
59        project_dir: impl AsRef<Path>,
60        fm: &StagingFrontmatter,
61        roots: HashMap<String, String>,
62        scan_cfg: &ScanConfig,
63    ) -> Result<Self> {
64        let scan_path = fm
65            .scan_path
66            .as_ref()
67            .context(
68                "E_RBT_BRONZE_SCAN_PATH_MISSING: frontmatter.scan_path is required for bronze scan",
69            )?
70            .clone();
71        let format = fm.resolve_format().with_context(|| {
72            format!("E_RBT_BRONZE_FORMAT: cannot resolve source_format for scan_path '{scan_path}'")
73        })?;
74        let path_glob = fm.path_glob.clone().unwrap_or_default();
75        validate_glob_patterns(&path_glob).with_context(|| {
76            format!("E_RBT_PATH_GLOB_INVALID: bad path_glob on scan_path '{scan_path}'")
77        })?;
78        Ok(Self {
79            project_dir: project_dir.as_ref().to_path_buf(),
80            scan_path,
81            format,
82            paths: fm.paths.clone().unwrap_or_default(),
83            toml_rows_key: fm.toml_rows_key.clone(),
84            partition_by: fm.partition_by.clone().unwrap_or_default(),
85            require_partitions: fm.require_partitions.clone().unwrap_or_default(),
86            path_glob,
87            inject_source_path: fm.inject_source_path.unwrap_or(false),
88            roots,
89            protobuf_max_payload_bytes: scan_cfg.protobuf_max_payload_bytes,
90        })
91    }
92
93    pub fn resolved_path(&self) -> Result<PathBuf> {
94        resolve_project_path(&self.project_dir, &self.scan_path, &self.roots)
95    }
96}
97
98/// Multi-format lake scanner.
99pub struct LakeScanner {
100    pub paths: Vec<String>,
101}
102
103impl LakeScanner {
104    pub fn new(paths: Vec<String>) -> Self {
105        Self { paths }
106    }
107
108    pub fn from_request(req: &ScanRequest) -> Self {
109        Self {
110            paths: req.paths.clone(),
111        }
112    }
113
114    /// Resolve scan root and list bronze files after partition/glob filters.
115    pub fn list_files(&self, req: &ScanRequest) -> Result<(PathBuf, Vec<PathBuf>)> {
116        let root = req.resolved_path().with_context(|| {
117            format!(
118                "E_RBT_BRONZE_PATH: failed resolving scan_path '{}' \
119                 (project_dir={}). Check absolute paths and `roots:` templates.",
120                req.scan_path,
121                req.project_dir.display()
122            )
123        })?;
124        if !root.exists() {
125            bail!(
126                "E_RBT_BRONZE_SCAN_PATH_NOT_FOUND: bronze scan_path does not exist: {} \
127                 (resolved from '{}'). Hint: verify the lake path and `$root` expansion.",
128                root.display(),
129                req.scan_path
130            );
131        }
132
133        let mut files = Vec::new();
134        collect_files_for_format(&root, req.format, &mut files)?;
135        if !req.require_partitions.is_empty() {
136            files.retain(|f| path_matches_require_partitions(f, &root, &req.require_partitions));
137        }
138        if !req.path_glob.is_empty() {
139            let glob_set = PathGlobSet::compile(&req.path_glob)?;
140            files.retain(|f| glob_set.matches(f, &root));
141        }
142        if files.is_empty() {
143            bail!(
144                "E_RBT_BRONZE_SCAN_EMPTY: no {} files under {} after filters \
145                 (require_partitions={:?}, path_glob={:?}). \
146                 Hint: path_glob disables DataFusion listing pushdown and uses the \
147                 scan→MemTable/spill path; check filename patterns and hive partitions.",
148                req.format.as_str(),
149                root.display(),
150                req.require_partitions,
151                req.path_glob
152            );
153        }
154        Ok((root, files))
155    }
156
157    /// Scan using a full [`ScanRequest`] (format-aware). Loads all batches into memory.
158    pub async fn scan(&self, req: &ScanRequest) -> Result<Vec<RecordBatch>> {
159        let (root, files) = self.list_files(req)?;
160        tracing::info!(
161            "Bronze scan: {} file(s) under {} (format={}, globs={:?})",
162            files.len(),
163            root.display(),
164            req.format,
165            req.path_glob
166        );
167
168        let mut batches = Vec::new();
169        for file_path in files {
170            for batch in self.read_enriched(&file_path, &root, req)? {
171                batches.push(batch);
172            }
173        }
174        Ok(batches)
175    }
176
177    /// Stream bronze files into a single Parquet file (atomic publish).
178    ///
179    /// Peak memory ≈ one source file's batches + Parquet encoder — not the full hive tree.
180    /// Used for Arrow IPC multi-file bronze when `scan.spill_arrow_ipc` is true.
181    pub fn scan_spill_to_parquet(
182        &self,
183        req: &ScanRequest,
184        dest: &Path,
185        opts: &crate::materializer::MaterializeWriteOptions,
186    ) -> Result<crate::materializer::StreamWriteStats> {
187        use crate::materializer::{atomic_publish, partial_path_for};
188        use parquet::arrow::ArrowWriter;
189        use parquet::basic::Compression;
190        use parquet::file::properties::WriterProperties;
191        use std::fs::{self, File};
192        use std::io::BufWriter;
193
194        let (root, files) = self.list_files(req)?;
195        tracing::info!(
196            "Bronze spill→parquet: {} file(s) under {} → {} (format={})",
197            files.len(),
198            root.display(),
199            dest.display(),
200            req.format
201        );
202
203        let partial = partial_path_for(dest);
204        if partial.exists() {
205            let _ = fs::remove_file(&partial);
206        }
207        if let Some(parent) = partial.parent() {
208            fs::create_dir_all(parent)?;
209        }
210
211        let props = WriterProperties::builder()
212            .set_max_row_group_row_count(Some(opts.max_row_group_rows.max(1)))
213            .set_compression(Compression::SNAPPY)
214            .build();
215
216        let mut writer: Option<ArrowWriter<BufWriter<File>>> = None;
217        let mut schema: Option<SchemaRef> = None;
218        let mut rows = 0usize;
219        let mut batches_n = 0usize;
220
221        let mut write_loop = || -> Result<()> {
222            for file_path in &files {
223                let file_batches = self.read_enriched(file_path, &root, req).with_context(|| {
224                    format!(
225                        "E_RBT_BRONZE_SPILL: read {} for spill",
226                        file_path.display()
227                    )
228                })?;
229                for batch in file_batches {
230                    if batch.num_rows() == 0 && batch.num_columns() == 0 {
231                        continue;
232                    }
233                    if writer.is_none() {
234                        schema = Some(batch.schema());
235                        let file = File::create(&partial).with_context(|| {
236                            format!(
237                                "E_RBT_BRONZE_SPILL: create partial {}",
238                                partial.display()
239                            )
240                        })?;
241                        let buf = BufWriter::with_capacity(8 * 1024 * 1024, file);
242                        writer = Some(
243                            ArrowWriter::try_new(buf, batch.schema(), Some(props.clone()))
244                                .with_context(|| {
245                                    format!(
246                                        "E_RBT_BRONZE_SPILL: ArrowWriter for {}",
247                                        partial.display()
248                                    )
249                                })?,
250                        );
251                    }
252                    let w = writer.as_mut().unwrap();
253                    w.write(&batch).with_context(|| {
254                        format!(
255                            "E_RBT_BRONZE_SPILL: write batch from {}",
256                            file_path.display()
257                        )
258                    })?;
259                    rows += batch.num_rows();
260                    batches_n += 1;
261                    if w.in_progress_size() >= opts.max_row_group_bytes {
262                        w.flush()?;
263                    }
264                    // batch dropped
265                }
266            }
267            Ok(())
268        };
269
270        if let Err(e) = write_loop() {
271            if let Some(w) = writer.take() {
272                let _ = w.close();
273            }
274            let _ = fs::remove_file(&partial);
275            return Err(e);
276        }
277
278        if let Some(w) = writer.take() {
279            w.close().with_context(|| {
280                format!("E_RBT_BRONZE_SPILL: close writer {}", partial.display())
281            })?;
282            atomic_publish(&partial, dest)?;
283        } else {
284            // Zero rows but known schema is rare; fail if we never saw a batch.
285            let _ = schema;
286            bail!(
287                "E_RBT_BRONZE_SPILL: no batches produced for spill to {} \
288                 ({} candidate files)",
289                dest.display(),
290                files.len()
291            );
292        }
293
294        let bytes_written = std::fs::metadata(dest).map(|m| m.len()).unwrap_or(0);
295        tracing::info!(
296            "Bronze spill complete: {} rows, {} batches, {} bytes → {}",
297            rows,
298            batches_n,
299            bytes_written,
300            dest.display()
301        );
302        Ok(crate::materializer::StreamWriteStats {
303            rows,
304            batches: batches_n,
305            path: dest.to_path_buf(),
306            bytes_written,
307            validation: crate::testing::ValidationResult {
308                total_rows: rows,
309                passed_assertions: 0,
310                failed_assertions: 0,
311                errors: Vec::new(),
312            },
313        })
314    }
315
316    fn read_enriched(
317        &self,
318        file_path: &Path,
319        root: &Path,
320        req: &ScanRequest,
321    ) -> Result<Vec<RecordBatch>> {
322        let file_batches = self
323            .read_file(file_path, req)
324            .with_context(|| format!("Failed reading bronze file {}", file_path.display()))?;
325        let mut out = Vec::with_capacity(file_batches.len());
326        for batch in file_batches {
327            let mut batch = if req.partition_by.is_empty() {
328                batch
329            } else {
330                inject_hive_partitions(batch, file_path, root, &req.partition_by)?
331            };
332            if req.inject_source_path {
333                batch = inject_source_path_column(batch, file_path)?;
334            }
335            out.push(batch);
336        }
337        Ok(out)
338    }
339
340    /// Legacy helper: recurse path and read by extension with an explicit schema
341    /// (used by unit tests and jshift-projected JSONL).
342    pub async fn scan_path(
343        &self,
344        path: impl AsRef<Path>,
345        schema: SchemaRef,
346    ) -> Result<Vec<RecordBatch>> {
347        let mut files = Vec::new();
348        collect_files(path.as_ref(), &mut files)?;
349        let mut batches = Vec::new();
350
351        for file_path in files {
352            let ext = file_path
353                .extension()
354                .and_then(|e| e.to_str())
355                .unwrap_or("")
356                .to_ascii_lowercase();
357            match ext.as_str() {
358                "parquet" => batches.extend(read_parquet(&file_path, Some(schema.clone()))?),
359                "json" | "jsonl" | "ndjson" => {
360                    let bytes = std::fs::read(&file_path)?;
361                    let extractor = crate::json::JShiftExtractor::new(self.paths.clone());
362                    batches.push(extractor.extract_jsonl(&bytes, schema.clone())?);
363                }
364                "csv" | "tsv" => batches.extend(read_csv(&file_path, Some(schema.clone()))?),
365                "log" | "txt" | "text" | "md" => {
366                    batches.push(read_line_oriented(&file_path)?);
367                }
368                "toml" => batches.push(read_toml(&file_path, None)?),
369                "arrow" | "arrows" | "ipc" | "feather" => {
370                    batches.extend(read_arrow_ipc_auto(&file_path)?);
371                }
372                _ => {
373                    tracing::debug!("Skipping unsupported file: {:?}", file_path);
374                }
375            }
376        }
377        Ok(batches)
378    }
379
380    fn read_file(&self, file_path: &Path, req: &ScanRequest) -> Result<Vec<RecordBatch>> {
381        match req.format {
382            SourceFormat::Parquet => read_parquet(file_path, None),
383            SourceFormat::Protobuf => Ok(vec![read_protobuf_opaque(
384                file_path,
385                req.protobuf_max_payload_bytes,
386            )?]),
387            SourceFormat::Csv => read_csv(file_path, None),
388            SourceFormat::Jsonl | SourceFormat::Json => {
389                if !self.paths.is_empty() {
390                    let schema = utf8_schema_from_paths(&self.paths);
391                    let bytes = std::fs::read(file_path)?;
392                    let extractor = crate::json::JShiftExtractor::new(self.paths.clone());
393                    // JSONL extractor works line-wise; for single JSON array, expand first
394                    if req.format == SourceFormat::Json {
395                        let expanded = expand_json_document_to_jsonl(&bytes)?;
396                        Ok(vec![extractor.extract_jsonl(&expanded, schema)?])
397                    } else {
398                        Ok(vec![extractor.extract_jsonl(&bytes, schema)?])
399                    }
400                } else {
401                    // Schema-free path: use Arrow JSON reader (line-delimited)
402                    read_json_arrow(file_path, req.format)
403                }
404            }
405            // Real lakes often write stream IPC with a `.arrow` extension.
406            SourceFormat::ArrowIpc | SourceFormat::ArrowIpcStream => read_arrow_ipc_auto(file_path),
407            SourceFormat::Log | SourceFormat::Txt => Ok(vec![read_line_oriented(file_path)?]),
408            SourceFormat::Toml => Ok(vec![read_toml(file_path, req.toml_rows_key.as_deref())?]),
409        }
410    }
411}
412
413fn utf8_schema_from_paths(paths: &[String]) -> SchemaRef {
414    Arc::new(Schema::new(
415        paths
416            .iter()
417            .map(|p| Field::new(p.as_str(), DataType::Utf8, true))
418            .collect::<Vec<_>>(),
419    ))
420}
421
422fn expand_json_document_to_jsonl(bytes: &[u8]) -> Result<Vec<u8>> {
423    let v: serde_json::Value =
424        serde_json::from_slice(bytes).map_err(|e| anyhow!("Invalid JSON document: {}", e))?;
425    match v {
426        serde_json::Value::Array(items) => {
427            let mut out = Vec::new();
428            for item in items {
429                out.extend(serde_json::to_vec(&item)?);
430                out.push(b'\n');
431            }
432            Ok(out)
433        }
434        other => {
435            let mut out = serde_json::to_vec(&other)?;
436            out.push(b'\n');
437            Ok(out)
438        }
439    }
440}
441
442fn collect_files(path: &Path, files: &mut Vec<PathBuf>) -> Result<()> {
443    if path.is_file() {
444        files.push(path.to_path_buf());
445    } else if path.is_dir() {
446        for entry in std::fs::read_dir(path)? {
447            let entry = entry?;
448            collect_files(&entry.path(), files)?;
449        }
450    }
451    Ok(())
452}
453
454fn extension_matches(format: SourceFormat, ext: &str) -> bool {
455    let ext = ext.to_ascii_lowercase();
456    match format {
457        SourceFormat::Jsonl => matches!(ext.as_str(), "jsonl" | "ndjson"),
458        SourceFormat::Json => ext == "json",
459        SourceFormat::Parquet => matches!(ext.as_str(), "parquet" | "pq"),
460        SourceFormat::Csv => matches!(ext.as_str(), "csv" | "tsv"),
461        SourceFormat::ArrowIpc => matches!(ext.as_str(), "arrow" | "arrows" | "ipc" | "feather"),
462        SourceFormat::ArrowIpcStream => {
463            matches!(
464                ext.as_str(),
465                "arrow" | "arrows" | "ipc" | "arrows_stream" | "ipc_stream"
466            )
467        }
468        SourceFormat::Log => ext == "log",
469        SourceFormat::Txt => matches!(ext.as_str(), "txt" | "text" | "md"),
470        SourceFormat::Toml => ext == "toml",
471        SourceFormat::Protobuf => matches!(ext.as_str(), "pb" | "protobuf" | "protobin"),
472    }
473}
474
475/// Opaque protobuf bronze: one row per file with path + raw bytes.
476///
477/// Typed message decode is intentionally deferred (schema registry / Rust models).
478/// `max_bytes` defaults to 1 GiB via project `scan.protobuf_max_payload_bytes`.
479fn read_protobuf_opaque(path: &Path, max_bytes: u64) -> Result<RecordBatch> {
480    let meta = std::fs::metadata(path).with_context(|| {
481        format!(
482            "E_RBT_PROTOBUF_IO: cannot stat protobuf file {}",
483            path.display()
484        )
485    })?;
486    let len = meta.len();
487    if len > max_bytes {
488        bail!(
489            "E_RBT_PROTOBUF_TOO_LARGE: file {} is {len} bytes, exceeds \
490             scan.protobuf_max_payload_bytes={max_bytes} (default {} = 1 GiB). \
491             Raise `scan.protobuf_max_payload_bytes` in rbt_project.yml only if intentional.",
492            path.display(),
493            DEFAULT_PROTOBUF_MAX_PAYLOAD_BYTES
494        );
495    }
496
497    let mut file = File::open(path)
498        .with_context(|| format!("E_RBT_PROTOBUF_IO: open protobuf file {}", path.display()))?;
499    let mut payload = Vec::with_capacity(len as usize);
500    file.read_to_end(&mut payload)
501        .with_context(|| format!("E_RBT_PROTOBUF_IO: read protobuf file {}", path.display()))?;
502
503    let schema = Arc::new(Schema::new(vec![
504        Field::new("_source_path", DataType::Utf8, false),
505        Field::new("payload", DataType::Binary, false),
506        Field::new("payload_len", DataType::Int64, false),
507    ]));
508
509    let mut path_b = StringBuilder::new();
510    let mut bin_b = BinaryBuilder::new();
511    let mut len_b = Int64Builder::new();
512    path_b.append_value(path.to_string_lossy().as_ref());
513    bin_b.append_value(&payload);
514    len_b.append_value(payload.len() as i64);
515
516    Ok(RecordBatch::try_new(
517        schema,
518        vec![
519            Arc::new(path_b.finish()),
520            Arc::new(bin_b.finish()),
521            Arc::new(len_b.finish()),
522        ],
523    )?)
524}
525
526fn collect_files_for_format(
527    path: &Path,
528    format: SourceFormat,
529    files: &mut Vec<PathBuf>,
530) -> Result<()> {
531    if path.is_file() {
532        files.push(path.to_path_buf());
533        return Ok(());
534    }
535    if path.is_dir() {
536        let mut all = Vec::new();
537        collect_files(path, &mut all)?;
538        for f in all {
539            let ext = f.extension().and_then(|e| e.to_str()).unwrap_or("");
540            if extension_matches(format, ext) {
541                files.push(f);
542            }
543        }
544        // If directory has files but none matched extension, and format is Log/Txt,
545        // still allow extensionless? Skip for now.
546        return Ok(());
547    }
548    bail!(
549        "scan path is neither file nor directory: {}",
550        path.display()
551    );
552}
553
554fn read_parquet(path: &Path, projection: Option<SchemaRef>) -> Result<Vec<RecordBatch>> {
555    let file = File::open(path)?;
556    let builder = parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder::try_new(file)?;
557    let reader = if let Some(schema) = projection {
558        let schema_descr = builder.metadata().file_metadata().schema_descr_ptr();
559        let arrow_schema = builder.schema();
560        let mut indices = Vec::new();
561        for field in schema.fields() {
562            if let Some((idx, _)) =
563                parquet::arrow::parquet_column(&schema_descr, arrow_schema, field.name())
564            {
565                indices.push(idx);
566            }
567        }
568        let mask = parquet::arrow::ProjectionMask::leaves(&schema_descr, indices);
569        builder.with_projection(mask).build()?
570    } else {
571        builder.build()?
572    };
573    let mut batches = Vec::new();
574    for batch in reader {
575        batches.push(batch?);
576    }
577    Ok(batches)
578}
579
580fn read_csv(path: &Path, schema: Option<SchemaRef>) -> Result<Vec<RecordBatch>> {
581    let file = File::open(path)?;
582    let mut batches = Vec::new();
583    if let Some(schema) = schema {
584        let reader = arrow::csv::ReaderBuilder::new(schema)
585            .with_header(true)
586            .build(file)?;
587        for batch in reader {
588            batches.push(batch?);
589        }
590    } else {
591        // Infer schema from first rows
592        let format = arrow::csv::reader::Format::default().with_header(true);
593        let (inferred, _) = format.infer_schema(File::open(path)?, Some(1024))?;
594        let schema = Arc::new(inferred);
595        let file = File::open(path)?;
596        let reader = arrow::csv::ReaderBuilder::new(schema)
597            .with_header(true)
598            .build(file)?;
599        for batch in reader {
600            batches.push(batch?);
601        }
602    }
603    Ok(batches)
604}
605
606fn read_json_arrow(path: &Path, format: SourceFormat) -> Result<Vec<RecordBatch>> {
607    let data = std::fs::read(path)?;
608    let data = if format == SourceFormat::Json {
609        expand_json_document_to_jsonl(&data)?
610    } else {
611        data
612    };
613    read_json_infer_from_bytes(&data)
614}
615
616fn read_json_infer_from_bytes(data: &[u8]) -> Result<Vec<RecordBatch>> {
617    use arrow::json::reader::infer_json_schema_from_seekable;
618    let mut cursor = std::io::Cursor::new(data);
619    let (schema, _n) = infer_json_schema_from_seekable(&mut cursor, Some(1024))?;
620    cursor.set_position(0);
621    let schema = Arc::new(schema);
622    let reader = arrow::json::ReaderBuilder::new(schema).build(cursor)?;
623    let mut batches = Vec::new();
624    for batch in reader {
625        batches.push(batch?);
626    }
627    Ok(batches)
628}
629
630fn read_arrow_ipc_file(path: &Path) -> Result<Vec<RecordBatch>> {
631    let file = File::open(path)?;
632    let reader = arrow::ipc::reader::FileReader::try_new(file, None)?;
633    let mut batches = Vec::new();
634    for batch in reader {
635        batches.push(batch?);
636    }
637    Ok(batches)
638}
639
640fn read_arrow_ipc_stream(path: &Path) -> Result<Vec<RecordBatch>> {
641    let file = File::open(path)?;
642    let reader = arrow::ipc::reader::StreamReader::try_new(file, None)?;
643    let mut batches = Vec::new();
644    for batch in reader {
645        batches.push(batch?);
646    }
647    Ok(batches)
648}
649
650/// Prefer random-access IPC file; fall back to stream (common for `.arrow` lake dumps).
651fn read_arrow_ipc_auto(path: &Path) -> Result<Vec<RecordBatch>> {
652    match read_arrow_ipc_file(path) {
653        Ok(batches) => Ok(batches),
654        Err(file_err) => read_arrow_ipc_stream(path).with_context(|| {
655            format!(
656                "Arrow IPC file and stream readers both failed for {} (file error: {file_err})",
657                path.display()
658            )
659        }),
660    }
661}
662
663/// Parse `key=value` segments from `file` relative to `root`.
664pub fn parse_hive_partitions(
665    file: &Path,
666    root: &Path,
667) -> std::collections::HashMap<String, String> {
668    let mut out = std::collections::HashMap::new();
669    let rel = file.strip_prefix(root).unwrap_or(file);
670    for comp in rel.components() {
671        if let std::path::Component::Normal(os) = comp {
672            let s = os.to_string_lossy();
673            if let Some((k, v)) = s.split_once('=') {
674                if !k.is_empty() {
675                    out.insert(k.to_string(), v.to_string());
676                }
677            }
678        }
679    }
680    out
681}
682
683fn path_matches_require_partitions(
684    file: &Path,
685    root: &Path,
686    require: &std::collections::HashMap<String, String>,
687) -> bool {
688    if require.is_empty() {
689        return true;
690    }
691    let parts = parse_hive_partitions(file, root);
692    require
693        .iter()
694        .all(|(k, v)| parts.get(k).map(|pv| pv == v).unwrap_or(false))
695}
696
697/// Append hive partition columns (Utf8) for keys in `partition_by` that are not already present.
698fn inject_hive_partitions(
699    batch: RecordBatch,
700    file: &Path,
701    root: &Path,
702    partition_by: &[String],
703) -> Result<RecordBatch> {
704    if partition_by.is_empty() {
705        return Ok(batch);
706    }
707    let parts = parse_hive_partitions(file, root);
708    let n = batch.num_rows();
709    let mut fields: Vec<arrow::datatypes::Field> = batch
710        .schema()
711        .fields()
712        .iter()
713        .map(|f| f.as_ref().clone())
714        .collect();
715    let mut columns: Vec<ArrayRef> = batch.columns().to_vec();
716
717    for key in partition_by {
718        if batch.schema().index_of(key).is_ok() {
719            // Column already in payload (e.g. symbol); keep file data, do not overwrite.
720            continue;
721        }
722        let val = parts.get(key).map(|s| s.as_str());
723        let mut b = StringBuilder::with_capacity(n, n * 8);
724        for _ in 0..n {
725            match val {
726                Some(v) => b.append_value(v),
727                None => b.append_null(),
728            }
729        }
730        fields.push(Field::new(key.as_str(), DataType::Utf8, true));
731        columns.push(Arc::new(b.finish()) as ArrayRef);
732    }
733
734    let schema = Arc::new(Schema::new(fields));
735    Ok(RecordBatch::try_new(schema, columns)?)
736}
737
738/// Append `_source_path` Utf8 with the absolute bronze file path (for latest-wins dedupe).
739fn inject_source_path_column(batch: RecordBatch, file: &Path) -> Result<RecordBatch> {
740    if batch.schema().index_of("_source_path").is_ok() {
741        return Ok(batch);
742    }
743    let n = batch.num_rows();
744    let path_str = file.to_string_lossy();
745    let mut b = StringBuilder::with_capacity(n, n * path_str.len().max(16));
746    for _ in 0..n {
747        b.append_value(path_str.as_ref());
748    }
749    let mut fields: Vec<Field> = batch
750        .schema()
751        .fields()
752        .iter()
753        .map(|f| f.as_ref().clone())
754        .collect();
755    fields.push(Field::new("_source_path", DataType::Utf8, false));
756    let mut columns = batch.columns().to_vec();
757    columns.push(Arc::new(b.finish()) as ArrayRef);
758    Ok(RecordBatch::try_new(
759        Arc::new(Schema::new(fields)),
760        columns,
761    )?)
762}
763
764/// Line-oriented bronze: `.log`, `.txt`, llms.txt-style docs.
765/// Schema: `line_no` Int64, `content` Utf8.
766fn read_line_oriented(path: &Path) -> Result<RecordBatch> {
767    let file = File::open(path)?;
768    let reader = BufReader::new(file);
769    let mut line_nos = Int64Builder::new();
770    let mut contents = StringBuilder::new();
771    let mut n: i64 = 0;
772    for line in reader.lines() {
773        let line = line?;
774        n += 1;
775        line_nos.append_value(n);
776        contents.append_value(line);
777    }
778    let schema = Arc::new(Schema::new(vec![
779        Field::new("line_no", DataType::Int64, false),
780        Field::new("content", DataType::Utf8, false),
781    ]));
782    Ok(RecordBatch::try_new(
783        schema,
784        vec![
785            Arc::new(line_nos.finish()) as ArrayRef,
786            Arc::new(contents.finish()) as ArrayRef,
787        ],
788    )?)
789}
790
791fn read_toml(path: &Path, rows_key: Option<&str>) -> Result<RecordBatch> {
792    let text = std::fs::read_to_string(path)?;
793    let value: toml::Value = text
794        .parse()
795        .with_context(|| format!("Invalid TOML: {}", path.display()))?;
796
797    let rows: Vec<toml::map::Map<String, toml::Value>> = match &value {
798        toml::Value::Table(table) => {
799            if let Some(key) = rows_key {
800                match table.get(key) {
801                    Some(toml::Value::Array(arr)) => array_of_tables(arr)?,
802                    other => bail!(
803                        "toml_rows_key '{}' is not an array of tables (got {:?})",
804                        key,
805                        other.map(|v| v.type_str())
806                    ),
807                }
808            } else if let Some((_k, toml::Value::Array(arr))) = table
809                .iter()
810                .find(|(_, v)| matches!(v, toml::Value::Array(a) if a.iter().all(|x| x.is_table())))
811            {
812                array_of_tables(arr)?
813            } else {
814                // Single-row: top-level scalars / nested stringified
815                vec![table.clone()]
816            }
817        }
818        toml::Value::Array(arr) => array_of_tables(arr)?,
819        other => bail!("Unsupported top-level TOML type: {}", other.type_str()),
820    };
821
822    if rows.is_empty() {
823        let schema = Arc::new(Schema::new(vec![Field::new("empty", DataType::Utf8, true)]));
824        return Ok(RecordBatch::try_new(
825            schema,
826            vec![Arc::new(StringBuilder::new().finish()) as ArrayRef],
827        )?);
828    }
829
830    // Column union
831    let mut col_names: Vec<String> = Vec::new();
832    for row in &rows {
833        for k in row.keys() {
834            if !col_names.iter().any(|c| c == k) {
835                col_names.push(k.clone());
836            }
837        }
838    }
839
840    let mut builders: Vec<StringBuilder> = col_names
841        .iter()
842        .map(|_| StringBuilder::with_capacity(rows.len(), rows.len() * 16))
843        .collect();
844
845    for row in &rows {
846        for (i, col) in col_names.iter().enumerate() {
847            match row.get(col) {
848                Some(v) => builders[i].append_value(toml_value_to_string(v)),
849                None => builders[i].append_null(),
850            }
851        }
852    }
853
854    let fields: Vec<Field> = col_names
855        .iter()
856        .map(|c| Field::new(c.as_str(), DataType::Utf8, true))
857        .collect();
858    let schema = Arc::new(Schema::new(fields));
859    let arrays: Vec<ArrayRef> = builders
860        .into_iter()
861        .map(|mut b| Arc::new(b.finish()) as ArrayRef)
862        .collect();
863    Ok(RecordBatch::try_new(schema, arrays)?)
864}
865
866fn array_of_tables(arr: &[toml::Value]) -> Result<Vec<toml::map::Map<String, toml::Value>>> {
867    let mut rows = Vec::with_capacity(arr.len());
868    for (i, item) in arr.iter().enumerate() {
869        match item {
870            toml::Value::Table(t) => rows.push(t.clone()),
871            other => bail!(
872                "TOML array element {} is not a table (got {})",
873                i,
874                other.type_str()
875            ),
876        }
877    }
878    Ok(rows)
879}
880
881fn toml_value_to_string(v: &toml::Value) -> String {
882    match v {
883        toml::Value::String(s) => s.clone(),
884        toml::Value::Integer(i) => i.to_string(),
885        toml::Value::Float(f) => f.to_string(),
886        toml::Value::Boolean(b) => b.to_string(),
887        toml::Value::Datetime(d) => d.to_string(),
888        other => other.to_string(),
889    }
890}
891
892#[cfg(test)]
893mod tests {
894    use super::*;
895    use arrow::array::{Array, Int64Array, StringArray};
896    use arrow::datatypes::{DataType, Field, Schema};
897    use std::sync::Arc;
898
899    #[tokio::test]
900    async fn test_lake_scanner_multi_format() -> Result<()> {
901        let temp_dir = tempfile::tempdir()?;
902        let dir_path = temp_dir.path();
903
904        let schema = Arc::new(Schema::new(vec![
905            Field::new("id", DataType::Int64, true),
906            Field::new("name", DataType::Utf8, true),
907        ]));
908
909        std::fs::write(
910            dir_path.join("file1.jsonl"),
911            b"{\"id\": 1, \"name\": \"Alice\"}\n{\"id\": 2, \"name\": \"Bob\"}\n",
912        )?;
913        std::fs::write(dir_path.join("file2.csv"), b"id,name\n3,Charlie\n4,Dave\n")?;
914
915        let batch = RecordBatch::try_new(
916            schema.clone(),
917            vec![
918                Arc::new(Int64Array::from(vec![5, 6])),
919                Arc::new(StringArray::from(vec!["Eve", "Frank"])),
920            ],
921        )?;
922        let file = File::create(dir_path.join("file3.parquet"))?;
923        let mut writer =
924            parquet::arrow::arrow_writer::ArrowWriter::try_new(file, schema.clone(), None)?;
925        writer.write(&batch)?;
926        writer.close()?;
927
928        std::fs::write(dir_path.join("app.log"), "info boot\nwarn disk\n")?;
929        std::fs::write(
930            dir_path.join("meta.toml"),
931            r#"
932[[records]]
933id = "1"
934name = "toml_a"
935[[records]]
936id = "2"
937name = "toml_b"
938"#,
939        )?;
940
941        let scanner = LakeScanner::new(vec!["id".to_string(), "name".to_string()]);
942        let batches = scanner.scan_path(dir_path, schema.clone()).await?;
943        let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
944        // 2 jsonl + 2 csv + 2 parquet + 2 log + 2 toml = 10
945        assert_eq!(total_rows, 10);
946        Ok(())
947    }
948
949    #[tokio::test]
950    async fn test_scan_request_jsonl() -> Result<()> {
951        let temp = tempfile::tempdir()?;
952        let path = temp.path().join("trades.jsonl");
953        std::fs::write(
954            &path,
955            r#"{"ticker":"NVDA","price":1.0}
956{"ticker":"AAPL","price":2.0}
957"#,
958        )?;
959
960        let fm = StagingFrontmatter {
961            source_format: Some(SourceFormat::Jsonl),
962            scan_path: Some(path.file_name().unwrap().to_string_lossy().into()),
963            ..Default::default()
964        };
965        let req = ScanRequest::from_frontmatter(temp.path(), &fm)?;
966        let scanner = LakeScanner::from_request(&req);
967        let batches = scanner.scan(&req).await?;
968        let rows: usize = batches.iter().map(|b| b.num_rows()).sum();
969        assert_eq!(rows, 2);
970        Ok(())
971    }
972
973    #[tokio::test]
974    async fn test_hive_partition_inject_and_filter() -> Result<()> {
975        let temp = tempfile::tempdir()?;
976        let root = temp.path();
977        let dir = root.join("symbol=NVDA").join("timeframe=1m");
978        std::fs::create_dir_all(&dir)?;
979        // minimal IPC stream with one utf8 column
980        let schema = Arc::new(Schema::new(vec![Field::new(
981            "close",
982            DataType::Float64,
983            true,
984        )]));
985        let batch = RecordBatch::try_new(
986            schema,
987            vec![Arc::new(arrow::array::Float64Array::from(vec![1.0, 2.0])) as ArrayRef],
988        )?;
989        let path = dir.join("chunk.arrow");
990        {
991            let file = File::create(&path)?;
992            let mut writer = arrow::ipc::writer::StreamWriter::try_new(file, &batch.schema())?;
993            writer.write(&batch)?;
994            writer.finish()?;
995        }
996        // noise partition that should be filtered out
997        let other = root.join("symbol=AAPL").join("timeframe=1d");
998        std::fs::create_dir_all(&other)?;
999        {
1000            let file = File::create(other.join("chunk.arrow"))?;
1001            let mut writer = arrow::ipc::writer::StreamWriter::try_new(file, &batch.schema())?;
1002            writer.write(&batch)?;
1003            writer.finish()?;
1004        }
1005
1006        let mut require = std::collections::HashMap::new();
1007        require.insert("timeframe".into(), "1m".into());
1008        let fm = StagingFrontmatter {
1009            source_format: Some(SourceFormat::ArrowIpcStream),
1010            scan_path: Some(".".into()),
1011            partition_by: Some(vec!["symbol".into(), "timeframe".into()]),
1012            require_partitions: Some(require),
1013            ..Default::default()
1014        };
1015        let req = ScanRequest::from_frontmatter(root, &fm)?;
1016        let scanner = LakeScanner::from_request(&req);
1017        let batches = scanner.scan(&req).await?;
1018        let rows: usize = batches.iter().map(|b| b.num_rows()).sum();
1019        assert_eq!(rows, 2);
1020        let schema = batches[0].schema();
1021        assert!(schema.index_of("timeframe").is_ok());
1022        assert!(schema.index_of("symbol").is_ok());
1023        Ok(())
1024    }
1025
1026    #[test]
1027    fn test_line_oriented_and_toml() -> Result<()> {
1028        let temp = tempfile::tempdir()?;
1029        let txt = temp.path().join("llms.txt");
1030        std::fs::write(&txt, "# Title\n\nSome doc line\n")?;
1031        let batch = read_line_oriented(&txt)?;
1032        assert_eq!(batch.num_rows(), 3);
1033        assert_eq!(batch.schema().field(0).name(), "line_no");
1034        assert_eq!(batch.schema().field(1).name(), "content");
1035
1036        let toml_path = temp.path().join("cfg.toml");
1037        std::fs::write(
1038            &toml_path,
1039            r#"
1040[[items]]
1041k = "a"
1042[[items]]
1043k = "b"
1044"#,
1045        )?;
1046        let tbatch = read_toml(&toml_path, Some("items"))?;
1047        assert_eq!(tbatch.num_rows(), 2);
1048        let col = tbatch
1049            .column(0)
1050            .as_any()
1051            .downcast_ref::<StringArray>()
1052            .unwrap();
1053        assert!(col.value(0) == "a" || col.value(1) == "a");
1054        Ok(())
1055    }
1056
1057    #[tokio::test]
1058    async fn test_path_glob_filters_artifacts() -> Result<()> {
1059        let temp = tempfile::tempdir()?;
1060        let root = temp.path();
1061        // Multi-artifact hive (same layout pattern as multi-domain landing zones)
1062        let d1 = root
1063            .join("domain=x.com")
1064            .join("report_date=2026-07-29")
1065            .join("run_id=r1")
1066            .join("raw_snoop");
1067        std::fs::create_dir_all(&d1)?;
1068        // two different parquet artifact types
1069        let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)]));
1070        for (name, ids) in [
1071            ("crawlplan.parquet", vec![1i64]),
1072            ("other.parquet", vec![9i64]),
1073        ] {
1074            let batch =
1075                RecordBatch::try_new(schema.clone(), vec![Arc::new(Int64Array::from(ids))])?;
1076            let f = File::create(d1.join(name))?;
1077            let mut w =
1078                parquet::arrow::arrow_writer::ArrowWriter::try_new(f, schema.clone(), None)?;
1079            w.write(&batch)?;
1080            w.close()?;
1081        }
1082
1083        let fm = StagingFrontmatter {
1084            source_format: Some(SourceFormat::Parquet),
1085            scan_path: Some(".".into()),
1086            path_glob: Some(vec!["**/crawlplan.parquet".into()]),
1087            partition_by: Some(vec!["domain".into()]),
1088            inject_source_path: Some(true),
1089            ..Default::default()
1090        };
1091        let req = ScanRequest::from_frontmatter(root, &fm)?;
1092        let scanner = LakeScanner::from_request(&req);
1093        let batches = scanner.scan(&req).await?;
1094        let rows: usize = batches.iter().map(|b| b.num_rows()).sum();
1095        assert_eq!(rows, 1, "only crawlplan.parquet rows");
1096        assert!(batches[0].schema().index_of("domain").is_ok());
1097        Ok(())
1098    }
1099
1100    #[tokio::test]
1101    async fn test_protobuf_opaque_scan() -> Result<()> {
1102        let temp = tempfile::tempdir()?;
1103        let pb = temp.path().join("msg.pb");
1104        std::fs::write(&pb, b"\x08\x96\x01")?; // arbitrary bytes
1105        let fm = StagingFrontmatter {
1106            source_format: Some(SourceFormat::Protobuf),
1107            scan_path: Some(pb.file_name().unwrap().to_string_lossy().into()),
1108            ..Default::default()
1109        };
1110        let req = ScanRequest::from_frontmatter(temp.path(), &fm)?;
1111        assert_eq!(
1112            req.protobuf_max_payload_bytes,
1113            DEFAULT_PROTOBUF_MAX_PAYLOAD_BYTES
1114        );
1115        let scanner = LakeScanner::from_request(&req);
1116        let batches = scanner.scan(&req).await?;
1117        assert_eq!(batches.len(), 1);
1118        assert_eq!(batches[0].num_rows(), 1);
1119        assert_eq!(batches[0].schema().field(0).name(), "_source_path");
1120        assert_eq!(batches[0].schema().field(1).name(), "payload");
1121        Ok(())
1122    }
1123
1124    #[tokio::test]
1125    async fn test_protobuf_respects_max_payload_bytes() -> Result<()> {
1126        let temp = tempfile::tempdir()?;
1127        let pb = temp.path().join("big.pb");
1128        std::fs::write(&pb, vec![0u8; 64])?;
1129        let fm = StagingFrontmatter {
1130            source_format: Some(SourceFormat::Protobuf),
1131            scan_path: Some("big.pb".into()),
1132            ..Default::default()
1133        };
1134        let scan_cfg = ScanConfig {
1135            protobuf_max_payload_bytes: 16,
1136            ..Default::default()
1137        };
1138        let req = ScanRequest::from_frontmatter_with_config(
1139            temp.path(),
1140            &fm,
1141            HashMap::new(),
1142            &scan_cfg,
1143        )?;
1144        let scanner = LakeScanner::from_request(&req);
1145        let err = scanner.scan(&req).await.unwrap_err();
1146        // anyhow chains: outer "Failed reading bronze file …" + inner E_RBT_PROTOBUF_TOO_LARGE
1147        let full = format!("{err:#}");
1148        assert!(
1149            full.contains("E_RBT_PROTOBUF_TOO_LARGE"),
1150            "expected size-cap error in chain, got: {full}"
1151        );
1152        assert!(full.contains("protobuf_max_payload_bytes"));
1153        Ok(())
1154    }
1155
1156    #[tokio::test]
1157    async fn test_path_glob_single_star_is_not_recursive() -> Result<()> {
1158        let temp = tempfile::tempdir()?;
1159        let root = temp.path();
1160        let deep = root.join("a").join("b").join("raw");
1161        std::fs::create_dir_all(&deep)?;
1162        let shallow = root.join("raw");
1163        std::fs::create_dir_all(&shallow)?;
1164        let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)]));
1165        for (dir, id) in [(&deep, 1i64), (&shallow, 2i64)] {
1166            let batch =
1167                RecordBatch::try_new(schema.clone(), vec![Arc::new(Int64Array::from(vec![id]))])?;
1168            let f = File::create(dir.join("crawlplan.parquet"))?;
1169            let mut w =
1170                parquet::arrow::arrow_writer::ArrowWriter::try_new(f, schema.clone(), None)?;
1171            w.write(&batch)?;
1172            w.close()?;
1173        }
1174        let fm = StagingFrontmatter {
1175            source_format: Some(SourceFormat::Parquet),
1176            scan_path: Some(".".into()),
1177            // one segment only — should hit raw/crawlplan, not a/b/raw/crawlplan
1178            path_glob: Some(vec!["*/crawlplan.parquet".into()]),
1179            ..Default::default()
1180        };
1181        let req = ScanRequest::from_frontmatter(root, &fm)?;
1182        let scanner = LakeScanner::from_request(&req);
1183        let batches = scanner.scan(&req).await?;
1184        let rows: usize = batches.iter().map(|b| b.num_rows()).sum();
1185        assert_eq!(rows, 1, "single-star glob must not match deep hive path");
1186        Ok(())
1187    }
1188
1189    #[tokio::test]
1190    async fn test_absolute_scan_path() -> Result<()> {
1191        let temp = tempfile::tempdir()?;
1192        let abs = temp.path().join("abs.jsonl");
1193        std::fs::write(&abs, b"{\"a\":1}\n{\"a\":2}\n")?;
1194        let fm = StagingFrontmatter {
1195            source_format: Some(SourceFormat::Jsonl),
1196            scan_path: Some(abs.to_string_lossy().into()),
1197            ..Default::default()
1198        };
1199        // project_dir is different from parent of abs
1200        let proj = tempfile::tempdir()?;
1201        let req = ScanRequest::from_frontmatter(proj.path(), &fm)?;
1202        let scanner = LakeScanner::from_request(&req);
1203        let batches = scanner.scan(&req).await?;
1204        assert_eq!(batches.iter().map(|b| b.num_rows()).sum::<usize>(), 2);
1205        Ok(())
1206    }
1207}