swiftide_integrations/parquet/
loader.rs

1use anyhow::{Context as _, Result};
2use fs_err::tokio::File;
3use futures_util::StreamExt as _;
4use parquet::arrow::{ParquetRecordBatchStreamBuilder, ProjectionMask};
5use parquet_arrow_array::{LargeStringArray, StringArray, StringViewArray};
6use swiftide_core::{
7    Loader,
8    indexing::{IndexingStream, TextNode},
9};
10use tokio::runtime::Handle;
11
12use super::Parquet;
13
14impl Loader for Parquet {
15    type Output = String;
16
17    fn into_stream(self) -> IndexingStream<String> {
18        let mut builder = tokio::task::block_in_place(|| {
19            Handle::current().block_on(async {
20                let file = File::open(self.path).await.expect("Failed to open file");
21
22                ParquetRecordBatchStreamBuilder::new(file)
23                    .await
24                    .context("Failed to load builder")
25                    .unwrap()
26                    .with_batch_size(self.batch_size)
27            })
28        });
29
30        let file_metadata = builder.metadata().file_metadata().clone();
31        dbg!(file_metadata.schema_descr().columns());
32        let column_idx = file_metadata
33            .schema()
34            .get_fields()
35            .iter()
36            .enumerate()
37            .find_map(|(pos, column)| {
38                if self.column_name == column.name() {
39                    Some(pos)
40                } else {
41                    None
42                }
43            })
44            .unwrap_or_else(|| panic!("Column {} not found in dataset", self.column_name));
45
46        let mask = ProjectionMask::roots(file_metadata.schema_descr(), [column_idx]);
47        builder = builder.with_projection(mask);
48
49        let stream = builder.build().expect("Failed to build parquet builder");
50
51        let swiftide_stream = stream.flat_map_unordered(None, move |result_batch| {
52            let Ok(batch) = result_batch else {
53                let new_result: Result<TextNode> = Err(anyhow::anyhow!(result_batch.unwrap_err()));
54
55                return vec![new_result].into();
56            };
57            assert!(batch.num_columns() == 1, "Number of columns _must_ be 1");
58
59            let column = batch.column(0); // Should only have one column at this point
60            let node_values = if let Some(values) = column.as_any().downcast_ref::<StringArray>() {
61                values
62                    .iter()
63                    .flatten()
64                    .map(TextNode::from)
65                    .map(Ok)
66                    .collect::<Vec<_>>()
67            } else if let Some(values) = column.as_any().downcast_ref::<LargeStringArray>() {
68                values
69                    .iter()
70                    .flatten()
71                    .map(TextNode::from)
72                    .map(Ok)
73                    .collect::<Vec<_>>()
74            } else if let Some(values) = column.as_any().downcast_ref::<StringViewArray>() {
75                values
76                    .iter()
77                    .flatten()
78                    .map(TextNode::from)
79                    .map(Ok)
80                    .collect::<Vec<_>>()
81            } else {
82                let new_result: Result<TextNode> = Err(anyhow::anyhow!(
83                    "Parquet column is not a string array (got {:?})",
84                    column.data_type()
85                ));
86
87                return vec![new_result].into();
88            };
89
90            IndexingStream::iter(node_values)
91        });
92
93        swiftide_stream.boxed().into()
94
95        // let mask = ProjectionMask::
96    }
97
98    fn into_stream_boxed(self: Box<Self>) -> IndexingStream<String> {
99        self.into_stream()
100    }
101}
102
103#[cfg(test)]
104mod tests {
105    use std::path::PathBuf;
106
107    use futures_util::TryStreamExt as _;
108
109    use super::*;
110
111    #[test_log::test(tokio::test(flavor = "multi_thread"))]
112    async fn test_parquet_loader() {
113        let mut path = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
114        path.push("src/parquet/test.parquet");
115        dbg!(&path);
116
117        let loader = Parquet::builder()
118            .path(path)
119            .column_name("chunk")
120            .build()
121            .unwrap();
122
123        let result = loader.into_stream().try_collect::<Vec<_>>().await.unwrap();
124
125        let expected = [TextNode::new("hello"), TextNode::new("world")];
126        assert_eq!(result, expected);
127    }
128}
swiftide_integrations/parquet/loader.rs

swiftide_integrations/parquet/
loader.rs
