feat: Add slice API for TableOperations (#6272)

- In support of #6059 - Fixes #6279, Fixes #6280
deephaven · Oct 25, 2024 · e33e9d7 · e33e9d7
1 parent 5da544b
commit e33e9d7
Show file tree

Hide file tree

Showing 13 changed files with 213 additions and 29 deletions.
diff --git a/engine/api/src/main/java/io/deephaven/engine/table/Table.java b/engine/api/src/main/java/io/deephaven/engine/table/Table.java
@@ -453,35 +453,6 @@ <DATA_TYPE> CloseableIterator<DATA_TYPE> objectColumnIterator(@NotNull String co
     // Slice Operations
     // -----------------------------------------------------------------------------------------------------------------
 
-    /**
-     * Extracts a subset of a table by row position.
-     * <p>
-     * If both firstPosition and lastPosition are positive, then the rows are counted from the beginning of the table.
-     * The firstPosition is inclusive, and the lastPosition is exclusive. The {@link #head}(N) call is equivalent to
-     * slice(0, N). The firstPosition must be less than or equal to the lastPosition.
-     * <p>
-     * If firstPosition is positive and lastPosition is negative, then the firstRow is counted from the beginning of the
-     * table, inclusively. The lastPosition is counted from the end of the table. For example, slice(1, -1) includes all
-     * rows but the first and last. If the lastPosition would be before the firstRow, the result is an emptyTable.
-     * <p>
-     * If firstPosition is negative, and lastPosition is zero, then the firstRow is counted from the end of the table,
-     * and the end of the slice is the size of the table. slice(-N, 0) is equivalent to {@link #tail}(N).
-     * <p>
-     * If the firstPosition is negative and the lastPosition is negative, they are both counted from the end of the
-     * table. For example, slice(-2, -1) returns the second to last row of the table.
-     * <p>
-     * If firstPosition is negative and lastPosition is positive, then firstPosition is counted from the end of the
-     * table, inclusively. The lastPosition is counted from the beginning of the table, exclusively. For example,
-     * slice(-3, 5) returns all rows starting from the third-last row to the fifth row of the table. If there are no
-     * rows between these positions, the function will return an empty table.
-     *
-     * @param firstPositionInclusive the first position to include in the result
-     * @param lastPositionExclusive the last position to include in the result
-     * @return a new Table, which is the request subset of rows from the original table
-     */
-    @ConcurrentMethod
-    Table slice(long firstPositionInclusive, long lastPositionExclusive);
-
     /**
      * Extracts a subset of a table by row percentages.
      * <p>

diff --git a/...e/src/main/java/io/deephaven/engine/table/impl/partitioned/PartitionedTableProxyImpl.java b/...e/src/main/java/io/deephaven/engine/table/impl/partitioned/PartitionedTableProxyImpl.java
@@ -430,6 +430,11 @@ public PartitionedTable.Proxy tail(long size) {
         return basicTransform(ct -> ct.tail(size));
     }
 
+    @Override
+    public PartitionedTable.Proxy slice(long firstPositionInclusive, long lastPositionExclusive) {
+        return basicTransform(ct -> ct.slice(firstPositionInclusive, lastPositionExclusive));
+    }
+
     @Override
     public PartitionedTable.Proxy reverse() {
         return basicTransform(Table::reverse);

diff --git a/java-client/session-dagger/src/test/java/io/deephaven/client/SliceTest.java b/java-client/session-dagger/src/test/java/io/deephaven/client/SliceTest.java
@@ -0,0 +1,83 @@
+//
+// Copyright (c) 2016-2024 Deephaven Data Labs and Patent Pending
+//
+package io.deephaven.client;
+
+import io.deephaven.client.impl.TableHandle;
+import io.deephaven.client.impl.TableHandle.TableHandleException;
+import io.deephaven.qst.table.TableSpec;
+import io.deephaven.qst.table.TimeTable;
+import org.junit.Test;
+
+import java.time.Duration;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.failBecauseExceptionWasNotThrown;
+
+public class SliceTest extends DeephavenSessionTestBase {
+
+    private static final TableSpec STATIC_BASE = TableSpec.empty(100).view("I=ii");
+    private static final TableSpec TICKING_BASE = TimeTable.of(Duration.ofMillis(100)).view("I=ii");
+
+    @Test
+    public void bothNonNegativeStartBeforeEnd() throws InterruptedException, TableHandleException {
+        allow(0, 50);
+        allow(25, 75);
+    }
+
+    @Test
+    public void bothNonNegativeStartAfterEnd() throws IllegalArgumentException {
+        disallow(50, 0);
+    }
+
+    @Test
+    public void bothNegativeStartBeforeEnd() throws InterruptedException, TableHandleException {
+        allow(-50, -25);
+    }
+
+    @Test
+    public void bothNegativeStartAfterEnd() throws IllegalArgumentException {
+        disallow(-25, -50);
+    }
+
+    @Test
+    public void diffSignStartBeforeEnd() throws InterruptedException, TableHandleException {
+        allow(-25, 25);
+    }
+
+    @Test
+    public void diffSignStartAfterEnd() throws InterruptedException, TableHandleException {
+        allow(25, -25);
+    }
+
+    @Test
+    public void startZeroEndNegative() throws InterruptedException, TableHandleException {
+        allow(0, -25);
+    }
+
+    private void allow(long start, long end)
+            throws InterruptedException, TableHandleException {
+        try (final TableHandle handle = session.batch().execute(STATIC_BASE.slice(start, end))) {
+            assertThat(handle.isSuccessful()).isTrue();
+        }
+        try (final TableHandle handle = session.batch().execute(TICKING_BASE.slice(start, end))) {
+            assertThat(handle.isSuccessful()).isTrue();
+        }
+    }
+
+    private void disallow(long start, long end) throws IllegalArgumentException {
+        try {
+            STATIC_BASE.slice(start, end);
+            failBecauseExceptionWasNotThrown(TableHandle.TableHandleException.class);
+        } catch (IllegalArgumentException e) {
+            // expected
+        }
+        try {
+            TICKING_BASE.slice(start, end);
+            failBecauseExceptionWasNotThrown(TableHandle.TableHandleException.class);
+        } catch (IllegalArgumentException e) {
+            // expected
+        }
+    }
+
+}
diff --git a/java-client/session/src/main/java/io/deephaven/client/impl/BatchTableRequestBuilder.java b/java-client/session/src/main/java/io/deephaven/client/impl/BatchTableRequestBuilder.java
@@ -61,6 +61,7 @@
 import io.deephaven.proto.backplane.grpc.Reference;
 import io.deephaven.proto.backplane.grpc.SelectDistinctRequest;
 import io.deephaven.proto.backplane.grpc.SelectOrUpdateRequest;
+import io.deephaven.proto.backplane.grpc.SliceRequest;
 import io.deephaven.proto.backplane.grpc.SnapshotTableRequest;
 import io.deephaven.proto.backplane.grpc.SnapshotWhenTableRequest;
 import io.deephaven.proto.backplane.grpc.SortDescriptor;
@@ -100,6 +101,7 @@
 import io.deephaven.qst.table.SelectDistinctTable;
 import io.deephaven.qst.table.SelectTable;
 import io.deephaven.qst.table.SingleParentTable;
+import io.deephaven.qst.table.SliceTable;
 import io.deephaven.qst.table.SnapshotTable;
 import io.deephaven.qst.table.SnapshotWhenTable;
 import io.deephaven.qst.table.SortTable;
@@ -236,6 +238,15 @@ public Operation visit(TailTable tailTable) {
                     .setSourceId(ref(tailTable.parent())).setNumRows(tailTable.size()));
         }
 
+        @Override
+        public Operation visit(SliceTable sliceTable) {
+            return op(Builder::setSlice, SliceRequest.newBuilder().setResultId(ticket)
+                    .setSourceId(ref(sliceTable.parent()))
+                    .setFirstPositionInclusive(sliceTable.firstPositionInclusive())
+                    .setLastPositionExclusive(sliceTable.lastPositionExclusive())
+                    .build());
+        }
+
         @Override
         public Operation visit(ReverseTable reverseTable) {
             // a bit hacky at the proto level, but this is how to specify a reverse

diff --git a/qst/src/main/java/io/deephaven/qst/TableAdapterImpl.java b/qst/src/main/java/io/deephaven/qst/TableAdapterImpl.java
@@ -28,6 +28,7 @@
 import io.deephaven.qst.table.SelectDistinctTable;
 import io.deephaven.qst.table.SelectTable;
 import io.deephaven.qst.table.SingleParentTable;
+import io.deephaven.qst.table.SliceTable;
 import io.deephaven.qst.table.SnapshotTable;
 import io.deephaven.qst.table.SnapshotWhenTable;
 import io.deephaven.qst.table.SortTable;
@@ -168,6 +169,13 @@ public Void visit(TailTable tailTable) {
         return null;
     }
 
+    @Override
+    public Void visit(SliceTable sliceTable) {
+        addOp(sliceTable,
+                parentOps(sliceTable).slice(sliceTable.firstPositionInclusive(), sliceTable.lastPositionExclusive()));
+        return null;
+    }
+
     @Override
     public Void visit(ReverseTable reverseTable) {
         addOp(reverseTable, parentOps(reverseTable).reverse());

diff --git a/qst/src/main/java/io/deephaven/qst/table/ParentsVisitor.java b/qst/src/main/java/io/deephaven/qst/table/ParentsVisitor.java
@@ -167,6 +167,11 @@ public Stream<TableSpec> visit(TailTable tailTable) {
         return single(tailTable);
     }
 
+    @Override
+    public Stream<TableSpec> visit(SliceTable sliceTable) {
+        return single(sliceTable);
+    }
+
     @Override
     public Stream<TableSpec> visit(ReverseTable reverseTable) {
         return single(reverseTable);

diff --git a/qst/src/main/java/io/deephaven/qst/table/SliceTable.java b/qst/src/main/java/io/deephaven/qst/table/SliceTable.java
@@ -0,0 +1,50 @@
+//
+// Copyright (c) 2016-2024 Deephaven Data Labs and Patent Pending
+//
+package io.deephaven.qst.table;
+
+import io.deephaven.annotations.NodeStyle;
+import org.immutables.value.Value.Check;
+import org.immutables.value.Value.Immutable;
+import org.immutables.value.Value.Parameter;
+
+@Immutable
+@NodeStyle
+public abstract class SliceTable extends TableBase implements SingleParentTable {
+
+    public static SliceTable of(TableSpec parent, long firstPositionInclusive, long lastPositionExclusive) {
+        return ImmutableSliceTable.of(parent, firstPositionInclusive, lastPositionExclusive);
+    }
+
+    @Parameter
+    public abstract TableSpec parent();
+
+    @Parameter
+    public abstract long firstPositionInclusive();
+
+    @Parameter
+    public abstract long lastPositionExclusive();
+
+    @Override
+    public final <T> T walk(Visitor<T> visitor) {
+        return visitor.visit(this);
+    }
+
+    @Check
+    final void checkPositions() {
+        if (firstPositionInclusive() >= 0 && lastPositionExclusive() >= 0
+                && lastPositionExclusive() < firstPositionInclusive()) {
+            throw new IllegalArgumentException(
+                    String.format(
+                            "Cannot slice with a non-negative start position (%d) that is after a non-negative end position (%d).",
+                            firstPositionInclusive(), lastPositionExclusive()));
+        }
+        if (firstPositionInclusive() < 0 && lastPositionExclusive() < 0
+                && lastPositionExclusive() < firstPositionInclusive()) {
+            throw new IllegalArgumentException(
+                    String.format(
+                            "Cannot slice with a negative start position (%d) that is after a negative end position (%d).",
+                            firstPositionInclusive(), lastPositionExclusive()));
+        }
+    }
+}
diff --git a/qst/src/main/java/io/deephaven/qst/table/TableBase.java b/qst/src/main/java/io/deephaven/qst/table/TableBase.java
@@ -33,6 +33,11 @@ public final TableSpec tail(long size) {
         return TailTable.of(this, size);
     }
 
+    @Override
+    public final TableSpec slice(long firstPositionInclusive, long lastPositionExclusive) {
+        return SliceTable.of(this, firstPositionInclusive, lastPositionExclusive);
+    }
+
     @Override
     public final TableSpec reverse() {
         return ReverseTable.of(this);

diff --git a/qst/src/main/java/io/deephaven/qst/table/TableLabelVisitor.java b/qst/src/main/java/io/deephaven/qst/table/TableLabelVisitor.java
@@ -64,6 +64,11 @@ public String visit(TailTable tailTable) {
         return "tail(" + tailTable.size() + ")";
     }
 
+    @Override
+    public String visit(SliceTable sliceTable) {
+        return String.format("slice(%d,%d)", sliceTable.firstPositionInclusive(), sliceTable.lastPositionExclusive());
+    }
+
     @Override
     public String visit(NaturalJoinTable naturalJoinTable) {
         return join("naturalJoin", naturalJoinTable);

diff --git a/qst/src/main/java/io/deephaven/qst/table/TableSpec.java b/qst/src/main/java/io/deephaven/qst/table/TableSpec.java
@@ -102,6 +102,8 @@ interface Visitor<T> {
 
         T visit(TailTable tailTable);
 
+        T visit(SliceTable sliceTable);
+
         T visit(ReverseTable reverseTable);
 
         T visit(SortTable sortTable);

diff --git a/qst/src/main/java/io/deephaven/qst/table/TableVisitorGeneric.java b/qst/src/main/java/io/deephaven/qst/table/TableVisitorGeneric.java
@@ -37,6 +37,11 @@ public T visit(TailTable tailTable) {
         return accept(tailTable);
     }
 
+    @Override
+    public T visit(SliceTable sliceTable) {
+        return accept(sliceTable);
+    }
+
     @Override
     public T visit(ReverseTable reverseTable) {
         return accept(reverseTable);

diff --git a/table-api/src/main/java/io/deephaven/api/TableOperations.java b/table-api/src/main/java/io/deephaven/api/TableOperations.java
@@ -160,6 +160,35 @@ public interface TableOperations<TOPS extends TableOperations<TOPS, TABLE>, TABL
      */
     TOPS whereNotIn(TABLE rightTable, Collection<? extends JoinMatch> columnsToMatch);
 
+    /**
+     * Extracts a subset of a table by row position.
+     * <p>
+     * If both firstPosition and lastPosition are positive, then the rows are counted from the beginning of the table.
+     * The firstPosition is inclusive, and the lastPosition is exclusive. The {@link #head}(N) call is equivalent to
+     * slice(0, N). The firstPosition must be less than or equal to the lastPosition.
+     * <p>
+     * If firstPosition is positive and lastPosition is negative, then the firstRow is counted from the beginning of the
+     * table, inclusively. The lastPosition is counted from the end of the table. For example, slice(1, -1) includes all
+     * rows but the first and last. If the lastPosition would be before the firstRow, the result is an emptyTable.
+     * <p>
+     * If firstPosition is negative, and lastPosition is zero, then the firstRow is counted from the end of the table,
+     * and the end of the slice is the size of the table. slice(-N, 0) is equivalent to {@link #tail}(N).
+     * <p>
+     * If the firstPosition is negative and the lastPosition is negative, they are both counted from the end of the
+     * table. For example, slice(-2, -1) returns the second to last row of the table.
+     * <p>
+     * If firstPosition is negative and lastPosition is positive, then firstPosition is counted from the end of the
+     * table, inclusively. The lastPosition is counted from the beginning of the table, exclusively. For example,
+     * slice(-3, 5) returns all rows starting from the third-last row to the fifth row of the table. If there are no
+     * rows between these positions, the function will return an empty table.
+     *
+     * @param firstPositionInclusive the first position to include in the result
+     * @param lastPositionExclusive the last position to include in the result
+     * @return a new Table, which is the request subset of rows from the original table
+     */
+    @ConcurrentMethod
+    TOPS slice(long firstPositionInclusive, long lastPositionExclusive);
+
     // -------------------------------------------------------------------------------------------
 
     @ConcurrentMethod

diff --git a/table-api/src/main/java/io/deephaven/api/TableOperationsAdapter.java b/table-api/src/main/java/io/deephaven/api/TableOperationsAdapter.java
@@ -41,6 +41,11 @@ public final TOPS_1 tail(long size) {
         return adapt(delegate.tail(size));
     }
 
+    @Override
+    public final TOPS_1 slice(long firstPositionInclusive, long lastPositionExclusive) {
+        return adapt(delegate.slice(firstPositionInclusive, lastPositionExclusive));
+    }
+
     @Override
     public final TOPS_1 reverse() {
         return adapt(delegate.reverse());