tass 0.10.1

A pager for tabular data
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68

use super::DataSource;
use anyhow::anyhow;
use arrow::record_batch::RecordBatch;
use parquet::arrow::arrow_reader::RowSelector;
use parquet::file::reader::FileReader;
use parquet::file::serialized_reader::SerializedFileReader;
use std::fs::File;
use std::time::Instant;
use tracing::debug;

pub struct ParquetFile {
    file: File,
    n_rows: usize,
}

impl ParquetFile {
    pub fn new(file: File) -> anyhow::Result<ParquetFile> {
        // We don't support live-updating parquet files, so we may as well cache
        // the row count
        let n_rows = count_rows(&file)?;
        Ok(ParquetFile { file, n_rows })
    }
}

impl DataSource for ParquetFile {
    fn check_for_new_rows(&mut self) -> anyhow::Result<usize> {
        Ok(0)
    }

    fn row_count(&self) -> usize {
        self.n_rows
    }

    fn fetch_batch(&mut self, offset: usize, len: usize) -> anyhow::Result<RecordBatch> {
        let file = self.file.try_clone()?;
        let mut rdr = parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder::try_new(file)?
            .with_batch_size(len)
            .with_row_selection(
                vec![
                    RowSelector {
                        row_count: offset,
                        skip: true,
                    },
                    RowSelector {
                        row_count: len,
                        skip: false,
                    },
                ]
                .into(),
            )
            .build()?;
        let batch = rdr.next().unwrap()?;
        Ok(batch)
    }

    fn search(&self, _needle: &str) -> anyhow::Result<Vec<usize>> {
        Err(anyhow!("Searching parquet not supported yet"))
    }
}

fn count_rows(file: &File) -> anyhow::Result<usize> {
    let start = Instant::now();
    let file = file.try_clone()?;
    let rdr = SerializedFileReader::new(file)?;
    let total_rows = rdr.metadata().file_metadata().num_rows() as usize;
    debug!("Counted {total_rows} rows (took {:?})", start.elapsed());
    Ok(total_rows)
}