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::{
@@ -675,6 +675,39 @@ 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).
+    ///
+    /// The returned array contains raw `Binary` values, even for columns
+    /// annotated as strings. Callers can compare byte slices directly or
+    /// convert values to UTF-8 explicitly.
+    ///
+    /// 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_column_chunk_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`]
@@ -972,6 +1005,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::PageIndexPolicy;
     use crate::file::metadata::ParquetMetaDataReader;
     use crate::file::metadata::page_index::PageIndex;
@@ -982,8 +1016,8 @@ mod tests {
     use arrow_array::cast::AsArray;
     use arrow_array::types::Int32Type;
     use arrow_array::{
-        Array, ArrayRef, BooleanArray, Int32Array, RecordBatchReader, Scalar, StringArray,
-        StructArray, UInt64Array,
+        Array, ArrayRef, BinaryArray, BooleanArray, Int32Array, RecordBatchReader, Scalar,
+        StringArray, StructArray, UInt64Array,
     };
     use arrow_schema::{DataType, Field, Schema};
     use futures::{StreamExt, TryStreamExt};
@@ -1126,6 +1160,157 @@ mod tests {
         );
     }
 
+    #[tokio::test]
+    async fn test_get_column_chunk_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 direct_metadata = ParquetMetaDataReader::new()
+            .parse_and_finish(&data)
+            .unwrap();
+        let mut direct_reader = TestReader::new(data.clone());
+        let direct = ParquetMetaDataReader::read_column_dictionary_async(
+            &mut direct_reader,
+            &direct_metadata,
+            0,
+            0,
+        )
+        .await
+        .unwrap()
+        .unwrap();
+        let direct = direct.as_any().downcast_ref::<BinaryArray>().unwrap();
+        assert_eq!(direct.value(0), b"alpha");
+
+        let async_reader = TestReader::new(data);
+        let mut builder = ParquetRecordBatchStreamBuilder::new(async_reader)
+            .await
+            .unwrap();
+
+        let dictionary = builder
+            .get_column_chunk_dictionary(0, 0)
+            .await
+            .unwrap()
+            .unwrap();
+        let dictionary = dictionary.as_any().downcast_ref::<BinaryArray>().unwrap();
+        let dictionary_values: Vec<&[u8]> = dictionary.iter().map(|v| v.unwrap()).collect();
+        assert_eq!(
+            dictionary_values,
+            vec![b"alpha".as_slice(), b"beta", b"gamma"]
+        );
+    }
+
+    // This test demonstrates row group pruning using dictionary pages,
+    // verifying that data pages of skipped row groups are not read
+    #[tokio::test]
+    async fn test_dictionary_selects_row_groups_without_reading_skipped_data() {
+        // Write two row groups, each dictionary-encoded and containing a single
+        // distinct string repeated 30 times: row group 0 is all "skip", row
+        // group 1 is all "target".
+        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));
+        // `requests` records the byte ranges fetched from the underlying reader,
+        // which we later use to check that we read only the needed page
+        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);
+
+        // For each row group, fetch and decode just the dictionary page
+        // to decide whether to read the whole row group
+        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_column_chunk_dictionary(row_group_idx, 0)
+                .await
+                .unwrap()
+                .unwrap();
+            let dictionary = dictionary.as_binary::<i32>();
+            if dictionary
+                .iter()
+                .any(|value| value == Some(b"target".as_slice()))
+            {
+                selected_row_groups.push(row_group_idx);
+            }
+        }
+        assert_eq!(selected_row_groups, vec![1]);
+
+        // Read the filtered row group and verify the values
+        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]);
+
+        // Finally, verify none of the requests overlapped the data pages
+        // of the skipped row group
+        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

@@ -186,9 +186,13 @@
 pub mod array_reader;
 #[cfg(not(feature = "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")]