tustvold commented on code in PR #4332: URL: https://github.com/apache/arrow-rs/pull/4332#discussion_r1212952603
########## parquet/examples/read_with_rowgroup.rs: ########## @@ -0,0 +1,197 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use arrow::util::pretty::print_batches; +use bytes::{Buf, Bytes}; +use parquet::arrow::arrow_reader::{ParquetRecordBatchReader, RowGroups, RowSelection}; +use parquet::arrow::async_reader::AsyncFileReader; +use parquet::arrow::{parquet_to_arrow_field_levels, ProjectionMask}; +use parquet::column::page::{PageIterator, PageReader}; +use parquet::errors::{ParquetError, Result}; +use parquet::file::metadata::RowGroupMetaData; +use parquet::file::reader::{ChunkReader, Length}; +use parquet::file::serialized_reader::SerializedPageReader; +use std::io::{Read, Seek}; +use std::sync::Arc; +use tokio::fs::File; + +#[tokio::main(flavor = "current_thread")] +async fn main() -> Result<()> { + let testdata = arrow::util::test_util::parquet_test_data(); + let path = format!("{testdata}/alltypes_plain.parquet"); + let mut file = File::open(&path).await.unwrap(); + + // The metadata could be cached in other places, this example only shows how to read + let metadata = file.get_metadata().await?; + + for rg in metadata.row_groups() { + let mut rowgroup = InMemoryRowGroup::create(rg.clone(), ProjectionMask::all()); + let mut reader = std::fs::File::open(&path).unwrap(); + rowgroup.fetch_data(&mut reader, None)?; + let reader = rowgroup.build_reader(1024, None)?; + + for batch in reader { + let batch = batch?; + print_batches(&[batch])?; + } + } + + Ok(()) +} + +/// Implements [`PageIterator`] for a single column chunk, yielding a single [`PageReader`] +struct ColumnChunkIterator { + reader: Option<Result<Box<dyn PageReader>>>, +} + +impl Iterator for ColumnChunkIterator { + type Item = Result<Box<dyn PageReader>>; + + fn next(&mut self) -> Option<Self::Item> { + self.reader.take() + } +} + +impl PageIterator for ColumnChunkIterator {} + +/// An in-memory column chunk +#[derive(Clone)] +pub struct ColumnChunkData { + offset: usize, + data: Bytes, +} + +impl ColumnChunkData { + fn get(&self, start: u64) -> Result<Bytes> { + let start = start as usize - self.offset; + Ok(self.data.slice(start..)) + } +} + +impl Length for ColumnChunkData { + fn len(&self) -> u64 { + self.data.len() as u64 + } +} + +impl ChunkReader for ColumnChunkData { + type T = bytes::buf::Reader<Bytes>; + + fn get_read(&self, start: u64) -> Result<Self::T> { + Ok(self.get(start)?.reader()) + } + + fn get_bytes(&self, start: u64, length: usize) -> Result<Bytes> { + Ok(self.get(start)?.slice(..length)) + } +} + +#[derive(Clone)] +pub struct InMemoryRowGroup { + pub metadata: RowGroupMetaData, + mask: ProjectionMask, + column_chunks: Vec<Option<Arc<ColumnChunkData>>>, +} + +impl RowGroups for InMemoryRowGroup { + fn num_rows(&self) -> usize { + self.metadata.num_rows() as usize + } + + fn column_chunks(&self, i: usize) -> Result<Box<dyn PageIterator>> { + match &self.column_chunks[i] { + None => Err(ParquetError::General(format!( + "Invalid column index {i}, column was not fetched" + ))), + Some(data) => { + let page_reader: Box<dyn PageReader> = + Box::new(SerializedPageReader::new( + data.clone(), + self.metadata.column(i), + self.num_rows(), + None, + )?); + + Ok(Box::new(ColumnChunkIterator { + reader: Some(Ok(page_reader)), + })) + } + } + } +} + +impl InMemoryRowGroup { + pub fn create(metadata: RowGroupMetaData, mask: ProjectionMask) -> Self { + let column_chunks = metadata.columns().iter().map(|_| None).collect::<Vec<_>>(); + + Self { + metadata, + mask, + column_chunks, + } + } + + pub fn build_reader( + &self, + batch_size: usize, + selection: Option<RowSelection>, + ) -> Result<ParquetRecordBatchReader> { + let levels = parquet_to_arrow_field_levels( + &self.metadata.schema_descr_ptr(), + self.mask.clone(), + None, + )?; + + ParquetRecordBatchReader::try_new_with_row_groups( + &levels, self, batch_size, selection, + ) + } + + /// fetch data from a reader in sync mode + pub fn fetch_data<R: Read + Seek>( + &mut self, + reader: &mut R, + _selection: Option<&RowSelection>, + ) -> Result<()> { + let mut vs = std::mem::take(&mut self.column_chunks); + for (leaf_idx, meta) in self.metadata.columns().iter().enumerate() { + if self.mask.leaf_included(leaf_idx) { + let (start, len) = meta.byte_range(); + reader.seek(std::io::SeekFrom::Start(start))?; + let mut chunk = vec![0; len as usize]; + reader.read_exact(&mut chunk)?; + + vs[leaf_idx] = Some(Arc::new(ColumnChunkData { + offset: start as usize, + data: Bytes::from(chunk), + })); + } + } + self.column_chunks = std::mem::take(&mut vs); + Ok(()) + } + + /// fetch data from a reader in async mode + pub fn async_fetch_data<R: Read + Seek>( Review Comment: Perhaps we could either consistently use sync or async IO instead of mixing the two together? -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: github-unsubscr...@arrow.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org
