From 064aa6f555eec4806f7950c9065b75fb801f3967 Mon Sep 17 00:00:00 2001
From: Colin Alworth <colinalworth@deephaven.io>
Date: Tue, 15 Oct 2024 21:40:37 -0500
Subject: [PATCH] fix: Capture QueryCompiler for plotting PartitionedTable
 transform (#6199)

Also removes other dead/deprecated code.

Fixes #6187
---
 .../main/java/io/deephaven/plot/AxesImpl.java  | 18 +++---------------
 .../plot/filters/SelectableDataSet.java        | 14 --------------
 .../plot/example_plots/SimpleXYTable.java      |  2 --
 3 files changed, 3 insertions(+), 31 deletions(-)

diff --git a/Plot/src/main/java/io/deephaven/plot/AxesImpl.java b/Plot/src/main/java/io/deephaven/plot/AxesImpl.java
index 235c18b943c..2535736ba56 100644
--- a/Plot/src/main/java/io/deephaven/plot/AxesImpl.java
+++ b/Plot/src/main/java/io/deephaven/plot/AxesImpl.java
@@ -5,6 +5,7 @@
 
 import io.deephaven.api.ColumnName;
 import io.deephaven.api.agg.Aggregation;
+import io.deephaven.engine.context.ExecutionContext;
 import io.deephaven.engine.table.impl.MemoizedOperationKey;
 import io.deephaven.plot.axisformatters.AxisFormat;
 import io.deephaven.plot.axisformatters.NanosAxisFormat;
@@ -277,24 +278,11 @@ private static SelectableDataSet getAggregatedSelectableDataSet(final Selectable
         }
         final Collection<? extends Aggregation> aggs = aggSupplier.get();
         final Collection<? extends ColumnName> columnNames = ColumnName.from(cols);
-        final Function<Table, Table> applyAggs = t -> t.aggBy(aggs, columnNames);
+        ExecutionContext ctx = ExecutionContext.newBuilder().captureQueryCompiler().build();
+        final Function<Table, Table> applyAggs = t -> ctx.apply(() -> t.aggBy(aggs, columnNames));
         return sds.transform(MemoizedOperationKey.aggBy(aggs, false, null, columnNames), applyAggs);
     }
 
-    private static SelectableDataSet getLastBySelectableDataSet(final SelectableDataSet sds, final String... columns) {
-        final List<String> cols = new ArrayList<>();
-        Collections.addAll(cols, columns);
-        return getLastBySelectableDataSet(sds, cols);
-    }
-
-    private static SelectableDataSet getLastBySelectableDataSet(final SelectableDataSet sds,
-            final Collection<String> columns) {
-        if (sds instanceof SelectableDataSetOneClick) {
-            Collections.addAll(columns, ((SelectableDataSetOneClick) sds).getByColumns());
-        }
-        return sds.transform(columns, t -> ((Table) t).lastBy(columns));
-    }
-
     public Set<SwappableTable> getSwappableTables() {
         final Set<SwappableTable> result = new HashSet<>();
 
diff --git a/Plot/src/main/java/io/deephaven/plot/filters/SelectableDataSet.java b/Plot/src/main/java/io/deephaven/plot/filters/SelectableDataSet.java
index 702b7d81d39..7db0829920f 100644
--- a/Plot/src/main/java/io/deephaven/plot/filters/SelectableDataSet.java
+++ b/Plot/src/main/java/io/deephaven/plot/filters/SelectableDataSet.java
@@ -9,7 +9,6 @@
 import io.deephaven.engine.table.TableDefinition;
 import org.jetbrains.annotations.NotNull;
 
-import java.util.Collection;
 import java.util.function.Function;
 
 /**
@@ -45,19 +44,6 @@ default SwappableTable getSwappableTable(final Comparable seriesName, final Char
         return getSwappableTable(seriesName, chart, null, cols);
     }
 
-    /**
-     * Gets a version of the SelectableDataSet with a lastBy applied to the tables.
-     *
-     * @param groupByColumns The grouping columns for the lastBy
-     * @return a new SelectableDataSet with lastBy applied
-     *
-     * @deprecated This method will be removed in a future release, use {@link #transform(Object, Function)} instead.
-     */
-    @Deprecated()
-    default SelectableDataSet<KEY_TYPE, VALUE_TYPE> getLastBy(final Collection<String> groupByColumns) {
-        return transform(groupByColumns, t -> t.lastBy(groupByColumns));
-    }
-
     /**
      * Produces a derivative {@link SelectableDataSet} with the specified transformation applied to all internal tables.
      *
diff --git a/Plot/src/test/java/io/deephaven/plot/example_plots/SimpleXYTable.java b/Plot/src/test/java/io/deephaven/plot/example_plots/SimpleXYTable.java
index 77cc49fe474..27065a57dd6 100644
--- a/Plot/src/test/java/io/deephaven/plot/example_plots/SimpleXYTable.java
+++ b/Plot/src/test/java/io/deephaven/plot/example_plots/SimpleXYTable.java
@@ -9,8 +9,6 @@
 import io.deephaven.engine.table.Table;
 import io.deephaven.engine.util.TableTools;
 
-import static io.deephaven.plot.filters.Selectables.oneClick;
-
 public class SimpleXYTable {
 
     public static void main(String[] args) {