adriangb commented on code in PR #11031:
URL: https://github.com/apache/arrow-rs/pull/11031#discussion_r4077685255
##########
parquet/src/file/metadata/writer.rs:
##########
@@ -91,10 +98,11 @@ impl<'a, W: Write> ThriftMetadataWriter<'a, W> {
// set offset and index for offset index
column_metadata.offset_index_offset = Some(start_offset as
i64);
column_metadata.offset_index_length = Some((end_offset -
start_offset) as i32);
+ offidx_vec[row_group_idx][column_idx] =
Some(offset_index.clone());
}
Review Comment:
**Bug** (A1). If the provider has no offset index for a chunk, the chunk
keeps `offset_index_offset/length` from the input metadata. These offsets point
into the source file, not into the new buffer. A partial provider (for example
Example 2 in `examples/custom_page_index.rs`) writes metadata that fails to
read back:
```text
source file 4386 B; provider returns column 0 only (custom provider or
PageIndexBuilder)
written metadata: 890 bytes
rg0 col0: column_index_range=Some(0..23)
offset_index_range=Some(46..58)
rg0 col1: column_index_range=Some(3207..3230)
offset_index_range=Some(3334..3346) <- source file
read back with page index: Err(EOF: Parquet file too small. Range 0..3394 is
beyond file bounds 890)
```
With this suggestion (and the same change at L145), col 1 and col 2 have no
index ranges and the read back is `Ok`. On `main`, a partial `PageIndex` fails
the same way, so this is not a regression. But custom providers are usually
partial, so this PR is the right place to fix it.
Trade-off: clearing removes offsets that are still correct only when the
output is appended to the source file. If that use matters, return an error
instead.
```suggestion
offidx_vec[row_group_idx][column_idx] =
Some(offset_index.clone());
} else {
// no offset index for this chunk: do not keep an offset
into the source file
column_metadata.offset_index_offset = None;
column_metadata.offset_index_length = None;
}
```
##########
parquet/src/file/metadata/writer.rs:
##########
@@ -71,14 +71,21 @@ impl<'a, W: Write> ThriftMetadataWriter<'a, W> {
/// of the serialized offset indexes.
fn write_offset_indexes(
&mut self,
- offset_indexes: &[Vec<Option<OffsetIndexMetaData>>],
- ) -> Result<()> {
+ page_index: &Arc<dyn PageIndexProvider>,
+ ) -> Result<Option<Vec<Vec<Option<OffsetIndexMetaData>>>>> {
+ let mut offset_indexes =
+ PageIndexBuilder::empty_index(self.row_groups.len(),
self.schema_descr.num_columns());
+ let offidx_vec = offset_indexes.as_mut().unwrap();
Review Comment:
**Perf** (A2). These lines start a deep copy of every index. Only the
returned `ParquetMetaData` uses the copy:
- `SerializedFileWriter` passes the `PageIndex` that it built from the same
vectors. On `main` these vectors were moved, not copied.
- `ParquetMetaDataWriter::finish` drops the returned metadata, so it never
uses the copy (also true on `main`).
Measured with a counting global allocator (500 Int32 columns × 40 row
groups, page statistics):
| | `main` | this PR | PR + sketch below |
|---|---|---|---|
| `ArrowWriter::close`: bytes allocated | 0.83 MB | 7.26 MB | 0.83 MB |
| `ArrowWriter::close`: allocations | 120,006 | 220,089 | 120,006 |
| `ArrowWriter::close`: peak | 35 kB | 6.46 MB | 35 kB |
| `ParquetMetaDataWriter::finish`: bytes allocated (without the 8 MiB output
buffer) | 15.7 MB | 15.7 MB | 9.3 MB |
Suggestion: do not copy. Write from the provider and return the same `Arc`.
No downcast is necessary, and `empty_index` can stay private.
```diff
- fn write_offset_indexes(
- &mut self,
- page_index: &Arc<dyn PageIndexProvider>,
- ) -> Result<Option<Vec<Vec<Option<OffsetIndexMetaData>>>>> {
+ fn write_offset_indexes(&mut self, page_index: &dyn PageIndexProvider)
-> Result<()> {
// same loop without `offidx_vec`; same for write_column_indexes;
remove finalize_*
pub fn finish(mut self) -> Result<ParquetMetaData> {
let page_index = self.page_index.take();
- let column_indexes =
self.finalize_column_indexes(page_index.as_ref())?;
- let offset_indexes =
self.finalize_offset_indexes(page_index.as_ref())?;
+ if let Some(pi) = page_index.as_deref() {
+ if pi.has_column_indexes() {
+ self.write_column_indexes(pi)?;
+ }
+ if pi.has_offset_indexes() {
+ self.write_offset_indexes(pi)?;
+ }
+ }
...
- let builder =
ParquetMetaDataBuilder::new(file_metadata).set_page_index(Some(Arc::new(
- PageIndex::new(column_indexes, offset_indexes),
- )));
+ let builder =
ParquetMetaDataBuilder::new(file_metadata).set_page_index(page_index);
```
The metadata from `SerializedFileWriter::close` changes in two edge cases
(observed). No existing test depends on the old result:
| case | this PR: `page_index()` | with sketch |
|---|---|---|
| no statistics, offset index disabled | `Some`, both `has_*` false | `None`
(same as the reader) |
| 0 row groups | `Some`, `has_offset_indexes()` false | `Some`,
`has_offset_indexes()` true |
##########
parquet/src/file/metadata/writer.rs:
##########
@@ -127,21 +142,28 @@ impl<'a, W: Write> ThriftMetadataWriter<'a, W> {
column_metadata.column_index_length =
Some((end_offset - start_offset) as i32);
}
+ colidx_vec[row_group_idx][column_idx] =
Some(column_index.clone());
}
Review Comment:
**Bug** (A1, continued). Same change for the column index.
```suggestion
colidx_vec[row_group_idx][column_idx] =
Some(column_index.clone());
} else {
// no column index for this chunk: do not keep an offset
into the source file
column_metadata.column_index_offset = None;
column_metadata.column_index_length = None;
}
```
##########
parquet/src/file/metadata/writer.rs:
##########
@@ -333,17 +350,6 @@ impl<'a, W: Write> ThriftMetadataWriter<'a, W> {
/// metadata writer. Then set the corresponding `bloom_filter_offset` and
/// `bloom_filter_length` on [`ColumnChunkMetaData`] passed to this writer.
///
Review Comment:
**Docs** (A5). The PR removes the warning, but the docs do not say that
custom providers are now serialized. The `[PageIndexProvider]` link definition
at L367 is now unused. Proposed text (the second sentence assumes A1):
```suggestion
/// `bloom_filter_length` on [`ColumnChunkMetaData`] passed to this writer.
///
/// The writer serializes the page index of any [`PageIndexProvider`]. If the
/// provider has no index for a column chunk, the writer clears the offset
and
/// length of that index in the [`ColumnChunkMetaData`].
///
```
##########
parquet/src/file/writer.rs:
##########
@@ -377,10 +378,27 @@ impl<W: Write + Send> SerializedFileWriter<W> {
encoder = encoder.with_key_value_metadata(key_value_metadata)
}
- encoder = encoder.with_column_indexes(column_indexes);
- if !self.props.offset_index_disabled() {
- encoder = encoder.with_offset_indexes(offset_indexes);
+ // check for empty column index
+ let column_indexes = if column_indexes.is_empty()
+ || column_indexes
+ .iter()
+ .all(|cis| cis.iter().all(|ci| ci.is_none()))
+ {
Review Comment:
**Nit** (A6). `all()` returns `true` for an empty iterator, so the
`is_empty()` check is redundant.
```suggestion
let column_indexes = if
column_indexes.iter().flatten().all(Option::is_none) {
```
##########
parquet/src/file/metadata/writer.rs:
##########
@@ -464,21 +470,10 @@ impl<'a, W: Write> ParquetMetaDataWriter<'a, W> {
self.write_path_in_schema,
);
- // Downcast to PageIndex to access raw index structures for
serialization.
- // Page indexes from custom PageIndexProviders are not written. See
- // <https://github.com/apache/arrow-rs/issues/11030>
if let Some(page_index_arc) = self.metadata.page_index.as_ref()
- && let Some(page_index) = page_index_arc
- .as_any()
- .downcast_ref::<crate::file::metadata::PageIndex>()
+ && (page_index_arc.has_column_indexes() ||
page_index_arc.has_offset_indexes())
{
- if let Some(column_indexes) = page_index.column_indexes_raw() {
- encoder = encoder.with_column_indexes(column_indexes.clone());
- }
-
- if let Some(offset_indexes) = page_index.offset_indexes_raw() {
- encoder = encoder.with_offset_indexes(offset_indexes.clone());
- }
+ encoder = encoder.with_page_index(page_index_arc.clone());
}
Review Comment:
**Test gap** (A3). `finalize_column_indexes` and `finalize_offset_indexes`
already check `has_*()`, so this guard has no effect. But the mutant `||` →
`&&` here survives the full suite (`--lib`, `arrow_reader`, `arrow_writer`,
`encryption`). With the mutant, an offset-index-only page index
(`EnabledStatistics::Chunk`) is not written, and the read back fails. Suggest:
remove the guard, and add
`test_metadata_read_write_roundtrip_offset_index_only` from the review body.
That test kills the mutant.
```suggestion
if let Some(page_index) = self.metadata.page_index.clone() {
encoder = encoder.with_page_index(page_index);
}
```
--
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]
