feat: support refreshing Iceberg tables

Util/src/main/java/io/deephaven/util/datastructures/SubscriptionSet.java

@@ -7,6 +7,7 @@
 import io.deephaven.base.reference.WeakReferenceWrapper;
 import io.deephaven.base.verify.Assert;
 import io.deephaven.base.verify.Require;
+import org.apache.commons.lang3.function.TriConsumer;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
@@ -219,6 +220,36 @@ public final <NOTIFICATION_TYPE> boolean deliverNotification(
         return initialSize > 0 && size == 0;
     }
 
+    /**
+     * Dispatch a notification to all subscribers. Clean up any GC'd subscriptions.
+     *
+     * @param procedure The notification procedure to invoke
+     * @param firstNotification The first item to deliver
+     * @param secondNotification The second item to deliver (must be of the same type as {@code firstNotification})
+     * @param activeOnly Whether to restrict this notification to active subscriptions only
+     * @return Whether this operation caused the set to become <b>empty</b>
+     */
+    public final <NOTIFICATION_TYPE> boolean deliverNotification(
+            @NotNull final TriConsumer<LISTENER_TYPE, NOTIFICATION_TYPE, NOTIFICATION_TYPE> procedure,
+            @Nullable final NOTIFICATION_TYPE firstNotification,
+            @Nullable final NOTIFICATION_TYPE secondNotification,
+            final boolean activeOnly) {
+        final int initialSize = size;
+        for (int si = 0; si < size;) {
+            final Entry currentEntry = subscriptions[si];
+            final LISTENER_TYPE currentListener = currentEntry.getListener();
+            if (currentListener == null) {
+                removeAt(si);
+                continue; // si is not incremented in this case - we'll reconsider the same slot if necessary.
+            }
+            if (!activeOnly || currentEntry.isActive()) {
+                procedure.accept(currentListener, firstNotification, secondNotification);
+            }
+            ++si;
+        }
+        return initialSize > 0 && size == 0;
+    }
+
     private void removeAt(final int subscriptionIndex) {
         final int lastSubscriptionIndex = --size;
         subscriptions[subscriptionIndex] = subscriptions[lastSubscriptionIndex];

engine/table/src/main/java/io/deephaven/engine/table/impl/ColumnSourceManager.java

@@ -3,14 +3,10 @@
 //
 package io.deephaven.engine.table.impl;
 
-import io.deephaven.engine.liveness.LivenessReferent;
+import io.deephaven.engine.liveness.LivenessNode;
 import io.deephaven.engine.rowset.RowSet;
 import io.deephaven.engine.rowset.TrackingWritableRowSet;
-import io.deephaven.engine.rowset.WritableRowSet;
-import io.deephaven.engine.table.ColumnSource;
-import io.deephaven.engine.table.DataIndex;
-import io.deephaven.engine.table.Table;
-import io.deephaven.engine.table.TableListener;
+import io.deephaven.engine.table.*;
 import io.deephaven.engine.table.impl.locations.ImmutableTableLocationKey;
 import io.deephaven.engine.table.impl.locations.TableLocation;
 import org.jetbrains.annotations.NotNull;
@@ -22,7 +18,7 @@
 /**
  * Manager for ColumnSources in a Table.
  */
-public interface ColumnSourceManager extends LivenessReferent {
+public interface ColumnSourceManager extends LivenessNode {
 
     /**
      * Get a map of name to {@link ColumnSource} for the column sources maintained by this manager.
@@ -53,7 +49,7 @@ public interface ColumnSourceManager extends LivenessReferent {
      *
      * @return The set of added row keys, to be owned by the caller
      */
-    WritableRowSet refresh();
+    TableUpdate refresh();
 
     /**
      * Advise this ColumnSourceManager that an error has occurred, and that it will no longer be {@link #refresh()
@@ -116,8 +112,18 @@ public interface ColumnSourceManager extends LivenessReferent {
     /**
      * Remove a table location key from the sources.
      *
-     * @return true if the location key was actually removed
      * @param tableLocationKey the location key being removed
      */
-    boolean removeLocationKey(@NotNull ImmutableTableLocationKey tableLocationKey);
+    void removeLocationKey(@NotNull ImmutableTableLocationKey tableLocationKey);
+
+    /**
+     * Get a map of Table attributes that can be applied to the output source table, given the update modes of the
+     * underlying table location provider.
+     *
+     * @param tableUpdateMode The update mode of the table location set
+     * @param tableLocationUpdateMode The update mode of the table location rows
+     */
+    Map<String, Object> getTableAttributes(
+            @NotNull TableUpdateMode tableUpdateMode,
+            @NotNull TableUpdateMode tableLocationUpdateMode);
 }

engine/table/src/main/java/io/deephaven/engine/table/impl/PartitionAwareSourceTable.java

@@ -6,6 +6,7 @@
 import io.deephaven.api.Selectable;
 import io.deephaven.api.filter.Filter;
 import io.deephaven.base.verify.Assert;
+import io.deephaven.engine.liveness.LiveSupplier;
 import io.deephaven.engine.table.*;
 import io.deephaven.engine.table.impl.select.analyzers.SelectAndViewAnalyzer;
 import io.deephaven.engine.updategraph.UpdateSourceRegistrar;
@@ -204,7 +205,7 @@ protected final Table redefine(TableDefinition newDefinitionExternal, TableDefin
                 reference, null, viewColumns, null);
     }
 
-    private static final String LOCATION_KEY_COLUMN_NAME = "__PartitionAwareSourceTable_TableLocationKey__";
+    private static final String KEY_SUPPLIER_COLUMN_NAME = "__PartitionAwareSourceTable_KeySupplier__";
 
     private static <T> ColumnSource<? super T> makePartitionSource(@NotNull final ColumnDefinition<T> columnDefinition,
             @NotNull final Collection<ImmutableTableLocationKey> locationKeys) {
@@ -221,30 +222,42 @@ private static <T> ColumnSource<? super T> makePartitionSource(@NotNull final Co
     }
 
     @Override
-    protected final Collection<ImmutableTableLocationKey> filterLocationKeys(
-            @NotNull final Collection<ImmutableTableLocationKey> foundLocationKeys) {
+    protected final Collection<LiveSupplier<ImmutableTableLocationKey>> filterLocationKeys(
+            @NotNull final Collection<LiveSupplier<ImmutableTableLocationKey>> foundLocationKeys) {
         if (partitioningColumnFilters.length == 0) {
             return foundLocationKeys;
         }
+
+        final Collection<ImmutableTableLocationKey> immutableTableLocationKeys = foundLocationKeys.stream()
+                .map(LiveSupplier::get)
+                .collect(Collectors.toList());
+
         // TODO (https://github.com/deephaven/deephaven-core/issues/867): Refactor around a ticking partition table
         final List<String> partitionTableColumnNames = Stream.concat(
                 partitioningColumnDefinitions.keySet().stream(),
-                Stream.of(LOCATION_KEY_COLUMN_NAME)).collect(Collectors.toList());
+                Stream.of(KEY_SUPPLIER_COLUMN_NAME)).collect(Collectors.toList());
         final List<ColumnSource<?>> partitionTableColumnSources =
                 new ArrayList<>(partitioningColumnDefinitions.size() + 1);
         for (final ColumnDefinition<?> columnDefinition : partitioningColumnDefinitions.values()) {
-            partitionTableColumnSources.add(makePartitionSource(columnDefinition, foundLocationKeys));
+            partitionTableColumnSources.add(makePartitionSource(columnDefinition, immutableTableLocationKeys));
         }
-        partitionTableColumnSources.add(ArrayBackedColumnSource.getMemoryColumnSource(foundLocationKeys,
-                ImmutableTableLocationKey.class, null));
+        // Add the key suppliers to the table
+        // noinspection unchecked,rawtypes
+        partitionTableColumnSources.add(ArrayBackedColumnSource.getMemoryColumnSource(
+                (Collection<LiveSupplier>) (Collection) foundLocationKeys,
+                LiveSupplier.class,
+                null));
+
         final Table filteredColumnPartitionTable = TableTools
                 .newTable(foundLocationKeys.size(), partitionTableColumnNames, partitionTableColumnSources)
                 .where(Filter.and(partitioningColumnFilters));
         if (filteredColumnPartitionTable.size() == foundLocationKeys.size()) {
             return foundLocationKeys;
         }
-        final Iterable<ImmutableTableLocationKey> iterable =
-                () -> filteredColumnPartitionTable.columnIterator(LOCATION_KEY_COLUMN_NAME);
+
+        // Return the filtered keys
+        final Iterable<LiveSupplier<ImmutableTableLocationKey>> iterable =
+                () -> filteredColumnPartitionTable.columnIterator(KEY_SUPPLIER_COLUMN_NAME);
         return StreamSupport.stream(iterable.spliterator(), false).collect(Collectors.toList());
     }
 

engine/table/src/main/java/io/deephaven/engine/table/impl/SourcePartitionedTable.java

@@ -4,6 +4,7 @@
 package io.deephaven.engine.table.impl;
 
 import io.deephaven.base.verify.Assert;
+import io.deephaven.engine.liveness.*;
 import io.deephaven.engine.primitive.iterator.CloseableIterator;
 import io.deephaven.engine.rowset.*;
 import io.deephaven.engine.table.*;
@@ -15,7 +16,9 @@
 import io.deephaven.engine.table.impl.sources.ArrayBackedColumnSource;
 import io.deephaven.engine.table.impl.sources.regioned.RegionedTableComponentFactoryImpl;
 import io.deephaven.engine.table.iterators.ChunkedObjectColumnIterator;
+import io.deephaven.engine.updategraph.NotificationQueue;
 import io.deephaven.engine.updategraph.UpdateCommitter;
+import io.deephaven.engine.updategraph.UpdateGraph;
 import io.deephaven.engine.updategraph.UpdateSourceCombiner;
 import io.deephaven.util.datastructures.linked.IntrusiveDoublyLinkedNode;
 import io.deephaven.util.datastructures.linked.IntrusiveDoublyLinkedQueue;
@@ -73,7 +76,8 @@ public SourcePartitionedTable(
                 false);
     }
 
-    private static final class UnderlyingTableMaintainer {
+    private static final class UnderlyingTableMaintainer extends ReferenceCountedLivenessNode
+            implements NotificationQueue.Dependency {
 
         private final TableDefinition constituentDefinition;
         private final UnaryOperator<Table> applyTablePermissions;
@@ -103,13 +107,14 @@ private UnderlyingTableMaintainer(
                 final boolean refreshLocations,
                 final boolean refreshSizes,
                 @NotNull final Predicate<ImmutableTableLocationKey> locationKeyMatcher) {
+            super(false);
+
             this.constituentDefinition = constituentDefinition;
             this.applyTablePermissions = applyTablePermissions;
             this.tableLocationProvider = tableLocationProvider;
             this.refreshSizes = refreshSizes;
             this.locationKeyMatcher = locationKeyMatcher;
 
-            // noinspection resource
             resultRows = RowSetFactory.empty().toTracking();
             resultTableLocationKeys = ArrayBackedColumnSource.getMemoryColumnSource(TableLocationKey.class, null);
             resultLocationTables = ArrayBackedColumnSource.getMemoryColumnSource(Table.class, null);
@@ -123,13 +128,19 @@ private UnderlyingTableMaintainer(
             if (needToRefreshLocations || refreshSizes) {
                 result.setRefreshing(true);
                 refreshCombiner = new UpdateSourceCombiner(result.getUpdateGraph());
-                result.addParentReference(refreshCombiner);
+                result.addParentReference(this);
+                manage(refreshCombiner);
             } else {
                 refreshCombiner = null;
             }
 
             if (needToRefreshLocations) {
+                resultTableLocationKeys.startTrackingPrevValues();
+                resultLocationTables.startTrackingPrevValues();
+
                 subscriptionBuffer = new TableLocationSubscriptionBuffer(tableLocationProvider);
+                manage(subscriptionBuffer);
+
                 pendingLocationStates = new IntrusiveDoublyLinkedQueue<>(
                         IntrusiveDoublyLinkedNode.Adapter.<PendingLocationState>getInstance());
                 readyLocationStates = new IntrusiveDoublyLinkedQueue<>(
@@ -143,7 +154,6 @@ protected void instrumentedRefresh() {
                         processPendingLocations(true);
                     }
                 };
-                result.addParentReference(processNewLocationsUpdateRoot);
                 refreshCombiner.addSource(processNewLocationsUpdateRoot);
 
                 this.removedLocationsComitter = new UpdateCommitter<>(
@@ -154,7 +164,6 @@ protected void instrumentedRefresh() {
                             removedConstituents.forEach(result::unmanage);
                             removedConstituents = null;
                         });
-
                 processPendingLocations(false);
             } else {
                 subscriptionBuffer = null;
@@ -163,9 +172,12 @@ protected void instrumentedRefresh() {
                 processNewLocationsUpdateRoot = null;
                 removedLocationsComitter = null;
                 tableLocationProvider.refresh();
-                try (final RowSet added = sortAndAddLocations(tableLocationProvider.getTableLocationKeys().stream()
-                        .filter(locationKeyMatcher)
-                        .map(tableLocationProvider::getTableLocation))) {
+
+                final Collection<TableLocation> locations = new ArrayList<>();
+                tableLocationProvider.getTableLocationKeys(
+                        tlk -> locations.add(tableLocationProvider.getTableLocation(tlk.get())),
+                        locationKeyMatcher);
+                try (final RowSet added = sortAndAddLocations(locations.stream())) {
                     resultRows.insert(added);
                 }
             }
@@ -204,18 +216,30 @@ private Table makeConstituentTable(@NotNull final TableLocation tableLocation) {
                     constituentDefinition,
                     "SingleLocationSourceTable-" + tableLocation,
                     RegionedTableComponentFactoryImpl.INSTANCE,
-                    new SingleTableLocationProvider(tableLocation),
+                    new SingleTableLocationProvider(tableLocation, refreshSizes
+                            ? tableLocationProvider.getLocationUpdateMode()
+                            : TableUpdateMode.STATIC),
                     refreshSizes ? refreshCombiner : null);
 
+            // Transfer management to the constituent CSM. NOTE: this is likely to end up double-managed
+            // after the CSM adds the location to the table, but that's acceptable.
+            constituent.columnSourceManager.manage(tableLocation);
+            unmanage(tableLocation);
+
             // Be careful to propagate the systemic attribute properly to child tables
             constituent.setAttribute(Table.SYSTEMIC_TABLE_ATTRIBUTE, result.isSystemicObject());
             return applyTablePermissions.apply(constituent);
         }
 
         private void processPendingLocations(final boolean notifyListeners) {
-            final TableLocationSubscriptionBuffer.LocationUpdate locationUpdate = subscriptionBuffer.processPending();
-            final RowSet removed = processRemovals(locationUpdate);
-            final RowSet added = processAdditions(locationUpdate);
+            final RowSet removed;
+            final RowSet added;
+
+            try (final TableLocationSubscriptionBuffer.LocationUpdate locationUpdate =
+                    subscriptionBuffer.processPending()) {
+                removed = processRemovals(locationUpdate);
+                added = processAdditions(locationUpdate);
+            }
 
             resultRows.update(added, removed);
             if (notifyListeners) {
@@ -242,8 +266,10 @@ private RowSet processAdditions(final TableLocationSubscriptionBuffer.LocationUp
              */
             // TODO (https://github.com/deephaven/deephaven-core/issues/867): Refactor around a ticking partition table
             locationUpdate.getPendingAddedLocationKeys().stream()
+                    .map(LiveSupplier::get)
                     .filter(locationKeyMatcher)
                     .map(tableLocationProvider::getTableLocation)
+                    .peek(this::manage)
                     .map(PendingLocationState::new)
                     .forEach(pendingLocationStates::offer);
             for (final Iterator<PendingLocationState> iter = pendingLocationStates.iterator(); iter.hasNext();) {
@@ -263,13 +289,24 @@ private RowSet processRemovals(final TableLocationSubscriptionBuffer.LocationUpd
             final Set<ImmutableTableLocationKey> relevantRemovedLocations =
                     locationUpdate.getPendingRemovedLocationKeys()
                             .stream()
+                            .map(LiveSupplier::get)
                             .filter(locationKeyMatcher)
                             .collect(Collectors.toSet());
 
             if (relevantRemovedLocations.isEmpty()) {
                 return RowSetFactory.empty();
             }
 
+            // Iterate through the pending locations and remove any that are in the removed set.
+            for (final Iterator<PendingLocationState> iter = pendingLocationStates.iterator(); iter.hasNext();) {
+                final PendingLocationState pendingLocationState = iter.next();
+                if (relevantRemovedLocations.contains(pendingLocationState.location.getKey())) {
+                    iter.remove();
+                    // Release the state and unmanage the location
+                    unmanage(pendingLocationState.release());
+                }
+            }
+
             // At the end of the cycle we need to make sure we unmanage any removed constituents.
             this.removedConstituents = new ArrayList<>(relevantRemovedLocations.size());
             final RowSetBuilderSequential deleteBuilder = RowSetFactory.builderSequential();
@@ -306,6 +343,22 @@ private RowSet processRemovals(final TableLocationSubscriptionBuffer.LocationUpd
             resultLocationTables.setNull(deletedRows);
             return deletedRows;
         }
+
+        @Override
+        public boolean satisfied(final long step) {
+            if (refreshCombiner == null) {
+                throw new UnsupportedOperationException("This method should not be called when result is static");
+            }
+            return refreshCombiner.satisfied(step);
+        }
+
+        @Override
+        public UpdateGraph getUpdateGraph() {
+            if (refreshCombiner == null) {
+                throw new UnsupportedOperationException("This method should not be called when result is static");
+            }
+            return refreshCombiner.getUpdateGraph();
+        }
     }
 
     private static final class PendingLocationState extends IntrusiveDoublyLinkedNode.Impl<PendingLocationState> {