Repository: apache/arrow-rs
Pull request #6626 (2024-10-24 21:15:43)

Fix test feature selection so all feature combinations work as expected
# Which issue does this PR close?

Some combinations of features cause compilation failures or test failures where not necessary. For example, in the `parquet/src/file/metadata/reader.rs` file, the `MetadataFetch` trait is only imported when `feature = "async"`, but that trait exists under a module that's marked `feature = "arrow"`, so we need to enable both those features to get it to work.

# Rationale for this change
 
There's nothing stopping people from using any combination of features from this crate, so we should make sure they're all supported

# What changes are included in this PR?

This just adds extra `#[cfg()]` attributes where necessary to fix compilation and test failures.

# Are there any user-facing changes?

No. It looks, on the surface, like we could be breaking the parquet crate for people who have the `async` flag enabled with the `arrow` flag disabled, but this crate doesn't compile in its current form for that combination, so this won't actually adversely affect anyone.



Patch:
diff --git a/arrow-select/src/take.rs b/arrow-select/src/take.rs
index 7e17e61e1..07630a49f 100644
--- a/arrow-select/src/take.rs
+++ b/arrow-select/src/take.rs
@@ -50,6 +50,10 @@ use num::{One, Zero};
 ///
 /// For selecting values by index from multiple arrays see [`crate::interleave`]
 ///
+/// Note that this kernel, similar to other kernels in this crate,
+/// will avoid allocating where not necessary. Consequently
+/// the returned array may share buffers with the inputs
+///
 /// # Errors
 /// This function errors whenever:
 /// * An index cannot be casted to `usize` (typically 32 bit architectures)
diff --git a/parquet/Cargo.toml b/parquet/Cargo.toml
index 8ab78efe0..32bc13b62 100644
--- a/parquet/Cargo.toml
+++ b/parquet/Cargo.toml
@@ -222,5 +222,10 @@ harness = false
 name = "metadata"
 harness = false
 
+[[bench]]
+name = "row_selector"
+harness = false
+required-features = ["arrow"]
+
 [lib]
 bench = false
diff --git a/parquet/benches/row_selector.rs b/parquet/benches/row_selector.rs
new file mode 100644
index 000000000..32f0d6a56
--- /dev/null
+++ b/parquet/benches/row_selector.rs
@@ -0,0 +1,87 @@
+// 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::BooleanArray;
+use criterion::*;
+use parquet::arrow::arrow_reader::RowSelection;
+use rand::Rng;
+
+/// Generates a random RowSelection with a specified selection ratio.
+///
+/// # Arguments
+///
+/// * `total_rows` - The total number of rows in the selection.
+/// * `selection_ratio` - The ratio of rows to select (e.g., 1/3 for ~33% selection).
+///
+/// # Returns
+///
+/// * A `BooleanArray` instance with randomly selected rows based on the provided ratio.
+fn generate_random_row_selection(total_rows: usize, selection_ratio: f64) -> BooleanArray {
+    let mut rng = rand::thread_rng();
+    let bools: Vec<bool> = (0..total_rows)
+        .map(|_| rng.gen_bool(selection_ratio))
+        .collect();
+    BooleanArray::from(bools)
+}
+
+fn criterion_benchmark(c: &mut Criterion) {
+    let total_rows = 300_000;
+    let selection_ratio = 1.0 / 3.0;
+
+    // Generate two random RowSelections with approximately 1/3 of the rows selected.
+    let row_selection_a =
+        RowSelection::from_filters(&[generate_random_row_selection(total_rows, selection_ratio)]);
+    let row_selection_b =
+        RowSelection::from_filters(&[generate_random_row_selection(total_rows, selection_ratio)]);
+
+    // Benchmark the intersection of the two RowSelections.
+    c.bench_function("intersection", |b| {
+        b.iter(|| {
+            let intersection = row_selection_a.intersection(&row_selection_b);
+            criterion::black_box(intersection);
+        })
+    });
+
+    c.bench_function("union", |b| {
+        b.iter(|| {
+            let union = row_selection_a.union(&row_selection_b);
+            criterion::black_box(union);
+        })
+    });
+
+    c.bench_function("from_filters", |b| {
+        let boolean_array = generate_random_row_selection(total_rows, selection_ratio);
+        b.iter(|| {
+            let array = boolean_array.clone();
+            let selection = RowSelection::from_filters(&[array]);
+            criterion::black_box(selection);
+        })
+    });
+
+    c.bench_function("and_then", |b| {
+        let selected = row_selection_a.row_count();
+        let sub_selection =
+            RowSelection::from_filters(&[generate_random_row_selection(selected, selection_ratio)]);
+        b.iter(|| {
+            let result = row_selection_a.and_then(&sub_selection);
+            criterion::black_box(result);
+        })
+    });
+}
+
+criterion_group!(benches, criterion_benchmark);
+criterion_main!(benches);
diff --git a/parquet/src/file/metadata/reader.rs b/parquet/src/file/metadata/reader.rs
index c803d6013..2a927f15f 100644
--- a/parquet/src/file/metadata/reader.rs
+++ b/parquet/src/file/metadata/reader.rs
@@ -31,7 +31,7 @@ use crate::schema::types;
 use crate::schema::types::SchemaDescriptor;
 use crate::thrift::{TCompactSliceInputProtocol, TSerializable};
 
-#[cfg(feature = "async")]
+#[cfg(all(feature = "async", feature = "arrow"))]
 use crate::arrow::async_reader::MetadataFetch;
 
 /// Reads the [`ParquetMetaData`] from a byte stream.
@@ -321,7 +321,7 @@ impl ParquetMetaDataReader {
     ///
     /// See [`Self::with_prefetch_hint`] for a discussion of how to reduce the number of fetches
     /// performed by this function.
-    #[cfg(feature = "async")]
+    #[cfg(all(feature = "async", feature = "arrow"))]
     pub async fn load_and_finish<F: MetadataFetch>(
         mut self,
         fetch: F,
@@ -336,7 +336,7 @@ impl ParquetMetaDataReader {
     ///
     /// See [`Self::with_prefetch_hint`] for a discussion of how to reduce the number of fetches
     /// performed by this function.
-    #[cfg(feature = "async")]
+    #[cfg(all(feature = "async", feature = "arrow"))]
     pub async fn try_load<F: MetadataFetch>(
         &mut self,
         mut fetch: F,
@@ -357,12 +357,12 @@ impl ParquetMetaDataReader {
 
     /// Asynchronously fetch the page index structures when a [`ParquetMetaData`] has already
     /// been obtained. See [`Self::new_with_metadata()`].
-    #[cfg(feature = "async")]
+    #[cfg(all(feature = "async", feature = "arrow"))]
     pub async fn load_page_index<F: MetadataFetch>(&mut self, fetch: F) -> Result<()> {
         self.load_page_index_with_remainder(fetch, None).await
     }
 
-    #[cfg(feature = "async")]
+    #[cfg(all(feature = "async", feature = "arrow"))]
     async fn load_page_index_with_remainder<F: MetadataFetch>(
         &mut self,
         mut fetch: F,
@@ -513,7 +513,7 @@ impl ParquetMetaDataReader {
     /// Return the number of bytes to read in the initial pass. If `prefetch_size` has
     /// been provided, then return that value if it is larger than the size of the Parquet
     /// file footer (8 bytes). Otherwise returns `8`.
-    #[cfg(feature = "async")]
+    #[cfg(all(feature = "async", feature = "arrow"))]
     fn get_prefetch_size(&self) -> usize {
         if let Some(prefetch) = self.prefetch_hint {
             if prefetch > FOOTER_SIZE {
@@ -523,7 +523,7 @@ impl ParquetMetaDataReader {
         FOOTER_SIZE
     }
 
-    #[cfg(feature = "async")]
+    #[cfg(all(feature = "async", feature = "arrow"))]
     async fn load_metadata<F: MetadataFetch>(
         fetch: &mut F,
         file_size: usize,
@@ -851,8 +851,7 @@ mod tests {
     }
 }
 
-#[cfg(feature = "async")]
-#[cfg(test)]
+#[cfg(all(feature = "async", feature = "arrow", test))]
 mod async_tests {
     use super::*;
     use bytes::Bytes;
diff --git a/parquet/src/file/writer.rs b/parquet/src/file/writer.rs
index 95ff109a3..b84c57a60 100644
--- a/parquet/src/file/writer.rs
+++ b/parquet/src/file/writer.rs
@@ -378,7 +378,12 @@ fn write_bloom_filters<W: Write + Send>(
         .ordinal()
         .expect("Missing row group ordinal")
         .try_into()
-        .expect("Negative row group ordinal");
+        .map_err(|_| {
+            ParquetError::General(format!(
+                "Negative row group ordinal: {})",
+                row_group.ordinal().unwrap()
+            ))
+        })?;
     let row_group_idx = row_group_idx as usize;
     for (column_idx, column_chunk) in row_group.columns_mut().iter_mut().enumerate() {
         if let Some(bloom_filter) = bloom_filters[row_group_idx][column_idx].take() {
@@ -1892,6 +1897,44 @@ mod tests {
         assert_eq!(page_sizes[0], unenc_size);
     }
 
+    #[test]
+    fn test_too_many_rowgroups() {
+        let message_type = "
+            message test_schema {
+                REQUIRED BYTE_ARRAY a (UTF8);
+            }
+        ";
+        let schema = Arc::new(parse_message_type(message_type).unwrap());
+        let file: File = tempfile::tempfile().unwrap();
+        let props = Arc::new(
+            WriterProperties::builder()
+                .set_statistics_enabled(EnabledStatistics::None)
+                .set_max_row_group_size(1)
+                .build(),
+        );
+        let mut writer = SerializedFileWriter::new(&file, schema, props).unwrap();
+
+        // Create 32k empty rowgroups. Should error when i == 32768.
+        for i in 0..0x8001 {
+            match writer.next_row_group() {
+                Ok(mut row_group_writer) => {
+                    assert_ne!(i, 0x8000);
+                    let col_writer = row_group_writer.next_column().unwrap().unwrap();
+                    col_writer.close().unwrap();
+                    row_group_writer.close().unwrap();
+                }
+                Err(e) => {
+                    assert_eq!(i, 0x8000);
+                    assert_eq!(
+                        e.to_string(),
+                        "Parquet error: Parquet does not support more than 32767 row groups per file (currently: 32768)"
+                    );
+                }
+            }
+        }
+        writer.close().unwrap();
+    }
+
     #[test]
     fn test_size_statistics_with_repetition_and_nulls() {
         let message_type = "


Test patch:
diff --git a/parquet/tests/arrow_reader/bad_data.rs b/parquet/tests/arrow_reader/bad_data.rs
index e2975c17c..743420314 100644
--- a/parquet/tests/arrow_reader/bad_data.rs
+++ b/parquet/tests/arrow_reader/bad_data.rs
@@ -111,6 +111,7 @@ fn test_arrow_rs_gh_6229_dict_header() {
 }
 
 #[test]
+#[cfg(feature = "snap")]
 fn test_arrow_rs_gh_6229_dict_levels() {
     let err = read_file("ARROW-RS-GH-6229-LEVELS.parquet").unwrap_err();
     assert_eq!(
diff --git a/parquet/tests/arrow_reader/checksum.rs b/parquet/tests/arrow_reader/checksum.rs
index c60908d8b..b500b7cb1 100644
--- a/parquet/tests/arrow_reader/checksum.rs
+++ b/parquet/tests/arrow_reader/checksum.rs
@@ -26,10 +26,10 @@ use parquet::arrow::arrow_reader::ArrowReaderBuilder;
 fn test_datapage_v1_corrupt_checksum() {
     let errors = read_file_batch_errors("datapage_v1-corrupt-checksum.parquet");
     assert_eq!(errors, [
-        Err("Parquet argument error: Parquet error: Page CRC checksum mismatch".to_string()), 
+        Err("Parquet argument error: Parquet error: Page CRC checksum mismatch".to_string()),
         Ok(()),
         Ok(()),
-        Err("Parquet argument error: Parquet error: Page CRC checksum mismatch".to_string()), 
+        Err("Parquet argument error: Parquet error: Page CRC checksum mismatch".to_string()),
         Err("Parquet argument error: Parquet error: Not all children array length are the same!".to_string())
     ]);
 }
@@ -41,6 +41,7 @@ fn test_datapage_v1_uncompressed_checksum() {
 }
 
 #[test]
+#[cfg(feature = "snap")]
 fn test_datapage_v1_snappy_compressed_checksum() {
     let errors = read_file_batch_errors("datapage_v1-snappy-compressed-checksum.parquet");
     assert_eq!(errors, [Ok(()), Ok(()), Ok(()), Ok(()), Ok(())]);
@@ -52,6 +53,7 @@ fn test_plain_dict_uncompressed_checksum() {
     assert_eq!(errors, [Ok(())]);
 }
 #[test]
+#[cfg(feature = "snap")]
 fn test_rle_dict_snappy_checksum() {
     let errors = read_file_batch_errors("rle-dict-snappy-checksum.parquet");
     assert_eq!(errors, [Ok(())]);