You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@druid.apache.org by "gianm (via GitHub)" <gi...@apache.org> on 2025/12/29 18:23:47 UTC
[PR] MSQ: Use VSF, give stages control over inputs. (druid)
gianm opened a new pull request, #18873:
URL: https://github.com/apache/druid/pull/18873
This patch integrates MSQ with virtual storage. It also refactors how MSQ reads inputs to give stages more control over how inputs are read and merged. In particular, stages are now able to fully control merging logic.
The main changes:
1) Integrate with virtual storage: merge the two DataSegmentProvider impls (Dart and Task) into DataSegmentProviderImpl that relies on SegmentManager.
2) Give stages control over input merging: rework InputSliceReader to return ReadablePartitions directly, without embedding any merging logic. Break out StandardPartitionReader as a separate class.
Other changes:
1) Move ReadableInput to the querykit package. It is no longer specific to the MSQ framework.
2) Remove StandardStageProcessor, refactoring dependent code to not require it.
3) Remove ExternalColumnSelectorFactory wrapper. Type casting is now handled directly by RowBasedColumnSelectorFactory.
4) Include full query context in worker context, rather than just a subset.
Includes #18871.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For additional commands, e-mail: commits-help@druid.apache.org
Re: [PR] MSQ: Use VSF, give stages control over inputs. (druid)
Posted by "gianm (via GitHub)" <gi...@apache.org>.
gianm merged PR #18873:
URL: https://github.com/apache/druid/pull/18873
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For additional commands, e-mail: commits-help@druid.apache.org
Re: [PR] MSQ: Use VSF, give stages control over inputs. (druid)
Posted by "gianm (via GitHub)" <gi...@apache.org>.
gianm commented on code in PR #18873:
URL: https://github.com/apache/druid/pull/18873#discussion_r2677677172
##########
multi-stage-query/src/main/java/org/apache/druid/msq/querykit/ReadableInputQueue.java:
##########
@@ -0,0 +1,366 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.msq.querykit;
+
+import com.google.common.collect.Sets;
+import com.google.common.util.concurrent.Futures;
+import com.google.common.util.concurrent.ListenableFuture;
+import com.google.errorprone.annotations.concurrent.GuardedBy;
+import org.apache.druid.common.guava.FutureUtils;
+import org.apache.druid.error.DruidException;
+import org.apache.druid.frame.channel.ReadableFrameChannel;
+import org.apache.druid.java.util.common.logger.Logger;
+import org.apache.druid.msq.exec.DataServerQueryHandler;
+import org.apache.druid.msq.exec.std.StandardPartitionReader;
+import org.apache.druid.msq.input.LoadableSegment;
+import org.apache.druid.msq.input.PhysicalInputSlice;
+import org.apache.druid.msq.input.stage.ReadablePartition;
+import org.apache.druid.msq.kernel.StageId;
+import org.apache.druid.msq.kernel.StagePartition;
+import org.apache.druid.segment.SegmentReference;
+import org.apache.druid.segment.loading.AcquireSegmentAction;
+import org.apache.druid.utils.CloseableUtils;
+
+import javax.annotation.Nullable;
+import java.io.Closeable;
+import java.io.IOException;
+import java.util.ArrayDeque;
+import java.util.LinkedHashSet;
+import java.util.List;
+import java.util.Queue;
+import java.util.Set;
+import java.util.concurrent.atomic.AtomicBoolean;
+
+/**
+ * Queue for returning {@link ReadableInput} from a list of {@link PhysicalInputSlice}.
+ *
+ * When closed, this object cancels all pending segment loads and releases all segments that have not yet been
+ * acquired by callers through {@link SegmentReferenceHolder#getSegmentReferenceOnce()}. Callers that have acquired
+ * segment references are responsible for closing those references, they will not be closed by this class.
+ */
+public class ReadableInputQueue implements Closeable
+{
+ private static final Logger log = new Logger(ReadableInputQueue.class);
+
+ /**
+ * Partitions to be read.
+ */
+ @GuardedBy("this")
+ private final Queue<ReadablePartition> readablePartitions = new ArrayDeque<>();
+
+ /**
+ * Segments to be loaded. Some may already be cached locally.
+ */
+ @GuardedBy("this")
+ private final Queue<LoadableSegment> loadableSegments = new ArrayDeque<>();
+
+ /**
+ * Realtime servers to be queried.
+ */
+ @GuardedBy("this")
+ private final Queue<DataServerQueryHandler> queryableServers = new ArrayDeque<>();
+
+ /**
+ * Segments currently being loaded.
+ */
+ @GuardedBy("this")
+ private final Set<AcquireSegmentAction> loadingSegments = new LinkedHashSet<>();
+
+ /**
+ * Segments that have been loaded.
+ */
+ @GuardedBy("this")
+ private final Set<SegmentReferenceHolder> loadedSegments = new LinkedHashSet<>();
+
+ /**
+ * Futures that are generated by loadahead, prior to a call to {@link #nextInput()} being made.
+ */
+ @GuardedBy("this")
+ private final Set<ListenableFuture<ReadableInput>> loadaheadFutures = Sets.newIdentityHashSet();
+
+ private final String queryId;
+ private final StandardPartitionReader partitionReader;
+ private final int loadahead;
+ private final AtomicBoolean loadaheadStarted = new AtomicBoolean(false);
+
+ public ReadableInputQueue(
+ final String queryId,
+ final StandardPartitionReader partitionReader,
+ final List<PhysicalInputSlice> slices,
+ final int loadahead
+ )
+ {
+ this.queryId = queryId;
+ this.partitionReader = partitionReader;
+ this.loadahead = loadahead;
+
+ for (final PhysicalInputSlice slice : slices) {
+ // First add all locally-cached segments, before any non-locally-cached segments. This promotes processing
+ // the locally-cached ones first. It also improves efficiency of loadahead, since once we get to the
+ // set of non-locally-cached segments, the "next" one is always one that is worth loading ahead.
+ for (final LoadableSegment loadableSegment : slice.getLoadableSegments()) {
+ if (loadableSegment.isLikelyCached()) {
+ loadableSegments.add(loadableSegment);
+ }
+ }
+
+ for (final LoadableSegment loadableSegment : slice.getLoadableSegments()) {
+ if (!loadableSegment.isLikelyCached()) {
+ loadableSegments.add(loadableSegment);
+ }
+ }
+
+ queryableServers.addAll(slice.getQueryableServers());
+ slice.getReadablePartitions().forEach(readablePartitions::add);
+ }
+ }
+
+ /**
+ * Starts loading up to {@link #loadahead} segments for future calls to {@link #nextInput()}, if this method
+ * has not yet been called. If it has previously been called, subsequent calls do nothing.
+ */
+ public void startLoadaheadIfNeeded()
+ {
+ if (loadaheadStarted.compareAndSet(false, true)) {
+ for (int i = 0; i < loadahead; i++) {
+ if (!addLoadaheadFuture()) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * Returns the number of remaining inputs that can be returned by calls to {@link #nextInput()}.
+ */
+ public int remaining()
+ {
+ synchronized (this) {
+ return readablePartitions.size() + loadableSegments.size() + queryableServers.size() + loadaheadFutures.size();
+ }
+ }
+
+ /**
+ * Returns the next {@link ReadableInput}. The future resolves when the input is ready to read.
+ */
+ @Nullable
+ public ListenableFuture<ReadableInput> nextInput()
+ {
+ ListenableFuture<ReadableInput> future;
+
+ future = nextServerInput();
+ if (future != null) {
+ return future;
+ }
+
+ future = nextChannelInput();
+ if (future != null) {
+ return future;
+ }
+
+ future = nextSegmentInput();
+ if (future != null) {
+ return future;
+ }
+
+ return null;
+ }
+
+ /**
+ * Returns the next input from {@link #queryableServers}, if any. Returns null if none remain.
+ */
+ @Nullable
+ private ListenableFuture<ReadableInput> nextServerInput()
+ {
+ final DataServerQueryHandler handler;
+ synchronized (this) {
+ handler = queryableServers.poll();
+ }
+
+ if (handler == null) {
+ return null;
+ }
+
+ return Futures.immediateFuture(ReadableInput.dataServerQuery(handler));
+ }
+
+ /**
+ * Returns the next input from {@link #readablePartitions}, if any. Returns null if none remain.
+ */
+ @Nullable
+ private ListenableFuture<ReadableInput> nextChannelInput()
+ {
+ final ReadablePartition readablePartition;
+ synchronized (this) {
+ readablePartition = readablePartitions.poll();
+ }
+
+ if (readablePartition == null) {
+ return null;
+ }
+
+ ReadableFrameChannel channel = null;
+ try {
+ channel = partitionReader.openChannel(readablePartition);
+ return Futures.immediateFuture(
+ ReadableInput.channel(
+ channel,
+ partitionReader.frameReader(readablePartition.getStageNumber()),
+ new StagePartition(
+ new StageId(queryId, readablePartition.getStageNumber()),
+ readablePartition.getPartitionNumber()
+ )
+ )
+ );
+ }
+ catch (IOException e) {
+ throw CloseableUtils.closeAndWrapInCatch(e, channel);
+ }
+ }
+
+ /**
+ * Returns the next input from {@link #loadableSegments}, if any. Returns null if none remain.
+ */
+ @Nullable
+ private ListenableFuture<ReadableInput> nextSegmentInput()
+ {
+ // Pick a loadahead future, preferring ones that are already loaded.
+ ListenableFuture<ReadableInput> selectedLoadaheadFuture = null;
+ synchronized (this) {
+ for (ListenableFuture<ReadableInput> f : loadaheadFutures) {
+ if (selectedLoadaheadFuture == null || f.isDone()) {
+ selectedLoadaheadFuture = f;
+ if (f.isDone()) {
+ break;
+ }
+ }
+ }
+
+ if (selectedLoadaheadFuture != null) {
+ loadaheadFutures.remove(selectedLoadaheadFuture);
+ addLoadaheadFuture(); // Replace the one we just took out.
+ return selectedLoadaheadFuture;
+ }
+ }
+
+ return loadNextSegment();
+ }
+
+ /**
+ * Load the next segment from {@link #loadableSegments} and return a future to its reference. Returns null
+ * if {@link #loadableSegments} is empty.
+ */
+ @Nullable
+ private ListenableFuture<ReadableInput> loadNextSegment()
+ {
+ synchronized (this) {
+ final LoadableSegment nextLoadableSegment = loadableSegments.poll();
+ if (nextLoadableSegment == null) {
+ return null;
+ }
+
+ final AcquireSegmentAction acquireSegmentAction = nextLoadableSegment.acquire();
+ loadingSegments.add(acquireSegmentAction);
+ return FutureUtils.transform(
+ acquireSegmentAction.getSegmentFuture(),
+ segment -> {
+ synchronized (ReadableInputQueue.this) {
+ // Transfer segment from "loadingSegments" to "loadedSegments" and return a reference to it.
+ if (loadingSegments.remove(acquireSegmentAction)) {
+ try {
+ final SegmentReferenceHolder referenceHolder = new SegmentReferenceHolder(
+ new SegmentReference(
+ nextLoadableSegment.descriptor(),
+ segment.getReferenceProvider().acquireReference(),
+ acquireSegmentAction // Release the hold when the SegmentReference is closed.
+ ),
+ nextLoadableSegment.inputCounters(),
+ nextLoadableSegment.description()
+ );
+ loadedSegments.add(referenceHolder);
+ return ReadableInput.segment(referenceHolder);
+ }
+ catch (Throwable e) {
+ // Javadoc for segment.acquireReference() suggests it can throw exceptions; handle that here
+ // by closing the original AcquireSegmentAction.
+ throw CloseableUtils.closeAndWrapInCatch(e, acquireSegmentAction);
+ }
+ } else {
+ throw DruidException.defensive(
+ "Segment[%s] removed from loadingSegments before loading complete. It is possible that close() "
+ + "was called with futures in flight.",
+ nextLoadableSegment.descriptor()
+ );
+ }
+ }
+ }
+ );
+ }
+ }
+
+ /**
+ * Calls {@link #nextSegmentInput()} and adds the future to {@link #loadaheadFutures}. Returns whether a future
+ * was added.
+ */
+ private boolean addLoadaheadFuture()
+ {
+ final ListenableFuture<ReadableInput> nextFuture = loadNextSegment();
+ if (nextFuture != null) {
+ synchronized (this) {
+ loadaheadFutures.add(nextFuture);
+ }
+ return true;
+ } else {
+ return false;
+ }
+ }
+
+ @Override
+ public void close()
+ {
+ synchronized (this) {
+ readablePartitions.clear();
+ queryableServers.clear();
+ loadableSegments.clear();
+
+ // Cancel all pending segment loads.
+ for (AcquireSegmentAction acquireSegmentAction : loadingSegments) {
+ CloseableUtils.closeAndSuppressExceptions(
+ acquireSegmentAction,
+ e -> log.warn(e, "Failed to close loadingSegment[%s]", acquireSegmentAction)
Review Comment:
I thought about that but didn't want to add a field to `AcquireSegmentAction` just for this. I was going to leave it the way it is.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For additional commands, e-mail: commits-help@druid.apache.org
Re: [PR] MSQ: Use VSF, give stages control over inputs. (druid)
Posted by "github-advanced-security[bot] (via GitHub)" <gi...@apache.org>.
github-advanced-security[bot] commented on code in PR #18873:
URL: https://github.com/apache/druid/pull/18873#discussion_r2667311941
##########
multi-stage-query/src/test/java/org/apache/druid/msq/input/RegularLoadableSegmentTest.java:
##########
@@ -0,0 +1,530 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.msq.input;
+
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonProperty;
+import com.fasterxml.jackson.annotation.JsonTypeName;
+import com.fasterxml.jackson.databind.InjectableValues;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.google.common.collect.ImmutableList;
+import com.google.common.util.concurrent.Futures;
+import com.google.common.util.concurrent.ListenableFuture;
+import com.google.common.util.concurrent.ListeningExecutorService;
+import com.google.common.util.concurrent.MoreExecutors;
+import org.apache.druid.client.coordinator.NoopCoordinatorClient;
+import org.apache.druid.common.guava.FutureUtils;
+import org.apache.druid.jackson.SegmentizerModule;
+import org.apache.druid.java.util.common.DateTimes;
+import org.apache.druid.java.util.common.FileUtils;
+import org.apache.druid.java.util.common.ISE;
+import org.apache.druid.java.util.common.Intervals;
+import org.apache.druid.java.util.common.Pair;
+import org.apache.druid.java.util.common.concurrent.Execs;
+import org.apache.druid.java.util.common.jackson.JacksonUtils;
+import org.apache.druid.java.util.emitter.EmittingLogger;
+import org.apache.druid.math.expr.ExprMacroTable;
+import org.apache.druid.query.expression.TestExprMacroTable;
+import org.apache.druid.segment.IndexIO;
+import org.apache.druid.segment.IndexSpec;
+import org.apache.druid.segment.PhysicalSegmentInspector;
+import org.apache.druid.segment.Segment;
+import org.apache.druid.segment.TestHelper;
+import org.apache.druid.segment.TestIndex;
+import org.apache.druid.segment.loading.AcquireSegmentAction;
+import org.apache.druid.segment.loading.AcquireSegmentResult;
+import org.apache.druid.segment.loading.LeastBytesUsedStorageLocationSelectorStrategy;
+import org.apache.druid.segment.loading.LoadSpec;
+import org.apache.druid.segment.loading.SegmentLoaderConfig;
+import org.apache.druid.segment.loading.SegmentLoadingException;
+import org.apache.druid.segment.loading.SegmentLocalCacheManager;
+import org.apache.druid.segment.loading.StorageLocation;
+import org.apache.druid.segment.loading.StorageLocationConfig;
+import org.apache.druid.server.SegmentManager;
+import org.apache.druid.server.metrics.NoopServiceEmitter;
+import org.apache.druid.testing.InitializedNullHandlingTest;
+import org.apache.druid.timeline.DataSegment;
+import org.apache.druid.timeline.partition.NumberedShardSpec;
+import org.apache.druid.utils.CompressionUtils;
+import org.joda.time.DateTime;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.io.TempDir;
+
+import java.io.File;
+import java.io.IOException;
+import java.nio.file.Path;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.Optional;
+import java.util.concurrent.TimeUnit;
+
+/**
+ * Tests for {@link RegularLoadableSegment}.
+ */
+class RegularLoadableSegmentTest extends InitializedNullHandlingTest
+{
+ private static final String DATASOURCE = "foo";
+ private static final int NUM_SEGMENTS = 10;
+ private static final int THREADS = 8;
+ private static File SEGMENT_ZIP_FILE;
+
+ @TempDir
+ public Path tempDir;
+
+ private List<DataSegment> segments;
+ private File cacheDir;
+ private File preLoadCacheDir;
+ private SegmentManager segmentManagerDynamic;
+ private SegmentManager segmentManagerPreLoad;
+ private ListeningExecutorService exec;
+
+ @BeforeAll
+ public static void setupStatic(@TempDir Path tempDir) throws IOException
+ {
+ EmittingLogger.registerEmitter(new NoopServiceEmitter());
+ final File segmentFile = TestIndex.persist(
+ TestIndex.getIncrementalTestIndex(),
+ IndexSpec.getDefault(),
+ tempDir.resolve("segment").toFile()
+ );
+ final File zipPath = tempDir.resolve("zip").toFile();
+ FileUtils.mkdirp(zipPath);
+ SEGMENT_ZIP_FILE = new File(zipPath, "index.zip");
+ CompressionUtils.zip(segmentFile, SEGMENT_ZIP_FILE);
+ }
+
+ @BeforeEach
+ public void setUp() throws Exception
+ {
+ final ObjectMapper jsonMapper = TestHelper.makeJsonMapper();
+ jsonMapper.registerSubtypes(TestLoadSpec.class);
+ jsonMapper.registerModule(new SegmentizerModule());
+ jsonMapper.setInjectableValues(
+ new InjectableValues.Std()
+ .addValue(ExprMacroTable.class.getName(), TestExprMacroTable.INSTANCE)
+ .addValue(ObjectMapper.class.getName(), jsonMapper)
+ .addValue(DataSegment.PruneSpecsHolder.class, DataSegment.PruneSpecsHolder.DEFAULT)
+ .addValue(IndexIO.class, TestIndex.INDEX_IO)
+ );
+
+ segments = new ArrayList<>();
+
+ for (int i = 0; i < NUM_SEGMENTS; i++) {
+ // Two segments per interval; helps verify that directory creation + deletion does not include races.
+ final DateTime startTime = DateTimes.of("2000").plusDays(i / 2);
+ final int partitionNum = i % 2;
+
+ segments.add(
+ DataSegment.builder()
Review Comment:
## Deprecated method or constructor invocation
Invoking [DataSegment.builder](1) should be avoided because it has been deprecated.
[Show more details](https://github.com/apache/druid/security/code-scanning/10632)
##########
multi-stage-query/src/main/java/org/apache/druid/msq/querykit/ReadableInputQueue.java:
##########
@@ -0,0 +1,386 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.msq.querykit;
+
+import com.google.common.collect.Sets;
+import com.google.common.util.concurrent.Futures;
+import com.google.common.util.concurrent.ListenableFuture;
+import com.google.errorprone.annotations.concurrent.GuardedBy;
+import org.apache.druid.common.guava.FutureUtils;
+import org.apache.druid.error.DruidException;
+import org.apache.druid.frame.channel.ReadableFrameChannel;
+import org.apache.druid.java.util.common.logger.Logger;
+import org.apache.druid.msq.exec.DataServerQueryHandler;
+import org.apache.druid.msq.exec.std.StandardPartitionReader;
+import org.apache.druid.msq.input.LoadableSegment;
+import org.apache.druid.msq.input.PhysicalInputSlice;
+import org.apache.druid.msq.input.stage.ReadablePartition;
+import org.apache.druid.msq.kernel.StageId;
+import org.apache.druid.msq.kernel.StagePartition;
+import org.apache.druid.segment.Segment;
+import org.apache.druid.segment.SegmentReference;
+import org.apache.druid.segment.loading.AcquireSegmentAction;
+import org.apache.druid.utils.CloseableUtils;
+
+import javax.annotation.Nullable;
+import java.io.Closeable;
+import java.io.IOException;
+import java.util.ArrayDeque;
+import java.util.ArrayList;
+import java.util.LinkedHashSet;
+import java.util.List;
+import java.util.Optional;
+import java.util.Queue;
+import java.util.Set;
+import java.util.concurrent.atomic.AtomicBoolean;
+
+/**
+ * Queue for returning {@link ReadableInput} from a list of {@link PhysicalInputSlice}.
+ *
+ * When closed, this object cancels all pending segment loads and releases all segments that have not yet been
+ * acquired by callers through {@link SegmentReferenceHolder#getSegmentReferenceOnce()}. Callers that have acquired
+ * segment references are responsible for closing those references, they will not be closed by this class.
+ */
+public class ReadableInputQueue implements Closeable
+{
+ private static final Logger log = new Logger(ReadableInputQueue.class);
+
+ /**
+ * Partitions to be read.
+ */
+ @GuardedBy("this")
+ private final Queue<ReadablePartition> readablePartitions = new ArrayDeque<>();
+
+ /**
+ * Segments to be loaded.
+ */
+ @GuardedBy("this")
+ private final Queue<LoadableSegment> loadableSegments = new ArrayDeque<>();
+
+ /**
+ * Realtime servers to be queried.
+ */
+ @GuardedBy("this")
+ private final Queue<DataServerQueryHandler> queryableServers = new ArrayDeque<>();
+
+ /**
+ * Segments currently being loaded.
+ */
+ @GuardedBy("this")
+ private final Set<AcquireSegmentAction> loadingSegments = new LinkedHashSet<>();
+
+ /**
+ * Segments that have been loaded. These are tracked here so we can close them if needed.
+ */
+ @GuardedBy("this")
+ private final Set<SegmentReferenceHolder> loadedSegments = new LinkedHashSet<>();
+
+ /**
+ * Futures that are sitting ready to be handed out by a call to {@link #nextInput()}.
+ */
+ @GuardedBy("this")
+ private final Set<ListenableFuture<ReadableInput>> pendingNextInputs = Sets.newIdentityHashSet();
+
+ private final String queryId;
+ private final StandardPartitionReader partitionReader;
+ private final int loadahead;
+ private final AtomicBoolean started = new AtomicBoolean(false);
+
+ public ReadableInputQueue(
+ final String queryId,
+ final StandardPartitionReader partitionReader,
+ final List<PhysicalInputSlice> slices,
+ final int loadahead
+ )
+ {
+ this.queryId = queryId;
+ this.partitionReader = partitionReader;
+ this.loadahead = loadahead;
+
+ for (final PhysicalInputSlice slice : slices) {
+ loadableSegments.addAll(slice.getLoadableSegments());
+ queryableServers.addAll(slice.getQueryableServers());
+ slice.getReadablePartitions().forEach(readablePartitions::add);
+ }
+ }
+
+ /**
+ * If this method has not yet been called, then:
+ * (1) transition all locally-cached segments out of {@link #loadableSegments}
+ * (2) start loading up to {@link #loadahead} additional segments for future calls to {@link #nextInput()}
+ * If this method has previously been called, subsequent calls do nothing.
+ * This is separated from the constructor because we don't want to acquire resources immediately on construction.
+ */
+ public void start()
+ {
+ if (started.compareAndSet(false, true)) {
+ // (1) acquire all locally-cached segments
+ synchronized (this) {
+ final List<LoadableSegment> toLoad = new ArrayList<>(); // Temporarily store all non-cached segments
+ LoadableSegment loadableSegment;
+ while ((loadableSegment = loadableSegments.poll()) != null) {
+ final Optional<Segment> cachedSegment = loadableSegment.acquireIfCached();
+ if (cachedSegment.isPresent()) {
+ final SegmentReferenceHolder holder = new SegmentReferenceHolder(
+ new SegmentReference(loadableSegment.descriptor(), cachedSegment, null),
+ loadableSegment.inputCounters(),
+ loadableSegment.description()
+ );
+ loadedSegments.add(holder);
+ pendingNextInputs.add(Futures.immediateFuture(ReadableInput.segment(holder)));
+ } else {
+ toLoad.add(loadableSegment);
+ }
+ }
+ loadableSegments.addAll(toLoad); // Put non-cached segments back into loadableSegments
+ }
+
+ // (2) start loading up to "loadahead" additional segments
+ for (int i = 0; i < loadahead; i++) {
+ if (!addLoadaheadFuture()) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * Returns the number of remaining inputs that can be returned by calls to {@link #nextInput()}.
+ */
+ public int remaining()
+ {
+ synchronized (this) {
+ return readablePartitions.size() + loadableSegments.size() + queryableServers.size() + pendingNextInputs.size();
+ }
+ }
+
+ /**
+ * Returns the next {@link ReadableInput}. The future resolves when the input is ready to read.
+ */
+ @Nullable
+ public ListenableFuture<ReadableInput> nextInput()
+ {
+ if (!started.get()) {
+ throw DruidException.defensive("Not started, must call start() first");
+ }
+
+ ListenableFuture<ReadableInput> future;
+
+ future = nextServerInput();
+ if (future != null) {
+ return future;
+ }
+
+ future = nextChannelInput();
+ if (future != null) {
+ return future;
+ }
+
+ future = nextSegmentInput();
+ if (future != null) {
+ return future;
+ }
+
+ return null;
+ }
+
+ /**
+ * Returns the next input from {@link #queryableServers}, if any. Returns null if none remain.
+ */
+ @Nullable
+ private ListenableFuture<ReadableInput> nextServerInput()
+ {
+ final DataServerQueryHandler handler;
+ synchronized (this) {
+ handler = queryableServers.poll();
+ }
+
+ if (handler == null) {
+ return null;
+ }
+
+ return Futures.immediateFuture(ReadableInput.dataServerQuery(handler));
+ }
+
+ /**
+ * Returns the next input from {@link #readablePartitions}, if any. Returns null if none remain.
+ */
+ @Nullable
+ private ListenableFuture<ReadableInput> nextChannelInput()
+ {
+ final ReadablePartition readablePartition;
+ synchronized (this) {
+ readablePartition = readablePartitions.poll();
+ }
+
+ if (readablePartition == null) {
+ return null;
+ }
+
+ ReadableFrameChannel channel = null;
+ try {
+ channel = partitionReader.openChannel(readablePartition);
+ return Futures.immediateFuture(
+ ReadableInput.channel(
+ channel,
+ partitionReader.frameReader(readablePartition.getStageNumber()),
+ new StagePartition(
+ new StageId(queryId, readablePartition.getStageNumber()),
+ readablePartition.getPartitionNumber()
+ )
+ )
+ );
+ }
+ catch (IOException e) {
+ throw CloseableUtils.closeAndWrapInCatch(e, channel);
+ }
+ }
+
+ /**
+ * Returns the next input from {@link #loadableSegments}, if any. Returns null if none remain.
+ */
+ @Nullable
+ private ListenableFuture<ReadableInput> nextSegmentInput()
+ {
+ // Pick a loadahead future, preferring ones that are already loaded.
+ ListenableFuture<ReadableInput> selectedLoadaheadFuture = null;
+ synchronized (this) {
+ for (ListenableFuture<ReadableInput> f : pendingNextInputs) {
+ if (selectedLoadaheadFuture == null || f.isDone()) {
+ selectedLoadaheadFuture = f;
+ if (f.isDone()) {
+ break;
+ }
+ }
+ }
+
+ if (selectedLoadaheadFuture != null) {
+ pendingNextInputs.remove(selectedLoadaheadFuture);
+ if (pendingNextInputs.size() < loadahead) {
+ addLoadaheadFuture(); // Replace the one we just took out.
+ }
+ return selectedLoadaheadFuture;
+ }
+ }
+
+ return loadNextSegment();
+ }
+
+ /**
+ * Load the next segment from {@link #loadableSegments} and return a future to its reference. Returns null
+ * if {@link #loadableSegments} is empty.
+ */
+ @Nullable
+ private ListenableFuture<ReadableInput> loadNextSegment()
+ {
+ synchronized (this) {
+ final LoadableSegment nextLoadableSegment = loadableSegments.poll();
+ if (nextLoadableSegment == null) {
+ return null;
+ }
+
+ final AcquireSegmentAction acquireSegmentAction = nextLoadableSegment.acquire();
+ loadingSegments.add(acquireSegmentAction);
+ return FutureUtils.transform(
+ acquireSegmentAction.getSegmentFuture(),
+ segment -> {
+ synchronized (ReadableInputQueue.this) {
+ // Transfer segment from "loadingSegments" to "loadedSegments" and return a reference to it.
+ if (loadingSegments.remove(acquireSegmentAction)) {
+ try {
+ final SegmentReferenceHolder referenceHolder = new SegmentReferenceHolder(
+ new SegmentReference(
+ nextLoadableSegment.descriptor(),
+ segment.getReferenceProvider().acquireReference(),
+ acquireSegmentAction // Release the hold when the SegmentReference is closed.
+ ),
+ nextLoadableSegment.inputCounters(),
+ nextLoadableSegment.description()
+ );
+ loadedSegments.add(referenceHolder);
+ return ReadableInput.segment(referenceHolder);
+ }
+ catch (Throwable e) {
+ // Javadoc for segment.acquireReference() suggests it can throw exceptions; handle that here
+ // by closing the original AcquireSegmentAction.
+ throw CloseableUtils.closeAndWrapInCatch(e, acquireSegmentAction);
+ }
+ } else {
+ throw DruidException.defensive(
+ "Segment[%s] removed from loadingSegments before loading complete. It is possible that close() "
+ + "was called with futures in flight.",
+ nextLoadableSegment.description()
+ );
+ }
+ }
+ }
+ );
+ }
+ }
+
+ /**
+ * Calls {@link #nextSegmentInput()} and adds the future to {@link #pendingNextInputs}. Returns whether a future
+ * was added.
+ */
+ private boolean addLoadaheadFuture()
+ {
+ final ListenableFuture<ReadableInput> nextFuture = loadNextSegment();
+ if (nextFuture != null) {
+ synchronized (this) {
+ pendingNextInputs.add(nextFuture);
+ }
+ return true;
+ } else {
+ return false;
+ }
+ }
+
+ @Override
+ public void close()
+ {
+ synchronized (this) {
+ readablePartitions.clear();
+ queryableServers.clear();
+ loadableSegments.clear();
+
+ // Cancel all pending segment loads.
+ for (AcquireSegmentAction acquireSegmentAction : loadingSegments) {
+ CloseableUtils.closeAndSuppressExceptions(
+ acquireSegmentAction,
+ e -> log.warn(e, "Failed to close loadingSegment[%s]", acquireSegmentAction)
Review Comment:
## Use of default toString()
Default toString(): AcquireSegmentAction inherits toString() from Object, and so is not suitable for printing.
[Show more details](https://github.com/apache/druid/security/code-scanning/10631)
##########
server/src/test/java/org/apache/druid/test/utils/TestSegmentManager.java:
##########
@@ -0,0 +1,126 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.test.utils;
+
+import org.apache.druid.java.util.common.ISE;
+import org.apache.druid.segment.Segment;
+import org.apache.druid.segment.loading.SegmentLoadingException;
+import org.apache.druid.server.SegmentManager;
+import org.apache.druid.timeline.DataSegment;
+import org.apache.druid.timeline.SegmentId;
+import org.apache.druid.timeline.partition.NumberedShardSpec;
+
+import javax.annotation.Nullable;
+import java.io.IOException;
+import java.util.Collection;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentMap;
+
+/**
+ * Wrapper around {@link SegmentManager} using {@link TestSegmentCacheManager} that helps with adding in-memory
+ * test segments.
+ */
+public class TestSegmentManager
+{
+ private final SegmentManager segmentManager;
+ private final TestSegmentCacheManager cacheManager;
+ private final ConcurrentMap<SegmentId, DataSegment> testGeneratedDataSegments = new ConcurrentHashMap<>();
+ private final ConcurrentMap<SegmentId, DataSegment> addedSegments = new ConcurrentHashMap<>();
+
+ public TestSegmentManager()
+ {
+ this.cacheManager = new TestSegmentCacheManager();
+ this.segmentManager = new SegmentManager(cacheManager);
+ }
+
+ /**
+ * Creates a minimal DataSegment for a test Segment.
+ */
+ public static DataSegment createDataSegmentForTest(final SegmentId segmentId)
+ {
+ return DataSegment.builder()
Review Comment:
## Deprecated method or constructor invocation
Invoking [DataSegment.builder](1) should be avoided because it has been deprecated.
[Show more details](https://github.com/apache/druid/security/code-scanning/10633)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For additional commands, e-mail: commits-help@druid.apache.org
Re: [PR] MSQ: Use VSF, give stages control over inputs. (druid)
Posted by "gianm (via GitHub)" <gi...@apache.org>.
gianm commented on code in PR #18873:
URL: https://github.com/apache/druid/pull/18873#discussion_r2677677966
##########
multi-stage-query/src/main/java/org/apache/druid/msq/querykit/ReadableInputQueue.java:
##########
@@ -0,0 +1,366 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.msq.querykit;
+
+import com.google.common.collect.Sets;
+import com.google.common.util.concurrent.Futures;
+import com.google.common.util.concurrent.ListenableFuture;
+import com.google.errorprone.annotations.concurrent.GuardedBy;
+import org.apache.druid.common.guava.FutureUtils;
+import org.apache.druid.error.DruidException;
+import org.apache.druid.frame.channel.ReadableFrameChannel;
+import org.apache.druid.java.util.common.logger.Logger;
+import org.apache.druid.msq.exec.DataServerQueryHandler;
+import org.apache.druid.msq.exec.std.StandardPartitionReader;
+import org.apache.druid.msq.input.LoadableSegment;
+import org.apache.druid.msq.input.PhysicalInputSlice;
+import org.apache.druid.msq.input.stage.ReadablePartition;
+import org.apache.druid.msq.kernel.StageId;
+import org.apache.druid.msq.kernel.StagePartition;
+import org.apache.druid.segment.SegmentReference;
+import org.apache.druid.segment.loading.AcquireSegmentAction;
+import org.apache.druid.utils.CloseableUtils;
+
+import javax.annotation.Nullable;
+import java.io.Closeable;
+import java.io.IOException;
+import java.util.ArrayDeque;
+import java.util.LinkedHashSet;
+import java.util.List;
+import java.util.Queue;
+import java.util.Set;
+import java.util.concurrent.atomic.AtomicBoolean;
+
+/**
+ * Queue for returning {@link ReadableInput} from a list of {@link PhysicalInputSlice}.
+ *
+ * When closed, this object cancels all pending segment loads and releases all segments that have not yet been
+ * acquired by callers through {@link SegmentReferenceHolder#getSegmentReferenceOnce()}. Callers that have acquired
+ * segment references are responsible for closing those references, they will not be closed by this class.
+ */
+public class ReadableInputQueue implements Closeable
+{
+ private static final Logger log = new Logger(ReadableInputQueue.class);
+
+ /**
+ * Partitions to be read.
+ */
+ @GuardedBy("this")
+ private final Queue<ReadablePartition> readablePartitions = new ArrayDeque<>();
+
+ /**
+ * Segments to be loaded. Some may already be cached locally.
+ */
+ @GuardedBy("this")
+ private final Queue<LoadableSegment> loadableSegments = new ArrayDeque<>();
+
+ /**
+ * Realtime servers to be queried.
+ */
+ @GuardedBy("this")
+ private final Queue<DataServerQueryHandler> queryableServers = new ArrayDeque<>();
+
+ /**
+ * Segments currently being loaded.
+ */
+ @GuardedBy("this")
+ private final Set<AcquireSegmentAction> loadingSegments = new LinkedHashSet<>();
+
+ /**
+ * Segments that have been loaded.
+ */
+ @GuardedBy("this")
+ private final Set<SegmentReferenceHolder> loadedSegments = new LinkedHashSet<>();
+
+ /**
+ * Futures that are generated by loadahead, prior to a call to {@link #nextInput()} being made.
+ */
+ @GuardedBy("this")
+ private final Set<ListenableFuture<ReadableInput>> loadaheadFutures = Sets.newIdentityHashSet();
+
+ private final String queryId;
+ private final StandardPartitionReader partitionReader;
+ private final int loadahead;
+ private final AtomicBoolean loadaheadStarted = new AtomicBoolean(false);
+
+ public ReadableInputQueue(
+ final String queryId,
+ final StandardPartitionReader partitionReader,
+ final List<PhysicalInputSlice> slices,
+ final int loadahead
+ )
+ {
+ this.queryId = queryId;
+ this.partitionReader = partitionReader;
+ this.loadahead = loadahead;
+
+ for (final PhysicalInputSlice slice : slices) {
+ // First add all locally-cached segments, before any non-locally-cached segments. This promotes processing
+ // the locally-cached ones first. It also improves efficiency of loadahead, since once we get to the
+ // set of non-locally-cached segments, the "next" one is always one that is worth loading ahead.
+ for (final LoadableSegment loadableSegment : slice.getLoadableSegments()) {
+ if (loadableSegment.isLikelyCached()) {
+ loadableSegments.add(loadableSegment);
+ }
+ }
+
+ for (final LoadableSegment loadableSegment : slice.getLoadableSegments()) {
+ if (!loadableSegment.isLikelyCached()) {
+ loadableSegments.add(loadableSegment);
+ }
+ }
+
+ queryableServers.addAll(slice.getQueryableServers());
+ slice.getReadablePartitions().forEach(readablePartitions::add);
+ }
+ }
+
+ /**
+ * Starts loading up to {@link #loadahead} segments for future calls to {@link #nextInput()}, if this method
+ * has not yet been called. If it has previously been called, subsequent calls do nothing.
+ */
+ public void startLoadaheadIfNeeded()
+ {
+ if (loadaheadStarted.compareAndSet(false, true)) {
+ for (int i = 0; i < loadahead; i++) {
+ if (!addLoadaheadFuture()) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * Returns the number of remaining inputs that can be returned by calls to {@link #nextInput()}.
+ */
+ public int remaining()
+ {
+ synchronized (this) {
+ return readablePartitions.size() + loadableSegments.size() + queryableServers.size() + loadaheadFutures.size();
+ }
+ }
+
+ /**
+ * Returns the next {@link ReadableInput}. The future resolves when the input is ready to read.
+ */
+ @Nullable
+ public ListenableFuture<ReadableInput> nextInput()
+ {
+ ListenableFuture<ReadableInput> future;
+
+ future = nextServerInput();
+ if (future != null) {
+ return future;
+ }
+
+ future = nextChannelInput();
+ if (future != null) {
+ return future;
+ }
+
+ future = nextSegmentInput();
+ if (future != null) {
+ return future;
+ }
+
+ return null;
+ }
+
+ /**
+ * Returns the next input from {@link #queryableServers}, if any. Returns null if none remain.
+ */
+ @Nullable
+ private ListenableFuture<ReadableInput> nextServerInput()
+ {
+ final DataServerQueryHandler handler;
+ synchronized (this) {
+ handler = queryableServers.poll();
+ }
+
+ if (handler == null) {
+ return null;
+ }
+
+ return Futures.immediateFuture(ReadableInput.dataServerQuery(handler));
+ }
+
+ /**
+ * Returns the next input from {@link #readablePartitions}, if any. Returns null if none remain.
+ */
+ @Nullable
+ private ListenableFuture<ReadableInput> nextChannelInput()
+ {
+ final ReadablePartition readablePartition;
+ synchronized (this) {
+ readablePartition = readablePartitions.poll();
+ }
+
+ if (readablePartition == null) {
+ return null;
+ }
+
+ ReadableFrameChannel channel = null;
+ try {
+ channel = partitionReader.openChannel(readablePartition);
+ return Futures.immediateFuture(
+ ReadableInput.channel(
+ channel,
+ partitionReader.frameReader(readablePartition.getStageNumber()),
+ new StagePartition(
+ new StageId(queryId, readablePartition.getStageNumber()),
+ readablePartition.getPartitionNumber()
+ )
+ )
+ );
+ }
+ catch (IOException e) {
+ throw CloseableUtils.closeAndWrapInCatch(e, channel);
+ }
+ }
+
+ /**
+ * Returns the next input from {@link #loadableSegments}, if any. Returns null if none remain.
+ */
+ @Nullable
+ private ListenableFuture<ReadableInput> nextSegmentInput()
+ {
+ // Pick a loadahead future, preferring ones that are already loaded.
+ ListenableFuture<ReadableInput> selectedLoadaheadFuture = null;
+ synchronized (this) {
+ for (ListenableFuture<ReadableInput> f : loadaheadFutures) {
+ if (selectedLoadaheadFuture == null || f.isDone()) {
+ selectedLoadaheadFuture = f;
+ if (f.isDone()) {
+ break;
+ }
+ }
+ }
+
+ if (selectedLoadaheadFuture != null) {
+ loadaheadFutures.remove(selectedLoadaheadFuture);
+ addLoadaheadFuture(); // Replace the one we just took out.
+ return selectedLoadaheadFuture;
+ }
+ }
+
+ return loadNextSegment();
+ }
+
+ /**
+ * Load the next segment from {@link #loadableSegments} and return a future to its reference. Returns null
+ * if {@link #loadableSegments} is empty.
+ */
+ @Nullable
+ private ListenableFuture<ReadableInput> loadNextSegment()
+ {
+ synchronized (this) {
+ final LoadableSegment nextLoadableSegment = loadableSegments.poll();
+ if (nextLoadableSegment == null) {
+ return null;
+ }
+
+ final AcquireSegmentAction acquireSegmentAction = nextLoadableSegment.acquire();
+ loadingSegments.add(acquireSegmentAction);
+ return FutureUtils.transform(
+ acquireSegmentAction.getSegmentFuture(),
+ segment -> {
+ synchronized (ReadableInputQueue.this) {
+ // Transfer segment from "loadingSegments" to "loadedSegments" and return a reference to it.
+ if (loadingSegments.remove(acquireSegmentAction)) {
+ try {
+ final SegmentReferenceHolder referenceHolder = new SegmentReferenceHolder(
+ new SegmentReference(
+ nextLoadableSegment.descriptor(),
+ segment.getReferenceProvider().acquireReference(),
+ acquireSegmentAction // Release the hold when the SegmentReference is closed.
+ ),
+ nextLoadableSegment.inputCounters(),
+ nextLoadableSegment.description()
+ );
+ loadedSegments.add(referenceHolder);
+ return ReadableInput.segment(referenceHolder);
+ }
+ catch (Throwable e) {
+ // Javadoc for segment.acquireReference() suggests it can throw exceptions; handle that here
+ // by closing the original AcquireSegmentAction.
+ throw CloseableUtils.closeAndWrapInCatch(e, acquireSegmentAction);
+ }
+ } else {
+ throw DruidException.defensive(
+ "Segment[%s] removed from loadingSegments before loading complete. It is possible that close() "
+ + "was called with futures in flight.",
+ nextLoadableSegment.descriptor()
+ );
+ }
+ }
+ }
+ );
+ }
+ }
+
+ /**
+ * Calls {@link #nextSegmentInput()} and adds the future to {@link #loadaheadFutures}. Returns whether a future
+ * was added.
+ */
+ private boolean addLoadaheadFuture()
+ {
+ final ListenableFuture<ReadableInput> nextFuture = loadNextSegment();
+ if (nextFuture != null) {
+ synchronized (this) {
+ loadaheadFutures.add(nextFuture);
+ }
+ return true;
+ } else {
+ return false;
+ }
+ }
+
+ @Override
+ public void close()
+ {
+ synchronized (this) {
+ readablePartitions.clear();
+ queryableServers.clear();
+ loadableSegments.clear();
+
+ // Cancel all pending segment loads.
+ for (AcquireSegmentAction acquireSegmentAction : loadingSegments) {
+ CloseableUtils.closeAndSuppressExceptions(
+ acquireSegmentAction,
+ e -> log.warn(e, "Failed to close loadingSegment[%s]", acquireSegmentAction)
Review Comment:
I'll add a comment about this though.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For additional commands, e-mail: commits-help@druid.apache.org
Re: [PR] MSQ: Use VSF, give stages control over inputs. (druid)
Posted by "clintropolis (via GitHub)" <gi...@apache.org>.
clintropolis commented on code in PR #18873:
URL: https://github.com/apache/druid/pull/18873#discussion_r2667697963
##########
indexing-service/src/main/java/org/apache/druid/indexing/common/SegmentCacheManagerFactory.java:
##########
@@ -52,11 +52,12 @@ public SegmentCacheManagerFactory(
this.jsonMapper = mapper;
}
- public SegmentCacheManager manufacturate(File storageDir)
+ public SegmentCacheManager manufacturate(File storageDir, boolean virtualStorage)
{
- final SegmentLoaderConfig loaderConfig = new SegmentLoaderConfig().withLocations(
- Collections.singletonList(new StorageLocationConfig(storageDir, null, null))
- );
+ final SegmentLoaderConfig loaderConfig =
+ new SegmentLoaderConfig()
+ .setLocations(Collections.singletonList(new StorageLocationConfig(storageDir, null, null)))
Review Comment:
i know this isn't new, but should this have a size so like... there is a guard rail to not just wreck things if something tries to load too much? There is `tmpStorageBytesPerTask` already on the task config, but that is used for tmp sorter output it looks like so.. maybe those files should be managed with StorageLocation too. This doesn't need to be done in this PR, just thinking
##########
multi-stage-query/src/main/java/org/apache/druid/msq/input/external/ExternalInputSliceReader.java:
##########
@@ -66,48 +65,63 @@
public class ExternalInputSliceReader implements InputSliceReader
{
public static final String SEGMENT_ID = "__external";
+ public static final SegmentDescriptor SEGMENT_DESCRIPTOR = SegmentId.dummy(SEGMENT_ID).toDescriptor();
private final File temporaryDirectory;
public ExternalInputSliceReader(final File temporaryDirectory)
{
this.temporaryDirectory = temporaryDirectory;
}
- @Override
- public int numReadableInputs(InputSlice slice)
+ public static boolean isFileBasedInputSource(final InputSource inputSource)
{
- final ExternalInputSlice externalInputSlice = (ExternalInputSlice) slice;
- return externalInputSlice.getInputSources().size();
+ return !(inputSource instanceof NilInputSource) && !(inputSource instanceof InlineInputSource);
}
@Override
- public ReadableInputs attach(
+ public PhysicalInputSlice attach(
final int inputNumber,
final InputSlice slice,
final CounterTracker counters,
final Consumer<Throwable> warningPublisher
)
{
final ExternalInputSlice externalInputSlice = (ExternalInputSlice) slice;
+ final ChannelCounters inputCounters = counters.channel(CounterNames.inputChannel(inputNumber))
+ .setTotalFiles(slice.fileCount());
+ final List<LoadableSegment> loadableSegments = new ArrayList<>();
- return ReadableInputs.segments(
- () -> Iterators.transform(
- inputSourceSegmentIterator(
- externalInputSlice.getInputSources(),
- externalInputSlice.getInputFormat(),
- externalInputSlice.getSignature(),
- new File(temporaryDirectory, String.valueOf(inputNumber)),
- counters.channel(CounterNames.inputChannel(inputNumber)).setTotalFiles(slice.fileCount()),
- counters.warnings(),
- warningPublisher
- ),
- ReadableInput::segment
- )
- );
+ for (final InputSource inputSource : externalInputSlice.getInputSources()) {
+ // The LoadableSegment generated here does not acquire a real hold, and ends up loading the external data in a
+ // processing thread (when the cursor is created). Ideally, this would be better integrated with the virtual
+ // storage system, giving us storage holds and the ability to load data outside of a processing thread.
Review Comment:
:+1: it feels like this should be possible, `StorageLocation` is pretty generic, but agree we can worry about this later
##########
processing/src/main/java/org/apache/druid/segment/SegmentReference.java:
##########
@@ -52,23 +57,60 @@ public SegmentReference(
)
{
this.segmentDescriptor = segmentDescriptor;
- closer.register(cleanupHold);
- this.segmentReference = segmentReference.map(closer::register);
+ this.segmentReference = segmentReference;
+ this.cleanupHold = cleanupHold;
}
public SegmentDescriptor getSegmentDescriptor()
{
return segmentDescriptor;
}
+ /**
+ * Returns the actual segment. Do not close the Segment when you are done with it, only close the SegmentReference.
+ */
public Optional<Segment> getSegmentReference()
{
return segmentReference;
}
+ /**
+ * Maps the wrapped segment and returns a reference to it. Closes the reference if the {@link SegmentMapFunction}
+ * throws an exception. Regardless of success or failure of this method, the old reference should be discarded.
+ * Do not call {@link #close()} on the old reference, only call it on the new one.
+ */
+ public SegmentReference map(final SegmentMapFunction segmentMapFn)
Review Comment:
could you update the javadoc at the top of this class to indicate that callers can call this if they don't have a mapFn to apply to the segment at the time of creation of the `SegmentReference`?
##########
multi-stage-query/src/main/java/org/apache/druid/msq/input/PhysicalInputSlice.java:
##########
@@ -0,0 +1,63 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.msq.input;
+
+import com.google.common.base.Preconditions;
+import org.apache.druid.msq.exec.DataServerQueryHandler;
+import org.apache.druid.msq.input.stage.ReadablePartitions;
+
+import java.util.List;
+
+/**
+ * An {@link InputSlice} that has been prepared for reading by an {@link InputSliceReader}. Nothing contained in
+ * this class references open resources, so this class is not closeable and does not need to be closed.
+ */
+public class PhysicalInputSlice
Review Comment:
what does 'physical' mean here? obligatory 'naming is hard' comment. Is it basically 'slice that gets processed to kick everything off by a leaf processor'? If so, should it be `LeafInputSlice` to match other terminology? though i guess i could see how that could be confusing too
##########
multi-stage-query/src/main/java/org/apache/druid/msq/querykit/ReadableInputQueue.java:
##########
@@ -0,0 +1,366 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.msq.querykit;
+
+import com.google.common.collect.Sets;
+import com.google.common.util.concurrent.Futures;
+import com.google.common.util.concurrent.ListenableFuture;
+import com.google.errorprone.annotations.concurrent.GuardedBy;
+import org.apache.druid.common.guava.FutureUtils;
+import org.apache.druid.error.DruidException;
+import org.apache.druid.frame.channel.ReadableFrameChannel;
+import org.apache.druid.java.util.common.logger.Logger;
+import org.apache.druid.msq.exec.DataServerQueryHandler;
+import org.apache.druid.msq.exec.std.StandardPartitionReader;
+import org.apache.druid.msq.input.LoadableSegment;
+import org.apache.druid.msq.input.PhysicalInputSlice;
+import org.apache.druid.msq.input.stage.ReadablePartition;
+import org.apache.druid.msq.kernel.StageId;
+import org.apache.druid.msq.kernel.StagePartition;
+import org.apache.druid.segment.SegmentReference;
+import org.apache.druid.segment.loading.AcquireSegmentAction;
+import org.apache.druid.utils.CloseableUtils;
+
+import javax.annotation.Nullable;
+import java.io.Closeable;
+import java.io.IOException;
+import java.util.ArrayDeque;
+import java.util.LinkedHashSet;
+import java.util.List;
+import java.util.Queue;
+import java.util.Set;
+import java.util.concurrent.atomic.AtomicBoolean;
+
+/**
+ * Queue for returning {@link ReadableInput} from a list of {@link PhysicalInputSlice}.
+ *
+ * When closed, this object cancels all pending segment loads and releases all segments that have not yet been
+ * acquired by callers through {@link SegmentReferenceHolder#getSegmentReferenceOnce()}. Callers that have acquired
+ * segment references are responsible for closing those references, they will not be closed by this class.
+ */
+public class ReadableInputQueue implements Closeable
+{
+ private static final Logger log = new Logger(ReadableInputQueue.class);
+
+ /**
+ * Partitions to be read.
+ */
+ @GuardedBy("this")
+ private final Queue<ReadablePartition> readablePartitions = new ArrayDeque<>();
+
+ /**
+ * Segments to be loaded. Some may already be cached locally.
+ */
+ @GuardedBy("this")
+ private final Queue<LoadableSegment> loadableSegments = new ArrayDeque<>();
+
+ /**
+ * Realtime servers to be queried.
+ */
+ @GuardedBy("this")
+ private final Queue<DataServerQueryHandler> queryableServers = new ArrayDeque<>();
+
+ /**
+ * Segments currently being loaded.
+ */
+ @GuardedBy("this")
+ private final Set<AcquireSegmentAction> loadingSegments = new LinkedHashSet<>();
+
+ /**
+ * Segments that have been loaded.
+ */
+ @GuardedBy("this")
+ private final Set<SegmentReferenceHolder> loadedSegments = new LinkedHashSet<>();
+
+ /**
+ * Futures that are generated by loadahead, prior to a call to {@link #nextInput()} being made.
+ */
+ @GuardedBy("this")
+ private final Set<ListenableFuture<ReadableInput>> loadaheadFutures = Sets.newIdentityHashSet();
+
+ private final String queryId;
+ private final StandardPartitionReader partitionReader;
+ private final int loadahead;
+ private final AtomicBoolean loadaheadStarted = new AtomicBoolean(false);
+
+ public ReadableInputQueue(
+ final String queryId,
+ final StandardPartitionReader partitionReader,
+ final List<PhysicalInputSlice> slices,
+ final int loadahead
+ )
+ {
+ this.queryId = queryId;
+ this.partitionReader = partitionReader;
+ this.loadahead = loadahead;
+
+ for (final PhysicalInputSlice slice : slices) {
+ // First add all locally-cached segments, before any non-locally-cached segments. This promotes processing
+ // the locally-cached ones first. It also improves efficiency of loadahead, since once we get to the
+ // set of non-locally-cached segments, the "next" one is always one that is worth loading ahead.
+ for (final LoadableSegment loadableSegment : slice.getLoadableSegments()) {
+ if (loadableSegment.isLikelyCached()) {
+ loadableSegments.add(loadableSegment);
+ }
+ }
+
+ for (final LoadableSegment loadableSegment : slice.getLoadableSegments()) {
+ if (!loadableSegment.isLikelyCached()) {
+ loadableSegments.add(loadableSegment);
+ }
+ }
+
+ queryableServers.addAll(slice.getQueryableServers());
+ slice.getReadablePartitions().forEach(readablePartitions::add);
+ }
+ }
+
+ /**
+ * Starts loading up to {@link #loadahead} segments for future calls to {@link #nextInput()}, if this method
+ * has not yet been called. If it has previously been called, subsequent calls do nothing.
+ */
+ public void startLoadaheadIfNeeded()
+ {
+ if (loadaheadStarted.compareAndSet(false, true)) {
+ for (int i = 0; i < loadahead; i++) {
+ if (!addLoadaheadFuture()) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * Returns the number of remaining inputs that can be returned by calls to {@link #nextInput()}.
+ */
+ public int remaining()
+ {
+ synchronized (this) {
+ return readablePartitions.size() + loadableSegments.size() + queryableServers.size() + loadaheadFutures.size();
+ }
+ }
+
+ /**
+ * Returns the next {@link ReadableInput}. The future resolves when the input is ready to read.
+ */
+ @Nullable
+ public ListenableFuture<ReadableInput> nextInput()
+ {
+ ListenableFuture<ReadableInput> future;
+
+ future = nextServerInput();
+ if (future != null) {
+ return future;
+ }
+
+ future = nextChannelInput();
+ if (future != null) {
+ return future;
+ }
+
+ future = nextSegmentInput();
+ if (future != null) {
+ return future;
+ }
+
+ return null;
+ }
+
+ /**
+ * Returns the next input from {@link #queryableServers}, if any. Returns null if none remain.
+ */
+ @Nullable
+ private ListenableFuture<ReadableInput> nextServerInput()
+ {
+ final DataServerQueryHandler handler;
+ synchronized (this) {
+ handler = queryableServers.poll();
+ }
+
+ if (handler == null) {
+ return null;
+ }
+
+ return Futures.immediateFuture(ReadableInput.dataServerQuery(handler));
+ }
+
+ /**
+ * Returns the next input from {@link #readablePartitions}, if any. Returns null if none remain.
+ */
+ @Nullable
+ private ListenableFuture<ReadableInput> nextChannelInput()
+ {
+ final ReadablePartition readablePartition;
+ synchronized (this) {
+ readablePartition = readablePartitions.poll();
+ }
+
+ if (readablePartition == null) {
+ return null;
+ }
+
+ ReadableFrameChannel channel = null;
+ try {
+ channel = partitionReader.openChannel(readablePartition);
+ return Futures.immediateFuture(
+ ReadableInput.channel(
+ channel,
+ partitionReader.frameReader(readablePartition.getStageNumber()),
+ new StagePartition(
+ new StageId(queryId, readablePartition.getStageNumber()),
+ readablePartition.getPartitionNumber()
+ )
+ )
+ );
+ }
+ catch (IOException e) {
+ throw CloseableUtils.closeAndWrapInCatch(e, channel);
+ }
+ }
+
+ /**
+ * Returns the next input from {@link #loadableSegments}, if any. Returns null if none remain.
+ */
+ @Nullable
+ private ListenableFuture<ReadableInput> nextSegmentInput()
+ {
+ // Pick a loadahead future, preferring ones that are already loaded.
+ ListenableFuture<ReadableInput> selectedLoadaheadFuture = null;
+ synchronized (this) {
+ for (ListenableFuture<ReadableInput> f : loadaheadFutures) {
+ if (selectedLoadaheadFuture == null || f.isDone()) {
+ selectedLoadaheadFuture = f;
+ if (f.isDone()) {
+ break;
+ }
+ }
+ }
+
+ if (selectedLoadaheadFuture != null) {
+ loadaheadFutures.remove(selectedLoadaheadFuture);
+ addLoadaheadFuture(); // Replace the one we just took out.
+ return selectedLoadaheadFuture;
+ }
+ }
+
+ return loadNextSegment();
+ }
+
+ /**
+ * Load the next segment from {@link #loadableSegments} and return a future to its reference. Returns null
+ * if {@link #loadableSegments} is empty.
+ */
+ @Nullable
+ private ListenableFuture<ReadableInput> loadNextSegment()
+ {
+ synchronized (this) {
+ final LoadableSegment nextLoadableSegment = loadableSegments.poll();
+ if (nextLoadableSegment == null) {
+ return null;
+ }
+
+ final AcquireSegmentAction acquireSegmentAction = nextLoadableSegment.acquire();
+ loadingSegments.add(acquireSegmentAction);
+ return FutureUtils.transform(
+ acquireSegmentAction.getSegmentFuture(),
+ segment -> {
+ synchronized (ReadableInputQueue.this) {
+ // Transfer segment from "loadingSegments" to "loadedSegments" and return a reference to it.
+ if (loadingSegments.remove(acquireSegmentAction)) {
+ try {
+ final SegmentReferenceHolder referenceHolder = new SegmentReferenceHolder(
+ new SegmentReference(
+ nextLoadableSegment.descriptor(),
+ segment.getReferenceProvider().acquireReference(),
+ acquireSegmentAction // Release the hold when the SegmentReference is closed.
+ ),
+ nextLoadableSegment.inputCounters(),
+ nextLoadableSegment.description()
+ );
+ loadedSegments.add(referenceHolder);
+ return ReadableInput.segment(referenceHolder);
+ }
+ catch (Throwable e) {
+ // Javadoc for segment.acquireReference() suggests it can throw exceptions; handle that here
+ // by closing the original AcquireSegmentAction.
+ throw CloseableUtils.closeAndWrapInCatch(e, acquireSegmentAction);
+ }
+ } else {
+ throw DruidException.defensive(
+ "Segment[%s] removed from loadingSegments before loading complete. It is possible that close() "
+ + "was called with futures in flight.",
+ nextLoadableSegment.descriptor()
+ );
+ }
+ }
+ }
+ );
+ }
+ }
+
+ /**
+ * Calls {@link #nextSegmentInput()} and adds the future to {@link #loadaheadFutures}. Returns whether a future
+ * was added.
+ */
+ private boolean addLoadaheadFuture()
+ {
+ final ListenableFuture<ReadableInput> nextFuture = loadNextSegment();
+ if (nextFuture != null) {
+ synchronized (this) {
+ loadaheadFutures.add(nextFuture);
+ }
+ return true;
+ } else {
+ return false;
+ }
+ }
+
+ @Override
+ public void close()
+ {
+ synchronized (this) {
+ readablePartitions.clear();
+ queryableServers.clear();
+ loadableSegments.clear();
+
+ // Cancel all pending segment loads.
+ for (AcquireSegmentAction acquireSegmentAction : loadingSegments) {
+ CloseableUtils.closeAndSuppressExceptions(
+ acquireSegmentAction,
+ e -> log.warn(e, "Failed to close loadingSegment[%s]", acquireSegmentAction)
Review Comment:
this looks legit, maybe we should decorate `AcquireSegmentAction` to have a segmentId or something and implement toString
##########
multi-stage-query/src/main/java/org/apache/druid/msq/input/RegularLoadableSegment.java:
##########
@@ -0,0 +1,236 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.msq.input;
+
+import com.google.common.base.Supplier;
+import com.google.common.base.Suppliers;
+import com.google.common.util.concurrent.Futures;
+import com.google.common.util.concurrent.ListenableFuture;
+import com.google.errorprone.annotations.concurrent.GuardedBy;
+import org.apache.druid.client.coordinator.CoordinatorClient;
+import org.apache.druid.common.guava.FutureUtils;
+import org.apache.druid.error.DruidException;
+import org.apache.druid.java.util.common.io.Closer;
+import org.apache.druid.msq.counters.ChannelCounters;
+import org.apache.druid.query.SegmentDescriptor;
+import org.apache.druid.query.TableDataSource;
+import org.apache.druid.segment.Segment;
+import org.apache.druid.segment.loading.AcquireSegmentAction;
+import org.apache.druid.server.SegmentManager;
+import org.apache.druid.timeline.DataSegment;
+import org.apache.druid.timeline.SegmentId;
+import org.apache.druid.timeline.VersionedIntervalTimeline;
+import org.apache.druid.timeline.partition.PartitionChunk;
+
+import javax.annotation.Nullable;
+import java.util.Optional;
+
+/**
+ * Implementation of {@link LoadableSegment} for regular Druid segments loaded via {@link SegmentManager}.
+ * Created by {@link org.apache.druid.msq.input.table.SegmentsInputSliceReader}.
+ */
+public class RegularLoadableSegment implements LoadableSegment
+{
+ private final SegmentManager segmentManager;
+ private final SegmentId segmentId;
+ private final SegmentDescriptor descriptor;
+ @Nullable
+ private final ChannelCounters inputCounters;
+ @Nullable
+ private final CoordinatorClient coordinatorClient;
+ private final boolean isReindex;
+
+ @GuardedBy("this")
+ private boolean acquired;
+
+ /**
+ * Cached DataSegment from local timeline, if available. Null if not in local timeline or if isReindex is true.
+ */
+ @Nullable
+ private final DataSegment cachedDataSegment;
+
+ /**
+ * Memoized supplier for the DataSegment future.
+ */
+ private final Supplier<ListenableFuture<DataSegment>> dataSegmentFutureSupplier;
+
+ /**
+ * Create a new RegularLoadableSegment.
+ *
+ * @param segmentManager segment manager for loading and caching segments
+ * @param segmentId the segment ID to load
+ * @param descriptor segment descriptor for querying
+ * @param inputCounters optional counters for tracking input
+ * @param coordinatorClient optional client for fetching DataSegment from Coordinator when not available locally
+ * @param isReindex true if this is a DML command writing to the same table it's reading from
+ */
+ public RegularLoadableSegment(
+ final SegmentManager segmentManager,
+ final SegmentId segmentId,
+ final SegmentDescriptor descriptor,
+ @Nullable final ChannelCounters inputCounters,
+ @Nullable final CoordinatorClient coordinatorClient,
+ final boolean isReindex
+ )
+ {
+ if (isReindex && coordinatorClient == null) {
+ throw DruidException.defensive("Got isReindex[%s], cannot respect this without a coordinatorClient", isReindex);
+ }
+
+ this.segmentManager = segmentManager;
+ this.segmentId = segmentId;
+ this.descriptor = descriptor;
+ this.inputCounters = inputCounters;
+ this.coordinatorClient = coordinatorClient;
+ this.isReindex = isReindex;
+
+ // Can't rely on local timeline if isReindex; always need to check the Coordinator to confirm the segment
+ // is still active.
+ this.cachedDataSegment = isReindex ? null : getDataSegmentFromLocalTimeline();
+ this.dataSegmentFutureSupplier = Suppliers.memoize(this::fetchDataSegment);
+ }
+
+ @Override
+ public ListenableFuture<DataSegment> dataSegmentFuture()
+ {
+ return Futures.nonCancellationPropagating(dataSegmentFutureSupplier.get());
+ }
+
+ @Override
+ public SegmentDescriptor descriptor()
+ {
+ return descriptor;
+ }
+
+ @Override
+ @Nullable
+ public ChannelCounters inputCounters()
+ {
+ return inputCounters;
+ }
+
+ @Override
+ @Nullable
+ public String description()
+ {
+ return segmentId.toString();
+ }
+
+ @Override
+ public synchronized Optional<Segment> acquireIfCached()
+ {
+ if (acquired) {
+ throw DruidException.defensive("Segment with descriptor[%s] is already acquired", descriptor);
+ }
+
+ if (cachedDataSegment != null) {
+ final Optional<Segment> cachedSegment = segmentManager.acquireCachedSegment(cachedDataSegment);
+ if (cachedSegment.isPresent()) {
+ acquired = true;
+ }
+ return cachedSegment;
+ }
+
+ return Optional.empty();
+ }
+
+ @Override
+ public synchronized AcquireSegmentAction acquire()
+ {
+ if (acquired) {
+ throw DruidException.defensive("Segment with descriptor[%s] is already acquired", descriptor);
+ }
+
+ acquired = true;
+
+ if (cachedDataSegment != null) {
+ return segmentManager.acquireSegment(cachedDataSegment);
+ } else {
+ // Create a shim AcquireSegmentAction that doesn't acquire a hold (yet). We can't make a real
+ // AcquireSegmentAction yet because we don't have the DataSegment object. It needs to be fetched
+ // from the Coordinator. That call is deferred until we're actually ready to load the segment, because
+ // we don't make the calls all at once when loading a lot of segments.
+
+ final Closer closer = Closer.create();
+ return new AcquireSegmentAction(
+ Suppliers.memoize(() -> FutureUtils.transformAsync(
+ dataSegmentFutureSupplier.get(),
+ dataSegment -> closer.register(segmentManager.acquireSegment(dataSegment)).getSegmentFuture()
+ )),
+ closer
+ );
+ }
+ }
+
+ /**
+ * Fetches the {@link DataSegment}, either returning it immediately if cached or fetching from the Coordinator.
+ */
+ private ListenableFuture<DataSegment> fetchDataSegment()
+ {
+ if (cachedDataSegment != null) {
+ return Futures.immediateFuture(cachedDataSegment);
+ } else if (coordinatorClient != null) {
+ return coordinatorClient.fetchSegment(
Review Comment:
i find myself wondering if maybe we can just do something to make `SegmentManager` more flexible (like say backed with a coordinator client to download datasegments on the fly) to make it so we can just do all the stuff with a segmentManager, but haven't really thought about this and maybe is too weird, either way, i think this is fine for this PR
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For additional commands, e-mail: commits-help@druid.apache.org
Re: [PR] MSQ: Use VSF, give stages control over inputs. (druid)
Posted by "gianm (via GitHub)" <gi...@apache.org>.
gianm commented on code in PR #18873:
URL: https://github.com/apache/druid/pull/18873#discussion_r2677686137
##########
multi-stage-query/src/main/java/org/apache/druid/msq/input/RegularLoadableSegment.java:
##########
@@ -0,0 +1,236 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.msq.input;
+
+import com.google.common.base.Supplier;
+import com.google.common.base.Suppliers;
+import com.google.common.util.concurrent.Futures;
+import com.google.common.util.concurrent.ListenableFuture;
+import com.google.errorprone.annotations.concurrent.GuardedBy;
+import org.apache.druid.client.coordinator.CoordinatorClient;
+import org.apache.druid.common.guava.FutureUtils;
+import org.apache.druid.error.DruidException;
+import org.apache.druid.java.util.common.io.Closer;
+import org.apache.druid.msq.counters.ChannelCounters;
+import org.apache.druid.query.SegmentDescriptor;
+import org.apache.druid.query.TableDataSource;
+import org.apache.druid.segment.Segment;
+import org.apache.druid.segment.loading.AcquireSegmentAction;
+import org.apache.druid.server.SegmentManager;
+import org.apache.druid.timeline.DataSegment;
+import org.apache.druid.timeline.SegmentId;
+import org.apache.druid.timeline.VersionedIntervalTimeline;
+import org.apache.druid.timeline.partition.PartitionChunk;
+
+import javax.annotation.Nullable;
+import java.util.Optional;
+
+/**
+ * Implementation of {@link LoadableSegment} for regular Druid segments loaded via {@link SegmentManager}.
+ * Created by {@link org.apache.druid.msq.input.table.SegmentsInputSliceReader}.
+ */
+public class RegularLoadableSegment implements LoadableSegment
+{
+ private final SegmentManager segmentManager;
+ private final SegmentId segmentId;
+ private final SegmentDescriptor descriptor;
+ @Nullable
+ private final ChannelCounters inputCounters;
+ @Nullable
+ private final CoordinatorClient coordinatorClient;
+ private final boolean isReindex;
+
+ @GuardedBy("this")
+ private boolean acquired;
+
+ /**
+ * Cached DataSegment from local timeline, if available. Null if not in local timeline or if isReindex is true.
+ */
+ @Nullable
+ private final DataSegment cachedDataSegment;
+
+ /**
+ * Memoized supplier for the DataSegment future.
+ */
+ private final Supplier<ListenableFuture<DataSegment>> dataSegmentFutureSupplier;
+
+ /**
+ * Create a new RegularLoadableSegment.
+ *
+ * @param segmentManager segment manager for loading and caching segments
+ * @param segmentId the segment ID to load
+ * @param descriptor segment descriptor for querying
+ * @param inputCounters optional counters for tracking input
+ * @param coordinatorClient optional client for fetching DataSegment from Coordinator when not available locally
+ * @param isReindex true if this is a DML command writing to the same table it's reading from
+ */
+ public RegularLoadableSegment(
+ final SegmentManager segmentManager,
+ final SegmentId segmentId,
+ final SegmentDescriptor descriptor,
+ @Nullable final ChannelCounters inputCounters,
+ @Nullable final CoordinatorClient coordinatorClient,
+ final boolean isReindex
+ )
+ {
+ if (isReindex && coordinatorClient == null) {
+ throw DruidException.defensive("Got isReindex[%s], cannot respect this without a coordinatorClient", isReindex);
+ }
+
+ this.segmentManager = segmentManager;
+ this.segmentId = segmentId;
+ this.descriptor = descriptor;
+ this.inputCounters = inputCounters;
+ this.coordinatorClient = coordinatorClient;
+ this.isReindex = isReindex;
+
+ // Can't rely on local timeline if isReindex; always need to check the Coordinator to confirm the segment
+ // is still active.
+ this.cachedDataSegment = isReindex ? null : getDataSegmentFromLocalTimeline();
+ this.dataSegmentFutureSupplier = Suppliers.memoize(this::fetchDataSegment);
+ }
+
+ @Override
+ public ListenableFuture<DataSegment> dataSegmentFuture()
+ {
+ return Futures.nonCancellationPropagating(dataSegmentFutureSupplier.get());
+ }
+
+ @Override
+ public SegmentDescriptor descriptor()
+ {
+ return descriptor;
+ }
+
+ @Override
+ @Nullable
+ public ChannelCounters inputCounters()
+ {
+ return inputCounters;
+ }
+
+ @Override
+ @Nullable
+ public String description()
+ {
+ return segmentId.toString();
+ }
+
+ @Override
+ public synchronized Optional<Segment> acquireIfCached()
+ {
+ if (acquired) {
+ throw DruidException.defensive("Segment with descriptor[%s] is already acquired", descriptor);
+ }
+
+ if (cachedDataSegment != null) {
+ final Optional<Segment> cachedSegment = segmentManager.acquireCachedSegment(cachedDataSegment);
+ if (cachedSegment.isPresent()) {
+ acquired = true;
+ }
+ return cachedSegment;
+ }
+
+ return Optional.empty();
+ }
+
+ @Override
+ public synchronized AcquireSegmentAction acquire()
+ {
+ if (acquired) {
+ throw DruidException.defensive("Segment with descriptor[%s] is already acquired", descriptor);
+ }
+
+ acquired = true;
+
+ if (cachedDataSegment != null) {
+ return segmentManager.acquireSegment(cachedDataSegment);
+ } else {
+ // Create a shim AcquireSegmentAction that doesn't acquire a hold (yet). We can't make a real
+ // AcquireSegmentAction yet because we don't have the DataSegment object. It needs to be fetched
+ // from the Coordinator. That call is deferred until we're actually ready to load the segment, because
+ // we don't make the calls all at once when loading a lot of segments.
+
+ final Closer closer = Closer.create();
+ return new AcquireSegmentAction(
+ Suppliers.memoize(() -> FutureUtils.transformAsync(
+ dataSegmentFutureSupplier.get(),
+ dataSegment -> closer.register(segmentManager.acquireSegment(dataSegment)).getSegmentFuture()
+ )),
+ closer
+ );
+ }
+ }
+
+ /**
+ * Fetches the {@link DataSegment}, either returning it immediately if cached or fetching from the Coordinator.
+ */
+ private ListenableFuture<DataSegment> fetchDataSegment()
+ {
+ if (cachedDataSegment != null) {
+ return Futures.immediateFuture(cachedDataSegment);
+ } else if (coordinatorClient != null) {
+ return coordinatorClient.fetchSegment(
Review Comment:
Along with the ability to manage non-segment data, like external files from cloud storage.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For additional commands, e-mail: commits-help@druid.apache.org
Re: [PR] MSQ: Use VSF, give stages control over inputs. (druid)
Posted by "gianm (via GitHub)" <gi...@apache.org>.
gianm commented on code in PR #18873:
URL: https://github.com/apache/druid/pull/18873#discussion_r2677669338
##########
indexing-service/src/main/java/org/apache/druid/indexing/common/SegmentCacheManagerFactory.java:
##########
@@ -52,11 +52,12 @@ public SegmentCacheManagerFactory(
this.jsonMapper = mapper;
}
- public SegmentCacheManager manufacturate(File storageDir)
+ public SegmentCacheManager manufacturate(File storageDir, boolean virtualStorage)
{
- final SegmentLoaderConfig loaderConfig = new SegmentLoaderConfig().withLocations(
- Collections.singletonList(new StorageLocationConfig(storageDir, null, null))
- );
+ final SegmentLoaderConfig loaderConfig =
+ new SegmentLoaderConfig()
+ .setLocations(Collections.singletonList(new StorageLocationConfig(storageDir, null, null)))
Review Comment:
I was hoping to manage more stuff through the storage locations as time goes on and that would be a good time to unify things. The other thing that isn't currently tracked is temporary storage of external data that has been downloaded.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For additional commands, e-mail: commits-help@druid.apache.org
Re: [PR] MSQ: Use VSF, give stages control over inputs. (druid)
Posted by "github-advanced-security[bot] (via GitHub)" <gi...@apache.org>.
github-advanced-security[bot] commented on code in PR #18873:
URL: https://github.com/apache/druid/pull/18873#discussion_r2651556051
##########
multi-stage-query/src/main/java/org/apache/druid/msq/shuffle/output/StageOutputHolder.java:
##########
@@ -82,10 +83,11 @@
public void setChannel(final ReadableFrameChannel channel)
{
if (!channelFuture.set(channel)) {
- if (FutureUtils.getUncheckedImmediately(channelFuture) == null) {
- throw new ISE("Closed");
+ final ReadableFrameChannel existingChannel = FutureUtils.getUncheckedImmediately(channelFuture);
+ if (existingChannel == null) {
+ throw DruidException.defensive("Closed, cannot set to[%s]");
Review Comment:
## Missing format argument
This format call refers to 1 argument(s) but only supplies 0 argument(s).
[Show more details](https://github.com/apache/druid/security/code-scanning/10605)
##########
multi-stage-query/src/main/java/org/apache/druid/msq/querykit/ReadableInputQueue.java:
##########
@@ -0,0 +1,366 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.msq.querykit;
+
+import com.google.common.collect.Sets;
+import com.google.common.util.concurrent.Futures;
+import com.google.common.util.concurrent.ListenableFuture;
+import com.google.errorprone.annotations.concurrent.GuardedBy;
+import org.apache.druid.common.guava.FutureUtils;
+import org.apache.druid.error.DruidException;
+import org.apache.druid.frame.channel.ReadableFrameChannel;
+import org.apache.druid.java.util.common.logger.Logger;
+import org.apache.druid.msq.exec.DataServerQueryHandler;
+import org.apache.druid.msq.exec.std.StandardPartitionReader;
+import org.apache.druid.msq.input.LoadableSegment;
+import org.apache.druid.msq.input.PhysicalInputSlice;
+import org.apache.druid.msq.input.stage.ReadablePartition;
+import org.apache.druid.msq.kernel.StageId;
+import org.apache.druid.msq.kernel.StagePartition;
+import org.apache.druid.segment.SegmentReference;
+import org.apache.druid.segment.loading.AcquireSegmentAction;
+import org.apache.druid.utils.CloseableUtils;
+
+import javax.annotation.Nullable;
+import java.io.Closeable;
+import java.io.IOException;
+import java.util.ArrayDeque;
+import java.util.LinkedHashSet;
+import java.util.List;
+import java.util.Queue;
+import java.util.Set;
+import java.util.concurrent.atomic.AtomicBoolean;
+
+/**
+ * Queue for returning {@link ReadableInput} from a list of {@link PhysicalInputSlice}.
+ *
+ * When closed, this object cancels all pending segment loads and releases all segments that have not yet been
+ * acquired by callers through {@link SegmentReferenceHolder#getSegmentReferenceOnce()}. Callers that have acquired
+ * segment references are responsible for closing those references, they will not be closed by this class.
+ */
+public class ReadableInputQueue implements Closeable
+{
+ private static final Logger log = new Logger(ReadableInputQueue.class);
+
+ /**
+ * Partitions to be read.
+ */
+ @GuardedBy("this")
+ private final Queue<ReadablePartition> readablePartitions = new ArrayDeque<>();
+
+ /**
+ * Segments to be loaded. Some may already be cached locally.
+ */
+ @GuardedBy("this")
+ private final Queue<LoadableSegment> loadableSegments = new ArrayDeque<>();
+
+ /**
+ * Realtime servers to be queried.
+ */
+ @GuardedBy("this")
+ private final Queue<DataServerQueryHandler> queryableServers = new ArrayDeque<>();
+
+ /**
+ * Segments currently being loaded.
+ */
+ @GuardedBy("this")
+ private final Set<AcquireSegmentAction> loadingSegments = new LinkedHashSet<>();
+
+ /**
+ * Segments that have been loaded.
+ */
+ @GuardedBy("this")
+ private final Set<SegmentReferenceHolder> loadedSegments = new LinkedHashSet<>();
+
+ /**
+ * Futures that are generated by loadahead, prior to a call to {@link #nextInput()} being made.
+ */
+ @GuardedBy("this")
+ private final Set<ListenableFuture<ReadableInput>> loadaheadFutures = Sets.newIdentityHashSet();
+
+ private final String queryId;
+ private final StandardPartitionReader partitionReader;
+ private final int loadahead;
+ private final AtomicBoolean loadaheadStarted = new AtomicBoolean(false);
+
+ public ReadableInputQueue(
+ final String queryId,
+ final StandardPartitionReader partitionReader,
+ final List<PhysicalInputSlice> slices,
+ final int loadahead
+ )
+ {
+ this.queryId = queryId;
+ this.partitionReader = partitionReader;
+ this.loadahead = loadahead;
+
+ for (final PhysicalInputSlice slice : slices) {
+ // First add all locally-cached segments, before any non-locally-cached segments. This promotes processing
+ // the locally-cached ones first. It also improves efficiency of loadahead, since once we get to the
+ // set of non-locally-cached segments, the "next" one is always one that is worth loading ahead.
+ for (final LoadableSegment loadableSegment : slice.getLoadableSegments()) {
+ if (loadableSegment.isLikelyCached()) {
+ loadableSegments.add(loadableSegment);
+ }
+ }
+
+ for (final LoadableSegment loadableSegment : slice.getLoadableSegments()) {
+ if (!loadableSegment.isLikelyCached()) {
+ loadableSegments.add(loadableSegment);
+ }
+ }
+
+ queryableServers.addAll(slice.getQueryableServers());
+ slice.getReadablePartitions().forEach(readablePartitions::add);
+ }
+ }
+
+ /**
+ * Starts loading up to {@link #loadahead} segments for future calls to {@link #nextInput()}, if this method
+ * has not yet been called. If it has previously been called, subsequent calls do nothing.
+ */
+ public void startLoadaheadIfNeeded()
+ {
+ if (loadaheadStarted.compareAndSet(false, true)) {
+ for (int i = 0; i < loadahead; i++) {
+ if (!addLoadaheadFuture()) {
+ break;
+ }
+ }
+ }
+ }
+
+ /**
+ * Returns the number of remaining inputs that can be returned by calls to {@link #nextInput()}.
+ */
+ public int remaining()
+ {
+ synchronized (this) {
+ return readablePartitions.size() + loadableSegments.size() + queryableServers.size() + loadaheadFutures.size();
+ }
+ }
+
+ /**
+ * Returns the next {@link ReadableInput}. The future resolves when the input is ready to read.
+ */
+ @Nullable
+ public ListenableFuture<ReadableInput> nextInput()
+ {
+ ListenableFuture<ReadableInput> future;
+
+ future = nextServerInput();
+ if (future != null) {
+ return future;
+ }
+
+ future = nextChannelInput();
+ if (future != null) {
+ return future;
+ }
+
+ future = nextSegmentInput();
+ if (future != null) {
+ return future;
+ }
+
+ return null;
+ }
+
+ /**
+ * Returns the next input from {@link #queryableServers}, if any. Returns null if none remain.
+ */
+ @Nullable
+ private ListenableFuture<ReadableInput> nextServerInput()
+ {
+ final DataServerQueryHandler handler;
+ synchronized (this) {
+ handler = queryableServers.poll();
+ }
+
+ if (handler == null) {
+ return null;
+ }
+
+ return Futures.immediateFuture(ReadableInput.dataServerQuery(handler));
+ }
+
+ /**
+ * Returns the next input from {@link #readablePartitions}, if any. Returns null if none remain.
+ */
+ @Nullable
+ private ListenableFuture<ReadableInput> nextChannelInput()
+ {
+ final ReadablePartition readablePartition;
+ synchronized (this) {
+ readablePartition = readablePartitions.poll();
+ }
+
+ if (readablePartition == null) {
+ return null;
+ }
+
+ ReadableFrameChannel channel = null;
+ try {
+ channel = partitionReader.openChannel(readablePartition);
+ return Futures.immediateFuture(
+ ReadableInput.channel(
+ channel,
+ partitionReader.frameReader(readablePartition.getStageNumber()),
+ new StagePartition(
+ new StageId(queryId, readablePartition.getStageNumber()),
+ readablePartition.getPartitionNumber()
+ )
+ )
+ );
+ }
+ catch (IOException e) {
+ throw CloseableUtils.closeAndWrapInCatch(e, channel);
+ }
+ }
+
+ /**
+ * Returns the next input from {@link #loadableSegments}, if any. Returns null if none remain.
+ */
+ @Nullable
+ private ListenableFuture<ReadableInput> nextSegmentInput()
+ {
+ // Pick a loadahead future, preferring ones that are already loaded.
+ ListenableFuture<ReadableInput> selectedLoadaheadFuture = null;
+ synchronized (this) {
+ for (ListenableFuture<ReadableInput> f : loadaheadFutures) {
+ if (selectedLoadaheadFuture == null || f.isDone()) {
+ selectedLoadaheadFuture = f;
+ if (f.isDone()) {
+ break;
+ }
+ }
+ }
+
+ if (selectedLoadaheadFuture != null) {
+ loadaheadFutures.remove(selectedLoadaheadFuture);
+ addLoadaheadFuture(); // Replace the one we just took out.
+ return selectedLoadaheadFuture;
+ }
+ }
+
+ return loadNextSegment();
+ }
+
+ /**
+ * Load the next segment from {@link #loadableSegments} and return a future to its reference. Returns null
+ * if {@link #loadableSegments} is empty.
+ */
+ @Nullable
+ private ListenableFuture<ReadableInput> loadNextSegment()
+ {
+ synchronized (this) {
+ final LoadableSegment nextLoadableSegment = loadableSegments.poll();
+ if (nextLoadableSegment == null) {
+ return null;
+ }
+
+ final AcquireSegmentAction acquireSegmentAction = nextLoadableSegment.acquire();
+ loadingSegments.add(acquireSegmentAction);
+ return FutureUtils.transform(
+ acquireSegmentAction.getSegmentFuture(),
+ segment -> {
+ synchronized (ReadableInputQueue.this) {
+ // Transfer segment from "loadingSegments" to "loadedSegments" and return a reference to it.
+ if (loadingSegments.remove(acquireSegmentAction)) {
+ try {
+ final SegmentReferenceHolder referenceHolder = new SegmentReferenceHolder(
+ new SegmentReference(
+ nextLoadableSegment.descriptor(),
+ segment.getReferenceProvider().acquireReference(),
+ acquireSegmentAction // Release the hold when the SegmentReference is closed.
+ ),
+ nextLoadableSegment.inputCounters(),
+ nextLoadableSegment.description()
+ );
+ loadedSegments.add(referenceHolder);
+ return ReadableInput.segment(referenceHolder);
+ }
+ catch (Throwable e) {
+ // Javadoc for segment.acquireReference() suggests it can throw exceptions; handle that here
+ // by closing the original AcquireSegmentAction.
+ throw CloseableUtils.closeAndWrapInCatch(e, acquireSegmentAction);
+ }
+ } else {
+ throw DruidException.defensive(
+ "Segment[%s] removed from loadingSegments before loading complete. It is possible that close() "
+ + "was called with futures in flight.",
+ nextLoadableSegment.descriptor()
+ );
+ }
+ }
+ }
+ );
+ }
+ }
+
+ /**
+ * Calls {@link #nextSegmentInput()} and adds the future to {@link #loadaheadFutures}. Returns whether a future
+ * was added.
+ */
+ private boolean addLoadaheadFuture()
+ {
+ final ListenableFuture<ReadableInput> nextFuture = loadNextSegment();
+ if (nextFuture != null) {
+ synchronized (this) {
+ loadaheadFutures.add(nextFuture);
+ }
+ return true;
+ } else {
+ return false;
+ }
+ }
+
+ @Override
+ public void close()
+ {
+ synchronized (this) {
+ readablePartitions.clear();
+ queryableServers.clear();
+ loadableSegments.clear();
+
+ // Cancel all pending segment loads.
+ for (AcquireSegmentAction acquireSegmentAction : loadingSegments) {
+ CloseableUtils.closeAndSuppressExceptions(
+ acquireSegmentAction,
+ e -> log.warn(e, "Failed to close loadingSegment[%s]", acquireSegmentAction)
Review Comment:
## Use of default toString()
Default toString(): AcquireSegmentAction inherits toString() from Object, and so is not suitable for printing.
[Show more details](https://github.com/apache/druid/security/code-scanning/10606)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For additional commands, e-mail: commits-help@druid.apache.org
Re: [PR] MSQ: Use VSF, give stages control over inputs. (druid)
Posted by "gianm (via GitHub)" <gi...@apache.org>.
gianm commented on code in PR #18873:
URL: https://github.com/apache/druid/pull/18873#discussion_r2677685387
##########
multi-stage-query/src/main/java/org/apache/druid/msq/input/PhysicalInputSlice.java:
##########
@@ -0,0 +1,63 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.msq.input;
+
+import com.google.common.base.Preconditions;
+import org.apache.druid.msq.exec.DataServerQueryHandler;
+import org.apache.druid.msq.input.stage.ReadablePartitions;
+
+import java.util.List;
+
+/**
+ * An {@link InputSlice} that has been prepared for reading by an {@link InputSliceReader}. Nothing contained in
+ * this class references open resources, so this class is not closeable and does not need to be closed.
+ */
+public class PhysicalInputSlice
Review Comment:
It's definitely a "naming is hard" thing.
It's a somewhat-more-physical version of `InputSlice`. `InputSlice` is always purely logical (it just has the identifiers / paths of things that should be read, as strings and POJOs) but `PhysicalInputSlice` may contain the ability to actually read external data or actually load segments.
It's not just "leaf" since it includes partition inputs, which would be used by non-leaf stages.
##########
multi-stage-query/src/main/java/org/apache/druid/msq/input/RegularLoadableSegment.java:
##########
@@ -0,0 +1,236 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.msq.input;
+
+import com.google.common.base.Supplier;
+import com.google.common.base.Suppliers;
+import com.google.common.util.concurrent.Futures;
+import com.google.common.util.concurrent.ListenableFuture;
+import com.google.errorprone.annotations.concurrent.GuardedBy;
+import org.apache.druid.client.coordinator.CoordinatorClient;
+import org.apache.druid.common.guava.FutureUtils;
+import org.apache.druid.error.DruidException;
+import org.apache.druid.java.util.common.io.Closer;
+import org.apache.druid.msq.counters.ChannelCounters;
+import org.apache.druid.query.SegmentDescriptor;
+import org.apache.druid.query.TableDataSource;
+import org.apache.druid.segment.Segment;
+import org.apache.druid.segment.loading.AcquireSegmentAction;
+import org.apache.druid.server.SegmentManager;
+import org.apache.druid.timeline.DataSegment;
+import org.apache.druid.timeline.SegmentId;
+import org.apache.druid.timeline.VersionedIntervalTimeline;
+import org.apache.druid.timeline.partition.PartitionChunk;
+
+import javax.annotation.Nullable;
+import java.util.Optional;
+
+/**
+ * Implementation of {@link LoadableSegment} for regular Druid segments loaded via {@link SegmentManager}.
+ * Created by {@link org.apache.druid.msq.input.table.SegmentsInputSliceReader}.
+ */
+public class RegularLoadableSegment implements LoadableSegment
+{
+ private final SegmentManager segmentManager;
+ private final SegmentId segmentId;
+ private final SegmentDescriptor descriptor;
+ @Nullable
+ private final ChannelCounters inputCounters;
+ @Nullable
+ private final CoordinatorClient coordinatorClient;
+ private final boolean isReindex;
+
+ @GuardedBy("this")
+ private boolean acquired;
+
+ /**
+ * Cached DataSegment from local timeline, if available. Null if not in local timeline or if isReindex is true.
+ */
+ @Nullable
+ private final DataSegment cachedDataSegment;
+
+ /**
+ * Memoized supplier for the DataSegment future.
+ */
+ private final Supplier<ListenableFuture<DataSegment>> dataSegmentFutureSupplier;
+
+ /**
+ * Create a new RegularLoadableSegment.
+ *
+ * @param segmentManager segment manager for loading and caching segments
+ * @param segmentId the segment ID to load
+ * @param descriptor segment descriptor for querying
+ * @param inputCounters optional counters for tracking input
+ * @param coordinatorClient optional client for fetching DataSegment from Coordinator when not available locally
+ * @param isReindex true if this is a DML command writing to the same table it's reading from
+ */
+ public RegularLoadableSegment(
+ final SegmentManager segmentManager,
+ final SegmentId segmentId,
+ final SegmentDescriptor descriptor,
+ @Nullable final ChannelCounters inputCounters,
+ @Nullable final CoordinatorClient coordinatorClient,
+ final boolean isReindex
+ )
+ {
+ if (isReindex && coordinatorClient == null) {
+ throw DruidException.defensive("Got isReindex[%s], cannot respect this without a coordinatorClient", isReindex);
+ }
+
+ this.segmentManager = segmentManager;
+ this.segmentId = segmentId;
+ this.descriptor = descriptor;
+ this.inputCounters = inputCounters;
+ this.coordinatorClient = coordinatorClient;
+ this.isReindex = isReindex;
+
+ // Can't rely on local timeline if isReindex; always need to check the Coordinator to confirm the segment
+ // is still active.
+ this.cachedDataSegment = isReindex ? null : getDataSegmentFromLocalTimeline();
+ this.dataSegmentFutureSupplier = Suppliers.memoize(this::fetchDataSegment);
+ }
+
+ @Override
+ public ListenableFuture<DataSegment> dataSegmentFuture()
+ {
+ return Futures.nonCancellationPropagating(dataSegmentFutureSupplier.get());
+ }
+
+ @Override
+ public SegmentDescriptor descriptor()
+ {
+ return descriptor;
+ }
+
+ @Override
+ @Nullable
+ public ChannelCounters inputCounters()
+ {
+ return inputCounters;
+ }
+
+ @Override
+ @Nullable
+ public String description()
+ {
+ return segmentId.toString();
+ }
+
+ @Override
+ public synchronized Optional<Segment> acquireIfCached()
+ {
+ if (acquired) {
+ throw DruidException.defensive("Segment with descriptor[%s] is already acquired", descriptor);
+ }
+
+ if (cachedDataSegment != null) {
+ final Optional<Segment> cachedSegment = segmentManager.acquireCachedSegment(cachedDataSegment);
+ if (cachedSegment.isPresent()) {
+ acquired = true;
+ }
+ return cachedSegment;
+ }
+
+ return Optional.empty();
+ }
+
+ @Override
+ public synchronized AcquireSegmentAction acquire()
+ {
+ if (acquired) {
+ throw DruidException.defensive("Segment with descriptor[%s] is already acquired", descriptor);
+ }
+
+ acquired = true;
+
+ if (cachedDataSegment != null) {
+ return segmentManager.acquireSegment(cachedDataSegment);
+ } else {
+ // Create a shim AcquireSegmentAction that doesn't acquire a hold (yet). We can't make a real
+ // AcquireSegmentAction yet because we don't have the DataSegment object. It needs to be fetched
+ // from the Coordinator. That call is deferred until we're actually ready to load the segment, because
+ // we don't make the calls all at once when loading a lot of segments.
+
+ final Closer closer = Closer.create();
+ return new AcquireSegmentAction(
+ Suppliers.memoize(() -> FutureUtils.transformAsync(
+ dataSegmentFutureSupplier.get(),
+ dataSegment -> closer.register(segmentManager.acquireSegment(dataSegment)).getSegmentFuture()
+ )),
+ closer
+ );
+ }
+ }
+
+ /**
+ * Fetches the {@link DataSegment}, either returning it immediately if cached or fetching from the Coordinator.
+ */
+ private ListenableFuture<DataSegment> fetchDataSegment()
+ {
+ if (cachedDataSegment != null) {
+ return Futures.immediateFuture(cachedDataSegment);
+ } else if (coordinatorClient != null) {
+ return coordinatorClient.fetchSegment(
Review Comment:
That would be nice.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For additional commands, e-mail: commits-help@druid.apache.org
Re: [PR] MSQ: Use VSF, give stages control over inputs. (druid)
Posted by "gianm (via GitHub)" <gi...@apache.org>.
gianm commented on code in PR #18873:
URL: https://github.com/apache/druid/pull/18873#discussion_r2677675309
##########
processing/src/main/java/org/apache/druid/segment/SegmentReference.java:
##########
@@ -52,23 +57,60 @@ public SegmentReference(
)
{
this.segmentDescriptor = segmentDescriptor;
- closer.register(cleanupHold);
- this.segmentReference = segmentReference.map(closer::register);
+ this.segmentReference = segmentReference;
+ this.cleanupHold = cleanupHold;
}
public SegmentDescriptor getSegmentDescriptor()
{
return segmentDescriptor;
}
+ /**
+ * Returns the actual segment. Do not close the Segment when you are done with it, only close the SegmentReference.
+ */
public Optional<Segment> getSegmentReference()
{
return segmentReference;
}
+ /**
+ * Maps the wrapped segment and returns a reference to it. Closes the reference if the {@link SegmentMapFunction}
+ * throws an exception. Regardless of success or failure of this method, the old reference should be discarded.
+ * Do not call {@link #close()} on the old reference, only call it on the new one.
+ */
+ public SegmentReference map(final SegmentMapFunction segmentMapFn)
Review Comment:
Updated to:
```
/**
* Wrapper for a {@link SegmentDescriptor} and {@link Optional<Segment>}. For regular segments, the latter is typically
* created by a {@link ReferenceCountedSegmentProvider} that may have a {@link SegmentMapFunction} applied to it.
*
* If the {@link SegmentMapFunction} you want to apply is not available at the time the {@link SegmentReference}
* is created, use {@link #map(SegmentMapFunction)} to apply it.
*
* Closing this object closes both the {@link #getSegmentReference()} and any closeables attached from the process of
* creating this object, such as from {@link AcquireSegmentAction}. The object from {@link #getSegmentReference()}
* should not be closed directly by callers.
*/
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For additional commands, e-mail: commits-help@druid.apache.org