Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Use task actions to fetch used segments in MSQ #15284

Merged
merged 9 commits into from
Dec 1, 2023
Merged

Use task actions to fetch used segments in MSQ #15284

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
9fb703a
Use task actions to fetch used segments in MSQ
AmatyaAvadhanula Oct 31, 2023
3ea0e49
Merge remote-tracking branch 'upstream/master' into msq_retrieveSegme…
AmatyaAvadhanula Nov 6, 2023
95ab649
Fix tests
AmatyaAvadhanula Nov 6, 2023
5444b69
Fixing tests.
cryptoe Nov 21, 2023
6dba33e
Revert "Fix tests"
cryptoe Nov 21, 2023
bb14fe8
Merge remote-tracking branch 'apache/master' into msq_retrieveSegment…
cryptoe Nov 21, 2023
bc6eaf1
Removing conditional check in tests.
cryptoe Nov 23, 2023
411dbac
Merge remote-tracking branch 'apache/master' into msq_retrieveSegment…
cryptoe Nov 28, 2023
8812495
Pulling in latest changes.
cryptoe Nov 28, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 13 additions & 2 deletions ...nsions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/ControllerImpl.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -69,6 +69,7 @@
import org.apache.druid.indexing.common.actions.LockListAction;
import org.apache.druid.indexing.common.actions.LockReleaseAction;
import org.apache.druid.indexing.common.actions.MarkSegmentsAsUnusedAction;
import org.apache.druid.indexing.common.actions.RetrieveSegmentsToReplaceAction;
import org.apache.druid.indexing.common.actions.SegmentAllocateAction;
import org.apache.druid.indexing.common.actions.SegmentTransactionalAppendAction;
import org.apache.druid.indexing.common.actions.SegmentTransactionalInsertAction;
Expand Down Expand Up @@ -1200,8 +1201,18 @@ private DataSegmentTimelineView makeDataSegmentTimelineView()
}

// Fetch all published, used segments (all non-realtime segments) from the metadata store.
final Collection<DataSegment> publishedUsedSegments =
FutureUtils.getUnchecked(context.coordinatorClient().fetchUsedSegments(dataSource, intervals), true);
// If the task is operating with a REPLACE lock,
// any segment created after the lock was acquired for its interval will not be considered.
final Collection<DataSegment> publishedUsedSegments;
try {
publishedUsedSegments = context.taskActionClient().submit(new RetrieveSegmentsToReplaceAction(
dataSource,
intervals
));
}
catch (IOException e) {
throw new MSQException(e, UnknownFault.forException(e));
}

int realtimeCount = 0;

Expand Down
2 changes: 1 addition & 1 deletion .../multi-stage-query/src/test/java/org/apache/druid/msq/test/CalciteArraysQueryMSQTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -80,7 +80,7 @@ public SqlEngine createEngine(
final MSQTestOverlordServiceClient indexingServiceClient = new MSQTestOverlordServiceClient(
queryJsonMapper,
injector,
new MSQTestTaskActionClient(queryJsonMapper),
new MSQTestTaskActionClient(queryJsonMapper, injector),
workerMemoryParameters,
ImmutableList.of()
);
Expand Down
2 changes: 1 addition & 1 deletion ...ti-stage-query/src/test/java/org/apache/druid/msq/test/CalciteSelectJoinQueryMSQTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -128,7 +128,7 @@ public SqlEngine createEngine(
final MSQTestOverlordServiceClient indexingServiceClient = new MSQTestOverlordServiceClient(
queryJsonMapper,
injector,
new MSQTestTaskActionClient(queryJsonMapper),
new MSQTestTaskActionClient(queryJsonMapper, injector),
workerMemoryParameters,
ImmutableList.of()
);
Expand Down
2 changes: 1 addition & 1 deletion .../multi-stage-query/src/test/java/org/apache/druid/msq/test/CalciteSelectQueryMSQTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -84,7 +84,7 @@ public SqlEngine createEngine(
final MSQTestOverlordServiceClient indexingServiceClient = new MSQTestOverlordServiceClient(
queryJsonMapper,
injector,
new MSQTestTaskActionClient(queryJsonMapper),
new MSQTestTaskActionClient(queryJsonMapper, injector),
workerMemoryParameters,
ImmutableList.of()
);
Expand Down
2 changes: 1 addition & 1 deletion ...e/multi-stage-query/src/test/java/org/apache/druid/msq/test/CalciteUnionQueryMSQTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -94,7 +94,7 @@ public SqlEngine createEngine(
final MSQTestOverlordServiceClient indexingServiceClient = new MSQTestOverlordServiceClient(
queryJsonMapper,
injector,
new MSQTestTaskActionClient(queryJsonMapper),
new MSQTestTaskActionClient(queryJsonMapper, injector),
workerMemoryParameters,
ImmutableList.of()
);
Expand Down
2 changes: 1 addition & 1 deletion extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestBase.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -503,7 +503,7 @@ public String getFormatString()

doReturn(mock(Request.class)).when(brokerClient).makeRequest(any(), anyString());

testTaskActionClient = Mockito.spy(new MSQTestTaskActionClient(objectMapper));
testTaskActionClient = Mockito.spy(new MSQTestTaskActionClient(objectMapper, injector));
indexingServiceClient = new MSQTestOverlordServiceClient(
objectMapper,
injector,
Expand Down
15 changes: 0 additions & 15 deletions ...e/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestControllerContext.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -50,7 +50,6 @@
import org.apache.druid.msq.util.MultiStageQueryContext;
import org.apache.druid.query.QueryContext;
import org.apache.druid.server.DruidNode;
import org.apache.druid.server.SpecificSegmentsQuerySegmentWalker;
import org.apache.druid.server.metrics.NoopServiceEmitter;
import org.mockito.ArgumentMatchers;
import org.mockito.Mockito;
Expand Down Expand Up @@ -104,20 +103,6 @@ public MSQTestControllerContext(
this.injector = injector;
this.taskActionClient = taskActionClient;
coordinatorClient = Mockito.mock(CoordinatorClient.class);
Mockito.when(coordinatorClient.fetchUsedSegments(
ArgumentMatchers.anyString(),
ArgumentMatchers.anyList()
)
).thenAnswer(invocation ->
Futures.immediateFuture(
injector.getInstance(SpecificSegmentsQuerySegmentWalker.class)
.getSegments()
.stream()
.filter(dataSegment -> dataSegment.getDataSource()
.equals(invocation.getArguments()[0]))
.collect(Collectors.toList())
)
);

Mockito.when(coordinatorClient.fetchServerViewSegments(
ArgumentMatchers.anyString(),
Expand Down
16 changes: 15 additions & 1 deletion ...re/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestTaskActionClient.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,9 +23,11 @@
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableSet;
import com.google.inject.Injector;
import org.apache.druid.indexing.common.TaskLockType;
import org.apache.druid.indexing.common.TimeChunkLock;
import org.apache.druid.indexing.common.actions.LockListAction;
import org.apache.druid.indexing.common.actions.RetrieveSegmentsToReplaceAction;
import org.apache.druid.indexing.common.actions.RetrieveUsedSegmentsAction;
import org.apache.druid.indexing.common.actions.SegmentAllocateAction;
import org.apache.druid.indexing.common.actions.SegmentTransactionalAppendAction;
Expand All @@ -39,6 +41,7 @@
import org.apache.druid.java.util.common.granularity.PeriodGranularity;
import org.apache.druid.msq.indexing.error.InsertLockPreemptedFaultTest;
import org.apache.druid.segment.realtime.appenderator.SegmentIdWithShardSpec;
import org.apache.druid.server.SpecificSegmentsQuerySegmentWalker;
import org.apache.druid.timeline.DataSegment;
import org.apache.druid.timeline.SegmentId;
import org.joda.time.Interval;
Expand All @@ -62,12 +65,15 @@ public class MSQTestTaskActionClient implements TaskActionClient
"foo2", ImmutableList.of(Intervals.of("2000-01-01/P1D"))
);
private final Set<DataSegment> publishedSegments = new HashSet<>();
private final Injector injector;

public MSQTestTaskActionClient(
ObjectMapper mapper
ObjectMapper mapper,
Injector injector
)
{
this.mapper = mapper;
this.injector = injector;
}

@Override
Expand Down Expand Up @@ -122,6 +128,14 @@ public <RetType> RetType submit(TaskAction<RetType> taskAction)
.build()
).collect(Collectors.toSet());
}
} else if (taskAction instanceof RetrieveSegmentsToReplaceAction) {
String dataSource = ((RetrieveSegmentsToReplaceAction) taskAction).getDataSource();
return (RetType) injector.getInstance(SpecificSegmentsQuerySegmentWalker.class)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could you please add a comment here explaining why we need to use the QuerySegmentWalker here instead of how we handle RetrieveUsedSegmentsAction here or how we would respond to the coordinatorClient().fetchUsedSegments() in MSQ tests?

.getSegments()
.stream()
.filter(dataSegment -> dataSegment.getDataSource()
.equals(dataSource))
.collect(Collectors.toSet());
} else if (taskAction instanceof SegmentTransactionalInsertAction) {
final Set<DataSegment> segments = ((SegmentTransactionalInsertAction) taskAction).getSegments();
publishedSegments.addAll(segments);
Expand Down