etseidl commented on code in PR #10420:
URL: https://github.com/apache/arrow-rs/pull/10420#discussion_r4041379249
##########
parquet/src/arrow/async_reader/mod.rs:
##########
@@ -675,6 +675,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(
Review Comment:
Minor nit: maybe name this `get_column_chunk_dictionary`?
##########
parquet/src/file/metadata/reader.rs:
##########
@@ -841,6 +890,41 @@ fn parse_index_data(push_decoder: &mut
ParquetMetaDataPushDecoder) -> Result<Par
}
}
+/// Returns the `[start, end)` byte range of a column chunk's dictionary
+/// page, or `None` if the column chunk has no dictionary page or is not a
+/// `BYTE_ARRAY` column (the only physical type [`decode_dictionary_page`]
+/// currently supports).
+#[cfg(feature = "arrow")]
+fn dictionary_page_byte_range(
+ metadata: &ParquetMetaData,
+ row_group_idx: usize,
+ column_idx: usize,
+) -> Result<Option<(u64, u64)>> {
Review Comment:
Maybe instead return `Range<u64>`, although to be fair we're pretty
inconsistent here with ranges.
##########
parquet/src/arrow/async_reader/mod.rs:
##########
@@ -1126,6 +1156,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()
{
Review Comment:
Maybe add some comments to explain what's going on in each step, esp since
this appears to be in part demonstrating the motivation for this change.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
