exec: decode build rows of either layout in the reversed grace join

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<u32> = 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<ArrayRef> = 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<i64>, v: Vec<i64>| {
+            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<dyn crate::Predicate> = Arc::new(crate::FnPredicate(|b: &RecordBatch| {
+            let probe = b.column(1).as_primitive::<Int64Type>();
+            let build = b.column(3).as_primitive::<Int64Type>();
+            Ok(probe
+                .iter()
+                .zip(build.iter())
+                .map(|(p, b)| Some(b? > p?))
+                .collect::<BooleanArray>())
+        }));
+        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<usize> = (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::<Int64Type>()
+                .values()
+                .to_vec();
+            got.sort_unstable();
+            let want: Vec<i64> = (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.