diff --git a/crates/krishiv-exec/src/join/grace.rs b/crates/krishiv-exec/src/join/grace.rs index d24d49514..1be7e01c5 100644 --- a/crates/krishiv-exec/src/join/grace.rs +++ b/crates/krishiv-exec/src/join/grace.rs @@ -630,8 +630,8 @@ fn reverse_blocking( let arena = RowArena::concat(schema, std::slice::from_ref(&chunk))?; let rows = u32::try_from(arena.len()).map_err(|_| internal("build chunk too large".to_owned()))?; - // Payloads borrowed from this chunk, evaluated before it is dropped. - let mut pairs: Vec<(u32, &[u8])> = Vec::new(); + // (probe row, build row of this chunk), evaluated before it is dropped. + let mut pairs: Vec<(u32, u32)> = Vec::new(); for r in 0..rows { let Some(candidates) = index.get(&arena.hash(r)) else { continue; @@ -644,16 +644,32 @@ fn reverse_blocking( continue; } if shared.filter.is_some() { - pairs.push((p, arena.payload(r))); + pairs.push((p, r)); } else { matched[p as usize] = true; } } if pairs.len() >= PAIR_BATCH { - mark_pairs(shared, &probe, &mut pairs, &mut matched, partition)?; + mark_pairs( + shared, + hashing, + &probe, + &arena, + &mut pairs, + &mut matched, + partition, + )?; } } - mark_pairs(shared, &probe, &mut pairs, &mut matched, partition) + mark_pairs( + shared, + hashing, + &probe, + &arena, + &mut pairs, + &mut matched, + partition, + ) }; for_each( build, @@ -679,12 +695,16 @@ fn reverse_blocking( Ok(Reversed::Joined(Some(out))) } -/// Applies the residual condition to candidate pairs (probe row, build -/// payload) and marks the probe rows of those that hold. +/// Applies the residual condition to candidate pairs (probe row, row of the +/// build chunk `arena`) and marks the probe rows of those that hold. The +/// build columns are decoded in the arena's payload format: payload rows +/// in the row layout, arrow rows otherwise. fn mark_pairs( shared: &Shared, + hashing: &Hashing, probe: &Batch, - pairs: &mut Vec<(u32, &[u8])>, + arena: &RowArena, + pairs: &mut Vec<(u32, u32)>, matched: &mut [bool], partition: u32, ) -> Result<()> { @@ -699,7 +719,8 @@ fn mark_pairs( } let idx: Vec = pairs.iter().map(|(p, _)| *p).collect(); let taken = take_batch(probe, &UInt32Array::from(idx.clone()))?; - let build_cols = shared.decode(pairs.iter().map(|(_, b)| *b))?; + let build: Vec<(usize, u32)> = pairs.iter().map(|&(_, r)| (0, r)).collect(); + let build_cols = shared.build_columns(hashing.rows(), &build, |_| Some(arena))?; let m = idx.len(); let mut columns: Vec = taken.data().columns().to_vec(); columns.extend(build_cols); @@ -1711,6 +1732,131 @@ mod tests { assert_eq!(manager.used_bytes(), 0); } + /// A reversed semi join with a residual condition decodes the build + /// columns in the arena's payload format: payload rows when + /// `krishiv.exec.join.rowLayout` is on (it read the arrow-row payload, + /// empty in that layout, and panicked decoding it), arrow rows + /// otherwise. Both give the same rows. + #[tokio::test] + async fn a_reversal_with_a_condition_decodes_the_build_rows_of_either_layout() { + use arrow_array::BooleanArray; + use arrow_array::types::Int64Type; + + let schema = Arc::new(Schema::new(vec![ + Field::new("k", DataType::Int64, false), + Field::new("v", DataType::Int64, false), + ])); + let batch = |k: Vec, v: Vec| { + Batch::new( + RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int64Array::from(k)), Arc::new(Int64Array::from(v))], + ) + .unwrap(), + ) + }; + // probe.k = build.k AND build.v > probe.v + let filter: Arc = Arc::new(crate::FnPredicate(|b: &RecordBatch| { + let probe = b.column(1).as_primitive::(); + let build = b.column(3).as_primitive::(); + Ok(probe + .iter() + .zip(build.iter()) + .map(|(p, b)| Some(b? > p?)) + .collect::()) + })); + for row_layout in [false, true] { + let root = TreePool::new(&MemoryConfig::default()); + let config = ExecConfig { + join_row_layout: row_layout, + ..ExecConfig::default() + }; + let query = QueryContext::new( + QueryId(1), + WorkloadClass::Interactive, + config.clone(), + root.child("q", None), + ); + let exec = TokioExecutor::new(&ExecutorConfig::default()).unwrap(); + let ctx = TaskContext::new(query, exec, SpillOwner::new(), OperatorId(1), "probe"); + let join = super::super::HashJoin::try_new(super::super::JoinSpec { + join_type: super::super::JoinType::LeftSemi, + probe_schema: Arc::clone(&schema), + build_schema: Arc::clone(&schema), + probe_keys: vec![0], + build_keys: vec![0], + null_safe: vec![false], + filter: Some(Arc::clone(&filter)), + single_match: false, + spill: None, + }) + .unwrap(); + let shared = &join.0; + let hashing = shared.hashing(&config); + assert_eq!(hashing.rows(), row_layout); + + // Probe rows k = v = 0..100; build rows k = 0..100 with v = k + 1 + // for odd k and v = k for even k: odd probe rows match. + let probe = batch((0..100).collect(), (0..100).collect()); + let build = batch((0..100).collect(), (0..100).map(|k| k + (k % 2)).collect()); + let key_cols = vec![Arc::clone(build.data().column(0))]; + let keys = shared.keys.encode(&key_cols, &hashing).unwrap(); + let columns = build.data().columns(); + let payload = match (shared.row_payload(hashing.rows()), &shared.payload) { + (Some(layout), _) => { + super::super::build::Payload::Rows(layout.encode(columns).unwrap()) + } + (None, Some((conv, _))) => { + super::super::build::Payload::Arrow(conv.convert_columns(columns).unwrap()) + } + (None, None) => panic!("the condition reads build columns"), + }; + let all: Vec = (0..100).collect(); + let chunk = super::super::build::arena_chunk( + shared, + &hashing, + &keys, + Some(&payload), + build.weights(), + &all, + ) + .unwrap(); + let mut build_part = vec![Part::with_memory(ctx.consumer().reservation())]; + reserve_chunks(&mut build_part, &[(0, chunk.clone())]).unwrap(); + build_part[0].chunks.push(chunk); + let mut probe_part = vec![Part::with_memory(ctx.consumer().reservation())]; + reserve_chunks(&mut probe_part, &[(0, probe.clone())]).unwrap(); + probe_part[0].chunks.push(probe); + let task = Task { + depth: 0, + build: build_part.remove(0), + probe: probe_part.remove(0), + }; + let handle = Arc::clone(ctx.consumer().consumer()); + let Ok(Reversed::Joined(Some(out))) = reverse_blocking( + shared, + &hashing, + task, + &handle, + ctx.query(), + ctx.partition(), + ) else { + panic!("the reversal joins (row layout {row_layout})"); + }; + let mut got = out + .data() + .column(0) + .as_primitive::() + .values() + .to_vec(); + got.sort_unstable(); + let want: Vec = (0..100).filter(|k| k % 2 == 1).collect(); + assert_eq!(got, want, "row layout {row_layout}"); + drop(out); + assert_eq!(root.reserved(), 0); + } + } + /// A spill whose rows and the pool's free memory cannot fund its /// writer's scratch is declined: nothing is written or lost, no credit /// is left behind, and the query does not fail.