[SYSTEMDS-3891] Add StoreBackedStream

src/main/java/org/apache/sysds/runtime/ooc/store/CountingLiveness.java

@@ -0,0 +1,68 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.sysds.runtime.ooc.store;
+
+import java.util.concurrent.atomic.AtomicIntegerArray;
+
+public final class CountingLiveness implements MaterializedStore.Liveness {
+	private final AtomicIntegerArray _remaining;
+	private final AtomicIntegerArray _reservable;
+
+	public CountingLiveness(int size, int count) {
+		if(size < 0 || count < 0)
+			throw new IllegalArgumentException("Invalid args: size=" + size + ", count=" + count);
+		_remaining = new AtomicIntegerArray(size);
+		_reservable = new AtomicIntegerArray(size);
+		for(int i = 0; i < size; i++) {
+			_remaining.set(i, count);
+			_reservable.set(i, count);
+		}
+	}
+
+	@Override
+	public boolean needs(int index) {
+		return index >= 0 && index < _remaining.length() && _remaining.get(index) > 0;
+	}
+
+	@Override
+	public void consumed(int index) {
+		decrement(_remaining, index);
+	}
+
+	@Override
+	public boolean reserve(int index) {
+		return index >= 0 && index < _reservable.length() && decrement(_reservable, index);
+	}
+
+	@Override
+	public void unreserve(int index) {
+		_reservable.incrementAndGet(index);
+	}
+
+	private static boolean decrement(AtomicIntegerArray counters, int index) {
+		while(true) {
+			int current = counters.get(index);
+			if(current <= 0)
+				return false;
+			if(counters.compareAndSet(index, current, current - 1))
+				return true;
+		}
+	}
+}

src/main/java/org/apache/sysds/runtime/ooc/store/MaterializedCallback.java

@@ -21,21 +21,21 @@
 
 import org.apache.sysds.runtime.DMLRuntimeException;
 import org.apache.sysds.runtime.instructions.ooc.OOCStream;
-import org.apache.sysds.runtime.instructions.spark.data.IndexedMatrixValue;
 import org.apache.sysds.runtime.ooc.cache.BlockEntry;
+import org.apache.sysds.runtime.ooc.cache.io.SpillableObject;
 
 import java.util.concurrent.atomic.AtomicReference;
 
-public final class MaterializedCallback implements OOCStream.QueueCallback<IndexedMatrixValue> {
-	private final StoreLease<IndexedMatrixValue> _lease;
+public final class MaterializedCallback<T extends SpillableObject> implements OOCStream.QueueCallback<T> {
+	private final StoreLease<T> _lease;
 	private final AtomicReference<DMLRuntimeException> _failure;
 	private boolean _closed;
 
-	public MaterializedCallback(StoreLease<IndexedMatrixValue> lease) {
+	public MaterializedCallback(StoreLease<T> lease) {
 		this(lease, new AtomicReference<>());
 	}
 
-	private MaterializedCallback(StoreLease<IndexedMatrixValue> lease, AtomicReference<DMLRuntimeException> failure) {
+	private MaterializedCallback(StoreLease<T> lease, AtomicReference<DMLRuntimeException> failure) {
 		_lease = lease;
 		_failure = failure;
 	}
@@ -45,18 +45,18 @@ public BlockEntry pinnedEntry() {
 	}
 
 	@Override
-	public IndexedMatrixValue get() {
+	public T get() {
 		DMLRuntimeException failure = _failure.get();
 		if(failure != null)
 			throw failure;
 		return _lease.value();
 	}
 
 	@Override
-	public synchronized OOCStream.QueueCallback<IndexedMatrixValue> keepOpen() {
+	public synchronized OOCStream.QueueCallback<T> keepOpen() {
 		if(_closed)
 			throw new IllegalStateException("Cannot keep open a closed callback");
-		return new MaterializedCallback(_lease.retain(), _failure);
+		return new MaterializedCallback<>(_lease.retain(), _failure);
 	}
 
 	@Override

src/main/java/org/apache/sysds/runtime/ooc/store/SequentialAccessPattern.java

@@ -0,0 +1,64 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.sysds.runtime.ooc.store;
+
+import org.apache.sysds.runtime.ooc.cache.collections.ConcurrentBitSet;
+
+public final class SequentialAccessPattern implements MaterializedStore.AccessPattern {
+	private final int _size;
+	private final ConcurrentBitSet _consumed;
+	private int _next;
+	private volatile int _consumedThrough;
+
+	public SequentialAccessPattern(int size) {
+		if(size < 0)
+			throw new IllegalArgumentException("Size must not be negative: " + size);
+		_size = size;
+		_consumed = new ConcurrentBitSet(Math.max(1, size));
+		_next = 0;
+		_consumedThrough = -1;
+	}
+
+	@Override
+	public boolean hasNext() {
+		return _next < _size;
+	}
+
+	@Override
+	public int next() {
+		if(!hasNext())
+			throw new IllegalStateException("No remaining index");
+		return _next++;
+	}
+
+	@Override
+	public boolean needs(int index) {
+		return index >= 0 && index < _size && index > _consumedThrough && !_consumed.get(index);
+	}
+
+	@Override
+	public synchronized void consumed(int index) {
+		if(index < 0 || index >= _size)
+			throw new IndexOutOfBoundsException("Invalid consumed index: " + index);
+		_consumed.set(index);
+		while(_consumedThrough + 1 < _size && _consumed.get(_consumedThrough + 1))
+			_consumedThrough++;
+	}
+}

src/main/java/org/apache/sysds/runtime/ooc/store/StoreBackedStream.java

@@ -0,0 +1,185 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.sysds.runtime.ooc.store;
+
+import java.util.function.Consumer;
+
+import org.apache.sysds.runtime.DMLRuntimeException;
+import org.apache.sysds.runtime.controlprogram.caching.CacheableData;
+import org.apache.sysds.runtime.instructions.ooc.CachingStream;
+import org.apache.sysds.runtime.instructions.ooc.OOCStream;
+import org.apache.sysds.runtime.meta.DataCharacteristics;
+import org.apache.sysds.runtime.ooc.cache.io.SpillableObject;
+
+public final class StoreBackedStream<T extends SpillableObject> implements OOCStream<T> {
+	private final OrderedMaterializedStoreReader<T> _reader;
+	private volatile DMLRuntimeException _failure;
+	private boolean _subscriberSet;
+	private OOCStream.QueueCallback<T> _lastDequeue;
+	private boolean _exhausted;
+	private CacheableData<?> _data;
+
+	public StoreBackedStream(OrderedMaterializedStoreReader<T> reader) {
+		_reader = reader;
+	}
+
+	@Override
+	public void enqueue(T value) {
+		throw new DMLRuntimeException("Cannot enqueue to a store-backed stream");
+	}
+
+	@Override
+	public void enqueue(QueueCallback<T> callback) {
+		throw new DMLRuntimeException("Cannot enqueue to a store-backed stream");
+	}
+
+	@Override
+	public void closeInput() {
+		throw new DMLRuntimeException("Cannot close the input of a store-backed stream");
+	}
+
+	@Override
+	public synchronized T dequeue() {
+		QueueCallback<T> callback = dequeueInternal();
+		return callback == null ? null : callback.get();
+	}
+
+	@Override
+	public synchronized QueueCallback<T> dequeueCB() {
+		return dequeueInternal();
+	}
+
+	private QueueCallback<T> dequeueInternal() {
+		if(_subscriberSet)
+			throw new IllegalStateException("Cannot dequeue after setting a subscriber");
+		if(_lastDequeue != null) {
+			_lastDequeue.close();
+			_lastDequeue = null;
+		}
+		if(_failure != null)
+			throw _failure;
+		if(_exhausted)
+			return null;
+		try {
+			if(!_reader.hasNext()) {
+				_exhausted = true;
+				_reader.close();
+				return null;
+			}
+			_lastDequeue = new MaterializedCallback<>(_reader.next());
+			return _lastDequeue;
+		}
+		catch(InterruptedException e) {
+			Thread.currentThread().interrupt();
+			throw recordFailure(e);
+		}
+		catch(RuntimeException e) {
+			throw recordFailure(e);
+		}
+	}
+
+	@Override
+	public synchronized void setSubscriber(Consumer<QueueCallback<T>> subscriber) {
+		if(subscriber == null)
+			throw new IllegalArgumentException("Cannot set subscriber to null");
+		if(_subscriberSet)
+			throw new IllegalStateException("Subscriber cannot be set multiple times");
+		_subscriberSet = true;
+		Thread driver = new Thread(() -> drive(subscriber), "ooc-store-replay");
+		driver.setDaemon(true);
+		driver.start();
+	}
+
+	private void drive(Consumer<QueueCallback<T>> subscriber) {
+		DMLRuntimeException failure = _failure;
+		if(failure != null) {
+			subscriber.accept(OOCStream.eos(failure));
+			return;
+		}
+		try {
+			while(_reader.hasNext()) {
+				try(QueueCallback<T> callback = new MaterializedCallback<>(_reader.next())) {
+					subscriber.accept(callback);
+				}
+			}
+			_reader.close();
+			subscriber.accept(OOCStream.eos(null));
+		}
+		catch(InterruptedException e) {
+			Thread.currentThread().interrupt();
+			subscriber.accept(OOCStream.eos(recordFailure(e)));
+		}
+		catch(RuntimeException e) {
+			subscriber.accept(OOCStream.eos(recordFailure(e)));
+		}
+	}
+
+	private synchronized DMLRuntimeException recordFailure(Exception error) {
+		if(_failure == null)
+			_failure = DMLRuntimeException.of(error);
+		_reader.close();
+		return _failure;
+	}
+
+	@Override
+	public void propagateFailure(DMLRuntimeException failure) {
+		recordFailure(failure);
+	}
+
+	@Override
+	public OOCStream<T> getReadStream() {
+		return this;
+	}
+
+	@Override
+	public OOCStream<T> getWriteStream() {
+		throw new UnsupportedOperationException("A store-backed stream has no write stream");
+	}
+
+	@Override
+	public boolean hasStreamCache() {
+		return false;
+	}
+
+	@Override
+	public CachingStream getStreamCache() {
+		return null;
+	}
+
+	@Override
+	public boolean isProcessed() {
+		return false;
+	}
+
+	@Override
+	public DataCharacteristics getDataCharacteristics() {
+		return _data == null ? null : _data.getDataCharacteristics();
+	}
+
+	@Override
+	public CacheableData<?> getData() {
+		return _data;
+	}
+
+	@Override
+	public void setData(CacheableData<?> data) {
+		_data = data;
+	}
+}