fix: DH-18567: Filter locations before removal allowed checks. (#6607)

Author: cpwright

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

@@ -28,6 +28,7 @@
 import javax.annotation.OverridingMethodsMustInvokeSuper;
 import java.util.ArrayList;
 import java.util.Collection;
+import java.util.stream.Stream;
 
 /**
  * Basic uncoalesced table that only adds keys.
@@ -153,7 +154,6 @@ private void initializeAvailableLocations() {
                     try (final TableLocationSubscriptionBuffer.LocationUpdate locationUpdate =
                             locationBuffer.processPending()) {
                         if (locationUpdate != null) {
-                            maybeRemoveLocations(locationUpdate.getPendingRemovedLocationKeys());
                             maybeAddLocations(locationUpdate.getPendingAddedLocationKeys());
                         }
                     }
@@ -188,14 +188,26 @@ private void maybeAddLocations(@NotNull final Collection<LiveSupplier<ImmutableT
                 .forEach(lk -> columnSourceManager.addLocation(locationProvider.getTableLocation(lk.get())));
     }
 
-    private void maybeRemoveLocations(@NotNull final Collection<LiveSupplier<ImmutableTableLocationKey>> removedKeys) {
+    private void maybeRemoveLocations(@NotNull final Collection<LiveSupplier<ImmutableTableLocationKey>> removedKeys,
+            final boolean removedAllowed) {
         if (removedKeys.isEmpty()) {
             return;
         }
 
-        filterLocationKeys(removedKeys).stream()
+        final Collection<LiveSupplier<ImmutableTableLocationKey>> filteredSuppliers = filterLocationKeys(removedKeys);
+        if (filteredSuppliers.isEmpty()) {
+            return;
+        }
+
+        if (removedAllowed) {
+            filteredSuppliers.stream().map(LiveSupplier::get).forEach(columnSourceManager::removeLocationKey);
+            return;
+        }
+
+        final ImmutableTableLocationKey[] keys = filteredSuppliers.stream()
                 .map(LiveSupplier::get)
-                .forEach(columnSourceManager::removeLocationKey);
+                .toArray(ImmutableTableLocationKey[]::new);
+        throw new TableLocationRemovedException("Source table does not support removed locations", keys);
     }
 
     private void initializeLocationSizes() {
@@ -238,16 +250,8 @@ protected void instrumentedRefresh() {
             try (final TableLocationSubscriptionBuffer.LocationUpdate locationUpdate =
                     locationBuffer.processPending()) {
                 if (locationUpdate != null) {
-                    if (!locationProvider.getUpdateMode().removeAllowed()
-                            && !locationUpdate.getPendingRemovedLocationKeys().isEmpty()) {
-                        // This TLP doesn't support removed locations, we need to throw an exception.
-                        final ImmutableTableLocationKey[] keys = locationUpdate.getPendingRemovedLocationKeys().stream()
-                                .map(LiveSupplier::get).toArray(ImmutableTableLocationKey[]::new);
-                        throw new TableLocationRemovedException(
-                                "Source table does not support removed locations", keys);
-                    }
-
-                    maybeRemoveLocations(locationUpdate.getPendingRemovedLocationKeys());
+                    maybeRemoveLocations(locationUpdate.getPendingRemovedLocationKeys(),
+                            locationProvider.getUpdateMode().removeAllowed());
                     maybeAddLocations(locationUpdate.getPendingAddedLocationKeys());
                 }
             }

engine/table/src/test/java/io/deephaven/engine/table/impl/TestPartitionAwareSourceTable.java

@@ -225,7 +225,6 @@ private void allowLivenessRelease() {
             {
                 allowing(locationProvider).supportsSubscriptions();
                 allowing(locationProvider).unsubscribe(with(any(TableLocationProvider.Listener.class)));
-                will(returnValue(true));
                 for (int li = 0; li < tableLocations.length; ++li) {
                     final TableLocation tableLocation = tableLocations[li];
                     allowing(tableLocation).supportsSubscriptions();