Value streaming for NdArrays (#15)

Author: karllessard

ndarray/src/main/java/org/tensorflow/ndarray/DoubleNdArray.java

@@ -1,5 +1,5 @@
 /*
- Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+ Copyright 2019-2023 The TensorFlow Authors. All Rights Reserved.
 
  Licensed under the Apache License, Version 2.0 (the "License");
  you may not use this file except in compliance with the License.
@@ -20,6 +20,9 @@
 import org.tensorflow.ndarray.buffer.DoubleDataBuffer;
 import org.tensorflow.ndarray.index.Index;
 
+import java.util.stream.DoubleStream;
+import java.util.stream.StreamSupport;
+
 /**
  * An {@link NdArray} of doubles.
  */
@@ -68,6 +71,18 @@ public interface DoubleNdArray extends NdArray<Double> {
    */
   DoubleNdArray setDouble(double value, long... coordinates);
 
+  /**
+   * Retrieve all scalar values of this array as a stream of doubles.
+   *
+   * <p>For {@code rank() > 1} arrays, all vectors of the last dimension are collated so that the scalar values are
+   * returned in sequential order.</p>
+   *
+   * @return scalar values as a stream
+   */
+  default DoubleStream streamOfDoubles() {
+    return StreamSupport.stream(scalars().spliterator(), false).mapToDouble(DoubleNdArray::getDouble);
+  }
+
   @Override
   DoubleNdArray slice(Index... indices);
 

ndarray/src/main/java/org/tensorflow/ndarray/IntNdArray.java

@@ -1,5 +1,5 @@
 /*
- Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+ Copyright 2019-2023 The TensorFlow Authors. All Rights Reserved.
 
  Licensed under the Apache License, Version 2.0 (the "License");
  you may not use this file except in compliance with the License.
@@ -20,6 +20,9 @@
 import org.tensorflow.ndarray.buffer.IntDataBuffer;
 import org.tensorflow.ndarray.index.Index;
 
+import java.util.stream.IntStream;
+import java.util.stream.StreamSupport;
+
 /**
  * An {@link NdArray} of integers.
  */
@@ -68,6 +71,18 @@ public interface IntNdArray extends NdArray<Integer> {
    */
   IntNdArray setInt(int value, long... coordinates);
 
+  /**
+   * Retrieve all scalar values of this array as a stream of integers.
+   *
+   * <p>For {@code rank() > 1} arrays, all vectors of the last dimension are collated so that the scalar values are
+   * returned in sequential order.</p>
+   *
+   * @return scalar values as a stream
+   */
+  default IntStream streamOfInts() {
+    return StreamSupport.stream(scalars().spliterator(), false).mapToInt(IntNdArray::getInt);
+  }
+
   @Override
   IntNdArray slice(Index... indices);
 

ndarray/src/main/java/org/tensorflow/ndarray/LongNdArray.java

@@ -1,5 +1,5 @@
 /*
- Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+ Copyright 2019-2023 The TensorFlow Authors. All Rights Reserved.
 
  Licensed under the Apache License, Version 2.0 (the "License");
  you may not use this file except in compliance with the License.
@@ -20,6 +20,9 @@
 import org.tensorflow.ndarray.buffer.LongDataBuffer;
 import org.tensorflow.ndarray.index.Index;
 
+import java.util.stream.LongStream;
+import java.util.stream.StreamSupport;
+
 /**
  * An {@link NdArray} of longs.
  */
@@ -68,6 +71,18 @@ public interface LongNdArray extends NdArray<Long> {
    */
   LongNdArray setLong(long value, long... coordinates);
 
+  /**
+   * Retrieve all scalar values of this array as a stream of longs.
+   *
+   * <p>For {@code rank() > 1} arrays, all vectors of the last dimension are collated so that the scalar values are
+   * returned in sequential order.</p>
+   *
+   * @return scalar values as a stream
+   */
+  default LongStream streamOfLongs() {
+    return StreamSupport.stream(scalars().spliterator(), false).mapToLong(LongNdArray::getLong);
+  }
+
   @Override
   LongNdArray slice(Index... indices);
 

ndarray/src/main/java/org/tensorflow/ndarray/NdArray.java

@@ -1,5 +1,5 @@
 /*
- Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+ Copyright 2019-2023 The TensorFlow Authors. All Rights Reserved.
 
  Licensed under the Apache License, Version 2.0 (the "License");
  you may not use this file except in compliance with the License.
@@ -18,6 +18,9 @@
 
 import java.util.function.BiConsumer;
 import java.util.function.Consumer;
+import java.util.stream.Stream;
+import java.util.stream.StreamSupport;
+
 import org.tensorflow.ndarray.buffer.DataBuffer;
 import org.tensorflow.ndarray.index.Index;
 
@@ -229,6 +232,18 @@ public interface NdArray<T> extends Shaped {
    */
   NdArray<T> setObject(T value, long... coordinates);
 
+  /**
+   * Retrieve all scalar values of this array as a stream of objects.
+   *
+   * <p>For {@code rank() > 1} arrays, all vectors of the last dimension are collated so that the scalar values are
+   * returned in sequential order.</p>
+   *
+   * @return scalar values as a stream
+   */
+  default Stream<T> streamOfObjects() {
+    return StreamSupport.stream(scalars().spliterator(), false).map(NdArray::getObject);
+  }
+
   /**
    * Copy the content of this array to the destination array.
    *

ndarray/src/test/java/org/tensorflow/ndarray/IntNdArrayTestBase.java

@@ -1,5 +1,5 @@
 /*
- Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+ Copyright 2019-2023 The TensorFlow Authors. All Rights Reserved.
 
  Licensed under the Apache License, Version 2.0 (the "License");
  you may not use this file except in compliance with the License.
@@ -16,10 +16,11 @@
  */
 package org.tensorflow.ndarray;
 
-import static org.junit.jupiter.api.Assertions.assertEquals;
-
 import org.junit.jupiter.api.Test;
 
+import static org.junit.jupiter.api.Assertions.assertArrayEquals;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
 public abstract class IntNdArrayTestBase extends NdArrayTestBase<Integer> {
 
     @Override
@@ -52,4 +53,29 @@ public void iteratePrimitiveElements() {
         assertEquals(9, matrix3d.getInt(0, 0, 4));
         assertEquals(7, matrix3d.getInt(0, 1, 2));
     }
+
+    @Test
+    public void streamingInts() {
+        IntNdArray scalar = allocate(Shape.scalar());
+        scalar.setInt(1);
+        var values = scalar.streamOfInts().toArray();
+        assertArrayEquals(new int[]{1}, values);
+
+        IntNdArray vector = allocate(Shape.of(5));
+        vector.setInt(1, 0);
+        vector.setInt(2, 1);
+        vector.setInt(3, 2);
+        vector.setInt(4, 3);
+        vector.setInt(5, 4);
+        values = vector.streamOfInts().toArray();
+        assertArrayEquals(new int[]{1, 2, 3, 4, 5}, values);
+
+        IntNdArray matrix = allocate(Shape.of(2, 2));
+        matrix.setInt(1, 0, 0);
+        matrix.setInt(2, 0, 1);
+        matrix.setInt(3, 1, 0);
+        matrix.setInt(4, 1, 1);
+        values = matrix.streamOfInts().toArray();
+        assertArrayEquals(new int[]{1, 2, 3, 4}, values);
+    }
 }

ndarray/src/test/java/org/tensorflow/ndarray/LongNdArrayTestBase.java

@@ -1,5 +1,5 @@
 /*
- Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+ Copyright 2019-2023 The TensorFlow Authors. All Rights Reserved.
 
  Licensed under the Apache License, Version 2.0 (the "License");
  you may not use this file except in compliance with the License.
@@ -16,6 +16,7 @@
  */
 package org.tensorflow.ndarray;
 
+import static org.junit.jupiter.api.Assertions.assertArrayEquals;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 
 import org.junit.jupiter.api.Test;
@@ -52,4 +53,29 @@ public void iteratePrimitiveElements() {
         assertEquals(9, matrix3d.getLong(0, 0, 4));
         assertEquals(7, matrix3d.getLong(0, 1, 2));
     }
+
+    @Test
+    public void streamingLongs() {
+        LongNdArray scalar = allocate(Shape.scalar());
+        scalar.setLong(1L);
+        var values = scalar.streamOfLongs().toArray();
+        assertArrayEquals(new long[]{1L}, values);
+
+        LongNdArray vector = allocate(Shape.of(5));
+        vector.setLong(1L, 0);
+        vector.setLong(2L, 1);
+        vector.setLong(3L, 2);
+        vector.setLong(4L, 3);
+        vector.setLong(5L, 4);
+        values = vector.streamOfLongs().toArray();
+        assertArrayEquals(new long[]{1L, 2L, 3L, 4L, 5L}, values);
+
+        LongNdArray matrix = allocate(Shape.of(2, 2));
+        matrix.setLong(1L, 0, 0);
+        matrix.setLong(2L, 0, 1);
+        matrix.setLong(3L, 1, 0);
+        matrix.setLong(4L, 1, 1);
+        values = matrix.streamOfLongs().toArray();
+        assertArrayEquals(new long[]{1L, 2L, 3L, 4L}, values);
+    }
 }

ndarray/src/test/java/org/tensorflow/ndarray/NdArrayTestBase.java

@@ -1,5 +1,5 @@
 /*
- Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+ Copyright 2019-2023 The TensorFlow Authors. All Rights Reserved.
 
  Licensed under the Apache License, Version 2.0 (the "License");
  you may not use this file except in compliance with the License.
@@ -16,9 +16,7 @@
  */
 package org.tensorflow.ndarray;
 
-import static org.junit.jupiter.api.Assertions.assertEquals;
-import static org.junit.jupiter.api.Assertions.assertNotEquals;
-import static org.junit.jupiter.api.Assertions.fail;
+import static org.junit.jupiter.api.Assertions.*;
 import static org.tensorflow.ndarray.NdArrays.vectorOfObjects;
 import static org.tensorflow.ndarray.index.Indices.all;
 import static org.tensorflow.ndarray.index.Indices.at;
@@ -32,6 +30,9 @@
 
 import java.nio.BufferOverflowException;
 import java.nio.BufferUnderflowException;
+import java.util.List;
+import java.util.stream.Collectors;
+
 import org.junit.jupiter.api.Test;
 import org.tensorflow.ndarray.buffer.DataBuffer;
 import org.tensorflow.ndarray.index.Indices;
@@ -358,4 +359,29 @@ public void iterateScalarsOnSegmentedElements() {
       });
     });
   }
+
+  @Test
+  public void streamingObjects() {
+    NdArray<T> scalar = allocate(Shape.scalar());
+    scalar.setObject(valueOf(1L));
+    var values = scalar.streamOfObjects().collect(Collectors.toList());
+    assertIterableEquals(List.of(valueOf(1L)), values);
+
+    NdArray<T> vector = allocate(Shape.of(5));
+    vector.setObject(valueOf(1L), 0);
+    vector.setObject(valueOf(2L), 1);
+    vector.setObject(valueOf(3L), 2);
+    vector.setObject(valueOf(4L), 3);
+    vector.setObject(valueOf(5L), 4);
+    values = vector.streamOfObjects().collect(Collectors.toList());
+    assertIterableEquals(List.of(valueOf(1L), valueOf(2L), valueOf(3L), valueOf(4L), valueOf(5L)), values);
+
+    NdArray<T> matrix = allocate(Shape.of(2, 2));
+    matrix.setObject(valueOf(1L), 0, 0);
+    matrix.setObject(valueOf(2L), 0, 1);
+    matrix.setObject(valueOf(3L), 1, 0);
+    matrix.setObject(valueOf(4L), 1, 1);
+    values = matrix.streamOfObjects().collect(Collectors.toList());
+    assertIterableEquals(List.of(valueOf(1L), valueOf(2L), valueOf(3L), valueOf(4L)), values);
+  }
 }