comphead opened a new issue, #3394:
URL: https://github.com/apache/iceberg-rust/issues/3394
### Is your feature request related to a problem or challenge?
`get_arrow_projection_mask_with_field_ids` in `arrow/reader/projection.rs`
runs once per `FileScanTask`. It fills a `HashMap<FieldRef, i32>` with every
Arrow leaf that has a field ID, then looks every leaf up in it again.
arrow-schema's `impl Hash for Field` hashes the name, the whole `DataType`, and
the metadata, and it collects the metadata keys into a new sorted `Vec` on
every call. So each leaf with an ID costs one allocation for the insert and one
for the lookup. The map starts empty, and every resize hashes all of its keys
again.
The field ID is already in the leaf's `PARQUET:field_id` metadata, which the
first pass parses anyway.
One call to `ArrowReader::get_arrow_projection_mask` on `main` at
`b2e90754b`, for a file whose columns are all `int` with field IDs (release
build, Apple M3 Max, 500 calls each):
| Columns | Projected | Allocations | Time per call |
| ---: | ---: | ---: | ---: |
| 100 | 10 | 486 | 34 to 46 µs |
| 1,000 | 10 | 3,976 | 318 µs |
| 1,000 | 1,000 | 15,915 | 897 µs |
About 3,800 of the allocations for 1,000 columns come from the map. That is
2,000 hashes for the inserts and lookups, and about 1,800 more from resizing.
<details>
<summary>Measurement test (local only, with a thread-local counting
allocator)</summary>
Added to the tests in `crates/iceberg/src/arrow/reader/projection.rs` and
run with `cargo test --release -p iceberg --lib measure_projection_mask --
--nocapture --test-threads=1`:
```rust
mod alloc_count {
use std::alloc::{GlobalAlloc, Layout, System};
use std::cell::Cell;
thread_local! {
pub static ALLOCATIONS: Cell<usize> = const { Cell::new(0) };
}
struct Counting;
unsafe impl GlobalAlloc for Counting {
unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
ALLOCATIONS.with(|count| count.set(count.get() + 1));
unsafe { System.alloc(layout) }
}
unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
unsafe { System.dealloc(ptr, layout) }
}
unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size:
usize) -> *mut u8 {
ALLOCATIONS.with(|count| count.set(count.get() + 1));
unsafe { System.realloc(ptr, layout, new_size) }
}
}
#[global_allocator]
static GLOBAL: Counting = Counting;
}
#[test]
fn measure_projection_mask() {
use parquet::arrow::ArrowSchemaConverter;
for (columns, projected) in [(100, 10), (1000, 10), (1000, 1000)] {
let arrow_schema = Arc::new(ArrowSchema::new(
(1..=columns)
.map(|id| {
Field::new(format!("c{id}"), DataType::Int32,
true).with_metadata(
HashMap::from([(PARQUET_FIELD_ID_META_KEY.to_string(), id.to_string())]),
)
})
.collect::<Vec<_>>(),
));
let iceberg_schema = Schema::builder()
.with_fields(
(1..=columns)
.map(|id| {
NestedField::optional(id, format!("c{id}"),
Type::Primitive(PrimitiveType::Int)).into()
})
.collect::<Vec<_>>(),
)
.build()
.unwrap();
let parquet_schema =
ArrowSchemaConverter::new().convert(&arrow_schema).unwrap();
let field_ids: Vec<i32> = (1..=projected).collect();
let call = || {
ArrowReader::get_arrow_projection_mask(&field_ids,
&iceberg_schema, &parquet_schema, &arrow_schema, false)
.unwrap()
};
let before = alloc_count::ALLOCATIONS.with(|count| count.get());
std::hint::black_box(call());
let allocations = alloc_count::ALLOCATIONS.with(|count| count.get())
- before;
for _ in 0..20 {
std::hint::black_box(call());
}
let iterations = 500;
let start = std::time::Instant::now();
for _ in 0..iterations {
std::hint::black_box(call());
}
let per_call = start.elapsed() / iterations;
println!("MEASURE columns={columns} projected={projected}:
{allocations} allocations, {per_call:?} per call");
}
}
```
</details>
Related: #3383 covers other per-task planning costs.
### Describe the solution you'd like
Read each leaf's field ID from its metadata in the second pass as well, and
drop the map. Behavior does not 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]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
