Source:
<code>public static HiveSplitSource bucketed( ConnectorSession session, String databaseName, String tableName, TupleDomain<? extends ColumnHandle> compactEffectivePredicate, int estimatedOutstandingSplitsPerBucket, int maxInitialSplits, DataSize maxOutstandingSplitsSize, HiveSplitLoader splitLoader, Executor executor, CounterStat highMemorySplitSourceCounter) { AtomicReference<State> stateReference = new AtomicReference<>(State.initial()); return new HiveSplitSource( session, databaseName, tableName, compactEffectivePredicate, new PerBucket() { private final Map<Integer, AsyncQueue<InternalHiveSplit>> queues = new ConcurrentHashMap<>(); private final AtomicBoolean finished = new AtomicBoolean(); @Override public ListenableFuture<?> offer(OptionalInt bucketNumber, InternalHiveSplit connectorSplit) { AsyncQueue<InternalHiveSplit> queue = queueFor(bucketNumber); queue.offer(connectorSplit); return immediateFuture(null); } @Override public ListenableFuture<List<ConnectorSplit>> borrowBatchAsync(OptionalInt bucketNumber, int maxSize, Function<List<InternalHiveSplit>, BorrowResult<InternalHiveSplit, List<ConnectorSplit>>> function) { return queueFor(bucketNumber).borrowBatchAsync(maxSize, function); } @Override public void finish() { if (finished.compareAndSet(false, true)) { queues.values().forEach(AsyncQueue::finish); } } @Override public boolean isFinished(OptionalInt bucketNumber) { return queueFor(bucketNumber).isFinished(); } @Override public int rewind(OptionalInt bucketNumber) { throw new UnsupportedOperationException("rewind is not supported for unrewindable split source"); } private AsyncQueue<InternalHiveSplit> queueFor(OptionalInt bucketNumber) { checkArgument(bucketNumber.isPresent()); AtomicBoolean isNew = new AtomicBoolean(); <START> <END> AsyncQueue<InternalHiveSplit> queue = queues.computeIfAbsent(bucketNumber.getAsInt(), ignored -> { isNew.set(true); return new AsyncQueue<>(estimatedOutstandingSplitsPerBucket, executor); }); if (isNew.get() && finished.get()) { queue.finish(); } return queue; } }, maxInitialSplits, maxOutstandingSplitsSize, splitLoader, stateReference, highMemorySplitSourceCounter); }</code><technical_language>nit: unneeded change</technical_language>

Target:
public static HiveSplitSource bucketed( ConnectorSession session, String databaseName, String tableName, TupleDomain<? extends ColumnHandle> compactEffectivePredicate, int estimatedOutstandingSplitsPerBucket, int maxInitialSplits, DataSize maxOutstandingSplitsSize, HiveSplitLoader splitLoader, Executor executor, CounterStat highMemorySplitSourceCounter) { AtomicReference<State> stateReference = new AtomicReference<>(State.initial()); return new HiveSplitSource( session, databaseName, tableName, compactEffectivePredicate, new PerBucket() { private final Map<Integer, AsyncQueue<InternalHiveSplit>> queues = new ConcurrentHashMap<>(); private final AtomicBoolean finished = new AtomicBoolean(); @Override public ListenableFuture<?> offer(OptionalInt bucketNumber, InternalHiveSplit connectorSplit) { AsyncQueue<InternalHiveSplit> queue = queueFor(bucketNumber); queue.offer(connectorSplit); return immediateFuture(null); } @Override public ListenableFuture<List<ConnectorSplit>> borrowBatchAsync(OptionalInt bucketNumber, int maxSize, Function<List<InternalHiveSplit>, BorrowResult<InternalHiveSplit, List<ConnectorSplit>>> function) { return queueFor(bucketNumber).borrowBatchAsync(maxSize, function); } @Override public void noMoreSplits() { if (finished.compareAndSet(false, true)) { queues.values().forEach(AsyncQueue::finish); } } @Override public boolean isFinished(OptionalInt bucketNumber) { return queueFor(bucketNumber).isFinished(); } @Override public int rewind(OptionalInt bucketNumber) { throw new UnsupportedOperationException("rewind is not supported for unrewindable split source"); } private AsyncQueue<InternalHiveSplit> queueFor(OptionalInt bucketNumber) { checkArgument(bucketNumber.isPresent()); AtomicBoolean isNew = new AtomicBoolean(); AsyncQueue<InternalHiveSplit> queue = queues.computeIfAbsent(bucketNumber.getAsInt(), ignored -> { isNew.set(true); return new AsyncQueue<>(estimatedOutstandingSplitsPerBucket, executor); }); if (isNew.get() && finished.get()) { queue.finish(); } return queue; } }, maxInitialSplits, maxOutstandingSplitsSize, splitLoader, stateReference, highMemorySplitSourceCounter); }