You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@systemds.apache.org by ja...@apache.org on 2026/08/08 10:39:52 UTC
(systemds) branch main updated: [SYSTEMDS-3891] OOC Non-blocking Join and Primitive Simplification
This is an automated email from the ASF dual-hosted git repository.
janniklinde pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/systemds.git
The following commit(s) were added to refs/heads/main by this push:
new 9dccbc3e60 [SYSTEMDS-3891] OOC Non-blocking Join and Primitive Simplification
9dccbc3e60 is described below
commit 9dccbc3e6013b91d8ddcb416cadae2d993957d7b
Author: Jannik Lindemann <52...@users.noreply.github.com>
AuthorDate: Sat Aug 8 12:39:46 2026 +0200
[SYSTEMDS-3891] OOC Non-blocking Join and Primitive Simplification
---
.../ooc/primitives/BroadcastOOCPrimitive.java | 13 +--
.../ooc/primitives/GroupedReduceOOCPrimitive.java | 15 +--
.../runtime/ooc/primitives/JoinOOCPrimitive.java | 117 ++++++++++++++-------
.../ooc/primitives/MaterializeOOCPrimitive.java | 6 --
.../sysds/runtime/ooc/primitives/OOCPrimitive.java | 15 +++
.../primitives/PlannableDataGenOOCPrimitive.java | 3 +-
.../runtime/ooc/util/OOCInstructionUtils.java | 19 ++++
.../component/ooc/OOCInstructionUtilsTest.java | 32 ++++++
8 files changed, 147 insertions(+), 73 deletions(-)
diff --git a/src/main/java/org/apache/sysds/runtime/ooc/primitives/BroadcastOOCPrimitive.java b/src/main/java/org/apache/sysds/runtime/ooc/primitives/BroadcastOOCPrimitive.java
index e2a9be63c2..49980c2895 100644
--- a/src/main/java/org/apache/sysds/runtime/ooc/primitives/BroadcastOOCPrimitive.java
+++ b/src/main/java/org/apache/sysds/runtime/ooc/primitives/BroadcastOOCPrimitive.java
@@ -51,7 +51,6 @@ public final class BroadcastOOCPrimitive extends OOCPrimitive {
private final Supplier<MaterializedStore.Liveness> _liveness;
private final BiFunction<IndexedMatrixValue, IndexedMatrixValue, IndexedMatrixValue> _operation;
private final AtomicBoolean _cleaned;
- private final AtomicBoolean _failed;
private final AtomicBoolean _sourceComplete;
private final AtomicInteger _active;
private MaterializedStore<IndexedMatrixValue> _store;
@@ -70,7 +69,6 @@ public final class BroadcastOOCPrimitive extends OOCPrimitive {
_liveness = liveness;
_operation = operation;
_cleaned = new AtomicBoolean();
- _failed = new AtomicBoolean();
_sourceComplete = new AtomicBoolean();
_active = new AtomicInteger(1);
}
@@ -100,7 +98,7 @@ public final class BroadcastOOCPrimitive extends OOCPrimitive {
_outputStream = _output.getWriteStream();
_ready = new SubscribableTaskQueue<>();
getContext().addOutStream(_outputStream, _ready);
- OOCInstructionUtils.submitOOCTasks(_ready, callback -> process(callback.get()), getContext())
+ OOCInstructionUtils.submitCloseableOOCTasks(_ready, this::process, getContext())
.whenComplete((ignored, error) -> {
try {
if(error != null)
@@ -232,7 +230,6 @@ public final class BroadcastOOCPrimitive extends OOCPrimitive {
fail(failure);
}
finally {
- work.close();
if(budget != null)
budget.close();
completeOne();
@@ -255,14 +252,6 @@ public final class BroadcastOOCPrimitive extends OOCPrimitive {
}
}
- private void fail(Throwable error) {
- if(!_failed.compareAndSet(false, true))
- return;
- DMLRuntimeException failure = DMLRuntimeException.of(error);
- _outputStream.propagateFailure(failure);
- getContext().failAll(failure);
- }
-
private void cleanup() {
if(!_cleaned.compareAndSet(false, true))
return;
diff --git a/src/main/java/org/apache/sysds/runtime/ooc/primitives/GroupedReduceOOCPrimitive.java b/src/main/java/org/apache/sysds/runtime/ooc/primitives/GroupedReduceOOCPrimitive.java
index aa67497206..b380848aa7 100644
--- a/src/main/java/org/apache/sysds/runtime/ooc/primitives/GroupedReduceOOCPrimitive.java
+++ b/src/main/java/org/apache/sysds/runtime/ooc/primitives/GroupedReduceOOCPrimitive.java
@@ -49,7 +49,6 @@ public final class GroupedReduceOOCPrimitive extends OOCPrimitive {
private final OOCStreamable<IndexedMatrixValue> _output;
private final BiFunction<MatrixBlock, MatrixBlock, MatrixBlock> _merge;
private final AtomicBoolean _cleaned;
- private final AtomicBoolean _failed;
private final AtomicBoolean _sourceComplete;
private final AtomicInteger _active;
private final AtomicInteger _finalizedGroups;
@@ -66,7 +65,6 @@ public final class GroupedReduceOOCPrimitive extends OOCPrimitive {
_output = output;
_merge = merge;
_cleaned = new AtomicBoolean();
- _failed = new AtomicBoolean();
_sourceComplete = new AtomicBoolean();
_active = new AtomicInteger(1);
_finalizedGroups = new AtomicInteger();
@@ -100,7 +98,7 @@ public final class GroupedReduceOOCPrimitive extends OOCPrimitive {
getContext().addInStream(input).addOutStream(_outputStream, _ready);
_table = new StateTable<>(OOCCacheManager.getGlobalCache(), CachingStream._streamSeq.getNextID());
- OOCInstructionUtils.submitOOCTasks(_ready, callback -> process(callback.get()), getContext())
+ OOCInstructionUtils.submitCloseableOOCTasks(_ready, this::process, getContext())
.whenComplete((ignored, error) -> {
try {
_outputStream.closeInput();
@@ -222,7 +220,6 @@ public final class GroupedReduceOOCPrimitive extends OOCPrimitive {
catch(Throwable failure) {
if(merged != null)
merged.release();
- work.close();
budget.close();
fail(failure);
completeOne();
@@ -276,7 +273,7 @@ public final class GroupedReduceOOCPrimitive extends OOCPrimitive {
int remaining = _active.decrementAndGet();
if(remaining != 0)
return;
- if(!_failed.get() && _finalizedGroups.get() != _numGroups)
+ if(!hasFailed() && _finalizedGroups.get() != _numGroups)
fail(new DMLRuntimeException(
"Grouped reduction completed " + _finalizedGroups.get() + " of " + _numGroups + " row groups."));
try {
@@ -287,14 +284,6 @@ public final class GroupedReduceOOCPrimitive extends OOCPrimitive {
}
}
- private void fail(Throwable error) {
- if(!_failed.compareAndSet(false, true))
- return;
- DMLRuntimeException failure = DMLRuntimeException.of(error);
- _outputStream.propagateFailure(failure);
- getContext().failAll(failure);
- }
-
private void cleanup() {
if(!_cleaned.compareAndSet(false, true))
return;
diff --git a/src/main/java/org/apache/sysds/runtime/ooc/primitives/JoinOOCPrimitive.java b/src/main/java/org/apache/sysds/runtime/ooc/primitives/JoinOOCPrimitive.java
index a4ef2c30ce..0ec802b181 100644
--- a/src/main/java/org/apache/sysds/runtime/ooc/primitives/JoinOOCPrimitive.java
+++ b/src/main/java/org/apache/sysds/runtime/ooc/primitives/JoinOOCPrimitive.java
@@ -19,7 +19,8 @@
package org.apache.sysds.runtime.ooc.primitives;
-import java.util.concurrent.ExecutionException;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.BiFunction;
import org.apache.sysds.runtime.DMLRuntimeException;
@@ -44,6 +45,9 @@ public class JoinOOCPrimitive extends OOCPrimitive {
private final OOCStreamable<IndexedMatrixValue> _right;
private final OOCStreamable<IndexedMatrixValue> _output;
private final BiFunction<MatrixBlock, MatrixBlock, MatrixBlock> _operation;
+ private final AtomicInteger _pending = new AtomicInteger(1);
+ private final AtomicInteger _unmatched = new AtomicInteger();
+ private final CompletableFuture<Void> _pendingCompletion = new CompletableFuture<>();
private StateTable<IndexedMatrixValue> _table;
public JoinOOCPrimitive(OOCStreamable<IndexedMatrixValue> left, OOCStreamable<IndexedMatrixValue> right,
@@ -87,16 +91,13 @@ public class JoinOOCPrimitive extends OOCPrimitive {
long taskBytes = outputBytes + 2 * inputBytes;
getContext().addOutStream(output);
- OOCInstructionUtils.submitOOCTasks(matches, callback -> {
- try(JoinWork work = callback.get()) {
- IndexedMatrixValue mleft = work._left.get();
- IndexedMatrixValue mright = work._right.get();
- OOCUtils.enqueueExact(output,
- new IndexedMatrixValue(mleft.getIndexes(),
- _operation.apply((MatrixBlock) mleft.getValue(), (MatrixBlock) mright.getValue())),
- work._budget);
- }
- }, callback -> true, (index, callback) -> callback.get().close(), getContext()).thenRun(() -> {
+ CompletableFuture<Void> processing = OOCInstructionUtils.submitCloseableOOCTasks(matches, (JoinWork work) -> {
+ IndexedMatrixValue mleft = work._left.get();
+ IndexedMatrixValue mright = work._right.get();
+ OOCUtils.enqueueExact(output, new IndexedMatrixValue(mleft.getIndexes(),
+ _operation.apply((MatrixBlock) mleft.getValue(), (MatrixBlock) mright.getValue())), work._budget);
+ }, getContext());
+ CompletableFuture.allOf(processing, _pendingCompletion).thenRun(() -> {
try {
_table.close();
onComplete();
@@ -113,7 +114,6 @@ public class JoinOOCPrimitive extends OOCPrimitive {
private void drive(OOCStream<IndexedMatrixValue> leftInput, OOCStream<IndexedMatrixValue> rightInput,
OOCStream<JoinWork> matches, long taskBytes) {
long cols = _right.getDataCharacteristics().getNumColBlocks();
- int unmatched = 0;
try {
while(true) {
OOCStream.QueueCallback<IndexedMatrixValue> left = leftInput.dequeueCB();
@@ -129,23 +129,26 @@ public class JoinOOCPrimitive extends OOCPrimitive {
throw new DMLRuntimeException("Join inputs contain a different number of blocks");
break;
}
- unmatched += accept(left, true, cols, taskBytes, matches);
- unmatched += accept(right, false, cols, taskBytes, matches);
+ accept(left, true, cols, taskBytes, matches);
+ accept(right, false, cols, taskBytes, matches);
}
- if(unmatched != 0)
- throw new DMLRuntimeException("Join inputs contain " + unmatched + " unmatched blocks");
+ }
+ catch(Throwable failure) {
+ fail(failure);
+ throw DMLRuntimeException.of(failure);
}
finally {
- matches.closeInput();
+ completePending(matches);
}
}
- private int accept(OOCStream.QueueCallback<IndexedMatrixValue> callback, boolean left, long cols, long taskBytes,
+ private void accept(OOCStream.QueueCallback<IndexedMatrixValue> callback, boolean left, long cols, long taskBytes,
OOCStream<JoinWork> matches) {
if(callback == null)
- return 0;
+ return;
OOCStream.QueueCallback<IndexedMatrixValue> owned = null;
ReservationBudget budget = null;
+ boolean pending = false;
try {
owned = callback.keepOpen();
callback.close();
@@ -155,25 +158,18 @@ public class JoinOOCPrimitive extends OOCPrimitive {
long row = value.getIndexes().getRowIndex() - 1;
long col = value.getIndexes().getColumnIndex() - 1;
int slot = Math.toIntExact(row * cols + col);
+ _pending.incrementAndGet();
+ pending = true;
OOCFuture<StateTableUtils.Match> future = StateTableUtils.putOrTake(_table, slot, owned, budget);
owned = null;
- StateTableUtils.Match match = await(future);
- if(match == null)
- return 1;
- JoinWork work = left ? new JoinWork(match.left(), match.right(), budget) : new JoinWork(match.right(),
- match.left(), budget);
+ ReservationBudget pendingBudget = budget;
budget = null;
- try {
- matches.enqueue(work);
- work = null;
- }
- finally {
- if(work != null)
- work.close();
- }
- return -1;
+ future.whenComplete((match, error) -> matchReady(match, left, pendingBudget, error, matches));
+ pending = false;
}
finally {
+ if(pending)
+ completePending(matches);
if(callback != null)
callback.close();
if(owned != null)
@@ -183,16 +179,57 @@ public class JoinOOCPrimitive extends OOCPrimitive {
}
}
- private static StateTableUtils.Match await(OOCFuture<StateTableUtils.Match> future) {
+ private void matchReady(StateTableUtils.Match match, boolean left, ReservationBudget budget, Throwable error,
+ OOCStream<JoinWork> matches) {
+ JoinWork work = null;
try {
- return future.get();
+ if(error != null)
+ throw DMLRuntimeException.of(error);
+ if(match == null) {
+ _unmatched.incrementAndGet();
+ return;
+ }
+ _unmatched.decrementAndGet();
+ work = left ? new JoinWork(match.left(), match.right(), budget) : new JoinWork(match.right(), match.left(),
+ budget);
+ match = null;
+ budget = null;
+ matches.enqueue(work);
+ work = null;
+ }
+ catch(Throwable failure) {
+ fail(failure);
+ }
+ finally {
+ if(work != null)
+ work.close();
+ if(match != null) {
+ match.left().close();
+ match.right().close();
+ }
+ if(budget != null)
+ budget.close();
+ completePending(matches);
}
- catch(InterruptedException error) {
- Thread.currentThread().interrupt();
- throw new DMLRuntimeException(error);
+ }
+
+ private void completePending(OOCStream<JoinWork> matches) {
+ if(_pending.decrementAndGet() != 0)
+ return;
+ try {
+ int unmatched = _unmatched.get();
+ if(unmatched != 0)
+ fail(new DMLRuntimeException("Join inputs contain " + unmatched + " unmatched blocks"));
+ else {
+ try {
+ matches.closeInput();
+ }
+ catch(Exception ignored) {
+ }
+ }
}
- catch(ExecutionException error) {
- throw DMLRuntimeException.of(error.getCause());
+ finally {
+ _pendingCompletion.complete(null);
}
}
diff --git a/src/main/java/org/apache/sysds/runtime/ooc/primitives/MaterializeOOCPrimitive.java b/src/main/java/org/apache/sysds/runtime/ooc/primitives/MaterializeOOCPrimitive.java
index 92bc28f853..4c5842fb60 100644
--- a/src/main/java/org/apache/sysds/runtime/ooc/primitives/MaterializeOOCPrimitive.java
+++ b/src/main/java/org/apache/sysds/runtime/ooc/primitives/MaterializeOOCPrimitive.java
@@ -23,7 +23,6 @@ import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.ToIntFunction;
-import org.apache.sysds.runtime.DMLRuntimeException;
import org.apache.sysds.runtime.instructions.ooc.CachingStream;
import org.apache.sysds.runtime.instructions.ooc.OOCStream;
import org.apache.sysds.runtime.instructions.ooc.OOCStreamable;
@@ -130,11 +129,6 @@ public final class MaterializeOOCPrimitive extends OOCPrimitive {
}
}
- private void fail(Throwable error) {
- if(getContext() != null)
- getContext().failAll(DMLRuntimeException.of(error));
- }
-
private void finish() {
if(_finished.compareAndSet(false, true))
onComplete();
diff --git a/src/main/java/org/apache/sysds/runtime/ooc/primitives/OOCPrimitive.java b/src/main/java/org/apache/sysds/runtime/ooc/primitives/OOCPrimitive.java
index 12d2d3a6d1..45b7a3ff03 100644
--- a/src/main/java/org/apache/sysds/runtime/ooc/primitives/OOCPrimitive.java
+++ b/src/main/java/org/apache/sysds/runtime/ooc/primitives/OOCPrimitive.java
@@ -26,6 +26,7 @@ import java.util.Objects;
import java.util.Set;
import java.util.concurrent.atomic.AtomicBoolean;
+import org.apache.sysds.runtime.DMLRuntimeException;
import org.apache.sysds.runtime.instructions.ooc.OOCStream;
import org.apache.sysds.runtime.instructions.ooc.OOCStreamable;
import org.apache.sysds.runtime.instructions.spark.data.IndexedMatrixValue;
@@ -46,6 +47,7 @@ public abstract class OOCPrimitive {
private final List<InputSlot> _inputs;
private final AtomicBoolean _started;
private final AtomicBoolean _executionStarted;
+ private final AtomicBoolean _failed;
protected OOCAccessPattern _pattern;
protected MemoryAllowance _allowance;
@@ -71,6 +73,7 @@ public abstract class OOCPrimitive {
_inputs = new ArrayList<>();
_started = new AtomicBoolean();
_executionStarted = new AtomicBoolean();
+ _failed = new AtomicBoolean();
_pattern = OOCAccessPattern.UNSET;
}
@@ -165,6 +168,18 @@ public abstract class OOCPrimitive {
}
}
+ protected final boolean fail(Throwable error) {
+ if(!_failed.compareAndSet(false, true))
+ return false;
+ if(_context != null)
+ _context.failAll(DMLRuntimeException.of(error));
+ return true;
+ }
+
+ protected final boolean hasFailed() {
+ return _failed.get();
+ }
+
public final void onComplete() {
for(int i = 0; i < _inputs.size(); i++)
discardInputHandle(i);
diff --git a/src/main/java/org/apache/sysds/runtime/ooc/primitives/PlannableDataGenOOCPrimitive.java b/src/main/java/org/apache/sysds/runtime/ooc/primitives/PlannableDataGenOOCPrimitive.java
index b604e66061..55d2a31808 100644
--- a/src/main/java/org/apache/sysds/runtime/ooc/primitives/PlannableDataGenOOCPrimitive.java
+++ b/src/main/java/org/apache/sysds/runtime/ooc/primitives/PlannableDataGenOOCPrimitive.java
@@ -21,7 +21,6 @@ package org.apache.sysds.runtime.ooc.primitives;
import java.util.function.Function;
-import org.apache.sysds.runtime.DMLRuntimeException;
import org.apache.sysds.runtime.instructions.ooc.OOCStream;
import org.apache.sysds.runtime.instructions.ooc.OOCStreamable;
import org.apache.sysds.runtime.instructions.ooc.SubscribableTaskQueue;
@@ -77,7 +76,7 @@ public class PlannableDataGenOOCPrimitive extends OOCPrimitive {
budget.close();
}
}, getContext()).thenRun(output::closeInput).exceptionally(error -> {
- output.propagateFailure(DMLRuntimeException.of(error));
+ fail(error);
return null;
}).thenRun(this::onComplete);
diff --git a/src/main/java/org/apache/sysds/runtime/ooc/util/OOCInstructionUtils.java b/src/main/java/org/apache/sysds/runtime/ooc/util/OOCInstructionUtils.java
index f86c236060..e93876dfc0 100644
--- a/src/main/java/org/apache/sysds/runtime/ooc/util/OOCInstructionUtils.java
+++ b/src/main/java/org/apache/sysds/runtime/ooc/util/OOCInstructionUtils.java
@@ -119,6 +119,25 @@ public final class OOCInstructionUtils {
return submitOOCTasks(List.of(queue), (i, callback) -> consumer.accept(callback), null, null, context);
}
+ public static <T extends AutoCloseable> CompletableFuture<Void> submitCloseableOOCTasks(OOCStream<T> queue,
+ Consumer<T> consumer, StreamContext context) {
+ return submitOOCTasks(List.of(queue), (index, callback) -> {
+ try(T value = callback.get()) {
+ consumer.accept(value);
+ }
+ catch(Exception error) {
+ throw DMLRuntimeException.of(error);
+ }
+ }, null, (index, callback) -> {
+ try {
+ callback.get().close();
+ }
+ catch(Exception error) {
+ throw DMLRuntimeException.of(error);
+ }
+ }, context);
+ }
+
public static CompletableFuture<Void> submitAdmittedOOCTasks(OOCStream<IndexedMatrixValue> in,
OOCStream<IndexedMatrixValue> out, Function<IndexedMatrixValue, IndexedMatrixValue> operation,
MemoryAllowance allowance, StreamContext context) {
diff --git a/src/test/java/org/apache/sysds/test/component/ooc/OOCInstructionUtilsTest.java b/src/test/java/org/apache/sysds/test/component/ooc/OOCInstructionUtilsTest.java
index a889cc73e8..7737f45e9f 100644
--- a/src/test/java/org/apache/sysds/test/component/ooc/OOCInstructionUtilsTest.java
+++ b/src/test/java/org/apache/sysds/test/component/ooc/OOCInstructionUtilsTest.java
@@ -64,6 +64,38 @@ public class OOCInstructionUtilsTest {
Assert.assertEquals(1, released.get());
}
+ @Test
+ public void testSubmitCloseableOOCTasks() throws Exception {
+ SubscribableTaskQueue<OwnedTask> source = new SubscribableTaskQueue<>();
+ AtomicInteger processed = new AtomicInteger();
+ AtomicInteger closed = new AtomicInteger();
+ CompletableFuture<Void> completion = OOCInstructionUtils.submitCloseableOOCTasks(source,
+ (OwnedTask work) -> processed.addAndGet(work._value), new StreamContext().addOutStream());
+
+ source.enqueue(new OwnedTask(1, closed));
+ source.enqueue(new OwnedTask(2, closed));
+ source.closeInput();
+ completion.get(10, TimeUnit.SECONDS);
+
+ Assert.assertEquals(3, processed.get());
+ Assert.assertEquals(2, closed.get());
+ }
+
+ private static final class OwnedTask implements AutoCloseable {
+ private final int _value;
+ private final AtomicInteger _closed;
+
+ private OwnedTask(int value, AtomicInteger closed) {
+ _value = value;
+ _closed = closed;
+ }
+
+ @Override
+ public void close() {
+ _closed.incrementAndGet();
+ }
+ }
+
@Test
public void testSubmitTasksWaitsForAllStreams() throws Exception {
SubscribableTaskQueue<Integer> first = new SubscribableTaskQueue<>();