Set bloom filter on byte array

parquet/src/arrow/arrow_writer/byte_array.rs

@@ -429,6 +429,7 @@ struct ByteArrayEncoder {
     dict_encoder: Option<DictEncoder>,
     min_value: Option<ByteArray>,
     max_value: Option<ByteArray>,
+    bloom_filter: Option<Sbbf>,
 }
 
 impl ColumnValueEncoder for ByteArrayEncoder {
@@ -453,8 +454,7 @@ impl ColumnValueEncoder for ByteArrayEncoder {
     }
 
     fn flush_bloom_filter(&mut self) -> Option<Sbbf> {
-        // TODO FIX ME need to handle bloom filter in arrow writer
-        None
+        self.bloom_filter.take()
     }
 
     fn try_new(descr: &ColumnDescPtr, props: &WriterProperties) -> Result<Self>
@@ -467,11 +467,17 @@ impl ColumnValueEncoder for ByteArrayEncoder {
 
         let fallback = FallbackEncoder::new(descr, props)?;
 
+        let bloom_filter = props
+            .bloom_filter_properties(descr.path())
+            .map(|props| Sbbf::new_with_ndv_fpp(props.ndv, props.fpp))
+            .transpose()?;
+
         Ok(Self {
             fallback,
             dict_encoder: dictionary,
             min_value: None,
             max_value: None,
+            bloom_filter,
         })
     }
 
@@ -543,7 +549,7 @@ impl ColumnValueEncoder for ByteArrayEncoder {
 fn encode<T>(values: T, indices: &[usize], encoder: &mut ByteArrayEncoder)
 where
     T: ArrayAccessor + Copy,
-    T::Item: Copy + Ord + AsRef<[u8]>,
+    T::Item: Copy + Ord + AsRef<[u8]> + AsBytes,
 {
     if let Some((min, max)) = compute_min_max(values, indices.iter().cloned()) {
         if encoder.min_value.as_ref().map_or(true, |m| m > &min) {
@@ -555,6 +561,13 @@ where
         }
     }
 
+    // encode the values into bloom filter if enabled
+    if let Some(bloom_filter) = &mut encoder.bloom_filter {
+        for idx in 0..values.len() {
+            bloom_filter.insert(&values.value(idx));
+        }
+    }
+
     match &mut encoder.dict_encoder {
         Some(dict_encoder) => dict_encoder.encode(values, indices),
         None => encoder.fallback.encode(values, indices),

parquet/src/arrow/arrow_writer/mod.rs

@@ -622,7 +622,8 @@ mod tests {
     use crate::basic::Encoding;
     use crate::file::metadata::ParquetMetaData;
     use crate::file::page_index::index_reader::read_pages_locations;
-    use crate::file::properties::WriterVersion;
+    use crate::file::properties::{ReaderProperties, WriterVersion};
+    use crate::file::serialized_reader::ReadOptionsBuilder;
     use crate::file::{
         reader::{FileReader, SerializedFileReader},
         statistics::Statistics,
@@ -1269,6 +1270,7 @@ mod tests {
                             .set_dictionary_enabled(dictionary_size != 0)
                             .set_dictionary_pagesize_limit(dictionary_size.max(1))
                             .set_encoding(*encoding)
+                            .set_bloom_filter_enabled(true)
                             .build();
 
                         files.push(roundtrip_opts(&expected_batch, props))
@@ -1279,14 +1281,14 @@ mod tests {
         files
     }
 
-    fn values_required<A, I>(iter: I)
+    fn values_required<A, I>(iter: I) -> Vec<File>
     where
         A: From<Vec<I::Item>> + Array + 'static,
         I: IntoIterator,
     {
         let raw_values: Vec<_> = iter.into_iter().collect();
         let values = Arc::new(A::from(raw_values));
-        one_column_roundtrip(values, false);
+        one_column_roundtrip(values, false)
     }
 
     fn values_optional<A, I>(iter: I)
@@ -1312,6 +1314,70 @@ mod tests {
         values_optional::<A, I>(iter);
     }
 
+    fn check_bloom_filter<T: AsBytes>(
+        files: Vec<File>,
+        file_column: String,
+        positive_values: Vec<T>,
+        negative_values: Vec<T>,
+    ) {
+        files.into_iter().for_each(|file| {
+            let file_reader = SerializedFileReader::new_with_options(
+                file,
+                ReadOptionsBuilder::new()
+                    .with_reader_properties(
+                        ReaderProperties::builder()
+                            .set_read_bloom_filter(true)
+                            .build(),
+                    )
+                    .build(),
+            )
+            .expect("Unable to open file as Parquet");
+            let metadata = file_reader.metadata();
+            for (ri, row_group) in metadata.row_groups().iter().enumerate() {
+                if let Some((column_index, _)) = row_group
+                    .columns()
+                    .iter()
+                    .enumerate()
+                    .find(|(_, column)| column.column_path().string() == file_column)
+                {
+                    let row_group_reader = file_reader
+                        .get_row_group(ri)
+                        .expect("Unable to read row group");
+                    if let Some(sbbf) =
+                        row_group_reader.get_column_bloom_filter(column_index)
+                    {
+                        if row_group.num_rows() >= positive_values.len() as i64 {
+                            positive_values.iter().for_each(|value| {
+                                assert!(
+                                    sbbf.check(value),
+                                    "{}",
+                                    format!(
+                                        "Value {:?} should be in bloom filter",
+                                        value.as_bytes()
+                                    )
+                                );
+                            });
+                        }
+                        negative_values.iter().for_each(|value| {
+                            assert!(
+                                !sbbf.check(value),
+                                "{}",
+                                format!(
+                                    "Value {:?} should not be in bloom filter",
+                                    value.as_bytes()
+                                )
+                            );
+                        });
+                    } else {
+                        panic!("No bloom filter for column named {} found", file_column);
+                    }
+                } else {
+                    panic!("No column named {} found", file_column);
+                }
+            }
+        });
+    }
+
     #[test]
     fn all_null_primitive_single_column() {
         let values = Arc::new(Int32Array::from(vec![None; SMALL_SIZE]));
@@ -1528,6 +1594,33 @@ mod tests {
         values_required::<BinaryArray, _>(many_vecs_iter);
     }
 
+    #[test]
+    fn i32_column_bloom_filter() {
+        let positive_values: Vec<i32> = (0..SMALL_SIZE as i32).collect();
+        let files = values_required::<Int32Array, _>(positive_values);
+        check_bloom_filter(
+            files,
+            "col".to_string(),
+            (0..SMALL_SIZE as i32).collect(),
+            (SMALL_SIZE as i32 + 1..SMALL_SIZE as i32 + 10).collect(),
+        );
+    }
+
+    #[test]
+    fn binary_column_bloom_filter() {
+        let one_vec: Vec<u8> = (0..SMALL_SIZE as u8).collect();
+        let many_vecs: Vec<_> = std::iter::repeat(one_vec).take(SMALL_SIZE).collect();
+        let many_vecs_iter = many_vecs.iter().map(|v| v.as_slice());
+
+        let files = values_required::<BinaryArray, _>(many_vecs_iter);
+        check_bloom_filter(
+            files,
+            "col".to_string(),
+            many_vecs,
+            vec![vec![(SMALL_SIZE + 1) as u8]],
+        );
+    }
+
     #[test]
     fn large_binary_single_column() {
         let one_vec: Vec<u8> = (0..SMALL_SIZE as u8).collect();

parquet/src/data_type.rs

@@ -448,6 +448,12 @@ impl AsBytes for [u8] {
     }
 }
 
+impl AsBytes for &[u8] {
+    fn as_bytes(&self) -> &[u8] {
+        self
+    }
+}
+
 macro_rules! gen_as_bytes {
     ($source_ty:ident) => {
         impl AsBytes for $source_ty {