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 async fn scan(&self, req: &ScanRequest) -> Result<Vec<RecordBatch>> {
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 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
155 tracing::info!(
156 "Bronze scan: {} file(s) under {} (format={}, globs={:?})",
157 files.len(),
158 root.display(),
159 req.format,
160 req.path_glob
161 );
162
163 let mut batches = Vec::new();
164 for file_path in files {
165 let file_batches = self
166 .read_file(&file_path, req)
167 .with_context(|| format!("Failed reading bronze file {}", file_path.display()))?;
168 for batch in file_batches {
169 let mut batch = if req.partition_by.is_empty() {
170 batch
171 } else {
172 inject_hive_partitions(batch, &file_path, &root, &req.partition_by)?
173 };
174 if req.inject_source_path {
175 batch = inject_source_path_column(batch, &file_path)?;
176 }
177 batches.push(batch);
178 }
179 }
180 Ok(batches)
181 }
182
183 pub async fn scan_path(
186 &self,
187 path: impl AsRef<Path>,
188 schema: SchemaRef,
189 ) -> Result<Vec<RecordBatch>> {
190 let mut files = Vec::new();
191 collect_files(path.as_ref(), &mut files)?;
192 let mut batches = Vec::new();
193
194 for file_path in files {
195 let ext = file_path
196 .extension()
197 .and_then(|e| e.to_str())
198 .unwrap_or("")
199 .to_ascii_lowercase();
200 match ext.as_str() {
201 "parquet" => batches.extend(read_parquet(&file_path, Some(schema.clone()))?),
202 "json" | "jsonl" | "ndjson" => {
203 let bytes = std::fs::read(&file_path)?;
204 let extractor = crate::json::JShiftExtractor::new(self.paths.clone());
205 batches.push(extractor.extract_jsonl(&bytes, schema.clone())?);
206 }
207 "csv" | "tsv" => batches.extend(read_csv(&file_path, Some(schema.clone()))?),
208 "log" | "txt" | "text" | "md" => {
209 batches.push(read_line_oriented(&file_path)?);
210 }
211 "toml" => batches.push(read_toml(&file_path, None)?),
212 "arrow" | "arrows" | "ipc" | "feather" => {
213 batches.extend(read_arrow_ipc_auto(&file_path)?);
214 }
215 _ => {
216 tracing::debug!("Skipping unsupported file: {:?}", file_path);
217 }
218 }
219 }
220 Ok(batches)
221 }
222
223 fn read_file(&self, file_path: &Path, req: &ScanRequest) -> Result<Vec<RecordBatch>> {
224 match req.format {
225 SourceFormat::Parquet => read_parquet(file_path, None),
226 SourceFormat::Protobuf => Ok(vec![read_protobuf_opaque(
227 file_path,
228 req.protobuf_max_payload_bytes,
229 )?]),
230 SourceFormat::Csv => read_csv(file_path, None),
231 SourceFormat::Jsonl | SourceFormat::Json => {
232 if !self.paths.is_empty() {
233 let schema = utf8_schema_from_paths(&self.paths);
234 let bytes = std::fs::read(file_path)?;
235 let extractor = crate::json::JShiftExtractor::new(self.paths.clone());
236 if req.format == SourceFormat::Json {
238 let expanded = expand_json_document_to_jsonl(&bytes)?;
239 Ok(vec![extractor.extract_jsonl(&expanded, schema)?])
240 } else {
241 Ok(vec![extractor.extract_jsonl(&bytes, schema)?])
242 }
243 } else {
244 read_json_arrow(file_path, req.format)
246 }
247 }
248 SourceFormat::ArrowIpc | SourceFormat::ArrowIpcStream => read_arrow_ipc_auto(file_path),
250 SourceFormat::Log | SourceFormat::Txt => Ok(vec![read_line_oriented(file_path)?]),
251 SourceFormat::Toml => Ok(vec![read_toml(file_path, req.toml_rows_key.as_deref())?]),
252 }
253 }
254}
255
256fn utf8_schema_from_paths(paths: &[String]) -> SchemaRef {
257 Arc::new(Schema::new(
258 paths
259 .iter()
260 .map(|p| Field::new(p.as_str(), DataType::Utf8, true))
261 .collect::<Vec<_>>(),
262 ))
263}
264
265fn expand_json_document_to_jsonl(bytes: &[u8]) -> Result<Vec<u8>> {
266 let v: serde_json::Value =
267 serde_json::from_slice(bytes).map_err(|e| anyhow!("Invalid JSON document: {}", e))?;
268 match v {
269 serde_json::Value::Array(items) => {
270 let mut out = Vec::new();
271 for item in items {
272 out.extend(serde_json::to_vec(&item)?);
273 out.push(b'\n');
274 }
275 Ok(out)
276 }
277 other => {
278 let mut out = serde_json::to_vec(&other)?;
279 out.push(b'\n');
280 Ok(out)
281 }
282 }
283}
284
285fn collect_files(path: &Path, files: &mut Vec<PathBuf>) -> Result<()> {
286 if path.is_file() {
287 files.push(path.to_path_buf());
288 } else if path.is_dir() {
289 for entry in std::fs::read_dir(path)? {
290 let entry = entry?;
291 collect_files(&entry.path(), files)?;
292 }
293 }
294 Ok(())
295}
296
297fn extension_matches(format: SourceFormat, ext: &str) -> bool {
298 let ext = ext.to_ascii_lowercase();
299 match format {
300 SourceFormat::Jsonl => matches!(ext.as_str(), "jsonl" | "ndjson"),
301 SourceFormat::Json => ext == "json",
302 SourceFormat::Parquet => matches!(ext.as_str(), "parquet" | "pq"),
303 SourceFormat::Csv => matches!(ext.as_str(), "csv" | "tsv"),
304 SourceFormat::ArrowIpc => matches!(ext.as_str(), "arrow" | "arrows" | "ipc" | "feather"),
305 SourceFormat::ArrowIpcStream => {
306 matches!(
307 ext.as_str(),
308 "arrow" | "arrows" | "ipc" | "arrows_stream" | "ipc_stream"
309 )
310 }
311 SourceFormat::Log => ext == "log",
312 SourceFormat::Txt => matches!(ext.as_str(), "txt" | "text" | "md"),
313 SourceFormat::Toml => ext == "toml",
314 SourceFormat::Protobuf => matches!(ext.as_str(), "pb" | "protobuf" | "protobin"),
315 }
316}
317
318fn read_protobuf_opaque(path: &Path, max_bytes: u64) -> Result<RecordBatch> {
323 let meta = std::fs::metadata(path).with_context(|| {
324 format!(
325 "E_RBT_PROTOBUF_IO: cannot stat protobuf file {}",
326 path.display()
327 )
328 })?;
329 let len = meta.len();
330 if len > max_bytes {
331 bail!(
332 "E_RBT_PROTOBUF_TOO_LARGE: file {} is {len} bytes, exceeds \
333 scan.protobuf_max_payload_bytes={max_bytes} (default {} = 1 GiB). \
334 Raise `scan.protobuf_max_payload_bytes` in rbt_project.yml only if intentional.",
335 path.display(),
336 DEFAULT_PROTOBUF_MAX_PAYLOAD_BYTES
337 );
338 }
339
340 let mut file = File::open(path)
341 .with_context(|| format!("E_RBT_PROTOBUF_IO: open protobuf file {}", path.display()))?;
342 let mut payload = Vec::with_capacity(len as usize);
343 file.read_to_end(&mut payload)
344 .with_context(|| format!("E_RBT_PROTOBUF_IO: read protobuf file {}", path.display()))?;
345
346 let schema = Arc::new(Schema::new(vec![
347 Field::new("_source_path", DataType::Utf8, false),
348 Field::new("payload", DataType::Binary, false),
349 Field::new("payload_len", DataType::Int64, false),
350 ]));
351
352 let mut path_b = StringBuilder::new();
353 let mut bin_b = BinaryBuilder::new();
354 let mut len_b = Int64Builder::new();
355 path_b.append_value(path.to_string_lossy().as_ref());
356 bin_b.append_value(&payload);
357 len_b.append_value(payload.len() as i64);
358
359 Ok(RecordBatch::try_new(
360 schema,
361 vec![
362 Arc::new(path_b.finish()),
363 Arc::new(bin_b.finish()),
364 Arc::new(len_b.finish()),
365 ],
366 )?)
367}
368
369fn collect_files_for_format(
370 path: &Path,
371 format: SourceFormat,
372 files: &mut Vec<PathBuf>,
373) -> Result<()> {
374 if path.is_file() {
375 files.push(path.to_path_buf());
376 return Ok(());
377 }
378 if path.is_dir() {
379 let mut all = Vec::new();
380 collect_files(path, &mut all)?;
381 for f in all {
382 let ext = f.extension().and_then(|e| e.to_str()).unwrap_or("");
383 if extension_matches(format, ext) {
384 files.push(f);
385 }
386 }
387 return Ok(());
390 }
391 bail!(
392 "scan path is neither file nor directory: {}",
393 path.display()
394 );
395}
396
397fn read_parquet(path: &Path, projection: Option<SchemaRef>) -> Result<Vec<RecordBatch>> {
398 let file = File::open(path)?;
399 let builder = parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder::try_new(file)?;
400 let reader = if let Some(schema) = projection {
401 let schema_descr = builder.metadata().file_metadata().schema_descr_ptr();
402 let arrow_schema = builder.schema();
403 let mut indices = Vec::new();
404 for field in schema.fields() {
405 if let Some((idx, _)) =
406 parquet::arrow::parquet_column(&schema_descr, arrow_schema, field.name())
407 {
408 indices.push(idx);
409 }
410 }
411 let mask = parquet::arrow::ProjectionMask::leaves(&schema_descr, indices);
412 builder.with_projection(mask).build()?
413 } else {
414 builder.build()?
415 };
416 let mut batches = Vec::new();
417 for batch in reader {
418 batches.push(batch?);
419 }
420 Ok(batches)
421}
422
423fn read_csv(path: &Path, schema: Option<SchemaRef>) -> Result<Vec<RecordBatch>> {
424 let file = File::open(path)?;
425 let mut batches = Vec::new();
426 if let Some(schema) = schema {
427 let reader = arrow::csv::ReaderBuilder::new(schema)
428 .with_header(true)
429 .build(file)?;
430 for batch in reader {
431 batches.push(batch?);
432 }
433 } else {
434 let format = arrow::csv::reader::Format::default().with_header(true);
436 let (inferred, _) = format.infer_schema(File::open(path)?, Some(1024))?;
437 let schema = Arc::new(inferred);
438 let file = File::open(path)?;
439 let reader = arrow::csv::ReaderBuilder::new(schema)
440 .with_header(true)
441 .build(file)?;
442 for batch in reader {
443 batches.push(batch?);
444 }
445 }
446 Ok(batches)
447}
448
449fn read_json_arrow(path: &Path, format: SourceFormat) -> Result<Vec<RecordBatch>> {
450 let data = std::fs::read(path)?;
451 let data = if format == SourceFormat::Json {
452 expand_json_document_to_jsonl(&data)?
453 } else {
454 data
455 };
456 read_json_infer_from_bytes(&data)
457}
458
459fn read_json_infer_from_bytes(data: &[u8]) -> Result<Vec<RecordBatch>> {
460 use arrow::json::reader::infer_json_schema_from_seekable;
461 let mut cursor = std::io::Cursor::new(data);
462 let (schema, _n) = infer_json_schema_from_seekable(&mut cursor, Some(1024))?;
463 cursor.set_position(0);
464 let schema = Arc::new(schema);
465 let reader = arrow::json::ReaderBuilder::new(schema).build(cursor)?;
466 let mut batches = Vec::new();
467 for batch in reader {
468 batches.push(batch?);
469 }
470 Ok(batches)
471}
472
473fn read_arrow_ipc_file(path: &Path) -> Result<Vec<RecordBatch>> {
474 let file = File::open(path)?;
475 let reader = arrow::ipc::reader::FileReader::try_new(file, None)?;
476 let mut batches = Vec::new();
477 for batch in reader {
478 batches.push(batch?);
479 }
480 Ok(batches)
481}
482
483fn read_arrow_ipc_stream(path: &Path) -> Result<Vec<RecordBatch>> {
484 let file = File::open(path)?;
485 let reader = arrow::ipc::reader::StreamReader::try_new(file, None)?;
486 let mut batches = Vec::new();
487 for batch in reader {
488 batches.push(batch?);
489 }
490 Ok(batches)
491}
492
493fn read_arrow_ipc_auto(path: &Path) -> Result<Vec<RecordBatch>> {
495 match read_arrow_ipc_file(path) {
496 Ok(batches) => Ok(batches),
497 Err(file_err) => read_arrow_ipc_stream(path).with_context(|| {
498 format!(
499 "Arrow IPC file and stream readers both failed for {} (file error: {file_err})",
500 path.display()
501 )
502 }),
503 }
504}
505
506pub fn parse_hive_partitions(
508 file: &Path,
509 root: &Path,
510) -> std::collections::HashMap<String, String> {
511 let mut out = std::collections::HashMap::new();
512 let rel = file.strip_prefix(root).unwrap_or(file);
513 for comp in rel.components() {
514 if let std::path::Component::Normal(os) = comp {
515 let s = os.to_string_lossy();
516 if let Some((k, v)) = s.split_once('=') {
517 if !k.is_empty() {
518 out.insert(k.to_string(), v.to_string());
519 }
520 }
521 }
522 }
523 out
524}
525
526fn path_matches_require_partitions(
527 file: &Path,
528 root: &Path,
529 require: &std::collections::HashMap<String, String>,
530) -> bool {
531 if require.is_empty() {
532 return true;
533 }
534 let parts = parse_hive_partitions(file, root);
535 require
536 .iter()
537 .all(|(k, v)| parts.get(k).map(|pv| pv == v).unwrap_or(false))
538}
539
540fn inject_hive_partitions(
542 batch: RecordBatch,
543 file: &Path,
544 root: &Path,
545 partition_by: &[String],
546) -> Result<RecordBatch> {
547 if partition_by.is_empty() {
548 return Ok(batch);
549 }
550 let parts = parse_hive_partitions(file, root);
551 let n = batch.num_rows();
552 let mut fields: Vec<arrow::datatypes::Field> = batch
553 .schema()
554 .fields()
555 .iter()
556 .map(|f| f.as_ref().clone())
557 .collect();
558 let mut columns: Vec<ArrayRef> = batch.columns().to_vec();
559
560 for key in partition_by {
561 if batch.schema().index_of(key).is_ok() {
562 continue;
564 }
565 let val = parts.get(key).map(|s| s.as_str());
566 let mut b = StringBuilder::with_capacity(n, n * 8);
567 for _ in 0..n {
568 match val {
569 Some(v) => b.append_value(v),
570 None => b.append_null(),
571 }
572 }
573 fields.push(Field::new(key.as_str(), DataType::Utf8, true));
574 columns.push(Arc::new(b.finish()) as ArrayRef);
575 }
576
577 let schema = Arc::new(Schema::new(fields));
578 Ok(RecordBatch::try_new(schema, columns)?)
579}
580
581fn inject_source_path_column(batch: RecordBatch, file: &Path) -> Result<RecordBatch> {
583 if batch.schema().index_of("_source_path").is_ok() {
584 return Ok(batch);
585 }
586 let n = batch.num_rows();
587 let path_str = file.to_string_lossy();
588 let mut b = StringBuilder::with_capacity(n, n * path_str.len().max(16));
589 for _ in 0..n {
590 b.append_value(path_str.as_ref());
591 }
592 let mut fields: Vec<Field> = batch
593 .schema()
594 .fields()
595 .iter()
596 .map(|f| f.as_ref().clone())
597 .collect();
598 fields.push(Field::new("_source_path", DataType::Utf8, false));
599 let mut columns = batch.columns().to_vec();
600 columns.push(Arc::new(b.finish()) as ArrayRef);
601 Ok(RecordBatch::try_new(
602 Arc::new(Schema::new(fields)),
603 columns,
604 )?)
605}
606
607fn read_line_oriented(path: &Path) -> Result<RecordBatch> {
610 let file = File::open(path)?;
611 let reader = BufReader::new(file);
612 let mut line_nos = Int64Builder::new();
613 let mut contents = StringBuilder::new();
614 let mut n: i64 = 0;
615 for line in reader.lines() {
616 let line = line?;
617 n += 1;
618 line_nos.append_value(n);
619 contents.append_value(line);
620 }
621 let schema = Arc::new(Schema::new(vec![
622 Field::new("line_no", DataType::Int64, false),
623 Field::new("content", DataType::Utf8, false),
624 ]));
625 Ok(RecordBatch::try_new(
626 schema,
627 vec![
628 Arc::new(line_nos.finish()) as ArrayRef,
629 Arc::new(contents.finish()) as ArrayRef,
630 ],
631 )?)
632}
633
634fn read_toml(path: &Path, rows_key: Option<&str>) -> Result<RecordBatch> {
635 let text = std::fs::read_to_string(path)?;
636 let value: toml::Value = text
637 .parse()
638 .with_context(|| format!("Invalid TOML: {}", path.display()))?;
639
640 let rows: Vec<toml::map::Map<String, toml::Value>> = match &value {
641 toml::Value::Table(table) => {
642 if let Some(key) = rows_key {
643 match table.get(key) {
644 Some(toml::Value::Array(arr)) => array_of_tables(arr)?,
645 other => bail!(
646 "toml_rows_key '{}' is not an array of tables (got {:?})",
647 key,
648 other.map(|v| v.type_str())
649 ),
650 }
651 } else if let Some((_k, toml::Value::Array(arr))) = table
652 .iter()
653 .find(|(_, v)| matches!(v, toml::Value::Array(a) if a.iter().all(|x| x.is_table())))
654 {
655 array_of_tables(arr)?
656 } else {
657 vec![table.clone()]
659 }
660 }
661 toml::Value::Array(arr) => array_of_tables(arr)?,
662 other => bail!("Unsupported top-level TOML type: {}", other.type_str()),
663 };
664
665 if rows.is_empty() {
666 let schema = Arc::new(Schema::new(vec![Field::new("empty", DataType::Utf8, true)]));
667 return Ok(RecordBatch::try_new(
668 schema,
669 vec![Arc::new(StringBuilder::new().finish()) as ArrayRef],
670 )?);
671 }
672
673 let mut col_names: Vec<String> = Vec::new();
675 for row in &rows {
676 for k in row.keys() {
677 if !col_names.iter().any(|c| c == k) {
678 col_names.push(k.clone());
679 }
680 }
681 }
682
683 let mut builders: Vec<StringBuilder> = col_names
684 .iter()
685 .map(|_| StringBuilder::with_capacity(rows.len(), rows.len() * 16))
686 .collect();
687
688 for row in &rows {
689 for (i, col) in col_names.iter().enumerate() {
690 match row.get(col) {
691 Some(v) => builders[i].append_value(toml_value_to_string(v)),
692 None => builders[i].append_null(),
693 }
694 }
695 }
696
697 let fields: Vec<Field> = col_names
698 .iter()
699 .map(|c| Field::new(c.as_str(), DataType::Utf8, true))
700 .collect();
701 let schema = Arc::new(Schema::new(fields));
702 let arrays: Vec<ArrayRef> = builders
703 .into_iter()
704 .map(|mut b| Arc::new(b.finish()) as ArrayRef)
705 .collect();
706 Ok(RecordBatch::try_new(schema, arrays)?)
707}
708
709fn array_of_tables(arr: &[toml::Value]) -> Result<Vec<toml::map::Map<String, toml::Value>>> {
710 let mut rows = Vec::with_capacity(arr.len());
711 for (i, item) in arr.iter().enumerate() {
712 match item {
713 toml::Value::Table(t) => rows.push(t.clone()),
714 other => bail!(
715 "TOML array element {} is not a table (got {})",
716 i,
717 other.type_str()
718 ),
719 }
720 }
721 Ok(rows)
722}
723
724fn toml_value_to_string(v: &toml::Value) -> String {
725 match v {
726 toml::Value::String(s) => s.clone(),
727 toml::Value::Integer(i) => i.to_string(),
728 toml::Value::Float(f) => f.to_string(),
729 toml::Value::Boolean(b) => b.to_string(),
730 toml::Value::Datetime(d) => d.to_string(),
731 other => other.to_string(),
732 }
733}
734
735#[cfg(test)]
736mod tests {
737 use super::*;
738 use arrow::array::{Array, Int64Array, StringArray};
739 use arrow::datatypes::{DataType, Field, Schema};
740 use std::sync::Arc;
741
742 #[tokio::test]
743 async fn test_lake_scanner_multi_format() -> Result<()> {
744 let temp_dir = tempfile::tempdir()?;
745 let dir_path = temp_dir.path();
746
747 let schema = Arc::new(Schema::new(vec![
748 Field::new("id", DataType::Int64, true),
749 Field::new("name", DataType::Utf8, true),
750 ]));
751
752 std::fs::write(
753 dir_path.join("file1.jsonl"),
754 b"{\"id\": 1, \"name\": \"Alice\"}\n{\"id\": 2, \"name\": \"Bob\"}\n",
755 )?;
756 std::fs::write(dir_path.join("file2.csv"), b"id,name\n3,Charlie\n4,Dave\n")?;
757
758 let batch = RecordBatch::try_new(
759 schema.clone(),
760 vec![
761 Arc::new(Int64Array::from(vec![5, 6])),
762 Arc::new(StringArray::from(vec!["Eve", "Frank"])),
763 ],
764 )?;
765 let file = File::create(dir_path.join("file3.parquet"))?;
766 let mut writer =
767 parquet::arrow::arrow_writer::ArrowWriter::try_new(file, schema.clone(), None)?;
768 writer.write(&batch)?;
769 writer.close()?;
770
771 std::fs::write(dir_path.join("app.log"), "info boot\nwarn disk\n")?;
772 std::fs::write(
773 dir_path.join("meta.toml"),
774 r#"
775[[records]]
776id = "1"
777name = "toml_a"
778[[records]]
779id = "2"
780name = "toml_b"
781"#,
782 )?;
783
784 let scanner = LakeScanner::new(vec!["id".to_string(), "name".to_string()]);
785 let batches = scanner.scan_path(dir_path, schema.clone()).await?;
786 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
787 assert_eq!(total_rows, 10);
789 Ok(())
790 }
791
792 #[tokio::test]
793 async fn test_scan_request_jsonl() -> Result<()> {
794 let temp = tempfile::tempdir()?;
795 let path = temp.path().join("trades.jsonl");
796 std::fs::write(
797 &path,
798 r#"{"ticker":"NVDA","price":1.0}
799{"ticker":"AAPL","price":2.0}
800"#,
801 )?;
802
803 let fm = StagingFrontmatter {
804 source_format: Some(SourceFormat::Jsonl),
805 scan_path: Some(path.file_name().unwrap().to_string_lossy().into()),
806 ..Default::default()
807 };
808 let req = ScanRequest::from_frontmatter(temp.path(), &fm)?;
809 let scanner = LakeScanner::from_request(&req);
810 let batches = scanner.scan(&req).await?;
811 let rows: usize = batches.iter().map(|b| b.num_rows()).sum();
812 assert_eq!(rows, 2);
813 Ok(())
814 }
815
816 #[tokio::test]
817 async fn test_hive_partition_inject_and_filter() -> Result<()> {
818 let temp = tempfile::tempdir()?;
819 let root = temp.path();
820 let dir = root.join("symbol=NVDA").join("timeframe=1m");
821 std::fs::create_dir_all(&dir)?;
822 let schema = Arc::new(Schema::new(vec![Field::new(
824 "close",
825 DataType::Float64,
826 true,
827 )]));
828 let batch = RecordBatch::try_new(
829 schema,
830 vec![Arc::new(arrow::array::Float64Array::from(vec![1.0, 2.0])) as ArrayRef],
831 )?;
832 let path = dir.join("chunk.arrow");
833 {
834 let file = File::create(&path)?;
835 let mut writer = arrow::ipc::writer::StreamWriter::try_new(file, &batch.schema())?;
836 writer.write(&batch)?;
837 writer.finish()?;
838 }
839 let other = root.join("symbol=AAPL").join("timeframe=1d");
841 std::fs::create_dir_all(&other)?;
842 {
843 let file = File::create(other.join("chunk.arrow"))?;
844 let mut writer = arrow::ipc::writer::StreamWriter::try_new(file, &batch.schema())?;
845 writer.write(&batch)?;
846 writer.finish()?;
847 }
848
849 let mut require = std::collections::HashMap::new();
850 require.insert("timeframe".into(), "1m".into());
851 let fm = StagingFrontmatter {
852 source_format: Some(SourceFormat::ArrowIpcStream),
853 scan_path: Some(".".into()),
854 partition_by: Some(vec!["symbol".into(), "timeframe".into()]),
855 require_partitions: Some(require),
856 ..Default::default()
857 };
858 let req = ScanRequest::from_frontmatter(root, &fm)?;
859 let scanner = LakeScanner::from_request(&req);
860 let batches = scanner.scan(&req).await?;
861 let rows: usize = batches.iter().map(|b| b.num_rows()).sum();
862 assert_eq!(rows, 2);
863 let schema = batches[0].schema();
864 assert!(schema.index_of("timeframe").is_ok());
865 assert!(schema.index_of("symbol").is_ok());
866 Ok(())
867 }
868
869 #[test]
870 fn test_line_oriented_and_toml() -> Result<()> {
871 let temp = tempfile::tempdir()?;
872 let txt = temp.path().join("llms.txt");
873 std::fs::write(&txt, "# Title\n\nSome doc line\n")?;
874 let batch = read_line_oriented(&txt)?;
875 assert_eq!(batch.num_rows(), 3);
876 assert_eq!(batch.schema().field(0).name(), "line_no");
877 assert_eq!(batch.schema().field(1).name(), "content");
878
879 let toml_path = temp.path().join("cfg.toml");
880 std::fs::write(
881 &toml_path,
882 r#"
883[[items]]
884k = "a"
885[[items]]
886k = "b"
887"#,
888 )?;
889 let tbatch = read_toml(&toml_path, Some("items"))?;
890 assert_eq!(tbatch.num_rows(), 2);
891 let col = tbatch
892 .column(0)
893 .as_any()
894 .downcast_ref::<StringArray>()
895 .unwrap();
896 assert!(col.value(0) == "a" || col.value(1) == "a");
897 Ok(())
898 }
899
900 #[tokio::test]
901 async fn test_path_glob_filters_artifacts() -> Result<()> {
902 let temp = tempfile::tempdir()?;
903 let root = temp.path();
904 let d1 = root
906 .join("domain=x.com")
907 .join("report_date=2026-07-29")
908 .join("run_id=r1")
909 .join("raw_snoop");
910 std::fs::create_dir_all(&d1)?;
911 let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)]));
913 for (name, ids) in [
914 ("crawlplan.parquet", vec![1i64]),
915 ("other.parquet", vec![9i64]),
916 ] {
917 let batch =
918 RecordBatch::try_new(schema.clone(), vec![Arc::new(Int64Array::from(ids))])?;
919 let f = File::create(d1.join(name))?;
920 let mut w =
921 parquet::arrow::arrow_writer::ArrowWriter::try_new(f, schema.clone(), None)?;
922 w.write(&batch)?;
923 w.close()?;
924 }
925
926 let fm = StagingFrontmatter {
927 source_format: Some(SourceFormat::Parquet),
928 scan_path: Some(".".into()),
929 path_glob: Some(vec!["**/crawlplan.parquet".into()]),
930 partition_by: Some(vec!["domain".into()]),
931 inject_source_path: Some(true),
932 ..Default::default()
933 };
934 let req = ScanRequest::from_frontmatter(root, &fm)?;
935 let scanner = LakeScanner::from_request(&req);
936 let batches = scanner.scan(&req).await?;
937 let rows: usize = batches.iter().map(|b| b.num_rows()).sum();
938 assert_eq!(rows, 1, "only crawlplan.parquet rows");
939 assert!(batches[0].schema().index_of("domain").is_ok());
940 Ok(())
941 }
942
943 #[tokio::test]
944 async fn test_protobuf_opaque_scan() -> Result<()> {
945 let temp = tempfile::tempdir()?;
946 let pb = temp.path().join("msg.pb");
947 std::fs::write(&pb, b"\x08\x96\x01")?; let fm = StagingFrontmatter {
949 source_format: Some(SourceFormat::Protobuf),
950 scan_path: Some(pb.file_name().unwrap().to_string_lossy().into()),
951 ..Default::default()
952 };
953 let req = ScanRequest::from_frontmatter(temp.path(), &fm)?;
954 assert_eq!(
955 req.protobuf_max_payload_bytes,
956 DEFAULT_PROTOBUF_MAX_PAYLOAD_BYTES
957 );
958 let scanner = LakeScanner::from_request(&req);
959 let batches = scanner.scan(&req).await?;
960 assert_eq!(batches.len(), 1);
961 assert_eq!(batches[0].num_rows(), 1);
962 assert_eq!(batches[0].schema().field(0).name(), "_source_path");
963 assert_eq!(batches[0].schema().field(1).name(), "payload");
964 Ok(())
965 }
966
967 #[tokio::test]
968 async fn test_protobuf_respects_max_payload_bytes() -> Result<()> {
969 let temp = tempfile::tempdir()?;
970 let pb = temp.path().join("big.pb");
971 std::fs::write(&pb, vec![0u8; 64])?;
972 let fm = StagingFrontmatter {
973 source_format: Some(SourceFormat::Protobuf),
974 scan_path: Some("big.pb".into()),
975 ..Default::default()
976 };
977 let scan_cfg = ScanConfig {
978 protobuf_max_payload_bytes: 16,
979 };
980 let req = ScanRequest::from_frontmatter_with_config(
981 temp.path(),
982 &fm,
983 HashMap::new(),
984 &scan_cfg,
985 )?;
986 let scanner = LakeScanner::from_request(&req);
987 let err = scanner.scan(&req).await.unwrap_err();
988 let full = format!("{err:#}");
990 assert!(
991 full.contains("E_RBT_PROTOBUF_TOO_LARGE"),
992 "expected size-cap error in chain, got: {full}"
993 );
994 assert!(full.contains("protobuf_max_payload_bytes"));
995 Ok(())
996 }
997
998 #[tokio::test]
999 async fn test_path_glob_single_star_is_not_recursive() -> Result<()> {
1000 let temp = tempfile::tempdir()?;
1001 let root = temp.path();
1002 let deep = root.join("a").join("b").join("raw");
1003 std::fs::create_dir_all(&deep)?;
1004 let shallow = root.join("raw");
1005 std::fs::create_dir_all(&shallow)?;
1006 let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)]));
1007 for (dir, id) in [(&deep, 1i64), (&shallow, 2i64)] {
1008 let batch =
1009 RecordBatch::try_new(schema.clone(), vec![Arc::new(Int64Array::from(vec![id]))])?;
1010 let f = File::create(dir.join("crawlplan.parquet"))?;
1011 let mut w =
1012 parquet::arrow::arrow_writer::ArrowWriter::try_new(f, schema.clone(), None)?;
1013 w.write(&batch)?;
1014 w.close()?;
1015 }
1016 let fm = StagingFrontmatter {
1017 source_format: Some(SourceFormat::Parquet),
1018 scan_path: Some(".".into()),
1019 path_glob: Some(vec!["*/crawlplan.parquet".into()]),
1021 ..Default::default()
1022 };
1023 let req = ScanRequest::from_frontmatter(root, &fm)?;
1024 let scanner = LakeScanner::from_request(&req);
1025 let batches = scanner.scan(&req).await?;
1026 let rows: usize = batches.iter().map(|b| b.num_rows()).sum();
1027 assert_eq!(rows, 1, "single-star glob must not match deep hive path");
1028 Ok(())
1029 }
1030
1031 #[tokio::test]
1032 async fn test_absolute_scan_path() -> Result<()> {
1033 let temp = tempfile::tempdir()?;
1034 let abs = temp.path().join("abs.jsonl");
1035 std::fs::write(&abs, b"{\"a\":1}\n{\"a\":2}\n")?;
1036 let fm = StagingFrontmatter {
1037 source_format: Some(SourceFormat::Jsonl),
1038 scan_path: Some(abs.to_string_lossy().into()),
1039 ..Default::default()
1040 };
1041 let proj = tempfile::tempdir()?;
1043 let req = ScanRequest::from_frontmatter(proj.path(), &fm)?;
1044 let scanner = LakeScanner::from_request(&req);
1045 let batches = scanner.scan(&req).await?;
1046 assert_eq!(batches.iter().map(|b| b.num_rows()).sum::<usize>(), 2);
1047 Ok(())
1048 }
1049}