Repository: apache/arrow-rs
Pull request #6308 (2024-08-25 18:01:41)

Parquet: add `union` method to `RowSelection`
Complements the existing `RowSelection::intersection` method. Useful for OR-ing row selections together, complementing 
 the use of `RowSelection::intersection` for AND-ing row selections

# Which issue does this PR close?

Closes https://github.com/apache/arrow-rs/issues/6307.

# Are there any user-facing changes?

Yes. Doc comments for `RowSelection::union` have been added accordingly as part of this PR.

<!--
If there are user-facing changes then we may require documentation to be updated before approving the PR.
-->



Patch:
diff --git a/parquet/src/arrow/arrow_reader/selection.rs b/parquet/src/arrow/arrow_reader/selection.rs
index 0287e5b42..ce3fbbf4f 100644
--- a/parquet/src/arrow/arrow_reader/selection.rs
+++ b/parquet/src/arrow/arrow_reader/selection.rs
@@ -343,6 +343,16 @@ impl RowSelection {
         intersect_row_selections(&self.selectors, &other.selectors)
     }
 
+    /// Compute the union of two [`RowSelection`]
+    /// For example:
+    /// self:      NNYYYYNNYYNYN
+    /// other:     NYNNNNNNN
+    ///
+    /// returned:  NYYYYYNNYYNYN
+    pub fn union(&self, other: &Self) -> Self {
+        union_row_selections(&self.selectors, &other.selectors)
+    }
+
     /// Returns `true` if this [`RowSelection`] selects any rows
     pub fn selects_any(&self) -> bool {
         self.selectors.iter().any(|x| !x.skip)
@@ -536,6 +546,92 @@ fn intersect_row_selections(left: &[RowSelector], right: &[RowSelector]) -> RowS
     iter.collect()
 }
 
+/// Combine two lists of `RowSelector` return the union of them
+/// For example:
+/// self:      NNYYYYNNYYNYN
+/// other:     NYNNNNNNY
+///
+/// returned:  NYYYYYNNYYNYN
+///
+/// This can be removed from here once RowSelection::union is in parquet::arrow
+fn union_row_selections(left: &[RowSelector], right: &[RowSelector]) -> RowSelection {
+    let mut l_iter = left.iter().copied().peekable();
+    let mut r_iter = right.iter().copied().peekable();
+
+    let iter = std::iter::from_fn(move || {
+        loop {
+            let l = l_iter.peek_mut();
+            let r = r_iter.peek_mut();
+
+            match (l, r) {
+                (Some(a), _) if a.row_count == 0 => {
+                    l_iter.next().unwrap();
+                }
+                (_, Some(b)) if b.row_count == 0 => {
+                    r_iter.next().unwrap();
+                }
+                (Some(l), Some(r)) => {
+                    return match (l.skip, r.skip) {
+                        // Skip both ranges
+                        (true, true) => {
+                            if l.row_count < r.row_count {
+                                let skip = l.row_count;
+                                r.row_count -= l.row_count;
+                                l_iter.next();
+                                Some(RowSelector::skip(skip))
+                            } else {
+                                let skip = r.row_count;
+                                l.row_count -= skip;
+                                r_iter.next();
+                                Some(RowSelector::skip(skip))
+                            }
+                        }
+                        // Keep rows from left
+                        (false, true) => {
+                            if l.row_count < r.row_count {
+                                r.row_count -= l.row_count;
+                                l_iter.next()
+                            } else {
+                                let r_row_count = r.row_count;
+                                l.row_count -= r_row_count;
+                                r_iter.next();
+                                Some(RowSelector::select(r_row_count))
+                            }
+                        }
+                        // Keep rows from right
+                        (true, false) => {
+                            if l.row_count < r.row_count {
+                                let l_row_count = l.row_count;
+                                r.row_count -= l_row_count;
+                                l_iter.next();
+                                Some(RowSelector::select(l_row_count))
+                            } else {
+                                l.row_count -= r.row_count;
+                                r_iter.next()
+                            }
+                        }
+                        // Keep at least one
+                        _ => {
+                            if l.row_count < r.row_count {
+                                r.row_count -= l.row_count;
+                                l_iter.next()
+                            } else {
+                                l.row_count -= r.row_count;
+                                r_iter.next()
+                            }
+                        }
+                    };
+                }
+                (Some(_), None) => return l_iter.next(),
+                (None, Some(_)) => return r_iter.next(),
+                (None, None) => return None,
+            }
+        }
+    });
+
+    iter.collect()
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -1213,4 +1309,40 @@ mod tests {
             ]
         );
     }
+
+    #[test]
+    fn test_union() {
+        let selection = RowSelection::from(vec![RowSelector::select(1048576)]);
+        let result = selection.union(&selection);
+        assert_eq!(result, selection);
+
+        // NYNYY
+        let a = RowSelection::from(vec![
+            RowSelector::skip(10),
+            RowSelector::select(10),
+            RowSelector::skip(10),
+            RowSelector::select(20),
+        ]);
+
+        // NNYYNYN
+        let b = RowSelection::from(vec![
+            RowSelector::skip(20),
+            RowSelector::select(20),
+            RowSelector::skip(10),
+            RowSelector::select(10),
+            RowSelector::skip(10),
+        ]);
+
+        let result = a.union(&b);
+
+        // NYYYYYN
+        assert_eq!(
+            result.iter().collect::<Vec<_>>(),
+            vec![
+                &RowSelector::skip(10),
+                &RowSelector::select(50),
+                &RowSelector::skip(10),
+            ]
+        );
+    }
 }
diff --git a/parquet/src/encodings/rle.rs b/parquet/src/encodings/rle.rs
index 581f14b3c..97a122941 100644
--- a/parquet/src/encodings/rle.rs
+++ b/parquet/src/encodings/rle.rs
@@ -20,7 +20,6 @@ use std::{cmp, mem::size_of};
 use bytes::Bytes;
 
 use crate::errors::{ParquetError, Result};
-use crate::util::bit_util::from_le_slice;
 use crate::util::bit_util::{self, BitReader, BitWriter, FromBytes};
 
 /// Rle/Bit-Packing Hybrid Encoding
@@ -356,13 +355,13 @@ impl RleDecoder {
         }
 
         let value = if self.rle_left > 0 {
-            let rle_value = from_le_slice(
+            let rle_value = T::try_from_le_slice(
                 &self
                     .current_value
                     .as_mut()
                     .expect("current_value should be Some")
                     .to_ne_bytes(),
-            );
+            )?;
             self.rle_left -= 1;
             rle_value
         } else {
@@ -388,9 +387,9 @@ impl RleDecoder {
                 let num_values =
                     cmp::min(buffer.len() - values_read, self.rle_left as usize);
                 for i in 0..num_values {
-                    let repeated_value = from_le_slice(
+                    let repeated_value = T::try_from_le_slice(
                         &self.current_value.as_mut().unwrap().to_ne_bytes(),
-                    );
+                    )?;
                     buffer[values_read + i] = repeated_value;
                 }
                 self.rle_left -= num_values as u32;
diff --git a/parquet/src/file/page_index/index.rs b/parquet/src/file/page_index/index.rs
index 0c23e4aa3..662ba4562 100644
--- a/parquet/src/file/page_index/index.rs
+++ b/parquet/src/file/page_index/index.rs
@@ -23,7 +23,6 @@ use crate::data_type::{AsBytes, ByteArray, FixedLenByteArray, Int96};
 use crate::errors::ParquetError;
 use crate::file::metadata::LevelHistogram;
 use crate::format::{BoundaryOrder, ColumnIndex};
-use crate::util::bit_util::from_le_slice;
 use std::fmt::Debug;
 
 /// Typed statistics for one data page
@@ -192,7 +191,7 @@ impl<T: ParquetValueType> NativeIndex<T> {
         let indexes = index
             .min_values
             .iter()
-            .zip(index.max_values.into_iter())
+            .zip(index.max_values.iter())
             .zip(index.null_pages.into_iter())
             .zip(null_counts.into_iter())
             .zip(rep_hists.into_iter())
@@ -205,9 +204,10 @@ impl<T: ParquetValueType> NativeIndex<T> {
                     let (min, max) = if is_null {
                         (None, None)
                     } else {
-                        let min = min.as_slice();
-                        let max = max.as_slice();
-                        (Some(from_le_slice::<T>(min)), Some(from_le_slice::<T>(max)))
+                        (
+                            Some(T::try_from_le_slice(min)?),
+                            Some(T::try_from_le_slice(max)?),
+                        )
                     };
                     Ok(PageIndex {
                         min,
@@ -321,4 +321,29 @@ mod tests {
         assert_eq!(page_index.repetition_level_histogram(), None);
         assert_eq!(page_index.definition_level_histogram(), None);
     }
+
+    #[test]
+    fn test_invalid_column_index() {
+        let column_index = ColumnIndex {
+            null_pages: vec![true, false],
+            min_values: vec![
+                vec![],
+                vec![], // this shouldn't be empty as null_pages[1] is false
+            ],
+            max_values: vec![
+                vec![],
+                vec![], // this shouldn't be empty as null_pages[1] is false
+            ],
+            null_counts: None,
+            repetition_level_histograms: None,
+            definition_level_histograms: None,
+            boundary_order: BoundaryOrder::UNORDERED,
+        };
+
+        let err = NativeIndex::<i32>::try_new(column_index).unwrap_err();
+        assert_eq!(
+            err.to_string(),
+            "Parquet error: error converting value, expected 4 bytes got 0"
+        );
+    }
 }
diff --git a/parquet/src/file/serialized_reader.rs b/parquet/src/file/serialized_reader.rs
index 0a3e51931..b8ee4001a 100644
--- a/parquet/src/file/serialized_reader.rs
+++ b/parquet/src/file/serialized_reader.rs
@@ -781,7 +781,6 @@ mod tests {
     use crate::file::writer::SerializedFileWriter;
     use crate::record::RowAccessor;
     use crate::schema::parser::parse_message_type;
-    use crate::util::bit_util::from_le_slice;
     use crate::util::test_common::file_util::{get_test_file, get_test_path};
 
     use super::*;
@@ -1537,8 +1536,8 @@ mod tests {
         assert_eq!(row_group_index.indexes.len(), page_size);
         assert_eq!(row_group_index.boundary_order, boundary_order);
         row_group_index.indexes.iter().all(|x| {
-            x.min.as_ref().unwrap() >= &from_le_slice::<T>(min_max.0)
-                && x.max.as_ref().unwrap() <= &from_le_slice::<T>(min_max.1)
+            x.min.as_ref().unwrap() >= &T::try_from_le_slice(min_max.0).unwrap()
+                && x.max.as_ref().unwrap() <= &T::try_from_le_slice(min_max.1).unwrap()
         });
     }
 
diff --git a/parquet/src/file/statistics.rs b/parquet/src/file/statistics.rs
index 680c75d6b..854900f1e 100644
--- a/parquet/src/file/statistics.rs
+++ b/parquet/src/file/statistics.rs
@@ -47,7 +47,7 @@ use crate::basic::Type;
 use crate::data_type::private::ParquetValueType;
 use crate::data_type::*;
 use crate::errors::{ParquetError, Result};
-use crate::util::bit_util::from_le_slice;
+use crate::util::bit_util::FromBytes;
 
 pub(crate) mod private {
     use super::*;
@@ -186,14 +186,18 @@ pub fn from_thrift(
                     // INT96 statistics may not be correct, because comparison is signed
                     // byte-wise, not actual timestamps. It is recommended to ignore
                     // min/max statistics for INT96 columns.
-                    let min = min.map(|data| {
+                    let min = if let Some(data) = min {
                         assert_eq!(data.len(), 12);
-                        from_le_slice::<Int96>(&data)
-                    });
-                    let max = max.map(|data| {
+                        Some(Int96::try_from_le_slice(&data)?)
+                    } else {
+                        None
+                    };
+                    let max = if let Some(data) = max {
                         assert_eq!(data.len(), 12);
-                        from_le_slice::<Int96>(&data)
-                    });
+                        Some(Int96::try_from_le_slice(&data)?)
+                    } else {
+                        None
+                    };
                     Statistics::int96(min, max, distinct_count, null_count, old_format)
                 }
                 Type::FLOAT => Statistics::float(
diff --git a/parquet/src/util/bit_util.rs b/parquet/src/util/bit_util.rs
index adbf45014..a17202254 100644
--- a/parquet/src/util/bit_util.rs
+++ b/parquet/src/util/bit_util.rs
@@ -23,12 +23,6 @@ use crate::data_type::{AsBytes, ByteArray, FixedLenByteArray, Int96};
 use crate::errors::{ParquetError, Result};
 use crate::util::bit_pack::{unpack16, unpack32, unpack64, unpack8};
 
-#[inline]
-pub fn from_le_slice<T: FromBytes>(bs: &[u8]) -> T {
-    // TODO: propagate the error (#3577)
-    T::try_from_le_slice(bs).unwrap()
-}
-
 #[inline]
 fn array_from_slice<const N: usize>(bs: &[u8]) -> Result<[u8; N]> {
     // Need to slice as may be called with zero-padded values
@@ -91,15 +85,22 @@ unsafe impl FromBytes for Int96 {
     type Buffer = [u8; 12];
 
     fn try_from_le_slice(b: &[u8]) -> Result<Self> {
-        Ok(Self::from_le_bytes(array_from_slice(b)?))
+        let bs: [u8; 12] = array_from_slice(b)?;
+        let mut i = Int96::new();
+        i.set_data(
+            u32::try_from_le_slice(&bs[0..4])?,
+            u32::try_from_le_slice(&bs[4..8])?,
+            u32::try_from_le_slice(&bs[8..12])?,
+        );
+        Ok(i)
     }
 
     fn from_le_bytes(bs: Self::Buffer) -> Self {
         let mut i = Int96::new();
         i.set_data(
-            from_le_slice(&bs[0..4]),
-            from_le_slice(&bs[4..8]),
-            from_le_slice(&bs[8..12]),
+            u32::try_from_le_slice(&bs[0..4]).unwrap(),
+            u32::try_from_le_slice(&bs[4..8]).unwrap(),
+            u32::try_from_le_slice(&bs[8..12]).unwrap(),
         );
         i
     }
@@ -438,7 +439,7 @@ impl BitReader {
         }
 
         // TODO: better to avoid copying here
-        Some(from_le_slice(v.as_bytes()))
+        T::try_from_le_slice(v.as_bytes()).ok()
     }
 
     /// Read multiple values from their packed representation where each element is represented
@@ -1026,7 +1027,7 @@ mod tests {
             .collect();
 
         // Generic values used to check against actual values read from `get_batch`.
-        let expected_values: Vec<T> = values.iter().map(|v| from_le_slice(v.as_bytes())).collect();
+        let expected_values: Vec<T> = values.iter().map(|v| T::try_from_le_slice(v.as_bytes()).unwrap()).collect();
 
         (0..total).for_each(|i| writer.put_value(values[i], num_bits));
 


Test patch:
diff --git a/parquet/tests/arrow_reader/bad_data.rs b/parquet/tests/arrow_reader/bad_data.rs
index 6e325f119..cbd5d4d3b 100644
--- a/parquet/tests/arrow_reader/bad_data.rs
+++ b/parquet/tests/arrow_reader/bad_data.rs
@@ -134,3 +134,28 @@ fn read_file(name: &str) -> Result<usize, ParquetError> {
     }
     Ok(num_rows)
 }
+
+#[cfg(feature = "async")]
+#[tokio::test]
+async fn bad_metadata_err() {
+    use bytes::Bytes;
+    use parquet::arrow::async_reader::MetadataLoader;
+
+    let metadata_buffer = Bytes::from_static(include_bytes!("bad_raw_metadata.bin"));
+
+    let metadata_length = metadata_buffer.len();
+
+    let mut reader = std::io::Cursor::new(&metadata_buffer);
+    let mut loader = MetadataLoader::load(&mut reader, metadata_length, None)
+        .await
+        .unwrap();
+    loader.load_page_index(false, false).await.unwrap();
+    loader.load_page_index(false, true).await.unwrap();
+
+    let err = loader.load_page_index(true, false).await.unwrap_err();
+
+    assert_eq!(
+        err.to_string(),
+        "Parquet error: error converting value, expected 4 bytes got 0"
+    );
+}