feat(parquet): provide direct access to dictionary pages

parquet/src/arrow/array_reader/mod.rs

@@ -55,6 +55,9 @@ pub(crate) mod test_util;
 use crate::file::metadata::RowGroupMetaData;
 pub use builder::{ArrayReaderBuilder, CacheOptions, CacheOptionsBuilder};
 pub use byte_array::make_byte_array_reader;
+// Re-exported (beyond the `experimental` feature) so `file::metadata::dictionary`
+// can PLAIN-decode a raw dictionary page without duplicating this logic.
+pub(crate) use byte_array::ByteArrayDecoderPlain;
 pub use byte_array_dictionary::make_byte_array_dictionary_reader;
 #[cfg_attr(not(feature = "experimental"), expect(unused_imports))]
 pub use byte_view_array::make_byte_view_array_reader;

parquet/src/arrow/async_reader/mod.rs

@@ -33,7 +33,7 @@ use futures::future::{BoxFuture, FutureExt};
 use futures::stream::Stream;
 use tokio::io::{AsyncRead, AsyncReadExt, AsyncSeek, AsyncSeekExt};
 
-use arrow_array::RecordBatch;
+use arrow_array::{ArrayRef, RecordBatch};
 use arrow_schema::{Schema, SchemaRef};
 
 use crate::arrow::arrow_reader::{
@@ -670,6 +670,35 @@ impl<T: AsyncFileReader + Send + 'static> ParquetRecordBatchStreamBuilder<T> {
         self
     }
 
+    /// Read and decode the dictionary page for a column in a row group, if any.
+    ///
+    /// Returns `Ok(None)` if the column chunk has no dictionary page, or if
+    /// its physical type is not `BYTE_ARRAY` (the only physical type
+    /// currently supported).
+    ///
+    /// This can be used to inspect dictionary values when selecting or pruning
+    /// row groups before passing the selected indices to
+    /// [`ParquetRecordBatchStreamBuilder::with_row_groups`].
+    ///
+    /// Note this does not verify that the *entire* column chunk is
+    /// dictionary-encoded -- callers that need that guarantee (e.g. to treat
+    /// the dictionary as an exhaustive set of the column's values) should
+    /// check
+    /// [`crate::file::metadata::ColumnChunkMetaData::page_encoding_stats_mask`].
+    pub async fn get_row_group_column_dictionary(
+        &mut self,
+        row_group_idx: usize,
+        column_idx: usize,
+    ) -> Result<Option<ArrayRef>> {
+        ParquetMetaDataReader::read_column_dictionary_async(
+            &mut self.input.0,
+            &self.metadata,
+            row_group_idx,
+            column_idx,
+        )
+        .await
+    }
+
     /// Build a new [`ParquetRecordBatchStream`]
     ///
     /// See examples on [`ParquetRecordBatchStreamBuilder::new`]
@@ -967,6 +996,7 @@ mod tests {
     use crate::arrow::arrow_reader::{ArrowReaderMetadata, ArrowReaderOptions};
     use crate::arrow::schema::virtual_type::RowNumber;
     use crate::arrow::{ArrowWriter, AsyncArrowWriter, ProjectionMask};
+    use crate::basic::Encoding;
     use crate::file::metadata::ParquetMetaDataReader;
     use crate::file::metadata::{PageIndex, PageIndexPolicy};
     use crate::file::properties::WriterProperties;
@@ -1120,6 +1150,123 @@ mod tests {
         );
     }
 
+    #[tokio::test]
+    async fn test_get_row_group_column_dictionary() {
+        let schema = Arc::new(Schema::new(vec![Field::new("s", DataType::Utf8, false)]));
+        let values: Vec<&str> = ["alpha", "beta", "gamma"]
+            .iter()
+            .copied()
+            .cycle()
+            .take(30)
+            .collect();
+        let array: ArrayRef = Arc::new(StringArray::from(values));
+        let batch = RecordBatch::try_new(schema.clone(), vec![array]).unwrap();
+
+        let props = WriterProperties::builder()
+            .set_dictionary_enabled(true)
+            .build();
+        let mut buf = Vec::new();
+        {
+            let mut writer = ArrowWriter::try_new(&mut buf, schema, Some(props)).unwrap();
+            writer.write(&batch).unwrap();
+            writer.close().unwrap();
+        }
+        let data = Bytes::from(buf);
+
+        let async_reader = TestReader::new(data);
+        let mut builder = ParquetRecordBatchStreamBuilder::new(async_reader)
+            .await
+            .unwrap();
+
+        let dictionary = builder
+            .get_row_group_column_dictionary(0, 0)
+            .await
+            .unwrap()
+            .unwrap();
+        let dictionary = dictionary.as_any().downcast_ref::<StringArray>().unwrap();
+        let dictionary_values: Vec<&str> = dictionary.iter().map(|v| v.unwrap()).collect();
+        assert_eq!(dictionary_values, vec!["alpha", "beta", "gamma"]);
+    }
+
+    #[tokio::test]
+    async fn test_dictionary_selects_row_groups_without_reading_skipped_data() {
+        let schema = Arc::new(Schema::new(vec![Field::new("s", DataType::Utf8, false)]));
+        let row_group_values = ["skip", "target"];
+        let props = WriterProperties::builder()
+            .set_dictionary_enabled(true)
+            .build();
+        let mut buf = Vec::new();
+        {
+            let mut writer = ArrowWriter::try_new(&mut buf, schema.clone(), Some(props)).unwrap();
+            for value in row_group_values {
+                let array: ArrayRef = Arc::new(StringArray::from(vec![value; 30]));
+                let batch = RecordBatch::try_new(schema.clone(), vec![array]).unwrap();
+                writer.write(&batch).unwrap();
+                writer.flush().unwrap();
+            }
+            writer.close().unwrap();
+        }
+
+        let async_reader = TestReader::new(Bytes::from(buf));
+        let requests = async_reader.requests.clone();
+        let mut builder = ParquetRecordBatchStreamBuilder::new(async_reader)
+            .await
+            .unwrap();
+        let metadata = builder.metadata().clone();
+        assert_eq!(metadata.num_row_groups(), 2);
+
+        let mut selected_row_groups = Vec::new();
+        let mut dictionary_ranges = Vec::new();
+        for row_group_idx in 0..metadata.num_row_groups() {
+            let column = metadata.row_group(row_group_idx).column(0);
+            let encoding_mask = column.page_encoding_stats_mask().unwrap();
+            assert!(
+                encoding_mask.is_only(Encoding::PLAIN_DICTIONARY)
+                    || encoding_mask.is_only(Encoding::RLE_DICTIONARY)
+            );
+
+            let dictionary_start = column.dictionary_page_offset().unwrap() as usize;
+            let data_start = column.data_page_offset() as usize;
+            dictionary_ranges.push(dictionary_start..data_start);
+
+            let dictionary = builder
+                .get_row_group_column_dictionary(row_group_idx, 0)
+                .await
+                .unwrap()
+                .unwrap();
+            let dictionary = dictionary.as_string::<i32>();
+            if dictionary.iter().any(|value| value == Some("target")) {
+                selected_row_groups.push(row_group_idx);
+            }
+        }
+        assert_eq!(selected_row_groups, vec![1]);
+
+        let batches: Vec<_> = builder
+            .with_row_groups(selected_row_groups)
+            .build()
+            .unwrap()
+            .try_collect()
+            .await
+            .unwrap();
+        let values: Vec<_> = batches
+            .iter()
+            .flat_map(|batch| batch.column(0).as_string::<i32>().iter())
+            .collect();
+        assert_eq!(values, vec![Some("target"); 30]);
+
+        let skipped_column = metadata.row_group(0).column(0);
+        let (skipped_start, skipped_len) = skipped_column.byte_range();
+        let skipped_data_range =
+            skipped_column.data_page_offset() as usize..(skipped_start + skipped_len) as usize;
+        let requests = requests.lock().unwrap();
+        for dictionary_range in dictionary_ranges {
+            assert!(requests.contains(&dictionary_range));
+        }
+        assert!(requests.iter().all(|request| {
+            request.end <= skipped_data_range.start || request.start >= skipped_data_range.end
+        }));
+    }
+
     #[tokio::test]
     async fn test_async_reader_with_next_row_group() {
         let testdata = arrow::util::test_util::parquet_test_data();

parquet/src/arrow/mod.rs

@@ -180,9 +180,13 @@
 //! ```
 
 experimental!(mod array_reader);
+// Re-exported (beyond the `experimental` feature) so `file::metadata::dictionary`
+// can PLAIN-decode a raw dictionary page without duplicating this logic.
+pub(crate) use array_reader::ByteArrayDecoderPlain;
 pub mod arrow_reader;
 pub mod arrow_writer;
 mod buffer;
+pub(crate) use buffer::offset_buffer::OffsetBuffer;
 mod decoder;
 
 #[cfg(feature = "async")]