rbt/scan/
mod.rs

1//! `rbt::scan`: Multi-format bronze lake scanners producing Arrow `RecordBatch`es.
2//!
3//! Formats: JSONL (jshift), JSON, Parquet, CSV, Arrow IPC (file/stream), log, txt, TOML.
4
5use anyhow::{anyhow, bail, Context, Result};
6use arrow::array::{ArrayRef, Int64Builder, StringBuilder};
7use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
8use arrow::record_batch::RecordBatch;
9use crate::core::frontmatter::{resolve_scan_path, SourceFormat, StagingFrontmatter};
10use std::fs::File;
11use std::io::{BufRead, BufReader};
12use std::path::{Path, PathBuf};
13use std::sync::Arc;
14
15/// Scan request derived from staging frontmatter + project root.
16#[derive(Debug, Clone)]
17pub struct ScanRequest {
18    pub project_dir: PathBuf,
19    pub scan_path: String,
20    pub format: SourceFormat,
21    /// jshift / projection paths (JSONL selective extract).
22    pub paths: Vec<String>,
23    /// TOML array-of-tables key (optional).
24    pub toml_rows_key: Option<String>,
25    /// Hive-style partition keys to inject as Utf8 columns from the file path.
26    pub partition_by: Vec<String>,
27    /// Keep only files whose hive path matches these partition values.
28    pub require_partitions: std::collections::HashMap<String, String>,
29    /// Inject `_source_path` column (absolute path of the bronze file).
30    pub inject_source_path: bool,
31}
32
33impl ScanRequest {
34    pub fn from_frontmatter(
35        project_dir: impl AsRef<Path>,
36        fm: &StagingFrontmatter,
37    ) -> Result<Self> {
38        let scan_path = fm
39            .scan_path
40            .as_ref()
41            .context("frontmatter.scan_path is required for bronze scan")?
42            .clone();
43        let format = fm.resolve_format()?;
44        Ok(Self {
45            project_dir: project_dir.as_ref().to_path_buf(),
46            scan_path,
47            format,
48            paths: fm.paths.clone().unwrap_or_default(),
49            toml_rows_key: fm.toml_rows_key.clone(),
50            partition_by: fm.partition_by.clone().unwrap_or_default(),
51            require_partitions: fm.require_partitions.clone().unwrap_or_default(),
52            inject_source_path: fm.inject_source_path.unwrap_or(false),
53        })
54    }
55
56    pub fn resolved_path(&self) -> PathBuf {
57        resolve_scan_path(&self.project_dir, &self.scan_path)
58    }
59}
60
61/// Multi-format lake scanner.
62pub struct LakeScanner {
63    pub paths: Vec<String>,
64}
65
66impl LakeScanner {
67    pub fn new(paths: Vec<String>) -> Self {
68        Self { paths }
69    }
70
71    pub fn from_request(req: &ScanRequest) -> Self {
72        Self {
73            paths: req.paths.clone(),
74        }
75    }
76
77    /// Scan using a full [`ScanRequest`] (format-aware).
78    pub async fn scan(&self, req: &ScanRequest) -> Result<Vec<RecordBatch>> {
79        let root = req.resolved_path();
80        if !root.exists() {
81            bail!(
82                "Bronze scan_path does not exist: {} (resolved from '{}')",
83                root.display(),
84                req.scan_path
85            );
86        }
87
88        let mut files = Vec::new();
89        collect_files_for_format(&root, req.format, &mut files)?;
90        if !req.require_partitions.is_empty() {
91            files.retain(|f| path_matches_require_partitions(f, &root, &req.require_partitions));
92        }
93        if files.is_empty() {
94            bail!(
95                "No {} files found under {} (after partition filters)",
96                req.format.as_str(),
97                root.display()
98            );
99        }
100
101        tracing::info!(
102            "Bronze scan: {} file(s) under {} (format={})",
103            files.len(),
104            root.display(),
105            req.format
106        );
107
108        let mut batches = Vec::new();
109        for file_path in files {
110            let file_batches = self
111                .read_file(&file_path, req)
112                .with_context(|| format!("Failed reading bronze file {}", file_path.display()))?;
113            for batch in file_batches {
114                let mut batch = if req.partition_by.is_empty() {
115                    batch
116                } else {
117                    inject_hive_partitions(batch, &file_path, &root, &req.partition_by)?
118                };
119                if req.inject_source_path {
120                    batch = inject_source_path_column(batch, &file_path)?;
121                }
122                batches.push(batch);
123            }
124        }
125        Ok(batches)
126    }
127
128    /// Legacy helper: recurse path and read by extension with an explicit schema
129    /// (used by unit tests and jshift-projected JSONL).
130    pub async fn scan_path(
131        &self,
132        path: impl AsRef<Path>,
133        schema: SchemaRef,
134    ) -> Result<Vec<RecordBatch>> {
135        let mut files = Vec::new();
136        collect_files(path.as_ref(), &mut files)?;
137        let mut batches = Vec::new();
138
139        for file_path in files {
140            let ext = file_path
141                .extension()
142                .and_then(|e| e.to_str())
143                .unwrap_or("")
144                .to_ascii_lowercase();
145            match ext.as_str() {
146                "parquet" => batches.extend(read_parquet(&file_path, Some(schema.clone()))?),
147                "json" | "jsonl" | "ndjson" => {
148                    let bytes = std::fs::read(&file_path)?;
149                    let extractor = crate::json::JShiftExtractor::new(self.paths.clone());
150                    batches.push(extractor.extract_jsonl(&bytes, schema.clone())?);
151                }
152                "csv" | "tsv" => batches.extend(read_csv(&file_path, Some(schema.clone()))?),
153                "log" | "txt" | "text" | "md" => {
154                    batches.push(read_line_oriented(&file_path)?);
155                }
156                "toml" => batches.push(read_toml(&file_path, None)?),
157                "arrow" | "arrows" | "ipc" | "feather" => {
158                    batches.extend(read_arrow_ipc_auto(&file_path)?);
159                }
160                _ => {
161                    tracing::debug!("Skipping unsupported file: {:?}", file_path);
162                }
163            }
164        }
165        Ok(batches)
166    }
167
168    fn read_file(&self, file_path: &Path, req: &ScanRequest) -> Result<Vec<RecordBatch>> {
169        match req.format {
170            SourceFormat::Parquet => read_parquet(file_path, None),
171            SourceFormat::Csv => read_csv(file_path, None),
172            SourceFormat::Jsonl | SourceFormat::Json => {
173                if !self.paths.is_empty() {
174                    let schema = utf8_schema_from_paths(&self.paths);
175                    let bytes = std::fs::read(file_path)?;
176                    let extractor = crate::json::JShiftExtractor::new(self.paths.clone());
177                    // JSONL extractor works line-wise; for single JSON array, expand first
178                    if req.format == SourceFormat::Json {
179                        let expanded = expand_json_document_to_jsonl(&bytes)?;
180                        Ok(vec![extractor.extract_jsonl(&expanded, schema)?])
181                    } else {
182                        Ok(vec![extractor.extract_jsonl(&bytes, schema)?])
183                    }
184                } else {
185                    // Schema-free path: use Arrow JSON reader (line-delimited)
186                    read_json_arrow(file_path, req.format)
187                }
188            }
189            // Real lakes often write stream IPC with a `.arrow` extension.
190            SourceFormat::ArrowIpc | SourceFormat::ArrowIpcStream => {
191                read_arrow_ipc_auto(file_path)
192            }
193            SourceFormat::Log | SourceFormat::Txt => Ok(vec![read_line_oriented(file_path)?]),
194            SourceFormat::Toml => Ok(vec![read_toml(
195                file_path,
196                req.toml_rows_key.as_deref(),
197            )?]),
198        }
199    }
200}
201
202fn utf8_schema_from_paths(paths: &[String]) -> SchemaRef {
203    Arc::new(Schema::new(
204        paths
205            .iter()
206            .map(|p| Field::new(p.as_str(), DataType::Utf8, true))
207            .collect::<Vec<_>>(),
208    ))
209}
210
211fn expand_json_document_to_jsonl(bytes: &[u8]) -> Result<Vec<u8>> {
212    let v: serde_json::Value = serde_json::from_slice(bytes)
213        .map_err(|e| anyhow!("Invalid JSON document: {}", e))?;
214    match v {
215        serde_json::Value::Array(items) => {
216            let mut out = Vec::new();
217            for item in items {
218                out.extend(serde_json::to_vec(&item)?);
219                out.push(b'\n');
220            }
221            Ok(out)
222        }
223        other => {
224            let mut out = serde_json::to_vec(&other)?;
225            out.push(b'\n');
226            Ok(out)
227        }
228    }
229}
230
231fn collect_files(path: &Path, files: &mut Vec<PathBuf>) -> Result<()> {
232    if path.is_file() {
233        files.push(path.to_path_buf());
234    } else if path.is_dir() {
235        for entry in std::fs::read_dir(path)? {
236            let entry = entry?;
237            collect_files(&entry.path(), files)?;
238        }
239    }
240    Ok(())
241}
242
243fn extension_matches(format: SourceFormat, ext: &str) -> bool {
244    let ext = ext.to_ascii_lowercase();
245    match format {
246        SourceFormat::Jsonl => matches!(ext.as_str(), "jsonl" | "ndjson"),
247        SourceFormat::Json => ext == "json",
248        SourceFormat::Parquet => matches!(ext.as_str(), "parquet" | "pq"),
249        SourceFormat::Csv => matches!(ext.as_str(), "csv" | "tsv"),
250        SourceFormat::ArrowIpc => matches!(ext.as_str(), "arrow" | "arrows" | "ipc" | "feather"),
251        SourceFormat::ArrowIpcStream => {
252            matches!(ext.as_str(), "arrow" | "arrows" | "ipc" | "arrows_stream" | "ipc_stream")
253        }
254        SourceFormat::Log => ext == "log",
255        SourceFormat::Txt => matches!(ext.as_str(), "txt" | "text" | "md"),
256        SourceFormat::Toml => ext == "toml",
257    }
258}
259
260fn collect_files_for_format(
261    path: &Path,
262    format: SourceFormat,
263    files: &mut Vec<PathBuf>,
264) -> Result<()> {
265    if path.is_file() {
266        files.push(path.to_path_buf());
267        return Ok(());
268    }
269    if path.is_dir() {
270        let mut all = Vec::new();
271        collect_files(path, &mut all)?;
272        for f in all {
273            let ext = f.extension().and_then(|e| e.to_str()).unwrap_or("");
274            if extension_matches(format, ext) {
275                files.push(f);
276            }
277        }
278        // If directory has files but none matched extension, and format is Log/Txt,
279        // still allow extensionless? Skip for now.
280        return Ok(());
281    }
282    bail!("scan path is neither file nor directory: {}", path.display());
283}
284
285fn read_parquet(path: &Path, projection: Option<SchemaRef>) -> Result<Vec<RecordBatch>> {
286    let file = File::open(path)?;
287    let builder = parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder::try_new(file)?;
288    let reader = if let Some(schema) = projection {
289        let schema_descr = builder.metadata().file_metadata().schema_descr_ptr();
290        let arrow_schema = builder.schema();
291        let mut indices = Vec::new();
292        for field in schema.fields() {
293            if let Some((idx, _)) =
294                parquet::arrow::parquet_column(&schema_descr, arrow_schema, field.name())
295            {
296                indices.push(idx);
297            }
298        }
299        let mask = parquet::arrow::ProjectionMask::leaves(&schema_descr, indices);
300        builder.with_projection(mask).build()?
301    } else {
302        builder.build()?
303    };
304    let mut batches = Vec::new();
305    for batch in reader {
306        batches.push(batch?);
307    }
308    Ok(batches)
309}
310
311fn read_csv(path: &Path, schema: Option<SchemaRef>) -> Result<Vec<RecordBatch>> {
312    let file = File::open(path)?;
313    let mut batches = Vec::new();
314    if let Some(schema) = schema {
315        let reader = arrow::csv::ReaderBuilder::new(schema)
316            .with_header(true)
317            .build(file)?;
318        for batch in reader {
319            batches.push(batch?);
320        }
321    } else {
322        // Infer schema from first rows
323        let format = arrow::csv::reader::Format::default().with_header(true);
324        let (inferred, _) = format.infer_schema(File::open(path)?, Some(1024))?;
325        let schema = Arc::new(inferred);
326        let file = File::open(path)?;
327        let reader = arrow::csv::ReaderBuilder::new(schema)
328            .with_header(true)
329            .build(file)?;
330        for batch in reader {
331            batches.push(batch?);
332        }
333    }
334    Ok(batches)
335}
336
337fn read_json_arrow(path: &Path, format: SourceFormat) -> Result<Vec<RecordBatch>> {
338    let data = std::fs::read(path)?;
339    let data = if format == SourceFormat::Json {
340        expand_json_document_to_jsonl(&data)?
341    } else {
342        data
343    };
344    read_json_infer_from_bytes(&data)
345}
346
347fn read_json_infer_from_bytes(data: &[u8]) -> Result<Vec<RecordBatch>> {
348    use arrow::json::reader::infer_json_schema_from_seekable;
349    let mut cursor = std::io::Cursor::new(data);
350    let (schema, _n) = infer_json_schema_from_seekable(&mut cursor, Some(1024))?;
351    cursor.set_position(0);
352    let schema = Arc::new(schema);
353    let reader = arrow::json::ReaderBuilder::new(schema).build(cursor)?;
354    let mut batches = Vec::new();
355    for batch in reader {
356        batches.push(batch?);
357    }
358    Ok(batches)
359}
360
361fn read_arrow_ipc_file(path: &Path) -> Result<Vec<RecordBatch>> {
362    let file = File::open(path)?;
363    let reader = arrow::ipc::reader::FileReader::try_new(file, None)?;
364    let mut batches = Vec::new();
365    for batch in reader {
366        batches.push(batch?);
367    }
368    Ok(batches)
369}
370
371fn read_arrow_ipc_stream(path: &Path) -> Result<Vec<RecordBatch>> {
372    let file = File::open(path)?;
373    let reader = arrow::ipc::reader::StreamReader::try_new(file, None)?;
374    let mut batches = Vec::new();
375    for batch in reader {
376        batches.push(batch?);
377    }
378    Ok(batches)
379}
380
381/// Prefer random-access IPC file; fall back to stream (common for `.arrow` lake dumps).
382fn read_arrow_ipc_auto(path: &Path) -> Result<Vec<RecordBatch>> {
383    match read_arrow_ipc_file(path) {
384        Ok(batches) => Ok(batches),
385        Err(file_err) => read_arrow_ipc_stream(path).with_context(|| {
386            format!(
387                "Arrow IPC file and stream readers both failed for {} (file error: {file_err})",
388                path.display()
389            )
390        }),
391    }
392}
393
394/// Parse `key=value` segments from `file` relative to `root`.
395pub fn parse_hive_partitions(file: &Path, root: &Path) -> std::collections::HashMap<String, String> {
396    let mut out = std::collections::HashMap::new();
397    let rel = file.strip_prefix(root).unwrap_or(file);
398    for comp in rel.components() {
399        if let std::path::Component::Normal(os) = comp {
400            let s = os.to_string_lossy();
401            if let Some((k, v)) = s.split_once('=') {
402                if !k.is_empty() {
403                    out.insert(k.to_string(), v.to_string());
404                }
405            }
406        }
407    }
408    out
409}
410
411fn path_matches_require_partitions(
412    file: &Path,
413    root: &Path,
414    require: &std::collections::HashMap<String, String>,
415) -> bool {
416    if require.is_empty() {
417        return true;
418    }
419    let parts = parse_hive_partitions(file, root);
420    require.iter().all(|(k, v)| parts.get(k).map(|pv| pv == v).unwrap_or(false))
421}
422
423/// Append hive partition columns (Utf8) for keys in `partition_by` that are not already present.
424fn inject_hive_partitions(
425    batch: RecordBatch,
426    file: &Path,
427    root: &Path,
428    partition_by: &[String],
429) -> Result<RecordBatch> {
430    if partition_by.is_empty() {
431        return Ok(batch);
432    }
433    let parts = parse_hive_partitions(file, root);
434    let n = batch.num_rows();
435    let mut fields: Vec<arrow::datatypes::Field> =
436        batch.schema().fields().iter().map(|f| f.as_ref().clone()).collect();
437    let mut columns: Vec<ArrayRef> = batch.columns().to_vec();
438
439    for key in partition_by {
440        if batch.schema().index_of(key).is_ok() {
441            // Column already in payload (e.g. symbol); keep file data, do not overwrite.
442            continue;
443        }
444        let val = parts.get(key).map(|s| s.as_str());
445        let mut b = StringBuilder::with_capacity(n, n * 8);
446        for _ in 0..n {
447            match val {
448                Some(v) => b.append_value(v),
449                None => b.append_null(),
450            }
451        }
452        fields.push(Field::new(key.as_str(), DataType::Utf8, true));
453        columns.push(Arc::new(b.finish()) as ArrayRef);
454    }
455
456    let schema = Arc::new(Schema::new(fields));
457    Ok(RecordBatch::try_new(schema, columns)?)
458}
459
460/// Append `_source_path` Utf8 with the absolute bronze file path (for latest-wins dedupe).
461fn inject_source_path_column(batch: RecordBatch, file: &Path) -> Result<RecordBatch> {
462    if batch.schema().index_of("_source_path").is_ok() {
463        return Ok(batch);
464    }
465    let n = batch.num_rows();
466    let path_str = file.to_string_lossy();
467    let mut b = StringBuilder::with_capacity(n, n * path_str.len().max(16));
468    for _ in 0..n {
469        b.append_value(path_str.as_ref());
470    }
471    let mut fields: Vec<Field> = batch
472        .schema()
473        .fields()
474        .iter()
475        .map(|f| f.as_ref().clone())
476        .collect();
477    fields.push(Field::new("_source_path", DataType::Utf8, false));
478    let mut columns = batch.columns().to_vec();
479    columns.push(Arc::new(b.finish()) as ArrayRef);
480    Ok(RecordBatch::try_new(Arc::new(Schema::new(fields)), columns)?)
481}
482
483/// Line-oriented bronze: `.log`, `.txt`, llms.txt-style docs.
484/// Schema: `line_no` Int64, `content` Utf8.
485fn read_line_oriented(path: &Path) -> Result<RecordBatch> {
486    let file = File::open(path)?;
487    let reader = BufReader::new(file);
488    let mut line_nos = Int64Builder::new();
489    let mut contents = StringBuilder::new();
490    let mut n: i64 = 0;
491    for line in reader.lines() {
492        let line = line?;
493        n += 1;
494        line_nos.append_value(n);
495        contents.append_value(line);
496    }
497    let schema = Arc::new(Schema::new(vec![
498        Field::new("line_no", DataType::Int64, false),
499        Field::new("content", DataType::Utf8, false),
500    ]));
501    Ok(RecordBatch::try_new(
502        schema,
503        vec![
504            Arc::new(line_nos.finish()) as ArrayRef,
505            Arc::new(contents.finish()) as ArrayRef,
506        ],
507    )?)
508}
509
510fn read_toml(path: &Path, rows_key: Option<&str>) -> Result<RecordBatch> {
511    let text = std::fs::read_to_string(path)?;
512    let value: toml::Value = text.parse().with_context(|| format!("Invalid TOML: {}", path.display()))?;
513
514    let rows: Vec<toml::map::Map<String, toml::Value>> = match &value {
515        toml::Value::Table(table) => {
516            if let Some(key) = rows_key {
517                match table.get(key) {
518                    Some(toml::Value::Array(arr)) => array_of_tables(arr)?,
519                    other => bail!(
520                        "toml_rows_key '{}' is not an array of tables (got {:?})",
521                        key,
522                        other.map(|v| v.type_str())
523                    ),
524                }
525            } else if let Some((_k, toml::Value::Array(arr))) = table
526                .iter()
527                .find(|(_, v)| matches!(v, toml::Value::Array(a) if a.iter().all(|x| x.is_table())))
528            {
529                array_of_tables(arr)?
530            } else {
531                // Single-row: top-level scalars / nested stringified
532                vec![table.clone()]
533            }
534        }
535        toml::Value::Array(arr) => array_of_tables(arr)?,
536        other => bail!("Unsupported top-level TOML type: {}", other.type_str()),
537    };
538
539    if rows.is_empty() {
540        let schema = Arc::new(Schema::new(vec![Field::new("empty", DataType::Utf8, true)]));
541        return Ok(RecordBatch::try_new(schema, vec![Arc::new(StringBuilder::new().finish()) as ArrayRef])?);
542    }
543
544    // Column union
545    let mut col_names: Vec<String> = Vec::new();
546    for row in &rows {
547        for k in row.keys() {
548            if !col_names.iter().any(|c| c == k) {
549                col_names.push(k.clone());
550            }
551        }
552    }
553
554    let mut builders: Vec<StringBuilder> =
555        col_names.iter().map(|_| StringBuilder::with_capacity(rows.len(), rows.len() * 16)).collect();
556
557    for row in &rows {
558        for (i, col) in col_names.iter().enumerate() {
559            match row.get(col) {
560                Some(v) => builders[i].append_value(toml_value_to_string(v)),
561                None => builders[i].append_null(),
562            }
563        }
564    }
565
566    let fields: Vec<Field> = col_names
567        .iter()
568        .map(|c| Field::new(c.as_str(), DataType::Utf8, true))
569        .collect();
570    let schema = Arc::new(Schema::new(fields));
571    let arrays: Vec<ArrayRef> = builders
572        .into_iter()
573        .map(|mut b| Arc::new(b.finish()) as ArrayRef)
574        .collect();
575    Ok(RecordBatch::try_new(schema, arrays)?)
576}
577
578fn array_of_tables(arr: &[toml::Value]) -> Result<Vec<toml::map::Map<String, toml::Value>>> {
579    let mut rows = Vec::with_capacity(arr.len());
580    for (i, item) in arr.iter().enumerate() {
581        match item {
582            toml::Value::Table(t) => rows.push(t.clone()),
583            other => bail!("TOML array element {} is not a table (got {})", i, other.type_str()),
584        }
585    }
586    Ok(rows)
587}
588
589fn toml_value_to_string(v: &toml::Value) -> String {
590    match v {
591        toml::Value::String(s) => s.clone(),
592        toml::Value::Integer(i) => i.to_string(),
593        toml::Value::Float(f) => f.to_string(),
594        toml::Value::Boolean(b) => b.to_string(),
595        toml::Value::Datetime(d) => d.to_string(),
596        other => other.to_string(),
597    }
598}
599
600#[cfg(test)]
601mod tests {
602    use super::*;
603    use arrow::array::{Array, Int64Array, StringArray};
604    use arrow::datatypes::{DataType, Field, Schema};
605    use std::sync::Arc;
606
607    #[tokio::test]
608    async fn test_lake_scanner_multi_format() -> Result<()> {
609        let temp_dir = tempfile::tempdir()?;
610        let dir_path = temp_dir.path();
611
612        let schema = Arc::new(Schema::new(vec![
613            Field::new("id", DataType::Int64, true),
614            Field::new("name", DataType::Utf8, true),
615        ]));
616
617        std::fs::write(
618            dir_path.join("file1.jsonl"),
619            b"{\"id\": 1, \"name\": \"Alice\"}\n{\"id\": 2, \"name\": \"Bob\"}\n",
620        )?;
621        std::fs::write(dir_path.join("file2.csv"), b"id,name\n3,Charlie\n4,Dave\n")?;
622
623        let batch = RecordBatch::try_new(
624            schema.clone(),
625            vec![
626                Arc::new(Int64Array::from(vec![5, 6])),
627                Arc::new(StringArray::from(vec!["Eve", "Frank"])),
628            ],
629        )?;
630        let file = File::create(dir_path.join("file3.parquet"))?;
631        let mut writer =
632            parquet::arrow::arrow_writer::ArrowWriter::try_new(file, schema.clone(), None)?;
633        writer.write(&batch)?;
634        writer.close()?;
635
636        std::fs::write(dir_path.join("app.log"), "info boot\nwarn disk\n")?;
637        std::fs::write(
638            dir_path.join("meta.toml"),
639            r#"
640[[records]]
641id = "1"
642name = "toml_a"
643[[records]]
644id = "2"
645name = "toml_b"
646"#,
647        )?;
648
649        let scanner = LakeScanner::new(vec!["id".to_string(), "name".to_string()]);
650        let batches = scanner.scan_path(dir_path, schema.clone()).await?;
651        let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
652        // 2 jsonl + 2 csv + 2 parquet + 2 log + 2 toml = 10
653        assert_eq!(total_rows, 10);
654        Ok(())
655    }
656
657    #[tokio::test]
658    async fn test_scan_request_jsonl() -> Result<()> {
659        let temp = tempfile::tempdir()?;
660        let path = temp.path().join("trades.jsonl");
661        std::fs::write(
662            &path,
663            r#"{"ticker":"NVDA","price":1.0}
664{"ticker":"AAPL","price":2.0}
665"#,
666        )?;
667
668        let fm = StagingFrontmatter {
669            source_format: Some(SourceFormat::Jsonl),
670            scan_path: Some(path.file_name().unwrap().to_string_lossy().into()),
671            ..Default::default()
672        };
673        let req = ScanRequest::from_frontmatter(temp.path(), &fm)?;
674        let scanner = LakeScanner::from_request(&req);
675        let batches = scanner.scan(&req).await?;
676        let rows: usize = batches.iter().map(|b| b.num_rows()).sum();
677        assert_eq!(rows, 2);
678        Ok(())
679    }
680
681    #[tokio::test]
682    async fn test_hive_partition_inject_and_filter() -> Result<()> {
683        let temp = tempfile::tempdir()?;
684        let root = temp.path();
685        let dir = root.join("symbol=NVDA").join("timeframe=1m");
686        std::fs::create_dir_all(&dir)?;
687        // minimal IPC stream with one utf8 column
688        let schema = Arc::new(Schema::new(vec![Field::new("close", DataType::Float64, true)]));
689        let batch = RecordBatch::try_new(
690            schema,
691            vec![Arc::new(arrow::array::Float64Array::from(vec![1.0, 2.0])) as ArrayRef],
692        )?;
693        let path = dir.join("chunk.arrow");
694        {
695            let file = File::create(&path)?;
696            let mut writer = arrow::ipc::writer::StreamWriter::try_new(file, &batch.schema())?;
697            writer.write(&batch)?;
698            writer.finish()?;
699        }
700        // noise partition that should be filtered out
701        let other = root.join("symbol=AAPL").join("timeframe=1d");
702        std::fs::create_dir_all(&other)?;
703        {
704            let file = File::create(other.join("chunk.arrow"))?;
705            let mut writer = arrow::ipc::writer::StreamWriter::try_new(file, &batch.schema())?;
706            writer.write(&batch)?;
707            writer.finish()?;
708        }
709
710        let mut require = std::collections::HashMap::new();
711        require.insert("timeframe".into(), "1m".into());
712        let fm = StagingFrontmatter {
713            source_format: Some(SourceFormat::ArrowIpcStream),
714            scan_path: Some(".".into()),
715            partition_by: Some(vec!["symbol".into(), "timeframe".into()]),
716            require_partitions: Some(require),
717            ..Default::default()
718        };
719        let req = ScanRequest::from_frontmatter(root, &fm)?;
720        let scanner = LakeScanner::from_request(&req);
721        let batches = scanner.scan(&req).await?;
722        let rows: usize = batches.iter().map(|b| b.num_rows()).sum();
723        assert_eq!(rows, 2);
724        let schema = batches[0].schema();
725        assert!(schema.index_of("timeframe").is_ok());
726        assert!(schema.index_of("symbol").is_ok());
727        Ok(())
728    }
729
730    #[test]
731    fn test_line_oriented_and_toml() -> Result<()> {
732        let temp = tempfile::tempdir()?;
733        let txt = temp.path().join("llms.txt");
734        std::fs::write(&txt, "# Title\n\nSome doc line\n")?;
735        let batch = read_line_oriented(&txt)?;
736        assert_eq!(batch.num_rows(), 3);
737        assert_eq!(batch.schema().field(0).name(), "line_no");
738        assert_eq!(batch.schema().field(1).name(), "content");
739
740        let toml_path = temp.path().join("cfg.toml");
741        std::fs::write(
742            &toml_path,
743            r#"
744[[items]]
745k = "a"
746[[items]]
747k = "b"
748"#,
749        )?;
750        let tbatch = read_toml(&toml_path, Some("items"))?;
751        assert_eq!(tbatch.num_rows(), 2);
752        let col = tbatch
753            .column(0)
754            .as_any()
755            .downcast_ref::<StringArray>()
756            .unwrap();
757        assert!(col.value(0) == "a" || col.value(1) == "a");
758        Ok(())
759    }
760}
rbt/scan/mod.rs

rbt/scan/
mod.rs
