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;
 #[allow(unused_imports)] // Only used for benchmarks
 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::{
@@ -570,6 +570,30 @@ impl<T: AsyncFileReader + Send + 'static> ParquetRecordBatchStreamBuilder<T> {
         Ok(Some(Sbbf::new(&bitset)))
     }
 
+    /// 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).
+    ///
+    /// 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 the column chunk's page encoding statistics themselves.
+    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`]
@@ -984,6 +1008,44 @@ 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_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")]