From 4d0b82f3dbef8b6d3b8ea832d0627fcc94bb162d Mon Sep 17 00:00:00 2001
From: Ryan Caudy <ryan@deephaven.io>
Date: Thu, 15 Aug 2024 14:30:05 -0400
Subject: [PATCH] fix: Fix PartitioningColumnDataIndex: region RowSet
 provenance and index RowSet tracking (#5942)

---
 .../regioned/PartitioningColumnDataIndex.java    |  2 +-
 .../regioned/RegionedColumnSourceManager.java    | 16 ++++++++++------
 2 files changed, 11 insertions(+), 7 deletions(-)

diff --git a/engine/table/src/main/java/io/deephaven/engine/table/impl/sources/regioned/PartitioningColumnDataIndex.java b/engine/table/src/main/java/io/deephaven/engine/table/impl/sources/regioned/PartitioningColumnDataIndex.java
index fb972aa92e9..ae95d86b8ee 100644
--- a/engine/table/src/main/java/io/deephaven/engine/table/impl/sources/regioned/PartitioningColumnDataIndex.java
+++ b/engine/table/src/main/java/io/deephaven/engine/table/impl/sources/regioned/PartitioningColumnDataIndex.java
@@ -216,7 +216,7 @@ private void handleKey(
             indexKeySource.set(addedKeyPos, locationKey);
 
             indexRowSetSource.ensureCapacity(addedKeyPos + 1);
-            indexRowSetSource.set(addedKeyPos, regionRowSet.shift(regionFirstRowKey));
+            indexRowSetSource.set(addedKeyPos, regionRowSet.shift(regionFirstRowKey).toTracking());
         } else {
             // noinspection DataFlowIssue
             final WritableRowSet existingRowSet = indexRowSetSource.get(pos).writableCast();
diff --git a/engine/table/src/main/java/io/deephaven/engine/table/impl/sources/regioned/RegionedColumnSourceManager.java b/engine/table/src/main/java/io/deephaven/engine/table/impl/sources/regioned/RegionedColumnSourceManager.java
index 0344e117d75..1889f5dbb25 100644
--- a/engine/table/src/main/java/io/deephaven/engine/table/impl/sources/regioned/RegionedColumnSourceManager.java
+++ b/engine/table/src/main/java/io/deephaven/engine/table/impl/sources/regioned/RegionedColumnSourceManager.java
@@ -291,11 +291,10 @@ private WritableRowSet update(final boolean initializing) {
             if (entry.pollUpdates(addedRowSetBuilder)) {
                 // Changes were detected, update the row set in the table and mark the row/column as modified.
                 /*
-                 * Since TableLocationState.getRowSet() returns a copy(), we should consider adding an UpdateCommitter
-                 * to close() the previous row sets for modified locations. This is not important for current
-                 * implementations, since they always allocate new, flat RowSets.
+                 * Since TableLocationState.getRowSet() returns a copy(), we own entry.rowSetAtLastUpdate and can
+                 * propagate it without making another copy().
                  */
-                rowSetSource.set(entry.regionIndex, entry.location.getRowSet());
+                rowSetSource.set(entry.regionIndex, entry.rowSetAtLastUpdate);
                 if (modifiedRegionBuilder != null) {
                     modifiedRegionBuilder.appendKey(entry.regionIndex);
                 }
@@ -346,7 +345,7 @@ private WritableRowSet update(final boolean initializing) {
                                 wcs.set(entry.regionIndex, entry.location.getKey().getPartitionValue(key)));
                 // @formatter:on
                 locationSource.set(entry.regionIndex, entry.location);
-                rowSetSource.set(entry.regionIndex, entry.location.getRowSet());
+                rowSetSource.set(entry.regionIndex, entry.rowSetAtLastUpdate);
             });
         }
 
@@ -574,7 +573,12 @@ private boolean pollUpdates(final RowSetBuilderSequential addedRowSetBuilder) {
                             .appendRange(regionFirstKey + subRegionFirstKey, regionFirstKey + subRegionLastKey));
                 }
             } finally {
-                rowSetAtLastUpdate.close();
+                /*
+                 * Since we record rowSetAtLastUpdate in the RowSet column of our includedLocationsTable, we must not
+                 * close() the old rowSetAtLastUpdate here. We should instead consider adding an UpdateCommitter to
+                 * close() the previous RowSets for modified locations, but this is not important for current
+                 * implementations since they always allocate new, flat RowSets.
+                 */
                 rowSetAtLastUpdate = updateRowSet;
             }
             // There was a change to the row set.