1use std::ffi::c_void;
4use std::sync::{atomic::Ordering, Arc};
5
6use crate::{
7 cuda_graph::{CapturedCudaGraph, CsmCudaGraphKey, CudaGraphNodeKind},
8 AsKernelParam, DeviceSlice, LaunchAsync, LaunchConfig,
9};
10use xlog_core::{Result, ScalarType, Schema, XlogError};
11
12use super::{
13 dedup_kernels, filter_kernels, ilp_kernels, join_kernels, pack_kernels, scan_kernels,
14 set_ops_kernels, sort_kernels, CsmCudaGraphEntry, CsmCudaGraphNodes, HashTableU64,
15 JoinHashTableV2, JoinIndexV2, JoinType, PackedKeyData, RadixSortScratch, DEDUP_MODULE,
16 DEFAULT_JOIN_MAX_OUTPUT, FILTER_MODULE, ILP_MODULE, JOIN_MODULE, NESTED_LOOP_TOTAL_THRESHOLD,
17 PACK_MODULE, SCAN_MODULE, SET_OPS_MODULE, SORT_MODULE,
18};
19use crate::device_runtime::{StreamId, XlogDeviceRuntime};
20use crate::launch::{LaunchEnqueueError, LaunchRecorder};
21use crate::memory::{CudaColumn, DeviceRead, TrackedCudaSlice};
22use crate::CudaBuffer;
23
24const XLOG_TY_U32: u8 = 0;
29const XLOG_TY_U64: u8 = 1;
30const XLOG_TY_I32: u8 = 2;
31const XLOG_TY_I64: u8 = 3;
32const XLOG_TY_F32: u8 = 4;
33const XLOG_TY_F64: u8 = 5;
34const XLOG_TY_BOOL: u8 = 6;
35const XLOG_TY_SYMBOL: u8 = 7;
36const SMALL_FULL_ROW_SORT_MAX_ROWS: usize = 1024;
37const UNION_MANY_CHUNK_BYTES: usize = 1 << 30;
43
44#[inline]
45fn scalar_type_code_dedup(ty: ScalarType) -> u8 {
46 match ty {
47 ScalarType::U32 => XLOG_TY_U32,
48 ScalarType::U64 => XLOG_TY_U64,
49 ScalarType::I32 => XLOG_TY_I32,
50 ScalarType::I64 => XLOG_TY_I64,
51 ScalarType::F32 => XLOG_TY_F32,
52 ScalarType::F64 => XLOG_TY_F64,
53 ScalarType::Bool => XLOG_TY_BOOL,
54 ScalarType::Symbol => XLOG_TY_SYMBOL,
55 }
56}
57
58impl super::CudaKernelProvider {
59 pub fn hash_join(
77 &self,
78 left: &CudaBuffer,
79 right: &CudaBuffer,
80 left_keys: &[usize],
81 right_keys: &[usize],
82 ) -> Result<CudaBuffer> {
83 self.hash_join_with_limit(left, right, left_keys, right_keys, None)
84 }
85
86 pub fn hash_join_with_limit(
105 &self,
106 left: &CudaBuffer,
107 right: &CudaBuffer,
108 left_keys: &[usize],
109 right_keys: &[usize],
110 max_output: Option<usize>,
111 ) -> Result<CudaBuffer> {
112 let max_output_limit = max_output.unwrap_or(DEFAULT_JOIN_MAX_OUTPUT);
113
114 if left_keys.is_empty() || right_keys.is_empty() {
116 return Err(XlogError::Kernel(
117 "Join requires at least one key column".to_string(),
118 ));
119 }
120 if left_keys.len() != right_keys.len() {
121 return Err(XlogError::Kernel(
122 "Left and right key columns must have same length".to_string(),
123 ));
124 }
125 for (&left_idx, &right_idx) in left_keys.iter().zip(right_keys.iter()) {
126 if left_idx >= left.arity() {
127 return Err(XlogError::Kernel(format!(
128 "Left key column index {} out of bounds (arity {})",
129 left_idx,
130 left.arity()
131 )));
132 }
133 if right_idx >= right.arity() {
134 return Err(XlogError::Kernel(format!(
135 "Right key column index {} out of bounds (arity {})",
136 right_idx,
137 right.arity()
138 )));
139 }
140 }
141
142 let right_key_set: std::collections::HashSet<usize> = right_keys.iter().copied().collect();
144 let mut result_columns_schema = left.schema().columns.clone();
145 let mut result_sort_labels = left.schema().sort_labels().to_vec();
146 for (idx, col) in right.schema().columns.iter().enumerate() {
147 if !right_key_set.contains(&idx) {
148 result_columns_schema.push(col.clone());
149 result_sort_labels.push(
150 right
151 .schema()
152 .column_sort_label(idx)
153 .unwrap_or(&col.0)
154 .to_string(),
155 );
156 }
157 }
158 let result_schema = Schema::new(result_columns_schema)
159 .with_sort_labels(result_sort_labels)
160 .expect("natural join sort labels match result schema arity");
161
162 if left.is_empty() || right.is_empty() {
164 return self.create_empty_buffer(result_schema);
165 }
166
167 let combined = self.hash_join_v2_with_limit(
169 left,
170 right,
171 left_keys,
172 right_keys,
173 JoinType::Inner,
174 Some(max_output_limit),
175 )?;
176
177 if combined.is_empty() {
178 return self.create_empty_buffer(result_schema);
179 }
180
181 let left_arity = left.arity();
182 let right_arity = right.arity();
183
184 let CudaBuffer {
185 columns: combined_columns,
186 row_cap,
187 d_num_rows,
188 ..
189 } = combined;
190
191 if combined_columns.len() != left_arity + right_arity {
192 return Err(XlogError::Kernel(format!(
193 "Join internal error: expected {} columns, got {}",
194 left_arity + right_arity,
195 combined_columns.len()
196 )));
197 }
198
199 let mut output_columns = Vec::with_capacity(result_schema.arity());
200 let mut it = combined_columns.into_iter();
201
202 for _ in 0..left_arity {
204 let col = it.next().ok_or_else(|| {
205 XlogError::Kernel("Join internal error: missing left columns".to_string())
206 })?;
207 output_columns.push(col);
208 }
209
210 for (right_col_idx, col) in it.enumerate() {
212 if !right_key_set.contains(&right_col_idx) {
213 output_columns.push(col);
214 }
215 }
216
217 Ok(CudaBuffer::from_columns(
218 output_columns,
219 row_cap,
220 d_num_rows,
221 result_schema,
222 ))
223 }
224 pub fn dedup(&self, input: &CudaBuffer, key_cols: &[usize]) -> Result<CudaBuffer> {
238 if !input.canonical_full_row_set_certified() {
239 self.validated_logical_row_count(input)?;
240 }
241 if input.is_empty() {
242 return self.create_empty_buffer(input.schema().clone());
243 }
244
245 if key_cols.is_empty() {
246 if input.arity() == 0 {
247 let rows = self.device_row_count(input)?;
250 if rows == 0 {
251 return self.create_empty_buffer(input.schema().clone());
252 }
253 let mut result = self.buffer_from_columns(Vec::new(), 1, input.schema().clone())?;
254 result.certify_canonical_full_row_set();
255 return Ok(result);
256 }
257 return Err(XlogError::Kernel(
258 "Dedup requires at least one key column".to_string(),
259 ));
260 }
261
262 if Self::is_full_row_key(key_cols, input.arity()) && input.arity() > 1 {
263 let mut result = self.dedup_full_row_deterministic(input)?;
264 result.certify_canonical_full_row_set();
265 return Ok(result);
266 }
267
268 let sorted = self.sort(input, key_cols)?;
269 let mut result = self.dedup_sorted(&sorted, key_cols)?;
270 if Self::is_full_row_key(key_cols, input.arity()) {
271 result.certify_canonical_full_row_set();
272 }
273 Ok(result)
274 }
275
276 pub fn dedup_sorted(&self, input: &CudaBuffer, key_cols: &[usize]) -> Result<CudaBuffer> {
288 if !input.canonical_full_row_set_certified() {
289 self.validated_logical_row_count(input)?;
290 }
291 if input.is_empty() {
292 return self.create_empty_buffer(input.schema().clone());
293 }
294
295 if key_cols.is_empty() {
296 if input.arity() == 0 {
297 let rows = self.device_row_count(input)?;
298 if rows == 0 {
299 return self.create_empty_buffer(input.schema().clone());
300 }
301 return self.buffer_from_columns(Vec::new(), 1, input.schema().clone());
302 }
303 return Err(XlogError::Kernel(
304 "Dedup requires at least one key column".to_string(),
305 ));
306 }
307
308 if Self::is_full_row_key(key_cols, input.arity()) && input.arity() > 1 {
309 return self.dedup_full_row_deterministic(input);
310 }
311
312 if input.num_rows() <= 1 {
313 return self.clone_buffer(input);
314 }
315
316 if input.num_rows() > u32::MAX as u64 {
317 return Err(XlogError::Kernel(format!(
318 "Dedup supports at most {} rows, got {}",
319 u32::MAX,
320 input.num_rows()
321 )));
322 }
323
324 let scalar_type_code = scalar_type_code_dedup;
328
329 let device = self.device.inner();
330 let num_rows = input.num_rows() as u32;
331
332 let mut col_ptrs_host: Vec<u64> = Vec::with_capacity(key_cols.len());
333 let mut col_sizes_host: Vec<u32> = Vec::with_capacity(key_cols.len());
334 let mut col_types_host: Vec<u8> = Vec::with_capacity(key_cols.len());
335
336 for &key_col in key_cols {
337 if key_col >= input.arity() {
338 return Err(XlogError::Kernel(format!(
339 "Key column {} out of bounds (arity {})",
340 key_col,
341 input.arity()
342 )));
343 }
344
345 let col = input
346 .column(key_col)
347 .ok_or_else(|| XlogError::Kernel(format!("Key column {} not found", key_col)))?;
348 let ty = input.schema().column_type(key_col).ok_or_else(|| {
349 XlogError::Kernel(format!("Key column {} type not found in schema", key_col))
350 })?;
351
352 let elem_size = ty.size_bytes();
353 let expected_bytes = (num_rows as usize) * elem_size;
354 if col.num_bytes() != expected_bytes {
355 return Err(XlogError::Kernel(format!(
356 "Key column {} has {} bytes but expected {} (num_rows={}, elem_size={})",
357 key_col,
358 col.num_bytes(),
359 expected_bytes,
360 num_rows,
361 elem_size
362 )));
363 }
364
365 let ptr = *col.device_ptr();
366 col_ptrs_host.push(ptr);
367 col_sizes_host.push(elem_size as u32);
368 col_types_host.push(scalar_type_code(ty));
369 }
370
371 let num_key_cols = key_cols.len() as u32;
372 let mut d_col_ptrs = self.memory.alloc::<u64>(key_cols.len())?;
373 let mut d_col_sizes = self.memory.alloc::<u32>(key_cols.len())?;
374 let mut d_col_types = self.memory.alloc::<u8>(key_cols.len())?;
375
376 self.htod_launch_metadata_sync_copy_into(&col_ptrs_host, &mut d_col_ptrs)
377 .map_err(|e| XlogError::Kernel(format!("Failed to upload key column ptrs: {}", e)))?;
378 self.htod_launch_metadata_sync_copy_into(&col_sizes_host, &mut d_col_sizes)
379 .map_err(|e| XlogError::Kernel(format!("Failed to upload key column sizes: {}", e)))?;
380 self.htod_launch_metadata_sync_copy_into(&col_types_host, &mut d_col_types)
381 .map_err(|e| XlogError::Kernel(format!("Failed to upload key column types: {}", e)))?;
382
383 let block_size = 256u32;
384 let num_blocks = num_rows.div_ceil(block_size);
385 let config = LaunchConfig {
386 grid_dim: (num_blocks, 1, 1),
387 block_dim: (block_size, 1, 1),
388 shared_mem_bytes: 0,
389 };
390
391 let d_unique_mask = self.memory.alloc::<u8>(num_rows as usize)?;
392 let d_prefix_sum = self.memory.alloc::<u32>(num_rows as usize)?;
393 let mut d_block_sums = self.memory.alloc::<u32>(num_blocks as usize)?;
394
395 let mark_and_scan_fn = device
396 .get_func(DEDUP_MODULE, dedup_kernels::MARK_UNIQUE_AND_SCAN_COLUMNAR)
397 .ok_or_else(|| {
398 XlogError::Kernel("mark_unique_and_scan_columnar kernel not found".to_string())
399 })?;
400
401 unsafe {
406 mark_and_scan_fn.clone().launch(
407 config,
408 (
409 &d_col_ptrs,
410 &d_col_sizes,
411 &d_col_types,
412 num_key_cols,
413 input.num_rows_device(),
414 num_rows,
415 &d_unique_mask,
416 &d_prefix_sum,
417 &d_block_sums,
418 ),
419 )
420 }
421 .map_err(|e| XlogError::Kernel(format!("mark_unique_and_scan_columnar failed: {}", e)))?;
422 self.device.synchronize()?;
423
424 if num_blocks > 1 {
425 self.multiblock_scan_u32_inplace(&mut d_block_sums, num_blocks)?;
426
427 let phase3_fn = device
428 .get_func(SCAN_MODULE, scan_kernels::MULTIBLOCK_SCAN_PHASE3)
429 .ok_or_else(|| {
430 XlogError::Kernel("Failed to get multiblock_scan_phase3 kernel".to_string())
431 })?;
432
433 unsafe {
435 phase3_fn.clone().launch(
436 LaunchConfig {
437 grid_dim: (num_blocks, 1, 1),
438 block_dim: (block_size, 1, 1),
439 shared_mem_bytes: 0,
440 },
441 (&d_prefix_sum, &d_block_sums, num_rows),
442 )
443 }
444 .map_err(|e| XlogError::Kernel(format!("multiblock_scan_phase3 failed: {}", e)))?;
445 self.device.synchronize()?;
446 }
447
448 self.device.synchronize()?;
449
450 let d_out_count = self.capture_compact_count(&d_prefix_sum, &d_unique_mask, num_rows)?;
451 self.compact_buffer_by_device_mask_device_count(
452 input,
453 &d_unique_mask,
454 &d_prefix_sum,
455 d_out_count,
456 )
457 }
458 pub fn union(&self, a: &CudaBuffer, b: &CudaBuffer) -> Result<CudaBuffer> {
470 self.union_gpu(a, b)
471 }
472
473 fn concat_buffers_gpu(&self, a: &CudaBuffer, b: &CudaBuffer) -> Result<CudaBuffer> {
474 if !self.schemas_type_compatible(a.schema(), b.schema()) {
475 return Err(XlogError::Kernel(format!(
476 "Concat requires compatible schemas: {:?} vs {:?}",
477 a.schema(),
478 b.schema()
479 )));
480 }
481
482 let schema = a.schema().clone();
483 let a_rows = self.device_row_count(a)? as u64;
484 let b_rows = self.device_row_count(b)? as u64;
485
486 if a_rows == 0 && b_rows == 0 {
487 return self.create_empty_buffer(schema);
488 }
489 if a_rows == 0 {
490 return self.clone_buffer(b);
491 }
492 if b_rows == 0 {
493 return self.clone_buffer(a);
494 }
495
496 let total_rows = a_rows + b_rows;
497 if total_rows > u32::MAX as u64 {
498 return Err(XlogError::Kernel(format!(
499 "Concat supports at most {} rows, got {}",
500 u32::MAX,
501 total_rows
502 )));
503 }
504
505 let device = self.device.inner();
506 let concat_fn = device
507 .get_func(SET_OPS_MODULE, set_ops_kernels::CONCAT_BYTES)
508 .ok_or_else(|| XlogError::Kernel("concat_bytes kernel not found".to_string()))?;
509
510 let block_size = 256u32;
511
512 let a_rows = usize::try_from(a_rows)
513 .map_err(|_| XlogError::Kernel(format!("Concat: a has too many rows: {}", a_rows)))?;
514 let b_rows = usize::try_from(b_rows)
515 .map_err(|_| XlogError::Kernel(format!("Concat: b has too many rows: {}", b_rows)))?;
516
517 let mut result_columns = Vec::with_capacity(schema.arity());
518 for col_idx in 0..schema.arity() {
519 let elem_size = schema
520 .column_type(col_idx)
521 .map(|t| t.size_bytes())
522 .unwrap_or(4);
523
524 let a_bytes = a_rows
525 .checked_mul(elem_size)
526 .ok_or_else(|| XlogError::Kernel("Concat: a_bytes overflow".to_string()))?;
527 let b_bytes = b_rows
528 .checked_mul(elem_size)
529 .ok_or_else(|| XlogError::Kernel("Concat: b_bytes overflow".to_string()))?;
530 let total_bytes = a_bytes
531 .checked_add(b_bytes)
532 .ok_or_else(|| XlogError::Kernel("Concat: total_bytes overflow".to_string()))?;
533
534 let a_bytes_u32 = u32::try_from(a_bytes).map_err(|_| {
535 XlogError::Kernel(format!("Concat: a_bytes too large: {}", a_bytes))
536 })?;
537 let b_bytes_u32 = u32::try_from(b_bytes).map_err(|_| {
538 XlogError::Kernel(format!("Concat: b_bytes too large: {}", b_bytes))
539 })?;
540 let total_bytes_u32 = u32::try_from(total_bytes).map_err(|_| {
541 XlogError::Kernel(format!("Concat: total_bytes too large: {}", total_bytes))
542 })?;
543
544 let a_col = a
545 .column(col_idx)
546 .ok_or_else(|| XlogError::Kernel(format!("A column {} not found", col_idx)))?;
547 let b_col = b
548 .column(col_idx)
549 .ok_or_else(|| XlogError::Kernel(format!("B column {} not found", col_idx)))?;
550
551 let mut out_col = self.memory.alloc::<u8>(total_bytes)?;
552
553 if total_bytes_u32 > 0 {
554 let grid_size = total_bytes_u32.div_ceil(block_size);
555 let config = LaunchConfig {
556 grid_dim: (grid_size, 1, 1),
557 block_dim: (block_size, 1, 1),
558 shared_mem_bytes: 0,
559 };
560
561 unsafe {
563 concat_fn.clone().launch(
564 config,
565 (a_col, a_bytes_u32, b_col, b_bytes_u32, &mut out_col),
566 )
567 }
568 .map_err(|e| XlogError::Kernel(format!("concat_bytes failed: {}", e)))?;
569 }
570
571 result_columns.push(out_col.into());
572 }
573
574 self.device.synchronize()?;
575
576 self.buffer_from_columns(result_columns, total_rows, schema)
577 }
578 pub fn diff(&self, a: &CudaBuffer, b: &CudaBuffer) -> Result<CudaBuffer> {
593 let num_a = self.device_row_count(a)?;
594 let num_b = self.device_row_count(b)?;
595 if num_a > u32::MAX as usize || num_b > u32::MAX as usize {
596 return Err(XlogError::Kernel(format!(
597 "Diff supports at most {} rows per side (a={}, b={})",
598 u32::MAX,
599 num_a,
600 num_b
601 )));
602 }
603
604 if num_a == 0 {
606 return self.create_empty_buffer(a.schema().clone());
607 }
608 if num_b == 0 {
609 return self.clone_buffer(a);
610 }
611
612 if !self.schemas_type_compatible(a.schema(), b.schema()) {
614 return Err(XlogError::Kernel(format!(
615 "Diff requires compatible schemas: {:?} vs {:?}",
616 a.schema(),
617 b.schema()
618 )));
619 }
620
621 if a.arity() == 0 {
623 return Err(XlogError::Kernel(
624 "Diff requires at least one column".to_string(),
625 ));
626 }
627
628 let num_b = num_b as u32;
629 let num_a = num_a as u32;
630
631 let hash_table_size = (num_b as usize * 2).max(1024) as u32;
633 let hash_table_alloc_size = (hash_table_size * 3) as usize;
634 let mut hash_table = self.memory.alloc::<u32>(hash_table_alloc_size)?;
635 let mut next_ptrs = self.memory.alloc::<u32>(num_b as usize)?;
636
637 let init_val = 0xFFFFFFFFu32;
639 self.device
640 .inner()
641 .htod_sync_copy_into(&vec![init_val; hash_table_alloc_size], &mut hash_table)
642 .map_err(|e| XlogError::Kernel(format!("Failed to init hash table: {}", e)))?;
643 self.device
644 .inner()
645 .htod_sync_copy_into(&vec![init_val; num_b as usize], &mut next_ptrs)
646 .map_err(|e| XlogError::Kernel(format!("Failed to init next pointers: {}", e)))?;
647
648 let build_func = self
650 .device
651 .inner()
652 .get_func(JOIN_MODULE, join_kernels::HASH_JOIN_BUILD)
653 .ok_or_else(|| XlogError::Kernel("hash_join_build kernel not found".to_string()))?;
654
655 let b_key_col = b
656 .column(0)
657 .ok_or_else(|| XlogError::Kernel("B key column not found".to_string()))?;
658 let b_keys_view = self.column_as_u32_view(b_key_col, num_b as usize)?;
659
660 let block_size = 256u32;
661 let build_grid = num_b.div_ceil(block_size);
662 let build_config = LaunchConfig {
663 grid_dim: (build_grid, 1, 1),
664 block_dim: (block_size, 1, 1),
665 shared_mem_bytes: 0,
666 };
667
668 unsafe {
670 build_func
671 .clone()
672 .launch(
673 build_config,
674 (
675 &b_keys_view,
676 &b_keys_view, num_b,
678 &hash_table,
679 &next_ptrs,
680 hash_table_size,
681 ),
682 )
683 .map_err(|e| XlogError::Kernel(format!("Build kernel failed: {}", e)))?;
684 }
685
686 self.device.synchronize()?;
688
689 let a_key_col = a
691 .column(0)
692 .ok_or_else(|| XlogError::Kernel("A key column not found".to_string()))?;
693
694 let mut a_keys_host = vec![0u8; (num_a as usize) * 4];
696 self.dtoh_sync_copy_into_tracked(a_key_col, &mut a_keys_host)
697 .map_err(|e| XlogError::Kernel(format!("Failed to read a keys: {}", e)))?;
698
699 let mut b_keys_host = vec![0u8; (num_b as usize) * 4];
700 self.dtoh_sync_copy_into_tracked(b_key_col, &mut b_keys_host)
701 .map_err(|e| XlogError::Kernel(format!("Failed to read b keys: {}", e)))?;
702
703 let b_keys_set: std::collections::HashSet<u32> = b_keys_host
705 .as_chunks::<4>()
706 .0
707 .iter()
708 .map(|chunk| u32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]))
709 .collect();
710
711 let diff_indices: Vec<usize> = a_keys_host
713 .as_chunks::<4>()
714 .0
715 .iter()
716 .enumerate()
717 .map(|(i, chunk)| {
718 (
719 i,
720 u32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]),
721 )
722 })
723 .filter(|(_, k)| !b_keys_set.contains(k))
724 .map(|(i, _)| i)
725 .collect();
726
727 let diff_count = diff_indices.len() as u64;
728
729 if diff_count == 0 {
730 return self.create_empty_buffer(a.schema().clone());
731 }
732
733 let schema = a.schema().clone();
735 let mut result_columns = Vec::with_capacity(schema.arity());
736
737 for col_idx in 0..schema.arity() {
738 let col_type_size = schema
739 .column_type(col_idx)
740 .map(|t| t.size_bytes())
741 .unwrap_or(4);
742 let result_bytes = (diff_count as usize) * col_type_size;
743
744 if let Some(a_col) = a.column(col_idx) {
745 let a_col_bytes = (num_a as usize) * col_type_size;
747 let mut a_col_host = vec![0u8; a_col_bytes];
748 self.dtoh_sync_copy_into_tracked(a_col, &mut a_col_host)
749 .map_err(|e| XlogError::Kernel(format!("Failed to read column: {}", e)))?;
750
751 let mut result_host = Vec::with_capacity(result_bytes);
753 for &idx in &diff_indices {
754 let start = idx * col_type_size;
755 let end = start + col_type_size;
756 result_host.extend_from_slice(&a_col_host[start..end]);
757 }
758
759 let mut result_col = self.memory.alloc::<u8>(result_bytes)?;
761 self.device
762 .inner()
763 .htod_sync_copy_into(&result_host, &mut result_col)
764 .map_err(|e| XlogError::Kernel(format!("Failed to upload result: {}", e)))?;
765
766 result_columns.push(result_col.into());
767 }
768 }
769
770 self.buffer_from_columns(result_columns, diff_count, schema)
771 }
772 fn ensure_column_bytes_kernel_indexable(&self, input: &CudaBuffer) -> Result<()> {
777 let rows = self.device_row_count(input)?;
778 for col_idx in 0..input.arity() {
779 let elem_size = input
780 .schema()
781 .column_type(col_idx)
782 .map(|t| t.size_bytes())
783 .unwrap_or(4);
784 let col_bytes = rows
785 .checked_mul(elem_size)
786 .ok_or_else(|| XlogError::Kernel("Sort: column byte size overflow".to_string()))?;
787 if u32::try_from(col_bytes).is_err() {
788 return Err(XlogError::Kernel(format!(
789 "Sort supports at most {} bytes per column, got {} (column {})",
790 u32::MAX,
791 col_bytes,
792 col_idx
793 )));
794 }
795 }
796 Ok(())
797 }
798
799 fn union_many_chunk_bytes() -> usize {
803 std::env::var("XLOG_UNION_CHUNK_BYTES")
804 .ok()
805 .and_then(|value| value.parse::<usize>().ok())
806 .filter(|&value| value > 0)
807 .unwrap_or(UNION_MANY_CHUNK_BYTES)
808 }
809
810 pub fn union_gpu(&self, a: &CudaBuffer, b: &CudaBuffer) -> Result<CudaBuffer> {
826 self.union_many_gpu(&[a, b])
827 }
828
829 pub fn union_many_gpu(&self, inputs: &[&CudaBuffer]) -> Result<CudaBuffer> {
863 let first = inputs.first().ok_or_else(|| {
864 XlogError::Kernel("union_many_gpu requires at least one input".to_string())
865 })?;
866 for other in &inputs[1..] {
868 if !self.schemas_type_compatible(first.schema(), other.schema()) {
869 return Err(XlogError::Kernel(format!(
870 "Union requires compatible schemas: {:?} vs {:?}",
871 first.schema(),
872 other.schema()
873 )));
874 }
875 }
876 for input in inputs {
877 if !input.canonical_full_row_set_certified() {
878 self.validated_logical_row_count(input)?;
879 }
880 }
881
882 let schema = first.schema().clone();
883
884 let possible_non_empty: Vec<&CudaBuffer> = inputs
885 .iter()
886 .copied()
887 .filter(|input| input.cached_row_count() != Some(0))
888 .collect();
889 if possible_non_empty.is_empty() {
890 return self.create_empty_buffer(schema);
891 }
892 if schema.arity() > 0
896 && possible_non_empty.len() == 1
897 && possible_non_empty[0].canonical_full_row_set_certified()
898 {
899 return self.clone_buffer(possible_non_empty[0]);
900 }
901
902 let mut non_empty: Vec<(&CudaBuffer, usize)> = Vec::with_capacity(possible_non_empty.len());
903 for input in possible_non_empty {
904 let rows = match input.cached_row_count() {
905 Some(rows) => rows as usize,
906 None => self.device_row_count(input)?,
907 };
908 if rows > 0 {
909 non_empty.push((input, rows));
910 }
911 }
912
913 if schema.arity() == 0 {
914 if non_empty.is_empty() {
916 return self.create_empty_buffer(schema);
917 }
918 return self.buffer_from_columns(Vec::new(), 1, schema);
919 }
920
921 if non_empty.is_empty() {
922 return self.create_empty_buffer(schema);
923 }
924
925 let key_cols: Vec<usize> = (0..schema.arity()).collect();
927 if non_empty.len() == 1 {
928 let input = non_empty[0].0;
929 if input.canonical_full_row_set_certified() {
930 return self.clone_buffer(input);
931 }
932 let mut result = self.dedup(input, &key_cols)?;
933 result.certify_canonical_full_row_set();
934 return Ok(result);
935 }
936
937 let row_bytes: usize = (0..schema.arity())
938 .map(|c| schema.column_type(c).map(|t| t.size_bytes()).unwrap_or(4))
939 .sum::<usize>()
940 .max(1);
941 let budget_rows = (Self::union_many_chunk_bytes() / row_bytes).max(1);
942
943 let mut acc: Option<CudaBuffer> = None;
948 let mut idx = 0usize;
949 while idx < non_empty.len() {
950 let mut chunk: Vec<&CudaBuffer> = Vec::new();
951 let mut chunk_rows = 0usize;
952 if let Some(acc_buf) = acc.as_ref() {
953 chunk_rows = self.device_row_count(acc_buf)?;
954 chunk.push(acc_buf);
955 }
956 let mut taken = 0usize;
959 while idx < non_empty.len() {
960 let (input, rows) = non_empty[idx];
961 if taken > 0 && chunk_rows.saturating_add(rows) > budget_rows {
962 break;
963 }
964 chunk.push(input);
965 chunk_rows = chunk_rows.saturating_add(rows);
966 idx += 1;
967 taken += 1;
968 }
969 acc = Some(self.union_chunk_gpu(&chunk, chunk_rows, &key_cols)?);
970 }
971 let mut result = acc.expect("at least one non-empty input was folded");
972 result.certify_canonical_full_row_set();
973 Ok(result)
974 }
975
976 fn union_chunk_gpu(
980 &self,
981 inputs: &[&CudaBuffer],
982 chunk_rows: usize,
983 key_cols: &[usize],
984 ) -> Result<CudaBuffer> {
985 if inputs.len() == 1 {
986 return self.dedup(inputs[0], key_cols);
987 }
988 let concat = if inputs.len() == 2 {
989 self.concat_buffers_gpu(inputs[0], inputs[1])?
990 } else {
991 self.concat_many_buffers_gpu(inputs, chunk_rows)?
992 };
993 if inputs[0].schema().arity() > 1 {
994 return self.dedup_full_row_deterministic(&concat);
999 }
1000 let sorted = self.sort(&concat, key_cols)?;
1001 self.dedup_sorted(&sorted, key_cols)
1002 }
1003
1004 fn concat_many_buffers_gpu(
1021 &self,
1022 inputs: &[&CudaBuffer],
1023 total_rows: usize,
1024 ) -> Result<CudaBuffer> {
1025 let schema = inputs[0].schema().clone();
1026 if total_rows > u32::MAX as usize {
1027 return Err(XlogError::Kernel(format!(
1028 "Concat supports at most {} rows, got {}",
1029 u32::MAX,
1030 total_rows
1031 )));
1032 }
1033
1034 let mut input_rows = Vec::with_capacity(inputs.len());
1035 for input in inputs {
1036 input_rows.push(self.device_row_count(input)?);
1037 }
1038
1039 let device = self.device.inner();
1040 let mut result_columns = Vec::with_capacity(schema.arity());
1041 let enqueue_result = (|| -> Result<()> {
1042 for col_idx in 0..schema.arity() {
1043 let elem_size = schema
1044 .column_type(col_idx)
1045 .map(|t| t.size_bytes())
1046 .unwrap_or(4);
1047 let total_bytes = total_rows
1048 .checked_mul(elem_size)
1049 .ok_or_else(|| XlogError::Kernel("Concat: total_bytes overflow".to_string()))?;
1050 u32::try_from(total_bytes).map_err(|_| {
1051 XlogError::Kernel(format!("Concat: total_bytes too large: {}", total_bytes))
1052 })?;
1053
1054 let mut out_col = self.memory.alloc::<u8>(total_bytes)?;
1055 let mut offset = 0usize;
1056 for (input, &rows) in inputs.iter().zip(&input_rows) {
1057 let col_bytes = rows.checked_mul(elem_size).ok_or_else(|| {
1058 XlogError::Kernel("Concat: col_bytes overflow".to_string())
1059 })?;
1060 u32::try_from(col_bytes).map_err(|_| {
1061 XlogError::Kernel(format!("Concat: col_bytes too large: {}", col_bytes))
1062 })?;
1063 let col = input.column(col_idx).ok_or_else(|| {
1064 XlogError::Kernel(format!("Concat: column {} not found", col_idx))
1065 })?;
1066 let src = self.column_bytes_view(col, col_bytes)?;
1067 let mut dst = out_col.slice_mut(offset..offset + col_bytes);
1068 device.dtod_copy_async(&src, &mut dst).map_err(|e| {
1069 XlogError::Kernel(format!("Concat: failed to copy column: {}", e))
1070 })?;
1071 offset += col_bytes;
1072 }
1073
1074 result_columns.push(out_col.into());
1075 }
1076 Ok(())
1077 })();
1078
1079 self.device.synchronize()?;
1083 enqueue_result?;
1084
1085 self.buffer_from_columns(result_columns, total_rows as u64, schema)
1086 }
1087
1088 pub fn diff_gpu(&self, a: &CudaBuffer, b: &CudaBuffer) -> Result<CudaBuffer> {
1106 let num_a = self.device_row_count(a)?;
1107 let num_b = self.device_row_count(b)?;
1108 if num_a > u32::MAX as usize || num_b > u32::MAX as usize {
1109 return Err(XlogError::Kernel(format!(
1110 "Diff supports at most {} rows per side (a={}, b={})",
1111 u32::MAX,
1112 num_a,
1113 num_b
1114 )));
1115 }
1116
1117 if num_a == 0 {
1118 return self.create_empty_buffer(a.schema().clone());
1119 }
1120
1121 if !self.schemas_type_compatible(a.schema(), b.schema()) {
1123 return Err(XlogError::Kernel(format!(
1124 "Diff requires compatible schemas: {:?} vs {:?}",
1125 a.schema(),
1126 b.schema()
1127 )));
1128 }
1129
1130 if a.arity() == 0 {
1131 if num_b == 0 {
1133 return self.buffer_from_columns(Vec::new(), 1, a.schema().clone());
1134 }
1135 return self.create_empty_buffer(a.schema().clone());
1136 }
1137
1138 let col_type = a
1139 .schema()
1140 .column_type(0)
1141 .ok_or_else(|| XlogError::Kernel("No columns".to_string()))?;
1142
1143 if a.arity() == 1 && matches!(col_type, ScalarType::U32) && num_b != 0 {
1146 return self.diff_gpu_u32(a, b);
1147 }
1148
1149 self.diff_via_deterministic_set(a, b)
1150 }
1151
1152 fn diff_gpu_u32(&self, a: &CudaBuffer, b: &CudaBuffer) -> Result<CudaBuffer> {
1154 if a.arity() != 1 {
1156 return self.diff_via_deterministic_set(a, b);
1157 }
1158
1159 let sorted_a = self.sort(a, &[0])?;
1161 let deduped_a = self.dedup_sorted(&sorted_a, &[0])?;
1162
1163 let sorted_b = self.sort(b, &[0])?;
1164 let deduped_b = self.dedup_sorted(&sorted_b, &[0])?;
1165
1166 let num_a = self.device_row_count(&deduped_a)?;
1167 let num_b = self.device_row_count(&deduped_b)?;
1168 if num_a > u32::MAX as usize || num_b > u32::MAX as usize {
1169 return Err(XlogError::Kernel(format!(
1170 "Diff supports at most {} rows per side (a={}, b={})",
1171 u32::MAX,
1172 num_a,
1173 num_b
1174 )));
1175 }
1176
1177 if num_a == 0 {
1178 return self.create_empty_buffer(a.schema().clone());
1179 }
1180
1181 let num_a = num_a as u32;
1182 let num_b = num_b as u32;
1183
1184 let diff_mark_fn = self
1186 .device
1187 .inner()
1188 .get_func(SET_OPS_MODULE, set_ops_kernels::SORTED_DIFF_MARK)
1189 .ok_or_else(|| XlogError::Kernel("sorted_diff_mark kernel not found".to_string()))?;
1190
1191 let a_col = deduped_a
1193 .column(0)
1194 .ok_or_else(|| XlogError::Kernel("A column 0 not found".to_string()))?;
1195 let b_col = deduped_b
1196 .column(0)
1197 .ok_or_else(|| XlogError::Kernel("B column 0 not found".to_string()))?;
1198
1199 let a_view = self.column_as_u32_view(a_col, num_a as usize)?;
1200 let b_view = self.column_as_u32_view(b_col, num_b as usize)?;
1201
1202 let diff_mask = self.memory.alloc::<u8>(num_a as usize)?;
1204
1205 let block_size = 256u32;
1207 let grid_size = num_a.div_ceil(block_size);
1208 let config = LaunchConfig {
1209 grid_dim: (grid_size, 1, 1),
1210 block_dim: (block_size, 1, 1),
1211 shared_mem_bytes: 0,
1212 };
1213
1214 unsafe {
1218 diff_mark_fn.clone().launch(
1219 config,
1220 (
1221 &a_view,
1222 deduped_a.num_rows_device(),
1223 num_a,
1224 &b_view,
1225 deduped_b.num_rows_device(),
1226 num_b,
1227 &diff_mask,
1228 ),
1229 )
1230 }
1231 .map_err(|e| XlogError::Kernel(format!("sorted_diff_mark failed: {}", e)))?;
1232
1233 let device = self.device.inner();
1235 let num_blocks = grid_size;
1236 let d_prefix_sum = self.memory.alloc::<u32>(num_a as usize)?;
1237 let mut d_block_sums = self.memory.alloc::<u32>(num_blocks as usize)?;
1238
1239 let phase1_fn = device
1240 .get_func(SCAN_MODULE, scan_kernels::MULTIBLOCK_SCAN_PHASE1)
1241 .ok_or_else(|| {
1242 XlogError::Kernel("Failed to get multiblock_scan_phase1 kernel".to_string())
1243 })?;
1244
1245 unsafe {
1247 phase1_fn.clone().launch(
1248 LaunchConfig {
1249 grid_dim: (num_blocks, 1, 1),
1250 block_dim: (block_size, 1, 1),
1251 shared_mem_bytes: 0,
1252 },
1253 (&diff_mask, &d_prefix_sum, &d_block_sums, num_a),
1254 )
1255 }
1256 .map_err(|e| XlogError::Kernel(format!("multiblock_scan_phase1 failed: {}", e)))?;
1257
1258 if num_blocks > 1 {
1259 self.multiblock_scan_u32_inplace(&mut d_block_sums, num_blocks)?;
1260
1261 let phase3_fn = device
1262 .get_func(SCAN_MODULE, scan_kernels::MULTIBLOCK_SCAN_PHASE3)
1263 .ok_or_else(|| {
1264 XlogError::Kernel("Failed to get multiblock_scan_phase3 kernel".to_string())
1265 })?;
1266
1267 unsafe {
1269 phase3_fn.clone().launch(
1270 LaunchConfig {
1271 grid_dim: (num_blocks, 1, 1),
1272 block_dim: (block_size, 1, 1),
1273 shared_mem_bytes: 0,
1274 },
1275 (&d_prefix_sum, &d_block_sums, num_a),
1276 )
1277 }
1278 .map_err(|e| XlogError::Kernel(format!("multiblock_scan_phase3 failed: {}", e)))?;
1279 }
1280
1281 self.device.synchronize()?;
1282
1283 let d_out_count = self.capture_compact_count(&d_prefix_sum, &diff_mask, num_a)?;
1284 self.compact_buffer_by_device_mask_device_count(
1285 &deduped_a,
1286 &diff_mask,
1287 &d_prefix_sum,
1288 d_out_count,
1289 )
1290 }
1291
1292 fn diff_via_deterministic_set(&self, a: &CudaBuffer, b: &CudaBuffer) -> Result<CudaBuffer> {
1315 let deduped_a = self.dedup_full_row_deterministic(a)?;
1323 let deduped_b = self.dedup_full_row_deterministic(b)?;
1324
1325 let a_rows = self.device_row_count(&deduped_a)? as u32;
1326 let b_rows = self.device_row_count(&deduped_b)? as u32;
1327 if a_rows == 0 {
1328 return self.create_empty_buffer(a.schema().clone());
1329 }
1330 if b_rows == 0 {
1331 return Ok(deduped_a);
1332 }
1333 let arity = deduped_a.arity();
1334 if arity == 0 {
1335 return self.create_empty_buffer(a.schema().clone());
1337 }
1338
1339 let sorted_b = deduped_b;
1345
1346 let schema = deduped_a.schema().clone();
1350 let device = self.device.inner();
1351
1352 let mut a_col_ptrs: Vec<u64> = Vec::with_capacity(arity);
1353 let mut b_col_ptrs: Vec<u64> = Vec::with_capacity(arity);
1354 let mut col_sizes: Vec<u32> = Vec::with_capacity(arity);
1355 let mut col_types: Vec<u8> = Vec::with_capacity(arity);
1356 for col_idx in 0..arity {
1357 let a_col = deduped_a.column(col_idx).ok_or_else(|| {
1358 XlogError::Kernel(format!("diff_full_row: a column {} missing", col_idx))
1359 })?;
1360 let b_col = sorted_b.column(col_idx).ok_or_else(|| {
1361 XlogError::Kernel(format!("diff_full_row: b column {} missing", col_idx))
1362 })?;
1363 let ty = schema.column_type(col_idx).ok_or_else(|| {
1364 XlogError::Kernel(format!("diff_full_row: column {} type missing", col_idx))
1365 })?;
1366 a_col_ptrs.push(*a_col.device_ptr());
1367 b_col_ptrs.push(*b_col.device_ptr());
1368 col_sizes.push(ty.size_bytes() as u32);
1369 col_types.push(scalar_type_code_dedup(ty));
1370 }
1371
1372 let mut d_a_ptrs = self.memory.alloc::<u64>(arity)?;
1373 let mut d_b_ptrs = self.memory.alloc::<u64>(arity)?;
1374 let mut d_sizes = self.memory.alloc::<u32>(arity)?;
1375 let mut d_types = self.memory.alloc::<u8>(arity)?;
1376 self.htod_launch_metadata_sync_copy_into(&a_col_ptrs, &mut d_a_ptrs)
1377 .map_err(|e| XlogError::Kernel(format!("diff_full_row a ptr upload: {}", e)))?;
1378 self.htod_launch_metadata_sync_copy_into(&b_col_ptrs, &mut d_b_ptrs)
1379 .map_err(|e| XlogError::Kernel(format!("diff_full_row b ptr upload: {}", e)))?;
1380 self.htod_launch_metadata_sync_copy_into(&col_sizes, &mut d_sizes)
1381 .map_err(|e| XlogError::Kernel(format!("diff_full_row size upload: {}", e)))?;
1382 self.htod_launch_metadata_sync_copy_into(&col_types, &mut d_types)
1383 .map_err(|e| XlogError::Kernel(format!("diff_full_row type upload: {}", e)))?;
1384
1385 let block_size = 256u32;
1386 let grid = a_rows.div_ceil(block_size);
1387 let cfg = LaunchConfig {
1388 grid_dim: (grid, 1, 1),
1389 block_dim: (block_size, 1, 1),
1390 shared_mem_bytes: 0,
1391 };
1392
1393 let d_keep_mask = self.memory.alloc::<u8>(a_rows as usize)?;
1394 let diff_fn = device
1395 .get_func(DEDUP_MODULE, dedup_kernels::MARK_DIFF_FULL_ROW_TYPED_SORTED)
1396 .ok_or_else(|| {
1397 XlogError::Kernel("mark_diff_full_row_typed_sorted kernel not found".to_string())
1398 })?;
1399 unsafe {
1403 diff_fn.clone().launch(
1404 cfg,
1405 (
1406 &d_a_ptrs,
1407 &d_b_ptrs,
1408 &d_sizes,
1409 &d_types,
1410 arity as u32,
1411 deduped_a.num_rows_device(),
1412 b_rows,
1413 a_rows,
1414 &d_keep_mask,
1415 ),
1416 )
1417 }
1418 .map_err(|e| XlogError::Kernel(format!("mark_diff_full_row_typed_sorted launch: {}", e)))?;
1419 self.device.synchronize()?;
1420
1421 let (d_prefix_sum, d_out_count) =
1423 self.scan_mask_to_prefix_with_count(&d_keep_mask, a_rows)?;
1424
1425 self.compact_buffer_by_device_mask_device_count(
1426 &deduped_a,
1427 &d_keep_mask,
1428 &d_prefix_sum,
1429 d_out_count,
1430 )
1431 }
1432
1433 pub fn dedup_full_row(&self, input: &CudaBuffer) -> Result<CudaBuffer> {
1462 if !input.canonical_full_row_set_certified() {
1463 self.validated_logical_row_count(input)?;
1464 }
1465 if Self::use_recorded_dedup_env()? && input.num_rows() > 1 && input.arity() > 0 {
1469 if let Some(launch_stream) = self.recorded_op_stream_or_init() {
1470 let recorded_compatible = (0..input.arity()).all(|c| {
1471 matches!(
1472 input.schema.column_type(c),
1473 Some(ScalarType::U32) | Some(ScalarType::Symbol)
1474 )
1475 });
1476 if recorded_compatible {
1477 let mut result = self.dedup_full_row_recorded(input, launch_stream)?;
1478 result.certify_canonical_full_row_set();
1479 return Ok(result);
1480 }
1481 }
1482 }
1483 let mut result = self.dedup_full_row_deterministic(input)?;
1484 result.certify_canonical_full_row_set();
1485 Ok(result)
1486 }
1487
1488 pub fn diff_full_row(&self, a: &CudaBuffer, b: &CudaBuffer) -> Result<CudaBuffer> {
1493 self.diff_gpu(a, b)
1497 }
1498
1499 fn read_join_output_count_metadata(&self, d_count: &TrackedCudaSlice<u32>) -> Result<u32> {
1521 self.dtoh_scalar_untracked::<u32>(d_count, 0)
1528 .map_err(|e| match e {
1529 XlogError::Kernel(message) => {
1530 XlogError::Kernel(format!("Failed to read output count: {}", message))
1531 }
1532 other => XlogError::Kernel(format!("Failed to read output count: {}", other)),
1533 })
1534 }
1535
1536 fn is_full_row_key(key_cols: &[usize], arity: usize) -> bool {
1537 key_cols.len() == arity
1538 && key_cols
1539 .iter()
1540 .copied()
1541 .enumerate()
1542 .all(|(expected, actual)| expected == actual)
1543 }
1544
1545 fn dedup_full_row_deterministic(&self, input: &CudaBuffer) -> Result<CudaBuffer> {
1562 self.ensure_column_bytes_kernel_indexable(input)?;
1566 let row_count = if input.canonical_full_row_set_certified() {
1567 self.device_row_count(input)?
1568 } else {
1569 self.validated_logical_row_count(input)?
1570 };
1571 if row_count == 0 {
1572 return self.create_empty_buffer(input.schema().clone());
1573 }
1574 if row_count == 1 {
1575 return self.clone_buffer(input);
1576 }
1577 if row_count > u32::MAX as usize {
1578 return Err(XlogError::Kernel(format!(
1579 "dedup_full_row supports at most {} rows, got {}",
1580 u32::MAX,
1581 row_count
1582 )));
1583 }
1584 let arity = input.arity();
1585 if arity == 0 {
1586 return self.buffer_from_columns(Vec::new(), 1, input.schema().clone());
1588 }
1589
1590 let sorted = if Self::use_csm_cuda_graph_env()? && row_count <= SMALL_FULL_ROW_SORT_MAX_ROWS
1593 {
1594 self.small_sort_full_row_deterministic(input, row_count)?
1595 } else {
1596 let all_cols: Vec<usize> = (0..arity).collect();
1597 self.sort(input, &all_cols)?
1598 };
1599
1600 let n = self.device_row_count(&sorted)? as u32;
1602 if n <= 1 {
1603 return Ok(sorted);
1604 }
1605
1606 let device = self.device.inner();
1607 let mut col_ptrs_host: Vec<u64> = Vec::with_capacity(arity);
1608 let mut col_sizes_host: Vec<u32> = Vec::with_capacity(arity);
1609 for col_idx in 0..arity {
1610 let col = sorted
1611 .column(col_idx)
1612 .ok_or_else(|| XlogError::Kernel(format!("Sorted column {} not found", col_idx)))?;
1613 let ty = sorted.schema().column_type(col_idx).ok_or_else(|| {
1614 XlogError::Kernel(format!("Sorted column {} type missing", col_idx))
1615 })?;
1616 col_ptrs_host.push(*col.device_ptr());
1617 col_sizes_host.push(ty.size_bytes() as u32);
1618 }
1619
1620 let mut d_col_ptrs = self.memory.alloc::<u64>(arity)?;
1621 let mut d_col_sizes = self.memory.alloc::<u32>(arity)?;
1622 self.htod_launch_metadata_sync_copy_into(&col_ptrs_host, &mut d_col_ptrs)
1623 .map_err(|e| XlogError::Kernel(format!("dedup_full_row_gpu col ptr upload: {}", e)))?;
1624 self.htod_launch_metadata_sync_copy_into(&col_sizes_host, &mut d_col_sizes)
1625 .map_err(|e| XlogError::Kernel(format!("dedup_full_row_gpu col size upload: {}", e)))?;
1626
1627 let block_size = 256u32;
1628 let grid = n.div_ceil(block_size);
1629 let cfg = LaunchConfig {
1630 grid_dim: (grid, 1, 1),
1631 block_dim: (block_size, 1, 1),
1632 shared_mem_bytes: 0,
1633 };
1634
1635 let d_unique_mask = self.memory.alloc::<u8>(n as usize)?;
1636 let mark_fn = device
1637 .get_func(DEDUP_MODULE, dedup_kernels::MARK_UNIQUE_FULL_ROW_BYTEWISE)
1638 .ok_or_else(|| {
1639 XlogError::Kernel("mark_unique_full_row_bytewise kernel not found".to_string())
1640 })?;
1641
1642 unsafe {
1646 mark_fn.clone().launch(
1647 cfg,
1648 (
1649 &d_col_ptrs,
1650 &d_col_sizes,
1651 arity as u32,
1652 sorted.num_rows_device(),
1653 n,
1654 &d_unique_mask,
1655 ),
1656 )
1657 }
1658 .map_err(|e| {
1659 XlogError::Kernel(format!(
1660 "mark_unique_full_row_bytewise launch failed: {}",
1661 e
1662 ))
1663 })?;
1664 self.device.synchronize()?;
1665
1666 let (d_prefix_sum, d_out_count) = self.scan_mask_to_prefix_with_count(&d_unique_mask, n)?;
1668
1669 let mut result = self.compact_buffer_by_device_mask_device_count(
1673 &sorted,
1674 &d_unique_mask,
1675 &d_prefix_sum,
1676 d_out_count,
1677 )?;
1678 result.certify_canonical_full_row_set();
1679 Ok(result)
1680 }
1681
1682 fn small_sort_full_row_deterministic(
1683 &self,
1684 input: &CudaBuffer,
1685 row_count: usize,
1686 ) -> Result<CudaBuffer> {
1687 if row_count > SMALL_FULL_ROW_SORT_MAX_ROWS {
1688 return Err(XlogError::Kernel(format!(
1689 "small full-row sort supports at most {} rows, got {}",
1690 SMALL_FULL_ROW_SORT_MAX_ROWS, row_count
1691 )));
1692 }
1693 if row_count == 0 {
1694 return self.create_empty_buffer(input.schema().clone());
1695 }
1696 if row_count == 1 {
1697 return self.clone_buffer(input);
1698 }
1699
1700 let arity = input.arity();
1701 let device = self.device.inner();
1702 let mut col_ptrs_host: Vec<u64> = Vec::with_capacity(arity);
1703 let mut col_sizes_host: Vec<u32> = Vec::with_capacity(arity);
1704 let mut col_types_host: Vec<u8> = Vec::with_capacity(arity);
1705 for col_idx in 0..arity {
1706 let col = input.column(col_idx).ok_or_else(|| {
1707 XlogError::Kernel(format!("small full-row sort: column {} missing", col_idx))
1708 })?;
1709 let ty = input.schema().column_type(col_idx).ok_or_else(|| {
1710 XlogError::Kernel(format!(
1711 "small full-row sort: column {} type missing",
1712 col_idx
1713 ))
1714 })?;
1715 let elem_size = ty.size_bytes();
1716 let expected_bytes_u64 =
1717 input
1718 .num_rows()
1719 .checked_mul(elem_size as u64)
1720 .ok_or_else(|| {
1721 XlogError::Kernel(
1722 "small full-row sort: column byte-size overflow".to_string(),
1723 )
1724 })?;
1725 let expected_bytes = usize::try_from(expected_bytes_u64).map_err(|_| {
1726 XlogError::Kernel(format!(
1727 "small full-row sort: expected byte size {} exceeds usize::MAX",
1728 expected_bytes_u64
1729 ))
1730 })?;
1731 if col.num_bytes() != expected_bytes {
1732 return Err(XlogError::Kernel(format!(
1733 "small full-row sort: column {} has {} bytes but expected {}",
1734 col_idx,
1735 col.num_bytes(),
1736 expected_bytes
1737 )));
1738 }
1739 col_ptrs_host.push(*col.device_ptr());
1740 col_sizes_host.push(elem_size as u32);
1741 col_types_host.push(scalar_type_code_dedup(ty));
1742 }
1743
1744 let mut d_col_ptrs = self.memory.alloc::<u64>(arity)?;
1745 let mut d_col_sizes = self.memory.alloc::<u32>(arity)?;
1746 let mut d_col_types = self.memory.alloc::<u8>(arity)?;
1747 self.htod_launch_metadata_sync_copy_into(&col_ptrs_host, &mut d_col_ptrs)
1748 .map_err(|e| XlogError::Kernel(format!("small full-row sort ptr upload: {}", e)))?;
1749 self.htod_launch_metadata_sync_copy_into(&col_sizes_host, &mut d_col_sizes)
1750 .map_err(|e| XlogError::Kernel(format!("small full-row sort size upload: {}", e)))?;
1751 self.htod_launch_metadata_sync_copy_into(&col_types_host, &mut d_col_types)
1752 .map_err(|e| XlogError::Kernel(format!("small full-row sort type upload: {}", e)))?;
1753
1754 let mut d_indices = self.memory.alloc::<u32>(row_count)?;
1755 let sort_fn = device
1756 .get_func(
1757 DEDUP_MODULE,
1758 dedup_kernels::SMALL_SORT_FULL_ROW_INDICES_TYPED,
1759 )
1760 .ok_or_else(|| {
1761 XlogError::Kernel("small_sort_full_row_indices_typed kernel not found".to_string())
1762 })?;
1763 let cfg = LaunchConfig {
1764 grid_dim: (1, 1, 1),
1765 block_dim: (SMALL_FULL_ROW_SORT_MAX_ROWS as u32, 1, 1),
1766 shared_mem_bytes: 0,
1767 };
1768
1769 unsafe {
1773 sort_fn.clone().launch(
1774 cfg,
1775 (
1776 &d_col_ptrs,
1777 &d_col_sizes,
1778 &d_col_types,
1779 arity as u32,
1780 input.num_rows_device(),
1781 row_count as u32,
1782 &mut d_indices,
1783 ),
1784 )
1785 }
1786 .map_err(|e| {
1787 XlogError::Kernel(format!(
1788 "small_sort_full_row_indices_typed launch failed: {}",
1789 e
1790 ))
1791 })?;
1792 self.device.synchronize()?;
1793 self.small_full_row_sort_invocations
1794 .fetch_add(1, Ordering::Relaxed);
1795
1796 self.gather_buffer_by_indices(input, &d_indices, row_count as u32)
1797 }
1798
1799 fn scan_mask_to_prefix_with_count(
1804 &self,
1805 d_mask: &impl DeviceRead<u8>,
1806 n: u32,
1807 ) -> Result<(
1808 crate::memory::TrackedCudaSlice<u32>,
1809 crate::memory::TrackedCudaSlice<u32>,
1810 )> {
1811 let d_mask = &d_mask.device_view();
1812 let device = self.device.inner();
1813 let block_size = 256u32;
1814 let num_blocks = n.div_ceil(block_size);
1815
1816 let d_prefix_sum = self.memory.alloc::<u32>(n as usize)?;
1817 let mut d_block_sums = self.memory.alloc::<u32>(num_blocks as usize)?;
1818
1819 let phase1_fn = device
1820 .get_func(SCAN_MODULE, scan_kernels::MULTIBLOCK_SCAN_PHASE1)
1821 .ok_or_else(|| {
1822 XlogError::Kernel("Failed to get multiblock_scan_phase1 kernel".to_string())
1823 })?;
1824 unsafe {
1826 phase1_fn.clone().launch(
1827 LaunchConfig {
1828 grid_dim: (num_blocks, 1, 1),
1829 block_dim: (block_size, 1, 1),
1830 shared_mem_bytes: 0,
1831 },
1832 (d_mask, &d_prefix_sum, &d_block_sums, n),
1833 )
1834 }
1835 .map_err(|e| XlogError::Kernel(format!("multiblock_scan_phase1 failed: {}", e)))?;
1836
1837 if num_blocks > 1 {
1838 self.multiblock_scan_u32_inplace(&mut d_block_sums, num_blocks)?;
1839
1840 let phase3_fn = device
1841 .get_func(SCAN_MODULE, scan_kernels::MULTIBLOCK_SCAN_PHASE3)
1842 .ok_or_else(|| {
1843 XlogError::Kernel("Failed to get multiblock_scan_phase3 kernel".to_string())
1844 })?;
1845 unsafe {
1847 phase3_fn.clone().launch(
1848 LaunchConfig {
1849 grid_dim: (num_blocks, 1, 1),
1850 block_dim: (block_size, 1, 1),
1851 shared_mem_bytes: 0,
1852 },
1853 (&d_prefix_sum, &d_block_sums, n),
1854 )
1855 }
1856 .map_err(|e| XlogError::Kernel(format!("multiblock_scan_phase3 failed: {}", e)))?;
1857 }
1858 self.device.synchronize()?;
1859
1860 let d_out_count = self.capture_compact_count(&d_prefix_sum, d_mask, n)?;
1861 Ok((d_prefix_sum, d_out_count))
1862 }
1863
1864 pub(super) const SORT_BLOCK_SIZE: u32 = 256;
1867
1868 pub fn sort(&self, input: &CudaBuffer, key_cols: &[usize]) -> Result<CudaBuffer> {
1886 self.ensure_column_bytes_kernel_indexable(input)?;
1891
1892 if Self::use_recorded_sort_env()? && !key_cols.is_empty() && input.num_rows() > 0 {
1897 if let Some(launch_stream) = self.recorded_op_stream_or_init() {
1898 let recorded_compatible = key_cols.iter().all(|&k| {
1899 matches!(
1900 input.schema.column_type(k),
1901 Some(ScalarType::U32) | Some(ScalarType::Symbol)
1902 )
1903 });
1904 if recorded_compatible {
1905 return self.sort_recorded(input, key_cols, launch_stream);
1906 }
1907 }
1908 }
1909
1910 if input.num_rows() == 0 {
1911 return self.create_empty_buffer(input.schema.clone());
1912 }
1913
1914 if key_cols.is_empty() {
1915 return Err(XlogError::Kernel(
1916 "Sort requires at least one key column".to_string(),
1917 ));
1918 }
1919
1920 if input.num_rows() > u32::MAX as u64 {
1921 return Err(XlogError::Kernel(format!(
1922 "Sort supports at most {} rows, got {}",
1923 u32::MAX,
1924 input.num_rows()
1925 )));
1926 }
1927
1928 for &key_col in key_cols {
1929 if key_col >= input.arity() {
1930 return Err(XlogError::Kernel(format!(
1931 "Key column index {} out of bounds (arity {})",
1932 key_col,
1933 input.arity()
1934 )));
1935 }
1936 }
1937
1938 let n = input.num_rows() as u32;
1939 let d_num_rows = input.num_rows_device();
1940 let device = self.device.inner();
1941
1942 let block_size = Self::SORT_BLOCK_SIZE;
1943 let grid_size = n.div_ceil(block_size);
1944 let launch_config = LaunchConfig {
1945 grid_dim: (grid_size, 1, 1),
1946 block_dim: (block_size, 1, 1),
1947 shared_mem_bytes: 0,
1948 };
1949
1950 let init_fn = device
1952 .get_func(SORT_MODULE, sort_kernels::INIT_INDICES)
1953 .ok_or_else(|| XlogError::Kernel("init_indices kernel not found".to_string()))?;
1954
1955 let mut indices_a = self.memory.alloc::<u32>(n as usize)?;
1956 let mut indices_b = self.memory.alloc::<u32>(n as usize)?;
1957
1958 unsafe {
1960 init_fn
1961 .clone()
1962 .launch(launch_config, (&mut indices_a, d_num_rows, n))
1963 }
1964 .map_err(|e| XlogError::Kernel(format!("init_indices failed: {}", e)))?;
1965 self.device.synchronize()?;
1966
1967 let mut keys_a = self.memory.alloc::<u32>(n as usize)?;
1969 let mut keys_b = self.memory.alloc::<u32>(n as usize)?;
1970
1971 let mut d_hist = self.memory.alloc::<u32>((grid_size as usize) * 16)?;
1973 let mut d_prefix = self.memory.alloc::<u32>(16)?;
1974 let mut d_ranks = self.memory.alloc::<u32>(n as usize)?;
1975
1976 for &col_idx in key_cols.iter().rev() {
1978 let ty = input.schema.column_type(col_idx).ok_or_else(|| {
1979 XlogError::Kernel(format!("Key column {} type not found in schema", col_idx))
1980 })?;
1981
1982 let col = input
1983 .column(col_idx)
1984 .ok_or_else(|| XlogError::Kernel(format!("Key column {} not found", col_idx)))?;
1985
1986 match ty {
1987 ScalarType::U32 | ScalarType::Symbol => {
1988 let col_view = self.column_as_u32_view(col, n as usize)?;
1989 let gather_fn = device
1990 .get_func(SORT_MODULE, sort_kernels::APPLY_PERMUTATION_U32)
1991 .ok_or_else(|| {
1992 XlogError::Kernel("apply_permutation_u32 kernel not found".to_string())
1993 })?;
1994
1995 unsafe {
1997 gather_fn.clone().launch(
1998 launch_config,
1999 (&col_view, &mut keys_a, &indices_a, d_num_rows, n),
2000 )
2001 }
2002 .map_err(|e| {
2003 XlogError::Kernel(format!("apply_permutation_u32 failed: {}", e))
2004 })?;
2005
2006 self.radix_sort_u32_pairs_with_scratch(
2007 &mut keys_a,
2008 &mut keys_b,
2009 &mut indices_a,
2010 &mut indices_b,
2011 &mut d_hist,
2012 &mut d_prefix,
2013 &mut d_ranks,
2014 d_num_rows,
2015 n,
2016 )?;
2017 }
2018 ScalarType::I32 => {
2019 let col_bits = self.column_as_u32_view(col, n as usize)?;
2020 let gather_fn = device
2021 .get_func(SORT_MODULE, sort_kernels::GATHER_KEYS_I32_ORDERED_U32)
2022 .ok_or_else(|| {
2023 XlogError::Kernel(
2024 "gather_keys_i32_ordered_u32 kernel not found".to_string(),
2025 )
2026 })?;
2027
2028 unsafe {
2030 gather_fn.clone().launch(
2031 launch_config,
2032 (&col_bits, &indices_a, d_num_rows, n, &mut keys_a),
2033 )
2034 }
2035 .map_err(|e| {
2036 XlogError::Kernel(format!("gather_keys_i32_ordered_u32 failed: {}", e))
2037 })?;
2038
2039 self.radix_sort_u32_pairs_with_scratch(
2040 &mut keys_a,
2041 &mut keys_b,
2042 &mut indices_a,
2043 &mut indices_b,
2044 &mut d_hist,
2045 &mut d_prefix,
2046 &mut d_ranks,
2047 d_num_rows,
2048 n,
2049 )?;
2050 }
2051 ScalarType::F32 => {
2052 let col_bits = self.column_as_u32_view(col, n as usize)?;
2053 let gather_fn = device
2054 .get_func(SORT_MODULE, sort_kernels::GATHER_KEYS_F32_ORDERED_U32)
2055 .ok_or_else(|| {
2056 XlogError::Kernel(
2057 "gather_keys_f32_ordered_u32 kernel not found".to_string(),
2058 )
2059 })?;
2060
2061 unsafe {
2063 gather_fn.clone().launch(
2064 launch_config,
2065 (&col_bits, &indices_a, d_num_rows, n, &mut keys_a),
2066 )
2067 }
2068 .map_err(|e| {
2069 XlogError::Kernel(format!("gather_keys_f32_ordered_u32 failed: {}", e))
2070 })?;
2071
2072 self.radix_sort_u32_pairs_with_scratch(
2073 &mut keys_a,
2074 &mut keys_b,
2075 &mut indices_a,
2076 &mut indices_b,
2077 &mut d_hist,
2078 &mut d_prefix,
2079 &mut d_ranks,
2080 d_num_rows,
2081 n,
2082 )?;
2083 }
2084 ScalarType::Bool => {
2085 if col.num_bytes() < n as usize {
2086 return Err(XlogError::Kernel(format!(
2087 "Bool column {} has {} bytes but expected {}",
2088 col_idx,
2089 col.num_bytes(),
2090 n
2091 )));
2092 }
2093
2094 let gather_fn = device
2095 .get_func(SORT_MODULE, sort_kernels::GATHER_KEYS_BOOL_ORDERED_U32)
2096 .ok_or_else(|| {
2097 XlogError::Kernel(
2098 "gather_keys_bool_ordered_u32 kernel not found".to_string(),
2099 )
2100 })?;
2101
2102 unsafe {
2104 gather_fn
2105 .clone()
2106 .launch(launch_config, (col, &indices_a, d_num_rows, n, &mut keys_a))
2107 }
2108 .map_err(|e| {
2109 XlogError::Kernel(format!("gather_keys_bool_ordered_u32 failed: {}", e))
2110 })?;
2111
2112 self.radix_sort_u32_pairs_with_scratch(
2113 &mut keys_a,
2114 &mut keys_b,
2115 &mut indices_a,
2116 &mut indices_b,
2117 &mut d_hist,
2118 &mut d_prefix,
2119 &mut d_ranks,
2120 d_num_rows,
2121 n,
2122 )?;
2123 }
2124 ScalarType::U64 => {
2125 let col_bits = self.column_as_u64_view(col, n as usize)?;
2126 for &word in &[
2127 sort_kernels::GATHER_KEYS_U64_LO_U32,
2128 sort_kernels::GATHER_KEYS_U64_HI_U32,
2129 ] {
2130 let gather_fn = device.get_func(SORT_MODULE, word).ok_or_else(|| {
2131 XlogError::Kernel(format!("{} kernel not found", word))
2132 })?;
2133
2134 unsafe {
2136 gather_fn.clone().launch(
2137 launch_config,
2138 (&col_bits, &indices_a, d_num_rows, n, &mut keys_a),
2139 )
2140 }
2141 .map_err(|e| XlogError::Kernel(format!("{} failed: {}", word, e)))?;
2142
2143 self.radix_sort_u32_pairs_with_scratch(
2144 &mut keys_a,
2145 &mut keys_b,
2146 &mut indices_a,
2147 &mut indices_b,
2148 &mut d_hist,
2149 &mut d_prefix,
2150 &mut d_ranks,
2151 d_num_rows,
2152 n,
2153 )?;
2154 }
2155 }
2156 ScalarType::I64 => {
2157 let col_bits = self.column_as_u64_view(col, n as usize)?;
2158 for &word in &[
2159 sort_kernels::GATHER_KEYS_I64_LO_U32,
2160 sort_kernels::GATHER_KEYS_I64_HI_U32,
2161 ] {
2162 let gather_fn = device.get_func(SORT_MODULE, word).ok_or_else(|| {
2163 XlogError::Kernel(format!("{} kernel not found", word))
2164 })?;
2165
2166 unsafe {
2168 gather_fn.clone().launch(
2169 launch_config,
2170 (&col_bits, &indices_a, d_num_rows, n, &mut keys_a),
2171 )
2172 }
2173 .map_err(|e| XlogError::Kernel(format!("{} failed: {}", word, e)))?;
2174
2175 self.radix_sort_u32_pairs_with_scratch(
2176 &mut keys_a,
2177 &mut keys_b,
2178 &mut indices_a,
2179 &mut indices_b,
2180 &mut d_hist,
2181 &mut d_prefix,
2182 &mut d_ranks,
2183 d_num_rows,
2184 n,
2185 )?;
2186 }
2187 }
2188 ScalarType::F64 => {
2189 let col_bits = self.column_as_u64_view(col, n as usize)?;
2190 for &word in &[
2191 sort_kernels::GATHER_KEYS_F64_LO_U32,
2192 sort_kernels::GATHER_KEYS_F64_HI_U32,
2193 ] {
2194 let gather_fn = device.get_func(SORT_MODULE, word).ok_or_else(|| {
2195 XlogError::Kernel(format!("{} kernel not found", word))
2196 })?;
2197
2198 unsafe {
2200 gather_fn.clone().launch(
2201 launch_config,
2202 (&col_bits, &indices_a, d_num_rows, n, &mut keys_a),
2203 )
2204 }
2205 .map_err(|e| XlogError::Kernel(format!("{} failed: {}", word, e)))?;
2206
2207 self.radix_sort_u32_pairs_with_scratch(
2208 &mut keys_a,
2209 &mut keys_b,
2210 &mut indices_a,
2211 &mut indices_b,
2212 &mut d_hist,
2213 &mut d_prefix,
2214 &mut d_ranks,
2215 d_num_rows,
2216 n,
2217 )?;
2218 }
2219 }
2220 }
2221 }
2222
2223 self.apply_permutation_gpu(input, &indices_a)
2224 }
2225
2226 #[expect(
2227 clippy::too_many_arguments,
2228 reason = "radix sorting keeps each caller-owned ping-pong and scan scratch buffer explicit"
2229 )]
2230 fn radix_sort_u32_pairs_with_scratch(
2231 &self,
2232 keys_a: &mut crate::memory::TrackedCudaSlice<u32>,
2233 keys_b: &mut crate::memory::TrackedCudaSlice<u32>,
2234 indices_a: &mut crate::memory::TrackedCudaSlice<u32>,
2235 indices_b: &mut crate::memory::TrackedCudaSlice<u32>,
2236 hist: &mut crate::memory::TrackedCudaSlice<u32>,
2237 prefix: &mut crate::memory::TrackedCudaSlice<u32>,
2238 ranks: &mut crate::memory::TrackedCudaSlice<u32>,
2239 num_rows_device: &crate::memory::TrackedCudaSlice<u32>,
2240 row_cap: u32,
2241 ) -> Result<()> {
2242 if row_cap == 0 {
2243 return Ok(());
2244 }
2245 self.device.synchronize()?;
2246
2247 let device = self.device.inner();
2248 let block_size = Self::SORT_BLOCK_SIZE;
2249 let grid_size = row_cap.div_ceil(block_size);
2250
2251 let sort_config = LaunchConfig {
2252 grid_dim: (grid_size, 1, 1),
2253 block_dim: (block_size, 1, 1),
2254 shared_mem_bytes: 0,
2255 };
2256
2257 let histogram_fn = device
2258 .get_func(SORT_MODULE, sort_kernels::RADIX_HISTOGRAM)
2259 .ok_or_else(|| XlogError::Kernel("radix_histogram kernel not found".to_string()))?;
2260 let prefix_fn = device
2261 .get_func(SORT_MODULE, sort_kernels::COMPUTE_DIGIT_PREFIX_SUMS)
2262 .ok_or_else(|| {
2263 XlogError::Kernel("compute_digit_prefix_sums kernel not found".to_string())
2264 })?;
2265 let ranks_fn = device
2266 .get_func(SORT_MODULE, sort_kernels::COMPUTE_RANKS)
2267 .ok_or_else(|| XlogError::Kernel("compute_ranks kernel not found".to_string()))?;
2268 let scatter_fn = device
2269 .get_func(SORT_MODULE, sort_kernels::RADIX_SCATTER_STABLE)
2270 .ok_or_else(|| {
2271 XlogError::Kernel("radix_scatter_stable kernel not found".to_string())
2272 })?;
2273
2274 let prefix_config = LaunchConfig {
2275 grid_dim: (1, 1, 1),
2276 block_dim: (256, 1, 1),
2277 shared_mem_bytes: 0,
2278 };
2279
2280 let mut in_a = true;
2281 for pass in 0..8u32 {
2282 let shift = pass * 4;
2283
2284 let (keys_in, indices_in, keys_out, indices_out) = if in_a {
2285 (&*keys_a, &*indices_a, &mut *keys_b, &mut *indices_b)
2286 } else {
2287 (&*keys_b, &*indices_b, &mut *keys_a, &mut *indices_a)
2288 };
2289
2290 unsafe {
2293 histogram_fn.clone().launch(
2294 sort_config,
2295 (keys_in, num_rows_device, row_cap, &mut *hist, shift),
2296 )
2297 }
2298 .map_err(|e| XlogError::Kernel(format!("radix_histogram failed: {}", e)))?;
2299 self.device.synchronize()?;
2300
2301 unsafe {
2304 prefix_fn
2305 .clone()
2306 .launch(prefix_config, (&*hist, grid_size, &mut *prefix))
2307 }
2308 .map_err(|e| XlogError::Kernel(format!("compute_digit_prefix_sums failed: {}", e)))?;
2309 self.device.synchronize()?;
2310
2311 for digit in 0..16u32 {
2313 let start = (digit * grid_size) as usize;
2314 let end = start + (grid_size as usize);
2315 let mut digit_slice = hist.slice_mut(start..end);
2316 self.multiblock_scan_u32_view_inplace(&mut digit_slice, grid_size)?;
2317 }
2318 self.device.synchronize()?;
2319
2320 unsafe {
2323 ranks_fn.clone().launch(
2324 sort_config,
2325 (keys_in, num_rows_device, row_cap, &mut *ranks, shift),
2326 )
2327 }
2328 .map_err(|e| XlogError::Kernel(format!("compute_ranks failed: {}", e)))?;
2329 self.device.synchronize()?;
2330
2331 unsafe {
2334 scatter_fn.clone().launch(
2335 sort_config,
2336 (
2337 keys_in,
2338 indices_in,
2339 &*ranks,
2340 keys_out,
2341 indices_out,
2342 &*prefix,
2343 &*hist,
2344 num_rows_device,
2345 row_cap,
2346 shift,
2347 ),
2348 )
2349 }
2350 .map_err(|e| XlogError::Kernel(format!("radix_scatter_stable failed: {}", e)))?;
2351 self.device.synchronize()?;
2352
2353 in_a = !in_a;
2354 }
2355
2356 if !in_a {
2358 return Err(XlogError::Kernel(
2359 "Unexpected radix-sort buffer parity (expected even number of passes)".to_string(),
2360 ));
2361 }
2362
2363 Ok(())
2364 }
2365 pub fn init_indices(
2367 &self,
2368 indices: &mut crate::memory::TrackedCudaSlice<u32>,
2369 n: u32,
2370 ) -> Result<()> {
2371 if n == 0 {
2372 return Ok(());
2373 }
2374 if n as usize > indices.len() {
2375 return Err(XlogError::Kernel(format!(
2376 "init_indices: n={} exceeds indices len={}",
2377 n,
2378 indices.len()
2379 )));
2380 }
2381 let device = self.device.inner();
2382 let block_size = Self::SORT_BLOCK_SIZE;
2383 let grid_size = n.div_ceil(block_size);
2384 let config = LaunchConfig {
2385 grid_dim: (grid_size, 1, 1),
2386 block_dim: (block_size, 1, 1),
2387 shared_mem_bytes: 0,
2388 };
2389 let init_fn = device
2390 .get_func(SORT_MODULE, sort_kernels::INIT_INDICES)
2391 .ok_or_else(|| XlogError::Kernel("init_indices kernel not found".to_string()))?;
2392 let d_num_rows = self.upload_device_row_count(n)?;
2393 unsafe {
2395 init_fn
2396 .clone()
2397 .launch(config, (&mut *indices, &d_num_rows, n))
2398 }
2399 .map_err(|e| XlogError::Kernel(format!("init_indices failed: {}", e)))?;
2400 Ok(())
2401 }
2402
2403 pub fn gather_u32_by_indices(
2405 &self,
2406 input: &crate::memory::TrackedCudaSlice<u32>,
2407 indices: &crate::memory::TrackedCudaSlice<u32>,
2408 output: &mut crate::memory::TrackedCudaSlice<u32>,
2409 n: u32,
2410 ) -> Result<()> {
2411 if n == 0 {
2412 return Ok(());
2413 }
2414 if n as usize > output.len() {
2415 return Err(XlogError::Kernel(format!(
2416 "gather_u32_by_indices: n={} exceeds output len={}",
2417 n,
2418 output.len()
2419 )));
2420 }
2421 let device = self.device.inner();
2422 let block_size = Self::SORT_BLOCK_SIZE;
2423 let grid_size = n.div_ceil(block_size);
2424 let config = LaunchConfig {
2425 grid_dim: (grid_size, 1, 1),
2426 block_dim: (block_size, 1, 1),
2427 shared_mem_bytes: 0,
2428 };
2429 let gather_fn = device
2430 .get_func(SORT_MODULE, sort_kernels::APPLY_PERMUTATION_U32)
2431 .ok_or_else(|| {
2432 XlogError::Kernel("apply_permutation_u32 kernel not found".to_string())
2433 })?;
2434 let d_num_rows = self.upload_device_row_count(n)?;
2435 unsafe {
2437 gather_fn
2438 .clone()
2439 .launch(config, (input, output, indices, &d_num_rows, n))
2440 }
2441 .map_err(|e| XlogError::Kernel(format!("gather_u32_by_indices failed: {}", e)))?;
2442 Ok(())
2443 }
2444
2445 pub fn gather_u8_by_indices(
2447 &self,
2448 input: &crate::memory::TrackedCudaSlice<u8>,
2449 indices: &crate::memory::TrackedCudaSlice<u32>,
2450 output: &mut crate::memory::TrackedCudaSlice<u8>,
2451 n: u32,
2452 ) -> Result<()> {
2453 if n == 0 {
2454 return Ok(());
2455 }
2456 if n as usize > output.len() {
2457 return Err(XlogError::Kernel(format!(
2458 "gather_u8_by_indices: n={} exceeds output len={}",
2459 n,
2460 output.len()
2461 )));
2462 }
2463 let device = self.device.inner();
2464 let block_size = Self::SORT_BLOCK_SIZE;
2465 let grid_size = n.div_ceil(block_size);
2466 let config = LaunchConfig {
2467 grid_dim: (grid_size, 1, 1),
2468 block_dim: (block_size, 1, 1),
2469 shared_mem_bytes: 0,
2470 };
2471 let gather_fn = device
2472 .get_func(SORT_MODULE, sort_kernels::APPLY_PERMUTATION_BYTES)
2473 .ok_or_else(|| {
2474 XlogError::Kernel("apply_permutation_bytes kernel not found".to_string())
2475 })?;
2476 let d_num_rows = self.upload_device_row_count(n)?;
2477 unsafe {
2479 gather_fn
2480 .clone()
2481 .launch(config, (input, output, indices, &d_num_rows, n, 1u32))
2482 }
2483 .map_err(|e| XlogError::Kernel(format!("gather_u8_by_indices failed: {}", e)))?;
2484 Ok(())
2485 }
2486
2487 pub fn gather_u64_lo_by_indices(
2489 &self,
2490 input: &crate::memory::TrackedCudaSlice<u64>,
2491 indices: &crate::memory::TrackedCudaSlice<u32>,
2492 output: &mut crate::memory::TrackedCudaSlice<u32>,
2493 n: u32,
2494 ) -> Result<()> {
2495 if n == 0 {
2496 return Ok(());
2497 }
2498 let device = self.device.inner();
2499 let block_size = Self::SORT_BLOCK_SIZE;
2500 let grid_size = n.div_ceil(block_size);
2501 let config = LaunchConfig {
2502 grid_dim: (grid_size, 1, 1),
2503 block_dim: (block_size, 1, 1),
2504 shared_mem_bytes: 0,
2505 };
2506 let gather_fn = device
2507 .get_func(SORT_MODULE, sort_kernels::GATHER_KEYS_U64_LO_U32)
2508 .ok_or_else(|| XlogError::Kernel("gather_keys_u64_lo_u32 not found".to_string()))?;
2509 let d_num_rows = self.upload_device_row_count(n)?;
2510 unsafe {
2512 gather_fn
2513 .clone()
2514 .launch(config, (input, indices, &d_num_rows, n, output))
2515 }
2516 .map_err(|e| XlogError::Kernel(format!("gather_u64_lo_by_indices failed: {}", e)))?;
2517 Ok(())
2518 }
2519
2520 pub fn gather_u64_hi_by_indices(
2522 &self,
2523 input: &crate::memory::TrackedCudaSlice<u64>,
2524 indices: &crate::memory::TrackedCudaSlice<u32>,
2525 output: &mut crate::memory::TrackedCudaSlice<u32>,
2526 n: u32,
2527 ) -> Result<()> {
2528 if n == 0 {
2529 return Ok(());
2530 }
2531 let device = self.device.inner();
2532 let block_size = Self::SORT_BLOCK_SIZE;
2533 let grid_size = n.div_ceil(block_size);
2534 let config = LaunchConfig {
2535 grid_dim: (grid_size, 1, 1),
2536 block_dim: (block_size, 1, 1),
2537 shared_mem_bytes: 0,
2538 };
2539 let gather_fn = device
2540 .get_func(SORT_MODULE, sort_kernels::GATHER_KEYS_U64_HI_U32)
2541 .ok_or_else(|| XlogError::Kernel("gather_keys_u64_hi_u32 not found".to_string()))?;
2542 let d_num_rows = self.upload_device_row_count(n)?;
2543 unsafe {
2545 gather_fn
2546 .clone()
2547 .launch(config, (input, indices, &d_num_rows, n, output))
2548 }
2549 .map_err(|e| XlogError::Kernel(format!("gather_u64_hi_by_indices failed: {}", e)))?;
2550 Ok(())
2551 }
2552
2553 pub fn radix_sort_u32_pairs(
2555 &self,
2556 keys: &mut crate::memory::TrackedCudaSlice<u32>,
2557 values: &mut crate::memory::TrackedCudaSlice<u32>,
2558 n: u32,
2559 scratch: &mut RadixSortScratch,
2560 ) -> Result<()> {
2561 if n == 0 {
2562 return Ok(());
2563 }
2564 scratch.ensure_capacity(self, n)?;
2565 let d_num_rows = self.upload_device_row_count(n)?;
2566 self.radix_sort_u32_pairs_with_scratch(
2567 keys,
2568 &mut scratch.keys_b,
2569 values,
2570 &mut scratch.values_b,
2571 &mut scratch.hist,
2572 &mut scratch.prefix,
2573 &mut scratch.ranks,
2574 &d_num_rows,
2575 n,
2576 )
2577 }
2578 pub fn scan_u8_mask_device(
2580 &self,
2581 mask: &crate::memory::TrackedCudaSlice<u8>,
2582 n: u32,
2583 ) -> Result<crate::memory::TrackedCudaSlice<u32>> {
2584 if n == 0 {
2585 return self.memory.alloc::<u32>(0);
2586 }
2587 if n as usize > mask.len() {
2588 return Err(XlogError::Kernel(format!(
2589 "scan_u8_mask_device: n={} exceeds mask len={}",
2590 n,
2591 mask.len()
2592 )));
2593 }
2594 let device = self.device.inner();
2595 let block_size = 256u32;
2596 let num_blocks = n.div_ceil(block_size);
2597
2598 let mut prefix_sum = self.memory.alloc::<u32>(n as usize)?;
2599 let mut block_sums = self.memory.alloc::<u32>(num_blocks as usize)?;
2600
2601 let phase1_fn = device
2602 .get_func(SCAN_MODULE, scan_kernels::MULTIBLOCK_SCAN_PHASE1)
2603 .ok_or_else(|| {
2604 XlogError::Kernel("multiblock_scan_phase1 kernel not found".to_string())
2605 })?;
2606
2607 unsafe {
2609 phase1_fn.clone().launch(
2610 LaunchConfig {
2611 grid_dim: (num_blocks, 1, 1),
2612 block_dim: (block_size, 1, 1),
2613 shared_mem_bytes: 0,
2614 },
2615 (mask, &mut prefix_sum, &mut block_sums, n),
2616 )
2617 }
2618 .map_err(|e| XlogError::Kernel(format!("multiblock_scan_phase1 failed: {}", e)))?;
2619
2620 if num_blocks > 1 {
2621 self.multiblock_scan_u32_inplace(&mut block_sums, num_blocks)?;
2622
2623 let phase3_fn = device
2624 .get_func(SCAN_MODULE, scan_kernels::MULTIBLOCK_SCAN_PHASE3)
2625 .ok_or_else(|| {
2626 XlogError::Kernel("multiblock_scan_phase3 kernel not found".to_string())
2627 })?;
2628
2629 unsafe {
2631 phase3_fn.clone().launch(
2632 LaunchConfig {
2633 grid_dim: (num_blocks, 1, 1),
2634 block_dim: (block_size, 1, 1),
2635 shared_mem_bytes: 0,
2636 },
2637 (&mut prefix_sum, &block_sums, n),
2638 )
2639 }
2640 .map_err(|e| XlogError::Kernel(format!("multiblock_scan_phase3 failed: {}", e)))?;
2641 }
2642
2643 Ok(prefix_sum)
2644 }
2645
2646 pub fn count_mask_device(
2650 &self,
2651 mask: &crate::memory::TrackedCudaSlice<u8>,
2652 n: u32,
2653 ) -> Result<crate::memory::TrackedCudaSlice<u32>> {
2654 let mut d_count = self.memory.alloc::<u32>(1)?;
2655 self.htod_launch_metadata_sync_copy_into(&[0u32], &mut d_count)
2656 .map_err(|e| {
2657 XlogError::Kernel(format!("count_mask_device: zero init failed: {}", e))
2658 })?;
2659
2660 if n == 0 {
2661 return Ok(d_count);
2662 }
2663
2664 let device = self.device.inner();
2665 let block_size = 256u32;
2666 let grid_size = n.div_ceil(block_size);
2667
2668 let count_fn = device
2669 .get_func(SCAN_MODULE, scan_kernels::COUNT_MASK)
2670 .ok_or_else(|| XlogError::Kernel("count_mask kernel not found".to_string()))?;
2671
2672 unsafe {
2674 count_fn.clone().launch(
2675 LaunchConfig {
2676 grid_dim: (grid_size, 1, 1),
2677 block_dim: (block_size, 1, 1),
2678 shared_mem_bytes: 0,
2679 },
2680 (mask, n, &mut d_count),
2681 )
2682 }
2683 .map_err(|e| XlogError::Kernel(format!("count_mask kernel failed: {}", e)))?;
2684
2685 self.device.synchronize()?;
2686
2687 Ok(d_count)
2688 }
2689
2690 pub fn count_mask_into_slot(
2699 &self,
2700 mask: &crate::memory::TrackedCudaSlice<u8>,
2701 n: u32,
2702 task_counts: &mut crate::memory::TrackedCudaSlice<u32>,
2703 slot_idx: usize,
2704 ) -> Result<()> {
2705 if n == 0 {
2706 return Ok(());
2708 }
2709 if slot_idx >= task_counts.len() {
2710 return Err(XlogError::Kernel(format!(
2711 "count_mask_into_slot: slot_idx={} >= len={}",
2712 slot_idx,
2713 task_counts.len()
2714 )));
2715 }
2716
2717 let device = self.device.inner();
2718 let block_size = 256u32;
2719 let grid_size = n.div_ceil(block_size);
2720
2721 let count_fn = device
2722 .get_func(SCAN_MODULE, scan_kernels::COUNT_MASK)
2723 .ok_or_else(|| XlogError::Kernel("count_mask kernel not found".to_string()))?;
2724
2725 let mut slot = task_counts.slice_mut(slot_idx..slot_idx + 1);
2727
2728 unsafe {
2732 count_fn.clone().launch(
2733 LaunchConfig {
2734 grid_dim: (grid_size, 1, 1),
2735 block_dim: (block_size, 1, 1),
2736 shared_mem_bytes: 0,
2737 },
2738 (mask, n, &mut slot),
2739 )
2740 }
2741 .map_err(|e| XlogError::Kernel(format!("count_mask_into_slot kernel failed: {}", e)))?;
2742
2743 Ok(())
2744 }
2745 fn apply_permutation_gpu(
2747 &self,
2748 input: &CudaBuffer,
2749 permutation: &impl DeviceRead<u32>,
2750 ) -> Result<CudaBuffer> {
2751 let permutation = &permutation.device_view();
2752 let row_cap = input.num_rows() as u32;
2753 let d_num_rows = input.num_rows_device();
2754 let device = self.device.inner();
2755
2756 let grid_size = row_cap.div_ceil(Self::SORT_BLOCK_SIZE);
2757 let launch_config = LaunchConfig {
2758 grid_dim: (grid_size, 1, 1),
2759 block_dim: (Self::SORT_BLOCK_SIZE, 1, 1),
2760 shared_mem_bytes: 0,
2761 };
2762
2763 let apply_perm_fn = device
2764 .get_func(SORT_MODULE, sort_kernels::APPLY_PERMUTATION_BYTES)
2765 .ok_or_else(|| {
2766 XlogError::Kernel("apply_permutation_bytes kernel not found".to_string())
2767 })?;
2768
2769 let mut new_columns = Vec::with_capacity(input.columns.len());
2770
2771 for col_idx in 0..input.columns.len() {
2772 let src_col = input
2773 .column(col_idx)
2774 .ok_or_else(|| XlogError::Kernel(format!("Column {} not found", col_idx)))?;
2775
2776 let elem_size = input
2777 .schema
2778 .column_type(col_idx)
2779 .ok_or_else(|| {
2780 XlogError::Kernel(format!("Schema type for column {} not found", col_idx))
2781 })?
2782 .size_bytes() as u32;
2783
2784 let output_bytes = (row_cap as usize) * (elem_size as usize);
2785 if src_col.num_bytes() != output_bytes {
2786 return Err(XlogError::Kernel(format!(
2787 "Column {} has {} bytes but expected {} (num_rows={}, elem_size={})",
2788 col_idx,
2789 src_col.num_bytes(),
2790 output_bytes,
2791 row_cap,
2792 elem_size
2793 )));
2794 }
2795 let dst_col = self.memory.alloc::<u8>(output_bytes)?;
2796
2797 unsafe {
2799 apply_perm_fn.clone().launch(
2800 launch_config,
2801 (
2802 src_col,
2803 &dst_col,
2804 permutation,
2805 d_num_rows,
2806 row_cap,
2807 elem_size,
2808 ),
2809 )
2810 }
2811 .map_err(|e| XlogError::Kernel(format!("apply_permutation_bytes failed: {}", e)))?;
2812
2813 new_columns.push(dst_col.into());
2814 }
2815
2816 self.device.synchronize()?;
2817
2818 self.buffer_from_columns_with_device_count(
2819 new_columns,
2820 input.num_rows(),
2821 input.schema.clone(),
2822 input,
2823 )
2824 }
2825
2826 fn gather_buffer_by_indices(
2831 &self,
2832 input: &CudaBuffer,
2833 indices: &impl DeviceRead<u32>,
2834 output_rows: u32,
2835 ) -> Result<CudaBuffer> {
2836 let indices = &indices.device_view();
2837 if output_rows == 0 {
2838 return self.create_empty_buffer(input.schema().clone());
2839 }
2840
2841 if input.num_rows() > u32::MAX as u64 {
2842 return Err(XlogError::Kernel(format!(
2843 "GPU gather supports at most {} input rows, got {}",
2844 u32::MAX,
2845 input.num_rows()
2846 )));
2847 }
2848
2849 let d_output_rows = self.upload_device_row_count(output_rows)?;
2850 let device = self.device.inner();
2851 let block_size = 256u32;
2852 let grid_size = output_rows.div_ceil(block_size);
2853 let launch_config = LaunchConfig {
2854 grid_dim: (grid_size, 1, 1),
2855 block_dim: (block_size, 1, 1),
2856 shared_mem_bytes: 0,
2857 };
2858
2859 let gather_fn = device
2860 .get_func(SORT_MODULE, sort_kernels::APPLY_PERMUTATION_BYTES)
2861 .ok_or_else(|| {
2862 XlogError::Kernel("apply_permutation_bytes kernel not found".to_string())
2863 })?;
2864
2865 let mut new_columns = Vec::with_capacity(input.columns.len());
2866 for col_idx in 0..input.columns.len() {
2867 let src_col = input
2868 .column(col_idx)
2869 .ok_or_else(|| XlogError::Kernel(format!("Column {} not found", col_idx)))?;
2870
2871 let elem_size = input
2872 .schema
2873 .column_type(col_idx)
2874 .ok_or_else(|| {
2875 XlogError::Kernel(format!("Schema type for column {} not found", col_idx))
2876 })?
2877 .size_bytes() as u32;
2878
2879 let expected_src_bytes = (input.num_rows() as usize) * (elem_size as usize);
2880 if src_col.num_bytes() != expected_src_bytes {
2881 return Err(XlogError::Kernel(format!(
2882 "Column {} has {} bytes but expected {} (num_rows={}, elem_size={})",
2883 col_idx,
2884 src_col.num_bytes(),
2885 expected_src_bytes,
2886 input.num_rows(),
2887 elem_size
2888 )));
2889 }
2890
2891 let dst_bytes = (output_rows as usize) * (elem_size as usize);
2892 let dst_col = self.memory.alloc::<u8>(dst_bytes)?;
2893
2894 unsafe {
2896 gather_fn.clone().launch(
2897 launch_config,
2898 (
2899 src_col,
2900 &dst_col,
2901 indices,
2902 &d_output_rows,
2903 output_rows,
2904 elem_size,
2905 ),
2906 )
2907 }
2908 .map_err(|e| XlogError::Kernel(format!("apply_permutation_bytes failed: {}", e)))?;
2909
2910 new_columns.push(dst_col.into());
2911 }
2912
2913 self.device.synchronize()?;
2914
2915 Ok(CudaBuffer::from_columns(
2916 new_columns,
2917 output_rows as u64,
2918 d_output_rows,
2919 input.schema.clone(),
2920 ))
2921 }
2922 pub fn hash_join_v2(
2936 &self,
2937 left: &CudaBuffer,
2938 right: &CudaBuffer,
2939 left_keys: &[usize],
2940 right_keys: &[usize],
2941 join_type: JoinType,
2942 ) -> Result<CudaBuffer> {
2943 self.hash_join_v2_with_limit(left, right, left_keys, right_keys, join_type, None)
2944 }
2945
2946 pub fn hash_join_v2_with_limit(
2962 &self,
2963 left: &CudaBuffer,
2964 right: &CudaBuffer,
2965 left_keys: &[usize],
2966 right_keys: &[usize],
2967 join_type: JoinType,
2968 max_output: Option<usize>,
2969 ) -> Result<CudaBuffer> {
2970 if Self::use_recorded_hash_join_env()?
2975 && !left_keys.is_empty()
2976 && left_keys.len() == right_keys.len()
2977 && left_keys.len() <= 4
2978 {
2979 if let Some(launch_stream) = self.recorded_op_stream_or_init() {
2980 return self.hash_join_v2_recorded(
2981 left,
2982 right,
2983 left_keys,
2984 right_keys,
2985 join_type,
2986 max_output,
2987 launch_stream,
2988 );
2989 }
2990 }
2991 match join_type {
2992 JoinType::Inner => {
2993 self.hash_join_inner_v2(left, right, left_keys, right_keys, max_output)
2994 }
2995 JoinType::Semi => self.hash_join_semi_impl(left, right, left_keys, right_keys),
2996 JoinType::Anti => self.hash_join_anti_impl(left, right, left_keys, right_keys),
2997 JoinType::LeftOuter => {
2998 self.hash_join_left_outer_impl(left, right, left_keys, right_keys, max_output)
2999 }
3000 }
3001 }
3002
3003 pub fn nested_loop_join_v2_inner_u32_1key(
3051 &self,
3052 left: &CudaBuffer,
3053 right: &CudaBuffer,
3054 left_key: usize,
3055 right_key: usize,
3056 ) -> Result<CudaBuffer> {
3057 let num_left = self.device_row_count(left)?;
3059 let num_right = self.device_row_count(right)?;
3060
3061 if num_left == 0 || num_right == 0 {
3063 let combined_schema = self.combine_schemas(left.schema(), right.schema());
3064 return self.create_empty_buffer(combined_schema);
3065 }
3066
3067 if left.arity() <= left_key {
3069 return Err(XlogError::Kernel(format!(
3070 "nested_loop: left_key={} out of bounds (arity={})",
3071 left_key,
3072 left.arity()
3073 )));
3074 }
3075 if right.arity() <= right_key {
3076 return Err(XlogError::Kernel(format!(
3077 "nested_loop: right_key={} out of bounds (arity={})",
3078 right_key,
3079 right.arity()
3080 )));
3081 }
3082 let lt = left.schema().column_type(left_key);
3083 let rt = right.schema().column_type(right_key);
3084 if lt != rt || !matches!(lt, Some(ScalarType::U32) | Some(ScalarType::Symbol)) {
3085 return Err(XlogError::Kernel(format!(
3086 "nested_loop: key types must be equal U32/Symbol; got left={:?} right={:?}",
3087 lt, rt
3088 )));
3089 }
3090 let left_col = left
3091 .column(left_key)
3092 .ok_or_else(|| XlogError::Kernel(format!("nested_loop: left.column({})", left_key)))?;
3093 let right_col = right.column(right_key).ok_or_else(|| {
3094 XlogError::Kernel(format!("nested_loop: right.column({})", right_key))
3095 })?;
3096 let required_left_bytes = num_left
3111 .checked_mul(4)
3112 .ok_or_else(|| XlogError::Kernel("nested_loop: left byte-count overflow".into()))?;
3113 let required_right_bytes = num_right
3114 .checked_mul(4)
3115 .ok_or_else(|| XlogError::Kernel("nested_loop: right byte-count overflow".into()))?;
3116 if left_col.num_bytes() < required_left_bytes {
3117 return Err(XlogError::Kernel(format!(
3118 "nested_loop: left key column has {} bytes; \
3119 require at least {} ({} rows × 4) — buffer allocation \
3120 is smaller than logical row count",
3121 left_col.num_bytes(),
3122 required_left_bytes,
3123 num_left
3124 )));
3125 }
3126 if right_col.num_bytes() < required_right_bytes {
3127 return Err(XlogError::Kernel(format!(
3128 "nested_loop: right key column has {} bytes; \
3129 require at least {} ({} rows × 4) — buffer allocation \
3130 is smaller than logical row count",
3131 right_col.num_bytes(),
3132 required_right_bytes,
3133 num_right
3134 )));
3135 }
3136
3137 let upper_bound: u64 = (num_left as u64)
3139 .checked_mul(num_right as u64)
3140 .ok_or_else(|| XlogError::Kernel("nested_loop: row-count product overflow".into()))?;
3141 if upper_bound > NESTED_LOOP_TOTAL_THRESHOLD {
3142 return Err(XlogError::Kernel(format!(
3143 "nested_loop: caller violated eligibility threshold: \
3144 num_left * num_right = {} > {} (NESTED_LOOP_TOTAL_THRESHOLD)",
3145 upper_bound, NESTED_LOOP_TOTAL_THRESHOLD
3146 )));
3147 }
3148
3149 let upper_bound_usize = upper_bound as usize;
3151 let mut d_output_left_idx = self.memory.alloc::<u32>(upper_bound_usize)?;
3152 let mut d_output_right_idx = self.memory.alloc::<u32>(upper_bound_usize)?;
3153 let mut d_output_count = self.memory.alloc::<u32>(1)?;
3154 self.device
3155 .inner()
3156 .memset_zeros(&mut d_output_count)
3157 .map_err(|e| XlogError::Kernel(format!("nested_loop: counter zero failed: {}", e)))?;
3158
3159 let func = self
3161 .device
3162 .inner()
3163 .get_func(
3164 JOIN_MODULE,
3165 join_kernels::NESTED_LOOP_JOIN_INNER_U32_1KEY_PAIRS,
3166 )
3167 .ok_or_else(|| {
3168 XlogError::Kernel("nested_loop_join_inner_u32_1key_pairs kernel not found".into())
3169 })?;
3170
3171 let num_left_u32 = num_left as u32;
3172 let num_right_u32 = num_right as u32;
3173 let upper_bound_u32 = upper_bound as u32;
3174 let block_size = 256u32;
3175 let grid_size = num_left_u32.div_ceil(block_size);
3176 let config = LaunchConfig {
3177 grid_dim: (grid_size, 1, 1),
3178 block_dim: (block_size, 1, 1),
3179 shared_mem_bytes: 0,
3180 };
3181
3182 unsafe {
3192 func.clone()
3193 .launch(
3194 config,
3195 (
3196 left_col,
3197 right_col,
3198 num_left_u32,
3199 num_right_u32,
3200 &mut d_output_left_idx,
3201 &mut d_output_right_idx,
3202 &mut d_output_count,
3203 upper_bound_u32,
3204 ),
3205 )
3206 .map_err(|e| XlogError::Kernel(format!("nested_loop launch failed: {}", e)))?;
3207 }
3208
3209 self.device.synchronize()?;
3210
3211 let output_rows = self.dtoh_scalar_untracked(&d_output_count, 0)?;
3213 if (output_rows as u64) > upper_bound {
3217 return Err(XlogError::Kernel(format!(
3218 "nested_loop: kernel reported {} output rows > upper_bound {}",
3219 output_rows, upper_bound
3220 )));
3221 }
3222
3223 let gathered_left = self.gather_buffer_by_indices(left, &d_output_left_idx, output_rows)?;
3225 let gathered_right =
3226 self.gather_buffer_by_indices(right, &d_output_right_idx, output_rows)?;
3227
3228 let combined_schema = self.combine_schemas(left.schema(), right.schema());
3230 let mut result_columns = Vec::with_capacity(combined_schema.arity());
3231 result_columns.extend(gathered_left.columns);
3232 result_columns.extend(gathered_right.columns);
3233 self.buffer_from_columns(result_columns, output_rows as u64, combined_schema)
3236 }
3237
3238 pub fn is_sorted_ascending_u32(&self, buf: &CudaBuffer, key_col: usize) -> Result<bool> {
3268 let n = self.device_row_count(buf)?;
3270 if n < 2 {
3271 return Ok(true);
3272 }
3273
3274 if buf.arity() <= key_col {
3276 return Err(XlogError::Kernel(format!(
3277 "is_sorted_ascending_u32: key_col={} out of bounds (arity={})",
3278 key_col,
3279 buf.arity()
3280 )));
3281 }
3282 let kt = buf.schema().column_type(key_col);
3283 if !matches!(kt, Some(ScalarType::U32) | Some(ScalarType::Symbol)) {
3284 return Err(XlogError::Kernel(format!(
3285 "is_sorted_ascending_u32: key column must be U32 or Symbol; got {:?}",
3286 kt
3287 )));
3288 }
3289 let key_column = buf.column(key_col).ok_or_else(|| {
3290 XlogError::Kernel(format!(
3291 "is_sorted_ascending_u32: column({}) missing",
3292 key_col
3293 ))
3294 })?;
3295 let required_bytes = n
3296 .checked_mul(4)
3297 .ok_or_else(|| XlogError::Kernel("is_sorted_ascending_u32: byte overflow".into()))?;
3298 if key_column.num_bytes() < required_bytes {
3299 return Err(XlogError::Kernel(format!(
3300 "is_sorted_ascending_u32: key column has {} bytes; require at least {} ({} rows × 4)",
3301 key_column.num_bytes(),
3302 required_bytes,
3303 n
3304 )));
3305 }
3306
3307 let mut d_result = self.memory.alloc::<u32>(1)?;
3311 self.htod_launch_metadata_sync_copy_into(&[1u32], &mut d_result)
3312 .map_err(|e| {
3313 XlogError::Kernel(format!("is_sorted_ascending_u32: htod result init: {}", e))
3314 })?;
3315
3316 let func = self
3318 .device
3319 .inner()
3320 .get_func(SORT_MODULE, sort_kernels::CHECK_ASCENDING_SORTED_U32)
3321 .ok_or_else(|| {
3322 XlogError::Kernel("check_ascending_sorted_u32 kernel not found".into())
3323 })?;
3324 let n_u32 = n as u32;
3325 let block_size = 256u32;
3326 let grid_size = n_u32.div_ceil(block_size);
3327 let config = LaunchConfig {
3328 grid_dim: (grid_size, 1, 1),
3329 block_dim: (block_size, 1, 1),
3330 shared_mem_bytes: 0,
3331 };
3332
3333 unsafe {
3341 func.clone()
3342 .launch(config, (key_column, n_u32, &mut d_result))
3343 .map_err(|e| {
3344 XlogError::Kernel(format!("check_ascending_sorted_u32 launch: {}", e))
3345 })?;
3346 }
3347
3348 self.device.synchronize()?;
3349 let result = self.dtoh_scalar_untracked(&d_result, 0)?;
3350 Ok(result == 1)
3351 }
3352
3353 pub fn sort_merge_join_v2_inner_u32_1key(
3411 &self,
3412 left: &CudaBuffer,
3413 right: &CudaBuffer,
3414 left_key: usize,
3415 right_key: usize,
3416 ) -> Result<CudaBuffer> {
3417 let num_left = self.device_row_count(left)?;
3419 let num_right = self.device_row_count(right)?;
3420
3421 if num_left == 0 || num_right == 0 {
3423 let combined_schema = self.combine_schemas(left.schema(), right.schema());
3424 return self.create_empty_buffer(combined_schema);
3425 }
3426
3427 if left.arity() <= left_key {
3429 return Err(XlogError::Kernel(format!(
3430 "sort_merge: left_key={} out of bounds (arity={})",
3431 left_key,
3432 left.arity()
3433 )));
3434 }
3435 if right.arity() <= right_key {
3436 return Err(XlogError::Kernel(format!(
3437 "sort_merge: right_key={} out of bounds (arity={})",
3438 right_key,
3439 right.arity()
3440 )));
3441 }
3442 let lt = left.schema().column_type(left_key);
3443 let rt = right.schema().column_type(right_key);
3444 if lt != rt || !matches!(lt, Some(ScalarType::U32) | Some(ScalarType::Symbol)) {
3445 return Err(XlogError::Kernel(format!(
3446 "sort_merge: key types must be equal U32/Symbol; got left={:?} right={:?}",
3447 lt, rt
3448 )));
3449 }
3450 let left_col = left
3451 .column(left_key)
3452 .ok_or_else(|| XlogError::Kernel(format!("sort_merge: left.column({})", left_key)))?;
3453 let right_col = right
3454 .column(right_key)
3455 .ok_or_else(|| XlogError::Kernel(format!("sort_merge: right.column({})", right_key)))?;
3456 let required_left_bytes = num_left
3457 .checked_mul(4)
3458 .ok_or_else(|| XlogError::Kernel("sort_merge: left byte overflow".into()))?;
3459 let required_right_bytes = num_right
3460 .checked_mul(4)
3461 .ok_or_else(|| XlogError::Kernel("sort_merge: right byte overflow".into()))?;
3462 if left_col.num_bytes() < required_left_bytes {
3463 return Err(XlogError::Kernel(format!(
3464 "sort_merge: left key column has {} bytes; \
3465 require at least {} ({} rows × 4)",
3466 left_col.num_bytes(),
3467 required_left_bytes,
3468 num_left
3469 )));
3470 }
3471 if right_col.num_bytes() < required_right_bytes {
3472 return Err(XlogError::Kernel(format!(
3473 "sort_merge: right key column has {} bytes; \
3474 require at least {} ({} rows × 4)",
3475 right_col.num_bytes(),
3476 required_right_bytes,
3477 num_right
3478 )));
3479 }
3480
3481 let upper_bound: u64 = (num_left as u64)
3483 .checked_mul(num_right as u64)
3484 .ok_or_else(|| XlogError::Kernel("sort_merge: row-count product overflow".into()))?;
3485 if upper_bound > NESTED_LOOP_TOTAL_THRESHOLD {
3486 return Err(XlogError::Kernel(format!(
3487 "sort_merge: caller violated eligibility threshold: \
3488 num_left * num_right = {} > {} (NESTED_LOOP_TOTAL_THRESHOLD)",
3489 upper_bound, NESTED_LOOP_TOTAL_THRESHOLD
3490 )));
3491 }
3492
3493 let upper_bound_usize = upper_bound as usize;
3495 let mut d_output_left_idx = self.memory.alloc::<u32>(upper_bound_usize)?;
3496 let mut d_output_right_idx = self.memory.alloc::<u32>(upper_bound_usize)?;
3497 let mut d_output_count = self.memory.alloc::<u32>(1)?;
3498 self.device
3499 .inner()
3500 .memset_zeros(&mut d_output_count)
3501 .map_err(|e| XlogError::Kernel(format!("sort_merge: counter zero: {}", e)))?;
3502
3503 let func = self
3505 .device
3506 .inner()
3507 .get_func(
3508 JOIN_MODULE,
3509 join_kernels::SORT_MERGE_JOIN_INNER_U32_1KEY_PAIRS,
3510 )
3511 .ok_or_else(|| {
3512 XlogError::Kernel("sort_merge_join_inner_u32_1key_pairs kernel not found".into())
3513 })?;
3514
3515 let num_left_u32 = num_left as u32;
3516 let num_right_u32 = num_right as u32;
3517 let upper_bound_u32 = upper_bound as u32;
3518 let block_size = 256u32;
3519 let grid_size = num_left_u32.div_ceil(block_size);
3520 let config = LaunchConfig {
3521 grid_dim: (grid_size, 1, 1),
3522 block_dim: (block_size, 1, 1),
3523 shared_mem_bytes: 0,
3524 };
3525
3526 unsafe {
3539 func.clone()
3540 .launch(
3541 config,
3542 (
3543 left_col,
3544 right_col,
3545 num_left_u32,
3546 num_right_u32,
3547 &mut d_output_left_idx,
3548 &mut d_output_right_idx,
3549 &mut d_output_count,
3550 upper_bound_u32,
3551 ),
3552 )
3553 .map_err(|e| XlogError::Kernel(format!("sort_merge launch: {}", e)))?;
3554 }
3555
3556 self.device.synchronize()?;
3557
3558 let output_rows = self.dtoh_scalar_untracked(&d_output_count, 0)?;
3560 if (output_rows as u64) > upper_bound {
3563 return Err(XlogError::Kernel(format!(
3564 "sort_merge: kernel reported {} output rows > upper_bound {}",
3565 output_rows, upper_bound
3566 )));
3567 }
3568
3569 let gathered_left = self.gather_buffer_by_indices(left, &d_output_left_idx, output_rows)?;
3571 let gathered_right =
3572 self.gather_buffer_by_indices(right, &d_output_right_idx, output_rows)?;
3573
3574 let combined_schema = self.combine_schemas(left.schema(), right.schema());
3576 let mut result_columns = Vec::with_capacity(combined_schema.arity());
3577 result_columns.extend(gathered_left.columns);
3578 result_columns.extend(gathered_right.columns);
3579 self.buffer_from_columns(result_columns, output_rows as u64, combined_schema)
3580 }
3581
3582 pub fn sort_merge_join_v2_inner_u32_1key_bounded(
3592 &self,
3593 left: &CudaBuffer,
3594 right: &CudaBuffer,
3595 left_key: usize,
3596 right_key: usize,
3597 output_capacity: usize,
3598 ) -> Result<CudaBuffer> {
3599 let num_left = self.device_row_count(left)?;
3600 let num_right = self.device_row_count(right)?;
3601
3602 if num_left == 0 || num_right == 0 {
3603 let combined_schema = self.combine_schemas(left.schema(), right.schema());
3604 return self.create_empty_buffer(combined_schema);
3605 }
3606 if num_left > u32::MAX as usize || num_right > u32::MAX as usize {
3607 return Err(XlogError::Kernel(format!(
3608 "sort_merge_bounded: row counts exceed u32 surface: left={} right={}",
3609 num_left, num_right
3610 )));
3611 }
3612 if output_capacity == 0 || output_capacity > u32::MAX as usize {
3613 return Err(XlogError::Kernel(format!(
3614 "sort_merge_bounded: invalid output capacity {}",
3615 output_capacity
3616 )));
3617 }
3618 if left.arity() <= left_key {
3619 return Err(XlogError::Kernel(format!(
3620 "sort_merge_bounded: left_key={} out of bounds (arity={})",
3621 left_key,
3622 left.arity()
3623 )));
3624 }
3625 if right.arity() <= right_key {
3626 return Err(XlogError::Kernel(format!(
3627 "sort_merge_bounded: right_key={} out of bounds (arity={})",
3628 right_key,
3629 right.arity()
3630 )));
3631 }
3632 let lt = left.schema().column_type(left_key);
3633 let rt = right.schema().column_type(right_key);
3634 if lt != rt || !matches!(lt, Some(ScalarType::U32) | Some(ScalarType::Symbol)) {
3635 return Err(XlogError::Kernel(format!(
3636 "sort_merge_bounded: key types must be equal U32/Symbol; got left={:?} right={:?}",
3637 lt, rt
3638 )));
3639 }
3640
3641 let left_col = left.column(left_key).ok_or_else(|| {
3642 XlogError::Kernel(format!("sort_merge_bounded: left.column({})", left_key))
3643 })?;
3644 let right_col = right.column(right_key).ok_or_else(|| {
3645 XlogError::Kernel(format!("sort_merge_bounded: right.column({})", right_key))
3646 })?;
3647 let required_left_bytes = num_left
3648 .checked_mul(4)
3649 .ok_or_else(|| XlogError::Kernel("sort_merge_bounded: left byte overflow".into()))?;
3650 let required_right_bytes = num_right
3651 .checked_mul(4)
3652 .ok_or_else(|| XlogError::Kernel("sort_merge_bounded: right byte overflow".into()))?;
3653 if left_col.num_bytes() < required_left_bytes {
3654 return Err(XlogError::Kernel(format!(
3655 "sort_merge_bounded: left key column has {} bytes; require at least {}",
3656 left_col.num_bytes(),
3657 required_left_bytes
3658 )));
3659 }
3660 if right_col.num_bytes() < required_right_bytes {
3661 return Err(XlogError::Kernel(format!(
3662 "sort_merge_bounded: right key column has {} bytes; require at least {}",
3663 right_col.num_bytes(),
3664 required_right_bytes
3665 )));
3666 }
3667
3668 let mut d_output_left_idx = self.memory.alloc::<u32>(output_capacity)?;
3669 let mut d_output_right_idx = self.memory.alloc::<u32>(output_capacity)?;
3670 let mut d_output_count = self.memory.alloc::<u32>(1)?;
3671 self.device
3672 .inner()
3673 .memset_zeros(&mut d_output_count)
3674 .map_err(|e| XlogError::Kernel(format!("sort_merge_bounded: counter zero: {}", e)))?;
3675
3676 let func = self
3677 .device
3678 .inner()
3679 .get_func(
3680 JOIN_MODULE,
3681 join_kernels::SORT_MERGE_JOIN_INNER_U32_1KEY_PAIRS,
3682 )
3683 .ok_or_else(|| {
3684 XlogError::Kernel("sort_merge_join_inner_u32_1key_pairs kernel not found".into())
3685 })?;
3686
3687 let num_left_u32 = num_left as u32;
3688 let num_right_u32 = num_right as u32;
3689 let output_capacity_u32 = output_capacity as u32;
3690 let block_size = 256u32;
3691 let grid_size = num_left_u32.div_ceil(block_size);
3692 let config = LaunchConfig {
3693 grid_dim: (grid_size, 1, 1),
3694 block_dim: (block_size, 1, 1),
3695 shared_mem_bytes: 0,
3696 };
3697
3698 unsafe {
3699 func.clone()
3700 .launch(
3701 config,
3702 (
3703 left_col,
3704 right_col,
3705 num_left_u32,
3706 num_right_u32,
3707 &mut d_output_left_idx,
3708 &mut d_output_right_idx,
3709 &mut d_output_count,
3710 output_capacity_u32,
3711 ),
3712 )
3713 .map_err(|e| XlogError::Kernel(format!("sort_merge_bounded launch: {}", e)))?;
3714 }
3715
3716 self.device.synchronize()?;
3717 let output_rows = self.dtoh_scalar_untracked(&d_output_count, 0)?;
3718 if output_rows as usize > output_capacity {
3719 return Err(XlogError::Kernel(format!(
3720 "sort_merge_bounded: output {} exceeded bounded capacity {}",
3721 output_rows, output_capacity
3722 )));
3723 }
3724
3725 let gathered_left = self.gather_buffer_by_indices(left, &d_output_left_idx, output_rows)?;
3726 let gathered_right =
3727 self.gather_buffer_by_indices(right, &d_output_right_idx, output_rows)?;
3728
3729 let combined_schema = self.combine_schemas(left.schema(), right.schema());
3730 let mut result_columns = Vec::with_capacity(combined_schema.arity());
3731 result_columns.extend(gathered_left.columns);
3732 result_columns.extend(gathered_right.columns);
3733 self.buffer_from_columns(result_columns, output_rows as u64, combined_schema)
3734 }
3735
3736 pub fn build_join_index_v2(
3738 &self,
3739 right: &CudaBuffer,
3740 right_keys: &[usize],
3741 ) -> Result<JoinIndexV2> {
3742 let num_right = self.device_row_count(right)?;
3743 if num_right == 0 {
3744 return Err(XlogError::Kernel(
3745 "Cannot build join index for empty relation".to_string(),
3746 ));
3747 }
3748 if num_right > u32::MAX as usize {
3749 return Err(XlogError::Kernel(format!(
3750 "Join index supports at most {} rows, got {}",
3751 u32::MAX,
3752 num_right
3753 )));
3754 }
3755 if right_keys.is_empty() {
3756 return Err(XlogError::Kernel(
3757 "Join requires at least one key column".to_string(),
3758 ));
3759 }
3760 for &k in right_keys {
3761 if k >= right.arity() {
3762 return Err(XlogError::Kernel(format!(
3763 "Right key column index {} out of bounds (arity {})",
3764 k,
3765 right.arity()
3766 )));
3767 }
3768 }
3769
3770 let num_right = num_right as u32;
3771 let right_packed = self.compute_hashes_and_pack_keys(right, right_keys)?;
3772 let table = self.build_hash_table_v2(&right_packed.hashes, num_right)?;
3773
3774 Ok(JoinIndexV2 {
3775 right_num_rows: num_right,
3776 right_keys: right_keys.to_vec(),
3777 key_bytes: right_packed.key_bytes,
3778 packed_keys: right_packed.packed_keys,
3779 table,
3780 })
3781 }
3782
3783 pub fn build_join_index_v2_background(
3790 &self,
3791 right: &CudaBuffer,
3792 right_keys: &[usize],
3793 ) -> Result<JoinIndexV2> {
3794 if Self::use_recorded_hash_join_env()?
3795 && !right_keys.is_empty()
3796 && right_keys.len() <= 4
3797 && right.num_rows() > 0
3798 {
3799 if let Some(launch_stream) = self.recorded_op_stream_or_init() {
3800 return self.build_join_index_v2_recorded(right, right_keys, launch_stream);
3801 }
3802 }
3803
3804 self.build_join_index_v2(right, right_keys)
3805 }
3806
3807 pub fn build_join_index_v2_recorded(
3813 &self,
3814 right: &CudaBuffer,
3815 right_keys: &[usize],
3816 launch_stream: StreamId,
3817 ) -> Result<JoinIndexV2> {
3818 let runtime = self.memory.runtime().ok_or_else(|| {
3819 XlogError::Kernel(
3820 "build_join_index_v2_recorded requires a runtime-backed GpuMemoryManager"
3821 .to_string(),
3822 )
3823 })?;
3824 let cu_stream = runtime
3825 .stream_pool()
3826 .resolve(launch_stream)
3827 .ok_or_else(|| {
3828 XlogError::Kernel(format!(
3829 "build_join_index_v2_recorded: launch_stream StreamId({}) does not resolve",
3830 launch_stream.0
3831 ))
3832 })?;
3833
3834 let num_right = self.device_row_count(right)?;
3835 if num_right == 0 {
3836 return Err(XlogError::Kernel(
3837 "Cannot build join index for empty relation".to_string(),
3838 ));
3839 }
3840 if num_right > u32::MAX as usize {
3841 return Err(XlogError::Kernel(format!(
3842 "Join index supports at most {} rows, got {}",
3843 u32::MAX,
3844 num_right
3845 )));
3846 }
3847 if right_keys.is_empty() {
3848 return Err(XlogError::Kernel(
3849 "Join requires at least one key column".to_string(),
3850 ));
3851 }
3852 if right_keys.len() > 4 {
3853 return Err(XlogError::Kernel(
3854 "build_join_index_v2_recorded: max 4 key columns supported".to_string(),
3855 ));
3856 }
3857 for &k in right_keys {
3858 if k >= right.arity() {
3859 return Err(XlogError::Kernel(format!(
3860 "Right key column index {} out of bounds (arity {})",
3861 k,
3862 right.arity()
3863 )));
3864 }
3865 }
3866
3867 let num_right = num_right as u32;
3868 let right_packed =
3869 self.pack_keys_gpu_on_stream(right, right_keys, &cu_stream, launch_stream, runtime)?;
3870 let table = self.build_hash_table_v2_on_stream(
3871 &right_packed.hashes,
3872 num_right,
3873 &cu_stream,
3874 launch_stream,
3875 runtime,
3876 )?;
3877
3878 Ok(JoinIndexV2 {
3879 right_num_rows: num_right,
3880 right_keys: right_keys.to_vec(),
3881 key_bytes: right_packed.key_bytes,
3882 packed_keys: right_packed.packed_keys,
3883 table,
3884 })
3885 }
3886
3887 #[expect(
3891 clippy::too_many_arguments,
3892 reason = "the cached join entry accepts both relations, their key projections, join semantics, cache identity, and output bound"
3893 )]
3894 pub fn hash_join_v2_with_index(
3895 &self,
3896 left: &CudaBuffer,
3897 right: &CudaBuffer,
3898 left_keys: &[usize],
3899 right_keys: &[usize],
3900 join_type: JoinType,
3901 index: &JoinIndexV2,
3902 max_output: Option<usize>,
3903 ) -> Result<CudaBuffer> {
3904 if Self::use_recorded_hash_join_env()?
3907 && !left_keys.is_empty()
3908 && left_keys.len() == right_keys.len()
3909 && left_keys.len() <= 4
3910 {
3911 if let Some(launch_stream) = self.recorded_op_stream_or_init() {
3912 return self.hash_join_v2_with_index_recorded(
3913 left,
3914 right,
3915 left_keys,
3916 right_keys,
3917 join_type,
3918 index,
3919 max_output,
3920 launch_stream,
3921 );
3922 }
3923 }
3924 let left_rows = self.device_row_count(left)?;
3925 let right_rows = self.device_row_count(right)?;
3926 if left_rows > u32::MAX as usize || right_rows > u32::MAX as usize {
3927 return Err(XlogError::Kernel(format!(
3928 "Join supports at most {} rows per side (left={}, right={})",
3929 u32::MAX,
3930 left_rows,
3931 right_rows
3932 )));
3933 }
3934
3935 if left_rows == 0 {
3937 return match join_type {
3938 JoinType::Inner | JoinType::LeftOuter => {
3939 let combined_schema = self.combine_schemas(left.schema(), right.schema());
3940 self.create_empty_buffer(combined_schema)
3941 }
3942 JoinType::Semi | JoinType::Anti => self.create_empty_buffer(left.schema().clone()),
3943 };
3944 }
3945 if right_rows == 0 {
3946 return match join_type {
3947 JoinType::Inner => {
3948 let combined_schema = self.combine_schemas(left.schema(), right.schema());
3949 self.create_empty_buffer(combined_schema)
3950 }
3951 JoinType::Semi => self.create_empty_buffer(left.schema().clone()),
3952 JoinType::Anti => self.clone_buffer(left),
3953 JoinType::LeftOuter => self.left_outer_with_nulls(left, right),
3954 };
3955 }
3956
3957 if left_keys.is_empty() || right_keys.is_empty() {
3959 return Err(XlogError::Kernel(
3960 "Join requires at least one key column".to_string(),
3961 ));
3962 }
3963 if left_keys.len() != right_keys.len() {
3964 return Err(XlogError::Kernel(
3965 "Left and right key columns must have same length".to_string(),
3966 ));
3967 }
3968 for (&left_idx, &right_idx) in left_keys.iter().zip(right_keys.iter()) {
3969 if left_idx >= left.arity() {
3970 return Err(XlogError::Kernel(format!(
3971 "Left key column index {} out of bounds (arity {})",
3972 left_idx,
3973 left.arity()
3974 )));
3975 }
3976 if right_idx >= right.arity() {
3977 return Err(XlogError::Kernel(format!(
3978 "Right key column index {} out of bounds (arity {})",
3979 right_idx,
3980 right.arity()
3981 )));
3982 }
3983 let left_type = left.schema().column_type(left_idx);
3984 let right_type = right.schema().column_type(right_idx);
3985 if left_type != right_type {
3986 return Err(XlogError::Kernel(format!(
3987 "Key column type mismatch: left[{}]={:?}, right[{}]={:?}",
3988 left_idx, left_type, right_idx, right_type
3989 )));
3990 }
3991 }
3992
3993 if index.right_num_rows != right_rows as u32 {
3995 return Err(XlogError::Kernel(
3996 "Join index row count does not match right relation".to_string(),
3997 ));
3998 }
3999 if index.right_keys.as_slice() != right_keys {
4000 return Err(XlogError::Kernel(
4001 "Join index key columns do not match requested right_keys".to_string(),
4002 ));
4003 }
4004
4005 match join_type {
4006 JoinType::Inner => {
4007 self.hash_join_inner_v2_indexed(left, right, left_keys, index, max_output)
4008 }
4009 JoinType::Semi => self.hash_join_semi_indexed(left, left_keys, index),
4010 JoinType::Anti => self.hash_join_anti_indexed(left, right, left_keys, index),
4011 JoinType::LeftOuter => {
4012 self.hash_join_left_outer_indexed(left, right, left_keys, index, max_output)
4013 }
4014 }
4015 }
4016
4017 fn pack_keys_gpu(&self, buffer: &CudaBuffer, key_cols: &[usize]) -> Result<PackedKeyData> {
4037 if key_cols.is_empty() {
4038 return Err(XlogError::Kernel(
4039 "pack_keys_gpu: no key columns specified".into(),
4040 ));
4041 }
4042 if key_cols.len() > 4 {
4043 return Err(XlogError::Kernel(
4044 "pack_keys_gpu: max 4 key columns supported".into(),
4045 ));
4046 }
4047
4048 let num_rows = self.device_row_count(buffer)?;
4049 if num_rows > u32::MAX as usize {
4050 return Err(XlogError::Kernel(format!(
4051 "pack_keys_gpu supports at most {} rows, got {}",
4052 u32::MAX,
4053 num_rows
4054 )));
4055 }
4056 let num_rows = num_rows as u32;
4057 if num_rows == 0 {
4058 return Ok(PackedKeyData {
4060 hashes: self.memory.alloc::<u64>(0)?,
4061 packed_keys: self.memory.alloc::<u8>(0)?,
4062 key_bytes: 0,
4063 });
4064 }
4065
4066 let mut col_sizes: Vec<u32> = Vec::with_capacity(key_cols.len());
4068 let mut row_size: u32 = 0;
4069 for &col_idx in key_cols {
4070 let col_type = buffer
4071 .schema()
4072 .column_type(col_idx)
4073 .ok_or_else(|| XlogError::Kernel(format!("Invalid column index: {}", col_idx)))?;
4074 let size = col_type.size_bytes() as u32;
4075 col_sizes.push(size);
4076 row_size += size;
4077 }
4078
4079 let packed_bytes = (num_rows as u64) * (row_size as u64);
4081 let packed_slice = self.memory.alloc::<u8>(packed_bytes as usize)?;
4082 let hash_slice = self.memory.alloc::<u64>(num_rows as usize)?;
4083
4084 let mut col_ptrs: [u64; 4] = [0; 4];
4087 for (i, &col_idx) in key_cols.iter().enumerate() {
4088 let col = buffer
4089 .column(col_idx)
4090 .ok_or_else(|| XlogError::Kernel(format!("Key column {} not found", col_idx)))?;
4091 col_ptrs[i] = *col.device_ptr();
4093 }
4094 let mut packed_col_sizes = 0u64;
4095 for (i, size) in col_sizes.iter().copied().enumerate() {
4096 if size > u16::MAX as u32 {
4097 return Err(XlogError::Kernel(format!(
4098 "pack_keys_gpu: column element size {} exceeds 16-bit kernel argument",
4099 size
4100 )));
4101 }
4102 packed_col_sizes |= (size as u64) << (i * 16);
4103 }
4104
4105 let func = self
4107 .device
4108 .inner()
4109 .get_func(PACK_MODULE, pack_kernels::PACK_AND_HASH_KEYS)
4110 .ok_or_else(|| XlogError::Kernel("pack_and_hash_keys kernel not found".to_string()))?;
4111
4112 let block_size = 256u32;
4114 let grid_size = num_rows.div_ceil(block_size);
4115 let config = LaunchConfig {
4116 grid_dim: (grid_size, 1, 1),
4117 block_dim: (block_size, 1, 1),
4118 shared_mem_bytes: 0,
4119 };
4120
4121 unsafe {
4128 func.clone()
4129 .launch(
4130 config,
4131 (
4132 col_ptrs[0],
4133 col_ptrs[1],
4134 col_ptrs[2],
4135 col_ptrs[3],
4136 packed_col_sizes,
4137 key_cols.len() as u32,
4138 num_rows,
4139 row_size,
4140 &packed_slice,
4141 &hash_slice,
4142 ),
4143 )
4144 .map_err(|e| {
4145 XlogError::Kernel(format!("pack_and_hash_keys launch failed: {}", e))
4146 })?;
4147 }
4148
4149 self.device.synchronize()?;
4150
4151 Ok(PackedKeyData {
4152 hashes: hash_slice,
4153 packed_keys: packed_slice,
4154 key_bytes: row_size,
4155 })
4156 }
4157
4158 fn pack_keys_gpu_generic(
4160 &self,
4161 buffer: &CudaBuffer,
4162 key_cols: &[usize],
4163 ) -> Result<PackedKeyData> {
4164 if key_cols.is_empty() {
4165 return Err(XlogError::Kernel(
4166 "pack_keys_gpu_generic: no key columns specified".into(),
4167 ));
4168 }
4169
4170 let num_rows = self.device_row_count(buffer)?;
4171 if num_rows > u32::MAX as usize {
4172 return Err(XlogError::Kernel(format!(
4173 "pack_keys_gpu_generic supports at most {} rows, got {}",
4174 u32::MAX,
4175 num_rows
4176 )));
4177 }
4178 let num_rows = num_rows as u32;
4179 if num_rows == 0 {
4180 return Ok(PackedKeyData {
4181 hashes: self.memory.alloc::<u64>(0)?,
4182 packed_keys: self.memory.alloc::<u8>(0)?,
4183 key_bytes: 0,
4184 });
4185 }
4186
4187 let mut col_sizes: Vec<u32> = Vec::with_capacity(key_cols.len());
4188 let mut col_ptrs: Vec<u64> = Vec::with_capacity(key_cols.len());
4189 let mut row_size: u32 = 0;
4190
4191 for &col_idx in key_cols {
4192 let col_type = buffer
4193 .schema()
4194 .column_type(col_idx)
4195 .ok_or_else(|| XlogError::Kernel(format!("Invalid column index: {}", col_idx)))?;
4196 let size = col_type.size_bytes() as u32;
4197 row_size = row_size
4198 .checked_add(size)
4199 .ok_or_else(|| XlogError::Kernel("Row size overflow".to_string()))?;
4200 col_sizes.push(size);
4201
4202 let col = buffer
4203 .column(col_idx)
4204 .ok_or_else(|| XlogError::Kernel(format!("Key column {} not found", col_idx)))?;
4205 col_ptrs.push(*col.device_ptr());
4206 }
4207
4208 let packed_bytes = (num_rows as u64)
4209 .checked_mul(row_size as u64)
4210 .ok_or_else(|| XlogError::Kernel("Packed key byte size overflow".to_string()))?;
4211 let packed_slice = self.memory.alloc::<u8>(packed_bytes as usize)?;
4212 let hash_slice = self.memory.alloc::<u64>(num_rows as usize)?;
4213
4214 let mut d_col_sizes = self.memory.alloc::<u32>(col_sizes.len())?;
4215 self.htod_sync_copy_into_tracked(&col_sizes, &mut d_col_sizes)
4216 .map_err(|e| XlogError::Kernel(format!("Failed to upload col_sizes: {}", e)))?;
4217
4218 let mut d_col_ptrs = self.memory.alloc::<u64>(col_ptrs.len())?;
4219 self.htod_sync_copy_into_tracked(&col_ptrs, &mut d_col_ptrs)
4220 .map_err(|e| XlogError::Kernel(format!("Failed to upload col_ptrs: {}", e)))?;
4221
4222 let func = self
4223 .device
4224 .inner()
4225 .get_func(PACK_MODULE, pack_kernels::PACK_AND_HASH_KEYS_GENERIC)
4226 .ok_or_else(|| {
4227 XlogError::Kernel("pack_and_hash_keys_generic kernel not found".to_string())
4228 })?;
4229
4230 let block_size = 256u32;
4231 let grid_size = num_rows.div_ceil(block_size);
4232 let config = LaunchConfig {
4233 grid_dim: (grid_size, 1, 1),
4234 block_dim: (block_size, 1, 1),
4235 shared_mem_bytes: 0,
4236 };
4237
4238 unsafe {
4240 func.clone()
4241 .launch(
4242 config,
4243 (
4244 &d_col_ptrs,
4245 &d_col_sizes,
4246 key_cols.len() as u32,
4247 num_rows,
4248 row_size,
4249 &packed_slice,
4250 &hash_slice,
4251 ),
4252 )
4253 .map_err(|e| {
4254 XlogError::Kernel(format!("pack_and_hash_keys_generic launch failed: {}", e))
4255 })?;
4256 }
4257
4258 self.device.synchronize()?;
4259
4260 Ok(PackedKeyData {
4261 hashes: hash_slice,
4262 packed_keys: packed_slice,
4263 key_bytes: row_size,
4264 })
4265 }
4266
4267 pub(super) fn compute_hashes_and_pack_keys(
4275 &self,
4276 buffer: &CudaBuffer,
4277 key_cols: &[usize],
4278 ) -> Result<PackedKeyData> {
4279 if key_cols.is_empty() {
4280 return Err(XlogError::Kernel(
4281 "compute_hashes_and_pack_keys: no key columns specified".to_string(),
4282 ));
4283 }
4284
4285 if key_cols.len() <= 4 {
4286 self.pack_keys_gpu(buffer, key_cols)
4287 } else {
4288 self.pack_keys_gpu_generic(buffer, key_cols)
4289 }
4290 }
4291
4292 fn build_hash_table_v2(
4297 &self,
4298 hashes: &impl DeviceRead<u64>,
4299 num_rows: u32,
4300 ) -> Result<JoinHashTableV2> {
4301 let hashes = &hashes.device_view();
4302 let device = self.device.inner();
4303
4304 let target = (num_rows as u64).saturating_mul(2).max(1024);
4306 let num_buckets_u64 = target.next_power_of_two();
4307 let num_buckets = u32::try_from(num_buckets_u64).map_err(|_| {
4308 XlogError::Kernel(format!(
4309 "Join hash table too large: num_buckets={}",
4310 num_buckets_u64
4311 ))
4312 })?;
4313 let bucket_mask = num_buckets
4314 .checked_sub(1)
4315 .ok_or_else(|| XlogError::Kernel("Join hash table size underflow".to_string()))?;
4316
4317 let mut bucket_counts = self.memory.alloc::<u32>(num_buckets as usize)?;
4318 if num_buckets > 0 {
4319 device
4320 .memset_zeros(&mut bucket_counts)
4321 .map_err(|e| XlogError::Kernel(format!("Failed to zero bucket_counts: {}", e)))?;
4322 self.device.synchronize()?;
4323 }
4324
4325 let block_size = 256u32;
4326 let grid_size = num_rows.div_ceil(block_size);
4327 let config = LaunchConfig {
4328 grid_dim: (grid_size, 1, 1),
4329 block_dim: (block_size, 1, 1),
4330 shared_mem_bytes: 0,
4331 };
4332
4333 let count_fn = device
4334 .get_func(JOIN_MODULE, join_kernels::HASH_JOIN_BUCKET_COUNT_V2)
4335 .ok_or_else(|| {
4336 XlogError::Kernel("hash_join_bucket_count_v2 kernel not found".to_string())
4337 })?;
4338
4339 unsafe {
4341 count_fn
4342 .clone()
4343 .launch(config, (hashes, num_rows, &bucket_counts, bucket_mask))
4344 .map_err(|e| {
4345 XlogError::Kernel(format!("hash_join_bucket_count_v2 failed: {}", e))
4346 })?;
4347 }
4348 self.device.synchronize()?;
4349
4350 let mut bucket_offsets = self.memory.alloc::<u32>(num_buckets as usize)?;
4352 if num_buckets > 0 {
4353 device
4354 .dtod_copy(&bucket_counts, &mut bucket_offsets)
4355 .map_err(|e| XlogError::Kernel(format!("Failed to copy bucket_counts: {}", e)))?;
4356 self.device.synchronize()?;
4357 self.multiblock_scan_u32_inplace(&mut bucket_offsets, num_buckets)?;
4358 self.device.synchronize()?;
4359 }
4360
4361 let mut bucket_cursors = self.memory.alloc::<u32>(num_buckets as usize)?;
4363 if num_buckets > 0 {
4364 device
4365 .dtod_copy(&bucket_offsets, &mut bucket_cursors)
4366 .map_err(|e| XlogError::Kernel(format!("Failed to copy bucket_offsets: {}", e)))?;
4367 self.device.synchronize()?;
4368 }
4369
4370 let bucket_entries = self.memory.alloc::<u32>(num_rows as usize)?;
4371 let bucket_entry_hashes = self.memory.alloc::<u64>(num_rows as usize)?;
4372
4373 let scatter_fn = device
4374 .get_func(JOIN_MODULE, join_kernels::HASH_JOIN_SCATTER_V2)
4375 .ok_or_else(|| {
4376 XlogError::Kernel("hash_join_scatter_v2 kernel not found".to_string())
4377 })?;
4378
4379 unsafe {
4381 scatter_fn
4382 .clone()
4383 .launch(
4384 config,
4385 (
4386 hashes,
4387 num_rows,
4388 &bucket_cursors,
4389 bucket_mask,
4390 &bucket_entries,
4391 &bucket_entry_hashes,
4392 ),
4393 )
4394 .map_err(|e| XlogError::Kernel(format!("hash_join_scatter_v2 failed: {}", e)))?;
4395 }
4396
4397 self.device.synchronize()?;
4398 Ok(JoinHashTableV2 {
4399 bucket_counts,
4400 bucket_offsets,
4401 bucket_entries,
4402 bucket_entry_hashes,
4403 bucket_mask,
4404 })
4405 }
4406
4407 pub fn build_hash_table_u64(
4409 &self,
4410 hashes: &crate::memory::TrackedCudaSlice<u64>,
4411 num_rows: u32,
4412 ) -> Result<HashTableU64> {
4413 let JoinHashTableV2 {
4414 bucket_counts,
4415 bucket_offsets,
4416 bucket_entries,
4417 bucket_entry_hashes,
4418 bucket_mask,
4419 } = self.build_hash_table_v2(hashes, num_rows)?;
4420 Ok(HashTableU64 {
4421 bucket_counts,
4422 bucket_offsets,
4423 bucket_entries,
4424 bucket_entry_hashes,
4425 bucket_mask,
4426 })
4427 }
4428
4429 fn hash_join_inner_v2(
4431 &self,
4432 left: &CudaBuffer,
4433 right: &CudaBuffer,
4434 left_keys: &[usize],
4435 right_keys: &[usize],
4436 max_output: Option<usize>,
4437 ) -> Result<CudaBuffer> {
4438 let num_left = self.device_row_count(left)?;
4439 let num_right = self.device_row_count(right)?;
4440 if num_left > u32::MAX as usize || num_right > u32::MAX as usize {
4441 return Err(XlogError::Kernel(format!(
4442 "Join supports at most {} rows per side (left={}, right={})",
4443 u32::MAX,
4444 num_left,
4445 num_right
4446 )));
4447 }
4448
4449 if num_left == 0 || num_right == 0 {
4451 let combined_schema = self.combine_schemas(left.schema(), right.schema());
4452 return self.create_empty_buffer(combined_schema);
4453 }
4454
4455 if left_keys.is_empty() || right_keys.is_empty() {
4457 return Err(XlogError::Kernel(
4458 "Join requires at least one key column".to_string(),
4459 ));
4460 }
4461 if left_keys.len() != right_keys.len() {
4462 return Err(XlogError::Kernel(
4463 "Left and right key columns must have same length".to_string(),
4464 ));
4465 }
4466
4467 for (&left_idx, &right_idx) in left_keys.iter().zip(right_keys.iter()) {
4469 let left_type = left.schema().column_type(left_idx);
4470 let right_type = right.schema().column_type(right_idx);
4471 if left_type != right_type {
4472 return Err(XlogError::Kernel(format!(
4473 "Key column type mismatch: left[{}]={:?}, right[{}]={:?}",
4474 left_idx, left_type, right_idx, right_type
4475 )));
4476 }
4477 }
4478
4479 let num_left = num_left as u32;
4480 let num_right = num_right as u32;
4481
4482 let left_packed = self.compute_hashes_and_pack_keys(left, left_keys)?;
4484 let right_packed = self.compute_hashes_and_pack_keys(right, right_keys)?;
4485
4486 let table = self.build_hash_table_v2(&right_packed.hashes, num_right)?;
4488
4489 let probe_func = self
4494 .device
4495 .inner()
4496 .get_func(JOIN_MODULE, join_kernels::HASH_JOIN_PROBE_V2)
4497 .ok_or_else(|| XlogError::Kernel("hash_join_probe_v2 kernel not found".to_string()))?;
4498
4499 let block_size = 256u32;
4500 let probe_grid = num_left.div_ceil(block_size);
4501 let probe_config = LaunchConfig {
4502 grid_dim: (probe_grid, 1, 1),
4503 block_dim: (block_size, 1, 1),
4504 shared_mem_bytes: 0,
4505 };
4506
4507 let mut d_count_only = self.memory.alloc::<u32>(1)?;
4508 self.device
4509 .inner()
4510 .memset_zeros(&mut d_count_only)
4511 .map_err(|e| XlogError::Kernel(format!("Failed to zero output count: {}", e)))?;
4512 self.device.synchronize()?;
4513 let d_dummy_left = self.memory.alloc::<u32>(1)?;
4514 let d_dummy_right = self.memory.alloc::<u32>(1)?;
4515 let max_output_count_only = 0u32;
4516
4517 unsafe {
4524 let mut params: Vec<*mut c_void> = vec![
4525 (&left_packed.hashes).as_kernel_param(),
4526 num_left.as_kernel_param(),
4527 (&table.bucket_offsets).as_kernel_param(),
4528 (&table.bucket_counts).as_kernel_param(),
4529 (&table.bucket_entries).as_kernel_param(),
4530 (&table.bucket_entry_hashes).as_kernel_param(),
4531 table.bucket_mask.as_kernel_param(),
4532 (&left_packed.packed_keys).as_kernel_param(),
4533 (&right_packed.packed_keys).as_kernel_param(),
4534 left_packed.key_bytes.as_kernel_param(),
4535 (&d_dummy_left).as_kernel_param(),
4536 (&d_dummy_right).as_kernel_param(),
4537 (&d_count_only).as_kernel_param(),
4538 max_output_count_only.as_kernel_param(),
4539 ];
4540 probe_func
4541 .clone()
4542 .launch(probe_config, &mut params)
4543 .map_err(|e| {
4544 XlogError::Kernel(format!("hash_join_probe_v2 (count) failed: {}", e))
4545 })?;
4546 }
4547
4548 self.device.synchronize()?;
4549
4550 let full_count = self.read_join_output_count_metadata(&d_count_only)? as u64;
4554 let requested = max_output
4555 .map(|limit| (limit as u64).min(full_count))
4556 .unwrap_or(full_count);
4557
4558 if requested == 0 {
4559 let combined_schema = self.combine_schemas(left.schema(), right.schema());
4560 return self.create_empty_buffer(combined_schema);
4561 }
4562
4563 if requested > u32::MAX as u64 {
4564 return Err(XlogError::Kernel(format!(
4565 "Join produced {} rows which exceeds the u32 index limit",
4566 requested
4567 )));
4568 }
4569
4570 let max_output = requested as u32;
4572 let d_output_left = self.memory.alloc::<u32>(max_output as usize)?;
4573 let d_output_right = self.memory.alloc::<u32>(max_output as usize)?;
4574 let mut d_output_count = self.memory.alloc::<u32>(1)?;
4575 self.device
4576 .inner()
4577 .memset_zeros(&mut d_output_count)
4578 .map_err(|e| XlogError::Kernel(format!("Failed to zero output count: {}", e)))?;
4579 self.device.synchronize()?;
4580
4581 unsafe {
4588 let mut params: Vec<*mut c_void> = vec![
4589 (&left_packed.hashes).as_kernel_param(),
4590 num_left.as_kernel_param(),
4591 (&table.bucket_offsets).as_kernel_param(),
4592 (&table.bucket_counts).as_kernel_param(),
4593 (&table.bucket_entries).as_kernel_param(),
4594 (&table.bucket_entry_hashes).as_kernel_param(),
4595 table.bucket_mask.as_kernel_param(),
4596 (&left_packed.packed_keys).as_kernel_param(),
4597 (&right_packed.packed_keys).as_kernel_param(),
4598 left_packed.key_bytes.as_kernel_param(),
4599 (&d_output_left).as_kernel_param(),
4600 (&d_output_right).as_kernel_param(),
4601 (&d_output_count).as_kernel_param(),
4602 max_output.as_kernel_param(),
4603 ];
4604 probe_func
4605 .clone()
4606 .launch(probe_config, &mut params)
4607 .map_err(|e| XlogError::Kernel(format!("hash_join_probe_v2 failed: {}", e)))?;
4608 }
4609
4610 self.device.synchronize()?;
4611
4612 let result_count =
4619 (self.read_join_output_count_metadata(&d_output_count)? as u64).min(max_output as u64);
4620
4621 if result_count == 0 {
4622 let combined_schema = self.combine_schemas(left.schema(), right.schema());
4623 return self.create_empty_buffer(combined_schema);
4624 }
4625
4626 let output_rows = result_count as u32;
4627
4628 let gathered_left = self.gather_buffer_by_indices(left, &d_output_left, output_rows)?;
4630 let gathered_right = self.gather_buffer_by_indices(right, &d_output_right, output_rows)?;
4631
4632 let combined_schema = self.combine_schemas(left.schema(), right.schema());
4633 let mut result_columns = Vec::with_capacity(combined_schema.arity());
4634 result_columns.extend(gathered_left.columns);
4635 result_columns.extend(gathered_right.columns);
4636
4637 self.buffer_from_columns(result_columns, result_count, combined_schema)
4638 }
4639
4640 fn hash_join_inner_v2_indexed(
4641 &self,
4642 left: &CudaBuffer,
4643 right: &CudaBuffer,
4644 left_keys: &[usize],
4645 index: &JoinIndexV2,
4646 max_output: Option<usize>,
4647 ) -> Result<CudaBuffer> {
4648 let num_left = self.device_row_count(left)?;
4649 let num_right = self.device_row_count(right)?;
4650 if num_left > u32::MAX as usize || num_right > u32::MAX as usize {
4651 return Err(XlogError::Kernel(format!(
4652 "Join supports at most {} rows per side (left={}, right={})",
4653 u32::MAX,
4654 num_left,
4655 num_right
4656 )));
4657 }
4658
4659 if num_left == 0 || num_right == 0 {
4661 let combined_schema = self.combine_schemas(left.schema(), right.schema());
4662 return self.create_empty_buffer(combined_schema);
4663 }
4664
4665 let num_left = num_left as u32;
4666
4667 let left_packed = self.compute_hashes_and_pack_keys(left, left_keys)?;
4669 if left_packed.key_bytes != index.key_bytes {
4670 return Err(XlogError::Kernel(
4671 "Join key byte width mismatch between probe and cached index".to_string(),
4672 ));
4673 }
4674
4675 let table = &index.table;
4676
4677 let probe_func = self
4679 .device
4680 .inner()
4681 .get_func(JOIN_MODULE, join_kernels::HASH_JOIN_PROBE_V2)
4682 .ok_or_else(|| XlogError::Kernel("hash_join_probe_v2 kernel not found".to_string()))?;
4683
4684 let block_size = 256u32;
4685 let probe_grid = num_left.div_ceil(block_size);
4686 let probe_config = LaunchConfig {
4687 grid_dim: (probe_grid, 1, 1),
4688 block_dim: (block_size, 1, 1),
4689 shared_mem_bytes: 0,
4690 };
4691
4692 let mut d_count_only = self.memory.alloc::<u32>(1)?;
4693 self.device
4694 .inner()
4695 .memset_zeros(&mut d_count_only)
4696 .map_err(|e| XlogError::Kernel(format!("Failed to zero output count: {}", e)))?;
4697 self.device.synchronize()?;
4698 let d_dummy_left = self.memory.alloc::<u32>(1)?;
4699 let d_dummy_right = self.memory.alloc::<u32>(1)?;
4700 let max_output_count_only = 0u32;
4701
4702 unsafe {
4704 let mut params: Vec<*mut c_void> = vec![
4705 (&left_packed.hashes).as_kernel_param(),
4706 num_left.as_kernel_param(),
4707 (&table.bucket_offsets).as_kernel_param(),
4708 (&table.bucket_counts).as_kernel_param(),
4709 (&table.bucket_entries).as_kernel_param(),
4710 (&table.bucket_entry_hashes).as_kernel_param(),
4711 table.bucket_mask.as_kernel_param(),
4712 (&left_packed.packed_keys).as_kernel_param(),
4713 (&index.packed_keys).as_kernel_param(),
4714 index.key_bytes.as_kernel_param(),
4715 (&d_dummy_left).as_kernel_param(),
4716 (&d_dummy_right).as_kernel_param(),
4717 (&d_count_only).as_kernel_param(),
4718 max_output_count_only.as_kernel_param(),
4719 ];
4720 probe_func
4721 .clone()
4722 .launch(probe_config, &mut params)
4723 .map_err(|e| {
4724 XlogError::Kernel(format!("hash_join_probe_v2 (count) failed: {}", e))
4725 })?;
4726 }
4727
4728 self.device.synchronize()?;
4729
4730 let full_count = self.read_join_output_count_metadata(&d_count_only)? as u64;
4734 let requested = max_output
4735 .map(|limit| (limit as u64).min(full_count))
4736 .unwrap_or(full_count);
4737
4738 if requested == 0 {
4739 let combined_schema = self.combine_schemas(left.schema(), right.schema());
4740 return self.create_empty_buffer(combined_schema);
4741 }
4742
4743 if requested > u32::MAX as u64 {
4744 return Err(XlogError::Kernel(format!(
4745 "Join produced {} rows which exceeds the u32 index limit",
4746 requested
4747 )));
4748 }
4749
4750 let max_output = requested as u32;
4752 let d_output_left = self.memory.alloc::<u32>(max_output as usize)?;
4753 let d_output_right = self.memory.alloc::<u32>(max_output as usize)?;
4754 let mut d_output_count = self.memory.alloc::<u32>(1)?;
4755 self.device
4756 .inner()
4757 .memset_zeros(&mut d_output_count)
4758 .map_err(|e| XlogError::Kernel(format!("Failed to zero output count: {}", e)))?;
4759 self.device.synchronize()?;
4760
4761 unsafe {
4763 let mut params: Vec<*mut c_void> = vec![
4764 (&left_packed.hashes).as_kernel_param(),
4765 num_left.as_kernel_param(),
4766 (&table.bucket_offsets).as_kernel_param(),
4767 (&table.bucket_counts).as_kernel_param(),
4768 (&table.bucket_entries).as_kernel_param(),
4769 (&table.bucket_entry_hashes).as_kernel_param(),
4770 table.bucket_mask.as_kernel_param(),
4771 (&left_packed.packed_keys).as_kernel_param(),
4772 (&index.packed_keys).as_kernel_param(),
4773 index.key_bytes.as_kernel_param(),
4774 (&d_output_left).as_kernel_param(),
4775 (&d_output_right).as_kernel_param(),
4776 (&d_output_count).as_kernel_param(),
4777 max_output.as_kernel_param(),
4778 ];
4779 probe_func
4780 .clone()
4781 .launch(probe_config, &mut params)
4782 .map_err(|e| XlogError::Kernel(format!("hash_join_probe_v2 failed: {}", e)))?;
4783 }
4784
4785 self.device.synchronize()?;
4786
4787 let result_count =
4790 (self.read_join_output_count_metadata(&d_output_count)? as u64).min(max_output as u64);
4791
4792 if result_count == 0 {
4793 let combined_schema = self.combine_schemas(left.schema(), right.schema());
4794 return self.create_empty_buffer(combined_schema);
4795 }
4796
4797 let output_rows = result_count as u32;
4798
4799 let gathered_left = self.gather_buffer_by_indices(left, &d_output_left, output_rows)?;
4800 let gathered_right = self.gather_buffer_by_indices(right, &d_output_right, output_rows)?;
4801
4802 let combined_schema = self.combine_schemas(left.schema(), right.schema());
4803 let mut result_columns = Vec::with_capacity(combined_schema.arity());
4804 result_columns.extend(gathered_left.columns);
4805 result_columns.extend(gathered_right.columns);
4806
4807 self.buffer_from_columns(result_columns, result_count, combined_schema)
4808 }
4809
4810 fn hash_join_semi_impl(
4812 &self,
4813 left: &CudaBuffer,
4814 right: &CudaBuffer,
4815 left_keys: &[usize],
4816 right_keys: &[usize],
4817 ) -> Result<CudaBuffer> {
4818 let num_left = self.device_row_count(left)?;
4819 let num_right = self.device_row_count(right)?;
4820 if num_left > u32::MAX as usize || num_right > u32::MAX as usize {
4821 return Err(XlogError::Kernel(format!(
4822 "Join supports at most {} rows per side (left={}, right={})",
4823 u32::MAX,
4824 num_left,
4825 num_right
4826 )));
4827 }
4828
4829 if num_left == 0 {
4831 return self.create_empty_buffer(left.schema().clone());
4832 }
4833 if num_right == 0 {
4834 return self.create_empty_buffer(left.schema().clone());
4836 }
4837
4838 if left_keys.is_empty() || right_keys.is_empty() {
4840 return Err(XlogError::Kernel(
4841 "Join requires at least one key column".to_string(),
4842 ));
4843 }
4844 if left_keys.len() != right_keys.len() {
4845 return Err(XlogError::Kernel(
4846 "Left and right key columns must have same length".to_string(),
4847 ));
4848 }
4849
4850 for (&left_idx, &right_idx) in left_keys.iter().zip(right_keys.iter()) {
4852 let left_type = left.schema().column_type(left_idx);
4853 let right_type = right.schema().column_type(right_idx);
4854 if left_type != right_type {
4855 return Err(XlogError::Kernel(format!(
4856 "Key column type mismatch: left[{}]={:?}, right[{}]={:?}",
4857 left_idx, left_type, right_idx, right_type
4858 )));
4859 }
4860 }
4861
4862 let num_left = num_left as u32;
4863 let num_right = num_right as u32;
4864
4865 let left_packed = self.compute_hashes_and_pack_keys(left, left_keys)?;
4867 let right_packed = self.compute_hashes_and_pack_keys(right, right_keys)?;
4868
4869 let table = self.build_hash_table_v2(&right_packed.hashes, num_right)?;
4871
4872 let d_has_match = self.memory.alloc::<u8>(num_left as usize)?;
4874
4875 let semi_func = self
4877 .device
4878 .inner()
4879 .get_func(JOIN_MODULE, join_kernels::HASH_JOIN_SEMI)
4880 .ok_or_else(|| XlogError::Kernel("hash_join_semi kernel not found".to_string()))?;
4881
4882 let block_size = 256u32;
4883 let grid_size = num_left.div_ceil(block_size);
4884 let config = LaunchConfig {
4885 grid_dim: (grid_size, 1, 1),
4886 block_dim: (block_size, 1, 1),
4887 shared_mem_bytes: 0,
4888 };
4889
4890 unsafe {
4895 semi_func
4896 .clone()
4897 .launch(
4898 config,
4899 (
4900 &left_packed.hashes,
4901 num_left,
4902 &table.bucket_offsets,
4903 &table.bucket_counts,
4904 &table.bucket_entries,
4905 &table.bucket_entry_hashes,
4906 table.bucket_mask,
4907 &left_packed.packed_keys,
4908 &right_packed.packed_keys,
4909 left_packed.key_bytes,
4910 &d_has_match,
4911 ),
4912 )
4913 .map_err(|e| XlogError::Kernel(format!("hash_join_semi failed: {}", e)))?;
4914 }
4915
4916 self.device.synchronize()?;
4917 self.filter_by_device_mask(left, &d_has_match)
4918 }
4919
4920 pub fn membership_mask_device(
4925 &self,
4926 probe: &CudaBuffer,
4927 build: &CudaBuffer,
4928 probe_keys: &[usize],
4929 build_keys: &[usize],
4930 ) -> Result<TrackedCudaSlice<u8>> {
4931 let num_probe = self.device_row_count(probe)?;
4932 let num_build = self.device_row_count(build)?;
4933
4934 if num_probe == 0 {
4936 return self.memory.alloc::<u8>(0);
4937 }
4938
4939 if num_build == 0 {
4941 let mut d_mask = self.memory.alloc::<u8>(num_probe)?;
4942 self.device.inner().memset_zeros(&mut d_mask).map_err(|e| {
4943 XlogError::Kernel(format!(
4944 "Failed to zero membership mask for empty build: {}",
4945 e
4946 ))
4947 })?;
4948 return Ok(d_mask);
4949 }
4950
4951 if num_probe > u32::MAX as usize || num_build > u32::MAX as usize {
4952 return Err(XlogError::Kernel(format!(
4953 "membership_mask supports at most {} rows per side (probe={}, build={})",
4954 u32::MAX,
4955 num_probe,
4956 num_build
4957 )));
4958 }
4959
4960 if probe_keys.is_empty() || build_keys.is_empty() {
4962 return Err(XlogError::Kernel(
4963 "membership_mask requires at least one key column".to_string(),
4964 ));
4965 }
4966 if probe_keys.len() != build_keys.len() {
4967 return Err(XlogError::Kernel(
4968 "Probe and build key columns must have same length".to_string(),
4969 ));
4970 }
4971
4972 for (&p_idx, &b_idx) in probe_keys.iter().zip(build_keys.iter()) {
4974 let p_type = probe.schema().column_type(p_idx);
4975 let b_type = build.schema().column_type(b_idx);
4976 if p_type != b_type {
4977 return Err(XlogError::Kernel(format!(
4978 "Key column type mismatch: probe[{}]={:?}, build[{}]={:?}",
4979 p_idx, p_type, b_idx, b_type
4980 )));
4981 }
4982 }
4983
4984 let num_probe_u32 = num_probe as u32;
4985 let num_build_u32 = num_build as u32;
4986
4987 let probe_packed = self.compute_hashes_and_pack_keys(probe, probe_keys)?;
4989 let build_packed = self.compute_hashes_and_pack_keys(build, build_keys)?;
4990
4991 let table = self.build_hash_table_v2(&build_packed.hashes, num_build_u32)?;
4993
4994 let d_has_match = self.memory.alloc::<u8>(num_probe)?;
4996
4997 let semi_func = self
4999 .device
5000 .inner()
5001 .get_func(JOIN_MODULE, join_kernels::HASH_JOIN_SEMI)
5002 .ok_or_else(|| XlogError::Kernel("hash_join_semi kernel not found".to_string()))?;
5003
5004 let block_size = 256u32;
5005 let grid_size = num_probe_u32.div_ceil(block_size);
5006 let config = LaunchConfig {
5007 grid_dim: (grid_size, 1, 1),
5008 block_dim: (block_size, 1, 1),
5009 shared_mem_bytes: 0,
5010 };
5011
5012 unsafe {
5018 semi_func
5019 .clone()
5020 .launch(
5021 config,
5022 (
5023 &probe_packed.hashes,
5024 num_probe_u32,
5025 &table.bucket_offsets,
5026 &table.bucket_counts,
5027 &table.bucket_entries,
5028 &table.bucket_entry_hashes,
5029 table.bucket_mask,
5030 &probe_packed.packed_keys,
5031 &build_packed.packed_keys,
5032 probe_packed.key_bytes,
5033 &d_has_match,
5034 ),
5035 )
5036 .map_err(|e| XlogError::Kernel(format!("hash_join_semi failed: {}", e)))?;
5037 }
5038
5039 Ok(d_has_match)
5040 }
5041
5042 pub fn membership_mask(
5052 &self,
5053 probe: &CudaBuffer,
5054 build: &CudaBuffer,
5055 probe_keys: &[usize],
5056 build_keys: &[usize],
5057 ) -> Result<Vec<bool>> {
5058 let d_has_match = self.membership_mask_device(probe, build, probe_keys, build_keys)?;
5059 let num_probe = d_has_match.len();
5060 if num_probe == 0 {
5061 return Ok(Vec::new());
5062 }
5063 let mut host_mask = vec![0u8; num_probe];
5064 self.dtoh_sync_copy_into_tracked(&d_has_match, &mut host_mask)?;
5065 Ok(host_mask.into_iter().map(|b| b != 0).collect())
5066 }
5067
5068 fn hash_join_semi_indexed(
5069 &self,
5070 left: &CudaBuffer,
5071 left_keys: &[usize],
5072 index: &JoinIndexV2,
5073 ) -> Result<CudaBuffer> {
5074 let num_left = self.device_row_count(left)?;
5075 if num_left > u32::MAX as usize {
5076 return Err(XlogError::Kernel(format!(
5077 "Join supports at most {} rows on left side (left={})",
5078 u32::MAX,
5079 num_left
5080 )));
5081 }
5082
5083 if num_left == 0 {
5085 return self.create_empty_buffer(left.schema().clone());
5086 }
5087 if index.right_num_rows == 0 {
5088 return self.create_empty_buffer(left.schema().clone());
5089 }
5090
5091 let num_left = num_left as u32;
5092
5093 let left_packed = self.compute_hashes_and_pack_keys(left, left_keys)?;
5094 if left_packed.key_bytes != index.key_bytes {
5095 return Err(XlogError::Kernel(
5096 "Join key byte width mismatch between probe and cached index".to_string(),
5097 ));
5098 }
5099
5100 let table = &index.table;
5101
5102 let d_has_match = self.memory.alloc::<u8>(num_left as usize)?;
5104
5105 let semi_func = self
5106 .device
5107 .inner()
5108 .get_func(JOIN_MODULE, join_kernels::HASH_JOIN_SEMI)
5109 .ok_or_else(|| XlogError::Kernel("hash_join_semi kernel not found".to_string()))?;
5110
5111 let block_size = 256u32;
5112 let grid_size = num_left.div_ceil(block_size);
5113 let config = LaunchConfig {
5114 grid_dim: (grid_size, 1, 1),
5115 block_dim: (block_size, 1, 1),
5116 shared_mem_bytes: 0,
5117 };
5118
5119 unsafe {
5121 semi_func
5122 .clone()
5123 .launch(
5124 config,
5125 (
5126 &left_packed.hashes,
5127 num_left,
5128 &table.bucket_offsets,
5129 &table.bucket_counts,
5130 &table.bucket_entries,
5131 &table.bucket_entry_hashes,
5132 table.bucket_mask,
5133 &left_packed.packed_keys,
5134 &index.packed_keys,
5135 index.key_bytes,
5136 &d_has_match,
5137 ),
5138 )
5139 .map_err(|e| XlogError::Kernel(format!("hash_join_semi failed: {}", e)))?;
5140 }
5141
5142 self.device.synchronize()?;
5143 self.filter_by_device_mask(left, &d_has_match)
5144 }
5145
5146 fn hash_join_anti_impl(
5148 &self,
5149 left: &CudaBuffer,
5150 right: &CudaBuffer,
5151 left_keys: &[usize],
5152 right_keys: &[usize],
5153 ) -> Result<CudaBuffer> {
5154 let num_left = self.device_row_count(left)?;
5155 let num_right = self.device_row_count(right)?;
5156 if num_left > u32::MAX as usize || num_right > u32::MAX as usize {
5157 return Err(XlogError::Kernel(format!(
5158 "Join supports at most {} rows per side (left={}, right={})",
5159 u32::MAX,
5160 num_left,
5161 num_right
5162 )));
5163 }
5164
5165 if num_left == 0 {
5167 return self.create_empty_buffer(left.schema().clone());
5168 }
5169 if num_right == 0 {
5170 return self.clone_buffer(left);
5172 }
5173
5174 if left_keys.is_empty() || right_keys.is_empty() {
5176 return Err(XlogError::Kernel(
5177 "Join requires at least one key column".to_string(),
5178 ));
5179 }
5180 if left_keys.len() != right_keys.len() {
5181 return Err(XlogError::Kernel(
5182 "Left and right key columns must have same length".to_string(),
5183 ));
5184 }
5185
5186 for (&left_idx, &right_idx) in left_keys.iter().zip(right_keys.iter()) {
5188 let left_type = left.schema().column_type(left_idx);
5189 let right_type = right.schema().column_type(right_idx);
5190 if left_type != right_type {
5191 return Err(XlogError::Kernel(format!(
5192 "Key column type mismatch: left[{}]={:?}, right[{}]={:?}",
5193 left_idx, left_type, right_idx, right_type
5194 )));
5195 }
5196 }
5197
5198 let num_left = num_left as u32;
5199 let num_right = num_right as u32;
5200
5201 let left_packed = self.compute_hashes_and_pack_keys(left, left_keys)?;
5203 let right_packed = self.compute_hashes_and_pack_keys(right, right_keys)?;
5204
5205 let table = self.build_hash_table_v2(&right_packed.hashes, num_right)?;
5207
5208 let d_no_match = self.memory.alloc::<u8>(num_left as usize)?;
5210
5211 let anti_func = self
5213 .device
5214 .inner()
5215 .get_func(JOIN_MODULE, join_kernels::HASH_JOIN_ANTI)
5216 .ok_or_else(|| XlogError::Kernel("hash_join_anti kernel not found".to_string()))?;
5217
5218 let block_size = 256u32;
5219 let grid_size = num_left.div_ceil(block_size);
5220 let config = LaunchConfig {
5221 grid_dim: (grid_size, 1, 1),
5222 block_dim: (block_size, 1, 1),
5223 shared_mem_bytes: 0,
5224 };
5225
5226 unsafe {
5231 anti_func
5232 .clone()
5233 .launch(
5234 config,
5235 (
5236 &left_packed.hashes,
5237 num_left,
5238 &table.bucket_offsets,
5239 &table.bucket_counts,
5240 &table.bucket_entries,
5241 &table.bucket_entry_hashes,
5242 table.bucket_mask,
5243 &left_packed.packed_keys,
5244 &right_packed.packed_keys,
5245 left_packed.key_bytes,
5246 &d_no_match,
5247 ),
5248 )
5249 .map_err(|e| XlogError::Kernel(format!("hash_join_anti failed: {}", e)))?;
5250 }
5251
5252 self.device.synchronize()?;
5253 self.filter_by_device_mask(left, &d_no_match)
5254 }
5255
5256 fn hash_join_anti_indexed(
5257 &self,
5258 left: &CudaBuffer,
5259 right: &CudaBuffer,
5260 left_keys: &[usize],
5261 index: &JoinIndexV2,
5262 ) -> Result<CudaBuffer> {
5263 let num_left = self.device_row_count(left)?;
5264 let num_right = self.device_row_count(right)?;
5265 if num_left > u32::MAX as usize || num_right > u32::MAX as usize {
5266 return Err(XlogError::Kernel(format!(
5267 "Join supports at most {} rows per side (left={}, right={})",
5268 u32::MAX,
5269 num_left,
5270 num_right
5271 )));
5272 }
5273 if num_left == 0 {
5274 return self.create_empty_buffer(left.schema().clone());
5275 }
5276 if num_right == 0 {
5277 return self.clone_buffer(left);
5278 }
5279
5280 let num_left = num_left as u32;
5281
5282 let left_packed = self.compute_hashes_and_pack_keys(left, left_keys)?;
5283 if left_packed.key_bytes != index.key_bytes {
5284 return Err(XlogError::Kernel(
5285 "Join key byte width mismatch between probe and cached index".to_string(),
5286 ));
5287 }
5288
5289 let table = &index.table;
5290
5291 let d_no_match = self.memory.alloc::<u8>(num_left as usize)?;
5292
5293 let anti_func = self
5294 .device
5295 .inner()
5296 .get_func(JOIN_MODULE, join_kernels::HASH_JOIN_ANTI)
5297 .ok_or_else(|| XlogError::Kernel("hash_join_anti kernel not found".to_string()))?;
5298
5299 let block_size = 256u32;
5300 let grid_size = num_left.div_ceil(block_size);
5301 let config = LaunchConfig {
5302 grid_dim: (grid_size, 1, 1),
5303 block_dim: (block_size, 1, 1),
5304 shared_mem_bytes: 0,
5305 };
5306
5307 unsafe {
5309 anti_func
5310 .clone()
5311 .launch(
5312 config,
5313 (
5314 &left_packed.hashes,
5315 num_left,
5316 &table.bucket_offsets,
5317 &table.bucket_counts,
5318 &table.bucket_entries,
5319 &table.bucket_entry_hashes,
5320 table.bucket_mask,
5321 &left_packed.packed_keys,
5322 &index.packed_keys,
5323 index.key_bytes,
5324 &d_no_match,
5325 ),
5326 )
5327 .map_err(|e| XlogError::Kernel(format!("hash_join_anti failed: {}", e)))?;
5328 }
5329
5330 self.device.synchronize()?;
5331 self.filter_by_device_mask(left, &d_no_match)
5332 }
5333
5334 fn hash_join_left_outer_indexed(
5335 &self,
5336 left: &CudaBuffer,
5337 right: &CudaBuffer,
5338 left_keys: &[usize],
5339 index: &JoinIndexV2,
5340 max_output: Option<usize>,
5341 ) -> Result<CudaBuffer> {
5342 let num_left = self.device_row_count(left)?;
5343 let num_right = self.device_row_count(right)?;
5344 if num_left > u32::MAX as usize || num_right > u32::MAX as usize {
5345 return Err(XlogError::Kernel(format!(
5346 "Join supports at most {} rows per side (left={}, right={})",
5347 u32::MAX,
5348 num_left,
5349 num_right
5350 )));
5351 }
5352
5353 if num_left == 0 {
5355 let combined_schema = self.combine_schemas(left.schema(), right.schema());
5356 return self.create_empty_buffer(combined_schema);
5357 }
5358 if num_right == 0 {
5360 return self.left_outer_with_nulls(left, right);
5361 }
5362
5363 let num_left = num_left as u32;
5364
5365 let left_packed = self.compute_hashes_and_pack_keys(left, left_keys)?;
5366 if left_packed.key_bytes != index.key_bytes {
5367 return Err(XlogError::Kernel(
5368 "Join key byte width mismatch between probe and cached index".to_string(),
5369 ));
5370 }
5371
5372 let table = &index.table;
5373
5374 let d_has_match = self.memory.alloc::<u8>(num_left as usize)?;
5376
5377 let semi_func = self
5378 .device
5379 .inner()
5380 .get_func(JOIN_MODULE, join_kernels::HASH_JOIN_SEMI)
5381 .ok_or_else(|| XlogError::Kernel("hash_join_semi kernel not found".to_string()))?;
5382
5383 let block_size = 256u32;
5384 let grid_size = num_left.div_ceil(block_size);
5385 let config = LaunchConfig {
5386 grid_dim: (grid_size, 1, 1),
5387 block_dim: (block_size, 1, 1),
5388 shared_mem_bytes: 0,
5389 };
5390
5391 unsafe {
5393 semi_func
5394 .clone()
5395 .launch(
5396 config,
5397 (
5398 &left_packed.hashes,
5399 num_left,
5400 &table.bucket_offsets,
5401 &table.bucket_counts,
5402 &table.bucket_entries,
5403 &table.bucket_entry_hashes,
5404 table.bucket_mask,
5405 &left_packed.packed_keys,
5406 &index.packed_keys,
5407 index.key_bytes,
5408 &d_has_match,
5409 ),
5410 )
5411 .map_err(|e| XlogError::Kernel(format!("hash_join_semi failed: {}", e)))?;
5412 }
5413
5414 let probe_func = self
5415 .device
5416 .inner()
5417 .get_func(JOIN_MODULE, join_kernels::HASH_JOIN_PROBE_V2)
5418 .ok_or_else(|| XlogError::Kernel("hash_join_probe_v2 kernel not found".to_string()))?;
5419
5420 let mut d_count_only = self.memory.alloc::<u32>(1)?;
5422 self.device
5423 .inner()
5424 .memset_zeros(&mut d_count_only)
5425 .map_err(|e| XlogError::Kernel(format!("Failed to zero output count: {}", e)))?;
5426 let d_dummy_left = self.memory.alloc::<u32>(1)?;
5427 let d_dummy_right = self.memory.alloc::<u32>(1)?;
5428 let max_output_count_only = 0u32;
5429
5430 unsafe {
5432 let mut params: Vec<*mut c_void> = vec![
5433 (&left_packed.hashes).as_kernel_param(),
5434 num_left.as_kernel_param(),
5435 (&table.bucket_offsets).as_kernel_param(),
5436 (&table.bucket_counts).as_kernel_param(),
5437 (&table.bucket_entries).as_kernel_param(),
5438 (&table.bucket_entry_hashes).as_kernel_param(),
5439 table.bucket_mask.as_kernel_param(),
5440 (&left_packed.packed_keys).as_kernel_param(),
5441 (&index.packed_keys).as_kernel_param(),
5442 index.key_bytes.as_kernel_param(),
5443 (&d_dummy_left).as_kernel_param(),
5444 (&d_dummy_right).as_kernel_param(),
5445 (&d_count_only).as_kernel_param(),
5446 max_output_count_only.as_kernel_param(),
5447 ];
5448 probe_func
5449 .clone()
5450 .launch(config, &mut params)
5451 .map_err(|e| {
5452 XlogError::Kernel(format!("hash_join_probe_v2 (count) failed: {}", e))
5453 })?;
5454 }
5455
5456 self.device.synchronize()?;
5457
5458 let full_inner = self.read_join_output_count_metadata(&d_count_only)? as u64;
5461 let requested_inner = max_output
5462 .map(|limit| (limit as u64).min(full_inner))
5463 .unwrap_or(full_inner);
5464
5465 if requested_inner > u32::MAX as u64 {
5466 return Err(XlogError::Kernel(format!(
5467 "Join produced {} rows which exceeds the u32 index limit",
5468 requested_inner
5469 )));
5470 }
5471
5472 let max_output = requested_inner as u32;
5473 let alloc_len = (requested_inner.max(1)) as usize;
5474 let d_output_left = self.memory.alloc::<u32>(alloc_len)?;
5475 let d_output_right = self.memory.alloc::<u32>(alloc_len)?;
5476 let mut d_output_count = self.memory.alloc::<u32>(1)?;
5477 self.device
5478 .inner()
5479 .memset_zeros(&mut d_output_count)
5480 .map_err(|e| XlogError::Kernel(format!("Failed to zero output count: {}", e)))?;
5481
5482 unsafe {
5484 let mut params: Vec<*mut c_void> = vec![
5485 (&left_packed.hashes).as_kernel_param(),
5486 num_left.as_kernel_param(),
5487 (&table.bucket_offsets).as_kernel_param(),
5488 (&table.bucket_counts).as_kernel_param(),
5489 (&table.bucket_entries).as_kernel_param(),
5490 (&table.bucket_entry_hashes).as_kernel_param(),
5491 table.bucket_mask.as_kernel_param(),
5492 (&left_packed.packed_keys).as_kernel_param(),
5493 (&index.packed_keys).as_kernel_param(),
5494 index.key_bytes.as_kernel_param(),
5495 (&d_output_left).as_kernel_param(),
5496 (&d_output_right).as_kernel_param(),
5497 (&d_output_count).as_kernel_param(),
5498 max_output.as_kernel_param(),
5499 ];
5500 probe_func
5501 .clone()
5502 .launch(config, &mut params)
5503 .map_err(|e| XlogError::Kernel(format!("hash_join_probe_v2 failed: {}", e)))?;
5504 }
5505
5506 self.device.synchronize()?;
5507
5508 let device = self.device.inner();
5509
5510 let inner_count = self
5515 .read_join_output_count_metadata(&d_output_count)?
5516 .min(max_output);
5517
5518 let mask_not_fn = device
5519 .get_func(FILTER_MODULE, filter_kernels::MASK_NOT)
5520 .ok_or_else(|| XlogError::Kernel("mask_not kernel not found".to_string()))?;
5521
5522 let mut d_no_match = self.memory.alloc::<u8>(num_left as usize)?;
5523
5524 unsafe {
5526 mask_not_fn
5527 .clone()
5528 .launch(config, (&d_has_match, &mut d_no_match, num_left))
5529 }
5530 .map_err(|e| XlogError::Kernel(format!("mask_not failed: {}", e)))?;
5531
5532 let unmatched_left = self.filter_by_device_mask(left, &d_no_match)?;
5533
5534 let unmatched_rows = self.device_row_count(&unmatched_left)? as u64;
5535 let total_rows = (inner_count as u64) + unmatched_rows;
5536
5537 let combined_schema = self.combine_schemas(left.schema(), right.schema());
5538
5539 if total_rows == 0 {
5540 return self.create_empty_buffer(combined_schema);
5541 }
5542
5543 let inner_left = self.gather_buffer_by_indices(left, &d_output_left, inner_count)?;
5544 let inner_right = self.gather_buffer_by_indices(right, &d_output_right, inner_count)?;
5545
5546 if unmatched_rows == 0 {
5547 let mut result_columns = Vec::with_capacity(combined_schema.arity());
5548 result_columns.extend(inner_left.columns);
5549 result_columns.extend(inner_right.columns);
5550 return self.buffer_from_columns(result_columns, inner_count as u64, combined_schema);
5551 }
5552
5553 if inner_count == 0 {
5554 let mut result_columns = Vec::with_capacity(combined_schema.arity());
5555 result_columns.extend(unmatched_left.columns);
5556
5557 for col_idx in 0..right.arity() {
5558 let elem_size = right
5559 .schema()
5560 .column_type(col_idx)
5561 .map(|t| t.size_bytes())
5562 .unwrap_or(4);
5563
5564 let bytes = (unmatched_rows as usize)
5565 .checked_mul(elem_size)
5566 .ok_or_else(|| {
5567 XlogError::Kernel(
5568 "Left outer join: right column byte size overflow".to_string(),
5569 )
5570 })?;
5571
5572 let mut dst_col = self.memory.alloc::<u8>(bytes)?;
5573 if bytes > 0 {
5574 device.memset_zeros(&mut dst_col).map_err(|e| {
5575 XlogError::Kernel(format!("Failed to zero null right column: {}", e))
5576 })?;
5577 }
5578 result_columns.push(dst_col.into());
5579 }
5580
5581 self.device.synchronize()?;
5582 return self.buffer_from_columns(result_columns, unmatched_rows, combined_schema);
5583 }
5584
5585 let mut result_columns = Vec::with_capacity(combined_schema.arity());
5586 let inner_rows = inner_count as u64;
5587
5588 for (col_idx, (inner_col, unmatched_col)) in inner_left
5589 .columns
5590 .into_iter()
5591 .zip(unmatched_left.columns)
5592 .enumerate()
5593 {
5594 let elem_size = left
5595 .schema()
5596 .column_type(col_idx)
5597 .map(|t| t.size_bytes())
5598 .unwrap_or(4);
5599
5600 let inner_bytes = (inner_rows as usize)
5601 .checked_mul(elem_size)
5602 .ok_or_else(|| {
5603 XlogError::Kernel("Left outer join: inner_bytes overflow".to_string())
5604 })?;
5605 let unmatched_bytes = (unmatched_rows as usize)
5606 .checked_mul(elem_size)
5607 .ok_or_else(|| {
5608 XlogError::Kernel("Left outer join: unmatched_bytes overflow".to_string())
5609 })?;
5610 let total_bytes = inner_bytes.checked_add(unmatched_bytes).ok_or_else(|| {
5611 XlogError::Kernel("Left outer join: total_bytes overflow".to_string())
5612 })?;
5613
5614 let mut out_col = self.memory.alloc::<u8>(total_bytes)?;
5615
5616 if inner_bytes > 0 {
5617 let mut out_view = out_col.slice_mut(0..inner_bytes);
5618 device.dtod_copy(&inner_col, &mut out_view).map_err(|e| {
5619 XlogError::Kernel(format!("Failed to copy inner left column: {}", e))
5620 })?;
5621 }
5622 if unmatched_bytes > 0 {
5623 let mut out_view = out_col.slice_mut(inner_bytes..total_bytes);
5624 let unmatched_view = self.column_bytes_view(&unmatched_col, unmatched_bytes)?;
5625 device
5626 .dtod_copy(&unmatched_view, &mut out_view)
5627 .map_err(|e| {
5628 XlogError::Kernel(format!("Failed to copy unmatched left column: {}", e))
5629 })?;
5630 }
5631
5632 result_columns.push(out_col.into());
5633 }
5634
5635 for (col_idx, inner_col) in inner_right.columns.into_iter().enumerate() {
5636 let elem_size = right
5637 .schema()
5638 .column_type(col_idx)
5639 .map(|t| t.size_bytes())
5640 .unwrap_or(4);
5641
5642 let inner_bytes = (inner_rows as usize)
5643 .checked_mul(elem_size)
5644 .ok_or_else(|| {
5645 XlogError::Kernel("Left outer join: inner_bytes overflow".to_string())
5646 })?;
5647 let unmatched_bytes = (unmatched_rows as usize)
5648 .checked_mul(elem_size)
5649 .ok_or_else(|| {
5650 XlogError::Kernel("Left outer join: unmatched_bytes overflow".to_string())
5651 })?;
5652 let total_bytes = inner_bytes.checked_add(unmatched_bytes).ok_or_else(|| {
5653 XlogError::Kernel("Left outer join: total_bytes overflow".to_string())
5654 })?;
5655
5656 let mut out_col = self.memory.alloc::<u8>(total_bytes)?;
5657
5658 if total_bytes > 0 {
5659 device.memset_zeros(&mut out_col).map_err(|e| {
5660 XlogError::Kernel(format!("Failed to zero right outer column: {}", e))
5661 })?;
5662 }
5663
5664 if inner_bytes > 0 {
5665 let mut out_view = out_col.slice_mut(0..inner_bytes);
5666 device.dtod_copy(&inner_col, &mut out_view).map_err(|e| {
5667 XlogError::Kernel(format!("Failed to copy inner right column: {}", e))
5668 })?;
5669 }
5670
5671 result_columns.push(out_col.into());
5672 }
5673
5674 self.device.synchronize()?;
5675
5676 self.buffer_from_columns(result_columns, total_rows, combined_schema)
5677 }
5678
5679 fn hash_join_left_outer_impl(
5681 &self,
5682 left: &CudaBuffer,
5683 right: &CudaBuffer,
5684 left_keys: &[usize],
5685 right_keys: &[usize],
5686 max_output: Option<usize>,
5687 ) -> Result<CudaBuffer> {
5688 let num_left = self.device_row_count(left)?;
5689 let num_right = self.device_row_count(right)?;
5690 if num_left > u32::MAX as usize || num_right > u32::MAX as usize {
5691 return Err(XlogError::Kernel(format!(
5692 "Join supports at most {} rows per side (left={}, right={})",
5693 u32::MAX,
5694 num_left,
5695 num_right
5696 )));
5697 }
5698
5699 if num_left == 0 {
5701 let combined_schema = self.combine_schemas(left.schema(), right.schema());
5702 return self.create_empty_buffer(combined_schema);
5703 }
5704
5705 if num_right == 0 {
5707 return self.left_outer_with_nulls(left, right);
5708 }
5709
5710 if left_keys.is_empty() || right_keys.is_empty() {
5712 return Err(XlogError::Kernel(
5713 "Join requires at least one key column".to_string(),
5714 ));
5715 }
5716 if left_keys.len() != right_keys.len() {
5717 return Err(XlogError::Kernel(
5718 "Left and right key columns must have same length".to_string(),
5719 ));
5720 }
5721
5722 for (&left_idx, &right_idx) in left_keys.iter().zip(right_keys.iter()) {
5724 let left_type = left.schema().column_type(left_idx);
5725 let right_type = right.schema().column_type(right_idx);
5726 if left_type != right_type {
5727 return Err(XlogError::Kernel(format!(
5728 "Key column type mismatch: left[{}]={:?}, right[{}]={:?}",
5729 left_idx, left_type, right_idx, right_type
5730 )));
5731 }
5732 }
5733
5734 let num_left = num_left as u32;
5735 let num_right = num_right as u32;
5736
5737 let left_packed = self.compute_hashes_and_pack_keys(left, left_keys)?;
5739 let right_packed = self.compute_hashes_and_pack_keys(right, right_keys)?;
5740
5741 let table = self.build_hash_table_v2(&right_packed.hashes, num_right)?;
5743
5744 let d_has_match = self.memory.alloc::<u8>(num_left as usize)?;
5746
5747 let semi_func = self
5749 .device
5750 .inner()
5751 .get_func(JOIN_MODULE, join_kernels::HASH_JOIN_SEMI)
5752 .ok_or_else(|| XlogError::Kernel("hash_join_semi kernel not found".to_string()))?;
5753
5754 let block_size = 256u32;
5755 let grid_size = num_left.div_ceil(block_size);
5756 let config = LaunchConfig {
5757 grid_dim: (grid_size, 1, 1),
5758 block_dim: (block_size, 1, 1),
5759 shared_mem_bytes: 0,
5760 };
5761
5762 unsafe {
5767 semi_func
5768 .clone()
5769 .launch(
5770 config,
5771 (
5772 &left_packed.hashes,
5773 num_left,
5774 &table.bucket_offsets,
5775 &table.bucket_counts,
5776 &table.bucket_entries,
5777 &table.bucket_entry_hashes,
5778 table.bucket_mask,
5779 &left_packed.packed_keys,
5780 &right_packed.packed_keys,
5781 left_packed.key_bytes,
5782 &d_has_match,
5783 ),
5784 )
5785 .map_err(|e| XlogError::Kernel(format!("hash_join_semi failed: {}", e)))?;
5786 }
5787
5788 let probe_func = self
5789 .device
5790 .inner()
5791 .get_func(JOIN_MODULE, join_kernels::HASH_JOIN_PROBE_V2)
5792 .ok_or_else(|| XlogError::Kernel("hash_join_probe_v2 kernel not found".to_string()))?;
5793
5794 let mut d_count_only = self.memory.alloc::<u32>(1)?;
5796 self.device
5797 .inner()
5798 .memset_zeros(&mut d_count_only)
5799 .map_err(|e| XlogError::Kernel(format!("Failed to zero output count: {}", e)))?;
5800 let d_dummy_left = self.memory.alloc::<u32>(1)?;
5801 let d_dummy_right = self.memory.alloc::<u32>(1)?;
5802 let max_output_count_only = 0u32;
5803
5804 unsafe {
5806 let mut params: Vec<*mut c_void> = vec![
5807 (&left_packed.hashes).as_kernel_param(),
5808 num_left.as_kernel_param(),
5809 (&table.bucket_offsets).as_kernel_param(),
5810 (&table.bucket_counts).as_kernel_param(),
5811 (&table.bucket_entries).as_kernel_param(),
5812 (&table.bucket_entry_hashes).as_kernel_param(),
5813 table.bucket_mask.as_kernel_param(),
5814 (&left_packed.packed_keys).as_kernel_param(),
5815 (&right_packed.packed_keys).as_kernel_param(),
5816 left_packed.key_bytes.as_kernel_param(),
5817 (&d_dummy_left).as_kernel_param(),
5818 (&d_dummy_right).as_kernel_param(),
5819 (&d_count_only).as_kernel_param(),
5820 max_output_count_only.as_kernel_param(),
5821 ];
5822 probe_func
5823 .clone()
5824 .launch(config, &mut params)
5825 .map_err(|e| {
5826 XlogError::Kernel(format!("hash_join_probe_v2 (count) failed: {}", e))
5827 })?;
5828 }
5829
5830 self.device.synchronize()?;
5831
5832 let full_inner = self.read_join_output_count_metadata(&d_count_only)? as u64;
5835 let requested_inner = max_output
5836 .map(|limit| (limit as u64).min(full_inner))
5837 .unwrap_or(full_inner);
5838
5839 if requested_inner > u32::MAX as u64 {
5840 return Err(XlogError::Kernel(format!(
5841 "Join produced {} rows which exceeds the u32 index limit",
5842 requested_inner
5843 )));
5844 }
5845
5846 let max_output = requested_inner as u32;
5847 let alloc_len = (requested_inner.max(1)) as usize;
5848 let d_output_left = self.memory.alloc::<u32>(alloc_len)?;
5849 let d_output_right = self.memory.alloc::<u32>(alloc_len)?;
5850 let mut d_output_count = self.memory.alloc::<u32>(1)?;
5851 self.device
5852 .inner()
5853 .memset_zeros(&mut d_output_count)
5854 .map_err(|e| XlogError::Kernel(format!("Failed to zero output count: {}", e)))?;
5855
5856 unsafe {
5863 let mut params: Vec<*mut c_void> = vec![
5864 (&left_packed.hashes).as_kernel_param(),
5865 num_left.as_kernel_param(),
5866 (&table.bucket_offsets).as_kernel_param(),
5867 (&table.bucket_counts).as_kernel_param(),
5868 (&table.bucket_entries).as_kernel_param(),
5869 (&table.bucket_entry_hashes).as_kernel_param(),
5870 table.bucket_mask.as_kernel_param(),
5871 (&left_packed.packed_keys).as_kernel_param(),
5872 (&right_packed.packed_keys).as_kernel_param(),
5873 left_packed.key_bytes.as_kernel_param(),
5874 (&d_output_left).as_kernel_param(),
5875 (&d_output_right).as_kernel_param(),
5876 (&d_output_count).as_kernel_param(),
5877 max_output.as_kernel_param(),
5878 ];
5879 probe_func
5880 .clone()
5881 .launch(config, &mut params)
5882 .map_err(|e| XlogError::Kernel(format!("hash_join_probe_v2 failed: {}", e)))?;
5883 }
5884
5885 self.device.synchronize()?;
5886
5887 let device = self.device.inner();
5888
5889 let inner_count = self
5894 .read_join_output_count_metadata(&d_output_count)?
5895 .min(max_output);
5896
5897 let mask_not_fn = device
5899 .get_func(FILTER_MODULE, filter_kernels::MASK_NOT)
5900 .ok_or_else(|| XlogError::Kernel("mask_not kernel not found".to_string()))?;
5901
5902 let mut d_no_match = self.memory.alloc::<u8>(num_left as usize)?;
5903
5904 unsafe {
5906 mask_not_fn
5907 .clone()
5908 .launch(config, (&d_has_match, &mut d_no_match, num_left))
5909 }
5910 .map_err(|e| XlogError::Kernel(format!("mask_not failed: {}", e)))?;
5911
5912 let unmatched_left = self.filter_by_device_mask(left, &d_no_match)?;
5913
5914 let unmatched_rows = self.device_row_count(&unmatched_left)? as u64;
5915 let total_rows = (inner_count as u64) + unmatched_rows;
5916
5917 let combined_schema = self.combine_schemas(left.schema(), right.schema());
5918
5919 if total_rows == 0 {
5920 return self.create_empty_buffer(combined_schema);
5921 }
5922
5923 let inner_left = self.gather_buffer_by_indices(left, &d_output_left, inner_count)?;
5925 let inner_right = self.gather_buffer_by_indices(right, &d_output_right, inner_count)?;
5926
5927 if unmatched_rows == 0 {
5928 let mut result_columns = Vec::with_capacity(combined_schema.arity());
5929 result_columns.extend(inner_left.columns);
5930 result_columns.extend(inner_right.columns);
5931 return self.buffer_from_columns(result_columns, inner_count as u64, combined_schema);
5932 }
5933
5934 if inner_count == 0 {
5935 let mut result_columns = Vec::with_capacity(combined_schema.arity());
5936 result_columns.extend(unmatched_left.columns);
5937
5938 for col_idx in 0..right.arity() {
5939 let elem_size = right
5940 .schema()
5941 .column_type(col_idx)
5942 .map(|t| t.size_bytes())
5943 .unwrap_or(4);
5944
5945 let bytes = (unmatched_rows as usize)
5946 .checked_mul(elem_size)
5947 .ok_or_else(|| {
5948 XlogError::Kernel(
5949 "Left outer join: right column byte size overflow".to_string(),
5950 )
5951 })?;
5952
5953 let mut dst_col = self.memory.alloc::<u8>(bytes)?;
5954 if bytes > 0 {
5955 device.memset_zeros(&mut dst_col).map_err(|e| {
5956 XlogError::Kernel(format!("Failed to zero null right column: {}", e))
5957 })?;
5958 }
5959 result_columns.push(dst_col.into());
5960 }
5961
5962 self.device.synchronize()?;
5963 return self.buffer_from_columns(result_columns, unmatched_rows, combined_schema);
5964 }
5965
5966 let mut result_columns = Vec::with_capacity(combined_schema.arity());
5968 let inner_rows = inner_count as u64;
5969
5970 for (col_idx, (inner_col, unmatched_col)) in inner_left
5972 .columns
5973 .into_iter()
5974 .zip(unmatched_left.columns)
5975 .enumerate()
5976 {
5977 let elem_size = left
5978 .schema()
5979 .column_type(col_idx)
5980 .map(|t| t.size_bytes())
5981 .unwrap_or(4);
5982
5983 let inner_bytes = (inner_rows as usize)
5984 .checked_mul(elem_size)
5985 .ok_or_else(|| {
5986 XlogError::Kernel("Left outer join: inner_bytes overflow".to_string())
5987 })?;
5988 let unmatched_bytes = (unmatched_rows as usize)
5989 .checked_mul(elem_size)
5990 .ok_or_else(|| {
5991 XlogError::Kernel("Left outer join: unmatched_bytes overflow".to_string())
5992 })?;
5993 let total_bytes = inner_bytes.checked_add(unmatched_bytes).ok_or_else(|| {
5994 XlogError::Kernel("Left outer join: total_bytes overflow".to_string())
5995 })?;
5996
5997 let mut out_col = self.memory.alloc::<u8>(total_bytes)?;
5998
5999 if inner_bytes > 0 {
6000 let mut out_view = out_col.slice_mut(0..inner_bytes);
6001 device.dtod_copy(&inner_col, &mut out_view).map_err(|e| {
6002 XlogError::Kernel(format!("Failed to copy inner left column: {}", e))
6003 })?;
6004 }
6005 if unmatched_bytes > 0 {
6006 let mut out_view = out_col.slice_mut(inner_bytes..total_bytes);
6007 let unmatched_view = self.column_bytes_view(&unmatched_col, unmatched_bytes)?;
6008 device
6009 .dtod_copy(&unmatched_view, &mut out_view)
6010 .map_err(|e| {
6011 XlogError::Kernel(format!("Failed to copy unmatched left column: {}", e))
6012 })?;
6013 }
6014
6015 result_columns.push(out_col.into());
6016 }
6017
6018 for (col_idx, inner_col) in inner_right.columns.into_iter().enumerate() {
6020 let elem_size = right
6021 .schema()
6022 .column_type(col_idx)
6023 .map(|t| t.size_bytes())
6024 .unwrap_or(4);
6025
6026 let inner_bytes = (inner_rows as usize)
6027 .checked_mul(elem_size)
6028 .ok_or_else(|| {
6029 XlogError::Kernel("Left outer join: inner_bytes overflow".to_string())
6030 })?;
6031 let unmatched_bytes = (unmatched_rows as usize)
6032 .checked_mul(elem_size)
6033 .ok_or_else(|| {
6034 XlogError::Kernel("Left outer join: unmatched_bytes overflow".to_string())
6035 })?;
6036 let total_bytes = inner_bytes.checked_add(unmatched_bytes).ok_or_else(|| {
6037 XlogError::Kernel("Left outer join: total_bytes overflow".to_string())
6038 })?;
6039
6040 let mut out_col = self.memory.alloc::<u8>(total_bytes)?;
6041
6042 if total_bytes > 0 {
6043 device.memset_zeros(&mut out_col).map_err(|e| {
6044 XlogError::Kernel(format!("Failed to zero right outer column: {}", e))
6045 })?;
6046 }
6047
6048 if inner_bytes > 0 {
6049 let mut out_view = out_col.slice_mut(0..inner_bytes);
6050 device.dtod_copy(&inner_col, &mut out_view).map_err(|e| {
6051 XlogError::Kernel(format!("Failed to copy inner right column: {}", e))
6052 })?;
6053 }
6054
6055 result_columns.push(out_col.into());
6056 }
6057
6058 self.device.synchronize()?;
6059
6060 self.buffer_from_columns(result_columns, total_rows, combined_schema)
6061 }
6062
6063 fn left_outer_with_nulls(&self, left: &CudaBuffer, right: &CudaBuffer) -> Result<CudaBuffer> {
6065 let combined_schema = self.combine_schemas(left.schema(), right.schema());
6066 let num_rows = self.device_row_count(left)? as u64;
6067 if num_rows == 0 {
6068 return self.create_empty_buffer(combined_schema);
6069 }
6070 let device = self.device.inner();
6071
6072 let mut result_columns = Vec::with_capacity(combined_schema.arity());
6073
6074 for col_idx in 0..left.arity() {
6076 let col = left
6077 .column(col_idx)
6078 .ok_or_else(|| XlogError::Kernel(format!("Left column {} not found", col_idx)))?;
6079
6080 let elem_size = left
6081 .schema()
6082 .column_type(col_idx)
6083 .map(|t| t.size_bytes())
6084 .unwrap_or(4);
6085
6086 let bytes = (num_rows as usize) * elem_size;
6087 let mut dst_col = self.memory.alloc::<u8>(bytes)?;
6088 if bytes > 0 {
6089 let src_view = self.column_bytes_view(col, bytes)?;
6090 device
6091 .dtod_copy(&src_view, &mut dst_col)
6092 .map_err(|e| XlogError::Kernel(format!("Failed to copy left column: {}", e)))?;
6093 }
6094
6095 result_columns.push(dst_col.into());
6096 }
6097
6098 for col_idx in 0..right.arity() {
6100 let elem_size = right
6101 .schema()
6102 .column_type(col_idx)
6103 .map(|t| t.size_bytes())
6104 .unwrap_or(4);
6105
6106 let bytes = (num_rows as usize) * elem_size;
6107 let mut dst_col = self.memory.alloc::<u8>(bytes)?;
6108 if bytes > 0 {
6109 device
6110 .memset_zeros(&mut dst_col)
6111 .map_err(|e| XlogError::Kernel(format!("Failed to zero null column: {}", e)))?;
6112 }
6113
6114 result_columns.push(dst_col.into());
6115 }
6116
6117 self.device.synchronize()?;
6118
6119 self.buffer_from_columns(result_columns, num_rows, combined_schema)
6120 }
6121
6122 pub fn clone_buffer(&self, buffer: &CudaBuffer) -> Result<CudaBuffer> {
6127 let verify = {
6132 static ENABLED: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6133 *ENABLED.get_or_init(|| {
6134 std::env::var("XLOG_DEBUG_VERIFY_CLONES").map(|v| v == "1") == Ok(true)
6135 })
6136 };
6137
6138 let mut result_columns = Vec::with_capacity(buffer.arity());
6139 let device = self.device.inner();
6140
6141 for col_idx in 0..buffer.arity() {
6142 let src_col = buffer
6143 .column(col_idx)
6144 .ok_or_else(|| XlogError::Kernel(format!("Column {} not found", col_idx)))?;
6145 let mut dst_col = self.memory.alloc::<u8>(src_col.len())?;
6146 if !src_col.is_empty() {
6147 device
6148 .dtod_copy(src_col, &mut dst_col)
6149 .map_err(|e| XlogError::Kernel(format!("Failed to clone column: {}", e)))?;
6150 }
6151 if verify && !src_col.is_empty() {
6152 self.device.synchronize()?;
6153 let mut src_host = vec![0u8; src_col.len()];
6154 let mut dst_host = vec![0u8; dst_col.len()];
6155 device
6156 .dtoh_sync_copy_into(src_col, &mut src_host)
6157 .map_err(|e| XlogError::Kernel(format!("verify src dtoh: {}", e)))?;
6158 device
6159 .dtoh_sync_copy_into(&dst_col, &mut dst_host)
6160 .map_err(|e| XlogError::Kernel(format!("verify dst dtoh: {}", e)))?;
6161 if src_host != dst_host {
6162 let first_diff = src_host
6163 .iter()
6164 .zip(dst_host.iter())
6165 .position(|(a, b)| a != b)
6166 .unwrap_or(0);
6167 return Err(XlogError::Kernel(format!(
6168 "CLONE VERIFY FAILED: column {} differs from source at byte {} of {} (clone is wrong at birth)",
6169 col_idx,
6170 first_diff,
6171 src_col.len(),
6172 )));
6173 }
6174 }
6175 result_columns.push(dst_col.into());
6176 }
6177
6178 let mut d_num_rows = self.memory.alloc::<u32>(1)?;
6179 device
6180 .dtod_copy(buffer.num_rows_device(), &mut d_num_rows)
6181 .map_err(|e| XlogError::Kernel(format!("Failed to clone row count: {}", e)))?;
6182
6183 let mut cloned = CudaBuffer::from_columns(
6184 result_columns,
6185 buffer.row_cap,
6186 d_num_rows,
6187 buffer.schema().clone(),
6188 );
6189 if let Some(cached) = buffer.cached_row_count() {
6192 cloned.set_cached_row_count_if_unset(cached);
6193 }
6194 if buffer.canonical_full_row_set_certified() {
6195 cloned.certify_canonical_full_row_set();
6196 }
6197 Ok(cloned)
6198 }
6199 pub fn extract_column(&self, buffer: &CudaBuffer, col_idx: usize) -> Result<CudaBuffer> {
6210 if buffer.is_empty() {
6211 let col_type = buffer
6212 .schema()
6213 .column_type(col_idx)
6214 .ok_or_else(|| XlogError::Kernel(format!("Column {} not found", col_idx)))?;
6215 let schema = Schema::new(vec![("col".to_string(), col_type)]);
6216 return self.create_empty_buffer(schema);
6217 }
6218
6219 let col_type = buffer
6220 .schema()
6221 .column_type(col_idx)
6222 .ok_or_else(|| XlogError::Kernel(format!("Column {} not found", col_idx)))?;
6223 let src_col = buffer
6224 .column(col_idx)
6225 .ok_or_else(|| XlogError::Kernel(format!("Column {} not found in buffer", col_idx)))?;
6226 let mut dst_col = self.memory.alloc::<u8>(src_col.len())?;
6227 let device = self.device.inner();
6228 if !src_col.is_empty() {
6229 device
6230 .dtod_copy(src_col, &mut dst_col)
6231 .map_err(|e| XlogError::Kernel(format!("Failed to copy column: {}", e)))?;
6232 }
6233
6234 let mut d_num_rows = self.memory.alloc::<u32>(1)?;
6235 device
6236 .dtod_copy(buffer.num_rows_device(), &mut d_num_rows)
6237 .map_err(|e| XlogError::Kernel(format!("Failed to copy row count: {}", e)))?;
6238 self.device.synchronize()?;
6239
6240 let schema = Schema::new(vec![("col".to_string(), col_type)]);
6241 Ok(CudaBuffer::from_columns(
6242 vec![dst_col.into()],
6243 buffer.row_cap,
6244 d_num_rows,
6245 schema,
6246 ))
6247 }
6248
6249 pub fn extract_active_rule_indices(
6252 &self,
6253 mask_hard: &CudaBuffer,
6254 mask_soft: &CudaBuffer,
6255 n: usize,
6256 max_active: usize,
6257 ) -> Result<Vec<(u32, u32, u32)>> {
6258 let total = n * n * n;
6259 let block_size = 256usize;
6260 let grid_size = total.div_ceil(block_size);
6261
6262 let mut out_i = self.memory().alloc::<u32>(total)?;
6263 let mut out_j = self.memory().alloc::<u32>(total)?;
6264 let mut out_k = self.memory().alloc::<u32>(total)?;
6265 let mut out_p = self.memory().alloc::<f32>(total)?;
6266 let mut count = self.memory().alloc::<u32>(1)?;
6267
6268 self.htod_launch_metadata_sync_copy_into(&[0u32], &mut count)
6269 .map_err(|e| XlogError::Kernel(format!("ILP htod count: {}", e)))?;
6270
6271 let hard_col = mask_hard
6272 .column(0)
6273 .ok_or_else(|| XlogError::Kernel("ILP hard mask has no column".into()))?;
6274 let soft_col = mask_soft
6275 .column(0)
6276 .ok_or_else(|| XlogError::Kernel("ILP soft mask has no column".into()))?;
6277
6278 let kernel = self
6279 .device()
6280 .inner()
6281 .get_func(ILP_MODULE, ilp_kernels::EXTRACT_NONZERO_INDICES)
6282 .ok_or_else(|| XlogError::Kernel("extract_nonzero_indices kernel not found".into()))?;
6283
6284 let hard_bytes = total * std::mem::size_of::<f32>();
6285 let soft_bytes = total * std::mem::size_of::<f32>();
6286 let hard_view = self.column_bytes_view(hard_col, hard_bytes)?;
6287 let soft_view = self.column_bytes_view(soft_col, soft_bytes)?;
6288
6289 unsafe {
6291 kernel
6292 .clone()
6293 .launch(
6294 cudarc::driver::LaunchConfig {
6295 grid_dim: (grid_size as u32, 1, 1),
6296 block_dim: (block_size as u32, 1, 1),
6297 shared_mem_bytes: 0,
6298 },
6299 (
6300 &hard_view, &soft_view, n as u32, &mut out_i, &mut out_j, &mut out_k,
6301 &mut out_p, &mut count,
6302 ),
6303 )
6304 .map_err(|e| {
6305 XlogError::Kernel(format!("Failed to launch extract_nonzero_indices: {}", e))
6306 })?;
6307 }
6308
6309 let mut count_host = [0u32];
6310 self.device()
6311 .inner()
6312 .dtoh_sync_copy_into(&count, &mut count_host)
6313 .map_err(|e| XlogError::Kernel(format!("ILP dtoh count: {}", e)))?;
6314 let active_count = count_host[0] as usize;
6315
6316 if active_count == 0 {
6317 return Ok(Vec::new());
6318 }
6319
6320 let mut i_host = vec![0u32; active_count];
6321 let mut j_host = vec![0u32; active_count];
6322 let mut k_host = vec![0u32; active_count];
6323 let mut p_host = vec![0f32; active_count];
6324
6325 let out_i_view = out_i
6326 .try_slice(0..active_count)
6327 .ok_or_else(|| XlogError::Kernel("ILP slice i out of bounds".into()))?;
6328 let out_j_view = out_j
6329 .try_slice(0..active_count)
6330 .ok_or_else(|| XlogError::Kernel("ILP slice j out of bounds".into()))?;
6331 let out_k_view = out_k
6332 .try_slice(0..active_count)
6333 .ok_or_else(|| XlogError::Kernel("ILP slice k out of bounds".into()))?;
6334 let out_p_view = out_p
6335 .try_slice(0..active_count)
6336 .ok_or_else(|| XlogError::Kernel("ILP slice p out of bounds".into()))?;
6337
6338 self.device()
6339 .inner()
6340 .dtoh_sync_copy_into(&out_i_view, &mut i_host)
6341 .map_err(|e| XlogError::Kernel(format!("ILP dtoh i: {}", e)))?;
6342 self.device()
6343 .inner()
6344 .dtoh_sync_copy_into(&out_j_view, &mut j_host)
6345 .map_err(|e| XlogError::Kernel(format!("ILP dtoh j: {}", e)))?;
6346 self.device()
6347 .inner()
6348 .dtoh_sync_copy_into(&out_k_view, &mut k_host)
6349 .map_err(|e| XlogError::Kernel(format!("ILP dtoh k: {}", e)))?;
6350 self.device()
6351 .inner()
6352 .dtoh_sync_copy_into(&out_p_view, &mut p_host)
6353 .map_err(|e| XlogError::Kernel(format!("ILP dtoh p: {}", e)))?;
6354
6355 let mut indices: Vec<(f32, u32, u32, u32)> = (0..active_count)
6356 .map(|idx| (p_host[idx], i_host[idx], j_host[idx], k_host[idx]))
6357 .collect();
6358 indices.sort_by(|a, b| b.0.partial_cmp(&a.0).unwrap_or(std::cmp::Ordering::Equal));
6359 indices.truncate(max_active);
6360
6361 Ok(indices.into_iter().map(|(_, i, j, k)| (i, j, k)).collect())
6362 }
6363
6364 #[expect(
6390 clippy::too_many_arguments,
6391 reason = "recorded radix sorting keeps scratch ownership plus the exact stream and runtime explicit"
6392 )]
6393 fn radix_sort_u32_pairs_with_scratch_on_stream(
6394 &self,
6395 keys_a: &mut TrackedCudaSlice<u32>,
6396 keys_b: &mut TrackedCudaSlice<u32>,
6397 indices_a: &mut TrackedCudaSlice<u32>,
6398 indices_b: &mut TrackedCudaSlice<u32>,
6399 hist: &mut TrackedCudaSlice<u32>,
6400 prefix: &mut TrackedCudaSlice<u32>,
6401 ranks: &mut TrackedCudaSlice<u32>,
6402 num_rows_device: &TrackedCudaSlice<u32>,
6403 row_cap: u32,
6404 enqueue: &crate::launch::CudaEnqueue<'_>,
6405 scan_scratch: &mut super::MultiblockScanScratchU32,
6406 ) -> Result<()> {
6407 if row_cap == 0 {
6408 return Ok(());
6409 }
6410 let device = self.device.inner();
6411 let block_size = Self::SORT_BLOCK_SIZE;
6412 let grid_size = row_cap.div_ceil(block_size);
6413 let sort_config = LaunchConfig {
6414 grid_dim: (grid_size, 1, 1),
6415 block_dim: (block_size, 1, 1),
6416 shared_mem_bytes: 0,
6417 };
6418
6419 let histogram_fn = device
6420 .get_func(SORT_MODULE, sort_kernels::RADIX_HISTOGRAM)
6421 .ok_or_else(|| XlogError::Kernel("radix_histogram kernel not found".to_string()))?;
6422 let prefix_fn = device
6423 .get_func(SORT_MODULE, sort_kernels::COMPUTE_DIGIT_PREFIX_SUMS)
6424 .ok_or_else(|| {
6425 XlogError::Kernel("compute_digit_prefix_sums kernel not found".to_string())
6426 })?;
6427 let ranks_fn = device
6428 .get_func(SORT_MODULE, sort_kernels::COMPUTE_RANKS)
6429 .ok_or_else(|| XlogError::Kernel("compute_ranks kernel not found".to_string()))?;
6430 let scatter_fn = device
6431 .get_func(SORT_MODULE, sort_kernels::RADIX_SCATTER_STABLE)
6432 .ok_or_else(|| {
6433 XlogError::Kernel("radix_scatter_stable kernel not found".to_string())
6434 })?;
6435 let prefix_config = LaunchConfig {
6436 grid_dim: (1, 1, 1),
6437 block_dim: (256, 1, 1),
6438 shared_mem_bytes: 0,
6439 };
6440
6441 let mut in_a = true;
6442 for pass in 0..8u32 {
6443 let shift = pass * 4;
6444 let (keys_in, indices_in, keys_out, indices_out) = if in_a {
6445 (&*keys_a, &*indices_a, &mut *keys_b, &mut *indices_b)
6446 } else {
6447 (&*keys_b, &*indices_b, &mut *keys_a, &mut *indices_a)
6448 };
6449
6450 unsafe {
6452 histogram_fn.clone().launch_in(
6453 enqueue,
6454 sort_config,
6455 (keys_in, num_rows_device, row_cap, &mut *hist, shift),
6456 )
6457 }
6458 .map_err(|e| XlogError::Kernel(format!("radix_histogram (on_stream) failed: {}", e)))?;
6459
6460 unsafe {
6462 prefix_fn.clone().launch_in(
6463 enqueue,
6464 prefix_config,
6465 (&*hist, grid_size, &mut *prefix),
6466 )
6467 }
6468 .map_err(|e| {
6469 XlogError::Kernel(format!(
6470 "compute_digit_prefix_sums (on_stream) failed: {}",
6471 e
6472 ))
6473 })?;
6474
6475 for digit in 0..16u32 {
6478 let start = (digit * grid_size) as usize;
6479 let end = start + (grid_size as usize);
6480 let mut digit_slice = hist.slice_mut(start..end);
6481 self.multiblock_scan_u32_inplace_on_stream(
6482 &mut digit_slice,
6483 grid_size,
6484 enqueue,
6485 scan_scratch,
6486 )?;
6487 }
6488
6489 unsafe {
6491 ranks_fn.clone().launch_in(
6492 enqueue,
6493 sort_config,
6494 (keys_in, num_rows_device, row_cap, &mut *ranks, shift),
6495 )
6496 }
6497 .map_err(|e| XlogError::Kernel(format!("compute_ranks (on_stream) failed: {}", e)))?;
6498
6499 unsafe {
6502 scatter_fn.clone().launch_in(
6503 enqueue,
6504 sort_config,
6505 (
6506 keys_in,
6507 indices_in,
6508 &*ranks,
6509 keys_out,
6510 indices_out,
6511 &*prefix,
6512 &*hist,
6513 num_rows_device,
6514 row_cap,
6515 shift,
6516 ),
6517 )
6518 }
6519 .map_err(|e| {
6520 XlogError::Kernel(format!("radix_scatter_stable (on_stream) failed: {}", e))
6521 })?;
6522
6523 in_a = !in_a;
6524 }
6525
6526 if !in_a {
6527 return Err(XlogError::Kernel(
6528 "Unexpected radix-sort buffer parity (expected even number of passes)".to_string(),
6529 ));
6530 }
6531 Ok(())
6532 }
6533
6534 fn apply_permutation_gpu_on_stream(
6540 &self,
6541 input: &CudaBuffer,
6542 permutation: &TrackedCudaSlice<u32>,
6543 dst_cols: &mut [TrackedCudaSlice<u8>],
6544 enqueue: &crate::launch::CudaEnqueue<'_>,
6545 ) -> Result<()> {
6546 let row_cap = input.num_rows() as u32;
6547 let d_num_rows = input.num_rows_device();
6548 let device = self.device.inner();
6549
6550 let grid_size = row_cap.div_ceil(Self::SORT_BLOCK_SIZE);
6551 let launch_config = LaunchConfig {
6552 grid_dim: (grid_size, 1, 1),
6553 block_dim: (Self::SORT_BLOCK_SIZE, 1, 1),
6554 shared_mem_bytes: 0,
6555 };
6556
6557 let apply_perm_fn = device
6558 .get_func(SORT_MODULE, sort_kernels::APPLY_PERMUTATION_BYTES)
6559 .ok_or_else(|| {
6560 XlogError::Kernel("apply_permutation_bytes kernel not found".to_string())
6561 })?;
6562
6563 if dst_cols.len() != input.columns.len() {
6564 return Err(XlogError::Kernel(format!(
6565 "apply_permutation_gpu_on_stream: dst_cols.len()={} mismatches input.cols={}",
6566 dst_cols.len(),
6567 input.columns.len()
6568 )));
6569 }
6570
6571 for (col_idx, dst_col) in dst_cols.iter_mut().enumerate() {
6572 let src_col = input
6573 .column(col_idx)
6574 .ok_or_else(|| XlogError::Kernel(format!("Column {} not found", col_idx)))?;
6575 let elem_size = input
6576 .schema
6577 .column_type(col_idx)
6578 .ok_or_else(|| {
6579 XlogError::Kernel(format!("Schema type for column {} not found", col_idx))
6580 })?
6581 .size_bytes() as u32;
6582 let output_bytes = (row_cap as usize) * (elem_size as usize);
6583 if src_col.num_bytes() != output_bytes {
6584 return Err(XlogError::Kernel(format!(
6585 "Column {} has {} bytes but expected {} (num_rows={}, elem_size={})",
6586 col_idx,
6587 src_col.num_bytes(),
6588 output_bytes,
6589 row_cap,
6590 elem_size
6591 )));
6592 }
6593 unsafe {
6596 apply_perm_fn.clone().launch_in(
6597 enqueue,
6598 launch_config,
6599 (
6600 src_col,
6601 &mut *dst_col,
6602 permutation,
6603 d_num_rows,
6604 row_cap,
6605 elem_size,
6606 ),
6607 )
6608 }
6609 .map_err(|e| {
6610 XlogError::Kernel(format!("apply_permutation_bytes (on_stream) failed: {}", e))
6611 })?;
6612 }
6613 Ok(())
6614 }
6615
6616 pub fn sort_recorded(
6634 &self,
6635 input: &CudaBuffer,
6636 key_cols: &[usize],
6637 launch_stream: StreamId,
6638 ) -> Result<CudaBuffer> {
6639 let runtime = self.memory.runtime().ok_or_else(|| {
6640 XlogError::Kernel(
6641 "sort_recorded requires a runtime-backed GpuMemoryManager (with_runtime)"
6642 .to_string(),
6643 )
6644 })?;
6645 let cu_stream = runtime
6646 .stream_pool()
6647 .resolve(launch_stream)
6648 .ok_or_else(|| {
6649 XlogError::Kernel(format!(
6650 "sort_recorded: launch_stream StreamId({}) does not resolve",
6651 launch_stream.0
6652 ))
6653 })?;
6654
6655 if input.num_rows() == 0 {
6656 return self.create_empty_buffer(input.schema.clone());
6657 }
6658 if key_cols.is_empty() {
6659 return Err(XlogError::Kernel(
6660 "Sort requires at least one key column".to_string(),
6661 ));
6662 }
6663 if input.num_rows() > u32::MAX as u64 {
6664 return Err(XlogError::Kernel(format!(
6665 "Sort supports at most {} rows, got {}",
6666 u32::MAX,
6667 input.num_rows()
6668 )));
6669 }
6670 for &k in key_cols {
6671 if k >= input.arity() {
6672 return Err(XlogError::Kernel(format!(
6673 "Key column index {} out of bounds (arity {})",
6674 k,
6675 input.arity()
6676 )));
6677 }
6678 let ty = input.schema.column_type(k).ok_or_else(|| {
6679 XlogError::Kernel(format!("Key column {} type not found in schema", k))
6680 })?;
6681 if !matches!(ty, ScalarType::U32 | ScalarType::Symbol | ScalarType::U64) {
6682 return Err(XlogError::Kernel(format!(
6683 "sort_recorded supports only U32 / Symbol / U64 key columns; \
6684 got {:?} for column {}",
6685 ty, k
6686 )));
6687 }
6688 }
6689
6690 let n = input.num_rows() as u32;
6691 let block_size = Self::SORT_BLOCK_SIZE;
6692 let grid_size = n.div_ceil(block_size);
6693 let device = self.device.inner();
6694 let launch_config = LaunchConfig {
6695 grid_dim: (grid_size, 1, 1),
6696 block_dim: (block_size, 1, 1),
6697 shared_mem_bytes: 0,
6698 };
6699
6700 let mut indices_a = self.memory.alloc::<u32>(n as usize)?;
6703 let mut indices_b = self.memory.alloc::<u32>(n as usize)?;
6704 let mut keys_a = self.memory.alloc::<u32>(n as usize)?;
6705 let mut keys_b = self.memory.alloc::<u32>(n as usize)?;
6706 let mut d_hist = self.memory.alloc::<u32>((grid_size as usize) * 16)?;
6707 let mut d_prefix = self.memory.alloc::<u32>(16)?;
6708 let mut d_ranks = self.memory.alloc::<u32>(n as usize)?;
6709 let output_d_num_rows = self.memory.alloc::<u32>(1)?;
6715
6716 let mut dst_cols: Vec<TrackedCudaSlice<u8>> = Vec::with_capacity(input.columns.len());
6717 for col_idx in 0..input.columns.len() {
6718 let elem_size = input
6719 .schema
6720 .column_type(col_idx)
6721 .ok_or_else(|| {
6722 XlogError::Kernel(format!("Schema type for column {} not found", col_idx))
6723 })?
6724 .size_bytes();
6725 dst_cols.push(self.memory.alloc::<u8>((n as usize) * elem_size)?);
6726 }
6727
6728 let mut scan_scratch = self.multiblock_scan_u32_scratch_for_len(grid_size)?;
6729 let mut rec = LaunchRecorder::new_strict(launch_stream);
6730 rec.read(input.num_rows_device());
6731 for col_idx in 0..input.columns.len() {
6732 let c = input
6733 .column(col_idx)
6734 .ok_or_else(|| XlogError::Kernel(format!("Column {} not found", col_idx)))?;
6735 rec.read_column(c);
6736 }
6737 rec.write(&indices_a);
6741 rec.write(&indices_b);
6742 rec.write(&keys_a);
6743 rec.write(&keys_b);
6744 rec.write(&d_hist);
6745 rec.write(&d_prefix);
6746 rec.write(&d_ranks);
6747 rec.write(&output_d_num_rows);
6748 for dst_col in &dst_cols {
6749 rec.write(dst_col);
6750 }
6751 for level in scan_scratch.levels() {
6752 rec.read_write(level);
6753 }
6754 rec.preflight(runtime)
6755 .map_err(|e| XlogError::Kernel(format!("sort_recorded: preflight failed: {}", e)))?;
6756
6757 let init_fn = device
6759 .get_func(SORT_MODULE, sort_kernels::INIT_INDICES)
6760 .ok_or_else(|| XlogError::Kernel("init_indices kernel not found".to_string()))?;
6761 let enqueue = |enqueue: &crate::launch::CudaEnqueue<'_>| -> Result<()> {
6763 unsafe {
6764 init_fn.clone().launch_in(
6765 enqueue,
6766 launch_config,
6767 (&mut indices_a, input.num_rows_device(), n),
6768 )
6769 }
6770 .map_err(|e| XlogError::Kernel(format!("init_indices (on_stream) failed: {}", e)))?;
6771
6772 for &col_idx in key_cols.iter().rev() {
6779 let col = input.column(col_idx).ok_or_else(|| {
6780 XlogError::Kernel(format!("Key column {} not found", col_idx))
6781 })?;
6782 let ty = input.schema.column_type(col_idx).ok_or_else(|| {
6783 XlogError::Kernel(format!("Key column {} type not found in schema", col_idx))
6784 })?;
6785 match ty {
6786 ScalarType::U32 | ScalarType::Symbol => {
6787 let col_view = self.column_as_u32_view(col, n as usize)?;
6788 let gather_fn = device
6789 .get_func(SORT_MODULE, sort_kernels::APPLY_PERMUTATION_U32)
6790 .ok_or_else(|| {
6791 XlogError::Kernel(
6792 "apply_permutation_u32 kernel not found".to_string(),
6793 )
6794 })?;
6795 unsafe {
6798 gather_fn.clone().launch_in(
6799 enqueue,
6800 launch_config,
6801 (
6802 &col_view,
6803 &mut keys_a,
6804 &indices_a,
6805 input.num_rows_device(),
6806 n,
6807 ),
6808 )
6809 }
6810 .map_err(|e| {
6811 XlogError::Kernel(format!(
6812 "apply_permutation_u32 (on_stream) failed: {}",
6813 e
6814 ))
6815 })?;
6816
6817 self.radix_sort_u32_pairs_with_scratch_on_stream(
6818 &mut keys_a,
6819 &mut keys_b,
6820 &mut indices_a,
6821 &mut indices_b,
6822 &mut d_hist,
6823 &mut d_prefix,
6824 &mut d_ranks,
6825 input.num_rows_device(),
6826 n,
6827 enqueue,
6828 &mut scan_scratch,
6829 )?;
6830 }
6831 ScalarType::U64 => {
6832 let col_view = self.column_as_u64_view(col, n as usize)?;
6833 for &word in &[
6834 sort_kernels::GATHER_KEYS_U64_LO_U32,
6835 sort_kernels::GATHER_KEYS_U64_HI_U32,
6836 ] {
6837 let gather_fn =
6838 device.get_func(SORT_MODULE, word).ok_or_else(|| {
6839 XlogError::Kernel(format!("{} kernel not found", word))
6840 })?;
6841 unsafe {
6844 gather_fn.clone().launch_in(
6845 enqueue,
6846 launch_config,
6847 (
6848 &col_view,
6849 &indices_a,
6850 input.num_rows_device(),
6851 n,
6852 &mut keys_a,
6853 ),
6854 )
6855 }
6856 .map_err(|e| {
6857 XlogError::Kernel(format!("{} (on_stream) failed: {}", word, e))
6858 })?;
6859
6860 self.radix_sort_u32_pairs_with_scratch_on_stream(
6861 &mut keys_a,
6862 &mut keys_b,
6863 &mut indices_a,
6864 &mut indices_b,
6865 &mut d_hist,
6866 &mut d_prefix,
6867 &mut d_ranks,
6868 input.num_rows_device(),
6869 n,
6870 enqueue,
6871 &mut scan_scratch,
6872 )?;
6873 }
6874 }
6875 other => {
6876 return Err(XlogError::Kernel(format!(
6877 "sort_recorded: column {} unexpected type {:?} after guard",
6878 col_idx, other
6879 )));
6880 }
6881 }
6882 }
6883
6884 self.apply_permutation_gpu_on_stream(input, &indices_a, &mut dst_cols, enqueue)?;
6886
6887 unsafe {
6894 let res = cudarc::driver::sys::cuMemcpyDtoDAsync_v2(
6895 *output_d_num_rows.device_ptr(),
6896 *input.num_rows_device().device_ptr(),
6897 std::mem::size_of::<u32>(),
6898 cu_stream.cu_stream(),
6899 );
6900 if res != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
6901 return Err(XlogError::Kernel(format!(
6902 "sort_recorded: cuMemcpyDtoDAsync (output_d_num_rows) failed: {:?}",
6903 res
6904 )));
6905 }
6906 }
6907 Ok(())
6908 };
6909 let rec = unsafe { rec.enqueue_prepared_with(&cu_stream, enqueue) }
6911 .map_err(LaunchEnqueueError::into_xlog_error)?;
6912
6913 rec.commit()
6914 .map_err(|e| XlogError::Kernel(format!("sort_recorded: commit failed: {}", e)))?;
6915
6916 let new_columns: Vec<CudaColumn> = dst_cols.into_iter().map(|s| s.into()).collect();
6917 Ok(CudaBuffer::from_columns(
6918 new_columns,
6919 input.num_rows(),
6920 output_d_num_rows,
6921 input.schema.clone(),
6922 ))
6923 }
6924
6925 pub fn dedup_full_row_recorded(
6936 &self,
6937 input: &CudaBuffer,
6938 launch_stream: StreamId,
6939 ) -> Result<CudaBuffer> {
6940 if !input.canonical_full_row_set_certified() {
6941 self.validated_logical_row_count(input)?;
6942 }
6943 let runtime = self.memory.runtime().ok_or_else(|| {
6944 XlogError::Kernel(
6945 "dedup_full_row_recorded requires a runtime-backed GpuMemoryManager".to_string(),
6946 )
6947 })?;
6948 let cu_stream = runtime
6949 .stream_pool()
6950 .resolve(launch_stream)
6951 .ok_or_else(|| {
6952 XlogError::Kernel(format!(
6953 "dedup_full_row_recorded: launch_stream StreamId({}) does not resolve",
6954 launch_stream.0
6955 ))
6956 })?;
6957
6958 let row_count = input.num_rows() as usize;
6959 if row_count == 0 {
6960 return self.create_empty_buffer(input.schema().clone());
6961 }
6962 if row_count == 1 {
6963 return self.clone_buffer(input);
6964 }
6965 if row_count > u32::MAX as usize {
6966 return Err(XlogError::Kernel(format!(
6967 "dedup_full_row_recorded supports at most {} rows, got {}",
6968 u32::MAX,
6969 row_count
6970 )));
6971 }
6972 let arity = input.arity();
6973 if arity == 0 {
6974 return self.buffer_from_columns(Vec::new(), 1, input.schema().clone());
6975 }
6976 for col_idx in 0..arity {
6977 let ty = input.schema.column_type(col_idx).ok_or_else(|| {
6978 XlogError::Kernel(format!("Column {} type not found in schema", col_idx))
6979 })?;
6980 if !matches!(ty, ScalarType::U32 | ScalarType::Symbol | ScalarType::U64) {
6981 return Err(XlogError::Kernel(format!(
6982 "dedup_full_row_recorded supports only U32 / Symbol / U64 columns; \
6983 got {:?} for column {}",
6984 ty, col_idx
6985 )));
6986 }
6987 }
6988
6989 let all_cols: Vec<usize> = (0..arity).collect();
6991 let sorted = self.sort_recorded(input, &all_cols, launch_stream)?;
6992 let n = sorted.num_rows() as u32;
6993 if n <= 1 {
6994 return Ok(sorted);
6995 }
6996
6997 let device = self.device.inner();
7003 let mut col_ptrs_host: Vec<u64> = Vec::with_capacity(arity);
7004 let mut col_sizes_host: Vec<u32> = Vec::with_capacity(arity);
7005 for col_idx in 0..arity {
7006 let c = sorted
7007 .column(col_idx)
7008 .ok_or_else(|| XlogError::Kernel(format!("Sorted column {} not found", col_idx)))?;
7009 let ty = sorted.schema().column_type(col_idx).ok_or_else(|| {
7010 XlogError::Kernel(format!("Sorted column {} type missing", col_idx))
7011 })?;
7012 col_ptrs_host.push(*c.device_ptr());
7013 col_sizes_host.push(ty.size_bytes() as u32);
7014 }
7015 let mut d_col_ptrs = self.memory.alloc::<u64>(arity)?;
7016 let mut d_col_sizes = self.memory.alloc::<u32>(arity)?;
7017 self.htod_launch_metadata_sync_copy_into(&col_ptrs_host, &mut d_col_ptrs)
7018 .map_err(|e| {
7019 XlogError::Kernel(format!("dedup_full_row_recorded col ptr upload: {}", e))
7020 })?;
7021 self.htod_launch_metadata_sync_copy_into(&col_sizes_host, &mut d_col_sizes)
7022 .map_err(|e| {
7023 XlogError::Kernel(format!("dedup_full_row_recorded col size upload: {}", e))
7024 })?;
7025 let d_unique_mask = self.memory.alloc::<u8>(n as usize)?;
7026
7027 let mut rec = LaunchRecorder::new_strict(launch_stream);
7028 for col_idx in 0..arity {
7029 let c = sorted
7030 .column(col_idx)
7031 .ok_or_else(|| XlogError::Kernel(format!("Sorted column {} not found", col_idx)))?;
7032 rec.read_column(c);
7033 }
7034 rec.read(sorted.num_rows_device());
7035 rec.write(&d_col_ptrs);
7036 rec.write(&d_col_sizes);
7037 rec.write(&d_unique_mask);
7038 rec.preflight(runtime).map_err(|e| {
7039 XlogError::Kernel(format!(
7040 "dedup_full_row_recorded: mark_unique preflight failed: {}",
7041 e
7042 ))
7043 })?;
7044
7045 let block_size = 256u32;
7046 let grid = n.div_ceil(block_size);
7047 let cfg = LaunchConfig {
7048 grid_dim: (grid, 1, 1),
7049 block_dim: (block_size, 1, 1),
7050 shared_mem_bytes: 0,
7051 };
7052 let mark_fn = device
7053 .get_func(DEDUP_MODULE, dedup_kernels::MARK_UNIQUE_FULL_ROW_BYTEWISE)
7054 .ok_or_else(|| {
7055 XlogError::Kernel("mark_unique_full_row_bytewise kernel not found".to_string())
7056 })?;
7057 let enqueue = |enqueue: &crate::launch::CudaEnqueue<'_>| -> Result<()> {
7060 unsafe {
7061 mark_fn.clone().launch_in(
7062 enqueue,
7063 cfg,
7064 (
7065 &d_col_ptrs,
7066 &d_col_sizes,
7067 arity as u32,
7068 sorted.num_rows_device(),
7069 n,
7070 &d_unique_mask,
7071 ),
7072 )
7073 }
7074 .map_err(|e| {
7075 XlogError::Kernel(format!(
7076 "mark_unique_full_row_bytewise (on_stream) failed: {}",
7077 e
7078 ))
7079 })
7080 };
7081 let rec = unsafe { rec.enqueue_prepared_with(&cu_stream, enqueue) }
7083 .map_err(LaunchEnqueueError::into_xlog_error)?;
7084
7085 rec.commit().map_err(|e| {
7086 XlogError::Kernel(format!(
7087 "dedup_full_row_recorded: mark_unique commit failed: {}",
7088 e
7089 ))
7090 })?;
7091
7092 let mut result = self.compact_buffer_by_device_mask_counted_recorded(
7094 &sorted,
7095 &d_unique_mask,
7096 launch_stream,
7097 )?;
7098 result.certify_canonical_full_row_set();
7099 Ok(result)
7100 }
7101
7102 fn build_hash_table_v2_on_stream(
7130 &self,
7131 hashes: &TrackedCudaSlice<u64>,
7132 num_rows: u32,
7133 cu_stream: &cudarc::driver::CudaStream,
7134 launch_stream: StreamId,
7135 runtime: &Arc<XlogDeviceRuntime>,
7136 ) -> Result<crate::provider::JoinHashTableV2> {
7137 let device = self.device.inner();
7138
7139 let target = (num_rows as u64).saturating_mul(2).max(1024);
7140 let num_buckets_u64 = target.next_power_of_two();
7141 let num_buckets = u32::try_from(num_buckets_u64).map_err(|_| {
7142 XlogError::Kernel(format!(
7143 "Join hash table too large: num_buckets={}",
7144 num_buckets_u64
7145 ))
7146 })?;
7147 let bucket_mask = num_buckets
7148 .checked_sub(1)
7149 .ok_or_else(|| XlogError::Kernel("Join hash table size underflow".to_string()))?;
7150
7151 let bucket_counts = self.memory.alloc::<u32>(num_buckets as usize)?;
7152 let mut bucket_offsets = self.memory.alloc::<u32>(num_buckets as usize)?;
7153 let bucket_cursors = self.memory.alloc::<u32>(num_buckets as usize)?;
7154 let bucket_entries = self.memory.alloc::<u32>(num_rows as usize)?;
7155 let bucket_entry_hashes = self.memory.alloc::<u64>(num_rows as usize)?;
7156
7157 let mut scan_scratch = self.multiblock_scan_u32_scratch_for_len(num_buckets)?;
7158 let mut rec = LaunchRecorder::new_strict(launch_stream);
7159 rec.read(hashes);
7160 rec.write(&bucket_counts);
7161 rec.write(&bucket_offsets);
7162 rec.write(&bucket_cursors);
7163 rec.write(&bucket_entries);
7164 rec.write(&bucket_entry_hashes);
7165 for level in scan_scratch.levels() {
7166 rec.read_write(level);
7167 }
7168 rec.preflight(runtime).map_err(|e| {
7169 XlogError::Kernel(format!(
7170 "build_hash_table_v2_on_stream: preflight failed: {}",
7171 e
7172 ))
7173 })?;
7174 let enqueue = |enqueue: &crate::launch::CudaEnqueue<'_>| -> Result<()> {
7175 if num_buckets > 0 {
7177 unsafe {
7181 let res = cudarc::driver::sys::cuMemsetD8Async(
7182 *bucket_counts.device_ptr(),
7183 0,
7184 (num_buckets as usize) * std::mem::size_of::<u32>(),
7185 cu_stream.cu_stream(),
7186 );
7187 if res != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
7188 return Err(XlogError::Kernel(format!(
7189 "cuMemsetD8Async (bucket_counts) failed: {:?}",
7190 res
7191 )));
7192 }
7193 }
7194 }
7195
7196 let block_size = 256u32;
7197 let grid_size = num_rows.div_ceil(block_size);
7198 let cfg = LaunchConfig {
7199 grid_dim: (grid_size, 1, 1),
7200 block_dim: (block_size, 1, 1),
7201 shared_mem_bytes: 0,
7202 };
7203
7204 let count_fn = device
7205 .get_func(JOIN_MODULE, join_kernels::HASH_JOIN_BUCKET_COUNT_V2)
7206 .ok_or_else(|| {
7207 XlogError::Kernel("hash_join_bucket_count_v2 kernel not found".to_string())
7208 })?;
7209 unsafe {
7211 count_fn.clone().launch_in(
7212 enqueue,
7213 cfg,
7214 (hashes, num_rows, &bucket_counts, bucket_mask),
7215 )
7216 }
7217 .map_err(|e| {
7218 XlogError::Kernel(format!(
7219 "hash_join_bucket_count_v2 (on_stream) failed: {}",
7220 e
7221 ))
7222 })?;
7223
7224 if num_buckets > 0 {
7225 unsafe {
7229 let res = cudarc::driver::sys::cuMemcpyDtoDAsync_v2(
7230 *bucket_offsets.device_ptr(),
7231 *bucket_counts.device_ptr(),
7232 (num_buckets as usize) * std::mem::size_of::<u32>(),
7233 cu_stream.cu_stream(),
7234 );
7235 if res != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
7236 return Err(XlogError::Kernel(format!(
7237 "cuMemcpyDtoDAsync (bucket_counts → bucket_offsets) failed: {:?}",
7238 res
7239 )));
7240 }
7241 }
7242 self.multiblock_scan_u32_inplace_on_stream(
7243 &mut bucket_offsets,
7244 num_buckets,
7245 enqueue,
7246 &mut scan_scratch,
7247 )?;
7248 }
7249
7250 if num_buckets > 0 {
7251 unsafe {
7254 let res = cudarc::driver::sys::cuMemcpyDtoDAsync_v2(
7255 *bucket_cursors.device_ptr(),
7256 *bucket_offsets.device_ptr(),
7257 (num_buckets as usize) * std::mem::size_of::<u32>(),
7258 cu_stream.cu_stream(),
7259 );
7260 if res != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
7261 return Err(XlogError::Kernel(format!(
7262 "cuMemcpyDtoDAsync (bucket_offsets → bucket_cursors) failed: {:?}",
7263 res
7264 )));
7265 }
7266 }
7267 }
7268
7269 let scatter_fn = device
7270 .get_func(JOIN_MODULE, join_kernels::HASH_JOIN_SCATTER_V2)
7271 .ok_or_else(|| {
7272 XlogError::Kernel("hash_join_scatter_v2 kernel not found".to_string())
7273 })?;
7274 unsafe {
7276 scatter_fn.clone().launch_in(
7277 enqueue,
7278 cfg,
7279 (
7280 hashes,
7281 num_rows,
7282 &bucket_cursors,
7283 bucket_mask,
7284 &bucket_entries,
7285 &bucket_entry_hashes,
7286 ),
7287 )
7288 }
7289 .map_err(|e| {
7290 XlogError::Kernel(format!("hash_join_scatter_v2 (on_stream) failed: {}", e))
7291 })?;
7292 Ok(())
7293 };
7294 let rec = unsafe { rec.enqueue_prepared_with(cu_stream, enqueue) }
7296 .map_err(LaunchEnqueueError::into_xlog_error)?;
7297 rec.commit().map_err(|e| {
7298 XlogError::Kernel(format!(
7299 "build_hash_table_v2_on_stream: commit failed: {}",
7300 e
7301 ))
7302 })?;
7303
7304 Ok(crate::provider::JoinHashTableV2 {
7305 bucket_counts,
7306 bucket_offsets,
7307 bucket_entries,
7308 bucket_entry_hashes,
7309 bucket_mask,
7310 })
7311 }
7312
7313 fn gather_buffer_by_indices_on_stream(
7319 &self,
7320 input: &CudaBuffer,
7321 indices: &TrackedCudaSlice<u32>,
7322 output_rows: u32,
7323 cu_stream: &cudarc::driver::CudaStream,
7324 launch_stream: StreamId,
7325 runtime: &Arc<XlogDeviceRuntime>,
7326 ) -> Result<CudaBuffer> {
7327 if output_rows == 0 {
7328 return self.create_empty_buffer(input.schema().clone());
7329 }
7330 if input.num_rows() > u32::MAX as u64 {
7331 return Err(XlogError::Kernel(format!(
7332 "GPU gather supports at most {} input rows, got {}",
7333 u32::MAX,
7334 input.num_rows()
7335 )));
7336 }
7337
7338 let d_output_rows = self.upload_device_row_count(output_rows)?;
7339 let device = self.device.inner();
7342 let block_size = 256u32;
7343 let grid_size = output_rows.div_ceil(block_size);
7344 let launch_config = LaunchConfig {
7345 grid_dim: (grid_size, 1, 1),
7346 block_dim: (block_size, 1, 1),
7347 shared_mem_bytes: 0,
7348 };
7349
7350 let gather_fn = device
7351 .get_func(SORT_MODULE, sort_kernels::APPLY_PERMUTATION_BYTES)
7352 .ok_or_else(|| {
7353 XlogError::Kernel("apply_permutation_bytes kernel not found".to_string())
7354 })?;
7355
7356 let mut dst_cols: Vec<TrackedCudaSlice<u8>> = Vec::with_capacity(input.columns.len());
7357 let mut elem_sizes: Vec<u32> = Vec::with_capacity(input.columns.len());
7358 for col_idx in 0..input.columns.len() {
7359 let elem_size = input
7360 .schema
7361 .column_type(col_idx)
7362 .ok_or_else(|| {
7363 XlogError::Kernel(format!("Schema type for column {} not found", col_idx))
7364 })?
7365 .size_bytes() as u32;
7366 input
7367 .column(col_idx)
7368 .ok_or_else(|| XlogError::Kernel(format!("Column {} not found", col_idx)))?;
7369 let dst_bytes = (output_rows as usize) * (elem_size as usize);
7370 let dst = self.memory.alloc::<u8>(dst_bytes)?;
7371 dst_cols.push(dst);
7372 elem_sizes.push(elem_size);
7373 }
7374
7375 let mut rec = LaunchRecorder::new_strict(launch_stream);
7376 for col_idx in 0..input.columns.len() {
7377 rec.read_column(
7378 input
7379 .column(col_idx)
7380 .expect("input column validated before preflight"),
7381 );
7382 }
7383 rec.read(indices);
7384 rec.read(&d_output_rows);
7385 for dst_col in &dst_cols {
7386 rec.write(dst_col);
7387 }
7388 rec.preflight(runtime).map_err(|e| {
7389 XlogError::Kernel(format!(
7390 "gather_buffer_by_indices_on_stream: preflight failed: {}",
7391 e
7392 ))
7393 })?;
7394 let enqueue = |enqueue: &crate::launch::CudaEnqueue<'_>| -> Result<()> {
7395 for (col_idx, dst_col) in dst_cols.iter_mut().enumerate() {
7396 let src_col = input
7397 .column(col_idx)
7398 .expect("input column validated before enqueue boundary");
7399 let elem_size = elem_sizes[col_idx];
7400 unsafe {
7403 gather_fn.clone().launch_in(
7404 enqueue,
7405 launch_config,
7406 (
7407 src_col,
7408 &mut *dst_col,
7409 indices,
7410 &d_output_rows,
7411 output_rows,
7412 elem_size,
7413 ),
7414 )
7415 }
7416 .map_err(|e| {
7417 XlogError::Kernel(format!("apply_permutation_bytes (on_stream) failed: {}", e))
7418 })?;
7419 }
7420 Ok(())
7421 };
7422 let rec = unsafe { rec.enqueue_prepared_with(cu_stream, enqueue) }
7424 .map_err(LaunchEnqueueError::into_xlog_error)?;
7425 rec.commit().map_err(|e| {
7426 XlogError::Kernel(format!(
7427 "gather_buffer_by_indices_on_stream: commit failed: {}",
7428 e
7429 ))
7430 })?;
7431
7432 let new_columns: Vec<CudaColumn> = dst_cols.into_iter().map(|s| s.into()).collect();
7433 Ok(CudaBuffer::from_columns(
7434 new_columns,
7435 output_rows as u64,
7436 d_output_rows,
7437 input.schema.clone(),
7438 ))
7439 }
7440
7441 #[expect(
7442 clippy::too_many_arguments,
7443 reason = "recorded left-outer concatenation needs both relation layouts, both matched buffers, the unmatched tail, row counts, and stream state"
7444 )]
7445 fn concatenate_left_outer_columns_on_stream(
7446 &self,
7447 left: &CudaBuffer,
7448 right: &CudaBuffer,
7449 inner_left: Option<&CudaBuffer>,
7450 inner_right: Option<&CudaBuffer>,
7451 unmatched_left: &CudaBuffer,
7452 inner_rows: u64,
7453 unmatched_rows: u64,
7454 cu_stream: &cudarc::driver::CudaStream,
7455 launch_stream: StreamId,
7456 runtime: &Arc<XlogDeviceRuntime>,
7457 context: &'static str,
7458 ) -> Result<Vec<CudaColumn>> {
7459 let inner_rows = usize::try_from(inner_rows).map_err(|_| {
7460 XlogError::Kernel(format!("{context}: matched row count exceeds usize::MAX"))
7461 })?;
7462 let unmatched_rows = usize::try_from(unmatched_rows).map_err(|_| {
7463 XlogError::Kernel(format!("{context}: unmatched row count exceeds usize::MAX"))
7464 })?;
7465
7466 let mut left_outputs = Vec::with_capacity(left.arity());
7467 for col_idx in 0..left.arity() {
7468 let elem_size = left
7469 .schema()
7470 .column_type(col_idx)
7471 .map(|ty| ty.size_bytes())
7472 .unwrap_or(4);
7473 let inner_bytes = inner_rows.checked_mul(elem_size).ok_or_else(|| {
7474 XlogError::Kernel(format!("{context}: matched left column size overflow"))
7475 })?;
7476 let unmatched_bytes = unmatched_rows.checked_mul(elem_size).ok_or_else(|| {
7477 XlogError::Kernel(format!("{context}: unmatched left column size overflow"))
7478 })?;
7479 let total_bytes = inner_bytes.checked_add(unmatched_bytes).ok_or_else(|| {
7480 XlogError::Kernel(format!("{context}: left output column size overflow"))
7481 })?;
7482 if inner_bytes > 0 {
7483 inner_left
7484 .and_then(|buffer| buffer.column(col_idx))
7485 .ok_or_else(|| {
7486 XlogError::Kernel(format!(
7487 "{context}: matched left column {col_idx} not found"
7488 ))
7489 })?;
7490 }
7491 if unmatched_bytes > 0 {
7492 unmatched_left.column(col_idx).ok_or_else(|| {
7493 XlogError::Kernel(format!(
7494 "{context}: unmatched left column {col_idx} not found"
7495 ))
7496 })?;
7497 }
7498 left_outputs.push((
7499 self.memory.alloc::<u8>(total_bytes)?,
7500 inner_bytes,
7501 unmatched_bytes,
7502 ));
7503 }
7504
7505 let mut right_outputs = Vec::with_capacity(right.arity());
7506 for col_idx in 0..right.arity() {
7507 let elem_size = right
7508 .schema()
7509 .column_type(col_idx)
7510 .map(|ty| ty.size_bytes())
7511 .unwrap_or(4);
7512 let inner_bytes = inner_rows.checked_mul(elem_size).ok_or_else(|| {
7513 XlogError::Kernel(format!("{context}: matched right column size overflow"))
7514 })?;
7515 let unmatched_bytes = unmatched_rows.checked_mul(elem_size).ok_or_else(|| {
7516 XlogError::Kernel(format!("{context}: unmatched right column size overflow"))
7517 })?;
7518 let total_bytes = inner_bytes.checked_add(unmatched_bytes).ok_or_else(|| {
7519 XlogError::Kernel(format!("{context}: right output column size overflow"))
7520 })?;
7521 if inner_bytes > 0 {
7522 inner_right
7523 .and_then(|buffer| buffer.column(col_idx))
7524 .ok_or_else(|| {
7525 XlogError::Kernel(format!(
7526 "{context}: matched right column {col_idx} not found"
7527 ))
7528 })?;
7529 }
7530 right_outputs.push((
7531 self.memory.alloc::<u8>(total_bytes)?,
7532 inner_bytes,
7533 total_bytes,
7534 ));
7535 }
7536
7537 let mut recorder = LaunchRecorder::new_strict(launch_stream);
7538 for (col_idx, (_, inner_bytes, unmatched_bytes)) in
7539 left_outputs.iter().enumerate().take(left.arity())
7540 {
7541 if *inner_bytes > 0 {
7542 recorder.read_column(
7543 inner_left
7544 .and_then(|buffer| buffer.column(col_idx))
7545 .expect("matched left column validated before recorder preflight"),
7546 );
7547 }
7548 if *unmatched_bytes > 0 {
7549 recorder.read_column(
7550 unmatched_left
7551 .column(col_idx)
7552 .expect("unmatched left column validated before recorder preflight"),
7553 );
7554 }
7555 }
7556 for (col_idx, (_, inner_bytes, _)) in right_outputs.iter().enumerate().take(right.arity()) {
7557 if *inner_bytes > 0 {
7558 recorder.read_column(
7559 inner_right
7560 .and_then(|buffer| buffer.column(col_idx))
7561 .expect("matched right column validated before recorder preflight"),
7562 );
7563 }
7564 }
7565 for (column, _, _) in &left_outputs {
7566 recorder.write(column);
7567 }
7568 for (column, _, _) in &right_outputs {
7569 recorder.write(column);
7570 }
7571 recorder.preflight(runtime).map_err(|error| {
7572 XlogError::Kernel(format!("{context}: concat preflight failed: {error}"))
7573 })?;
7574
7575 let has_enqueues = left_outputs
7576 .iter()
7577 .chain(&right_outputs)
7578 .any(|(_, inner_bytes, remaining_bytes)| *inner_bytes > 0 || *remaining_bytes > 0);
7579 if !has_enqueues {
7580 recorder.abort().map_err(|error| {
7581 XlogError::Kernel(format!("{context}: concat abort failed: {error}"))
7582 })?;
7583 } else {
7584 let enqueue = |enqueue: &crate::launch::CudaEnqueue<'_>| -> Result<()> {
7585 for (col_idx, (out_column, inner_bytes, unmatched_bytes)) in
7586 left_outputs.iter().enumerate()
7587 {
7588 let destination = *out_column.device_ptr();
7589 if *inner_bytes > 0 {
7590 let source = inner_left
7591 .and_then(|buffer| buffer.column(col_idx))
7592 .expect("matched left column validated before enqueue boundary");
7593 let result = unsafe {
7595 cudarc::driver::sys::cuMemcpyDtoDAsync_v2(
7596 destination,
7597 *source.device_ptr(),
7598 *inner_bytes,
7599 enqueue.stream().cu_stream(),
7600 )
7601 };
7602 if result != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
7603 return Err(XlogError::Kernel(format!(
7604 "{context}: matched left column {col_idx} copy failed: {result:?}"
7605 )));
7606 }
7607 }
7608 if *unmatched_bytes > 0 {
7609 let source = unmatched_left
7610 .column(col_idx)
7611 .expect("unmatched left column validated before enqueue boundary");
7612 let result = unsafe {
7614 cudarc::driver::sys::cuMemcpyDtoDAsync_v2(
7615 destination + *inner_bytes as u64,
7616 *source.device_ptr(),
7617 *unmatched_bytes,
7618 cu_stream.cu_stream(),
7619 )
7620 };
7621 if result != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
7622 return Err(XlogError::Kernel(format!(
7623 "{context}: unmatched left column {col_idx} copy failed: {result:?}"
7624 )));
7625 }
7626 }
7627 }
7628
7629 for (col_idx, (out_column, inner_bytes, total_bytes)) in
7630 right_outputs.iter().enumerate()
7631 {
7632 let destination = *out_column.device_ptr();
7633 if *total_bytes > 0 {
7634 let result = unsafe {
7636 cudarc::driver::sys::cuMemsetD8Async(
7637 destination,
7638 0,
7639 *total_bytes,
7640 cu_stream.cu_stream(),
7641 )
7642 };
7643 if result != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
7644 return Err(XlogError::Kernel(format!(
7645 "{context}: right column {col_idx} zero-fill failed: {result:?}"
7646 )));
7647 }
7648 }
7649 if *inner_bytes > 0 {
7650 let source = inner_right
7651 .and_then(|buffer| buffer.column(col_idx))
7652 .expect("matched right column validated before enqueue boundary");
7653 let result = unsafe {
7655 cudarc::driver::sys::cuMemcpyDtoDAsync_v2(
7656 destination,
7657 *source.device_ptr(),
7658 *inner_bytes,
7659 cu_stream.cu_stream(),
7660 )
7661 };
7662 if result != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
7663 return Err(XlogError::Kernel(format!(
7664 "{context}: matched right column {col_idx} copy failed: {result:?}"
7665 )));
7666 }
7667 }
7668 }
7669 Ok(())
7670 };
7671 let recorder = unsafe { recorder.enqueue_prepared_with(cu_stream, enqueue) }
7673 .map_err(LaunchEnqueueError::into_xlog_error)?;
7674 recorder.commit().map_err(|error| {
7675 XlogError::Kernel(format!("{context}: concat commit failed: {error}"))
7676 })?;
7677 }
7678
7679 Ok(left_outputs
7680 .into_iter()
7681 .chain(right_outputs)
7682 .map(|(column, _, _)| column.into())
7683 .collect())
7684 }
7685
7686 pub fn hash_join_inner_v2_recorded(
7693 &self,
7694 left: &CudaBuffer,
7695 right: &CudaBuffer,
7696 left_keys: &[usize],
7697 right_keys: &[usize],
7698 max_output: Option<usize>,
7699 launch_stream: StreamId,
7700 ) -> Result<CudaBuffer> {
7701 let runtime = self.memory.runtime().ok_or_else(|| {
7702 XlogError::Kernel(
7703 "hash_join_inner_v2_recorded requires a runtime-backed GpuMemoryManager"
7704 .to_string(),
7705 )
7706 })?;
7707 let cu_stream = runtime
7708 .stream_pool()
7709 .resolve(launch_stream)
7710 .ok_or_else(|| {
7711 XlogError::Kernel(format!(
7712 "hash_join_inner_v2_recorded: launch_stream StreamId({}) does not resolve",
7713 launch_stream.0
7714 ))
7715 })?;
7716
7717 let num_left = self.device_row_count(left)?;
7718 let num_right = self.device_row_count(right)?;
7719 if num_left > u32::MAX as usize || num_right > u32::MAX as usize {
7720 return Err(XlogError::Kernel(format!(
7721 "Join supports at most {} rows per side (left={}, right={})",
7722 u32::MAX,
7723 num_left,
7724 num_right
7725 )));
7726 }
7727 if num_left == 0 || num_right == 0 {
7728 let combined_schema = self.combine_schemas(left.schema(), right.schema());
7729 return self.create_empty_buffer(combined_schema);
7730 }
7731 if left_keys.is_empty() || right_keys.is_empty() {
7732 return Err(XlogError::Kernel(
7733 "Join requires at least one key column".to_string(),
7734 ));
7735 }
7736 if left_keys.len() != right_keys.len() {
7737 return Err(XlogError::Kernel(
7738 "Left and right key columns must have same length".to_string(),
7739 ));
7740 }
7741 if left_keys.len() > 4 {
7742 return Err(XlogError::Kernel(
7743 "hash_join_inner_v2_recorded: max 4 key columns supported (pack_keys constraint)"
7744 .to_string(),
7745 ));
7746 }
7747 for (&l, &r) in left_keys.iter().zip(right_keys.iter()) {
7748 let lt = left.schema().column_type(l);
7749 let rt = right.schema().column_type(r);
7750 if lt != rt {
7751 return Err(XlogError::Kernel(format!(
7752 "Key column type mismatch: left[{}]={:?}, right[{}]={:?}",
7753 l, lt, r, rt
7754 )));
7755 }
7756 }
7757
7758 let num_left = num_left as u32;
7759 let num_right = num_right as u32;
7760
7761 let left_packed =
7763 self.pack_keys_gpu_on_stream(left, left_keys, &cu_stream, launch_stream, runtime)?;
7764 let right_packed =
7765 self.pack_keys_gpu_on_stream(right, right_keys, &cu_stream, launch_stream, runtime)?;
7766
7767 let table = self.build_hash_table_v2_on_stream(
7769 &right_packed.hashes,
7770 num_right,
7771 &cu_stream,
7772 launch_stream,
7773 runtime,
7774 )?;
7775
7776 let probe_func = self
7777 .device
7778 .inner()
7779 .get_func(JOIN_MODULE, join_kernels::HASH_JOIN_PROBE_V2)
7780 .ok_or_else(|| XlogError::Kernel("hash_join_probe_v2 kernel not found".to_string()))?;
7781 let block_size = 256u32;
7782 let probe_grid = num_left.div_ceil(block_size);
7783 let probe_config = LaunchConfig {
7784 grid_dim: (probe_grid, 1, 1),
7785 block_dim: (block_size, 1, 1),
7786 shared_mem_bytes: 0,
7787 };
7788
7789 let d_count_only = self.memory.alloc::<u32>(1)?;
7794 let d_dummy_left = self.memory.alloc::<u32>(1)?;
7795 let d_dummy_right = self.memory.alloc::<u32>(1)?;
7796
7797 let max_output_count_only = 0u32;
7803 let mut rec_count = LaunchRecorder::new_strict(launch_stream);
7804 rec_count.read(&left_packed.hashes);
7805 rec_count.read(&left_packed.packed_keys);
7806 rec_count.read(&right_packed.packed_keys);
7807 rec_count.read(&table.bucket_offsets);
7808 rec_count.read(&table.bucket_counts);
7809 rec_count.read(&table.bucket_entries);
7810 rec_count.read(&table.bucket_entry_hashes);
7811 rec_count.write(&d_count_only);
7812 rec_count.write(&d_dummy_left);
7813 rec_count.write(&d_dummy_right);
7814 rec_count.preflight(runtime).map_err(|e| {
7815 XlogError::Kernel(format!(
7816 "hash_join_inner_v2_recorded: count-pass preflight failed: {}",
7817 e
7818 ))
7819 })?;
7820
7821 let enqueue = |enqueue: &crate::launch::CudaEnqueue<'_>| -> Result<()> {
7827 unsafe {
7828 let res = cudarc::driver::sys::cuMemsetD8Async(
7829 *d_count_only.device_ptr(),
7830 0,
7831 std::mem::size_of::<u32>(),
7832 cu_stream.cu_stream(),
7833 );
7834 if res != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
7835 return Err(XlogError::Kernel(format!(
7836 "cuMemsetD8Async (d_count_only) failed: {:?}",
7837 res
7838 )));
7839 }
7840 }
7841
7842 unsafe {
7847 let mut params: Vec<*mut c_void> = vec![
7848 (&left_packed.hashes).as_kernel_param(),
7849 num_left.as_kernel_param(),
7850 (&table.bucket_offsets).as_kernel_param(),
7851 (&table.bucket_counts).as_kernel_param(),
7852 (&table.bucket_entries).as_kernel_param(),
7853 (&table.bucket_entry_hashes).as_kernel_param(),
7854 table.bucket_mask.as_kernel_param(),
7855 (&left_packed.packed_keys).as_kernel_param(),
7856 (&right_packed.packed_keys).as_kernel_param(),
7857 left_packed.key_bytes.as_kernel_param(),
7858 (&d_dummy_left).as_kernel_param(),
7859 (&d_dummy_right).as_kernel_param(),
7860 (&d_count_only).as_kernel_param(),
7861 max_output_count_only.as_kernel_param(),
7862 ];
7863 probe_func
7864 .clone()
7865 .launch_in(enqueue, probe_config, &mut params)
7866 .map_err(|e| {
7867 XlogError::Kernel(format!(
7868 "hash_join_probe_v2 (count, on_stream) failed: {}",
7869 e
7870 ))
7871 })?;
7872 }
7873 Ok(())
7874 };
7875 let rec_count = unsafe { rec_count.enqueue_prepared_with(&cu_stream, enqueue) }
7877 .map_err(LaunchEnqueueError::into_xlog_error)?;
7878
7879 rec_count.commit().map_err(|e| {
7880 XlogError::Kernel(format!(
7881 "hash_join_inner_v2_recorded: count-pass commit failed: {}",
7882 e
7883 ))
7884 })?;
7885
7886 cu_stream.synchronize().map_err(|e| {
7888 XlogError::Kernel(format!(
7889 "hash_join_inner_v2_recorded: launch_stream sync (count read) failed: {}",
7890 e
7891 ))
7892 })?;
7893 let full_count = self.read_join_output_count_metadata(&d_count_only)? as u64;
7894 let requested = max_output
7895 .map(|limit| (limit as u64).min(full_count))
7896 .unwrap_or(full_count);
7897 if requested == 0 {
7898 let combined_schema = self.combine_schemas(left.schema(), right.schema());
7899 return self.create_empty_buffer(combined_schema);
7900 }
7901 if requested > u32::MAX as u64 {
7902 return Err(XlogError::Kernel(format!(
7903 "Join produced {} rows which exceeds the u32 index limit",
7904 requested
7905 )));
7906 }
7907 let max_output_u32 = requested as u32;
7908
7909 let d_output_left = self.memory.alloc::<u32>(max_output_u32 as usize)?;
7913 let d_output_right = self.memory.alloc::<u32>(max_output_u32 as usize)?;
7914 let d_output_count = self.memory.alloc::<u32>(1)?;
7915
7916 let mut rec_mat = LaunchRecorder::new_strict(launch_stream);
7917 rec_mat.read(&left_packed.hashes);
7918 rec_mat.read(&left_packed.packed_keys);
7919 rec_mat.read(&right_packed.packed_keys);
7920 rec_mat.read(&table.bucket_offsets);
7921 rec_mat.read(&table.bucket_counts);
7922 rec_mat.read(&table.bucket_entries);
7923 rec_mat.read(&table.bucket_entry_hashes);
7924 rec_mat.write(&d_output_left);
7925 rec_mat.write(&d_output_right);
7926 rec_mat.write(&d_output_count);
7927 rec_mat.preflight(runtime).map_err(|e| {
7928 XlogError::Kernel(format!(
7929 "hash_join_inner_v2_recorded: materialize-pass preflight failed: {}",
7930 e
7931 ))
7932 })?;
7933
7934 let enqueue = |enqueue: &crate::launch::CudaEnqueue<'_>| -> Result<()> {
7938 unsafe {
7939 let res = cudarc::driver::sys::cuMemsetD8Async(
7940 *d_output_count.device_ptr(),
7941 0,
7942 std::mem::size_of::<u32>(),
7943 cu_stream.cu_stream(),
7944 );
7945 if res != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
7946 return Err(XlogError::Kernel(format!(
7947 "cuMemsetD8Async (d_output_count) failed: {:?}",
7948 res
7949 )));
7950 }
7951 }
7952
7953 unsafe {
7955 let mut params: Vec<*mut c_void> = vec![
7956 (&left_packed.hashes).as_kernel_param(),
7957 num_left.as_kernel_param(),
7958 (&table.bucket_offsets).as_kernel_param(),
7959 (&table.bucket_counts).as_kernel_param(),
7960 (&table.bucket_entries).as_kernel_param(),
7961 (&table.bucket_entry_hashes).as_kernel_param(),
7962 table.bucket_mask.as_kernel_param(),
7963 (&left_packed.packed_keys).as_kernel_param(),
7964 (&right_packed.packed_keys).as_kernel_param(),
7965 left_packed.key_bytes.as_kernel_param(),
7966 (&d_output_left).as_kernel_param(),
7967 (&d_output_right).as_kernel_param(),
7968 (&d_output_count).as_kernel_param(),
7969 max_output_u32.as_kernel_param(),
7970 ];
7971 probe_func
7972 .clone()
7973 .launch_in(enqueue, probe_config, &mut params)
7974 .map_err(|e| {
7975 XlogError::Kernel(format!(
7976 "hash_join_probe_v2 (materialize, on_stream) failed: {}",
7977 e
7978 ))
7979 })?;
7980 }
7981 Ok(())
7982 };
7983 let rec_mat = unsafe { rec_mat.enqueue_prepared_with(&cu_stream, enqueue) }
7985 .map_err(LaunchEnqueueError::into_xlog_error)?;
7986
7987 rec_mat.commit().map_err(|e| {
7988 XlogError::Kernel(format!(
7989 "hash_join_inner_v2_recorded: materialize-pass commit failed: {}",
7990 e
7991 ))
7992 })?;
7993
7994 cu_stream.synchronize().map_err(|e| {
7996 XlogError::Kernel(format!(
7997 "hash_join_inner_v2_recorded: launch_stream sync (mat read) failed: {}",
7998 e
7999 ))
8000 })?;
8001 let result_count = (self.read_join_output_count_metadata(&d_output_count)? as u64)
8002 .min(max_output_u32 as u64);
8003 if result_count == 0 {
8004 let combined_schema = self.combine_schemas(left.schema(), right.schema());
8005 return self.create_empty_buffer(combined_schema);
8006 }
8007 let output_rows = result_count as u32;
8008
8009 let gathered_left = self.gather_buffer_by_indices_on_stream(
8012 left,
8013 &d_output_left,
8014 output_rows,
8015 &cu_stream,
8016 launch_stream,
8017 runtime,
8018 )?;
8019 let gathered_right = self.gather_buffer_by_indices_on_stream(
8020 right,
8021 &d_output_right,
8022 output_rows,
8023 &cu_stream,
8024 launch_stream,
8025 runtime,
8026 )?;
8027
8028 let combined_schema = self.combine_schemas(left.schema(), right.schema());
8029 let mut result_columns = Vec::with_capacity(combined_schema.arity());
8030 result_columns.extend(gathered_left.columns);
8031 result_columns.extend(gathered_right.columns);
8032 self.buffer_from_columns(result_columns, result_count, combined_schema)
8033 }
8034
8035 pub fn hash_join_inner_v2_count_scan_materialize_recorded(
8065 &self,
8066 left: &CudaBuffer,
8067 right: &CudaBuffer,
8068 left_keys: &[usize],
8069 right_keys: &[usize],
8070 max_output: Option<usize>,
8071 launch_stream: StreamId,
8072 ) -> Result<CudaBuffer> {
8073 if Self::use_csm_cuda_graph_env()? {
8074 if let Some(result) = self
8075 .hash_join_inner_v2_count_scan_materialize_cuda_graph_recorded(
8076 left,
8077 right,
8078 left_keys,
8079 right_keys,
8080 max_output,
8081 launch_stream,
8082 )?
8083 {
8084 return Ok(result);
8085 }
8086 self.csm_cuda_graph_fallbacks
8087 .fetch_add(1, Ordering::Relaxed);
8088 }
8089
8090 let runtime = self.memory.runtime().ok_or_else(|| {
8091 XlogError::Kernel(
8092 "hash_join_inner_v2_count_scan_materialize_recorded requires a \
8093 runtime-backed GpuMemoryManager"
8094 .to_string(),
8095 )
8096 })?;
8097 let cu_stream = runtime
8098 .stream_pool()
8099 .resolve(launch_stream)
8100 .ok_or_else(|| {
8101 XlogError::Kernel(format!(
8102 "hash_join_inner_v2_count_scan_materialize_recorded: launch_stream \
8103 StreamId({}) does not resolve",
8104 launch_stream.0
8105 ))
8106 })?;
8107
8108 let num_left = self.device_row_count(left)?;
8110 let num_right = self.device_row_count(right)?;
8111 if num_left > u32::MAX as usize || num_right > u32::MAX as usize {
8112 return Err(XlogError::Kernel(format!(
8113 "Join supports at most {} rows per side (left={}, right={})",
8114 u32::MAX,
8115 num_left,
8116 num_right
8117 )));
8118 }
8119 if num_left == 0 || num_right == 0 {
8120 let combined_schema = self.combine_schemas(left.schema(), right.schema());
8121 return self.create_empty_buffer(combined_schema);
8122 }
8123 if left_keys.is_empty() || right_keys.is_empty() {
8124 return Err(XlogError::Kernel(
8125 "Join requires at least one key column".to_string(),
8126 ));
8127 }
8128 if left_keys.len() != right_keys.len() {
8129 return Err(XlogError::Kernel(
8130 "Left and right key columns must have same length".to_string(),
8131 ));
8132 }
8133 if left_keys.len() > 4 {
8134 return Err(XlogError::Kernel(
8135 "hash_join_inner_v2_count_scan_materialize_recorded: max 4 key \
8136 columns supported (pack_keys constraint)"
8137 .to_string(),
8138 ));
8139 }
8140 for (&l, &r) in left_keys.iter().zip(right_keys.iter()) {
8141 let lt = left.schema().column_type(l);
8142 let rt = right.schema().column_type(r);
8143 if lt != rt {
8144 return Err(XlogError::Kernel(format!(
8145 "Key column type mismatch: left[{}]={:?}, right[{}]={:?}",
8146 l, lt, r, rt
8147 )));
8148 }
8149 }
8150
8151 let _num_left = num_left as u32;
8152 let probe_cap = left.num_rows() as u32;
8153
8154 let left_packed =
8156 self.pack_keys_gpu_on_stream(left, left_keys, &cu_stream, launch_stream, runtime)?;
8157 let right_packed =
8158 self.pack_keys_gpu_on_stream(right, right_keys, &cu_stream, launch_stream, runtime)?;
8159 let table = self.build_hash_table_v2_on_stream(
8160 &right_packed.hashes,
8161 num_right as u32,
8162 &cu_stream,
8163 launch_stream,
8164 runtime,
8165 )?;
8166
8167 let device = self.device.inner();
8168 let block_size = 256u32;
8169 let probe_grid = probe_cap.div_ceil(block_size);
8170 let probe_config = LaunchConfig {
8171 grid_dim: (probe_grid, 1, 1),
8172 block_dim: (block_size, 1, 1),
8173 shared_mem_bytes: 0,
8174 };
8175
8176 let per_probe_count = self.memory.alloc::<u32>(probe_cap as usize)?;
8178 let mut per_probe_offsets = self.memory.alloc::<u32>(probe_cap as usize)?;
8179 let d_logical_count = self.memory.alloc::<u32>(1)?;
8180 let d_overflow = self.memory.alloc::<u8>(1)?;
8181 let count_func = device
8187 .get_func(JOIN_MODULE, join_kernels::HASH_JOIN_PROBE_V2_COUNT_PER_ROW)
8188 .ok_or_else(|| {
8189 XlogError::Kernel("hash_join_probe_v2_count_per_row kernel not found".to_string())
8190 })?;
8191 let total_func = device
8192 .get_func(JOIN_MODULE, join_kernels::HASH_JOIN_TOTAL_FROM_SCAN)
8193 .ok_or_else(|| {
8194 XlogError::Kernel("hash_join_total_from_scan kernel not found".to_string())
8195 })?;
8196
8197 let mut scan_scratch = self.multiblock_scan_u32_scratch_for_len(probe_cap)?;
8198 let mut rec_count = LaunchRecorder::new_strict(launch_stream);
8199 rec_count.read(&left_packed.hashes);
8200 rec_count.read(&left_packed.packed_keys);
8201 rec_count.read(&right_packed.packed_keys);
8202 rec_count.read(&table.bucket_offsets);
8203 rec_count.read(&table.bucket_counts);
8204 rec_count.read(&table.bucket_entries);
8205 rec_count.read(&table.bucket_entry_hashes);
8206 rec_count.read(left.num_rows_device());
8207 rec_count.write(&per_probe_count);
8208 rec_count.write(&per_probe_offsets);
8209 rec_count.write(&d_logical_count);
8210 rec_count.write(&d_overflow);
8211 for level in scan_scratch.levels() {
8212 rec_count.read_write(level);
8213 }
8214 rec_count.preflight(runtime).map_err(|e| {
8215 XlogError::Kernel(format!("csm inner: count/scan preflight failed: {}", e))
8216 })?;
8217
8218 let enqueue = |enqueue: &crate::launch::CudaEnqueue<'_>| -> Result<()> {
8221 unsafe {
8223 let res = cudarc::driver::sys::cuMemsetD8Async(
8224 *d_overflow.device_ptr(),
8225 0,
8226 1,
8227 cu_stream.cu_stream(),
8228 );
8229 if res != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
8230 return Err(XlogError::Kernel(format!(
8231 "cuMemsetD8Async (d_overflow init) failed: {:?}",
8232 res
8233 )));
8234 }
8235 let res = cudarc::driver::sys::cuMemsetD8Async(
8236 *d_logical_count.device_ptr(),
8237 0,
8238 std::mem::size_of::<u32>(),
8239 cu_stream.cu_stream(),
8240 );
8241 if res != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
8242 return Err(XlogError::Kernel(format!(
8243 "cuMemsetD8Async (d_logical_count init) failed: {:?}",
8244 res
8245 )));
8246 }
8247 }
8248
8249 unsafe {
8250 count_func.clone().launch_in(
8251 enqueue,
8252 probe_config,
8253 (
8254 &left_packed.hashes,
8255 left.num_rows_device(),
8256 probe_cap,
8257 &table.bucket_offsets,
8258 &table.bucket_counts,
8259 &table.bucket_entries,
8260 &table.bucket_entry_hashes,
8261 table.bucket_mask,
8262 &left_packed.packed_keys,
8263 &right_packed.packed_keys,
8264 left_packed.key_bytes,
8265 &per_probe_count,
8266 ),
8267 )
8268 }
8269 .map_err(|e| {
8270 XlogError::Kernel(format!(
8271 "hash_join_probe_v2_count_per_row (on_stream) failed: {}",
8272 e
8273 ))
8274 })?;
8275
8276 unsafe {
8280 let res = cudarc::driver::sys::cuMemcpyDtoDAsync_v2(
8281 *per_probe_offsets.device_ptr(),
8282 *per_probe_count.device_ptr(),
8283 (probe_cap as usize) * std::mem::size_of::<u32>(),
8284 cu_stream.cu_stream(),
8285 );
8286 if res != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
8287 return Err(XlogError::Kernel(format!(
8288 "csm inner: cuMemcpyDtoDAsync (per_probe_count → offsets) failed: {:?}",
8289 res
8290 )));
8291 }
8292 }
8293 self.multiblock_scan_u32_inplace_on_stream(
8294 &mut per_probe_offsets,
8295 probe_cap,
8296 enqueue,
8297 &mut scan_scratch,
8298 )?;
8299
8300 let materialize_capacity_bound: u64 =
8307 (probe_cap as u64).saturating_mul(num_right as u64);
8308 let materialize_capacity_u32 = materialize_capacity_bound.min(u32::MAX as u64) as u32;
8309 unsafe {
8310 total_func.clone().launch_in(
8311 enqueue,
8312 LaunchConfig {
8313 grid_dim: (1, 1, 1),
8314 block_dim: (1, 1, 1),
8315 shared_mem_bytes: 0,
8316 },
8317 (
8318 &per_probe_offsets,
8319 &per_probe_count,
8320 left.num_rows_device(),
8321 probe_cap,
8322 materialize_capacity_u32,
8323 &d_logical_count,
8324 &d_overflow,
8325 ),
8326 )
8327 }
8328 .map_err(|e| {
8329 XlogError::Kernel(format!(
8330 "hash_join_total_from_scan (on_stream) failed: {}",
8331 e
8332 ))
8333 })?;
8334
8335 Ok(())
8336 };
8337 let rec_count = unsafe { rec_count.enqueue_prepared_with(&cu_stream, enqueue) }
8339 .map_err(LaunchEnqueueError::into_xlog_error)?;
8340
8341 rec_count.commit().map_err(|e| {
8342 XlogError::Kernel(format!("csm inner: count/scan commit failed: {}", e))
8343 })?;
8344
8345 cu_stream.synchronize().map_err(|e| {
8348 XlogError::Kernel(format!("csm inner: sync (total read) failed: {}", e))
8349 })?;
8350 let total = self.read_join_output_count_metadata(&d_logical_count)? as u64;
8351 let requested = max_output
8352 .map(|limit| (limit as u64).min(total))
8353 .unwrap_or(total);
8354 if requested == 0 {
8355 let combined_schema = self.combine_schemas(left.schema(), right.schema());
8356 return self.create_empty_buffer(combined_schema);
8357 }
8358 if requested > u32::MAX as u64 {
8359 return Err(XlogError::Kernel(format!(
8360 "Join produced {} rows which exceeds the u32 index limit",
8361 requested
8362 )));
8363 }
8364 let output_capacity = requested as u32;
8365
8366 let d_output_left = self.memory.alloc::<u32>(output_capacity as usize)?;
8375 let d_output_right = self.memory.alloc::<u32>(output_capacity as usize)?;
8376
8377 let mut rec_mat = LaunchRecorder::new_strict(launch_stream);
8378 rec_mat.read(&left_packed.hashes);
8379 rec_mat.read(&left_packed.packed_keys);
8380 rec_mat.read(&right_packed.packed_keys);
8381 rec_mat.read(&table.bucket_offsets);
8382 rec_mat.read(&table.bucket_counts);
8383 rec_mat.read(&table.bucket_entries);
8384 rec_mat.read(&table.bucket_entry_hashes);
8385 rec_mat.read(&per_probe_offsets);
8386 rec_mat.read(left.num_rows_device());
8387 rec_mat.write(&d_output_left);
8388 rec_mat.write(&d_output_right);
8389 rec_mat.write(&d_overflow);
8391 rec_mat.preflight(runtime).map_err(|e| {
8392 XlogError::Kernel(format!("csm inner: materialize preflight failed: {}", e))
8393 })?;
8394
8395 let materialize_func = device
8396 .get_func(JOIN_MODULE, join_kernels::HASH_JOIN_PROBE_V2_MATERIALIZE)
8397 .ok_or_else(|| {
8398 XlogError::Kernel("hash_join_probe_v2_materialize kernel not found".to_string())
8399 })?;
8400 let enqueue = |enqueue: &crate::launch::CudaEnqueue<'_>| -> Result<()> {
8403 unsafe {
8404 let mut params: Vec<*mut c_void> = vec![
8405 (&left_packed.hashes).as_kernel_param(),
8406 left.num_rows_device().as_kernel_param(),
8407 probe_cap.as_kernel_param(),
8408 (&table.bucket_offsets).as_kernel_param(),
8409 (&table.bucket_counts).as_kernel_param(),
8410 (&table.bucket_entries).as_kernel_param(),
8411 (&table.bucket_entry_hashes).as_kernel_param(),
8412 table.bucket_mask.as_kernel_param(),
8413 (&left_packed.packed_keys).as_kernel_param(),
8414 (&right_packed.packed_keys).as_kernel_param(),
8415 left_packed.key_bytes.as_kernel_param(),
8416 (&per_probe_offsets).as_kernel_param(),
8417 output_capacity.as_kernel_param(),
8418 (&d_output_left).as_kernel_param(),
8419 (&d_output_right).as_kernel_param(),
8420 (&d_overflow).as_kernel_param(),
8421 ];
8422 materialize_func
8423 .clone()
8424 .launch_in(enqueue, probe_config, &mut params)
8425 }
8426 .map_err(|e| {
8427 XlogError::Kernel(format!(
8428 "hash_join_probe_v2_materialize (on_stream) failed: {}",
8429 e
8430 ))
8431 })
8432 };
8433 let rec_mat = unsafe { rec_mat.enqueue_prepared_with(&cu_stream, enqueue) }
8435 .map_err(LaunchEnqueueError::into_xlog_error)?;
8436
8437 rec_mat.commit().map_err(|e| {
8438 XlogError::Kernel(format!("csm inner: materialize commit failed: {}", e))
8439 })?;
8440
8441 cu_stream.synchronize().map_err(|e| {
8442 XlogError::Kernel(format!("csm inner: sync (post-materialize) failed: {}", e))
8443 })?;
8444
8445 if requested == total && self.dtoh_scalar_untracked::<u8>(&d_overflow, 0)? != 0 {
8459 return Err(XlogError::Kernel(format!(
8460 "csm inner: join output truncated at {} rows (left_rows={}, \
8461 right_rows={}); refusing to return an incomplete relation",
8462 output_capacity, num_left, num_right
8463 )));
8464 }
8465
8466 let gathered_left = self.gather_buffer_by_indices_on_stream(
8469 left,
8470 &d_output_left,
8471 output_capacity,
8472 &cu_stream,
8473 launch_stream,
8474 runtime,
8475 )?;
8476 let gathered_right = self.gather_buffer_by_indices_on_stream(
8477 right,
8478 &d_output_right,
8479 output_capacity,
8480 &cu_stream,
8481 launch_stream,
8482 runtime,
8483 )?;
8484
8485 let combined_schema = self.combine_schemas(left.schema(), right.schema());
8486 let mut result_columns = Vec::with_capacity(combined_schema.arity());
8487 result_columns.extend(gathered_left.columns);
8488 result_columns.extend(gathered_right.columns);
8489 self.buffer_from_columns(result_columns, output_capacity as u64, combined_schema)
8490 }
8491
8492 fn hash_join_inner_v2_count_scan_materialize_cuda_graph_recorded(
8493 &self,
8494 left: &CudaBuffer,
8495 right: &CudaBuffer,
8496 left_keys: &[usize],
8497 right_keys: &[usize],
8498 max_output: Option<usize>,
8499 launch_stream: StreamId,
8500 ) -> Result<Option<CudaBuffer>> {
8501 let runtime = self.memory.runtime().ok_or_else(|| {
8502 XlogError::Kernel(
8503 "hash_join_inner_v2_count_scan_materialize_cuda_graph_recorded requires a \
8504 runtime-backed GpuMemoryManager"
8505 .to_string(),
8506 )
8507 })?;
8508 let cu_stream = runtime
8509 .stream_pool()
8510 .resolve(launch_stream)
8511 .ok_or_else(|| {
8512 XlogError::Kernel(format!(
8513 "hash_join_inner_v2_count_scan_materialize_cuda_graph_recorded: \
8514 launch_stream StreamId({}) does not resolve",
8515 launch_stream.0
8516 ))
8517 })?;
8518
8519 let num_left = self.device_row_count(left)?;
8520 let num_right = self.device_row_count(right)?;
8521 if num_left > u32::MAX as usize || num_right > u32::MAX as usize {
8522 return Err(XlogError::Kernel(format!(
8523 "Join supports at most {} rows per side (left={}, right={})",
8524 u32::MAX,
8525 num_left,
8526 num_right
8527 )));
8528 }
8529 if num_left == 0 || num_right == 0 || max_output == Some(0) {
8530 let combined_schema = self.combine_schemas(left.schema(), right.schema());
8531 return self.create_empty_buffer(combined_schema).map(Some);
8532 }
8533 if left_keys.is_empty() || right_keys.is_empty() {
8534 return Err(XlogError::Kernel(
8535 "Join requires at least one key column".to_string(),
8536 ));
8537 }
8538 if left_keys.len() != right_keys.len() {
8539 return Err(XlogError::Kernel(
8540 "Left and right key columns must have same length".to_string(),
8541 ));
8542 }
8543 if left_keys.len() > 4 {
8544 return Err(XlogError::Kernel(
8545 "hash_join_inner_v2_count_scan_materialize_cuda_graph_recorded: max 4 key \
8546 columns supported (pack_keys constraint)"
8547 .to_string(),
8548 ));
8549 }
8550 for (&l, &r) in left_keys.iter().zip(right_keys.iter()) {
8551 let lt = left.schema().column_type(l);
8552 let rt = right.schema().column_type(r);
8553 if lt != rt {
8554 return Err(XlogError::Kernel(format!(
8555 "Key column type mismatch: left[{}]={:?}, right[{}]={:?}",
8556 l, lt, r, rt
8557 )));
8558 }
8559 }
8560
8561 let logical_probe_cap = left.num_rows() as u32;
8562 let probe_cap = crate::cuda_graph::graph_capacity_class_u32(logical_probe_cap);
8563 let Some(output_capacity) =
8564 Self::csm_cuda_graph_output_capacity(logical_probe_cap, num_right as u32, max_output)?
8565 else {
8566 return Ok(None);
8567 };
8568 if output_capacity == 0 {
8569 let combined_schema = self.combine_schemas(left.schema(), right.schema());
8570 return self.create_empty_buffer(combined_schema).map(Some);
8571 }
8572
8573 let left_packed =
8574 self.pack_keys_gpu_on_stream(left, left_keys, &cu_stream, launch_stream, runtime)?;
8575 let right_packed =
8576 self.pack_keys_gpu_on_stream(right, right_keys, &cu_stream, launch_stream, runtime)?;
8577 let graph_key = CsmCudaGraphKey::inner(
8578 left_keys.len(),
8579 left_packed.key_bytes,
8580 probe_cap,
8581 output_capacity,
8582 )?;
8583 let table = self.build_hash_table_v2_on_stream(
8584 &right_packed.hashes,
8585 num_right as u32,
8586 &cu_stream,
8587 launch_stream,
8588 runtime,
8589 )?;
8590
8591 let device = self.device.inner();
8592 let block_size = 256u32;
8593 let probe_grid = probe_cap.div_ceil(block_size);
8594 let probe_config = LaunchConfig {
8595 grid_dim: (probe_grid, 1, 1),
8596 block_dim: (block_size, 1, 1),
8597 shared_mem_bytes: 0,
8598 };
8599
8600 let materialize_capacity_bound: u64 = (probe_cap as u64).saturating_mul(num_right as u64);
8601 let materialize_capacity_u32 = materialize_capacity_bound.min(u32::MAX as u64) as u32;
8602
8603 {
8604 let mut cache = self.csm_cuda_graph_cache.lock().map_err(|e| {
8605 XlogError::Kernel(format!("csm CUDA Graph cache lock poisoned: {}", e))
8606 })?;
8607 if let Some(entry) = cache.get_mut(&graph_key) {
8608 let result = self.launch_csm_cuda_graph_entry(
8609 entry,
8610 left,
8611 right,
8612 &left_packed,
8613 &right_packed,
8614 &table,
8615 max_output,
8616 materialize_capacity_u32,
8617 probe_config,
8618 &cu_stream,
8619 launch_stream,
8620 runtime,
8621 );
8622 if result.is_err() {
8623 let failed = cache.remove(&graph_key);
8624 drop(cache);
8625 drop(failed);
8626 return result.map(Some);
8627 }
8628 self.csm_cuda_graph_cache_hits
8629 .fetch_add(1, Ordering::Relaxed);
8630 return result.map(Some);
8631 }
8632 }
8633
8634 let per_probe_count = self.memory.alloc::<u32>(probe_cap as usize)?;
8635 let mut per_probe_offsets = self.memory.alloc::<u32>(probe_cap as usize)?;
8636 let d_logical_count = self.memory.alloc::<u32>(1)?;
8637 let d_overflow = self.memory.alloc::<u8>(1)?;
8638 let d_output_left = self.memory.alloc::<u32>(output_capacity as usize)?;
8639 let d_output_right = self.memory.alloc::<u32>(output_capacity as usize)?;
8640 let mut scan_scratch = self.multiblock_scan_u32_scratch_for_len(probe_cap)?;
8641
8642 let count_func = device
8643 .get_func(JOIN_MODULE, join_kernels::HASH_JOIN_PROBE_V2_COUNT_PER_ROW)
8644 .ok_or_else(|| {
8645 XlogError::Kernel("hash_join_probe_v2_count_per_row kernel not found".to_string())
8646 })?;
8647 let total_func = device
8648 .get_func(JOIN_MODULE, join_kernels::HASH_JOIN_TOTAL_FROM_SCAN)
8649 .ok_or_else(|| {
8650 XlogError::Kernel("hash_join_total_from_scan kernel not found".to_string())
8651 })?;
8652 let materialize_func = device
8653 .get_func(JOIN_MODULE, join_kernels::HASH_JOIN_PROBE_V2_MATERIALIZE)
8654 .ok_or_else(|| {
8655 XlogError::Kernel("hash_join_probe_v2_materialize kernel not found".to_string())
8656 })?;
8657
8658 let graph = CapturedCudaGraph::capture_on_stream(&cu_stream, || {
8659 let mut recorder = LaunchRecorder::new_strict(launch_stream);
8660 recorder.read(&left_packed.hashes);
8661 recorder.read(&left_packed.packed_keys);
8662 recorder.read(&right_packed.packed_keys);
8663 recorder.read(&table.bucket_offsets);
8664 recorder.read(&table.bucket_counts);
8665 recorder.read(&table.bucket_entries);
8666 recorder.read(&table.bucket_entry_hashes);
8667 recorder.read(left.num_rows_device());
8668 recorder.read_write(&per_probe_count);
8669 recorder.read_write(&per_probe_offsets);
8670 recorder.read_write(&d_logical_count);
8671 recorder.read_write(&d_overflow);
8672 recorder.write(&d_output_left);
8673 recorder.write(&d_output_right);
8674 for level in scan_scratch.levels() {
8675 recorder.read_write(level);
8676 }
8677 let capture_body = |enqueue: &crate::launch::CudaEnqueue<'_>| {
8680 unsafe {
8683 let res = cudarc::driver::sys::cuMemsetD8Async(
8684 *d_overflow.device_ptr(),
8685 0,
8686 1,
8687 cu_stream.cu_stream(),
8688 );
8689 if res != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
8690 return Err(XlogError::Kernel(format!(
8691 "csm inner graph: cuMemsetD8Async (d_overflow) failed: {:?}",
8692 res
8693 )));
8694 }
8695 let res = cudarc::driver::sys::cuMemsetD8Async(
8696 *d_logical_count.device_ptr(),
8697 0,
8698 std::mem::size_of::<u32>(),
8699 cu_stream.cu_stream(),
8700 );
8701 if res != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
8702 return Err(XlogError::Kernel(format!(
8703 "csm inner graph: cuMemsetD8Async (d_logical_count) failed: {:?}",
8704 res
8705 )));
8706 }
8707 }
8708
8709 unsafe {
8711 count_func.clone().launch_in(
8712 enqueue,
8713 probe_config,
8714 (
8715 &left_packed.hashes,
8716 left.num_rows_device(),
8717 probe_cap,
8718 &table.bucket_offsets,
8719 &table.bucket_counts,
8720 &table.bucket_entries,
8721 &table.bucket_entry_hashes,
8722 table.bucket_mask,
8723 &left_packed.packed_keys,
8724 &right_packed.packed_keys,
8725 left_packed.key_bytes,
8726 &per_probe_count,
8727 ),
8728 )
8729 }
8730 .map_err(|e| {
8731 XlogError::Kernel(format!("csm inner graph: count_per_row failed: {}", e))
8732 })?;
8733
8734 unsafe {
8736 let res = cudarc::driver::sys::cuMemcpyDtoDAsync_v2(
8737 *per_probe_offsets.device_ptr(),
8738 *per_probe_count.device_ptr(),
8739 (probe_cap as usize) * std::mem::size_of::<u32>(),
8740 cu_stream.cu_stream(),
8741 );
8742 if res != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
8743 return Err(XlogError::Kernel(format!(
8744 "csm inner graph: cuMemcpyDtoDAsync (count -> offsets) failed: {:?}",
8745 res
8746 )));
8747 }
8748 }
8749 self.multiblock_scan_u32_inplace_on_stream(
8750 &mut per_probe_offsets,
8751 probe_cap,
8752 enqueue,
8753 &mut scan_scratch,
8754 )?;
8755
8756 unsafe {
8758 total_func.clone().launch_in(
8759 enqueue,
8760 LaunchConfig {
8761 grid_dim: (1, 1, 1),
8762 block_dim: (1, 1, 1),
8763 shared_mem_bytes: 0,
8764 },
8765 (
8766 &per_probe_offsets,
8767 &per_probe_count,
8768 left.num_rows_device(),
8769 probe_cap,
8770 materialize_capacity_u32,
8771 &d_logical_count,
8772 &d_overflow,
8773 ),
8774 )
8775 }
8776 .map_err(|e| XlogError::Kernel(format!("csm inner graph: total failed: {}", e)))?;
8777
8778 unsafe {
8780 let mut params: Vec<*mut c_void> = vec![
8781 (&left_packed.hashes).as_kernel_param(),
8782 left.num_rows_device().as_kernel_param(),
8783 probe_cap.as_kernel_param(),
8784 (&table.bucket_offsets).as_kernel_param(),
8785 (&table.bucket_counts).as_kernel_param(),
8786 (&table.bucket_entries).as_kernel_param(),
8787 (&table.bucket_entry_hashes).as_kernel_param(),
8788 table.bucket_mask.as_kernel_param(),
8789 (&left_packed.packed_keys).as_kernel_param(),
8790 (&right_packed.packed_keys).as_kernel_param(),
8791 left_packed.key_bytes.as_kernel_param(),
8792 (&per_probe_offsets).as_kernel_param(),
8793 output_capacity.as_kernel_param(),
8794 (&d_output_left).as_kernel_param(),
8795 (&d_output_right).as_kernel_param(),
8796 (&d_overflow).as_kernel_param(),
8797 ];
8798 materialize_func
8799 .clone()
8800 .launch_in(enqueue, probe_config, &mut params)
8801 .map_err(|e| {
8802 XlogError::Kernel(format!("csm inner graph: materialize failed: {}", e))
8803 })?;
8804 }
8805 Ok(())
8806 };
8807 let enqueued = unsafe { recorder.enqueue(runtime, capture_body) }
8808 .map_err(LaunchEnqueueError::into_xlog_error)?;
8809 enqueued.commit().map_err(|error| {
8810 XlogError::Kernel(format!("graph capture memory commit failed: {error}"))
8811 })
8812 })?;
8813 let nodes = Self::csm_cuda_graph_nodes(&graph)?;
8814 let mut entry = CsmCudaGraphEntry {
8815 graph,
8816 nodes,
8817 per_probe_count,
8818 per_probe_offsets,
8819 d_logical_count,
8820 d_overflow,
8821 d_output_left,
8822 d_output_right,
8823 scan_scratch,
8824 probe_capacity: probe_cap,
8825 output_capacity,
8826 };
8827 self.csm_cuda_graph_captures.fetch_add(1, Ordering::Relaxed);
8828
8829 let result = self.launch_csm_cuda_graph_entry(
8830 &mut entry,
8831 left,
8832 right,
8833 &left_packed,
8834 &right_packed,
8835 &table,
8836 max_output,
8837 materialize_capacity_u32,
8838 probe_config,
8839 &cu_stream,
8840 launch_stream,
8841 runtime,
8842 )?;
8843 self.csm_cuda_graph_cache
8844 .lock()
8845 .map_err(|e| XlogError::Kernel(format!("csm CUDA Graph cache lock poisoned: {}", e)))?
8846 .insert(graph_key, entry);
8847 Ok(Some(result))
8848 }
8849
8850 #[expect(
8851 clippy::too_many_arguments,
8852 reason = "cached CUDA graph replay must update every input, table, capacity, launch, stream, and runtime binding explicitly"
8853 )]
8854 fn launch_csm_cuda_graph_entry(
8855 &self,
8856 entry: &mut CsmCudaGraphEntry,
8857 left: &CudaBuffer,
8858 right: &CudaBuffer,
8859 left_packed: &PackedKeyData,
8860 right_packed: &PackedKeyData,
8861 table: &JoinHashTableV2,
8862 max_output: Option<usize>,
8863 materialize_capacity_u32: u32,
8864 probe_config: LaunchConfig,
8865 cu_stream: &cudarc::driver::CudaStream,
8866 launch_stream: StreamId,
8867 runtime: &std::sync::Arc<crate::device_runtime::XlogDeviceRuntime>,
8868 ) -> Result<CudaBuffer> {
8869 let mut rec_graph = LaunchRecorder::new_strict(launch_stream);
8870 rec_graph.read(&left_packed.hashes);
8871 rec_graph.read(&left_packed.packed_keys);
8872 rec_graph.read(&right_packed.packed_keys);
8873 rec_graph.read(&table.bucket_offsets);
8874 rec_graph.read(&table.bucket_counts);
8875 rec_graph.read(&table.bucket_entries);
8876 rec_graph.read(&table.bucket_entry_hashes);
8877 rec_graph.read(left.num_rows_device());
8878 rec_graph.read_write(&entry.per_probe_count);
8879 rec_graph.read_write(&entry.per_probe_offsets);
8880 rec_graph.read_write(&entry.d_logical_count);
8881 rec_graph.read_write(&entry.d_overflow);
8882 rec_graph.write(&entry.d_output_left);
8883 rec_graph.write(&entry.d_output_right);
8884 for level in entry.scan_scratch.levels() {
8885 rec_graph.read_write(level);
8886 }
8887 rec_graph
8888 .preflight(runtime)
8889 .map_err(|e| XlogError::Kernel(format!("csm inner graph: preflight failed: {}", e)))?;
8890
8891 let probe_cap = entry.probe_capacity;
8892 let output_capacity = entry.output_capacity;
8893 if probe_config.grid_dim.0 != probe_cap.div_ceil(probe_config.block_dim.0) {
8894 return Err(XlogError::Kernel(format!(
8895 "csm CUDA Graph replay probe grid mismatch: graph probe_cap={}, grid={:?}",
8896 probe_cap, probe_config.grid_dim
8897 )));
8898 }
8899 if entry.nodes.node_count < 5 {
8900 return Err(XlogError::Kernel(format!(
8901 "csm CUDA Graph replay node inventory too small: {}",
8902 entry.nodes.node_count
8903 )));
8904 }
8905
8906 let mut count_params = entry.graph.kernel_node_params(entry.nodes.count)?;
8907 let mut total_params = entry.graph.kernel_node_params(entry.nodes.total)?;
8908 let mut materialize_params = entry.graph.kernel_node_params(entry.nodes.materialize)?;
8909 let mut count_args: Vec<*mut c_void> = vec![
8910 (&left_packed.hashes).as_kernel_param(),
8911 left.num_rows_device().as_kernel_param(),
8912 probe_cap.as_kernel_param(),
8913 (&table.bucket_offsets).as_kernel_param(),
8914 (&table.bucket_counts).as_kernel_param(),
8915 (&table.bucket_entries).as_kernel_param(),
8916 (&table.bucket_entry_hashes).as_kernel_param(),
8917 table.bucket_mask.as_kernel_param(),
8918 (&left_packed.packed_keys).as_kernel_param(),
8919 (&right_packed.packed_keys).as_kernel_param(),
8920 left_packed.key_bytes.as_kernel_param(),
8921 (&entry.per_probe_count).as_kernel_param(),
8922 ];
8923 let mut total_args: Vec<*mut c_void> = vec![
8924 (&entry.per_probe_offsets).as_kernel_param(),
8925 (&entry.per_probe_count).as_kernel_param(),
8926 left.num_rows_device().as_kernel_param(),
8927 probe_cap.as_kernel_param(),
8928 materialize_capacity_u32.as_kernel_param(),
8929 (&entry.d_logical_count).as_kernel_param(),
8930 (&entry.d_overflow).as_kernel_param(),
8931 ];
8932 let mut materialize_args: Vec<*mut c_void> = vec![
8933 (&left_packed.hashes).as_kernel_param(),
8934 left.num_rows_device().as_kernel_param(),
8935 probe_cap.as_kernel_param(),
8936 (&table.bucket_offsets).as_kernel_param(),
8937 (&table.bucket_counts).as_kernel_param(),
8938 (&table.bucket_entries).as_kernel_param(),
8939 (&table.bucket_entry_hashes).as_kernel_param(),
8940 table.bucket_mask.as_kernel_param(),
8941 (&left_packed.packed_keys).as_kernel_param(),
8942 (&right_packed.packed_keys).as_kernel_param(),
8943 left_packed.key_bytes.as_kernel_param(),
8944 (&entry.per_probe_offsets).as_kernel_param(),
8945 output_capacity.as_kernel_param(),
8946 (&entry.d_output_left).as_kernel_param(),
8947 (&entry.d_output_right).as_kernel_param(),
8948 (&entry.d_overflow).as_kernel_param(),
8949 ];
8950 count_params.kernelParams = count_args.as_mut_ptr();
8951 count_params.extra = std::ptr::null_mut();
8952 total_params.kernelParams = total_args.as_mut_ptr();
8953 total_params.extra = std::ptr::null_mut();
8954 materialize_params.kernelParams = materialize_args.as_mut_ptr();
8955 materialize_params.extra = std::ptr::null_mut();
8956 let enqueue = |enqueue: &crate::launch::CudaEnqueue<'_>| {
8957 unsafe {
8960 entry.graph.rebind_kernel_nodes_in(
8961 enqueue,
8962 &[
8963 crate::cuda_graph::KernelNodeUpdate {
8964 node: entry.nodes.count,
8965 params: &count_params,
8966 },
8967 crate::cuda_graph::KernelNodeUpdate {
8968 node: entry.nodes.total,
8969 params: &total_params,
8970 },
8971 crate::cuda_graph::KernelNodeUpdate {
8972 node: entry.nodes.materialize,
8973 params: &materialize_params,
8974 },
8975 ],
8976 )?;
8977 }
8978 entry.graph.launch_in(enqueue)
8979 };
8980 let rec_graph = unsafe { rec_graph.enqueue_prepared_with(cu_stream, enqueue) }
8982 .map_err(LaunchEnqueueError::into_xlog_error)?;
8983 self.csm_cuda_graph_launches.fetch_add(1, Ordering::Relaxed);
8984 rec_graph
8985 .commit()
8986 .map_err(|e| XlogError::Kernel(format!("csm inner graph: commit failed: {}", e)))?;
8987
8988 cu_stream.synchronize().map_err(|e| {
8989 XlogError::Kernel(format!("csm inner graph: sync (total read) failed: {}", e))
8990 })?;
8991 let total = self.read_join_output_count_metadata(&entry.d_logical_count)? as u64;
8992 let requested = max_output
8993 .map(|limit| (limit as u64).min(total))
8994 .unwrap_or(total);
8995 if requested == 0 {
8996 let combined_schema = self.combine_schemas(left.schema(), right.schema());
8997 return self.create_empty_buffer(combined_schema);
8998 }
8999 if requested > output_capacity as u64 {
9000 return Err(XlogError::Kernel(format!(
9001 "csm inner graph produced {} rows but graph output capacity is {}",
9002 requested, output_capacity
9003 )));
9004 }
9005 let output_rows = requested as u32;
9006
9007 let gathered_left = self.gather_buffer_by_indices_on_stream(
9008 left,
9009 &entry.d_output_left,
9010 output_rows,
9011 cu_stream,
9012 launch_stream,
9013 runtime,
9014 )?;
9015 let gathered_right = self.gather_buffer_by_indices_on_stream(
9016 right,
9017 &entry.d_output_right,
9018 output_rows,
9019 cu_stream,
9020 launch_stream,
9021 runtime,
9022 )?;
9023
9024 let combined_schema = self.combine_schemas(left.schema(), right.schema());
9025 let mut result_columns = Vec::with_capacity(combined_schema.arity());
9026 result_columns.extend(gathered_left.columns);
9027 result_columns.extend(gathered_right.columns);
9028 self.buffer_from_columns(result_columns, output_rows as u64, combined_schema)
9029 }
9030
9031 fn csm_cuda_graph_nodes(graph: &CapturedCudaGraph) -> Result<CsmCudaGraphNodes> {
9032 let nodes = graph.nodes()?;
9033 if nodes.len() < 5 {
9034 return Err(XlogError::Kernel(format!(
9035 "csm inner graph captured too few nodes: {}",
9036 nodes.len()
9037 )));
9038 }
9039 let kernel_nodes: Vec<_> = nodes
9040 .iter()
9041 .copied()
9042 .filter(|n| n.kind == CudaGraphNodeKind::Kernel)
9043 .collect();
9044 if kernel_nodes.len() < 3 {
9045 return Err(XlogError::Kernel(format!(
9046 "csm inner graph captured too few kernel nodes: {}",
9047 kernel_nodes.len()
9048 )));
9049 }
9050 Ok(CsmCudaGraphNodes {
9051 count: kernel_nodes[0],
9052 total: kernel_nodes[kernel_nodes.len() - 2],
9053 materialize: kernel_nodes[kernel_nodes.len() - 1],
9054 node_count: nodes.len(),
9055 })
9056 }
9057
9058 fn csm_cuda_graph_output_capacity(
9059 probe_cap: u32,
9060 num_right: u32,
9061 max_output: Option<usize>,
9062 ) -> Result<Option<u32>> {
9063 if let Some(limit) = max_output {
9064 let limit = u32::try_from(limit).map_err(|_| {
9065 XlogError::Kernel(format!(
9066 "csm CUDA Graph max_output {} exceeds u32::MAX",
9067 limit
9068 ))
9069 })?;
9070 return Ok(Some(crate::cuda_graph::graph_capacity_class_u32(limit)));
9071 }
9072
9073 let worst_case = (probe_cap as u64).saturating_mul(num_right as u64);
9074 if worst_case > u32::MAX as u64 {
9075 return Ok(None);
9076 }
9077 let auto_cap = std::env::var("XLOG_CSM_CUDA_GRAPH_AUTO_OUTPUT_CAP")
9078 .ok()
9079 .and_then(|v| v.parse::<u64>().ok())
9080 .unwrap_or(1_000_000);
9081 if worst_case <= auto_cap {
9082 Ok(Some(crate::cuda_graph::graph_capacity_class_u32(
9083 worst_case as u32,
9084 )))
9085 } else {
9086 Ok(None)
9087 }
9088 }
9089
9090 pub fn hash_join_left_outer_v2_count_scan_materialize_recorded(
9114 &self,
9115 left: &CudaBuffer,
9116 right: &CudaBuffer,
9117 left_keys: &[usize],
9118 right_keys: &[usize],
9119 max_output: Option<usize>,
9120 launch_stream: StreamId,
9121 ) -> Result<CudaBuffer> {
9122 let runtime = self.memory.runtime().ok_or_else(|| {
9123 XlogError::Kernel(
9124 "hash_join_left_outer_v2_count_scan_materialize_recorded requires a \
9125 runtime-backed GpuMemoryManager"
9126 .to_string(),
9127 )
9128 })?;
9129 let cu_stream = runtime
9130 .stream_pool()
9131 .resolve(launch_stream)
9132 .ok_or_else(|| {
9133 XlogError::Kernel(format!(
9134 "csm left_outer: launch_stream StreamId({}) does not resolve",
9135 launch_stream.0
9136 ))
9137 })?;
9138
9139 let num_left = self.device_row_count(left)?;
9141 let num_right = self.device_row_count(right)?;
9142 if num_left > u32::MAX as usize || num_right > u32::MAX as usize {
9143 return Err(XlogError::Kernel(format!(
9144 "Join supports at most {} rows per side (left={}, right={})",
9145 u32::MAX,
9146 num_left,
9147 num_right
9148 )));
9149 }
9150 if num_left == 0 {
9151 let combined_schema = self.combine_schemas(left.schema(), right.schema());
9152 return self.create_empty_buffer(combined_schema);
9153 }
9154 if num_right == 0 {
9155 return self.left_outer_with_nulls(left, right);
9160 }
9161 if left_keys.is_empty() || right_keys.is_empty() {
9162 return Err(XlogError::Kernel(
9163 "Join requires at least one key column".to_string(),
9164 ));
9165 }
9166 if left_keys.len() != right_keys.len() {
9167 return Err(XlogError::Kernel(
9168 "Left and right key columns must have same length".to_string(),
9169 ));
9170 }
9171 if left_keys.len() > 4 {
9172 return Err(XlogError::Kernel(
9173 "csm left_outer: max 4 key columns supported (pack_keys constraint)".to_string(),
9174 ));
9175 }
9176 for (&l, &r) in left_keys.iter().zip(right_keys.iter()) {
9177 let lt = left.schema().column_type(l);
9178 let rt = right.schema().column_type(r);
9179 if lt != rt {
9180 return Err(XlogError::Kernel(format!(
9181 "Key column type mismatch: left[{}]={:?}, right[{}]={:?}",
9182 l, lt, r, rt
9183 )));
9184 }
9185 }
9186
9187 let probe_cap = u32::try_from(num_left).map_err(|_| {
9195 XlogError::Kernel("csm left_outer: left row count exceeds u32::MAX".to_string())
9196 })?;
9197 let num_right_u32 = u32::try_from(num_right).map_err(|_| {
9198 XlogError::Kernel("csm left_outer: right row count exceeds u32::MAX".to_string())
9199 })?;
9200
9201 let left_packed =
9203 self.pack_keys_gpu_on_stream(left, left_keys, &cu_stream, launch_stream, runtime)?;
9204 let right_packed =
9205 self.pack_keys_gpu_on_stream(right, right_keys, &cu_stream, launch_stream, runtime)?;
9206 let table = self.build_hash_table_v2_on_stream(
9207 &right_packed.hashes,
9208 num_right_u32,
9209 &cu_stream,
9210 launch_stream,
9211 runtime,
9212 )?;
9213
9214 let device = self.device.inner();
9215 let block_size = 256u32;
9216 let probe_grid = probe_cap.div_ceil(block_size);
9217 let probe_config = LaunchConfig {
9218 grid_dim: (probe_grid, 1, 1),
9219 block_dim: (block_size, 1, 1),
9220 shared_mem_bytes: 0,
9221 };
9222
9223 let per_probe_count = self.memory.alloc::<u32>(probe_cap as usize)?;
9225 let mut per_probe_offsets = self.memory.alloc::<u32>(probe_cap as usize)?;
9226 let d_logical_count = self.memory.alloc::<u32>(1)?;
9227 let d_overflow = self.memory.alloc::<u8>(1)?;
9228 let count_func = device
9229 .get_func(JOIN_MODULE, join_kernels::HASH_JOIN_PROBE_V2_COUNT_PER_ROW)
9230 .ok_or_else(|| {
9231 XlogError::Kernel("hash_join_probe_v2_count_per_row kernel not found".to_string())
9232 })?;
9233 let total_func = device
9234 .get_func(JOIN_MODULE, join_kernels::HASH_JOIN_TOTAL_FROM_SCAN)
9235 .ok_or_else(|| {
9236 XlogError::Kernel("hash_join_total_from_scan kernel not found".to_string())
9237 })?;
9238
9239 let mut scan_scratch = self.multiblock_scan_u32_scratch_for_len(probe_cap)?;
9240 let mut rec_count = LaunchRecorder::new_strict(launch_stream);
9241 rec_count.read(&left_packed.hashes);
9242 rec_count.read(&left_packed.packed_keys);
9243 rec_count.read(&right_packed.packed_keys);
9244 rec_count.read(&table.bucket_offsets);
9245 rec_count.read(&table.bucket_counts);
9246 rec_count.read(&table.bucket_entries);
9247 rec_count.read(&table.bucket_entry_hashes);
9248 rec_count.read(left.num_rows_device());
9249 rec_count.write(&per_probe_count);
9250 rec_count.write(&per_probe_offsets);
9251 rec_count.write(&d_logical_count);
9252 rec_count.write(&d_overflow);
9253 for level in scan_scratch.levels() {
9254 rec_count.read_write(level);
9255 }
9256 rec_count.preflight(runtime).map_err(|e| {
9257 XlogError::Kernel(format!(
9258 "csm left_outer: count/scan preflight failed: {}",
9259 e
9260 ))
9261 })?;
9262
9263 let enqueue = |enqueue: &crate::launch::CudaEnqueue<'_>| -> Result<()> {
9266 unsafe {
9268 let res = cudarc::driver::sys::cuMemsetD8Async(
9269 *d_overflow.device_ptr(),
9270 0,
9271 1,
9272 cu_stream.cu_stream(),
9273 );
9274 if res != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
9275 return Err(XlogError::Kernel(format!(
9276 "csm left_outer: cuMemsetD8Async (d_overflow) failed: {:?}",
9277 res
9278 )));
9279 }
9280 let res = cudarc::driver::sys::cuMemsetD8Async(
9281 *d_logical_count.device_ptr(),
9282 0,
9283 std::mem::size_of::<u32>(),
9284 cu_stream.cu_stream(),
9285 );
9286 if res != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
9287 return Err(XlogError::Kernel(format!(
9288 "csm left_outer: cuMemsetD8Async (d_logical_count) failed: {:?}",
9289 res
9290 )));
9291 }
9292 }
9293
9294 unsafe {
9295 count_func.clone().launch_in(
9296 enqueue,
9297 probe_config,
9298 (
9299 &left_packed.hashes,
9300 left.num_rows_device(),
9301 probe_cap,
9302 &table.bucket_offsets,
9303 &table.bucket_counts,
9304 &table.bucket_entries,
9305 &table.bucket_entry_hashes,
9306 table.bucket_mask,
9307 &left_packed.packed_keys,
9308 &right_packed.packed_keys,
9309 left_packed.key_bytes,
9310 &per_probe_count,
9311 ),
9312 )
9313 }
9314 .map_err(|e| {
9315 XlogError::Kernel(format!(
9316 "hash_join_probe_v2_count_per_row (csm left_outer) failed: {}",
9317 e
9318 ))
9319 })?;
9320
9321 unsafe {
9325 let res = cudarc::driver::sys::cuMemcpyDtoDAsync_v2(
9326 *per_probe_offsets.device_ptr(),
9327 *per_probe_count.device_ptr(),
9328 (probe_cap as usize) * std::mem::size_of::<u32>(),
9329 cu_stream.cu_stream(),
9330 );
9331 if res != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
9332 return Err(XlogError::Kernel(format!(
9333 "csm left_outer: cuMemcpyDtoDAsync (count → offsets) failed: {:?}",
9334 res
9335 )));
9336 }
9337 }
9338 self.multiblock_scan_u32_inplace_on_stream(
9339 &mut per_probe_offsets,
9340 probe_cap,
9341 enqueue,
9342 &mut scan_scratch,
9343 )?;
9344
9345 let materialize_capacity_bound: u64 =
9347 (probe_cap as u64).saturating_mul(num_right as u64);
9348 let materialize_capacity_u32 = materialize_capacity_bound.min(u32::MAX as u64) as u32;
9349 unsafe {
9351 total_func.clone().launch_in(
9352 enqueue,
9353 LaunchConfig {
9354 grid_dim: (1, 1, 1),
9355 block_dim: (1, 1, 1),
9356 shared_mem_bytes: 0,
9357 },
9358 (
9359 &per_probe_offsets,
9360 &per_probe_count,
9361 left.num_rows_device(),
9362 probe_cap,
9363 materialize_capacity_u32,
9364 &d_logical_count,
9365 &d_overflow,
9366 ),
9367 )
9368 }
9369 .map_err(|e| {
9370 XlogError::Kernel(format!(
9371 "hash_join_total_from_scan (csm left_outer) failed: {}",
9372 e
9373 ))
9374 })?;
9375
9376 Ok(())
9377 };
9378 let rec_count = unsafe { rec_count.enqueue_prepared_with(&cu_stream, enqueue) }
9380 .map_err(LaunchEnqueueError::into_xlog_error)?;
9381
9382 rec_count.commit().map_err(|e| {
9383 XlogError::Kernel(format!("csm left_outer: count/scan commit failed: {}", e))
9384 })?;
9385
9386 cu_stream.synchronize().map_err(|e| {
9387 XlogError::Kernel(format!("csm left_outer: sync (count read) failed: {}", e))
9388 })?;
9389 let inner_total = self.read_join_output_count_metadata(&d_logical_count)? as u64;
9390 let inner_clamped = max_output
9391 .map(|limit| (limit as u64).min(inner_total))
9392 .unwrap_or(inner_total);
9393 if inner_clamped > u32::MAX as u64 {
9394 return Err(XlogError::Kernel(format!(
9395 "Join produced {} matched rows which exceeds the u32 index limit",
9396 inner_clamped
9397 )));
9398 }
9399 let inner_count_u32 = inner_clamped as u32;
9400
9401 let materialize_func = device
9403 .get_func(JOIN_MODULE, join_kernels::HASH_JOIN_PROBE_V2_MATERIALIZE)
9404 .ok_or_else(|| {
9405 XlogError::Kernel("hash_join_probe_v2_materialize kernel not found".to_string())
9406 })?;
9407 let d_output_left = self.memory.alloc::<u32>(inner_count_u32.max(1) as usize)?;
9408 let d_output_right = self.memory.alloc::<u32>(inner_count_u32.max(1) as usize)?;
9409
9410 let mut rec_mat = LaunchRecorder::new_strict(launch_stream);
9411 rec_mat.read(&left_packed.hashes);
9412 rec_mat.read(&left_packed.packed_keys);
9413 rec_mat.read(&right_packed.packed_keys);
9414 rec_mat.read(&table.bucket_offsets);
9415 rec_mat.read(&table.bucket_counts);
9416 rec_mat.read(&table.bucket_entries);
9417 rec_mat.read(&table.bucket_entry_hashes);
9418 rec_mat.read(&per_probe_offsets);
9419 rec_mat.read(left.num_rows_device());
9420 rec_mat.write(&d_output_left);
9421 rec_mat.write(&d_output_right);
9422 rec_mat.write(&d_overflow);
9424 rec_mat.preflight(runtime).map_err(|e| {
9425 XlogError::Kernel(format!(
9426 "csm left_outer: materialize preflight failed: {}",
9427 e
9428 ))
9429 })?;
9430 if inner_count_u32 > 0 {
9431 let enqueue = |enqueue: &crate::launch::CudaEnqueue<'_>| -> Result<()> {
9433 unsafe {
9434 let mut params: Vec<*mut c_void> = vec![
9435 (&left_packed.hashes).as_kernel_param(),
9436 left.num_rows_device().as_kernel_param(),
9437 probe_cap.as_kernel_param(),
9438 (&table.bucket_offsets).as_kernel_param(),
9439 (&table.bucket_counts).as_kernel_param(),
9440 (&table.bucket_entries).as_kernel_param(),
9441 (&table.bucket_entry_hashes).as_kernel_param(),
9442 table.bucket_mask.as_kernel_param(),
9443 (&left_packed.packed_keys).as_kernel_param(),
9444 (&right_packed.packed_keys).as_kernel_param(),
9445 left_packed.key_bytes.as_kernel_param(),
9446 (&per_probe_offsets).as_kernel_param(),
9447 inner_count_u32.as_kernel_param(),
9448 (&d_output_left).as_kernel_param(),
9449 (&d_output_right).as_kernel_param(),
9450 (&d_overflow).as_kernel_param(),
9451 ];
9452 materialize_func
9453 .clone()
9454 .launch_in(enqueue, probe_config, &mut params)
9455 }
9456 .map_err(|e| {
9457 XlogError::Kernel(format!(
9458 "hash_join_probe_v2_materialize (csm left_outer) failed: {}",
9459 e
9460 ))
9461 })
9462 };
9463 let rec_mat = unsafe { rec_mat.enqueue_prepared_with(&cu_stream, enqueue) }
9465 .map_err(LaunchEnqueueError::into_xlog_error)?;
9466 rec_mat.commit().map_err(|e| {
9467 XlogError::Kernel(format!("csm left_outer: materialize commit failed: {}", e))
9468 })?;
9469 } else {
9470 rec_mat.abort().map_err(|e| {
9471 XlogError::Kernel(format!("csm left_outer: materialize abort failed: {}", e))
9472 })?;
9473 }
9474
9475 let d_unmatched_mask = self.memory.alloc::<u8>(probe_cap as usize)?;
9477 let unmatched_mask_func = device
9478 .get_func(JOIN_MODULE, join_kernels::HASH_JOIN_CSM_UNMATCHED_MASK)
9479 .ok_or_else(|| {
9480 XlogError::Kernel("hash_join_csm_unmatched_mask kernel not found".to_string())
9481 })?;
9482 let mut rec_um = LaunchRecorder::new_strict(launch_stream);
9483 rec_um.read(&per_probe_count);
9484 rec_um.read(left.num_rows_device());
9485 rec_um.write(&d_unmatched_mask);
9486 rec_um.preflight(runtime).map_err(|e| {
9487 XlogError::Kernel(format!(
9488 "csm left_outer: unmatched mask preflight failed: {}",
9489 e
9490 ))
9491 })?;
9492 let enqueue = |enqueue: &crate::launch::CudaEnqueue<'_>| -> Result<()> {
9494 unsafe {
9495 unmatched_mask_func.clone().launch_in(
9496 enqueue,
9497 probe_config,
9498 (
9499 &per_probe_count,
9500 left.num_rows_device(),
9501 probe_cap,
9502 &d_unmatched_mask,
9503 ),
9504 )
9505 }
9506 .map_err(|e| {
9507 XlogError::Kernel(format!(
9508 "hash_join_csm_unmatched_mask (on_stream) failed: {}",
9509 e
9510 ))
9511 })
9512 };
9513 let rec_um = unsafe { rec_um.enqueue_prepared_with(&cu_stream, enqueue) }
9515 .map_err(LaunchEnqueueError::into_xlog_error)?;
9516 rec_um.commit().map_err(|e| {
9517 XlogError::Kernel(format!(
9518 "csm left_outer: unmatched mask commit failed: {}",
9519 e
9520 ))
9521 })?;
9522
9523 let unmatched_left = self.compact_buffer_by_device_mask_counted_recorded(
9524 left,
9525 &d_unmatched_mask,
9526 launch_stream,
9527 )?;
9528 let unmatched_rows = self.device_row_count(&unmatched_left)? as u64;
9529 let total_rows = (inner_count_u32 as u64) + unmatched_rows;
9530
9531 let combined_schema = self.combine_schemas(left.schema(), right.schema());
9532 if total_rows == 0 {
9533 return self.create_empty_buffer(combined_schema);
9534 }
9535
9536 let inner_left_buf;
9538 let inner_right_buf;
9539 if inner_count_u32 > 0 {
9540 inner_left_buf = Some(self.gather_buffer_by_indices_on_stream(
9541 left,
9542 &d_output_left,
9543 inner_count_u32,
9544 &cu_stream,
9545 launch_stream,
9546 runtime,
9547 )?);
9548 inner_right_buf = Some(self.gather_buffer_by_indices_on_stream(
9549 right,
9550 &d_output_right,
9551 inner_count_u32,
9552 &cu_stream,
9553 launch_stream,
9554 runtime,
9555 )?);
9556 } else {
9557 inner_left_buf = None;
9558 inner_right_buf = None;
9559 }
9560
9561 let inner_rows = inner_count_u32 as u64;
9565 let mut left_outputs = Vec::with_capacity(left.arity());
9566 for col_idx in 0..left.arity() {
9567 let elem_size = left
9568 .schema()
9569 .column_type(col_idx)
9570 .map(|t| t.size_bytes())
9571 .unwrap_or(4);
9572 let inner_bytes = (inner_rows as usize)
9573 .checked_mul(elem_size)
9574 .ok_or_else(|| XlogError::Kernel("csm left_outer: inner_bytes overflow".into()))?;
9575 let unmatched_bytes = (unmatched_rows as usize)
9576 .checked_mul(elem_size)
9577 .ok_or_else(|| {
9578 XlogError::Kernel("csm left_outer: unmatched_bytes overflow".into())
9579 })?;
9580 let total_bytes = inner_bytes
9581 .checked_add(unmatched_bytes)
9582 .ok_or_else(|| XlogError::Kernel("csm left_outer: total_bytes overflow".into()))?;
9583 left_outputs.push((
9584 self.memory.alloc::<u8>(total_bytes)?,
9585 inner_bytes,
9586 unmatched_bytes,
9587 ));
9588 }
9589
9590 let mut right_outputs = Vec::with_capacity(right.arity());
9591 for col_idx in 0..right.arity() {
9592 let elem_size = right
9593 .schema()
9594 .column_type(col_idx)
9595 .map(|t| t.size_bytes())
9596 .unwrap_or(4);
9597 let inner_bytes = (inner_rows as usize)
9598 .checked_mul(elem_size)
9599 .ok_or_else(|| {
9600 XlogError::Kernel("csm left_outer: right inner_bytes overflow".into())
9601 })?;
9602 let unmatched_bytes = (unmatched_rows as usize)
9603 .checked_mul(elem_size)
9604 .ok_or_else(|| {
9605 XlogError::Kernel("csm left_outer: right unmatched_bytes overflow".into())
9606 })?;
9607 let total_bytes = inner_bytes.checked_add(unmatched_bytes).ok_or_else(|| {
9608 XlogError::Kernel("csm left_outer: right total_bytes overflow".into())
9609 })?;
9610 right_outputs.push((
9611 self.memory.alloc::<u8>(total_bytes)?,
9612 inner_bytes,
9613 total_bytes,
9614 ));
9615 }
9616
9617 let mut rec_d = LaunchRecorder::new_strict(launch_stream);
9618 for col_idx in 0..unmatched_left.columns.len() {
9619 let c = unmatched_left.column(col_idx).ok_or_else(|| {
9620 XlogError::Kernel(format!("unmatched_left col {} not found", col_idx))
9621 })?;
9622 rec_d.read_column(c);
9623 }
9624 if let Some(b) = inner_left_buf.as_ref() {
9625 for col_idx in 0..b.columns.len() {
9626 let c = b.column(col_idx).ok_or_else(|| {
9627 XlogError::Kernel(format!("inner_left col {} not found", col_idx))
9628 })?;
9629 rec_d.read_column(c);
9630 }
9631 }
9632 if let Some(b) = inner_right_buf.as_ref() {
9633 for col_idx in 0..b.columns.len() {
9634 let c = b.column(col_idx).ok_or_else(|| {
9635 XlogError::Kernel(format!("inner_right col {} not found", col_idx))
9636 })?;
9637 rec_d.read_column(c);
9638 }
9639 }
9640 for (out_col, _, _) in &left_outputs {
9641 rec_d.write(out_col);
9642 }
9643 for (out_col, _, _) in &right_outputs {
9644 rec_d.write(out_col);
9645 }
9646 rec_d.preflight(runtime).map_err(|e| {
9647 XlogError::Kernel(format!("csm left_outer: concat preflight failed: {}", e))
9648 })?;
9649 if left_outputs.is_empty() && right_outputs.is_empty() {
9650 rec_d.abort().map_err(|e| {
9651 XlogError::Kernel(format!("csm left_outer: concat abort failed: {}", e))
9652 })?;
9653 } else {
9654 let enqueue = |enqueue: &crate::launch::CudaEnqueue<'_>| -> Result<()> {
9655 for (col_idx, (out_col, inner_bytes, unmatched_bytes)) in
9656 left_outputs.iter().enumerate()
9657 {
9658 let dst_ptr = *out_col.device_ptr();
9659 if *inner_bytes > 0 {
9660 let src_col = inner_left_buf
9661 .as_ref()
9662 .ok_or_else(|| XlogError::Kernel("inner_left buffer missing".into()))?
9663 .column(col_idx)
9664 .ok_or_else(|| XlogError::Kernel("inner_left col missing".into()))?;
9665 unsafe {
9667 let res = cudarc::driver::sys::cuMemcpyDtoDAsync_v2(
9668 dst_ptr,
9669 *src_col.device_ptr(),
9670 *inner_bytes,
9671 enqueue.stream().cu_stream(),
9672 );
9673 if res != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
9674 return Err(XlogError::Kernel(format!(
9675 "csm left_outer: dtod inner_left col {} failed: {:?}",
9676 col_idx, res
9677 )));
9678 }
9679 }
9680 }
9681 if *unmatched_bytes > 0 {
9682 let src_col = unmatched_left.column(col_idx).ok_or_else(|| {
9683 XlogError::Kernel(format!("unmatched_left col {} not found", col_idx))
9684 })?;
9685 unsafe {
9687 let res = cudarc::driver::sys::cuMemcpyDtoDAsync_v2(
9688 dst_ptr + *inner_bytes as u64,
9689 *src_col.device_ptr(),
9690 *unmatched_bytes,
9691 cu_stream.cu_stream(),
9692 );
9693 if res != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
9694 return Err(XlogError::Kernel(format!(
9695 "csm left_outer: dtod unmatched col {} failed: {:?}",
9696 col_idx, res
9697 )));
9698 }
9699 }
9700 }
9701 }
9702
9703 for (col_idx, (out_col, inner_bytes, total_bytes)) in
9704 right_outputs.iter().enumerate()
9705 {
9706 let dst_ptr = *out_col.device_ptr();
9707 unsafe {
9709 let res = cudarc::driver::sys::cuMemsetD8Async(
9710 dst_ptr,
9711 0,
9712 *total_bytes,
9713 cu_stream.cu_stream(),
9714 );
9715 if res != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
9716 return Err(XlogError::Kernel(format!(
9717 "csm left_outer: zero-fill right col {} failed: {:?}",
9718 col_idx, res
9719 )));
9720 }
9721 }
9722 if *inner_bytes > 0 {
9723 let src_col = inner_right_buf
9724 .as_ref()
9725 .ok_or_else(|| XlogError::Kernel("inner_right buffer missing".into()))?
9726 .column(col_idx)
9727 .ok_or_else(|| XlogError::Kernel("inner_right col missing".into()))?;
9728 unsafe {
9730 let res = cudarc::driver::sys::cuMemcpyDtoDAsync_v2(
9731 dst_ptr,
9732 *src_col.device_ptr(),
9733 *inner_bytes,
9734 cu_stream.cu_stream(),
9735 );
9736 if res != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
9737 return Err(XlogError::Kernel(format!(
9738 "csm left_outer: dtod inner_right col {} failed: {:?}",
9739 col_idx, res
9740 )));
9741 }
9742 }
9743 }
9744 }
9745 Ok(())
9746 };
9747 let rec_d = unsafe { rec_d.enqueue_prepared_with(&cu_stream, enqueue) }
9749 .map_err(LaunchEnqueueError::into_xlog_error)?;
9750 rec_d.commit().map_err(|e| {
9751 XlogError::Kernel(format!("csm left_outer: concat commit failed: {}", e))
9752 })?;
9753 }
9754
9755 let result_columns: Vec<CudaColumn> = left_outputs
9756 .into_iter()
9757 .chain(right_outputs)
9758 .map(|(column, _, _)| column.into())
9759 .collect();
9760
9761 if total_rows > u32::MAX as u64 {
9770 return Err(XlogError::Kernel(format!(
9771 "csm left_outer: output row count {} exceeds u32::MAX",
9772 total_rows
9773 )));
9774 }
9775 let total_rows_u32 = total_rows as u32;
9776 let d_num_rows = self.upload_device_row_count(total_rows_u32)?;
9777 Ok(CudaBuffer::from_columns_with_host_count(
9778 result_columns,
9779 total_rows,
9780 d_num_rows,
9781 combined_schema,
9782 total_rows_u32,
9783 ))
9784 }
9785
9786 #[expect(
9811 clippy::too_many_arguments,
9812 reason = "the indexed recorded inner join exposes both relation/key pairs, cache identity, output bound, and stream"
9813 )]
9814 pub fn hash_join_inner_v2_with_index_count_scan_materialize_recorded(
9815 &self,
9816 left: &CudaBuffer,
9817 right: &CudaBuffer,
9818 left_keys: &[usize],
9819 right_keys: &[usize],
9820 index: &crate::provider::JoinIndexV2,
9821 max_output: Option<usize>,
9822 launch_stream: StreamId,
9823 ) -> Result<CudaBuffer> {
9824 use crate::launch::LaunchRecorder;
9825
9826 let runtime = self.memory.runtime().ok_or_else(|| {
9827 XlogError::Kernel(
9828 "hash_join_inner_v2_with_index_count_scan_materialize_recorded requires \
9829 a runtime-backed GpuMemoryManager"
9830 .to_string(),
9831 )
9832 })?;
9833 let cu_stream = runtime
9834 .stream_pool()
9835 .resolve(launch_stream)
9836 .ok_or_else(|| {
9837 XlogError::Kernel(format!(
9838 "indexed CSM inner: launch_stream StreamId({}) does not resolve",
9839 launch_stream.0
9840 ))
9841 })?;
9842
9843 let left_rows = self.device_row_count(left)?;
9846 let right_rows = self.device_row_count(right)?;
9847 if left_rows > u32::MAX as usize || right_rows > u32::MAX as usize {
9848 return Err(XlogError::Kernel(format!(
9849 "Join supports at most {} rows per side (left={}, right={})",
9850 u32::MAX,
9851 left_rows,
9852 right_rows
9853 )));
9854 }
9855 if left_rows == 0 || right_rows == 0 {
9856 let combined_schema = self.combine_schemas(left.schema(), right.schema());
9857 return self.create_empty_buffer(combined_schema);
9858 }
9859 if left_keys.is_empty() || right_keys.is_empty() {
9860 return Err(XlogError::Kernel(
9861 "Join requires at least one key column".to_string(),
9862 ));
9863 }
9864 if left_keys.len() != right_keys.len() {
9865 return Err(XlogError::Kernel(
9866 "Left and right key columns must have same length".to_string(),
9867 ));
9868 }
9869 if left_keys.len() > 4 {
9870 return Err(XlogError::Kernel(
9871 "indexed CSM inner: max 4 key columns supported (pack_keys constraint)".to_string(),
9872 ));
9873 }
9874 for (&l, &r) in left_keys.iter().zip(right_keys.iter()) {
9875 if l >= left.arity() {
9876 return Err(XlogError::Kernel(format!(
9877 "Left key column index {} out of bounds (arity {})",
9878 l,
9879 left.arity()
9880 )));
9881 }
9882 if r >= right.arity() {
9883 return Err(XlogError::Kernel(format!(
9884 "Right key column index {} out of bounds (arity {})",
9885 r,
9886 right.arity()
9887 )));
9888 }
9889 let lt = left.schema().column_type(l);
9890 let rt = right.schema().column_type(r);
9891 if lt != rt {
9892 return Err(XlogError::Kernel(format!(
9893 "Key column type mismatch: left[{}]={:?}, right[{}]={:?}",
9894 l, lt, r, rt
9895 )));
9896 }
9897 }
9898 if index.right_num_rows() != right_rows as u32 {
9899 return Err(XlogError::Kernel(
9900 "Join index row count does not match right relation".to_string(),
9901 ));
9902 }
9903 if index.right_keys() != right_keys {
9904 return Err(XlogError::Kernel(
9905 "Join index key columns do not match requested right_keys".to_string(),
9906 ));
9907 }
9908
9909 let probe_cap = left.num_rows() as u32;
9910 let table = &index.table;
9911
9912 let left_packed =
9915 self.pack_keys_gpu_on_stream(left, left_keys, &cu_stream, launch_stream, runtime)?;
9916 if left_packed.key_bytes != index.key_bytes {
9917 return Err(XlogError::Kernel(
9918 "Join key byte width mismatch between probe and cached index".to_string(),
9919 ));
9920 }
9921
9922 let device = self.device.inner();
9923 let block_size = 256u32;
9924 let probe_grid = probe_cap.div_ceil(block_size);
9925 let probe_config = LaunchConfig {
9926 grid_dim: (probe_grid, 1, 1),
9927 block_dim: (block_size, 1, 1),
9928 shared_mem_bytes: 0,
9929 };
9930
9931 let per_probe_count = self.memory.alloc::<u32>(probe_cap as usize)?;
9933 let mut per_probe_offsets = self.memory.alloc::<u32>(probe_cap as usize)?;
9934 let d_logical_count = self.memory.alloc::<u32>(1)?;
9935 let d_overflow = self.memory.alloc::<u8>(1)?;
9936
9937 let count_func = device
9942 .get_func(JOIN_MODULE, join_kernels::HASH_JOIN_PROBE_V2_COUNT_PER_ROW)
9943 .ok_or_else(|| {
9944 XlogError::Kernel("hash_join_probe_v2_count_per_row kernel not found".to_string())
9945 })?;
9946 let total_func = device
9947 .get_func(JOIN_MODULE, join_kernels::HASH_JOIN_TOTAL_FROM_SCAN)
9948 .ok_or_else(|| {
9949 XlogError::Kernel("hash_join_total_from_scan kernel not found".to_string())
9950 })?;
9951
9952 let mut scan_scratch = self.multiblock_scan_u32_scratch_for_len(probe_cap)?;
9953 let mut rec_count = LaunchRecorder::new_strict(launch_stream);
9954 rec_count.read(&left_packed.hashes);
9955 rec_count.read(&left_packed.packed_keys);
9956 rec_count.read(&index.packed_keys);
9957 rec_count.read(&table.bucket_offsets);
9958 rec_count.read(&table.bucket_counts);
9959 rec_count.read(&table.bucket_entries);
9960 rec_count.read(&table.bucket_entry_hashes);
9961 rec_count.read(left.num_rows_device());
9962 rec_count.write(&per_probe_count);
9963 rec_count.write(&per_probe_offsets);
9964 rec_count.write(&d_logical_count);
9965 rec_count.write(&d_overflow);
9966 for level in scan_scratch.levels() {
9967 rec_count.read_write(level);
9968 }
9969 rec_count.preflight(runtime).map_err(|e| {
9970 XlogError::Kernel(format!(
9971 "indexed CSM inner: count/scan preflight failed: {}",
9972 e
9973 ))
9974 })?;
9975
9976 let enqueue = |enqueue: &crate::launch::CudaEnqueue<'_>| -> Result<()> {
9979 unsafe {
9981 let res = cudarc::driver::sys::cuMemsetD8Async(
9982 *d_overflow.device_ptr(),
9983 0,
9984 1,
9985 cu_stream.cu_stream(),
9986 );
9987 if res != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
9988 return Err(XlogError::Kernel(format!(
9989 "indexed CSM inner: cuMemsetD8Async (d_overflow) failed: {:?}",
9990 res
9991 )));
9992 }
9993 let res = cudarc::driver::sys::cuMemsetD8Async(
9994 *d_logical_count.device_ptr(),
9995 0,
9996 std::mem::size_of::<u32>(),
9997 cu_stream.cu_stream(),
9998 );
9999 if res != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
10000 return Err(XlogError::Kernel(format!(
10001 "indexed CSM inner: cuMemsetD8Async (d_logical_count) failed: {:?}",
10002 res
10003 )));
10004 }
10005 }
10006
10007 unsafe {
10008 count_func.clone().launch_in(
10009 enqueue,
10010 probe_config,
10011 (
10012 &left_packed.hashes,
10013 left.num_rows_device(),
10014 probe_cap,
10015 &table.bucket_offsets,
10016 &table.bucket_counts,
10017 &table.bucket_entries,
10018 &table.bucket_entry_hashes,
10019 table.bucket_mask,
10020 &left_packed.packed_keys,
10021 &index.packed_keys,
10022 index.key_bytes,
10023 &per_probe_count,
10024 ),
10025 )
10026 }
10027 .map_err(|e| {
10028 XlogError::Kernel(format!(
10029 "hash_join_probe_v2_count_per_row (on_stream, indexed) failed: {}",
10030 e
10031 ))
10032 })?;
10033
10034 unsafe {
10038 let res = cudarc::driver::sys::cuMemcpyDtoDAsync_v2(
10039 *per_probe_offsets.device_ptr(),
10040 *per_probe_count.device_ptr(),
10041 (probe_cap as usize) * std::mem::size_of::<u32>(),
10042 cu_stream.cu_stream(),
10043 );
10044 if res != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
10045 return Err(XlogError::Kernel(format!(
10046 "indexed CSM inner: cuMemcpyDtoDAsync (count → offsets) failed: {:?}",
10047 res
10048 )));
10049 }
10050 }
10051 self.multiblock_scan_u32_inplace_on_stream(
10052 &mut per_probe_offsets,
10053 probe_cap,
10054 enqueue,
10055 &mut scan_scratch,
10056 )?;
10057
10058 let materialize_capacity_bound: u64 =
10060 (probe_cap as u64).saturating_mul(right_rows as u64);
10061 let materialize_capacity_u32 = materialize_capacity_bound.min(u32::MAX as u64) as u32;
10062 unsafe {
10064 total_func.clone().launch_in(
10065 enqueue,
10066 LaunchConfig {
10067 grid_dim: (1, 1, 1),
10068 block_dim: (1, 1, 1),
10069 shared_mem_bytes: 0,
10070 },
10071 (
10072 &per_probe_offsets,
10073 &per_probe_count,
10074 left.num_rows_device(),
10075 probe_cap,
10076 materialize_capacity_u32,
10077 &d_logical_count,
10078 &d_overflow,
10079 ),
10080 )
10081 }
10082 .map_err(|e| {
10083 XlogError::Kernel(format!(
10084 "hash_join_total_from_scan (on_stream, indexed) failed: {}",
10085 e
10086 ))
10087 })?;
10088
10089 Ok(())
10090 };
10091 let rec_count = unsafe { rec_count.enqueue_prepared_with(&cu_stream, enqueue) }
10093 .map_err(LaunchEnqueueError::into_xlog_error)?;
10094
10095 rec_count.commit().map_err(|e| {
10096 XlogError::Kernel(format!(
10097 "indexed CSM inner: count/scan commit failed: {}",
10098 e
10099 ))
10100 })?;
10101
10102 cu_stream.synchronize().map_err(|e| {
10103 XlogError::Kernel(format!(
10104 "indexed CSM inner: sync (total read) failed: {}",
10105 e
10106 ))
10107 })?;
10108 let total = self.read_join_output_count_metadata(&d_logical_count)? as u64;
10109 let requested = max_output
10110 .map(|limit| (limit as u64).min(total))
10111 .unwrap_or(total);
10112 if requested == 0 {
10113 let combined_schema = self.combine_schemas(left.schema(), right.schema());
10114 return self.create_empty_buffer(combined_schema);
10115 }
10116 if requested > u32::MAX as u64 {
10117 return Err(XlogError::Kernel(format!(
10118 "Join produced {} rows which exceeds the u32 index limit",
10119 requested
10120 )));
10121 }
10122 let output_capacity = requested as u32;
10123
10124 let d_output_left = self.memory.alloc::<u32>(output_capacity as usize)?;
10126 let d_output_right = self.memory.alloc::<u32>(output_capacity as usize)?;
10127
10128 let materialize_func = device
10129 .get_func(JOIN_MODULE, join_kernels::HASH_JOIN_PROBE_V2_MATERIALIZE)
10130 .ok_or_else(|| {
10131 XlogError::Kernel("hash_join_probe_v2_materialize kernel not found".to_string())
10132 })?;
10133
10134 let mut rec_mat = LaunchRecorder::new_strict(launch_stream);
10135 rec_mat.read(&left_packed.hashes);
10136 rec_mat.read(&left_packed.packed_keys);
10137 rec_mat.read(&index.packed_keys);
10138 rec_mat.read(&table.bucket_offsets);
10139 rec_mat.read(&table.bucket_counts);
10140 rec_mat.read(&table.bucket_entries);
10141 rec_mat.read(&table.bucket_entry_hashes);
10142 rec_mat.read(&per_probe_offsets);
10143 rec_mat.read(left.num_rows_device());
10144 rec_mat.write(&d_output_left);
10145 rec_mat.write(&d_output_right);
10146 rec_mat.write(&d_overflow);
10148 rec_mat.preflight(runtime).map_err(|e| {
10149 XlogError::Kernel(format!(
10150 "indexed CSM inner: materialize preflight failed: {}",
10151 e
10152 ))
10153 })?;
10154
10155 let enqueue = |enqueue: &crate::launch::CudaEnqueue<'_>| -> Result<()> {
10157 unsafe {
10158 let mut params: Vec<*mut c_void> = vec![
10159 (&left_packed.hashes).as_kernel_param(),
10160 left.num_rows_device().as_kernel_param(),
10161 probe_cap.as_kernel_param(),
10162 (&table.bucket_offsets).as_kernel_param(),
10163 (&table.bucket_counts).as_kernel_param(),
10164 (&table.bucket_entries).as_kernel_param(),
10165 (&table.bucket_entry_hashes).as_kernel_param(),
10166 table.bucket_mask.as_kernel_param(),
10167 (&left_packed.packed_keys).as_kernel_param(),
10168 (&index.packed_keys).as_kernel_param(),
10169 index.key_bytes.as_kernel_param(),
10170 (&per_probe_offsets).as_kernel_param(),
10171 output_capacity.as_kernel_param(),
10172 (&d_output_left).as_kernel_param(),
10173 (&d_output_right).as_kernel_param(),
10174 (&d_overflow).as_kernel_param(),
10175 ];
10176 materialize_func
10177 .clone()
10178 .launch_in(enqueue, probe_config, &mut params)
10179 .map_err(|e| {
10180 XlogError::Kernel(format!(
10181 "hash_join_probe_v2_materialize (on_stream, indexed) failed: {}",
10182 e
10183 ))
10184 })
10185 }
10186 };
10187 let rec_mat = unsafe { rec_mat.enqueue_prepared_with(&cu_stream, enqueue) }
10189 .map_err(LaunchEnqueueError::into_xlog_error)?;
10190
10191 rec_mat.commit().map_err(|e| {
10192 XlogError::Kernel(format!(
10193 "indexed CSM inner: materialize commit failed: {}",
10194 e
10195 ))
10196 })?;
10197
10198 cu_stream.synchronize().map_err(|e| {
10199 XlogError::Kernel(format!(
10200 "indexed CSM inner: sync (post-materialize) failed: {}",
10201 e
10202 ))
10203 })?;
10204
10205 let gathered_left = self.gather_buffer_by_indices_on_stream(
10207 left,
10208 &d_output_left,
10209 output_capacity,
10210 &cu_stream,
10211 launch_stream,
10212 runtime,
10213 )?;
10214 let gathered_right = self.gather_buffer_by_indices_on_stream(
10215 right,
10216 &d_output_right,
10217 output_capacity,
10218 &cu_stream,
10219 launch_stream,
10220 runtime,
10221 )?;
10222
10223 let combined_schema = self.combine_schemas(left.schema(), right.schema());
10224 let mut result_columns = Vec::with_capacity(combined_schema.arity());
10225 result_columns.extend(gathered_left.columns);
10226 result_columns.extend(gathered_right.columns);
10227 self.buffer_from_columns(result_columns, output_capacity as u64, combined_schema)
10228 }
10229
10230 #[expect(
10261 clippy::too_many_arguments,
10262 reason = "the indexed recorded outer join exposes both relation/key pairs, cache identity, output bound, and stream"
10263 )]
10264 pub fn hash_join_left_outer_v2_with_index_count_scan_materialize_recorded(
10265 &self,
10266 left: &CudaBuffer,
10267 right: &CudaBuffer,
10268 left_keys: &[usize],
10269 right_keys: &[usize],
10270 index: &crate::provider::JoinIndexV2,
10271 max_output: Option<usize>,
10272 launch_stream: StreamId,
10273 ) -> Result<CudaBuffer> {
10274 use crate::launch::LaunchRecorder;
10275
10276 let runtime = self.memory.runtime().ok_or_else(|| {
10277 XlogError::Kernel(
10278 "hash_join_left_outer_v2_with_index_count_scan_materialize_recorded requires \
10279 a runtime-backed GpuMemoryManager"
10280 .to_string(),
10281 )
10282 })?;
10283 let cu_stream = runtime
10284 .stream_pool()
10285 .resolve(launch_stream)
10286 .ok_or_else(|| {
10287 XlogError::Kernel(format!(
10288 "indexed csm left_outer: launch_stream StreamId({}) does not resolve",
10289 launch_stream.0
10290 ))
10291 })?;
10292
10293 let num_left = self.device_row_count(left)?;
10296 let num_right = self.device_row_count(right)?;
10297 if num_left > u32::MAX as usize || num_right > u32::MAX as usize {
10298 return Err(XlogError::Kernel(format!(
10299 "Join supports at most {} rows per side (left={}, right={})",
10300 u32::MAX,
10301 num_left,
10302 num_right
10303 )));
10304 }
10305 if num_left == 0 {
10306 let combined_schema = self.combine_schemas(left.schema(), right.schema());
10307 return self.create_empty_buffer(combined_schema);
10308 }
10309 if num_right == 0 {
10310 return self.left_outer_with_nulls(left, right);
10315 }
10316 if left_keys.is_empty() || right_keys.is_empty() {
10317 return Err(XlogError::Kernel(
10318 "Join requires at least one key column".to_string(),
10319 ));
10320 }
10321 if left_keys.len() != right_keys.len() {
10322 return Err(XlogError::Kernel(
10323 "Left and right key columns must have same length".to_string(),
10324 ));
10325 }
10326 if left_keys.len() > 4 {
10327 return Err(XlogError::Kernel(
10328 "indexed csm left_outer: max 4 key columns supported (pack_keys constraint)"
10329 .to_string(),
10330 ));
10331 }
10332 for (&l, &r) in left_keys.iter().zip(right_keys.iter()) {
10333 if l >= left.arity() {
10334 return Err(XlogError::Kernel(format!(
10335 "Left key column index {} out of bounds (arity {})",
10336 l,
10337 left.arity()
10338 )));
10339 }
10340 if r >= right.arity() {
10341 return Err(XlogError::Kernel(format!(
10342 "Right key column index {} out of bounds (arity {})",
10343 r,
10344 right.arity()
10345 )));
10346 }
10347 let lt = left.schema().column_type(l);
10348 let rt = right.schema().column_type(r);
10349 if lt != rt {
10350 return Err(XlogError::Kernel(format!(
10351 "Key column type mismatch: left[{}]={:?}, right[{}]={:?}",
10352 l, lt, r, rt
10353 )));
10354 }
10355 }
10356 if index.right_num_rows() != num_right as u32 {
10357 return Err(XlogError::Kernel(
10358 "Join index row count does not match right relation".to_string(),
10359 ));
10360 }
10361 if index.right_keys() != right_keys {
10362 return Err(XlogError::Kernel(
10363 "Join index key columns do not match requested right_keys".to_string(),
10364 ));
10365 }
10366
10367 let probe_cap = u32::try_from(num_left).map_err(|_| {
10369 XlogError::Kernel("indexed csm left_outer: left row count exceeds u32::MAX".to_string())
10370 })?;
10371
10372 let table = &index.table;
10373
10374 let left_packed =
10377 self.pack_keys_gpu_on_stream(left, left_keys, &cu_stream, launch_stream, runtime)?;
10378 if left_packed.key_bytes != index.key_bytes {
10379 return Err(XlogError::Kernel(
10380 "Join key byte width mismatch between probe and cached index".to_string(),
10381 ));
10382 }
10383
10384 let device = self.device.inner();
10385 let block_size = 256u32;
10386 let probe_grid = probe_cap.div_ceil(block_size);
10387 let probe_config = LaunchConfig {
10388 grid_dim: (probe_grid, 1, 1),
10389 block_dim: (block_size, 1, 1),
10390 shared_mem_bytes: 0,
10391 };
10392
10393 let per_probe_count = self.memory.alloc::<u32>(probe_cap as usize)?;
10395 let mut per_probe_offsets = self.memory.alloc::<u32>(probe_cap as usize)?;
10396 let d_logical_count = self.memory.alloc::<u32>(1)?;
10397 let d_overflow = self.memory.alloc::<u8>(1)?;
10398
10399 let count_func = device
10400 .get_func(JOIN_MODULE, join_kernels::HASH_JOIN_PROBE_V2_COUNT_PER_ROW)
10401 .ok_or_else(|| {
10402 XlogError::Kernel("hash_join_probe_v2_count_per_row kernel not found".to_string())
10403 })?;
10404 let total_func = device
10405 .get_func(JOIN_MODULE, join_kernels::HASH_JOIN_TOTAL_FROM_SCAN)
10406 .ok_or_else(|| {
10407 XlogError::Kernel("hash_join_total_from_scan kernel not found".to_string())
10408 })?;
10409
10410 let mut scan_scratch = self.multiblock_scan_u32_scratch_for_len(probe_cap)?;
10411 let mut rec_count = LaunchRecorder::new_strict(launch_stream);
10412 rec_count.read(&left_packed.hashes);
10413 rec_count.read(&left_packed.packed_keys);
10414 rec_count.read(&index.packed_keys);
10415 rec_count.read(&table.bucket_offsets);
10416 rec_count.read(&table.bucket_counts);
10417 rec_count.read(&table.bucket_entries);
10418 rec_count.read(&table.bucket_entry_hashes);
10419 rec_count.read(left.num_rows_device());
10420 rec_count.write(&per_probe_count);
10421 rec_count.write(&per_probe_offsets);
10422 rec_count.write(&d_logical_count);
10423 rec_count.write(&d_overflow);
10424 for level in scan_scratch.levels() {
10425 rec_count.read_write(level);
10426 }
10427 rec_count.preflight(runtime).map_err(|e| {
10428 XlogError::Kernel(format!(
10429 "indexed csm left_outer: count/scan preflight failed: {}",
10430 e
10431 ))
10432 })?;
10433
10434 let enqueue = |enqueue: &crate::launch::CudaEnqueue<'_>| -> Result<()> {
10437 unsafe {
10439 let res = cudarc::driver::sys::cuMemsetD8Async(
10440 *d_overflow.device_ptr(),
10441 0,
10442 1,
10443 cu_stream.cu_stream(),
10444 );
10445 if res != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
10446 return Err(XlogError::Kernel(format!(
10447 "indexed csm left_outer: cuMemsetD8Async (d_overflow) failed: {:?}",
10448 res
10449 )));
10450 }
10451 let res = cudarc::driver::sys::cuMemsetD8Async(
10452 *d_logical_count.device_ptr(),
10453 0,
10454 std::mem::size_of::<u32>(),
10455 cu_stream.cu_stream(),
10456 );
10457 if res != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
10458 return Err(XlogError::Kernel(format!(
10459 "indexed csm left_outer: cuMemsetD8Async (d_logical_count) failed: {:?}",
10460 res
10461 )));
10462 }
10463 }
10464
10465 unsafe {
10466 count_func.clone().launch_in(
10467 enqueue,
10468 probe_config,
10469 (
10470 &left_packed.hashes,
10471 left.num_rows_device(),
10472 probe_cap,
10473 &table.bucket_offsets,
10474 &table.bucket_counts,
10475 &table.bucket_entries,
10476 &table.bucket_entry_hashes,
10477 table.bucket_mask,
10478 &left_packed.packed_keys,
10479 &index.packed_keys,
10480 index.key_bytes,
10481 &per_probe_count,
10482 ),
10483 )
10484 }
10485 .map_err(|e| {
10486 XlogError::Kernel(format!(
10487 "hash_join_probe_v2_count_per_row (indexed csm left_outer) failed: {}",
10488 e
10489 ))
10490 })?;
10491
10492 unsafe {
10495 let res = cudarc::driver::sys::cuMemcpyDtoDAsync_v2(
10496 *per_probe_offsets.device_ptr(),
10497 *per_probe_count.device_ptr(),
10498 (probe_cap as usize) * std::mem::size_of::<u32>(),
10499 cu_stream.cu_stream(),
10500 );
10501 if res != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
10502 return Err(XlogError::Kernel(format!(
10503 "indexed csm left_outer: cuMemcpyDtoDAsync (count → offsets) failed: {:?}",
10504 res
10505 )));
10506 }
10507 }
10508 self.multiblock_scan_u32_inplace_on_stream(
10509 &mut per_probe_offsets,
10510 probe_cap,
10511 enqueue,
10512 &mut scan_scratch,
10513 )?;
10514
10515 let materialize_capacity_bound: u64 =
10517 (probe_cap as u64).saturating_mul(num_right as u64);
10518 let materialize_capacity_u32 = materialize_capacity_bound.min(u32::MAX as u64) as u32;
10519 unsafe {
10521 total_func.clone().launch_in(
10522 enqueue,
10523 LaunchConfig {
10524 grid_dim: (1, 1, 1),
10525 block_dim: (1, 1, 1),
10526 shared_mem_bytes: 0,
10527 },
10528 (
10529 &per_probe_offsets,
10530 &per_probe_count,
10531 left.num_rows_device(),
10532 probe_cap,
10533 materialize_capacity_u32,
10534 &d_logical_count,
10535 &d_overflow,
10536 ),
10537 )
10538 }
10539 .map_err(|e| {
10540 XlogError::Kernel(format!(
10541 "hash_join_total_from_scan (indexed csm left_outer) failed: {}",
10542 e
10543 ))
10544 })?;
10545
10546 Ok(())
10547 };
10548 let rec_count = unsafe { rec_count.enqueue_prepared_with(&cu_stream, enqueue) }
10550 .map_err(LaunchEnqueueError::into_xlog_error)?;
10551
10552 rec_count.commit().map_err(|e| {
10553 XlogError::Kernel(format!(
10554 "indexed csm left_outer: count/scan commit failed: {}",
10555 e
10556 ))
10557 })?;
10558
10559 cu_stream.synchronize().map_err(|e| {
10560 XlogError::Kernel(format!(
10561 "indexed csm left_outer: sync (count read) failed: {}",
10562 e
10563 ))
10564 })?;
10565 let inner_total = self.read_join_output_count_metadata(&d_logical_count)? as u64;
10566 let inner_clamped = max_output
10567 .map(|limit| (limit as u64).min(inner_total))
10568 .unwrap_or(inner_total);
10569 if inner_clamped > u32::MAX as u64 {
10570 return Err(XlogError::Kernel(format!(
10571 "Join produced {} matched rows which exceeds the u32 index limit",
10572 inner_clamped
10573 )));
10574 }
10575 let inner_count_u32 = inner_clamped as u32;
10576
10577 let materialize_func = device
10579 .get_func(JOIN_MODULE, join_kernels::HASH_JOIN_PROBE_V2_MATERIALIZE)
10580 .ok_or_else(|| {
10581 XlogError::Kernel("hash_join_probe_v2_materialize kernel not found".to_string())
10582 })?;
10583 let d_output_left = self.memory.alloc::<u32>(inner_count_u32.max(1) as usize)?;
10584 let d_output_right = self.memory.alloc::<u32>(inner_count_u32.max(1) as usize)?;
10585
10586 let mut rec_mat = LaunchRecorder::new_strict(launch_stream);
10587 rec_mat.read(&left_packed.hashes);
10588 rec_mat.read(&left_packed.packed_keys);
10589 rec_mat.read(&index.packed_keys);
10590 rec_mat.read(&table.bucket_offsets);
10591 rec_mat.read(&table.bucket_counts);
10592 rec_mat.read(&table.bucket_entries);
10593 rec_mat.read(&table.bucket_entry_hashes);
10594 rec_mat.read(&per_probe_offsets);
10595 rec_mat.read(left.num_rows_device());
10596 rec_mat.write(&d_output_left);
10597 rec_mat.write(&d_output_right);
10598 rec_mat.write(&d_overflow);
10600 rec_mat.preflight(runtime).map_err(|e| {
10601 XlogError::Kernel(format!(
10602 "indexed csm left_outer: materialize preflight failed: {}",
10603 e
10604 ))
10605 })?;
10606 if inner_count_u32 > 0 {
10607 let enqueue = |enqueue: &crate::launch::CudaEnqueue<'_>| -> Result<()> {
10609 unsafe {
10610 let mut params: Vec<*mut c_void> = vec![
10611 (&left_packed.hashes).as_kernel_param(),
10612 left.num_rows_device().as_kernel_param(),
10613 probe_cap.as_kernel_param(),
10614 (&table.bucket_offsets).as_kernel_param(),
10615 (&table.bucket_counts).as_kernel_param(),
10616 (&table.bucket_entries).as_kernel_param(),
10617 (&table.bucket_entry_hashes).as_kernel_param(),
10618 table.bucket_mask.as_kernel_param(),
10619 (&left_packed.packed_keys).as_kernel_param(),
10620 (&index.packed_keys).as_kernel_param(),
10621 index.key_bytes.as_kernel_param(),
10622 (&per_probe_offsets).as_kernel_param(),
10623 inner_count_u32.as_kernel_param(),
10624 (&d_output_left).as_kernel_param(),
10625 (&d_output_right).as_kernel_param(),
10626 (&d_overflow).as_kernel_param(),
10627 ];
10628 materialize_func
10629 .clone()
10630 .launch_in(enqueue, probe_config, &mut params)
10631 .map_err(|e| {
10632 XlogError::Kernel(format!(
10633 "hash_join_probe_v2_materialize (indexed csm left_outer) failed: {}",
10634 e
10635 ))
10636 })
10637 }
10638 };
10639 let rec_mat = unsafe { rec_mat.enqueue_prepared_with(&cu_stream, enqueue) }
10641 .map_err(LaunchEnqueueError::into_xlog_error)?;
10642 rec_mat.commit().map_err(|e| {
10643 XlogError::Kernel(format!(
10644 "indexed csm left_outer: materialize commit failed: {}",
10645 e
10646 ))
10647 })?;
10648 } else {
10649 rec_mat.abort().map_err(|e| {
10650 XlogError::Kernel(format!(
10651 "indexed csm left_outer: materialize abort failed: {}",
10652 e
10653 ))
10654 })?;
10655 }
10656
10657 let d_unmatched_mask = self.memory.alloc::<u8>(probe_cap as usize)?;
10659 let unmatched_mask_func = device
10660 .get_func(JOIN_MODULE, join_kernels::HASH_JOIN_CSM_UNMATCHED_MASK)
10661 .ok_or_else(|| {
10662 XlogError::Kernel("hash_join_csm_unmatched_mask kernel not found".to_string())
10663 })?;
10664 let mut rec_um = LaunchRecorder::new_strict(launch_stream);
10665 rec_um.read(&per_probe_count);
10666 rec_um.read(left.num_rows_device());
10667 rec_um.write(&d_unmatched_mask);
10668 rec_um.preflight(runtime).map_err(|e| {
10669 XlogError::Kernel(format!(
10670 "indexed csm left_outer: unmatched mask preflight failed: {}",
10671 e
10672 ))
10673 })?;
10674 let enqueue = |enqueue: &crate::launch::CudaEnqueue<'_>| -> Result<()> {
10676 unsafe {
10677 unmatched_mask_func.clone().launch_in(
10678 enqueue,
10679 probe_config,
10680 (
10681 &per_probe_count,
10682 left.num_rows_device(),
10683 probe_cap,
10684 &d_unmatched_mask,
10685 ),
10686 )
10687 }
10688 .map_err(|e| {
10689 XlogError::Kernel(format!(
10690 "hash_join_csm_unmatched_mask (indexed csm left_outer) failed: {}",
10691 e
10692 ))
10693 })
10694 };
10695 let rec_um = unsafe { rec_um.enqueue_prepared_with(&cu_stream, enqueue) }
10697 .map_err(LaunchEnqueueError::into_xlog_error)?;
10698 rec_um.commit().map_err(|e| {
10699 XlogError::Kernel(format!(
10700 "indexed csm left_outer: unmatched mask commit failed: {}",
10701 e
10702 ))
10703 })?;
10704
10705 let unmatched_left = self.compact_buffer_by_device_mask_counted_recorded(
10706 left,
10707 &d_unmatched_mask,
10708 launch_stream,
10709 )?;
10710 let unmatched_rows = self.device_row_count(&unmatched_left)? as u64;
10711 let total_rows = (inner_count_u32 as u64) + unmatched_rows;
10712
10713 let combined_schema = self.combine_schemas(left.schema(), right.schema());
10714 if total_rows == 0 {
10715 return self.create_empty_buffer(combined_schema);
10716 }
10717
10718 let inner_left_buf;
10720 let inner_right_buf;
10721 if inner_count_u32 > 0 {
10722 inner_left_buf = Some(self.gather_buffer_by_indices_on_stream(
10723 left,
10724 &d_output_left,
10725 inner_count_u32,
10726 &cu_stream,
10727 launch_stream,
10728 runtime,
10729 )?);
10730 inner_right_buf = Some(self.gather_buffer_by_indices_on_stream(
10731 right,
10732 &d_output_right,
10733 inner_count_u32,
10734 &cu_stream,
10735 launch_stream,
10736 runtime,
10737 )?);
10738 } else {
10739 inner_left_buf = None;
10740 inner_right_buf = None;
10741 }
10742
10743 let result_columns = self.concatenate_left_outer_columns_on_stream(
10745 left,
10746 right,
10747 inner_left_buf.as_ref(),
10748 inner_right_buf.as_ref(),
10749 &unmatched_left,
10750 inner_count_u32 as u64,
10751 unmatched_rows,
10752 &cu_stream,
10753 launch_stream,
10754 runtime,
10755 "indexed csm left_outer",
10756 )?;
10757
10758 if total_rows > u32::MAX as u64 {
10760 return Err(XlogError::Kernel(format!(
10761 "indexed csm left_outer: output row count {} exceeds u32::MAX",
10762 total_rows
10763 )));
10764 }
10765 let total_rows_u32 = total_rows as u32;
10766 let d_num_rows = self.upload_device_row_count(total_rows_u32)?;
10767 Ok(CudaBuffer::from_columns_with_host_count(
10768 result_columns,
10769 total_rows,
10770 d_num_rows,
10771 combined_schema,
10772 total_rows_u32,
10773 ))
10774 }
10775
10776 #[expect(
10791 clippy::too_many_arguments,
10792 reason = "the recorded join dispatcher keeps both relation/key pairs, semantics, output bound, and stream explicit"
10793 )]
10794 pub fn hash_join_v2_recorded(
10795 &self,
10796 left: &CudaBuffer,
10797 right: &CudaBuffer,
10798 left_keys: &[usize],
10799 right_keys: &[usize],
10800 join_type: JoinType,
10801 max_output: Option<usize>,
10802 launch_stream: StreamId,
10803 ) -> Result<CudaBuffer> {
10804 let csm_on = Self::use_recorded_csm_env()?;
10805 match join_type {
10806 JoinType::Inner => {
10807 if csm_on {
10808 self.csm_invocations.fetch_add(1, Ordering::Relaxed);
10809 self.hash_join_inner_v2_count_scan_materialize_recorded(
10810 left,
10811 right,
10812 left_keys,
10813 right_keys,
10814 max_output,
10815 launch_stream,
10816 )
10817 } else {
10818 self.hash_join_inner_v2_recorded(
10819 left,
10820 right,
10821 left_keys,
10822 right_keys,
10823 max_output,
10824 launch_stream,
10825 )
10826 }
10827 }
10828 JoinType::Semi => self.hash_join_semi_or_anti_v2_recorded(
10829 left,
10830 right,
10831 left_keys,
10832 right_keys,
10833 false,
10834 launch_stream,
10835 ),
10836 JoinType::Anti => self.hash_join_semi_or_anti_v2_recorded(
10837 left,
10838 right,
10839 left_keys,
10840 right_keys,
10841 true,
10842 launch_stream,
10843 ),
10844 JoinType::LeftOuter => {
10845 if csm_on {
10846 self.csm_invocations.fetch_add(1, Ordering::Relaxed);
10847 self.hash_join_left_outer_v2_count_scan_materialize_recorded(
10848 left,
10849 right,
10850 left_keys,
10851 right_keys,
10852 max_output,
10853 launch_stream,
10854 )
10855 } else {
10856 self.hash_join_left_outer_v2_recorded(
10857 left,
10858 right,
10859 left_keys,
10860 right_keys,
10861 max_output,
10862 launch_stream,
10863 )
10864 }
10865 }
10866 }
10867 }
10868
10869 fn hash_join_left_outer_v2_recorded(
10890 &self,
10891 left: &CudaBuffer,
10892 right: &CudaBuffer,
10893 left_keys: &[usize],
10894 right_keys: &[usize],
10895 max_output: Option<usize>,
10896 launch_stream: StreamId,
10897 ) -> Result<CudaBuffer> {
10898 use crate::launch::LaunchRecorder;
10899
10900 let runtime = self.memory.runtime().ok_or_else(|| {
10901 XlogError::Kernel(
10902 "hash_join_v2_recorded (left_outer) requires a runtime-backed GpuMemoryManager"
10903 .to_string(),
10904 )
10905 })?;
10906 let cu_stream = runtime
10907 .stream_pool()
10908 .resolve(launch_stream)
10909 .ok_or_else(|| {
10910 XlogError::Kernel(format!(
10911 "hash_join_v2_recorded (left_outer): launch_stream StreamId({}) does not resolve",
10912 launch_stream.0
10913 ))
10914 })?;
10915
10916 let num_left = self.device_row_count(left)?;
10917 let num_right = self.device_row_count(right)?;
10918 if num_left > u32::MAX as usize || num_right > u32::MAX as usize {
10919 return Err(XlogError::Kernel(format!(
10920 "Join supports at most {} rows per side (left={}, right={})",
10921 u32::MAX,
10922 num_left,
10923 num_right
10924 )));
10925 }
10926 if num_left == 0 {
10927 let combined_schema = self.combine_schemas(left.schema(), right.schema());
10928 return self.create_empty_buffer(combined_schema);
10929 }
10930 if num_right == 0 {
10931 return self.left_outer_with_nulls(left, right);
10936 }
10937 if left_keys.is_empty() || right_keys.is_empty() {
10938 return Err(XlogError::Kernel(
10939 "Join requires at least one key column".to_string(),
10940 ));
10941 }
10942 if left_keys.len() != right_keys.len() {
10943 return Err(XlogError::Kernel(
10944 "Left and right key columns must have same length".to_string(),
10945 ));
10946 }
10947 if left_keys.len() > 4 {
10948 return Err(XlogError::Kernel(
10949 "hash_join_v2_recorded (left_outer): max 4 key columns supported \
10950 (pack_keys constraint)"
10951 .to_string(),
10952 ));
10953 }
10954 for (&l, &r) in left_keys.iter().zip(right_keys.iter()) {
10955 let lt = left.schema().column_type(l);
10956 let rt = right.schema().column_type(r);
10957 if lt != rt {
10958 return Err(XlogError::Kernel(format!(
10959 "Key column type mismatch: left[{}]={:?}, right[{}]={:?}",
10960 l, lt, r, rt
10961 )));
10962 }
10963 }
10964
10965 let num_left = num_left as u32;
10966 let num_right = num_right as u32;
10967
10968 let left_packed =
10969 self.pack_keys_gpu_on_stream(left, left_keys, &cu_stream, launch_stream, runtime)?;
10970 let right_packed =
10971 self.pack_keys_gpu_on_stream(right, right_keys, &cu_stream, launch_stream, runtime)?;
10972 let table = self.build_hash_table_v2_on_stream(
10973 &right_packed.hashes,
10974 num_right,
10975 &cu_stream,
10976 launch_stream,
10977 runtime,
10978 )?;
10979
10980 let device = self.device.inner();
10981 let block_size = 256u32;
10982 let grid_size = num_left.div_ceil(block_size);
10983 let cfg = LaunchConfig {
10984 grid_dim: (grid_size, 1, 1),
10985 block_dim: (block_size, 1, 1),
10986 shared_mem_bytes: 0,
10987 };
10988
10989 let d_has_match = self.memory.alloc::<u8>(num_left as usize)?;
10993 let d_count_only = self.memory.alloc::<u32>(1)?;
10994 let d_dummy_left = self.memory.alloc::<u32>(1)?;
10995 let d_dummy_right = self.memory.alloc::<u32>(1)?;
10996 let semi_func = device
10997 .get_func(JOIN_MODULE, join_kernels::HASH_JOIN_SEMI)
10998 .ok_or_else(|| XlogError::Kernel("hash_join_semi kernel not found".to_string()))?;
10999 let probe_func = device
11000 .get_func(JOIN_MODULE, join_kernels::HASH_JOIN_PROBE_V2)
11001 .ok_or_else(|| XlogError::Kernel("hash_join_probe_v2 kernel not found".to_string()))?;
11002
11003 let mut rec_a = LaunchRecorder::new_strict(launch_stream);
11004 rec_a.read(&left_packed.hashes);
11005 rec_a.read(&left_packed.packed_keys);
11006 rec_a.read(&right_packed.packed_keys);
11007 rec_a.read(&table.bucket_offsets);
11008 rec_a.read(&table.bucket_counts);
11009 rec_a.read(&table.bucket_entries);
11010 rec_a.read(&table.bucket_entry_hashes);
11011 rec_a.write(&d_has_match);
11012 rec_a.write(&d_count_only);
11013 rec_a.write(&d_dummy_left);
11014 rec_a.write(&d_dummy_right);
11015 rec_a.preflight(runtime).map_err(|e| {
11016 XlogError::Kernel(format!(
11017 "hash_join_v2_recorded (left_outer): semi/count preflight failed: {}",
11018 e
11019 ))
11020 })?;
11021
11022 let enqueue = |enqueue: &crate::launch::CudaEnqueue<'_>| -> Result<()> {
11024 unsafe {
11026 let res = cudarc::driver::sys::cuMemsetD8Async(
11027 *d_count_only.device_ptr(),
11028 0,
11029 std::mem::size_of::<u32>(),
11030 cu_stream.cu_stream(),
11031 );
11032 if res != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
11033 return Err(XlogError::Kernel(format!(
11034 "cuMemsetD8Async (left_outer d_count_only) failed: {:?}",
11035 res
11036 )));
11037 }
11038 }
11039
11040 unsafe {
11041 semi_func.clone().launch_in(
11042 enqueue,
11043 cfg,
11044 (
11045 &left_packed.hashes,
11046 num_left,
11047 &table.bucket_offsets,
11048 &table.bucket_counts,
11049 &table.bucket_entries,
11050 &table.bucket_entry_hashes,
11051 table.bucket_mask,
11052 &left_packed.packed_keys,
11053 &right_packed.packed_keys,
11054 left_packed.key_bytes,
11055 &d_has_match,
11056 ),
11057 )
11058 }
11059 .map_err(|e| XlogError::Kernel(format!("hash_join_semi (on_stream) failed: {}", e)))?;
11060
11061 let max_output_count_only = 0u32;
11062 unsafe {
11064 let mut params: Vec<*mut c_void> = vec![
11065 (&left_packed.hashes).as_kernel_param(),
11066 num_left.as_kernel_param(),
11067 (&table.bucket_offsets).as_kernel_param(),
11068 (&table.bucket_counts).as_kernel_param(),
11069 (&table.bucket_entries).as_kernel_param(),
11070 (&table.bucket_entry_hashes).as_kernel_param(),
11071 table.bucket_mask.as_kernel_param(),
11072 (&left_packed.packed_keys).as_kernel_param(),
11073 (&right_packed.packed_keys).as_kernel_param(),
11074 left_packed.key_bytes.as_kernel_param(),
11075 (&d_dummy_left).as_kernel_param(),
11076 (&d_dummy_right).as_kernel_param(),
11077 (&d_count_only).as_kernel_param(),
11078 max_output_count_only.as_kernel_param(),
11079 ];
11080 probe_func
11081 .clone()
11082 .launch_in(enqueue, cfg, &mut params)
11083 .map_err(|e| {
11084 XlogError::Kernel(format!(
11085 "hash_join_probe_v2 (count, on_stream, left_outer) failed: {}",
11086 e
11087 ))
11088 })?;
11089 }
11090
11091 Ok(())
11092 };
11093 let rec_a = unsafe { rec_a.enqueue_prepared_with(&cu_stream, enqueue) }
11095 .map_err(LaunchEnqueueError::into_xlog_error)?;
11096
11097 rec_a.commit().map_err(|e| {
11098 XlogError::Kernel(format!(
11099 "hash_join_v2_recorded (left_outer): semi/count commit failed: {}",
11100 e
11101 ))
11102 })?;
11103
11104 cu_stream.synchronize().map_err(|e| {
11106 XlogError::Kernel(format!(
11107 "hash_join_v2_recorded (left_outer): sync (count read) failed: {}",
11108 e
11109 ))
11110 })?;
11111 let full_inner = self.read_join_output_count_metadata(&d_count_only)? as u64;
11112 let requested_inner = max_output
11113 .map(|limit| (limit as u64).min(full_inner))
11114 .unwrap_or(full_inner);
11115 if requested_inner > u32::MAX as u64 {
11116 return Err(XlogError::Kernel(format!(
11117 "Join produced {} rows which exceeds the u32 index limit",
11118 requested_inner
11119 )));
11120 }
11121 let max_output_u32 = requested_inner as u32;
11122 let alloc_len = (requested_inner.max(1)) as usize;
11123
11124 let d_output_left = self.memory.alloc::<u32>(alloc_len)?;
11125 let d_output_right = self.memory.alloc::<u32>(alloc_len)?;
11126 let d_output_count = self.memory.alloc::<u32>(1)?;
11127
11128 let mut rec_b = LaunchRecorder::new_strict(launch_stream);
11129 rec_b.read(&left_packed.hashes);
11130 rec_b.read(&left_packed.packed_keys);
11131 rec_b.read(&right_packed.packed_keys);
11132 rec_b.read(&table.bucket_offsets);
11133 rec_b.read(&table.bucket_counts);
11134 rec_b.read(&table.bucket_entries);
11135 rec_b.read(&table.bucket_entry_hashes);
11136 rec_b.write(&d_output_left);
11137 rec_b.write(&d_output_right);
11138 rec_b.write(&d_output_count);
11139 rec_b.preflight(runtime).map_err(|e| {
11140 XlogError::Kernel(format!(
11141 "hash_join_v2_recorded (left_outer): materialize preflight failed: {}",
11142 e
11143 ))
11144 })?;
11145
11146 let enqueue = |enqueue: &crate::launch::CudaEnqueue<'_>| -> Result<()> {
11148 unsafe {
11150 let res = cudarc::driver::sys::cuMemsetD8Async(
11151 *d_output_count.device_ptr(),
11152 0,
11153 std::mem::size_of::<u32>(),
11154 cu_stream.cu_stream(),
11155 );
11156 if res != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
11157 return Err(XlogError::Kernel(format!(
11158 "cuMemsetD8Async (left_outer d_output_count) failed: {:?}",
11159 res
11160 )));
11161 }
11162 }
11163 unsafe {
11164 let mut params: Vec<*mut c_void> = vec![
11165 (&left_packed.hashes).as_kernel_param(),
11166 num_left.as_kernel_param(),
11167 (&table.bucket_offsets).as_kernel_param(),
11168 (&table.bucket_counts).as_kernel_param(),
11169 (&table.bucket_entries).as_kernel_param(),
11170 (&table.bucket_entry_hashes).as_kernel_param(),
11171 table.bucket_mask.as_kernel_param(),
11172 (&left_packed.packed_keys).as_kernel_param(),
11173 (&right_packed.packed_keys).as_kernel_param(),
11174 left_packed.key_bytes.as_kernel_param(),
11175 (&d_output_left).as_kernel_param(),
11176 (&d_output_right).as_kernel_param(),
11177 (&d_output_count).as_kernel_param(),
11178 max_output_u32.as_kernel_param(),
11179 ];
11180 probe_func
11181 .clone()
11182 .launch_in(enqueue, cfg, &mut params)
11183 .map_err(|e| {
11184 XlogError::Kernel(format!(
11185 "hash_join_probe_v2 (materialize, on_stream, left_outer) failed: {}",
11186 e
11187 ))
11188 })?;
11189 }
11190 Ok(())
11191 };
11192 let rec_b = unsafe { rec_b.enqueue_prepared_with(&cu_stream, enqueue) }
11194 .map_err(LaunchEnqueueError::into_xlog_error)?;
11195
11196 rec_b.commit().map_err(|e| {
11197 XlogError::Kernel(format!(
11198 "hash_join_v2_recorded (left_outer): materialize commit failed: {}",
11199 e
11200 ))
11201 })?;
11202
11203 cu_stream.synchronize().map_err(|e| {
11204 XlogError::Kernel(format!(
11205 "hash_join_v2_recorded (left_outer): sync (materialize read) failed: {}",
11206 e
11207 ))
11208 })?;
11209 let inner_count = self
11210 .read_join_output_count_metadata(&d_output_count)?
11211 .min(max_output_u32);
11212
11213 let d_no_match = self.memory.alloc::<u8>(num_left as usize)?;
11216 let mask_not_fn = device
11217 .get_func(FILTER_MODULE, filter_kernels::MASK_NOT)
11218 .ok_or_else(|| XlogError::Kernel("mask_not kernel not found".to_string()))?;
11219
11220 let mut rec_c = LaunchRecorder::new_strict(launch_stream);
11221 rec_c.read(&d_has_match);
11222 rec_c.write(&d_no_match);
11223 rec_c.preflight(runtime).map_err(|e| {
11224 XlogError::Kernel(format!(
11225 "hash_join_v2_recorded (left_outer): mask_not preflight failed: {}",
11226 e
11227 ))
11228 })?;
11229 let enqueue = |enqueue: &crate::launch::CudaEnqueue<'_>| -> Result<()> {
11231 unsafe {
11232 mask_not_fn
11233 .clone()
11234 .launch_in(enqueue, cfg, (&d_has_match, &d_no_match, num_left))
11235 }
11236 .map_err(|e| XlogError::Kernel(format!("mask_not (on_stream) failed: {}", e)))
11237 };
11238 let rec_c = unsafe { rec_c.enqueue_prepared_with(&cu_stream, enqueue) }
11240 .map_err(LaunchEnqueueError::into_xlog_error)?;
11241 rec_c.commit().map_err(|e| {
11242 XlogError::Kernel(format!(
11243 "hash_join_v2_recorded (left_outer): mask_not commit failed: {}",
11244 e
11245 ))
11246 })?;
11247
11248 let unmatched_left =
11249 self.compact_buffer_by_device_mask_counted_recorded(left, &d_no_match, launch_stream)?;
11250 let unmatched_rows = self.device_row_count(&unmatched_left)? as u64;
11251 let total_rows = (inner_count as u64) + unmatched_rows;
11252
11253 let combined_schema = self.combine_schemas(left.schema(), right.schema());
11254 if total_rows == 0 {
11255 return self.create_empty_buffer(combined_schema);
11256 }
11257
11258 let inner_count_u32 = inner_count;
11267 let inner_left_buf;
11268 let inner_right_buf;
11269 if inner_count > 0 {
11270 inner_left_buf = Some(self.gather_buffer_by_indices_on_stream(
11271 left,
11272 &d_output_left,
11273 inner_count_u32,
11274 &cu_stream,
11275 launch_stream,
11276 runtime,
11277 )?);
11278 inner_right_buf = Some(self.gather_buffer_by_indices_on_stream(
11279 right,
11280 &d_output_right,
11281 inner_count_u32,
11282 &cu_stream,
11283 launch_stream,
11284 runtime,
11285 )?);
11286 } else {
11287 inner_left_buf = None;
11288 inner_right_buf = None;
11289 }
11290
11291 let result_columns = self.concatenate_left_outer_columns_on_stream(
11308 left,
11309 right,
11310 inner_left_buf.as_ref(),
11311 inner_right_buf.as_ref(),
11312 &unmatched_left,
11313 inner_count as u64,
11314 unmatched_rows,
11315 &cu_stream,
11316 launch_stream,
11317 runtime,
11318 "hash_join_v2_recorded left_outer",
11319 )?;
11320
11321 let d_num_rows = self.upload_device_row_count(total_rows as u32)?;
11324 Ok(CudaBuffer::from_columns_with_host_count(
11325 result_columns,
11326 total_rows,
11327 d_num_rows,
11328 combined_schema,
11329 total_rows as u32,
11330 ))
11331 }
11332
11333 fn hash_join_semi_or_anti_v2_recorded(
11344 &self,
11345 left: &CudaBuffer,
11346 right: &CudaBuffer,
11347 left_keys: &[usize],
11348 right_keys: &[usize],
11349 anti: bool,
11350 launch_stream: StreamId,
11351 ) -> Result<CudaBuffer> {
11352 use crate::launch::LaunchRecorder;
11353
11354 let runtime = self.memory.runtime().ok_or_else(|| {
11355 XlogError::Kernel(
11356 "hash_join_v2_recorded (semi/anti) requires a runtime-backed GpuMemoryManager"
11357 .to_string(),
11358 )
11359 })?;
11360 let cu_stream = runtime
11361 .stream_pool()
11362 .resolve(launch_stream)
11363 .ok_or_else(|| {
11364 XlogError::Kernel(format!(
11365 "hash_join_v2_recorded (semi/anti): launch_stream StreamId({}) does not resolve",
11366 launch_stream.0
11367 ))
11368 })?;
11369
11370 let num_left = self.device_row_count(left)?;
11371 let num_right = self.device_row_count(right)?;
11372 if num_left > u32::MAX as usize || num_right > u32::MAX as usize {
11373 return Err(XlogError::Kernel(format!(
11374 "Join supports at most {} rows per side (left={}, right={})",
11375 u32::MAX,
11376 num_left,
11377 num_right
11378 )));
11379 }
11380 if num_left == 0 {
11381 return self.create_empty_buffer(left.schema().clone());
11382 }
11383 if num_right == 0 {
11384 return if anti {
11397 self.clone_buffer(left)
11398 } else {
11399 self.create_empty_buffer(left.schema().clone())
11400 };
11401 }
11402 if left_keys.is_empty() || right_keys.is_empty() {
11403 return Err(XlogError::Kernel(
11404 "Join requires at least one key column".to_string(),
11405 ));
11406 }
11407 if left_keys.len() != right_keys.len() {
11408 return Err(XlogError::Kernel(
11409 "Left and right key columns must have same length".to_string(),
11410 ));
11411 }
11412 if left_keys.len() > 4 {
11413 return Err(XlogError::Kernel(
11414 "hash_join_v2_recorded (semi/anti): max 4 key columns supported (pack_keys constraint)"
11415 .to_string(),
11416 ));
11417 }
11418 for (&l, &r) in left_keys.iter().zip(right_keys.iter()) {
11419 let lt = left.schema().column_type(l);
11420 let rt = right.schema().column_type(r);
11421 if lt != rt {
11422 return Err(XlogError::Kernel(format!(
11423 "Key column type mismatch: left[{}]={:?}, right[{}]={:?}",
11424 l, lt, r, rt
11425 )));
11426 }
11427 }
11428
11429 let num_left = num_left as u32;
11430 let num_right = num_right as u32;
11431
11432 let left_packed =
11433 self.pack_keys_gpu_on_stream(left, left_keys, &cu_stream, launch_stream, runtime)?;
11434 let right_packed =
11435 self.pack_keys_gpu_on_stream(right, right_keys, &cu_stream, launch_stream, runtime)?;
11436 let table = self.build_hash_table_v2_on_stream(
11437 &right_packed.hashes,
11438 num_right,
11439 &cu_stream,
11440 launch_stream,
11441 runtime,
11442 )?;
11443
11444 let d_mask = self.memory.alloc::<u8>(num_left as usize)?;
11445
11446 let kernel_name = if anti {
11447 join_kernels::HASH_JOIN_ANTI
11448 } else {
11449 join_kernels::HASH_JOIN_SEMI
11450 };
11451 let func = self
11452 .device
11453 .inner()
11454 .get_func(JOIN_MODULE, kernel_name)
11455 .ok_or_else(|| XlogError::Kernel(format!("{} kernel not found", kernel_name)))?;
11456
11457 let block_size = 256u32;
11458 let grid_size = num_left.div_ceil(block_size);
11459 let cfg = LaunchConfig {
11460 grid_dim: (grid_size, 1, 1),
11461 block_dim: (block_size, 1, 1),
11462 shared_mem_bytes: 0,
11463 };
11464
11465 let mut rec = LaunchRecorder::new_strict(launch_stream);
11466 rec.read(&left_packed.hashes);
11467 rec.read(&left_packed.packed_keys);
11468 rec.read(&right_packed.packed_keys);
11469 rec.read(&table.bucket_offsets);
11470 rec.read(&table.bucket_counts);
11471 rec.read(&table.bucket_entries);
11472 rec.read(&table.bucket_entry_hashes);
11473 rec.write(&d_mask);
11474 rec.preflight(runtime).map_err(|e| {
11475 XlogError::Kernel(format!(
11476 "hash_join_v2_recorded (semi/anti): preflight failed: {}",
11477 e
11478 ))
11479 })?;
11480
11481 let enqueue = |enqueue: &crate::launch::CudaEnqueue<'_>| -> Result<()> {
11486 unsafe {
11487 func.clone().launch_in(
11488 enqueue,
11489 cfg,
11490 (
11491 &left_packed.hashes,
11492 num_left,
11493 &table.bucket_offsets,
11494 &table.bucket_counts,
11495 &table.bucket_entries,
11496 &table.bucket_entry_hashes,
11497 table.bucket_mask,
11498 &left_packed.packed_keys,
11499 &right_packed.packed_keys,
11500 left_packed.key_bytes,
11501 &d_mask,
11502 ),
11503 )
11504 }
11505 .map_err(|e| XlogError::Kernel(format!("{} (on_stream) failed: {}", kernel_name, e)))
11506 };
11507 let rec = unsafe { rec.enqueue_prepared_with(&cu_stream, enqueue) }
11509 .map_err(LaunchEnqueueError::into_xlog_error)?;
11510
11511 rec.commit().map_err(|e| {
11512 XlogError::Kernel(format!(
11513 "hash_join_v2_recorded (semi/anti): commit failed: {}",
11514 e
11515 ))
11516 })?;
11517
11518 self.compact_buffer_by_device_mask_counted_recorded(left, &d_mask, launch_stream)
11525 }
11526
11527 #[expect(
11554 clippy::too_many_arguments,
11555 reason = "the indexed recorded dispatcher keeps both relation/key pairs, semantics, cache identity, output bound, and stream explicit"
11556 )]
11557 pub fn hash_join_v2_with_index_recorded(
11558 &self,
11559 left: &CudaBuffer,
11560 right: &CudaBuffer,
11561 left_keys: &[usize],
11562 right_keys: &[usize],
11563 join_type: JoinType,
11564 index: &crate::provider::JoinIndexV2,
11565 max_output: Option<usize>,
11566 launch_stream: StreamId,
11567 ) -> Result<CudaBuffer> {
11568 let runtime = self.memory.runtime().ok_or_else(|| {
11569 XlogError::Kernel(
11570 "hash_join_v2_with_index_recorded requires a runtime-backed GpuMemoryManager"
11571 .to_string(),
11572 )
11573 })?;
11574 runtime
11576 .stream_pool()
11577 .resolve(launch_stream)
11578 .ok_or_else(|| {
11579 XlogError::Kernel(format!(
11580 "hash_join_v2_with_index_recorded: launch_stream StreamId({}) does not resolve",
11581 launch_stream.0
11582 ))
11583 })?;
11584
11585 let left_rows = self.device_row_count(left)?;
11587 let right_rows = self.device_row_count(right)?;
11588 if left_rows > u32::MAX as usize || right_rows > u32::MAX as usize {
11589 return Err(XlogError::Kernel(format!(
11590 "Join supports at most {} rows per side (left={}, right={})",
11591 u32::MAX,
11592 left_rows,
11593 right_rows
11594 )));
11595 }
11596 if left_rows == 0 {
11597 return match join_type {
11598 JoinType::Inner | JoinType::LeftOuter => {
11599 let combined_schema = self.combine_schemas(left.schema(), right.schema());
11600 self.create_empty_buffer(combined_schema)
11601 }
11602 JoinType::Semi | JoinType::Anti => self.create_empty_buffer(left.schema().clone()),
11603 };
11604 }
11605 if right_rows == 0 {
11606 return match join_type {
11607 JoinType::Inner => {
11608 let combined_schema = self.combine_schemas(left.schema(), right.schema());
11609 self.create_empty_buffer(combined_schema)
11610 }
11611 JoinType::Semi => self.create_empty_buffer(left.schema().clone()),
11612 JoinType::Anti => self.clone_buffer(left),
11613 JoinType::LeftOuter => self.left_outer_with_nulls(left, right),
11614 };
11615 }
11616 if left_keys.is_empty() || right_keys.is_empty() {
11617 return Err(XlogError::Kernel(
11618 "Join requires at least one key column".to_string(),
11619 ));
11620 }
11621 if left_keys.len() != right_keys.len() {
11622 return Err(XlogError::Kernel(
11623 "Left and right key columns must have same length".to_string(),
11624 ));
11625 }
11626 if left_keys.len() > 4 {
11627 return Err(XlogError::Kernel(
11628 "hash_join_v2_with_index_recorded: max 4 key columns supported \
11629 (pack_keys constraint)"
11630 .to_string(),
11631 ));
11632 }
11633 for (&l, &r) in left_keys.iter().zip(right_keys.iter()) {
11634 if l >= left.arity() {
11635 return Err(XlogError::Kernel(format!(
11636 "Left key column index {} out of bounds (arity {})",
11637 l,
11638 left.arity()
11639 )));
11640 }
11641 if r >= right.arity() {
11642 return Err(XlogError::Kernel(format!(
11643 "Right key column index {} out of bounds (arity {})",
11644 r,
11645 right.arity()
11646 )));
11647 }
11648 let lt = left.schema().column_type(l);
11649 let rt = right.schema().column_type(r);
11650 if lt != rt {
11651 return Err(XlogError::Kernel(format!(
11652 "Key column type mismatch: left[{}]={:?}, right[{}]={:?}",
11653 l, lt, r, rt
11654 )));
11655 }
11656 }
11657 if index.right_num_rows() != right_rows as u32 {
11658 return Err(XlogError::Kernel(
11659 "Join index row count does not match right relation".to_string(),
11660 ));
11661 }
11662 if index.right_keys() != right_keys {
11663 return Err(XlogError::Kernel(
11664 "Join index key columns do not match requested right_keys".to_string(),
11665 ));
11666 }
11667
11668 let csm_on = Self::use_recorded_csm_env()?;
11669 match join_type {
11670 JoinType::Inner => {
11671 if csm_on {
11672 self.csm_invocations.fetch_add(1, Ordering::Relaxed);
11673 self.hash_join_inner_v2_with_index_count_scan_materialize_recorded(
11674 left,
11675 right,
11676 left_keys,
11677 right_keys,
11678 index,
11679 max_output,
11680 launch_stream,
11681 )
11682 } else {
11683 self.hash_join_inner_v2_with_index_recorded(
11684 left,
11685 right,
11686 left_keys,
11687 index,
11688 max_output,
11689 launch_stream,
11690 )
11691 }
11692 }
11693 JoinType::Semi => self.hash_join_semi_or_anti_v2_with_index_recorded(
11694 left,
11695 left_keys,
11696 index,
11697 false,
11698 launch_stream,
11699 ),
11700 JoinType::Anti => self.hash_join_semi_or_anti_v2_with_index_recorded(
11701 left,
11702 left_keys,
11703 index,
11704 true,
11705 launch_stream,
11706 ),
11707 JoinType::LeftOuter => {
11708 if csm_on {
11709 self.csm_invocations.fetch_add(1, Ordering::Relaxed);
11710 self.hash_join_left_outer_v2_with_index_count_scan_materialize_recorded(
11711 left,
11712 right,
11713 left_keys,
11714 right_keys,
11715 index,
11716 max_output,
11717 launch_stream,
11718 )
11719 } else {
11720 self.hash_join_left_outer_v2_with_index_recorded(
11721 left,
11722 right,
11723 left_keys,
11724 index,
11725 max_output,
11726 launch_stream,
11727 )
11728 }
11729 }
11730 }
11731 }
11732
11733 fn hash_join_inner_v2_with_index_recorded(
11739 &self,
11740 left: &CudaBuffer,
11741 right: &CudaBuffer,
11742 left_keys: &[usize],
11743 index: &crate::provider::JoinIndexV2,
11744 max_output: Option<usize>,
11745 launch_stream: StreamId,
11746 ) -> Result<CudaBuffer> {
11747 use crate::launch::LaunchRecorder;
11748
11749 let runtime = self.memory.runtime().ok_or_else(|| {
11750 XlogError::Kernel(
11751 "hash_join_v2_with_index_recorded (inner) requires runtime-backed manager"
11752 .to_string(),
11753 )
11754 })?;
11755 let cu_stream = runtime
11756 .stream_pool()
11757 .resolve(launch_stream)
11758 .ok_or_else(|| {
11759 XlogError::Kernel("indexed inner: launch_stream does not resolve".to_string())
11760 })?;
11761
11762 let num_left = left.num_rows() as u32;
11763 let table = &index.table;
11764
11765 let left_packed =
11767 self.pack_keys_gpu_on_stream(left, left_keys, &cu_stream, launch_stream, runtime)?;
11768 if left_packed.key_bytes != index.key_bytes {
11769 return Err(XlogError::Kernel(
11770 "Join key byte width mismatch between probe and cached index".to_string(),
11771 ));
11772 }
11773
11774 let probe_func = self
11775 .device
11776 .inner()
11777 .get_func(JOIN_MODULE, join_kernels::HASH_JOIN_PROBE_V2)
11778 .ok_or_else(|| XlogError::Kernel("hash_join_probe_v2 kernel not found".to_string()))?;
11779 let block_size = 256u32;
11780 let probe_grid = num_left.div_ceil(block_size);
11781 let probe_config = LaunchConfig {
11782 grid_dim: (probe_grid, 1, 1),
11783 block_dim: (block_size, 1, 1),
11784 shared_mem_bytes: 0,
11785 };
11786
11787 let d_count_only = self.memory.alloc::<u32>(1)?;
11789 let d_dummy_left = self.memory.alloc::<u32>(1)?;
11790 let d_dummy_right = self.memory.alloc::<u32>(1)?;
11791
11792 let max_output_count_only = 0u32;
11793 let mut rec_count = LaunchRecorder::new_strict(launch_stream);
11794 rec_count.read(&left_packed.hashes);
11795 rec_count.read(&left_packed.packed_keys);
11796 rec_count.read(&index.packed_keys);
11797 rec_count.read(&table.bucket_offsets);
11798 rec_count.read(&table.bucket_counts);
11799 rec_count.read(&table.bucket_entries);
11800 rec_count.read(&table.bucket_entry_hashes);
11801 rec_count.write(&d_count_only);
11802 rec_count.write(&d_dummy_left);
11803 rec_count.write(&d_dummy_right);
11804 rec_count.preflight(runtime).map_err(|e| {
11805 XlogError::Kernel(format!("indexed inner: count-pass preflight failed: {}", e))
11806 })?;
11807 let enqueue = |enqueue: &crate::launch::CudaEnqueue<'_>| -> Result<()> {
11809 unsafe {
11811 let res = cudarc::driver::sys::cuMemsetD8Async(
11812 *d_count_only.device_ptr(),
11813 0,
11814 std::mem::size_of::<u32>(),
11815 cu_stream.cu_stream(),
11816 );
11817 if res != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
11818 return Err(XlogError::Kernel(format!(
11819 "cuMemsetD8Async (indexed inner d_count_only) failed: {:?}",
11820 res
11821 )));
11822 }
11823 }
11824 unsafe {
11825 let mut params: Vec<*mut c_void> = vec![
11826 (&left_packed.hashes).as_kernel_param(),
11827 num_left.as_kernel_param(),
11828 (&table.bucket_offsets).as_kernel_param(),
11829 (&table.bucket_counts).as_kernel_param(),
11830 (&table.bucket_entries).as_kernel_param(),
11831 (&table.bucket_entry_hashes).as_kernel_param(),
11832 table.bucket_mask.as_kernel_param(),
11833 (&left_packed.packed_keys).as_kernel_param(),
11834 (&index.packed_keys).as_kernel_param(),
11835 index.key_bytes.as_kernel_param(),
11836 (&d_dummy_left).as_kernel_param(),
11837 (&d_dummy_right).as_kernel_param(),
11838 (&d_count_only).as_kernel_param(),
11839 max_output_count_only.as_kernel_param(),
11840 ];
11841 probe_func
11842 .clone()
11843 .launch_in(enqueue, probe_config, &mut params)
11844 .map_err(|e| {
11845 XlogError::Kernel(format!(
11846 "hash_join_probe_v2 (indexed count, on_stream) failed: {}",
11847 e
11848 ))
11849 })?;
11850 }
11851 Ok(())
11852 };
11853 let rec_count = unsafe { rec_count.enqueue_prepared_with(&cu_stream, enqueue) }
11855 .map_err(LaunchEnqueueError::into_xlog_error)?;
11856 rec_count.commit().map_err(|e| {
11857 XlogError::Kernel(format!("indexed inner: count-pass commit failed: {}", e))
11858 })?;
11859
11860 cu_stream.synchronize().map_err(|e| {
11861 XlogError::Kernel(format!("indexed inner: sync (count read) failed: {}", e))
11862 })?;
11863 let full_count = self.read_join_output_count_metadata(&d_count_only)? as u64;
11864 let requested = max_output
11865 .map(|limit| (limit as u64).min(full_count))
11866 .unwrap_or(full_count);
11867 if requested == 0 {
11868 let combined_schema = self.combine_schemas(left.schema(), right.schema());
11869 return self.create_empty_buffer(combined_schema);
11870 }
11871 if requested > u32::MAX as u64 {
11872 return Err(XlogError::Kernel(format!(
11873 "Join produced {} rows which exceeds the u32 index limit",
11874 requested
11875 )));
11876 }
11877 let max_output_u32 = requested as u32;
11878
11879 let d_output_left = self.memory.alloc::<u32>(max_output_u32 as usize)?;
11881 let d_output_right = self.memory.alloc::<u32>(max_output_u32 as usize)?;
11882 let d_output_count = self.memory.alloc::<u32>(1)?;
11883
11884 let mut rec_mat = LaunchRecorder::new_strict(launch_stream);
11885 rec_mat.read(&left_packed.hashes);
11886 rec_mat.read(&left_packed.packed_keys);
11887 rec_mat.read(&index.packed_keys);
11888 rec_mat.read(&table.bucket_offsets);
11889 rec_mat.read(&table.bucket_counts);
11890 rec_mat.read(&table.bucket_entries);
11891 rec_mat.read(&table.bucket_entry_hashes);
11892 rec_mat.write(&d_output_left);
11893 rec_mat.write(&d_output_right);
11894 rec_mat.write(&d_output_count);
11895 rec_mat.preflight(runtime).map_err(|e| {
11896 XlogError::Kernel(format!(
11897 "indexed inner: materialize preflight failed: {}",
11898 e
11899 ))
11900 })?;
11901 let enqueue = |enqueue: &crate::launch::CudaEnqueue<'_>| -> Result<()> {
11903 unsafe {
11905 let res = cudarc::driver::sys::cuMemsetD8Async(
11906 *d_output_count.device_ptr(),
11907 0,
11908 std::mem::size_of::<u32>(),
11909 cu_stream.cu_stream(),
11910 );
11911 if res != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
11912 return Err(XlogError::Kernel(format!(
11913 "cuMemsetD8Async (indexed inner d_output_count) failed: {:?}",
11914 res
11915 )));
11916 }
11917 }
11918 unsafe {
11919 let mut params: Vec<*mut c_void> = vec![
11920 (&left_packed.hashes).as_kernel_param(),
11921 num_left.as_kernel_param(),
11922 (&table.bucket_offsets).as_kernel_param(),
11923 (&table.bucket_counts).as_kernel_param(),
11924 (&table.bucket_entries).as_kernel_param(),
11925 (&table.bucket_entry_hashes).as_kernel_param(),
11926 table.bucket_mask.as_kernel_param(),
11927 (&left_packed.packed_keys).as_kernel_param(),
11928 (&index.packed_keys).as_kernel_param(),
11929 index.key_bytes.as_kernel_param(),
11930 (&d_output_left).as_kernel_param(),
11931 (&d_output_right).as_kernel_param(),
11932 (&d_output_count).as_kernel_param(),
11933 max_output_u32.as_kernel_param(),
11934 ];
11935 probe_func
11936 .clone()
11937 .launch_in(enqueue, probe_config, &mut params)
11938 .map_err(|e| {
11939 XlogError::Kernel(format!(
11940 "hash_join_probe_v2 (indexed mat, on_stream) failed: {}",
11941 e
11942 ))
11943 })?;
11944 }
11945 Ok(())
11946 };
11947 let rec_mat = unsafe { rec_mat.enqueue_prepared_with(&cu_stream, enqueue) }
11949 .map_err(LaunchEnqueueError::into_xlog_error)?;
11950 rec_mat.commit().map_err(|e| {
11951 XlogError::Kernel(format!("indexed inner: materialize commit failed: {}", e))
11952 })?;
11953
11954 cu_stream.synchronize().map_err(|e| {
11955 XlogError::Kernel(format!("indexed inner: sync (mat read) failed: {}", e))
11956 })?;
11957 let result_count = (self.read_join_output_count_metadata(&d_output_count)? as u64)
11958 .min(max_output_u32 as u64);
11959 if result_count == 0 {
11960 let combined_schema = self.combine_schemas(left.schema(), right.schema());
11961 return self.create_empty_buffer(combined_schema);
11962 }
11963 let output_rows = result_count as u32;
11964
11965 let gathered_left = self.gather_buffer_by_indices_on_stream(
11967 left,
11968 &d_output_left,
11969 output_rows,
11970 &cu_stream,
11971 launch_stream,
11972 runtime,
11973 )?;
11974 let gathered_right = self.gather_buffer_by_indices_on_stream(
11975 right,
11976 &d_output_right,
11977 output_rows,
11978 &cu_stream,
11979 launch_stream,
11980 runtime,
11981 )?;
11982
11983 let combined_schema = self.combine_schemas(left.schema(), right.schema());
11984 let mut result_columns = Vec::with_capacity(combined_schema.arity());
11985 result_columns.extend(gathered_left.columns);
11986 result_columns.extend(gathered_right.columns);
11987 self.buffer_from_columns(result_columns, result_count, combined_schema)
11988 }
11989
11990 fn hash_join_semi_or_anti_v2_with_index_recorded(
11996 &self,
11997 left: &CudaBuffer,
11998 left_keys: &[usize],
11999 index: &crate::provider::JoinIndexV2,
12000 anti: bool,
12001 launch_stream: StreamId,
12002 ) -> Result<CudaBuffer> {
12003 use crate::launch::LaunchRecorder;
12004
12005 let runtime = self.memory.runtime().ok_or_else(|| {
12006 XlogError::Kernel(
12007 "hash_join_v2_with_index_recorded (semi/anti) requires runtime-backed manager"
12008 .to_string(),
12009 )
12010 })?;
12011 let cu_stream = runtime
12012 .stream_pool()
12013 .resolve(launch_stream)
12014 .ok_or_else(|| {
12015 XlogError::Kernel("indexed semi/anti: launch_stream does not resolve".to_string())
12016 })?;
12017
12018 let num_left = left.num_rows() as u32;
12019 let table = &index.table;
12020
12021 let left_packed =
12022 self.pack_keys_gpu_on_stream(left, left_keys, &cu_stream, launch_stream, runtime)?;
12023 if left_packed.key_bytes != index.key_bytes {
12024 return Err(XlogError::Kernel(
12025 "Join key byte width mismatch between probe and cached index".to_string(),
12026 ));
12027 }
12028
12029 let d_mask = self.memory.alloc::<u8>(num_left as usize)?;
12030 let kernel_name = if anti {
12031 join_kernels::HASH_JOIN_ANTI
12032 } else {
12033 join_kernels::HASH_JOIN_SEMI
12034 };
12035 let func = self
12036 .device
12037 .inner()
12038 .get_func(JOIN_MODULE, kernel_name)
12039 .ok_or_else(|| XlogError::Kernel(format!("{} kernel not found", kernel_name)))?;
12040 let block_size = 256u32;
12041 let grid_size = num_left.div_ceil(block_size);
12042 let cfg = LaunchConfig {
12043 grid_dim: (grid_size, 1, 1),
12044 block_dim: (block_size, 1, 1),
12045 shared_mem_bytes: 0,
12046 };
12047
12048 let mut rec = LaunchRecorder::new_strict(launch_stream);
12049 rec.read(&left_packed.hashes);
12050 rec.read(&left_packed.packed_keys);
12051 rec.read(&index.packed_keys);
12052 rec.read(&table.bucket_offsets);
12053 rec.read(&table.bucket_counts);
12054 rec.read(&table.bucket_entries);
12055 rec.read(&table.bucket_entry_hashes);
12056 rec.write(&d_mask);
12057 rec.preflight(runtime).map_err(|e| {
12058 XlogError::Kernel(format!("indexed semi/anti: preflight failed: {}", e))
12059 })?;
12060 let enqueue = |enqueue: &crate::launch::CudaEnqueue<'_>| -> Result<()> {
12062 unsafe {
12063 func.clone().launch_in(
12064 enqueue,
12065 cfg,
12066 (
12067 &left_packed.hashes,
12068 num_left,
12069 &table.bucket_offsets,
12070 &table.bucket_counts,
12071 &table.bucket_entries,
12072 &table.bucket_entry_hashes,
12073 table.bucket_mask,
12074 &left_packed.packed_keys,
12075 &index.packed_keys,
12076 index.key_bytes,
12077 &d_mask,
12078 ),
12079 )
12080 }
12081 .map_err(|e| {
12082 XlogError::Kernel(format!(
12083 "{} (on_stream, indexed) failed: {}",
12084 kernel_name, e
12085 ))
12086 })
12087 };
12088 let rec = unsafe { rec.enqueue_prepared_with(&cu_stream, enqueue) }
12090 .map_err(LaunchEnqueueError::into_xlog_error)?;
12091 rec.commit()
12092 .map_err(|e| XlogError::Kernel(format!("indexed semi/anti: commit failed: {}", e)))?;
12093
12094 self.compact_buffer_by_device_mask_counted_recorded(left, &d_mask, launch_stream)
12095 }
12096
12097 fn hash_join_left_outer_v2_with_index_recorded(
12103 &self,
12104 left: &CudaBuffer,
12105 right: &CudaBuffer,
12106 left_keys: &[usize],
12107 index: &crate::provider::JoinIndexV2,
12108 max_output: Option<usize>,
12109 launch_stream: StreamId,
12110 ) -> Result<CudaBuffer> {
12111 use crate::launch::LaunchRecorder;
12112
12113 let runtime = self.memory.runtime().ok_or_else(|| {
12114 XlogError::Kernel(
12115 "hash_join_v2_with_index_recorded (left_outer) requires runtime-backed manager"
12116 .to_string(),
12117 )
12118 })?;
12119 let cu_stream = runtime
12120 .stream_pool()
12121 .resolve(launch_stream)
12122 .ok_or_else(|| {
12123 XlogError::Kernel("indexed left_outer: launch_stream does not resolve".to_string())
12124 })?;
12125
12126 let num_left = left.num_rows() as u32;
12127 let table = &index.table;
12128
12129 let left_packed =
12130 self.pack_keys_gpu_on_stream(left, left_keys, &cu_stream, launch_stream, runtime)?;
12131 if left_packed.key_bytes != index.key_bytes {
12132 return Err(XlogError::Kernel(
12133 "Join key byte width mismatch between probe and cached index".to_string(),
12134 ));
12135 }
12136
12137 let device = self.device.inner();
12138 let block_size = 256u32;
12139 let grid_size = num_left.div_ceil(block_size);
12140 let cfg = LaunchConfig {
12141 grid_dim: (grid_size, 1, 1),
12142 block_dim: (block_size, 1, 1),
12143 shared_mem_bytes: 0,
12144 };
12145
12146 let d_has_match = self.memory.alloc::<u8>(num_left as usize)?;
12148 let d_count_only = self.memory.alloc::<u32>(1)?;
12149 let d_dummy_left = self.memory.alloc::<u32>(1)?;
12150 let d_dummy_right = self.memory.alloc::<u32>(1)?;
12151 let semi_func = device
12152 .get_func(JOIN_MODULE, join_kernels::HASH_JOIN_SEMI)
12153 .ok_or_else(|| XlogError::Kernel("hash_join_semi kernel not found".to_string()))?;
12154 let probe_func = device
12155 .get_func(JOIN_MODULE, join_kernels::HASH_JOIN_PROBE_V2)
12156 .ok_or_else(|| XlogError::Kernel("hash_join_probe_v2 kernel not found".to_string()))?;
12157
12158 let mut rec_a = LaunchRecorder::new_strict(launch_stream);
12159 rec_a.read(&left_packed.hashes);
12160 rec_a.read(&left_packed.packed_keys);
12161 rec_a.read(&index.packed_keys);
12162 rec_a.read(&table.bucket_offsets);
12163 rec_a.read(&table.bucket_counts);
12164 rec_a.read(&table.bucket_entries);
12165 rec_a.read(&table.bucket_entry_hashes);
12166 rec_a.write(&d_has_match);
12167 rec_a.write(&d_count_only);
12168 rec_a.write(&d_dummy_left);
12169 rec_a.write(&d_dummy_right);
12170 rec_a.preflight(runtime).map_err(|e| {
12171 XlogError::Kernel(format!(
12172 "indexed left_outer: semi/count preflight failed: {}",
12173 e
12174 ))
12175 })?;
12176 let enqueue = |enqueue: &crate::launch::CudaEnqueue<'_>| -> Result<()> {
12178 unsafe {
12180 let res = cudarc::driver::sys::cuMemsetD8Async(
12181 *d_count_only.device_ptr(),
12182 0,
12183 std::mem::size_of::<u32>(),
12184 cu_stream.cu_stream(),
12185 );
12186 if res != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
12187 return Err(XlogError::Kernel(format!(
12188 "cuMemsetD8Async (indexed left_outer d_count_only) failed: {:?}",
12189 res
12190 )));
12191 }
12192 }
12193 unsafe {
12194 semi_func.clone().launch_in(
12195 enqueue,
12196 cfg,
12197 (
12198 &left_packed.hashes,
12199 num_left,
12200 &table.bucket_offsets,
12201 &table.bucket_counts,
12202 &table.bucket_entries,
12203 &table.bucket_entry_hashes,
12204 table.bucket_mask,
12205 &left_packed.packed_keys,
12206 &index.packed_keys,
12207 index.key_bytes,
12208 &d_has_match,
12209 ),
12210 )
12211 }
12212 .map_err(|e| {
12213 XlogError::Kernel(format!(
12214 "hash_join_semi (on_stream, indexed left_outer) failed: {}",
12215 e
12216 ))
12217 })?;
12218
12219 let max_output_count_only = 0u32;
12220 unsafe {
12222 let mut params: Vec<*mut c_void> = vec![
12223 (&left_packed.hashes).as_kernel_param(),
12224 num_left.as_kernel_param(),
12225 (&table.bucket_offsets).as_kernel_param(),
12226 (&table.bucket_counts).as_kernel_param(),
12227 (&table.bucket_entries).as_kernel_param(),
12228 (&table.bucket_entry_hashes).as_kernel_param(),
12229 table.bucket_mask.as_kernel_param(),
12230 (&left_packed.packed_keys).as_kernel_param(),
12231 (&index.packed_keys).as_kernel_param(),
12232 index.key_bytes.as_kernel_param(),
12233 (&d_dummy_left).as_kernel_param(),
12234 (&d_dummy_right).as_kernel_param(),
12235 (&d_count_only).as_kernel_param(),
12236 max_output_count_only.as_kernel_param(),
12237 ];
12238 probe_func
12239 .clone()
12240 .launch_in(enqueue, cfg, &mut params)
12241 .map_err(|e| {
12242 XlogError::Kernel(format!(
12243 "hash_join_probe_v2 (count, on_stream, indexed left_outer) failed: {}",
12244 e
12245 ))
12246 })?;
12247 }
12248 Ok(())
12249 };
12250 let rec_a = unsafe { rec_a.enqueue_prepared_with(&cu_stream, enqueue) }
12252 .map_err(LaunchEnqueueError::into_xlog_error)?;
12253 rec_a.commit().map_err(|e| {
12254 XlogError::Kernel(format!(
12255 "indexed left_outer: semi/count commit failed: {}",
12256 e
12257 ))
12258 })?;
12259
12260 cu_stream.synchronize().map_err(|e| {
12261 XlogError::Kernel(format!(
12262 "indexed left_outer: sync (count read) failed: {}",
12263 e
12264 ))
12265 })?;
12266 let full_inner = self.read_join_output_count_metadata(&d_count_only)? as u64;
12267 let requested_inner = max_output
12268 .map(|limit| (limit as u64).min(full_inner))
12269 .unwrap_or(full_inner);
12270 if requested_inner > u32::MAX as u64 {
12271 return Err(XlogError::Kernel(format!(
12272 "Join produced {} rows which exceeds the u32 index limit",
12273 requested_inner
12274 )));
12275 }
12276 let max_output_u32 = requested_inner as u32;
12277 let alloc_len = (requested_inner.max(1)) as usize;
12278
12279 let d_output_left = self.memory.alloc::<u32>(alloc_len)?;
12281 let d_output_right = self.memory.alloc::<u32>(alloc_len)?;
12282 let d_output_count = self.memory.alloc::<u32>(1)?;
12283
12284 let mut rec_b = LaunchRecorder::new_strict(launch_stream);
12285 rec_b.read(&left_packed.hashes);
12286 rec_b.read(&left_packed.packed_keys);
12287 rec_b.read(&index.packed_keys);
12288 rec_b.read(&table.bucket_offsets);
12289 rec_b.read(&table.bucket_counts);
12290 rec_b.read(&table.bucket_entries);
12291 rec_b.read(&table.bucket_entry_hashes);
12292 rec_b.write(&d_output_left);
12293 rec_b.write(&d_output_right);
12294 rec_b.write(&d_output_count);
12295 rec_b.preflight(runtime).map_err(|e| {
12296 XlogError::Kernel(format!(
12297 "indexed left_outer: materialize preflight failed: {}",
12298 e
12299 ))
12300 })?;
12301 let enqueue = |enqueue: &crate::launch::CudaEnqueue<'_>| -> Result<()> {
12303 unsafe {
12305 let res = cudarc::driver::sys::cuMemsetD8Async(
12306 *d_output_count.device_ptr(),
12307 0,
12308 std::mem::size_of::<u32>(),
12309 cu_stream.cu_stream(),
12310 );
12311 if res != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
12312 return Err(XlogError::Kernel(format!(
12313 "cuMemsetD8Async (indexed left_outer d_output_count) failed: {:?}",
12314 res
12315 )));
12316 }
12317 }
12318 unsafe {
12319 let mut params: Vec<*mut c_void> = vec![
12320 (&left_packed.hashes).as_kernel_param(),
12321 num_left.as_kernel_param(),
12322 (&table.bucket_offsets).as_kernel_param(),
12323 (&table.bucket_counts).as_kernel_param(),
12324 (&table.bucket_entries).as_kernel_param(),
12325 (&table.bucket_entry_hashes).as_kernel_param(),
12326 table.bucket_mask.as_kernel_param(),
12327 (&left_packed.packed_keys).as_kernel_param(),
12328 (&index.packed_keys).as_kernel_param(),
12329 index.key_bytes.as_kernel_param(),
12330 (&d_output_left).as_kernel_param(),
12331 (&d_output_right).as_kernel_param(),
12332 (&d_output_count).as_kernel_param(),
12333 max_output_u32.as_kernel_param(),
12334 ];
12335 probe_func
12336 .clone()
12337 .launch_in(enqueue, cfg, &mut params)
12338 .map_err(|e| {
12339 XlogError::Kernel(format!(
12340 "hash_join_probe_v2 (mat, on_stream, indexed left_outer) failed: {}",
12341 e
12342 ))
12343 })?;
12344 }
12345 Ok(())
12346 };
12347 let rec_b = unsafe { rec_b.enqueue_prepared_with(&cu_stream, enqueue) }
12349 .map_err(LaunchEnqueueError::into_xlog_error)?;
12350 rec_b.commit().map_err(|e| {
12351 XlogError::Kernel(format!(
12352 "indexed left_outer: materialize commit failed: {}",
12353 e
12354 ))
12355 })?;
12356
12357 cu_stream.synchronize().map_err(|e| {
12358 XlogError::Kernel(format!("indexed left_outer: sync (mat read) failed: {}", e))
12359 })?;
12360 let inner_count = self
12361 .read_join_output_count_metadata(&d_output_count)?
12362 .min(max_output_u32);
12363
12364 let d_no_match = self.memory.alloc::<u8>(num_left as usize)?;
12366 let mask_not_fn = device
12367 .get_func(FILTER_MODULE, filter_kernels::MASK_NOT)
12368 .ok_or_else(|| XlogError::Kernel("mask_not kernel not found".to_string()))?;
12369 let mut rec_c = LaunchRecorder::new_strict(launch_stream);
12370 rec_c.read(&d_has_match);
12371 rec_c.write(&d_no_match);
12372 rec_c.preflight(runtime).map_err(|e| {
12373 XlogError::Kernel(format!(
12374 "indexed left_outer: mask_not preflight failed: {}",
12375 e
12376 ))
12377 })?;
12378 let enqueue = |enqueue: &crate::launch::CudaEnqueue<'_>| -> Result<()> {
12380 unsafe {
12381 mask_not_fn
12382 .clone()
12383 .launch_in(enqueue, cfg, (&d_has_match, &d_no_match, num_left))
12384 }
12385 .map_err(|e| {
12386 XlogError::Kernel(format!(
12387 "mask_not (on_stream, indexed left_outer) failed: {}",
12388 e
12389 ))
12390 })
12391 };
12392 let rec_c = unsafe { rec_c.enqueue_prepared_with(&cu_stream, enqueue) }
12394 .map_err(LaunchEnqueueError::into_xlog_error)?;
12395 rec_c.commit().map_err(|e| {
12396 XlogError::Kernel(format!("indexed left_outer: mask_not commit failed: {}", e))
12397 })?;
12398
12399 let unmatched_left =
12400 self.compact_buffer_by_device_mask_counted_recorded(left, &d_no_match, launch_stream)?;
12401 let unmatched_rows = self.device_row_count(&unmatched_left)? as u64;
12402 let total_rows = (inner_count as u64) + unmatched_rows;
12403
12404 let combined_schema = self.combine_schemas(left.schema(), right.schema());
12405 if total_rows == 0 {
12406 return self.create_empty_buffer(combined_schema);
12407 }
12408
12409 let inner_count_u32 = inner_count;
12415 let inner_left_buf;
12416 let inner_right_buf;
12417 if inner_count > 0 {
12418 inner_left_buf = Some(self.gather_buffer_by_indices_on_stream(
12419 left,
12420 &d_output_left,
12421 inner_count_u32,
12422 &cu_stream,
12423 launch_stream,
12424 runtime,
12425 )?);
12426 inner_right_buf = Some(self.gather_buffer_by_indices_on_stream(
12427 right,
12428 &d_output_right,
12429 inner_count_u32,
12430 &cu_stream,
12431 launch_stream,
12432 runtime,
12433 )?);
12434 } else {
12435 inner_left_buf = None;
12436 inner_right_buf = None;
12437 }
12438
12439 let result_columns = self.concatenate_left_outer_columns_on_stream(
12445 left,
12446 right,
12447 inner_left_buf.as_ref(),
12448 inner_right_buf.as_ref(),
12449 &unmatched_left,
12450 inner_count as u64,
12451 unmatched_rows,
12452 &cu_stream,
12453 launch_stream,
12454 runtime,
12455 "indexed left_outer",
12456 )?;
12457
12458 let d_num_rows = self.upload_device_row_count(total_rows as u32)?;
12459 Ok(CudaBuffer::from_columns_with_host_count(
12460 result_columns,
12461 total_rows,
12462 d_num_rows,
12463 combined_schema,
12464 total_rows as u32,
12465 ))
12466 }
12467}