Skip to content

Commit 0fa56a1

Browse files
viiryaIsaac
viirya
and
Isaac
committed
fix: pair without the null-aware scope maps when memory is short
The scope maps only speed up pairing NULL-valued rows with the rows in their scope. When the memory pool declines one, pair with every row and compare the scope keys instead, as the join did before the tuple elements became scope keys, rather than fail. Co-authored-by: Isaac <[email protected]>
1 parent 44471f1 commit 0fa56a1

2 files changed

Lines changed: 188 additions & 81 deletions

File tree

‎datafusion/physical-plan/src/joins/hash_join/exec.rs‎

Lines changed: 102 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -10579,57 +10579,82 @@ mod tests {
1057910579
Ok(())
1058010580
}
1058110581

10582+
/// Runs a null-aware `LeftAnti` join with the value key `id` and the
10583+
/// scope key `grp`, over 10,000 build rows `(2 * i, i % 100)`, with a
10584+
/// NULL `id` for each multiple of `null_build_every`, and the probe rows
10585+
/// `(probe_ids[i], i)`, in a memory pool of `memory_limit` bytes or an
10586+
/// unbounded one. Returns the output rows and the `build_mem_used` metric.
10587+
async fn run_scope_key_anti_join(
10588+
null_build_every: Option<i32>,
10589+
probe_ids: Vec<Option<i32>>,
10590+
memory_limit: Option<usize>,
10591+
) -> Result<(usize, usize)> {
10592+
let build_rows = 10_000;
10593+
let build_id = |i: i32| match null_build_every {
10594+
Some(every) if i % every == 0 => None,
10595+
_ => Some(i * 2),
10596+
};
10597+
let left = build_table_two_cols(
10598+
("id", &(0..build_rows).map(build_id).collect()),
10599+
("grp", &(0..build_rows).map(|i| Some(i % 100)).collect()),
10600+
);
10601+
let probe_grps = (0..probe_ids.len() as i32).map(Some).collect();
10602+
let right = build_table_two_cols(("id", &probe_ids), ("grp", &probe_grps));
10603+
let on = vec![
10604+
(
10605+
Arc::new(Column::new_with_schema("id", &left.schema())?) as _,
10606+
Arc::new(Column::new_with_schema("id", &right.schema())?) as _,
10607+
),
10608+
(
10609+
Arc::new(Column::new_with_schema("grp", &left.schema())?) as _,
10610+
Arc::new(Column::new_with_schema("grp", &right.schema())?) as _,
10611+
),
10612+
];
10613+
let join = HashJoinExec::try_new(
10614+
left,
10615+
right,
10616+
on,
10617+
None,
10618+
&JoinType::LeftAnti,
10619+
None,
10620+
PartitionMode::CollectLeft,
10621+
NullEquality::NullEqualsNothing,
10622+
true,
10623+
)?;
10624+
let mut task_ctx = prepare_task_ctx(8192, false);
10625+
if let Some(limit) = memory_limit {
10626+
let runtime = RuntimeEnvBuilder::new()
10627+
.with_memory_limit(limit, 1.0)
10628+
.build_arc()?;
10629+
task_ctx = Arc::new(
10630+
TaskContext::default()
10631+
.with_session_config(task_ctx.session_config().clone())
10632+
.with_runtime(runtime),
10633+
);
10634+
}
10635+
let batches = common::collect(join.execute(0, task_ctx)?).await?;
10636+
let rows = batches.iter().map(|b| b.num_rows()).sum();
10637+
let build_mem_used = join
10638+
.metrics()
10639+
.unwrap()
10640+
.sum_by_name("build_mem_used")
10641+
.unwrap()
10642+
.as_usize();
10643+
Ok((rows, build_mem_used))
10644+
}
10645+
1058210646
/// The hash table over the build rows' correlation scope keys is only
1058310647
/// looked up for probe rows with a NULL value key, so a probe side
1058410648
/// without one never builds it, and the result is the same either way.
1058510649
#[tokio::test]
1058610650
async fn null_aware_scope_map_built_only_for_null_probe_values() -> Result<()> {
10587-
async fn run(probe_ids: Vec<Option<i32>>) -> Result<(usize, usize)> {
10588-
let build_rows = 10_000;
10589-
let left = build_table_two_cols(
10590-
("id", &(0..build_rows).map(|i| Some(i * 2)).collect()),
10591-
("grp", &(0..build_rows).map(|i| Some(i % 100)).collect()),
10592-
);
10593-
let probe_grps = (0..probe_ids.len() as i32).map(Some).collect();
10594-
let right = build_table_two_cols(("id", &probe_ids), ("grp", &probe_grps));
10595-
let on = vec![
10596-
(
10597-
Arc::new(Column::new_with_schema("id", &left.schema())?) as _,
10598-
Arc::new(Column::new_with_schema("id", &right.schema())?) as _,
10599-
),
10600-
(
10601-
Arc::new(Column::new_with_schema("grp", &left.schema())?) as _,
10602-
Arc::new(Column::new_with_schema("grp", &right.schema())?) as _,
10603-
),
10604-
];
10605-
let join = HashJoinExec::try_new(
10606-
left,
10607-
right,
10608-
on,
10609-
None,
10610-
&JoinType::LeftAnti,
10611-
None,
10612-
PartitionMode::CollectLeft,
10613-
NullEquality::NullEqualsNothing,
10614-
true,
10615-
)?;
10616-
let batches =
10617-
common::collect(join.execute(0, prepare_task_ctx(8192, false))?).await?;
10618-
let rows = batches.iter().map(|b| b.num_rows()).sum();
10619-
let build_mem_used = join
10620-
.metrics()
10621-
.unwrap()
10622-
.sum_by_name("build_mem_used")
10623-
.unwrap()
10624-
.as_usize();
10625-
Ok((rows, build_mem_used))
10626-
}
10627-
1062810651
// `(0, 0)` matches the build row `(0, 0)`; `(1, 1)` matches nothing.
10629-
let (rows, without_null) = run(vec![Some(0), Some(1)]).await?;
10652+
let (rows, without_null) =
10653+
run_scope_key_anti_join(None, vec![Some(0), Some(1)], None).await?;
1063010654
assert_eq!(rows, 9_999);
1063110655
// `(NULL, 1)` makes the 100 build rows with `grp = 1` UNKNOWN.
10632-
let (rows, with_null) = run(vec![Some(0), None]).await?;
10656+
let (rows, with_null) =
10657+
run_scope_key_anti_join(None, vec![Some(0), None], None).await?;
1063310658
assert_eq!(rows, 9_899);
1063410659
// The scope map holds a bucket and a chain entry for every build row.
1063510660
assert!(
@@ -10639,6 +10664,40 @@ mod tests {
1063910664
Ok(())
1064010665
}
1064110666

10667+
/// A memory pool that holds the build side but not the scope map: the
10668+
/// join pairs the NULL-valued probe rows with every build row instead,
10669+
/// and gives the same result.
10670+
#[tokio::test]
10671+
async fn null_aware_scope_map_falls_back_when_memory_is_short() -> Result<()> {
10672+
let (_, without_null) =
10673+
run_scope_key_anti_join(None, vec![Some(0), Some(1)], None).await?;
10674+
let limit = without_null + 10_000 * size_of::<u32>() / 2;
10675+
let (rows, build_mem_used) =
10676+
run_scope_key_anti_join(None, vec![Some(0), None], Some(limit)).await?;
10677+
assert_eq!(rows, 9_899);
10678+
assert_eq!(build_mem_used, without_null);
10679+
Ok(())
10680+
}
10681+
10682+
/// The same for the hash table over the scope keys of the NULL-valued
10683+
/// build rows, which is built with the build side.
10684+
#[tokio::test]
10685+
async fn null_aware_null_build_scope_map_falls_back_when_memory_is_short()
10686+
-> Result<()> {
10687+
// The build rows `(NULL, 0)` for each multiple of 100 are UNKNOWN for
10688+
// the probe row `(0, 0)`, and `(0, 0)` matches the build row `(0, 0)`
10689+
// no more, as its `id` is NULL.
10690+
let probe_ids = vec![Some(0), Some(1)];
10691+
let (rows, unbounded) =
10692+
run_scope_key_anti_join(Some(100), probe_ids.clone(), None).await?;
10693+
assert_eq!(rows, 9_900);
10694+
let (rows, build_mem_used) =
10695+
run_scope_key_anti_join(Some(100), probe_ids, Some(unbounded - 1)).await?;
10696+
assert_eq!(rows, 9_900);
10697+
assert!(build_mem_used < unbounded, "{build_mem_used} < {unbounded}");
10698+
Ok(())
10699+
}
10700+
1064210701
/// `left.z > right.z` over the second column of two-column tables: the
1064310702
/// non-equality correlation of
1064410703
/// `id NOT IN (SELECT r.id FROM r WHERE r.z < l.z)`.

‎datafusion/physical-plan/src/joins/hash_join/not_in.rs‎

Lines changed: 86 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,9 @@ use arrow::buffer::{BooleanBuffer, NullBuffer};
4545
use arrow::compute::{filter, filter_record_batch, take};
4646
use arrow::util::bit_util;
4747
use datafusion_common::hash_utils::{RandomState, create_hashes};
48-
use datafusion_common::{JoinSide, JoinType, NullEquality, Result, plan_err};
48+
use datafusion_common::{
49+
DataFusionError, JoinSide, JoinType, NullEquality, Result, plan_err,
50+
};
4951
use datafusion_execution::memory_pool::MemoryReservation;
5052
use datafusion_expr::{ColumnarValue, Operator};
5153
use datafusion_physical_expr::PhysicalExprRef;
@@ -629,24 +631,31 @@ fn correlated_row_state(
629631
reservation.try_grow(retained_size)?;
630632
metrics.build_mem_used.add(retained_size);
631633

632-
let null_scope_map = if scope_keys.is_empty() {
633-
None
634+
let scope = if scope_keys.is_empty() {
635+
NullValueScope::Unscoped
634636
} else {
635-
let scope_values = normalized[value_keys..]
636-
.iter()
637-
.map(|values| Ok(filter(values.as_ref(), &null_mask)?))
638-
.collect::<Result<Vec<_>>>()?;
639637
let null_rows = build_indices.len();
640-
let mut map = new_join_hashmap(null_rows, reservation, metrics)?;
641-
let mut hashes_buffer = vec![0; null_rows];
642-
create_hashes(&scope_values, random_state, &mut hashes_buffer)?;
643-
map.update_from_iter(Box::new(hashes_buffer.iter().enumerate().rev()), 0);
644-
Some(map)
638+
match new_join_hashmap(null_rows, reservation, metrics) {
639+
Ok(mut map) => {
640+
let scope_values = normalized[value_keys..]
641+
.iter()
642+
.map(|values| Ok(filter(values.as_ref(), &null_mask)?))
643+
.collect::<Result<Vec<_>>>()?;
644+
let mut hashes_buffer = vec![0; null_rows];
645+
create_hashes(&scope_values, random_state, &mut hashes_buffer)?;
646+
map.update_from_iter(Box::new(hashes_buffer.iter().enumerate().rev()), 0);
647+
NullValueScope::Map(map)
648+
}
649+
// The map only speeds up the pairing, so the join goes on without
650+
// it rather than fail.
651+
Err(DataFusionError::ResourcesExhausted(_)) => NullValueScope::Unindexed,
652+
Err(e) => return Err(e),
653+
}
645654
};
646655
Ok((
647656
scope_map,
648657
Some(NullValueBuildRows {
649-
scope_map: null_scope_map,
658+
scope,
650659
build_indices,
651660
}),
652661
))
@@ -658,8 +667,15 @@ fn correlated_row_state(
658667
/// Only the probe rows with a NULL value key look it up, so it is built on
659668
/// the first such probe row rather than with the build side: a probe side
660669
/// without NULLs never pays for a second hash table over the build side.
670+
///
671+
/// The map only speeds up the pairing of those probe rows with the build
672+
/// rows in their scope. So when the memory pool cannot hold it, the join
673+
/// pairs them with every build row instead, and compares their scope keys.
674+
/// That is the pairing of a join whose tuple elements are all value keys,
675+
/// which needs no map either.
661676
pub(super) struct BuildScopeMap {
662-
map: OnceLock<Box<dyn JoinHashMapType>>,
677+
/// `None` when the memory pool could not hold the map.
678+
map: OnceLock<Option<Box<dyn JoinHashMapType>>>,
663679
/// Charged for the map once it is built. Its lock also makes the probe
664680
/// partitions build the map only once.
665681
reservation: Mutex<MemoryReservation>,
@@ -669,23 +685,30 @@ pub(super) struct BuildScopeMap {
669685
impl BuildScopeMap {
670686
/// Returns the map, building it on first use from `scope_values`, the
671687
/// build side's scope keys (`JoinLeftData::values()[value_keys..]`).
688+
/// Returns `None` when the memory pool cannot hold it.
672689
fn get_or_try_init(
673690
&self,
674691
scope_values: &[ArrayRef],
675692
random_state: &RandomState,
676-
) -> Result<&dyn JoinHashMapType> {
693+
) -> Result<Option<&dyn JoinHashMapType>> {
677694
if let Some(map) = self.map.get() {
678-
return Ok(map.as_ref());
695+
return Ok(map.as_deref());
679696
}
680697
let mut reservation = self.reservation.lock();
681698
if let Some(map) = self.map.get() {
682-
return Ok(map.as_ref());
699+
return Ok(map.as_deref());
683700
}
684701
// Scope-only NULL marking uses a HashMap (the primary join map may
685702
// use ArrayMap for full-key matches, but scope keys have arbitrary
686703
// shape).
687704
let num_rows = scope_values[0].len();
688-
let mut map = new_join_hashmap(num_rows, &mut reservation, &self.metrics)?;
705+
let mut map = match new_join_hashmap(num_rows, &mut reservation, &self.metrics) {
706+
Ok(map) => map,
707+
Err(DataFusionError::ResourcesExhausted(_)) => {
708+
return Ok(self.map.get_or_init(|| None).as_deref());
709+
}
710+
Err(e) => return Err(e),
711+
};
689712
let mut hashes_buffer = vec![0; num_rows];
690713
create_hashes(scope_values, random_state, &mut hashes_buffer)?;
691714
map.extend_zero(num_rows);
@@ -696,7 +719,7 @@ impl BuildScopeMap {
696719
.enumerate()
697720
.filter(|(i, _)| valid.as_ref().is_none_or(|valid| valid.is_valid(*i)));
698721
map.update_from_iter(Box::new(hashes.rev()), 0);
699-
Ok(self.map.get_or_init(|| map).as_ref())
722+
Ok(self.map.get_or_init(|| Some(map)).as_deref())
700723
}
701724
}
702725

@@ -710,16 +733,26 @@ impl BuildScopeMap {
710733
/// proportional to the number of NULLs instead of enumerating every scope
711734
/// match of every probe row.
712735
pub(super) struct NullValueBuildRows {
713-
/// Hash table keyed by the correlation scope values of the NULL-valued
714-
/// build rows. Stored positions index into `build_indices`, not the full
715-
/// build batch. `None` when the join has no correlation scope keys, so
716-
/// every probe row is in scope.
717-
scope_map: Option<Box<dyn JoinHashMapType>>,
718-
/// Maps positions in `scope_map` back to row indices in the full build
719-
/// batch.
736+
scope: NullValueScope,
737+
/// The NULL-valued build rows. The positions stored in a
738+
/// [`NullValueScope::Map`] index into it.
720739
build_indices: UInt64Array,
721740
}
722741

742+
/// How the NULL-valued build rows find the probe rows in their correlation
743+
/// scope.
744+
enum NullValueScope {
745+
/// The join has no correlation scope keys, so every probe row is in scope.
746+
Unscoped,
747+
/// Hash table keyed by the correlation scope values of the NULL-valued
748+
/// build rows. Stored positions index into
749+
/// [`NullValueBuildRows::build_indices`], not the full build batch.
750+
Map(Box<dyn JoinHashMapType>),
751+
/// The memory pool could not hold the map. Every probe row is paired, and
752+
/// the scope keys of each pair are compared (see [`BuildScopeMap`]).
753+
Unindexed,
754+
}
755+
723756
/// Returns the rows where at least one of `value_keys` is NULL, or `None` when
724757
/// there is no such row.
725758
fn null_value_key_mask(value_keys: &[ArrayRef]) -> Option<BooleanArray> {
@@ -892,8 +925,8 @@ fn mark_null_candidates_for_probe_batch(
892925
// correlation scope.
893926
if let Some(null_rows) = null_value_build_rows {
894927
let probe_rows = (0..probe_batch.num_rows() as u32).filter(in_scope);
895-
match &null_rows.scope_map {
896-
Some(scope_map) => {
928+
match &null_rows.scope {
929+
NullValueScope::Map(scope_map) => {
897930
let mut pairs =
898931
UnmarkedPairs::new(marks, batch_size, &mut mark_scope_matches);
899932
// The map indexes only the NULL-valued build rows; translate
@@ -907,14 +940,23 @@ fn mark_null_candidates_for_probe_batch(
907940
)?;
908941
pairs.finish()?;
909942
}
910-
None => {
943+
NullValueScope::Unscoped => {
911944
let mut pairs = UnmarkedPairs::new(marks, batch_size, &mark);
912945
pairs.cross_product(
913946
null_rows.build_indices.values().iter().copied(),
914947
probe_rows,
915948
)?;
916949
pairs.finish()?;
917950
}
951+
NullValueScope::Unindexed => {
952+
let mut pairs =
953+
UnmarkedPairs::new(marks, batch_size, &mut mark_scope_matches);
954+
pairs.cross_product(
955+
null_rows.build_indices.values().iter().copied(),
956+
probe_rows,
957+
)?;
958+
pairs.finish()?;
959+
}
918960
}
919961
}
920962

@@ -933,17 +975,23 @@ fn mark_null_candidates_for_probe_batch(
933975
// so the scope map need not be built.
934976
Some(_) if null_probe_rows.peek().is_none() => {}
935977
Some(scope_map) => {
936-
let scope_map =
937-
scope_map.get_or_try_init(build_scope_values, random_state)?;
938978
let mut pairs =
939979
UnmarkedPairs::new(marks, batch_size, &mut mark_scope_matches);
940-
pairs.scope_matches(
941-
scope_map,
942-
|build_row| build_row,
943-
null_probe_rows,
944-
hashes_buffer,
945-
lookup,
946-
)?;
980+
match scope_map.get_or_try_init(build_scope_values, random_state)? {
981+
Some(scope_map) => pairs.scope_matches(
982+
scope_map,
983+
|build_row| build_row,
984+
null_probe_rows,
985+
hashes_buffer,
986+
lookup,
987+
)?,
988+
// Without the map, every build row is paired, and the
989+
// scope keys of each pair are compared.
990+
None => pairs.cross_product(
991+
0..left_data.batch().num_rows() as u64,
992+
null_probe_rows,
993+
)?,
994+
}
947995
pairs.finish()?;
948996
}
949997
None => {

0 commit comments

Comments
 (0)