mbutrovich commented on code in PR #11549:
URL: https://github.com/apache/arrow-rs/pull/11549#discussion_r4275632786
##########
parquet/src/arrow/arrow_writer/mod.rs:
##########
@@ -1812,6 +1608,241 @@ fn write_leaf(
}
}
+/// Returns the values [`ArrowWriter`] stores for `column` in an `INT32` column
+///
+/// These are also the values hashed into the column's bloom filter
+pub(crate) fn coerce_to_i32(column: &dyn arrow_array::Array) ->
Result<ScalarBuffer<i32>> {
Review Comment:
What do you think about not materializing a buffer here at all? As I read
it, the new `coerce_to_i32` and `coerce_to_i64` exist only so
`ArrowSbbf::check` can hash the converted values. I looked at whether
`arrow_cast::cast` could replace them, and it can't for the arms that matter,
because the writer's conversions are reinterpretations or truncations and
`cast` uses value-preserving semantics:
- `UInt32` to `INT32` reinterprets bits. `cast` for `(UInt32, Int32)` goes
through
[`cast_numeric_arrays`](https://github.com/apache/arrow-rs/blob/4dc6bc25f351dd0b8474339ca342d488723f1368/arrow-cast/src/cast/mod.rs#L1576)
and returns null on overflow by default.
- `Decimal128` to `INT32` is `v as i32` here. `cast` goes through
[`cast_decimal_to_integer`](https://github.com/apache/arrow-rs/blob/4dc6bc25f351dd0b8474339ca342d488723f1368/arrow-cast/src/cast/mod.rs#L2158),
which divides by the scale.
- `Date64` to `INT32` is `x / 86_400_000` here. `cast` for [`(Date64,
Int32)`](https://github.com/apache/arrow-rs/blob/4dc6bc25f351dd0b8474339ca342d488723f1368/arrow-cast/src/cast/mod.rs#L1722)
casts the raw milliseconds to `Int32`.
So the writer-specific match is needed. The part that looks avoidable is
that every arm now builds a `ScalarBuffer` (an `unary` allocation for the
narrowing arms) just to be iterated once. Would a per-value shape work better,
for example a function that takes the array and returns an iterator or calls a
closure on each converted `i32`/`i64`, used by both `write_leaf` and
`ArrowSbbf`? That would keep one copy of the conversion without the extra
buffers on the check path. If you tried that and it doesn't work, what got in
the way?
Separately, the `INT32` and `INT64` helpers are near-identical in structure.
Is there a reason to keep them as two functions with their own full matches
rather than one generic over the target type?
##########
parquet/src/arrow/bloom_filter.rs:
##########
@@ -0,0 +1,377 @@
+// 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.
+
+//! Check Arrow arrays against Parquet bloom filters
+
+use arrow_array::cast::AsArray;
+use arrow_array::types::{Float32Type, Float64Type};
+use arrow_array::{Array, BooleanArray};
+use arrow_buffer::BooleanBuffer;
+use arrow_schema::DataType;
+
+use crate::arrow::arrow_writer::{coerce_to_fixed_len_byte_array,
coerce_to_i32, coerce_to_i64};
+use crate::basic::Type as PhysicalType;
+use crate::bloom_filter::Sbbf;
+use crate::errors::{ParquetError, Result};
+use crate::schema::types::ColumnDescriptor;
+
+/// A Parquet bloom filter ([`Sbbf`]) that can be checked with Arrow arrays
+///
+/// [`ArrowWriter`] converts some Arrow types to a different Parquet physical
+/// type before writing them, for example `Int8` and `UInt16` are stored as
+/// `INT32`, and the bloom filter is built from the converted values. Checking
+/// such a filter with [`Sbbf::check`] and the original Arrow value hashes
+/// different bytes, and so reports values that are present as absent.
+///
+/// [`ArrowSbbf::check`] applies the same conversion as [`ArrowWriter`]
+/// before checking the filter.
+///
+/// ```
+/// # use std::sync::Arc;
+/// # use bytes::Bytes;
+/// # use arrow_array::{ArrayRef, Int8Array, RecordBatch};
+/// # use parquet::arrow::ArrowWriter;
+/// # use parquet::arrow::bloom_filter::ArrowSbbf;
+/// # use parquet::file::properties::{ReaderProperties, WriterProperties};
+/// # use parquet::file::reader::{FileReader, SerializedFileReader};
+/// # use parquet::file::serialized_reader::ReadOptionsBuilder;
+/// let array: ArrayRef = Arc::new(Int8Array::from(vec![1, 2, 3]));
+/// let batch = RecordBatch::try_from_iter([("c", array)]).unwrap();
+///
+/// let props = WriterProperties::builder()
+/// .set_bloom_filter_enabled(true)
+/// .build();
+/// let mut buffer = Vec::new();
+/// let mut writer = ArrowWriter::try_new(&mut buffer, batch.schema(),
Some(props)).unwrap();
+/// writer.write(&batch).unwrap();
+/// writer.close().unwrap();
+///
+/// let options = ReadOptionsBuilder::new()
+/// .with_reader_properties(
+/// ReaderProperties::builder()
+/// .set_read_bloom_filter(true)
+/// .build(),
+/// )
+/// .build();
+/// let reader = SerializedFileReader::new_with_options(Bytes::from(buffer),
options).unwrap();
+/// let row_group = reader.get_row_group(0).unwrap();
+/// let sbbf = row_group.get_column_bloom_filter(0).unwrap();
+/// let column = reader.metadata().row_group(0).column(0).column_descr();
+///
+/// // Checking the Arrow value directly hashes a single byte, not the INT32
that was written
+/// assert!(!sbbf.check(&1_i8));
+///
+/// let found = ArrowSbbf::new(sbbf, column)
+/// .check(&Int8Array::from(vec![1, 3, 4]))
+/// .unwrap();
+/// assert!(found.value(0));
+/// assert!(found.value(1));
+/// assert!(!found.value(2));
+/// ```
+///
+/// [`ArrowWriter`]: crate::arrow::ArrowWriter
+#[derive(Debug, Clone, Copy)]
+pub struct ArrowSbbf<'a> {
+ sbbf: &'a Sbbf,
+ physical_type: PhysicalType,
+}
+
+impl<'a> ArrowSbbf<'a> {
+ /// Create a new [`ArrowSbbf`] for the bloom filter of the column
described by `column`
+ pub fn new(sbbf: &'a Sbbf, column: &ColumnDescriptor) -> Self {
+ Self {
+ sbbf,
+ physical_type: column.physical_type(),
+ }
+ }
+
+ /// Check each value of `array` against the bloom filter
+ ///
+ /// Returns `true` if the value may be present, and `false` if it is
+ /// definitely absent. Nulls in `array` are null in the result.
+ ///
+ /// Returns an error if [`ArrowWriter`] cannot write `array`'s data type
to a
+ /// column of this physical type.
+ ///
+ /// [`ArrowWriter`]: crate::arrow::ArrowWriter
+ pub fn check(&self, array: &dyn Array) -> Result<BooleanArray> {
+ if let Some(dictionary) = array.as_any_dictionary_opt() {
+ let values = self.check(dictionary.values().as_ref())?;
+ let found = arrow_select::take::take(&values, dictionary.keys(),
None)?;
+ return Ok(found.as_boolean().clone());
+ }
+
+ let sbbf = self.sbbf;
+ let found = match (self.physical_type, array.data_type()) {
+ (PhysicalType::INT32, _) => coerce_to_i32(array)?
+ .iter()
+ .map(|v| sbbf.check(v))
+ .collect(),
+ (PhysicalType::INT64, _) => coerce_to_i64(array)?
+ .iter()
+ .map(|v| sbbf.check(v))
+ .collect(),
+ (PhysicalType::FLOAT, DataType::Float32) => array
+ .as_primitive::<Float32Type>()
+ .values()
+ .iter()
+ .map(|v| sbbf.check(v))
+ .collect(),
+ (PhysicalType::DOUBLE, DataType::Float64) => array
+ .as_primitive::<Float64Type>()
+ .values()
+ .iter()
+ .map(|v| sbbf.check(v))
+ .collect(),
+ (PhysicalType::BOOLEAN, DataType::Boolean) => array
+ .as_boolean()
+ .values()
+ .iter()
+ .map(|v| sbbf.check(&v))
+ .collect(),
+ (PhysicalType::BYTE_ARRAY, DataType::Utf8) => {
+ check_bytes(sbbf, array.as_string::<i32>())
+ }
+ (PhysicalType::BYTE_ARRAY, DataType::LargeUtf8) => {
+ check_bytes(sbbf, array.as_string::<i64>())
+ }
+ (PhysicalType::BYTE_ARRAY, DataType::Utf8View) => {
+ check_bytes(sbbf, array.as_string_view())
+ }
+ (PhysicalType::BYTE_ARRAY, DataType::Binary) => {
+ check_bytes(sbbf, array.as_binary::<i32>())
+ }
+ (PhysicalType::BYTE_ARRAY, DataType::LargeBinary) => {
+ check_bytes(sbbf, array.as_binary::<i64>())
+ }
+ (PhysicalType::BYTE_ARRAY, DataType::BinaryView) => {
+ check_bytes(sbbf, array.as_binary_view())
+ }
+ (PhysicalType::FIXED_LEN_BYTE_ARRAY, _) => {
+ coerce_to_fixed_len_byte_array(array, 0..array.len())?
+ .iter()
+ .map(|v| sbbf.check(v))
+ .collect()
+ }
Review Comment:
`ArrowSbbf::new` keeps only `column.physical_type()`, so the
`FIXED_LEN_BYTE_ARRAY` arm hashes `coerce_to_fixed_len_byte_array(array, ..)`,
whose decimal width comes from the precision of the array passed to `check`,
via `decimal_length_from_precision`. The width the writer used comes from the
column. What happens when the array's precision differs from the column's?
I wrote a `Decimal128` value `12345` at precision 30 (13 bytes) and checked
it on the head commit:
- precision 30 array: `true`
- precision 38 array: `false`
- precision 20 array: `false`
The value is in the file in all three cases, so a definite `false` is wrong.
Should `new` also keep `column.type_length()` and return an error from `check`
(or re-size the bytes) when the array's decimal width does not match? A test
for this case would help.
--
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]
