[SYSTEMDS-3891] OOC Non-blocking Join and Primitive Simplification

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 BroadcastOOCPrimitive(OOCStreamable<IndexedMatrixValue> streamed,
 		_liveness = liveness;
 		_operation = operation;
 		_cleaned = new AtomicBoolean();
-		_failed = new AtomicBoolean();
 		_sourceComplete = new AtomicBoolean();
 		_active = new AtomicInteger(1);
 	}
@@ -100,7 +98,7 @@ protected void startExecution() {
 		_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 @@ private void process(BroadcastWork work) {
 			fail(failure);
 		}
 		finally {
-			work.close();
 			if(budget != null)
 				budget.close();
 			completeOne();
@@ -255,14 +252,6 @@ private void completeOne() {
 		}
 	}
 
-	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;

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 GroupedReduceOOCPrimitive(OOCStreamable<IndexedMatrixValue> input, OOCStr
 		_output = output;
 		_merge = merge;
 		_cleaned = new AtomicBoolean();
-		_failed = new AtomicBoolean();
 		_sourceComplete = new AtomicBoolean();
 		_active = new AtomicInteger(1);
 		_finalizedGroups = new AtomicInteger();
@@ -100,7 +98,7 @@ protected void startExecution() {
 		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 @@ private void process(MergeWork work) {
 		catch(Throwable failure) {
 			if(merged != null)
 				merged.release();
-			work.close();
 			budget.close();
 			fail(failure);
 			completeOne();
@@ -276,7 +273,7 @@ private void completeOne() {
 		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 @@ private void completeOne() {
 		}
 	}
 
-	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;

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 @@ protected void startExecution() {
 		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 @@ protected void startExecution() {
 	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 @@ private void drive(OOCStream<IndexedMatrixValue> leftInput, OOCStream<IndexedMat
 						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 @@ private int accept(OOCStream.QueueCallback<IndexedMatrixValue> callback, boolean
 			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 @@ private int accept(OOCStream.QueueCallback<IndexedMatrixValue> callback, boolean
 		}
 	}
 
-	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);
 		}
 	}
 

src/main/java/org/apache/sysds/runtime/ooc/primitives/MaterializeOOCPrimitive.java

@@ -23,7 +23,6 @@
 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 @@ protected void startExecution() {
 		}
 	}
 
-	private void fail(Throwable error) {
-		if(getContext() != null)
-			getContext().failAll(DMLRuntimeException.of(error));
-	}
-
 	private void finish() {
 		if(_finished.compareAndSet(false, true))
 			onComplete();

src/main/java/org/apache/sysds/runtime/ooc/primitives/OOCPrimitive.java

@@ -26,6 +26,7 @@
 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 @@ protected OOCPrimitive(StreamContext context) {
 		_inputs = new ArrayList<>();
 		_started = new AtomicBoolean();
 		_executionStarted = new AtomicBoolean();
+		_failed = new AtomicBoolean();
 		_pattern = OOCAccessPattern.UNSET;
 	}
 
@@ -165,6 +168,18 @@ public final void tryStartExecution() {
 		}
 	}
 
+	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);

src/main/java/org/apache/sysds/runtime/ooc/primitives/PlannableDataGenOOCPrimitive.java

@@ -21,7 +21,6 @@
 
 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 @@ protected void startExecution() {
 					budget.close();
 			}
 		}, getContext()).thenRun(output::closeInput).exceptionally(error -> {
-			output.propagateFailure(DMLRuntimeException.of(error));
+			fail(error);
 			return null;
 		}).thenRun(this::onComplete);
 

src/main/java/org/apache/sysds/runtime/ooc/util/OOCInstructionUtils.java

@@ -119,6 +119,25 @@ public static <T> CompletableFuture<Void> submitOOCTasks(OOCStream<T> queue,
 		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) {