jayzhan211 commented on PR #25468:
URL: https://github.com/apache/datafusion/pull/25468#issuecomment-5740487563
@SubhamSinghal
## Emit is O(chunks × store batches × batch_size) for dictionary columns
`batch_refs` = every batch in the store, passed to every per-chunk
`interleave_record_batch`. `interleave_dictionaries` does per-input-array work
per call (key mask per dictionary + `indices.iter().filter(|(a, _)| *a ==
a_idx)`; the fallback concats all B value arrays into each output) → emit is
O(chunks × B × batch_size) for Dictionary columns; base was linear.
Measured (release, 8192-row batches, `Dictionary<Int32, Utf8>` with a
per-batch dictionary, 8 partitions × 4 ob values, all rows retained), emit only:
| batches | main | PR | PR + fix |
|---|---|---|---|
| 128 | 0.010s | 0.052s | 0.012s |
| 512 | 0.064s | 0.735s | 0.081s |
| 2048 | 0.263s | 16.1s | 0.335s |
Fix: build the refs per chunk from only the batches it touches (neutral on
the tiny-entry shape: 1.29s vs 1.29s dict, 0.58s vs 0.58s utf8).
```diff
- let mut batch_refs = Vec::with_capacity(store.len());
- let mut batch_id_pos = HashMap::with_capacity(store.len());
- for (array_pos, (batch_id, entry)) in
store.batches.iter().enumerate() {
- batch_refs.push(&entry.batch);
- batch_id_pos.insert(*batch_id, array_pos);
- }
+ // Rebuilt per chunk: `interleave` does per-input-array work for
+ // some types (dictionaries), so only pass batches the chunk uses.
+ let mut batch_refs: Vec<&RecordBatch> = Vec::new();
+ let mut batch_id_pos: HashMap<u32, usize> = HashMap::new();
```
```diff
- let array_pos = batch_id_pos[&entry.batch_id];
+ let batch = &store
+ .get(entry.batch_id)
+ .expect("retained batch_id present in store")
+ .batch;
+ let mut array_pos = None;
for row in entry.row_indices {
- indices.push((array_pos, row as usize));
+ let pos = *array_pos.get_or_insert_with(|| {
+
*batch_id_pos.entry(entry.batch_id).or_insert_with(|| {
+ batch_refs.push(batch);
+ batch_refs.len() - 1
+ })
+ });
+ indices.push((pos, row as usize));
if indices.len() == batch_size {
let b = interleave_record_batch(&batch_refs,
&indices)?;
(&b).record_output(&metrics.baseline);
out.push(Ok(b));
indices.clear();
+ batch_refs.clear();
+ batch_id_pos.clear();
+ array_pos = None;
}
```
Please also add an emit test with a Dictionary column spanning ≥2 chunks.
```
Tiny entries (1 row per GroupEntry, 1024 partitions), 512 batches:
┌────────────────────────────────────────┬────────────────────┬──────────────────┐
│ d type │ base insert / emit │ PR insert /
emit │
├────────────────────────────────────────┼────────────────────┼──────────────────┤
│ Utf8 │ 8.06s / 2.32s │ 1.01s /
0.63s │
├────────────────────────────────────────┼────────────────────┼──────────────────┤
│ Dict<Int32,Utf8>, per-batch dictionary │ 7.05s / 2.83s │ 0.71s /
1.23s │
└────────────────────────────────────────┴────────────────────┴──────────────────┘
Both PR claims hold on this shape: insert is ~8x faster (the size() walk is
gone) and emit is 2.3–3.7x faster.
Large entries (8 partitions × 4 ob values, 256 rows per entry), emit only:
┌─────────┬─────────────────┬─────────────────┬──────────────────────┐
│ batches │ Utf8 base / PR │ Dict base / PR │ Dict, PR + fix below │
├─────────┼─────────────────┼─────────────────┼──────────────────────┤
│ 128 │ 0.018s / 0.017s │ 0.010s / 0.052s │ 0.012s │
├─────────┼─────────────────┼─────────────────┼──────────────────────┤
│ 512 │ 0.106s / 0.109s │ 0.064s / 0.735s │ 0.081s │
├─────────┼─────────────────┼─────────────────┼──────────────────────┤
│ 2048 │ 0.417s / 0.481s │ 0.263s / 16.1s │ 0.335s │
└─────────┴─────────────────┴─────────────────┴──────────────────────┘
Fix vs PR on the tiny-entry shape (interleaved, 2 runs each), 512 batches:
┌────────┬───────────────┬───────────────┐
│ d type │ PR emit │ PR + fix emit │
├────────┼───────────────┼───────────────┤
│ Utf8 │ 0.59s / 0.58s │ 0.62s / 0.58s │
├────────┼───────────────┼───────────────┤
│ Dict │ 1.29s / 1.29s │ 1.32s / 1.26s │
└────────┴───────────────┴───────────────┘
```
--
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]
