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