alamb commented on code in PR #2335: URL: https://github.com/apache/arrow-rs/pull/2335#discussion_r939531214
########## parquet/src/arrow/arrow_reader/filter.rs: ########## @@ -0,0 +1,95 @@ +// 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 crate::arrow::ProjectionMask; +use arrow::array::BooleanArray; +use arrow::error::Result as ArrowResult; +use arrow::record_batch::RecordBatch; + +/// A predicate operating on [`RecordBatch`] +pub trait ArrowPredicate: Send + 'static { + /// Returns the projection mask for this predicate + fn projection(&self) -> &ProjectionMask; + + /// Called with a [`RecordBatch`] containing the columns identified by [`Self::mask`], + /// with `true` values in the returned [`BooleanArray`] indicating rows + /// matching the predicate. + /// + /// The returned [`BooleanArray`] must not contain any nulls Review Comment: ```suggestion /// The returned [`BooleanArray`] must not contain any nulls otherwise or else XXX will happen ``` What will happen if the array has nulls? ########## parquet/src/arrow/arrow_reader/filter.rs: ########## @@ -0,0 +1,95 @@ +// 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 crate::arrow::ProjectionMask; +use arrow::array::BooleanArray; +use arrow::error::Result as ArrowResult; +use arrow::record_batch::RecordBatch; + +/// A predicate operating on [`RecordBatch`] +pub trait ArrowPredicate: Send + 'static { + /// Returns the projection mask for this predicate + fn projection(&self) -> &ProjectionMask; + + /// Called with a [`RecordBatch`] containing the columns identified by [`Self::mask`], + /// with `true` values in the returned [`BooleanArray`] indicating rows + /// matching the predicate. + /// + /// The returned [`BooleanArray`] must not contain any nulls + fn filter(&mut self, batch: RecordBatch) -> ArrowResult<BooleanArray>; +} + +/// An [`ArrowPredicate`] created from an [`FnMut`] +pub struct ArrowPredicateFn<F> { + f: F, + projection: ProjectionMask, +} + +impl<F> ArrowPredicateFn<F> +where + F: FnMut(RecordBatch) -> ArrowResult<BooleanArray> + Send + 'static, +{ + /// Create a new [`ArrowPredicateFn`] + pub fn new(projection: ProjectionMask, f: F) -> Self { + Self { f, projection } + } +} + +impl<F> ArrowPredicate for ArrowPredicateFn<F> +where + F: FnMut(RecordBatch) -> ArrowResult<BooleanArray> + Send + 'static, +{ + fn projection(&self) -> &ProjectionMask { + &self.projection + } + + fn filter(&mut self, batch: RecordBatch) -> ArrowResult<BooleanArray> { + (self.f)(batch) + } +} + +/// A [`RowFilter`] allows pushing down a filter predicate to skip IO and decode +/// +/// This consists of a list of [`ArrowPredicate`] where only the rows that satisfy all +/// of the predicates will be returned. Any [`RowSelection`] will be applied prior +/// to the first predicate, and each predicate in turn will then be used to compute +/// a more refined [`RowSelection`] to use when evaluating the subsequent predicates. +/// +/// Once all predicates have been evaluated, the resulting [`RowSelection`] will be +/// used to return just the desired rows. +/// +/// This design has a couple of implications: +/// +/// * [`RowFilter`] can be used to skip fetching IO, in addition to decode overheads Review Comment: ```suggestion /// * [`RowFilter`] can be used to skip entire pages, and thus IO, in addition to CPU decode overheads ``` ########## parquet/src/arrow/arrow_reader/filter.rs: ########## @@ -0,0 +1,95 @@ +// 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 crate::arrow::ProjectionMask; +use arrow::array::BooleanArray; +use arrow::error::Result as ArrowResult; +use arrow::record_batch::RecordBatch; + +/// A predicate operating on [`RecordBatch`] +pub trait ArrowPredicate: Send + 'static { + /// Returns the projection mask for this predicate + fn projection(&self) -> &ProjectionMask; + + /// Called with a [`RecordBatch`] containing the columns identified by [`Self::mask`], + /// with `true` values in the returned [`BooleanArray`] indicating rows + /// matching the predicate. + /// + /// The returned [`BooleanArray`] must not contain any nulls + fn filter(&mut self, batch: RecordBatch) -> ArrowResult<BooleanArray>; +} + +/// An [`ArrowPredicate`] created from an [`FnMut`] +pub struct ArrowPredicateFn<F> { + f: F, + projection: ProjectionMask, +} + +impl<F> ArrowPredicateFn<F> +where + F: FnMut(RecordBatch) -> ArrowResult<BooleanArray> + Send + 'static, +{ + /// Create a new [`ArrowPredicateFn`] Review Comment: ```suggestion /// Create a new [`ArrowPredicateFn`]. `f` will be passed batches /// that contains the columns specified in `projection` /// and returns a [`BooleanArray`] that describes which rows should /// be passed along. ``` ########## parquet/src/arrow/arrow_reader/selection.rs: ########## @@ -0,0 +1,215 @@ +// 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::array::{Array, BooleanArray}; +use arrow::compute::SlicesIterator; +use std::cmp::Ordering; +use std::collections::VecDeque; +use std::ops::Range; + +/// [`RowSelection`] is a collection of [`RowSelect`] used to skip rows when Review Comment: ```suggestion /// [`RowSelector`] represents range of rows ``` ########## parquet/src/arrow/arrow_reader/filter.rs: ########## @@ -0,0 +1,95 @@ +// 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 crate::arrow::ProjectionMask; +use arrow::array::BooleanArray; +use arrow::error::Result as ArrowResult; +use arrow::record_batch::RecordBatch; + +/// A predicate operating on [`RecordBatch`] +pub trait ArrowPredicate: Send + 'static { + /// Returns the projection mask for this predicate + fn projection(&self) -> &ProjectionMask; + + /// Called with a [`RecordBatch`] containing the columns identified by [`Self::mask`], + /// with `true` values in the returned [`BooleanArray`] indicating rows + /// matching the predicate. + /// + /// The returned [`BooleanArray`] must not contain any nulls + fn filter(&mut self, batch: RecordBatch) -> ArrowResult<BooleanArray>; +} + +/// An [`ArrowPredicate`] created from an [`FnMut`] +pub struct ArrowPredicateFn<F> { + f: F, + projection: ProjectionMask, +} + +impl<F> ArrowPredicateFn<F> +where + F: FnMut(RecordBatch) -> ArrowResult<BooleanArray> + Send + 'static, +{ + /// Create a new [`ArrowPredicateFn`] + pub fn new(projection: ProjectionMask, f: F) -> Self { + Self { f, projection } + } +} + +impl<F> ArrowPredicate for ArrowPredicateFn<F> +where + F: FnMut(RecordBatch) -> ArrowResult<BooleanArray> + Send + 'static, +{ + fn projection(&self) -> &ProjectionMask { + &self.projection + } + + fn filter(&mut self, batch: RecordBatch) -> ArrowResult<BooleanArray> { + (self.f)(batch) + } +} + +/// A [`RowFilter`] allows pushing down a filter predicate to skip IO and decode +/// +/// This consists of a list of [`ArrowPredicate`] where only the rows that satisfy all +/// of the predicates will be returned. Any [`RowSelection`] will be applied prior +/// to the first predicate, and each predicate in turn will then be used to compute +/// a more refined [`RowSelection`] to use when evaluating the subsequent predicates. +/// +/// Once all predicates have been evaluated, the resulting [`RowSelection`] will be +/// used to return just the desired rows. Review Comment: ```suggestion /// Once all predicates have been evaluated, the final [`RowSelection`] is applied /// to the overall [`ColumnProjection`] to produce the final output [`RecordBatch`]es. ``` ########## parquet/src/arrow/arrow_reader/selection.rs: ########## @@ -0,0 +1,215 @@ +// 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::array::{Array, BooleanArray}; +use arrow::compute::SlicesIterator; +use std::cmp::Ordering; +use std::collections::VecDeque; +use std::ops::Range; + +/// [`RowSelection`] is a collection of [`RowSelect`] used to skip rows when +/// scanning a parquet file +#[derive(Debug, Clone, Copy)] +pub struct RowSelector { + /// The number of rows + pub row_count: usize, Review Comment: Alternately, perhaps we could represent Row selection by RLE runs of the rows to keep: (position, count) Again, using i32 or something smaller that might be much more efficient to process 🤷 ########## parquet/src/arrow/async_reader.rs: ########## @@ -253,6 +262,30 @@ impl<T: AsyncFileReader> ParquetRecordBatchStreamBuilder<T> { } } + /// Provide a [`RowSelection] to filter out rows, and avoid fetching their + /// data into memory + /// + /// Row group filtering is applied prior to this, and rows from skipped + /// row groups should not be included in the [`RowSelection`] + /// + /// TODO: Make public once stable (#1792) Review Comment: Probably also a good idea to link to the docs that describe the order of filter application in decoding (RowSelection followed by RowFilter) ########## parquet/src/arrow/async_reader.rs: ########## @@ -271,25 +304,122 @@ impl<T: AsyncFileReader> ParquetRecordBatchStreamBuilder<T> { None => (0..self.metadata.row_groups().len()).collect(), }; + let reader = ReaderFactory { + input: self.input, + filter: self.filter, + metadata: self.metadata.clone(), + schema: self.schema.clone(), + }; + Ok(ParquetRecordBatchStream { + metadata: self.metadata, + batch_size: self.batch_size, row_groups, projection: self.projection, - batch_size: self.batch_size, - metadata: self.metadata, + selection: self.selection, schema: self.schema, - input: Some(self.input), + reader: Some(reader), state: StreamState::Init, }) } } +type ReadResult<T> = Result<(ReaderFactory<T>, Option<ParquetRecordBatchReader>)>; + +/// [`ReaderFactory`] is used by [`ParquetRecordBatchStream`] to create +/// [`ParquetRecordBatchReader`] +struct ReaderFactory<T> { + metadata: Arc<ParquetMetaData>, + + schema: SchemaRef, + + input: T, + + filter: Option<RowFilter>, +} + +impl<T> ReaderFactory<T> +where + T: AsyncFileReader + Send, +{ + /// Reads the next row group with the provided `selection`, `projection` and `batch_size` + /// + /// Note: this captures self so that the resulting future has a static lifetime + async fn read_row_group( + mut self, + row_group_idx: usize, + mut selection: Option<RowSelection>, + projection: ProjectionMask, + batch_size: usize, + ) -> ReadResult<T> { + // TODO: calling build_array multiple times is wasteful + let selects_any = |selection: Option<&RowSelection>| { + selection.map(|x| x.selects_any()).unwrap_or(true) + }; + + let meta = self.metadata.row_group(row_group_idx); + let mut row_group = InMemoryRowGroup { + schema: meta.schema_descr_ptr(), + row_count: meta.num_rows() as usize, + column_chunks: vec![None; meta.columns().len()], + }; + + if let Some(filter) = self.filter.as_mut() { + for predicate in filter.predicates.iter_mut() { + if !selects_any(selection.as_ref()) { + return Ok((self, None)); + } + + let predicate_projection = predicate.projection().clone(); + row_group + .fetch( + &mut self.input, + meta, + &predicate_projection, + selection.as_ref(), + ) + .await?; + + let array_reader = build_array_reader( + self.schema.clone(), + predicate_projection, + &row_group, + )?; + + selection = Some(evaluate_predicate( + batch_size, + array_reader, + selection, + predicate.as_mut(), + )?); + } + } + + if !selects_any(selection.as_ref()) { + return Ok((self, None)); + } + + row_group + .fetch(&mut self.input, meta, &projection, selection.as_ref()) + .await?; + + let reader = ParquetRecordBatchReader::new( + batch_size, + build_array_reader(self.schema.clone(), projection, &row_group)?, + selection, + ); Review Comment: I think trying to eliminate redundant decoding is a good idea for the reasons that @thinkharderdev Conveniently, it seems like nothing in the API of this PR requires decoding multiple times, so I think we could also potentially implement the 'use take rather than redundant decode' in a follow on PR as well. In terms of "Eventually it might be possible to push simple predicates down to operate directly on the encoded data, which would avoid this." I agree it is a ways off however, I think it could fit into this API with something like adding a list of `ParquetFilter`s to apply during the decode itself that could be efficiently implemented ```rust /// Filter that is applied during decoding of a single column /// semantically takes the form <col> <op> <constant> struct ParquetFilter { op: ParquetOp left: ParquetConst, } enum ParquetFilterOp { EQ, NEQ } enum ParqeutConst { Int64(i64) Float(f64) ... } ``` ########## parquet/src/arrow/arrow_reader/filter.rs: ########## @@ -0,0 +1,95 @@ +// 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 crate::arrow::ProjectionMask; +use arrow::array::BooleanArray; +use arrow::error::Result as ArrowResult; +use arrow::record_batch::RecordBatch; + +/// A predicate operating on [`RecordBatch`] +pub trait ArrowPredicate: Send + 'static { + /// Returns the projection mask for this predicate Review Comment: ```suggestion /// Returns the [`ProjectionMask`] that describes the columns required /// to evaluate this predicate. All projected columns will be provided in the `batch` /// passed to [`filter`](Self::filter) ``` ########## parquet/src/arrow/arrow_reader/mod.rs: ########## @@ -375,11 +340,36 @@ impl ParquetRecordBatchReader { batch_size, array_reader, schema: Arc::new(schema), - selection, + selection: selection.map(Into::into), } } } +/// Evaluates an [`ArrowPredicate`] returning the [`RowSelection`] +/// +/// If this [`ParquetRecordBatchReader`] has a [`RowSelection`], the +/// returned [`RowSelection`] will be the conjunction of this and +/// the rows selected by `predicate` +pub(crate) fn evaluate_predicate( + batch_size: usize, + array_reader: Box<dyn ArrayReader>, + selection: Option<RowSelection>, Review Comment: It might make the code clearer if this were called `input_selection` to distinguish it from the output selection ```suggestion input_selection: Option<RowSelection>, ``` ########## parquet/src/arrow/arrow_reader/filter.rs: ########## @@ -0,0 +1,95 @@ +// 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 crate::arrow::ProjectionMask; +use arrow::array::BooleanArray; +use arrow::error::Result as ArrowResult; +use arrow::record_batch::RecordBatch; + +/// A predicate operating on [`RecordBatch`] +pub trait ArrowPredicate: Send + 'static { + /// Returns the projection mask for this predicate + fn projection(&self) -> &ProjectionMask; + + /// Called with a [`RecordBatch`] containing the columns identified by [`Self::mask`], + /// with `true` values in the returned [`BooleanArray`] indicating rows + /// matching the predicate. + /// + /// The returned [`BooleanArray`] must not contain any nulls + fn filter(&mut self, batch: RecordBatch) -> ArrowResult<BooleanArray>; +} + +/// An [`ArrowPredicate`] created from an [`FnMut`] +pub struct ArrowPredicateFn<F> { + f: F, + projection: ProjectionMask, +} + +impl<F> ArrowPredicateFn<F> +where + F: FnMut(RecordBatch) -> ArrowResult<BooleanArray> + Send + 'static, +{ + /// Create a new [`ArrowPredicateFn`] + pub fn new(projection: ProjectionMask, f: F) -> Self { + Self { f, projection } + } +} + +impl<F> ArrowPredicate for ArrowPredicateFn<F> +where + F: FnMut(RecordBatch) -> ArrowResult<BooleanArray> + Send + 'static, +{ + fn projection(&self) -> &ProjectionMask { + &self.projection + } + + fn filter(&mut self, batch: RecordBatch) -> ArrowResult<BooleanArray> { + (self.f)(batch) + } +} + +/// A [`RowFilter`] allows pushing down a filter predicate to skip IO and decode +/// +/// This consists of a list of [`ArrowPredicate`] where only the rows that satisfy all +/// of the predicates will be returned. Any [`RowSelection`] will be applied prior +/// to the first predicate, and each predicate in turn will then be used to compute +/// a more refined [`RowSelection`] to use when evaluating the subsequent predicates. +/// +/// Once all predicates have been evaluated, the resulting [`RowSelection`] will be +/// used to return just the desired rows. +/// +/// This design has a couple of implications: +/// +/// * [`RowFilter`] can be used to skip fetching IO, in addition to decode overheads +/// * Columns may be decoded multiple times if they appear in multiple [`ProjectionMask`] +/// * IO will be deferred until needed by a [`ProjectionMask`] +/// +/// As such there is a trade-off between a single large predicate, or multiple predicates, +/// that will depend on the shape of the data. Whilst multiple smaller predicates may +/// minimise the amount of data scanned/decoded, it may not be faster overall. +/// Review Comment: ```suggestion /// /// For example, if a predicate that needs a single column of data filters out all but /// 1% of the rows, applying it as one of the early `ArrowPredicateFn` will likely signficantly /// improve performance. /// /// As a counter example, if a predicate needs several columns of data to evaluate but /// leaves 99% of the rows, it may be better to not filter the data from parquet and /// apply the filter after the RecordBatch has been fully decoded. /// ``` ########## parquet/src/arrow/arrow_reader/selection.rs: ########## @@ -0,0 +1,215 @@ +// 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::array::{Array, BooleanArray}; +use arrow::compute::SlicesIterator; +use std::cmp::Ordering; +use std::collections::VecDeque; +use std::ops::Range; + +/// [`RowSelection`] is a collection of [`RowSelect`] used to skip rows when +/// scanning a parquet file +#[derive(Debug, Clone, Copy)] +pub struct RowSelector { + /// The number of rows + pub row_count: usize, Review Comment: I wonder if we could improve this to take less space in memory (and potentially be faster to process) by packing the row_count into a smaller size. As it is, I think this will take up 16 bytes: https://play.rust-lang.org/?version=stable&mode=debug&edition=2018&gist=3f7e94978d00ba227c594714b836e41d ########## parquet/src/arrow/arrow_reader/filter.rs: ########## @@ -0,0 +1,95 @@ +// 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 crate::arrow::ProjectionMask; +use arrow::array::BooleanArray; +use arrow::error::Result as ArrowResult; +use arrow::record_batch::RecordBatch; + +/// A predicate operating on [`RecordBatch`] +pub trait ArrowPredicate: Send + 'static { + /// Returns the projection mask for this predicate + fn projection(&self) -> &ProjectionMask; + + /// Called with a [`RecordBatch`] containing the columns identified by [`Self::mask`], + /// with `true` values in the returned [`BooleanArray`] indicating rows + /// matching the predicate. + /// + /// The returned [`BooleanArray`] must not contain any nulls Review Comment: I recommend a slightly different semantic, consistent with SQL, that would allow more flexibility from users ```suggestion /// All row that are `true` in returned [`BooleanArray`] will be returned to the reader. /// Any rows that are `false` or `Null` will not be ``` If performance is a concern, perhaps we could provide specialized paths for a `BooleanArray` that has no nulls and note that in the comments ########## parquet/src/arrow/arrow_reader/selection.rs: ########## @@ -0,0 +1,215 @@ +// 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::array::{Array, BooleanArray}; +use arrow::compute::SlicesIterator; +use std::cmp::Ordering; +use std::collections::VecDeque; +use std::ops::Range; + +/// [`RowSelection`] is a collection of [`RowSelect`] used to skip rows when +/// scanning a parquet file +#[derive(Debug, Clone, Copy)] +pub struct RowSelector { + /// The number of rows + pub row_count: usize, + + /// If true, skip `row_count` rows + pub skip: bool, +} + +impl RowSelector { + /// Select `row_count` rows + pub fn select(row_count: usize) -> Self { + Self { + row_count, + skip: false, + } + } + + /// Skip `row_count` rows + pub fn skip(row_count: usize) -> Self { + Self { + row_count, + skip: true, + } + } +} + +/// [`RowSelection`] allows selecting or skipping a provided number of rows +/// when scanning the parquet file. +/// +/// This is applied prior to reading column data, and can therefore +/// be used to skip IO to fetch data into memory +/// +/// A typical use-case would be using the [`PageIndex`] to filter out rows +/// that don't satisfy a predicate +/// +/// [`PageIndex`]: [crate::file::page_index::index::PageIndex] +#[derive(Debug, Clone, Default)] +pub struct RowSelection { + selectors: Vec<RowSelector>, +} + +impl RowSelection { + /// Creates a [`RowSelection`] from a slice of [`BooleanArray`] + /// + /// # Panic + /// + /// Panics if any of the [`BooleanArray`] contain nulls + pub fn from_filters(filters: &[BooleanArray]) -> Self { + let mut next_offset = 0; + let total_rows = filters.iter().map(|x| x.len()).sum(); + + let iter = filters.iter().flat_map(|filter| { + let offset = next_offset; + next_offset += filter.len(); + assert_eq!(filter.null_count(), 0); + SlicesIterator::new(filter) + .map(move |(start, end)| start + offset..end + offset) + }); + + Self::from_consecutive_ranges(iter, total_rows) + } + + /// Creates a [`RowSelection`] from an iterator of consecutive ranges to keep + fn from_consecutive_ranges<I: Iterator<Item = Range<usize>>>( + ranges: I, + total_rows: usize, + ) -> Self { + let mut selectors: Vec<RowSelector> = Vec::with_capacity(ranges.size_hint().0); + let mut last_end = 0; + for range in ranges { + let len = range.end - range.start; + + match range.start.cmp(&last_end) { + Ordering::Equal => match selectors.last_mut() { + Some(last) => last.row_count += len, + None => selectors.push(RowSelector::select(len)), + }, + Ordering::Greater => { + selectors.push(RowSelector::skip(range.start - last_end)); + selectors.push(RowSelector::select(len)) + } + Ordering::Less => panic!("out of order"), + } + last_end = range.end; + } + + if last_end != total_rows { + selectors.push(RowSelector::skip(total_rows - last_end)) + } + + Self { selectors } + } + + /// Splits off `row_count` from this [`RowSelection`] + pub fn split_off(&mut self, row_count: usize) -> Self { + let mut total_count = 0; + + // Find the index where the selector exceeds the row count + let find = self.selectors.iter().enumerate().find(|(_, selector)| { + total_count += selector.row_count; + total_count >= row_count + }); + + let split_idx = match find { + Some((idx, _)) => idx, + None => return Self::default(), + }; + + let mut remaining = self.selectors.split_off(split_idx); + if total_count != row_count { + let overflow = total_count - row_count; + let rem = remaining.first_mut().unwrap(); + rem.row_count -= overflow; + + self.selectors.push(RowSelector { + row_count, + skip: rem.skip, + }) + } + + std::mem::swap(&mut remaining, &mut self.selectors); + Self { + selectors: remaining, + } + } + + /// Given a [`RowSelection`] computed under `self` returns the [`RowSelection`] + /// representing their conjunction + pub fn and(&self, other: &Self) -> Self { + let mut selectors = vec![]; + let mut first = self.selectors.iter().cloned().peekable(); + let mut second = other.selectors.iter().cloned().peekable(); + + let mut to_skip = 0; + while let (Some(a), Some(b)) = (first.peek_mut(), second.peek_mut()) { + if a.row_count == 0 { + first.next().unwrap(); + continue; + } + + if b.row_count == 0 { + second.next().unwrap(); + continue; + } + + if a.skip { + // Records were skipped when producing second + to_skip += a.row_count; + first.next().unwrap(); + continue; + } + + let skip = b.skip; + let to_process = a.row_count.min(b.row_count); + + a.row_count -= to_process; + b.row_count -= to_process; + + match skip { + true => to_skip += to_process, + false => { + if to_skip != 0 { + selectors.push(RowSelector::skip(to_skip)); + to_skip = 0; + } + selectors.push(RowSelector::select(to_process)) + } + } + } + Self { selectors } + } + + /// Returns `true` if this [`RowSelection`] selects any rows + pub fn selects_any(&self) -> bool { + self.selectors.iter().any(|x| !x.skip) + } +} + +impl From<Vec<RowSelector>> for RowSelection { + fn from(selectors: Vec<RowSelector>) -> Self { + Self { selectors } + } +} + +impl Into<VecDeque<RowSelector>> for RowSelection { + fn into(self) -> VecDeque<RowSelector> { + self.selectors.into() + } +} Review Comment: I agree -- I didn't review the logic in detail either as I figured we are just at the "API feedback" design phase -- 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: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
