RWC Changes

engine/table/src/main/java/io/deephaven/engine/table/impl/dataindex/CompoundDataIndexBoxedKeySource.java

@@ -0,0 +1,120 @@
+/**
+ * Copyright (c) 2016-2022 Deephaven Data Labs and Patent Pending
+ */
+package io.deephaven.engine.table.impl.dataindex;
+
+import io.deephaven.chunk.*;
+import io.deephaven.chunk.attributes.Values;
+import io.deephaven.engine.rowset.RowSequence;
+import io.deephaven.engine.table.ChunkSource;
+import io.deephaven.engine.table.ColumnSource;
+import io.deephaven.engine.table.SharedContext;
+import io.deephaven.engine.table.impl.DefaultChunkSource;
+import io.deephaven.engine.table.impl.chunkboxer.ChunkBoxer;
+import io.deephaven.engine.table.impl.sources.ReinterpretUtils;
+import io.deephaven.util.SafeCloseable;
+import org.jetbrains.annotations.NotNull;
+
+import java.util.Arrays;
+import java.util.stream.Stream;
+
+/**
+ * {@link ChunkSource} that produces data index {@link io.deephaven.engine.table.DataIndex.RowSetLookup row set lookup}
+ * keys from multiple {@link ColumnSource sources}. This can be used to extract keys from a data index table, or from a
+ * table of probe values.
+ */
+final class CompoundDataIndexBoxedKeySource implements DefaultChunkSource.WithPrev<Values> {
+
+    private final ColumnSource<?>[] keySources;
+    private final int keyWidth;
+
+    /**
+     * Construct a new CompoundDataIndexBoxedKeySource backed by the supplied {@link ColumnSource column sources}.
+     *
+     * @param keySources Sources corresponding to the key columns
+     */
+    CompoundDataIndexBoxedKeySource(@NotNull final ColumnSource<?>... keySources) {
+        this.keySources = Arrays.stream(keySources)
+                .map(ReinterpretUtils::maybeConvertToPrimitive).toArray(ColumnSource[]::new);
+        keyWidth = keySources.length;
+    }
+
+    @Override
+    public ChunkType getChunkType() {
+        return ChunkType.Object;
+    }
+
+    @Override
+    public void fillChunk(
+            @NotNull final ChunkSource.FillContext context,
+            @NotNull final WritableChunk<? super Values> destination,
+            @NotNull final RowSequence rowSequence) {
+        fillChunkInternal(context, destination, rowSequence, false);
+    }
+
+    public void fillPrevChunk(
+            @NotNull final ChunkSource.FillContext context,
+            @NotNull final WritableChunk<? super Values> destination,
+            @NotNull final RowSequence rowSequence) {
+        fillChunkInternal(context, destination, rowSequence, true);
+    }
+
+    private void fillChunkInternal(
+            @NotNull final ChunkSource.FillContext context,
+            @NotNull final WritableChunk<? super Values> destination,
+            @NotNull final RowSequence rowSequence,
+            final boolean usePrev) {
+        if (rowSequence.isEmpty()) {
+            destination.setSize(0);
+            return;
+        }
+        final FillContext fc = (FillContext) context;
+        for (int ci = 0; ci < keyWidth; ++ci) {
+            fc.boxedKeys[ci] = fc.keyBoxers[ci].box(usePrev
+                    ? keySources[ci].getPrevChunk(fc.keyContexts[ci], rowSequence)
+                    : keySources[ci].getChunk(fc.keyContexts[ci], rowSequence));
+        }
+        final int size = rowSequence.intSize();
+        destination.setSize(size);
+        final WritableObjectChunk<Object, ? super Values> typedDestination = destination.asWritableObjectChunk();
+        for (int ri = 0; ri < size; ++ri) {
+            final Object[] columnValues = new Object[keyWidth];
+            for (int ci = 0; ci < keyWidth; ++ci) {
+                columnValues[ci] = fc.boxedKeys[ci].get(ri);
+            }
+            typedDestination.set(ri, columnValues);
+        }
+    }
+
+    private static class FillContext implements ChunkSource.FillContext {
+
+        private final GetContext[] keyContexts;
+        private final ChunkBoxer.BoxerKernel[] keyBoxers;
+        private final ObjectChunk<?, ? extends Values>[] boxedKeys;
+
+        private FillContext(
+                final int chunkCapacity,
+                @NotNull final ColumnSource<?>[] keySources,
+                final SharedContext sharedContext) {
+            keyContexts = Stream.of(keySources)
+                    .map(cs -> cs.makeGetContext(chunkCapacity, sharedContext))
+                    .toArray(GetContext[]::new);
+            keyBoxers = Stream.of(keySources)
+                    .map(cs -> ChunkBoxer.getBoxer(cs.getChunkType(), chunkCapacity))
+                    .toArray(ChunkBoxer.BoxerKernel[]::new);
+            // noinspection unchecked
+            boxedKeys = new ObjectChunk[keySources.length];
+        }
+
+        @Override
+        public void close() {
+            SafeCloseable.closeAll(keyContexts);
+            SafeCloseable.closeAll(keyBoxers);
+        }
+    }
+
+    @Override
+    public ChunkSource.FillContext makeFillContext(final int chunkCapacity, final SharedContext sharedContext) {
+        return new FillContext(chunkCapacity, keySources, sharedContext);
+    }
+}

engine/table/src/main/java/io/deephaven/engine/table/impl/dataindex/DataIndexBoxedKeyHashAdapter.java

@@ -0,0 +1,40 @@
+package io.deephaven.engine.table.impl.dataindex;
+
+import io.deephaven.hash.KeyedObjectKey;
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+
+import java.util.Arrays;
+import java.util.Objects;
+
+/**
+ * {@link KeyedObjectKey} implementation for {@link io.deephaven.engine.table.DataIndex} keys when boxed for hash-based
+ * lookup.
+ */
+public class DataIndexBoxedKeyHashAdapter implements KeyedObjectKey<Object, Object> {
+
+    public static final KeyedObjectKey<Object, Object> INSTANCE = new DataIndexBoxedKeyHashAdapter();
+
+    private DataIndexBoxedKeyHashAdapter() {}
+
+    @Override
+    public Object getKey(@NotNull final Object value) {
+        return value;
+    }
+
+    @Override
+    public int hashKey(@Nullable final Object key) {
+        if (key instanceof Object[]) {
+            return Arrays.hashCode((Object[]) key);
+        }
+        return Objects.hashCode(key);
+    }
+
+    @Override
+    public boolean equalKey(@Nullable final Object key, @NotNull final Object value) {
+        if (key instanceof Object[] && value instanceof Object[]) {
+            return Arrays.equals((Object[]) key, (Object[]) value);
+        }
+        return Objects.equals(key, value);
+    }
+}

engine/table/src/main/java/io/deephaven/engine/table/impl/dataindex/DataIndexUtils.java

@@ -0,0 +1,29 @@
+package io.deephaven.engine.table.impl.dataindex;
+
+import io.deephaven.chunk.attributes.Values;
+import io.deephaven.engine.table.ChunkSource;
+import io.deephaven.engine.table.ColumnSource;
+
+/**
+ * Tools for working with {@link io.deephaven.engine.table.DataIndex data indices}.
+ */
+public class DataIndexUtils {
+
+    /**
+     * Make a {@link ChunkSource} that produces data index {@link io.deephaven.engine.table.DataIndex.RowSetLookup row
+     * set lookup} keys from {@code keySources}.
+     * 
+     * @param keySources The individual key sources
+     * @return The boxed key source
+     */
+    public static ChunkSource.WithPrev<Values> makeBoxedKeySource(final ColumnSource<?>... keySources) {
+        switch (keySources.length) {
+            case 0:
+                throw new IllegalArgumentException("Data index must have at least one key column");
+            case 1:
+                return new SingleDataIndexBoxedKeySource(keySources[0]);
+            default:
+                return new CompoundDataIndexBoxedKeySource(keySources);
+        }
+    }
+}

engine/table/src/main/java/io/deephaven/engine/table/impl/dataindex/SingleDataIndexBoxedKeySource.java

@@ -0,0 +1,98 @@
+/**
+ * Copyright (c) 2016-2022 Deephaven Data Labs and Patent Pending
+ */
+package io.deephaven.engine.table.impl.dataindex;
+
+import io.deephaven.chunk.ChunkType;
+import io.deephaven.chunk.ObjectChunk;
+import io.deephaven.chunk.WritableChunk;
+import io.deephaven.chunk.attributes.Values;
+import io.deephaven.engine.rowset.RowSequence;
+import io.deephaven.engine.table.ChunkSource;
+import io.deephaven.engine.table.ColumnSource;
+import io.deephaven.engine.table.SharedContext;
+import io.deephaven.engine.table.impl.DefaultChunkSource;
+import io.deephaven.engine.table.impl.chunkboxer.ChunkBoxer;
+import io.deephaven.engine.table.impl.sources.ReinterpretUtils;
+import io.deephaven.util.SafeCloseable;
+import org.jetbrains.annotations.NotNull;
+
+/**
+ * {@link ChunkSource} that produces data index {@link io.deephaven.engine.table.DataIndex.RowSetLookup row set lookup}
+ * keys from a single {@link ColumnSource source}. This can be used to extract keys from a data index table, or from a
+ * table of probe values.
+ */
+final class SingleDataIndexBoxedKeySource implements DefaultChunkSource.WithPrev<Values> {
+
+    private final ColumnSource<?> keySource;
+
+    /**
+     * Construct a new SingleDataIndexBoxedKeySource backed by the supplied {@link ColumnSource column source}.
+     *
+     * @param keySource Source corresponding to the key column
+     */
+    SingleDataIndexBoxedKeySource(@NotNull final ColumnSource<?> keySource) {
+        this.keySource = ReinterpretUtils.maybeConvertToPrimitive(keySource);
+    }
+
+    @Override
+    public ChunkType getChunkType() {
+        return ChunkType.Object;
+    }
+
+    @Override
+    public void fillChunk(
+            @NotNull final ChunkSource.FillContext context,
+            @NotNull final WritableChunk<? super Values> destination,
+            @NotNull final RowSequence rowSequence) {
+        fillChunkInternal(context, destination, rowSequence, false);
+    }
+
+    public void fillPrevChunk(
+            @NotNull final ChunkSource.FillContext context,
+            @NotNull final WritableChunk<? super Values> destination,
+            @NotNull final RowSequence rowSequence) {
+        fillChunkInternal(context, destination, rowSequence, true);
+    }
+
+    private void fillChunkInternal(
+            @NotNull final ChunkSource.FillContext context,
+            @NotNull final WritableChunk<? super Values> destination,
+            @NotNull final RowSequence rowSequence,
+            final boolean usePrev) {
+        if (rowSequence.isEmpty()) {
+            destination.setSize(0);
+            return;
+        }
+        final FillContext fc = (FillContext) context;
+        final ObjectChunk<?, ? extends Values> boxedKeys = fc.keyBoxer.box(usePrev
+                ? keySource.getPrevChunk(fc.keyContext, rowSequence)
+                : keySource.getChunk(fc.keyContext, rowSequence));
+        destination.setSize(boxedKeys.size());
+        boxedKeys.copyToChunk(0, destination, 0, boxedKeys.size());
+    }
+
+    private static class FillContext implements ChunkSource.FillContext {
+
+        private final GetContext keyContext;
+        private final ChunkBoxer.BoxerKernel keyBoxer;
+
+        private FillContext(
+                final int chunkCapacity,
+                @NotNull final ColumnSource<?> keySource,
+                final SharedContext sharedContext) {
+            keyContext = keySource.makeGetContext(chunkCapacity, sharedContext);
+            keyBoxer = ChunkBoxer.getBoxer(keySource.getChunkType(), chunkCapacity);
+        }
+
+        @Override
+        public void close() {
+            SafeCloseable.closeAll(keyContext, keyBoxer);
+        }
+    }
+
+    @Override
+    public ChunkSource.FillContext makeFillContext(final int chunkCapacity, final SharedContext sharedContext) {
+        return new FillContext(chunkCapacity, keySource, sharedContext);
+    }
+}

engine/table/src/main/java/io/deephaven/engine/table/impl/hierarchical/RollupNodeKeySource.java

@@ -181,7 +181,7 @@ private static class FillContext implements ChunkSource.FillContext {
         private FillContext(
                 final int chunkCapacity,
                 @NotNull final ColumnSource<Integer> depthSource,
-                @NotNull final ColumnSource[] groupByValueSources,
+                @NotNull final ColumnSource<?>[] groupByValueSources,
                 final SharedContext sharedContext) {
             depthContext = depthSource.makeGetContext(chunkCapacity, sharedContext);
             groupByValueContexts = Stream.of(groupByValueSources)
@@ -201,7 +201,7 @@ public void close() {
     }
 
     @Override
-    public FillContext makeFillContext(final int chunkCapacity, final SharedContext sharedContext) {
+    public ChunkSource.FillContext makeFillContext(final int chunkCapacity, final SharedContext sharedContext) {
         return new FillContext(chunkCapacity, depthSource, groupByValueSources, sharedContext);
     }
 }