1use 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#[derive(Debug, Clone)]
20pub struct ScanRequest {
21 pub project_dir: PathBuf,
22 pub scan_path: String,
23 pub format: SourceFormat,
24 pub paths: Vec<String>,
26 pub toml_rows_key: Option<String>,
28 pub partition_by: Vec<String>,
30 pub require_partitions: std::collections::HashMap<String, String>,
32 pub path_glob: Vec<String>,
34 pub inject_source_path: bool,
36 pub roots: HashMap<String, String>,
38 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
98pub 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 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 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 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 }
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 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 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 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 read_json_arrow(file_path, req.format)
403 }
404 }
405 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
475fn 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 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 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
650fn 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
663pub 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
697fn 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 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
738fn 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
764fn 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 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 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 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 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 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 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 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")?; 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 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 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 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}