1pub 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#[derive(Debug, Clone)]
23pub struct ScanRequest {
24 pub project_dir: PathBuf,
25 pub scan_path: String,
26 pub format: SourceFormat,
27 pub paths: Vec<String>,
29 pub toml_rows_key: Option<String>,
31 pub partition_by: Vec<String>,
33 pub require_partitions: std::collections::HashMap<String, String>,
35 pub path_glob: Vec<String>,
37 pub inject_source_path: bool,
39 pub roots: HashMap<String, String>,
41 pub protobuf_max_payload_bytes: u64,
43 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
107pub 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 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 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 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 }
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 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 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 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 read_json_arrow(file_path, req.format)
431 }
432 }
433 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
503fn 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 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 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
678fn 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
691pub 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
725fn 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 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
766fn 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
792fn 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 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 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 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 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 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 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 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")?; 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 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 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 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}