SubhamSinghal commented on code in PR #24457:
URL: https://github.com/apache/datafusion/pull/24457#discussion_r3954234322
##########
datafusion/physical-plan/src/joins/piecewise_merge_join/exec.rs:
##########
@@ -941,6 +1070,147 @@ impl BufferedSideData {
}
}
+/// The entire buffered side of a right existence join, reduced to the one key
that decides
+/// every streamed row -- the minimum for `<`/`<=`, the maximum for `>`/`>=`.
+///
+/// Right existence joins never look at a buffered *row*, so unlike
[`BufferedSideData`] this
+/// holds no batch: the buffered input is folded away as it streams in and the
state is `O(1)`
+/// however large that side is. See `right_existence_join.rs`.
+pub(super) struct BufferedExtreme {
+ /// One-row array. The row is null exactly when no buffered key is
non-null (an empty or
+ /// all-NULL buffered side), in which case nothing can ever match.
+ extreme: ArrayRef,
+ _reservation: MemoryReservation,
+}
+
+impl BufferedExtreme {
+ pub(super) fn extreme(&self) -> &ArrayRef {
+ &self.extreme
+ }
+}
+
+/// Reduces one buffered partition to a single extreme key, dropping each
batch as it goes.
+/// `None` when the partition produced no batch at all.
+async fn partition_extreme(
+ mut buffered: SendableRecordBatchStream,
+ on_buffered: PhysicalExprRef,
+ metrics: BuildProbeJoinMetrics,
+ reservation: MemoryReservation,
+ descending: bool,
+) -> Result<Option<ArrayRef>> {
+ let mut extreme: Option<ArrayRef> = None;
+
+ while let Some(batch) = buffered.next().await.transpose()? {
+ metrics.build_input_batches.add(1);
+ metrics.build_input_rows.add(batch.num_rows());
+
+ let keys = on_buffered.evaluate(&batch)?.into_array(batch.num_rows())?;
+
+ // Reduced and dropped within this iteration, but as wide as the
batch, and every
+ // partition folds one of these at once. Resized rather than grown, so
what the pool
+ // sees is the widest key array in flight here and not their sum.
+ reservation.try_resize(keys.get_array_memory_size())?;
+
+ let batch_extreme = extreme_key(&keys, descending)?;
+
+ // Re-reduced as a pair rather than compared, so the running value is
ordered exactly
+ // as each single reduction was. `extreme_key` ignores nulls, so a
null from an
+ // all-NULL batch never displaces a real key.
+ extreme = Some(match extreme {
+ Some(running) => {
+ let pair = concat(&[running.as_ref(),
batch_extreme.as_ref()])?;
+ extreme_key(&pair, descending)?
+ }
+ None => batch_extreme,
+ });
+ }
+
+ Ok(extreme)
+}
+
+/// Folds every buffered partition down to one extreme key. `O(B)` time, and
the only state it
+/// retains is that one key -- no buffered batch is concatenated or held. The
transient cost is
+/// one key array per folding partition, which each task accounts against the
pool for as long
+/// as it holds it.
+///
+/// A min/max combines across partitions, so this side needs no
single-partition funnel --
+/// `input_distribution_requirements` asks for `UnspecifiedDistribution` here
and the partitions
+/// are folded independently, each in its own spawned task.
+async fn build_buffered_extreme(
+ buffered_streams: Vec<SendableRecordBatchStream>,
+ buffered_schema: SchemaRef,
+ on_buffered: PhysicalExprRef,
+ metrics: BuildProbeJoinMetrics,
+ reservation: MemoryReservation,
+ memory_pool: Arc<dyn MemoryPool>,
+ descending: bool,
+) -> Result<BufferedExtreme> {
+ let tasks: Vec<_> = buffered_streams
+ .into_iter()
+ .enumerate()
+ .map(|(partition, stream)| {
+ let on_buffered = Arc::clone(&on_buffered);
+ let metrics = metrics.clone();
+ // One reservation per task rather than a shared one:
`MemoryReservation` is not
+ // shareable, and each task's transient is freed as soon as it
finishes.
+ let reservation = MemoryConsumer::new(format!(
+ "PiecewiseMergeJoinBufferedFold[{partition}]"
+ ))
+ .register(&memory_pool);
+ SpawnedTask::spawn(partition_extreme(
+ stream,
+ on_buffered,
+ metrics,
+ reservation,
+ descending,
+ ))
+ })
+ .collect();
+
+ // The tasks run concurrently; this only collects them. Awaiting in order
is fine, and a
+ // failure propagates after the rest have been joined, since `SpawnedTask`
aborts on drop.
+ let mut extreme: Option<ArrayRef> = None;
+ for task in tasks {
+ let partition_extreme = task.join_unwind().await.map_err(|e| {
+ internal_datafusion_err!("buffered extreme task failed: {e}")
+ })??;
+ if let Some(partition_extreme) = partition_extreme {
+ // Re-reduced as a pair, exactly as the batches within a partition
were, so a value
+ // accumulated across partitions is ordered by the same rule.
+ extreme = Some(match extreme {
+ Some(running) => {
+ let pair = concat(&[running.as_ref(),
partition_extreme.as_ref()])?;
+ extreme_key(&pair, descending)?
+ }
+ None => partition_extreme,
+ });
Review Comment:
Addressed in
[1665517](https://github.com/apache/datafusion/pull/24457/commits/16655173354a996e7c16c3160cc4bbfc8421fc74)
--
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]
