Deprecate partitionHandle argument on ConnectorSplitSource::getNextBatch

It was always a `NOT_PARTITION` since the removal of grouped execution,
so it can be hardcoded at this point.
The ConnectorNodePartitioningProvider::listPartitionHandles method can also
be deprecated at this point

Author: skrzypo987

core/trino-main/src/main/java/io/trino/execution/scheduler/FixedSourcePartitionedScheduler.java

@@ -24,7 +24,6 @@
 import io.trino.metadata.InternalNode;
 import io.trino.metadata.Split;
 import io.trino.server.DynamicFilterService;
-import io.trino.spi.connector.ConnectorPartitionHandle;
 import io.trino.split.SplitSource;
 import io.trino.sql.planner.plan.PlanNodeId;
 
@@ -41,7 +40,6 @@
 import static com.google.common.base.Verify.verify;
 import static io.airlift.concurrent.MoreFutures.whenAnyComplete;
 import static io.trino.execution.scheduler.SourcePartitionedScheduler.newSourcePartitionedSchedulerAsSourceScheduler;
-import static io.trino.spi.connector.NotPartitionedPartitionHandle.NOT_PARTITIONED;
 import static java.util.Objects.requireNonNull;
 
 public class FixedSourcePartitionedScheduler
@@ -64,15 +62,13 @@ public FixedSourcePartitionedScheduler(
             BucketNodeMap bucketNodeMap,
             int splitBatchSize,
             NodeSelector nodeSelector,
-            List<ConnectorPartitionHandle> partitionHandles,
             DynamicFilterService dynamicFilterService,
             TableExecuteContextManager tableExecuteContextManager)
     {
         requireNonNull(stageExecution, "stageExecution is null");
         requireNonNull(splitSources, "splitSources is null");
         requireNonNull(bucketNodeMap, "bucketNodeMap is null");
         checkArgument(!requireNonNull(nodes, "nodes is null").isEmpty(), "nodes is empty");
-        requireNonNull(partitionHandles, "partitionHandles is null");
         requireNonNull(tableExecuteContextManager, "tableExecuteContextManager is null");
 
         this.stageExecution = stageExecution;
@@ -83,7 +79,6 @@ public FixedSourcePartitionedScheduler(
         BucketedSplitPlacementPolicy splitPlacementPolicy = new BucketedSplitPlacementPolicy(nodeSelector, nodes, bucketNodeMap, stageExecution::getAllTasks);
 
         ArrayList<SourceScheduler> sourceSchedulers = new ArrayList<>();
-        checkArgument(partitionHandles.equals(ImmutableList.of(NOT_PARTITIONED)), "PartitionHandles should be [NOT_PARTITIONED]");
 
         boolean firstPlanNode = true;
 

core/trino-main/src/main/java/io/trino/execution/scheduler/SqlQueryScheduler.java

@@ -61,7 +61,6 @@
 import io.trino.spi.ErrorCode;
 import io.trino.spi.QueryId;
 import io.trino.spi.TrinoException;
-import io.trino.spi.connector.ConnectorPartitionHandle;
 import io.trino.spi.exchange.Exchange;
 import io.trino.spi.exchange.ExchangeContext;
 import io.trino.spi.exchange.ExchangeId;
@@ -151,7 +150,6 @@
 import static io.trino.spi.StandardErrorCode.GENERIC_INTERNAL_ERROR;
 import static io.trino.spi.StandardErrorCode.NO_NODES_AVAILABLE;
 import static io.trino.spi.StandardErrorCode.REMOTE_TASK_FAILED;
-import static io.trino.spi.connector.NotPartitionedPartitionHandle.NOT_PARTITIONED;
 import static io.trino.sql.planner.SystemPartitioningHandle.FIXED_BROADCAST_DISTRIBUTION;
 import static io.trino.sql.planner.SystemPartitioningHandle.FIXED_HASH_DISTRIBUTION;
 import static io.trino.sql.planner.SystemPartitioningHandle.SCALED_WRITER_DISTRIBUTION;
@@ -1428,7 +1426,6 @@ else if (partitioningHandle.equals(SCALED_WRITER_DISTRIBUTION)) {
                     List<PlanNodeId> schedulingOrder = fragment.getPartitionedSources();
                     Optional<CatalogName> catalogName = partitioningHandle.getConnectorId();
                     checkArgument(catalogName.isPresent(), "No connector ID for partitioning handle: %s", partitioningHandle);
-                    List<ConnectorPartitionHandle> connectorPartitionHandles = ImmutableList.of(NOT_PARTITIONED);
 
                     BucketNodeMap bucketNodeMap;
                     List<InternalNode> stageNodeList;
@@ -1454,7 +1451,6 @@ else if (partitioningHandle.equals(SCALED_WRITER_DISTRIBUTION)) {
                             bucketNodeMap,
                             splitBatchSize,
                             nodeScheduler.createNodeSelector(session, catalogName),
-                            connectorPartitionHandles,
                             dynamicFilterService,
                             tableExecuteContextManager);
                 }

core/trino-main/src/main/java/io/trino/split/ConnectorAwareSplitSource.java

@@ -27,7 +27,6 @@
 
 import static com.google.common.util.concurrent.MoreExecutors.directExecutor;
 import static io.airlift.concurrent.MoreFutures.toListenableFuture;
-import static io.trino.spi.connector.NotPartitionedPartitionHandle.NOT_PARTITIONED;
 import static java.util.Objects.requireNonNull;
 
 public class ConnectorAwareSplitSource
@@ -51,7 +50,7 @@ public CatalogName getCatalogName()
     @Override
     public ListenableFuture<SplitBatch> getNextBatch(int maxSize)
     {
-        ListenableFuture<ConnectorSplitBatch> nextBatch = toListenableFuture(source.getNextBatch(NOT_PARTITIONED, maxSize));
+        ListenableFuture<ConnectorSplitBatch> nextBatch = toListenableFuture(source.getNextBatch(maxSize));
         return Futures.transform(nextBatch, splitBatch -> {
             ImmutableList.Builder<Split> result = ImmutableList.builder();
             for (ConnectorSplit connectorSplit : splitBatch.getSplits()) {

core/trino-main/src/main/java/io/trino/sql/planner/NodePartitioningManager.java

@@ -29,7 +29,6 @@
 import io.trino.spi.connector.BucketFunction;
 import io.trino.spi.connector.ConnectorBucketNodeMap;
 import io.trino.spi.connector.ConnectorNodePartitioningProvider;
-import io.trino.spi.connector.ConnectorPartitionHandle;
 import io.trino.spi.connector.ConnectorSplit;
 import io.trino.spi.type.Type;
 import io.trino.split.EmptySplit;
@@ -103,19 +102,6 @@ public BucketFunction getBucketFunction(Session session, PartitioningHandle part
         return bucketFunction;
     }
 
-    public List<ConnectorPartitionHandle> listPartitionHandles(
-            Session session,
-            PartitioningHandle partitioningHandle)
-    {
-        CatalogName catalogName = partitioningHandle.getConnectorId()
-                .orElseThrow(() -> new IllegalArgumentException("No connector ID for partitioning handle: " + partitioningHandle));
-        ConnectorNodePartitioningProvider partitioningProvider = getPartitioningProvider(catalogName);
-        return partitioningProvider.listPartitionHandles(
-                partitioningHandle.getTransactionHandle().orElseThrow(() -> new IllegalArgumentException("No transactionHandle for partitioning handle: " + partitioningHandle)),
-                session.toConnectorSession(catalogName),
-                partitioningHandle.getConnectorHandle());
-    }
-
     public NodePartitionMap getNodePartitioningMap(Session session, PartitioningHandle partitioningHandle)
     {
         requireNonNull(session, "session is null");

core/trino-main/src/test/java/io/trino/execution/scheduler/TestSourcePartitionedScheduler.java

@@ -42,7 +42,6 @@
 import io.trino.operator.RetryPolicy;
 import io.trino.server.DynamicFilterService;
 import io.trino.spi.QueryId;
-import io.trino.spi.connector.ConnectorPartitionHandle;
 import io.trino.spi.connector.ConnectorSplit;
 import io.trino.spi.connector.ConnectorSplitSource;
 import io.trino.spi.connector.DynamicFilter;
@@ -81,7 +80,6 @@
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.function.Supplier;
 
-import static com.google.common.base.Preconditions.checkArgument;
 import static io.airlift.concurrent.Threads.daemonThreadsNamed;
 import static io.trino.SessionTestUtils.TEST_SESSION;
 import static io.trino.execution.scheduler.NodeSchedulerConfig.SplitsBalancingPolicy.NODE;
@@ -94,7 +92,6 @@
 import static io.trino.metadata.FunctionManager.createTestingFunctionManager;
 import static io.trino.metadata.MetadataManager.createTestMetadataManager;
 import static io.trino.spi.StandardErrorCode.NO_NODES_AVAILABLE;
-import static io.trino.spi.connector.NotPartitionedPartitionHandle.NOT_PARTITIONED;
 import static io.trino.spi.type.BigintType.BIGINT;
 import static io.trino.spi.type.VarcharType.VARCHAR;
 import static io.trino.sql.DynamicFilters.createDynamicFilterExpression;
@@ -701,7 +698,7 @@ private static ConnectorSplitSource createBlockedSplitSource()
         return new ConnectorSplitSource()
         {
             @Override
-            public CompletableFuture<ConnectorSplitBatch> getNextBatch(ConnectorPartitionHandle partitionHandle, int maxSize)
+            public CompletableFuture<ConnectorSplitBatch> getNextBatch(int maxSize)
             {
                 return new CompletableFuture<>();
             }
@@ -781,9 +778,8 @@ synchronized void addSplits(int count)
         }
 
         @Override
-        public CompletableFuture<ConnectorSplitBatch> getNextBatch(ConnectorPartitionHandle partitionHandle, int maxSize)
+        public CompletableFuture<ConnectorSplitBatch> getNextBatch(int maxSize)
         {
-            checkArgument(partitionHandle.equals(NOT_PARTITIONED), "partitionHandle must be NOT_PARTITIONED");
             return notEmptyFuture
                     .thenApply(x -> getBatch(maxSize))
                     .thenApply(splits -> new ConnectorSplitBatch(splits, isFinished()));

core/trino-spi/src/main/java/io/trino/spi/connector/ConnectorNodePartitioningProvider.java

@@ -18,9 +18,6 @@
 import java.util.List;
 import java.util.function.ToIntFunction;
 
-import static io.trino.spi.connector.NotPartitionedPartitionHandle.NOT_PARTITIONED;
-import static java.util.Collections.singletonList;
-
 public interface ConnectorNodePartitioningProvider
 {
     // TODO: Use ConnectorPartitionHandle (instead of int) to represent individual buckets.
@@ -31,10 +28,13 @@ public interface ConnectorNodePartitioningProvider
      * <p>
      * This method must be implemented for connectors that support addressable split discovery.
      * The partitions return here will be used as address for the purpose of split discovery.
+     *
+     * @deprecated The method is not used. Implementations can be simply removed
      */
+    @Deprecated
     default List<ConnectorPartitionHandle> listPartitionHandles(ConnectorTransactionHandle transactionHandle, ConnectorSession session, ConnectorPartitioningHandle partitioningHandle)
     {
-        return singletonList(NOT_PARTITIONED);
+        throw new UnsupportedOperationException();
     }
 
     ConnectorBucketNodeMap getBucketNodeMap(ConnectorTransactionHandle transactionHandle, ConnectorSession session, ConnectorPartitioningHandle partitioningHandle);

core/trino-spi/src/main/java/io/trino/spi/connector/ConnectorPartitionHandle.java

@@ -14,6 +14,7 @@
 package io.trino.spi.connector;
 
 @SuppressWarnings("ClassMayBeInterface")
+@Deprecated
 public abstract class ConnectorPartitionHandle
 {
     @Override

core/trino-spi/src/main/java/io/trino/spi/connector/ConnectorSplitSource.java

@@ -18,12 +18,22 @@
 import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
 
+import static io.trino.spi.connector.NotPartitionedPartitionHandle.NOT_PARTITIONED;
 import static java.util.Objects.requireNonNull;
 
 public interface ConnectorSplitSource
         extends Closeable
 {
-    CompletableFuture<ConnectorSplitBatch> getNextBatch(ConnectorPartitionHandle partitionHandle, int maxSize);
+    @Deprecated
+    default CompletableFuture<ConnectorSplitBatch> getNextBatch(ConnectorPartitionHandle partitionHandle, int maxSize)
+    {
+        throw new UnsupportedOperationException();
+    }
+
+    default CompletableFuture<ConnectorSplitBatch> getNextBatch(int maxSize)
+    {
+        return getNextBatch(NOT_PARTITIONED, maxSize);
+    }
 
     @Override
     void close();

core/trino-spi/src/main/java/io/trino/spi/connector/FixedSplitSource.java

@@ -17,7 +17,6 @@
 import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
 
-import static io.trino.spi.connector.NotPartitionedPartitionHandle.NOT_PARTITIONED;
 import static java.util.Objects.requireNonNull;
 import static java.util.concurrent.CompletableFuture.completedFuture;
 import static java.util.stream.Collectors.toUnmodifiableList;
@@ -50,12 +49,8 @@ private FixedSplitSource(Iterable<? extends ConnectorSplit> splits, Optional<Lis
 
     @SuppressWarnings("ObjectEquality")
     @Override
-    public CompletableFuture<ConnectorSplitBatch> getNextBatch(ConnectorPartitionHandle partitionHandle, int maxSize)
+    public CompletableFuture<ConnectorSplitBatch> getNextBatch(int maxSize)
     {
-        if (!partitionHandle.equals(NOT_PARTITIONED)) {
-            throw new IllegalArgumentException("partitionHandle must be NOT_PARTITIONED");
-        }
-
         int remainingSplits = splits.size() - offset;
         int size = Math.min(remainingSplits, maxSize);
         List<ConnectorSplit> results = splits.subList(offset, offset + size);

core/trino-spi/src/main/java/io/trino/spi/connector/NotPartitionedPartitionHandle.java

@@ -13,6 +13,7 @@
  */
 package io.trino.spi.connector;
 
+@Deprecated
 public final class NotPartitionedPartitionHandle
         extends ConnectorPartitionHandle
 {