Skip to content

Commit 38b3d06

Browse files
committed
perf(repartition): compact sparse dictionary outputs
1 parent 3269f3f commit 38b3d06

3 files changed

Lines changed: 90 additions & 5 deletions

File tree

‎Cargo.lock‎

Lines changed: 1 addition & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎datafusion/physical-plan/Cargo.toml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -64,6 +64,7 @@ arrow-data = { workspace = true }
6464
arrow-ipc = { workspace = true, features = ["lz4", "zstd"] }
6565
arrow-ord = { workspace = true }
6666
arrow-schema = { workspace = true }
67+
arrow-select = { workspace = true }
6768
async-trait = { workspace = true }
6869
bytes = { workspace = true }
6970
datafusion-common = { workspace = true }

‎datafusion/physical-plan/src/repartition/mod.rs‎

Lines changed: 88 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -46,10 +46,14 @@ use crate::{
4646
PlanProperties, ReplaceChildrenOptions, Statistics, validate_child_count,
4747
};
4848

49-
use arrow::array::{PrimitiveArray, RecordBatch, RecordBatchOptions, UInt64Array};
49+
use arrow::array::{
50+
Array, ArrayRef, AsArray, PrimitiveArray, RecordBatch, RecordBatchOptions,
51+
UInt64Array,
52+
};
5053
use arrow::compute::take_arrays;
5154
use arrow::datatypes::{DataType, Schema, SchemaRef, UInt32Type};
5255
use arrow_schema::SortOptions;
56+
use arrow_select::dictionary::garbage_collect_any_dictionary;
5357
use datafusion_common::ScalarValue;
5458
use datafusion_common::config::ConfigOptions;
5559
use datafusion_common::stats::Precision;
@@ -1382,7 +1386,8 @@ impl BatchPartitioner {
13821386
partition_ranges
13831387
.into_iter()
13841388
.map(|(partition, start, len)| {
1385-
Ok((partition, reordered_batch.slice(start, len)))
1389+
let batch = reordered_batch.slice(start, len);
1390+
Ok((partition, compact_sparse_dictionaries(batch)?))
13861391
})
13871392
.collect()
13881393
};
@@ -1391,6 +1396,44 @@ impl BatchPartitioner {
13911396
}
13921397
}
13931398

1399+
/// Remove unreferenced values from dictionaries made sparse by repartitioning.
1400+
///
1401+
/// `take_arrays` reorders dictionary keys but preserves their complete values
1402+
/// arrays. Each output partition then references only a fraction of those
1403+
/// values, and coalescing the partitions would otherwise hash every retained
1404+
/// value while merging their dictionaries.
1405+
fn compact_sparse_dictionaries(batch: RecordBatch) -> Result<RecordBatch> {
1406+
let threshold = batch.num_rows().saturating_mul(4);
1407+
let mut changed = false;
1408+
let columns = batch
1409+
.columns()
1410+
.iter()
1411+
.map(|column| {
1412+
let DataType::Dictionary(_, _) = column.data_type() else {
1413+
return Ok(Arc::clone(column));
1414+
};
1415+
let dictionary = column.as_any_dictionary();
1416+
if dictionary.values().len() <= threshold {
1417+
return Ok(Arc::clone(column));
1418+
}
1419+
changed = true;
1420+
Ok(garbage_collect_any_dictionary(dictionary)?)
1421+
})
1422+
.collect::<Result<Vec<ArrayRef>>>()?;
1423+
1424+
if !changed {
1425+
return Ok(batch);
1426+
}
1427+
1428+
let mut options = RecordBatchOptions::new();
1429+
options = options.with_row_count(Some(batch.num_rows()));
1430+
Ok(RecordBatch::try_new_with_options(
1431+
batch.schema(),
1432+
columns,
1433+
&options,
1434+
)?)
1435+
}
1436+
13941437
/// Maps `N` input partitions to `M` output partitions based on a
13951438
/// [`Partitioning`] scheme.
13961439
///
@@ -2660,10 +2703,10 @@ mod tests {
26602703
};
26612704

26622705
use arrow::array::{
2663-
Array, ArrayRef, Decimal128Array, Int64Array, StringArray,
2664-
TimestampNanosecondArray, UInt32Array,
2706+
Array, ArrayRef, Decimal128Array, DictionaryArray, Int32Array, Int64Array,
2707+
StringArray, TimestampNanosecondArray, UInt32Array,
26652708
};
2666-
use arrow::datatypes::{DataType, Field, Schema, TimeUnit};
2709+
use arrow::datatypes::{DataType, Field, Int32Type, Schema, TimeUnit};
26672710
use datafusion_common::ScalarValue;
26682711
use datafusion_common::cast::{as_string_array, as_uint32_array};
26692712
use datafusion_common::config::SpillCompression;
@@ -2677,6 +2720,46 @@ mod tests {
26772720
};
26782721
use insta::assert_snapshot;
26792722

2723+
#[test]
2724+
fn compact_sparse_dictionaries_only_rebuilds_sparse_columns() -> Result<()> {
2725+
let values = Arc::new(StringArray::from_iter_values(
2726+
(0..32).map(|index| format!("value-{index}")),
2727+
));
2728+
let sparse = Arc::new(DictionaryArray::<Int32Type>::new(
2729+
Int32Array::from(vec![31, 2, 7, 2]),
2730+
values,
2731+
)) as ArrayRef;
2732+
let schema = Arc::new(Schema::new(vec![Field::new(
2733+
"dictionary",
2734+
sparse.data_type().clone(),
2735+
false,
2736+
)]));
2737+
let compacted = compact_sparse_dictionaries(RecordBatch::try_new(
2738+
Arc::clone(&schema),
2739+
vec![Arc::clone(&sparse)],
2740+
)?)?;
2741+
let dictionary = compacted.column(0).as_dictionary::<Int32Type>();
2742+
assert_eq!(dictionary.values().len(), 3);
2743+
let compacted_values = dictionary.values().as_string::<i32>();
2744+
assert_eq!(
2745+
(0..dictionary.len())
2746+
.map(|row| compacted_values.value(dictionary.key(row).unwrap()))
2747+
.collect::<Vec<_>>(),
2748+
vec!["value-31", "value-2", "value-7", "value-2"]
2749+
);
2750+
2751+
let dense_values = Arc::new(StringArray::from(vec!["a", "b", "c", "d"]));
2752+
let dense = Arc::new(DictionaryArray::<Int32Type>::new(
2753+
Int32Array::from(vec![0, 1, 2, 3]),
2754+
dense_values,
2755+
)) as ArrayRef;
2756+
let dense_batch = RecordBatch::try_new(schema, vec![Arc::clone(&dense)])?;
2757+
let dense_output = compact_sparse_dictionaries(dense_batch)?;
2758+
assert!(Arc::ptr_eq(dense_output.column(0), &dense));
2759+
2760+
Ok(())
2761+
}
2762+
26802763
#[derive(Debug)]
26812764
struct UnboundedTestPartition {
26822765
schema: SchemaRef,

0 commit comments

Comments
 (0)