Skip to content

Commit 4b4472c

Browse files
committed
perf: reuse dictionary hashes in repartition
1 parent 4ede583 commit 4b4472c

1 file changed

Lines changed: 174 additions & 11 deletions

File tree

  • datafusion/physical-plan/src/repartition

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

Lines changed: 174 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -21,8 +21,8 @@
2121
2222
use std::fmt::{Debug, Display, Formatter};
2323
use std::pin::Pin;
24-
use std::sync::Arc;
2524
use std::sync::atomic::{AtomicUsize, Ordering as AtomicOrdering};
25+
use std::sync::{Arc, Weak};
2626
use std::task::{Context, Poll};
2727
use std::vec;
2828

@@ -33,7 +33,9 @@ use super::{
3333
};
3434
use crate::coalesce::LimitedBatchCoalescer;
3535
use crate::execution_plan::{CardinalityEffect, EvaluationType, SchedulingType};
36-
use crate::hash_utils::create_hashes;
36+
#[cfg(not(feature = "force_hash_collisions"))]
37+
use crate::hash_utils::scatter_dictionary_hashes;
38+
use crate::hash_utils::{create_hashes, hash_single_array};
3739
use crate::metrics::{BaselineMetrics, SpillMetrics};
3840
use crate::projection::{ProjectionExec, all_columns, make_with_child, update_expr};
3941
use crate::sorts::streaming_merge::StreamingMergeBuilder;
@@ -46,9 +48,12 @@ use crate::{
4648
PlanProperties, ReplaceChildrenOptions, Statistics, validate_child_count,
4749
};
4850

49-
use arrow::array::{PrimitiveArray, RecordBatch, RecordBatchOptions, UInt64Array};
51+
use arrow::array::{
52+
Array, ArrayRef, PrimitiveArray, RecordBatch, RecordBatchOptions, UInt64Array,
53+
};
5054
use arrow::compute::take_arrays;
5155
use arrow::datatypes::{DataType, Schema, SchemaRef, UInt32Type};
56+
use arrow::downcast_dictionary_array;
5257
use arrow_schema::SortOptions;
5358
use datafusion_common::ScalarValue;
5459
use datafusion_common::config::ConfigOptions;
@@ -657,6 +662,7 @@ enum BatchPartitionerState {
657662
partition_reducer: StrengthReducedU64,
658663
hash_buffer: Vec<u64>,
659664
indices: Vec<Vec<u32>>,
665+
dictionary_hash_cache: Vec<DictionaryHashCache>,
660666
},
661667
RoundRobin {
662668
num_partitions: usize,
@@ -672,6 +678,40 @@ enum BatchPartitionerState {
672678
},
673679
}
674680

681+
#[derive(Default)]
682+
struct DictionaryHashCache {
683+
/// Dictionary values these hashes belong to.
684+
values: Option<Weak<dyn Array>>,
685+
/// Hash for each dictionary value.
686+
hashes: Vec<u64>,
687+
/// Number of times values were hashed.
688+
#[cfg(test)]
689+
rebuilds: usize,
690+
}
691+
692+
impl DictionaryHashCache {
693+
fn update(&mut self, values: &ArrayRef) -> Result<&[u64]> {
694+
let unchanged = self.values.as_ref().is_some_and(|cached| {
695+
std::ptr::addr_eq(cached.as_ptr(), Arc::as_ptr(values))
696+
});
697+
if !unchanged {
698+
self.hashes.clear();
699+
self.hashes.resize(values.len(), 0);
700+
create_hashes(
701+
[values],
702+
REPARTITION_RANDOM_STATE.random_state(),
703+
&mut self.hashes,
704+
)?;
705+
self.values = Some(Arc::downgrade(values));
706+
#[cfg(test)]
707+
{
708+
self.rebuilds += 1;
709+
}
710+
}
711+
Ok(&self.hashes)
712+
}
713+
}
714+
675715
/// Fixed RandomState used for hash repartitioning to ensure consistent behavior across
676716
/// executions and runs.
677717
pub const REPARTITION_RANDOM_STATE: SeededRandomState = SeededRandomState::with_seed(0);
@@ -1036,6 +1076,7 @@ impl BatchPartitioner {
10361076
partition_reducer: StrengthReducedU64::new(num_partitions as u64),
10371077
hash_buffer: vec![],
10381078
indices: vec![vec![]; num_partitions],
1079+
dictionary_hash_cache: vec![],
10391080
},
10401081
timer,
10411082
})
@@ -1250,6 +1291,7 @@ impl BatchPartitioner {
12501291
partition_reducer,
12511292
hash_buffer,
12521293
indices,
1294+
dictionary_hash_cache,
12531295
} => {
12541296
// Tracking time required for distributing indexes across output partitions
12551297
let timer = self.timer.timer();
@@ -1260,11 +1302,55 @@ impl BatchPartitioner {
12601302
hash_buffer.clear();
12611303
hash_buffer.resize(batch.num_rows(), 0);
12621304

1263-
create_hashes(
1264-
&arrays,
1265-
REPARTITION_RANDOM_STATE.random_state(),
1266-
hash_buffer,
1267-
)?;
1305+
dictionary_hash_cache
1306+
.resize_with(arrays.len(), DictionaryHashCache::default);
1307+
for (index, (array, cache)) in arrays
1308+
.iter()
1309+
.zip(dictionary_hash_cache.iter_mut())
1310+
.enumerate()
1311+
{
1312+
let array = array.as_ref();
1313+
if matches!(array.data_type(), DataType::Dictionary(_, _)) {
1314+
downcast_dictionary_array! {
1315+
array => {
1316+
let hashes = cache.update(array.values())?;
1317+
#[cfg(not(feature = "force_hash_collisions"))]
1318+
scatter_dictionary_hashes(
1319+
array,
1320+
hashes,
1321+
hash_buffer,
1322+
index != 0,
1323+
);
1324+
#[cfg(feature = "force_hash_collisions")]
1325+
let _ = hashes;
1326+
#[cfg(feature = "force_hash_collisions")]
1327+
hash_single_array(
1328+
array,
1329+
REPARTITION_RANDOM_STATE.random_state(),
1330+
hash_buffer,
1331+
index != 0,
1332+
)?;
1333+
},
1334+
_ => unreachable!(),
1335+
}
1336+
} else {
1337+
*cache = DictionaryHashCache::default();
1338+
hash_single_array(
1339+
array,
1340+
REPARTITION_RANDOM_STATE.random_state(),
1341+
hash_buffer,
1342+
index != 0,
1343+
)?;
1344+
}
1345+
}
1346+
1347+
if arrays.is_empty() {
1348+
create_hashes(
1349+
arrays,
1350+
REPARTITION_RANDOM_STATE.random_state(),
1351+
hash_buffer,
1352+
)?;
1353+
}
12681354

12691355
for v in indices.iter_mut() {
12701356
v.clear();
@@ -2660,10 +2746,10 @@ mod tests {
26602746
};
26612747

26622748
use arrow::array::{
2663-
Array, ArrayRef, Decimal128Array, Int64Array, StringArray,
2664-
TimestampNanosecondArray, UInt32Array,
2749+
Array, ArrayRef, Decimal128Array, DictionaryArray, Int32Array, Int64Array,
2750+
StringArray, TimestampNanosecondArray, UInt32Array,
26652751
};
2666-
use arrow::datatypes::{DataType, Field, Schema, TimeUnit};
2752+
use arrow::datatypes::{DataType, Field, Int32Type, Schema, TimeUnit};
26672753
use datafusion_common::ScalarValue;
26682754
use datafusion_common::cast::{as_string_array, as_uint32_array};
26692755
use datafusion_common::config::SpillCompression;
@@ -2918,6 +3004,83 @@ mod tests {
29183004
Ok(())
29193005
}
29203006

3007+
#[test]
3008+
fn dictionary_hash_cache_rebuilds_when_values_change() -> Result<()> {
3009+
let first: ArrayRef =
3010+
Arc::new(StringArray::from(vec![Some("a"), None, Some("ccc")]));
3011+
let second: ArrayRef =
3012+
Arc::new(StringArray::from(vec![Some("x"), None, Some("zzz")]));
3013+
let mut cache = DictionaryHashCache::default();
3014+
3015+
let first_hashes = cache.update(&first)?.to_vec();
3016+
cache.update(&first)?;
3017+
assert_eq!(cache.rebuilds, 1);
3018+
3019+
let second_hashes = cache.update(&second)?.to_vec();
3020+
assert_eq!(cache.rebuilds, 2);
3021+
assert_ne!(first_hashes, second_hashes);
3022+
Ok(())
3023+
}
3024+
3025+
#[tokio::test]
3026+
async fn hash_partition_reuses_dictionary_values_across_batches() -> Result<()> {
3027+
let dictionary_type =
3028+
DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8));
3029+
let schema = Arc::new(Schema::new(vec![
3030+
Field::new("key", dictionary_type, true),
3031+
Field::new("id", DataType::UInt32, false),
3032+
]));
3033+
let values: ArrayRef = Arc::new(StringArray::from(vec!["a", "bb", "ccc"]));
3034+
let make_batch = |keys: Vec<Option<i32>>, ids: Vec<u32>| -> Result<RecordBatch> {
3035+
let dictionary: ArrayRef = Arc::new(DictionaryArray::<Int32Type>::try_new(
3036+
Int32Array::from(keys),
3037+
Arc::clone(&values),
3038+
)?);
3039+
Ok(RecordBatch::try_new(
3040+
Arc::clone(&schema),
3041+
vec![dictionary, Arc::new(UInt32Array::from(ids))],
3042+
)?)
3043+
};
3044+
let batches = vec![
3045+
make_batch(vec![Some(0), Some(1), Some(2), None], vec![0, 1, 2, 3])?,
3046+
make_batch(vec![Some(2), Some(0), None, Some(1)], vec![4, 5, 6, 7])?,
3047+
];
3048+
3049+
let mut expected = vec![Vec::new(); 4];
3050+
for batch in &batches {
3051+
let mut hashes = vec![0; batch.num_rows()];
3052+
create_hashes(
3053+
[batch.column(0)],
3054+
REPARTITION_RANDOM_STATE.random_state(),
3055+
&mut hashes,
3056+
)?;
3057+
let ids = as_uint32_array(batch.column(1))?;
3058+
for (hash, id) in hashes.into_iter().zip(ids.values().iter().copied()) {
3059+
expected[(hash % 4) as usize].push(id);
3060+
}
3061+
}
3062+
3063+
let output = repartition(
3064+
&schema,
3065+
vec![batches],
3066+
Partitioning::Hash(vec![col("key", &schema)?], 4),
3067+
)
3068+
.await?;
3069+
let mut actual = vec![Vec::new(); 4];
3070+
for (partition, batches) in output.iter().enumerate() {
3071+
for batch in batches {
3072+
actual[partition]
3073+
.extend(as_uint32_array(batch.column(1))?.values().iter().copied());
3074+
}
3075+
}
3076+
for values in expected.iter_mut().chain(actual.iter_mut()) {
3077+
values.sort_unstable();
3078+
}
3079+
3080+
assert_eq!(actual, expected);
3081+
Ok(())
3082+
}
3083+
29213084
#[tokio::test]
29223085
async fn many_to_many_range_partition() -> Result<()> {
29233086
let schema = test_schema(false);

0 commit comments

Comments
 (0)