[OOC] Add NaryJoinOOCPrimitive and Generic Join

src/main/java/org/apache/sysds/runtime/instructions/ooc/TernaryOOCInstruction.java

@@ -26,13 +26,19 @@
 import org.apache.sysds.runtime.controlprogram.caching.MatrixObject;
 import org.apache.sysds.runtime.controlprogram.context.ExecutionContext;
 import org.apache.sysds.runtime.functionobjects.IfElse;
+import org.apache.sysds.runtime.functionobjects.Minus;
+import org.apache.sysds.runtime.functionobjects.MinusMultiply;
+import org.apache.sysds.runtime.functionobjects.Multiply;
+import org.apache.sysds.runtime.functionobjects.Plus;
+import org.apache.sysds.runtime.functionobjects.PlusMultiply;
 import org.apache.sysds.runtime.instructions.InstructionUtils;
 import org.apache.sysds.runtime.instructions.cp.CPOperand;
 import org.apache.sysds.runtime.instructions.cp.ScalarObject;
 import org.apache.sysds.runtime.instructions.cp.ScalarObjectFactory;
 import org.apache.sysds.runtime.instructions.cp.StringObject;
 import org.apache.sysds.runtime.instructions.spark.data.IndexedMatrixValue;
 import org.apache.sysds.runtime.matrix.data.MatrixBlock;
+import org.apache.sysds.runtime.matrix.operators.BinaryOperator;
 import org.apache.sysds.runtime.matrix.operators.Operator;
 import org.apache.sysds.runtime.matrix.operators.TernaryOperator;
 import org.apache.sysds.runtime.ooc.util.OOCInstructionUtils;
@@ -158,6 +164,32 @@ private void processThreeMatrixInstruction(ExecutionContext ec) {
 		OOCStream<IndexedMatrixValue> qOut = createWritableStream();
 		ec.getMatrixObject(output).setStreamHandle(qOut);
 
+		if(m1.getDataCharacteristics().dimsKnown() && m2.getDataCharacteristics().dimsKnown() &&
+			m3.getDataCharacteristics().dimsKnown()) {
+			TernaryOperator operator = (TernaryOperator) _optr;
+			if(operator.fn instanceof PlusMultiply || operator.fn instanceof MinusMultiply) {
+				OOCStream<IndexedMatrixValue> product = createWritableStream();
+				BinaryOperator multiply = new BinaryOperator(Multiply.getMultiplyFnObject());
+				BinaryOperator combine = operator.fn instanceof PlusMultiply ? new BinaryOperator(
+					Plus.getPlusFnObject()) : new BinaryOperator(Minus.getMinusFnObject());
+				OOCInstructionUtils.equiJoin(m2.getStreamable(), m3.getStreamable(), product,
+					(left, right) -> left.binaryOperations(multiply, right, new MatrixBlock()), getContext());
+				OOCInstructionUtils.equiJoin(m1.getStreamable(), product, qOut,
+					(left, right) -> left.binaryOperations(combine, right, new MatrixBlock()), getContext());
+				return;
+			}
+			if(operator.fn instanceof IfElse) {
+				OOCInstructionUtils.naryEquiJoin(List.of(m1.getStreamable(), m2.getStreamable(), m3.getStreamable()),
+					qOut,
+					blocks -> new IndexedMatrixValue(blocks.get(0).getIndexes(),
+						((MatrixBlock) blocks.get(0).getValue()).ternaryOperations(operator,
+							(MatrixBlock) blocks.get(1).getValue(), (MatrixBlock) blocks.get(2).getValue(),
+							new MatrixBlock())),
+					getContext());
+				return;
+			}
+		}
+
 		List<OOCStream<IndexedMatrixValue>> streams = List.of(
 			m1.getStreamHandle(), m2.getStreamHandle(), m3.getStreamHandle());
 

src/main/java/org/apache/sysds/runtime/instructions/spark/data/IndexedMatrixValue.java

@@ -97,6 +97,12 @@ public boolean tryWrite(DataOutput dataOutput) throws IOException {
 		return true;
 	}
 
+	@Override
+	public long size() {
+		MatrixBlock block = (MatrixBlock) _value;
+		return Math.max(block.getExactSerializedSize(), block.getInMemorySize());
+	}
+
 	@Override
 	public void discard() {
 		_value = null;

src/main/java/org/apache/sysds/runtime/ooc/cache/OOCFuture.java

@@ -19,10 +19,16 @@
 
 package org.apache.sysds.runtime.ooc.cache;
 
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.Objects;
 import java.util.concurrent.CompletionException;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.BiConsumer;
 import java.util.function.Consumer;
 import java.util.function.Function;
@@ -51,6 +57,52 @@ public static <T> OOCFuture<T> failed(Throwable error) {
 		return future;
 	}
 
+	public static <T> OOCFuture<List<T>> allOf(List<? extends OOCFuture<? extends T>> futures,
+		Consumer<? super T> failureCleanup) {
+		Objects.requireNonNull(futures);
+		Objects.requireNonNull(failureCleanup);
+		if(futures.isEmpty())
+			return completed(List.of());
+		OOCFuture<List<T>> result = new OOCFuture<>();
+		Object[] values = new Object[futures.size()];
+		AtomicInteger remaining = new AtomicInteger(futures.size());
+		AtomicReference<Throwable> firstError = new AtomicReference<>();
+		for(int i = 0; i < futures.size(); i++) {
+			int index = i;
+			futures.get(i).whenComplete((value, error) -> {
+				values[index] = value;
+				if(error != null)
+					firstError.compareAndSet(null, error);
+				if(remaining.decrementAndGet() != 0)
+					return;
+				Throwable failure = firstError.get();
+				List<T> completed = new ArrayList<>(values.length);
+				for(Object item : values) {
+					@SuppressWarnings("unchecked")
+					T typed = (T) item;
+					completed.add(typed);
+				}
+				if(failure == null) {
+					result.complete(Collections.unmodifiableList(completed));
+					return;
+				}
+				for(T item : completed) {
+					if(item == null)
+						continue;
+					try {
+						failureCleanup.accept(item);
+					}
+					catch(Throwable cleanupError) {
+						if(cleanupError != failure)
+							failure.addSuppressed(cleanupError);
+					}
+				}
+				result.completeExceptionally(failure);
+			});
+		}
+		return result;
+	}
+
 	public boolean complete(T value) {
 		return finish(value, null);
 	}

src/main/java/org/apache/sysds/runtime/ooc/cache/io/SpillableObject.java

@@ -26,6 +26,7 @@
 public interface SpillableObject {
 	boolean tryWrite(DataOutput out)  throws IOException;
 	void read(DataInput in)  throws IOException;
+	long size();
 
 	default void discard() {
 	}

src/main/java/org/apache/sysds/runtime/ooc/cache/packed/PackedBlock.java

@@ -57,6 +57,11 @@ public boolean tryWrite(DataOutput out) throws IOException {
 		return true;
 	}
 
+	@Override
+	public long size() {
+		return totalSize;
+	}
+
 	@Override
 	public void read(DataInput in) throws IOException {
 		int count = in.readInt();