viirya commented on code in PR #11276:
URL: https://github.com/apache/datafusion/pull/11276#discussion_r1667615450
##########
datafusion/physical-plan/src/joins/utils.rs:
##########
@@ -1411,6 +1424,63 @@ where
.collect::<PrimitiveArray<T>>()
}
+/// Appends probe indices in order by considering the given build indices.
+///
+/// This function constructs new build and probe indices by iterating through
+/// the provided indices, and appends any missing values between previous and
+/// current probe index with a corresponding null build index. It handles
various
+/// edge cases and returns an error if either index is `None`.
+///
+/// # Parameters
+/// - `build_indices`: `PrimitiveArray` of `UInt64Type` containing build
indices.
+/// - `probe_indices`: `PrimitiveArray` of `UInt32Type` containing probe
indices.
+/// - `count_probe_batch`: The number of elements in the probe batch, used for
+/// filling in any remaining indices.
+///
+/// # Returns
+/// A `Result` containing a tuple of two arrays:
+/// - A `PrimitiveArray` of `UInt64Type` with the newly constructed build
indices.
+/// - A `PrimitiveArray` of `UInt32Type` with the newly constructed probe
indices.
+///
+/// # Errors
+/// Returns an error if there is a failure in calculating probe indices.
+fn append_probe_indices_in_order(
+ build_indices: PrimitiveArray<UInt64Type>,
+ probe_indices: PrimitiveArray<UInt32Type>,
+ range: Range<usize>,
+) -> (PrimitiveArray<UInt64Type>, PrimitiveArray<UInt32Type>) {
+ // Builders for new indices:
+ let mut new_build_indices = UInt64Builder::new();
+ let mut new_probe_indices = UInt32Builder::new();
+ // Set previous index as zero for the initial loop:
+ let mut prev_index = range.start as u32;
+ // Zip the two iterators.
+ debug_assert!(build_indices.len() == probe_indices.len());
+ for (build_index, probe_index) in build_indices
+ .values()
+ .into_iter()
+ .zip(probe_indices.values().into_iter())
+ {
+ // Append values between previous and current left index with null
right index:
+ for value in prev_index..*probe_index {
+ new_probe_indices.append_value(value);
+ new_build_indices.append_null();
+ }
+ // Append current indices:
+ new_probe_indices.append_value(*probe_index);
+ new_build_indices.append_value(*build_index);
+ // Set current left index as previous for the next iteration:
Review Comment:
```suggestion
// Set current probe index as previous for the next iteration:
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
