Read a Parquet file from bytes and apply predicate pushdown via DataFusion. This is the query path for cold L2 data: the Parquet reader only reads row groups and columns that match the predicate, minimizing I/O.
(
parquet_bytes: &[u8],
projection: &[String],
)
| 67 | /// This is the query path for cold L2 data: the Parquet reader only reads |
| 68 | /// row groups and columns that match the predicate, minimizing I/O. |
| 69 | pub fn read_parquet_with_predicate( |
| 70 | parquet_bytes: &[u8], |
| 71 | projection: &[String], |
| 72 | ) -> crate::Result<Vec<RecordBatch>> { |
| 73 | use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; |
| 74 | |
| 75 | let reader = ParquetRecordBatchReaderBuilder::try_new(Bytes::copy_from_slice(parquet_bytes)) |
| 76 | .map_err(|e| crate::Error::ColdStorage { |
| 77 | detail: format!("parquet reader init: {e}"), |
| 78 | })?; |
| 79 | |
| 80 | // Apply column projection if specified. |
| 81 | let reader = if projection.is_empty() { |
| 82 | reader.build().map_err(|e| crate::Error::ColdStorage { |
| 83 | detail: format!("build reader: {e}"), |
| 84 | })? |
| 85 | } else { |
| 86 | let schema = reader.schema(); |
| 87 | let indices: Vec<usize> = projection |
| 88 | .iter() |
| 89 | .filter_map(|name| schema.index_of(name).ok()) |
| 90 | .collect(); |
| 91 | let mask = parquet::arrow::ProjectionMask::leaves(reader.parquet_schema(), indices); |
| 92 | reader |
| 93 | .with_projection(mask) |
| 94 | .build() |
| 95 | .map_err(|e| crate::Error::ColdStorage { |
| 96 | detail: format!("build projected reader: {e}"), |
| 97 | })? |
| 98 | }; |
| 99 | |
| 100 | let batches: Vec<RecordBatch> = |
| 101 | reader |
| 102 | .collect::<std::result::Result<_, _>>() |
| 103 | .map_err(|e| crate::Error::ColdStorage { |
| 104 | detail: format!("read batches: {e}"), |
| 105 | })?; |
| 106 | Ok(batches) |
| 107 | } |