From 747034b61e6e206f6bb475b3bdfe10a899df874b Mon Sep 17 00:00:00 2001 From: Zoltan Haindrich Date: Fri, 28 Feb 2025 10:04:56 +0000 Subject: [PATCH 01/50] bootstrap --- .../druid/msq/test/DartComponentSupplier.java | 75 +++++++++++++++++++ .../dart.iq | 54 +++++++++++++ 2 files changed, 129 insertions(+) create mode 100644 extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/DartComponentSupplier.java create mode 100644 extensions-core/multi-stage-query/src/test/quidem/org.apache.druid.msq.quidem.MSQQuidemTest/dart.iq diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/DartComponentSupplier.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/DartComponentSupplier.java new file mode 100644 index 000000000000..c207b41286d0 --- /dev/null +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/DartComponentSupplier.java @@ -0,0 +1,75 @@ +/* + * 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.test; + +import com.fasterxml.jackson.databind.ObjectMapper; +import com.google.common.collect.ImmutableList; +import com.google.inject.Binder; +import com.google.inject.Injector; +import com.google.inject.Provides; +import org.apache.druid.guice.LazySingleton; +import org.apache.druid.initialization.DruidModule; +import org.apache.druid.msq.exec.WorkerMemoryParameters; +import org.apache.druid.sql.calcite.TempDirProducer; +import org.apache.druid.sql.calcite.util.DruidModuleCollection; +import org.apache.druid.sql.calcite.util.SqlTestFramework.StandardComponentSupplier; + +public class DartComponentSupplier extends AbstractMSQComponentSupplierDelegate +{ + public DartComponentSupplier(TempDirProducer tempFolderProducer) + { + super(new StandardComponentSupplier(tempFolderProducer)); + } + + @Override + public DruidModule getOverrideModule() + { + return DruidModuleCollection.of(super.getOverrideModule(), + new LocalDartModule() + ); + } + + static class LocalDartModule implements DruidModule{ + + @Override + public void configure(Binder binder) + { + } + + @Provides + @LazySingleton + private MSQTestOverlordServiceClient makeOverlordServiceClient( + ObjectMapper queryJsonMapper, + Injector injector, + WorkerMemoryParameters workerMemoryParameters) + { + final MSQTestOverlordServiceClient indexingServiceClient = new MSQTestOverlordServiceClient( + queryJsonMapper, + injector, + new MSQTestTaskActionClient(queryJsonMapper, injector), + workerMemoryParameters, + ImmutableList.of() + ); + return indexingServiceClient; + } + + + } +} diff --git a/extensions-core/multi-stage-query/src/test/quidem/org.apache.druid.msq.quidem.MSQQuidemTest/dart.iq b/extensions-core/multi-stage-query/src/test/quidem/org.apache.druid.msq.quidem.MSQQuidemTest/dart.iq new file mode 100644 index 000000000000..a5ef2481f085 --- /dev/null +++ b/extensions-core/multi-stage-query/src/test/quidem/org.apache.druid.msq.quidem.MSQQuidemTest/dart.iq @@ -0,0 +1,54 @@ +!set plannerStrategy DECOUPLED +!use druidtest://?componentSupplier=DartComponentSupplier +!set outputformat mysql + +# This test validates that all output rows are flushed to the output channel even if frame writer's capacity is reached. + +select count(*) as actualNumRows +from ( + select countryName, cityName, channel, added, delta, row_number() over() as rowNumber + from wikipedia + group by countryName, cityName, channel, added, delta +); ++---------------+ +| actualNumRows | ++---------------+ +| 11631 | ++---------------+ +(1 row) + +!ok + +# Validate that all rows are outputted by window WindowOperatorQueryFrameProcessor layer for empty over() clause scenario. + +select count(*) as numRows, max(rowNumber) as maxRowNumber +from ( + select countryName, cityName, channel, added, delta, row_number() over() as rowNumber + from wikipedia + group by countryName, cityName, channel, added, delta +); ++---------+--------------+ +| numRows | maxRowNumber | ++---------+--------------+ +| 11631 | 11631 | ++---------+--------------+ +(1 row) + +!ok + +# Validate that all rows are outputted by window WindowOperatorQueryFrameProcessor layer for non-empty over() clause scenario. + +select rowNumber, count(rowNumber) as numRows +from ( + select countryName, cityName, channel, added, delta, row_number() over(partition by countryName, cityName, channel, added, delta) as rowNumber + from wikipedia + group by countryName, cityName, channel, added, delta +) group by rowNumber; ++-----------+---------+ +| rowNumber | numRows | ++-----------+---------+ +| 1 | 11631 | ++-----------+---------+ +(1 row) + +!ok From abe691653843e6d5d4cb1395ba9da0d6adb7e3f2 Mon Sep 17 00:00:00 2001 From: Zoltan Haindrich Date: Fri, 28 Feb 2025 10:20:57 +0000 Subject: [PATCH 02/50] use interface --- .../controller/DartControllerContext.java | 6 +- .../DartControllerContextFactoryImpl.java | 5 +- .../msq/dart/guice/DartControllerModule.java | 3 +- .../msq/test/CalciteSelectQueryMSQTest.java | 2 +- .../druid/msq/test/DartComponentSupplier.java | 52 +++- .../test/MSQTestOverlordServiceClient2.java | 293 ++++++++++++++++++ 6 files changed, 352 insertions(+), 9 deletions(-) create mode 100644 extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestOverlordServiceClient2.java diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartControllerContext.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartControllerContext.java index b7cf087a3f4d..c2257d160ed3 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartControllerContext.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartControllerContext.java @@ -21,7 +21,7 @@ import com.fasterxml.jackson.databind.ObjectMapper; import com.google.inject.Injector; -import org.apache.druid.client.BrokerServerView; +import org.apache.druid.client.TimelineServerView; import org.apache.druid.error.DruidException; import org.apache.druid.indexing.common.TaskLockType; import org.apache.druid.indexing.common.actions.TaskActionClient; @@ -84,7 +84,7 @@ public class DartControllerContext implements ControllerContext private final ObjectMapper jsonMapper; private final DruidNode selfNode; private final DartWorkerClient workerClient; - private final BrokerServerView serverView; + private final TimelineServerView serverView; private final MemoryIntrospector memoryIntrospector; private final ServiceMetricEvent.Builder metricBuilder; private final ServiceEmitter emitter; @@ -95,7 +95,7 @@ public DartControllerContext( final DruidNode selfNode, final DartWorkerClient workerClient, final MemoryIntrospector memoryIntrospector, - final BrokerServerView serverView, + final TimelineServerView serverView, final ServiceEmitter emitter ) { diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartControllerContextFactoryImpl.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartControllerContextFactoryImpl.java index 8cefb6af7ece..0a69a72fba23 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartControllerContextFactoryImpl.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartControllerContextFactoryImpl.java @@ -23,6 +23,7 @@ import com.google.inject.Inject; import com.google.inject.Injector; import org.apache.druid.client.BrokerServerView; +import org.apache.druid.client.TimelineServerView; import org.apache.druid.guice.annotations.EscalatedGlobal; import org.apache.druid.guice.annotations.Json; import org.apache.druid.guice.annotations.Self; @@ -41,7 +42,7 @@ public class DartControllerContextFactoryImpl implements DartControllerContextFa private final ObjectMapper smileMapper; private final DruidNode selfNode; private final ServiceClientFactory serviceClientFactory; - private final BrokerServerView serverView; + private final TimelineServerView serverView; private final MemoryIntrospector memoryIntrospector; private final ServiceEmitter emitter; @@ -53,7 +54,7 @@ public DartControllerContextFactoryImpl( @Self final DruidNode selfNode, @EscalatedGlobal final ServiceClientFactory serviceClientFactory, final MemoryIntrospector memoryIntrospector, - final BrokerServerView serverView, + final TimelineServerView serverView, final ServiceEmitter emitter ) { diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/guice/DartControllerModule.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/guice/DartControllerModule.java index 8a4b73bc9b0f..1a1f387a0eae 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/guice/DartControllerModule.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/guice/DartControllerModule.java @@ -65,7 +65,8 @@ public class DartControllerModule implements DruidModule @Override public void configure(Binder binder) { - if (DartModules.isDartEnabled(properties)) { + // FIXME + if (true || DartModules.isDartEnabled(properties)) { binder.install(new ActualModule()); } } diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/CalciteSelectQueryMSQTest.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/CalciteSelectQueryMSQTest.java index 33aa39e21a53..0d6683496cf5 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/CalciteSelectQueryMSQTest.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/CalciteSelectQueryMSQTest.java @@ -36,7 +36,7 @@ /** * Runs {@link CalciteQueryTest} but with MSQ engine */ -@SqlTestFrameworkConfig.ComponentSupplier(StandardMSQComponentSupplier.class) +@SqlTestFrameworkConfig.ComponentSupplier(DartComponentSupplier.class) public class CalciteSelectQueryMSQTest extends CalciteQueryTest { @Override diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/DartComponentSupplier.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/DartComponentSupplier.java index c207b41286d0..49649e9f98ce 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/DartComponentSupplier.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/DartComponentSupplier.java @@ -26,8 +26,13 @@ import com.google.inject.Provides; import org.apache.druid.guice.LazySingleton; import org.apache.druid.initialization.DruidModule; +import org.apache.druid.msq.dart.controller.sql.DartSqlEngine; +import org.apache.druid.msq.dart.guice.DartControllerModule; +import org.apache.druid.msq.dart.guice.DartWorkerModule; import org.apache.druid.msq.exec.WorkerMemoryParameters; +import org.apache.druid.server.QueryLifecycleFactory; import org.apache.druid.sql.calcite.TempDirProducer; +import org.apache.druid.sql.calcite.run.SqlEngine; import org.apache.druid.sql.calcite.util.DruidModuleCollection; import org.apache.druid.sql.calcite.util.SqlTestFramework.StandardComponentSupplier; @@ -46,6 +51,25 @@ public DruidModule getOverrideModule() ); } + @Override + public DruidModule getCoreModule() + { + return DruidModuleCollection.of(super.getCoreModule(), + new DartControllerModule(), + new DartWorkerModule() + ); + } + + @Override + public SqlEngine createEngine( + QueryLifecycleFactory qlf, + ObjectMapper queryJsonMapper, + Injector injector) + { + return injector.getInstance(DartSqlEngine.class); + } + + static class LocalDartModule implements DruidModule{ @Override @@ -53,14 +77,38 @@ public void configure(Binder binder) { } +// @Provides +// @LazySingleton +// public MSQTaskSqlEngine createEngine( +// ObjectMapper queryJsonMapper, +// MSQTestOverlordServiceClient2 indexingServiceClient) +// { +// return new DartSqlEngine(indexingServiceClient, queryJsonMapper, new SegmentGenerationTerminalStageSpecFactory()); +// } +// +// @Provides +// @LazySingleton +// public DartSqlEngine makeSqlEngine( +// DartControllerContextFactory controllerContextFactory, +// DartControllerRegistry controllerRegistry, +// DartControllerConfig controllerConfig +// ) +// { +// return new DartSqlEngine( +// controllerContextFactory, +// controllerRegistry, +// controllerConfig, +// Execs.multiThreaded(controllerConfig.getConcurrentQueries(), "dart-controller-%s") +// ); +// } @Provides @LazySingleton - private MSQTestOverlordServiceClient makeOverlordServiceClient( + private MSQTestOverlordServiceClient2 makeOverlordServiceClient( ObjectMapper queryJsonMapper, Injector injector, WorkerMemoryParameters workerMemoryParameters) { - final MSQTestOverlordServiceClient indexingServiceClient = new MSQTestOverlordServiceClient( + final MSQTestOverlordServiceClient2 indexingServiceClient = new MSQTestOverlordServiceClient2( queryJsonMapper, injector, new MSQTestTaskActionClient(queryJsonMapper, injector), diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestOverlordServiceClient2.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestOverlordServiceClient2.java new file mode 100644 index 000000000000..cc05e0ea146e --- /dev/null +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestOverlordServiceClient2.java @@ -0,0 +1,293 @@ +/* + * 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.test; + +import com.fasterxml.jackson.databind.ObjectMapper; +import com.google.common.base.Preconditions; +import com.google.common.collect.Iterables; +import com.google.common.util.concurrent.Futures; +import com.google.common.util.concurrent.ListenableFuture; +import com.google.common.util.concurrent.SettableFuture; +import com.google.inject.Injector; +import org.apache.calcite.sql.type.SqlTypeName; +import org.apache.druid.client.ImmutableSegmentLoadInfo; +import org.apache.druid.client.indexing.TaskPayloadResponse; +import org.apache.druid.client.indexing.TaskStatusResponse; +import org.apache.druid.indexer.TaskStatus; +import org.apache.druid.indexer.TaskStatusPlus; +import org.apache.druid.indexer.report.TaskReport; +import org.apache.druid.indexing.common.actions.TaskActionClient; +import org.apache.druid.java.util.common.DateTimes; +import org.apache.druid.java.util.common.ISE; +import org.apache.druid.msq.exec.Controller; +import org.apache.druid.msq.exec.ControllerImpl; +import org.apache.druid.msq.exec.QueryListener; +import org.apache.druid.msq.exec.ResultsContext; +import org.apache.druid.msq.exec.WorkerMemoryParameters; +import org.apache.druid.msq.indexing.MSQControllerTask; +import org.apache.druid.msq.indexing.destination.MSQDestination; +import org.apache.druid.msq.indexing.report.MSQResultsReport; +import org.apache.druid.msq.indexing.report.MSQStatusReport; +import org.apache.druid.msq.indexing.report.MSQTaskReport; +import org.apache.druid.msq.indexing.report.MSQTaskReportPayload; +import org.apache.druid.rpc.indexing.NoopOverlordClient; +import org.joda.time.DateTime; + +import javax.annotation.Nullable; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +public class MSQTestOverlordServiceClient2 extends NoopOverlordClient +{ + private final Injector injector; + private final ObjectMapper objectMapper; + private final TaskActionClient taskActionClient; + private final WorkerMemoryParameters workerMemoryParameters; + private final List loadedSegmentMetadata; + private final Map inMemoryControllers = new HashMap<>(); + private final Map reports = new HashMap<>(); + private final Map inMemoryControllerTask = new HashMap<>(); + private final Map inMemoryTaskStatus = new HashMap<>(); + + public static final DateTime CREATED_TIME = DateTimes.of("2023-05-31T12:00Z"); + public static final DateTime QUEUE_INSERTION_TIME = DateTimes.of("2023-05-31T12:01Z"); + + public static final long DURATION = 100L; + + public MSQTestOverlordServiceClient2( + ObjectMapper objectMapper, + Injector injector, + TaskActionClient taskActionClient, + WorkerMemoryParameters workerMemoryParameters, + List loadedSegmentMetadata + ) + { + this.objectMapper = objectMapper; + this.injector = injector; + this.taskActionClient = taskActionClient; + this.workerMemoryParameters = workerMemoryParameters; + this.loadedSegmentMetadata = loadedSegmentMetadata; + } + + @Override + public ListenableFuture runTask(String taskId, Object taskObject) + { + TestQueryListener queryListener = null; + ControllerImpl controller = null; + MSQTestControllerContext msqTestControllerContext; + try { + MSQControllerTask cTask = objectMapper.convertValue(taskObject, MSQControllerTask.class); + + msqTestControllerContext = new MSQTestControllerContext( + objectMapper, + injector, + taskActionClient, + workerMemoryParameters, + loadedSegmentMetadata, + cTask.getTaskLockType(), + cTask.getQuerySpec().getContext() + ); + + inMemoryControllerTask.put(cTask.getId(), cTask); + + controller = new ControllerImpl( + cTask.getId(), + cTask.getQuerySpec(), + new ResultsContext(cTask.getSqlTypeNames(), cTask.getSqlResultsContext()), + msqTestControllerContext + ); + + inMemoryControllers.put(controller.queryId(), controller); + + queryListener = + new TestQueryListener( + cTask.getId(), + cTask.getQuerySpec().getDestination() + ); + + try { + controller.run(queryListener); + inMemoryTaskStatus.put(taskId, queryListener.getStatusReport().toTaskStatus(cTask.getId())); + } + catch (Exception e) { + inMemoryTaskStatus.put(taskId, TaskStatus.failure(cTask.getId(), e.toString())); + } + return Futures.immediateFuture(null); + } + catch (Exception e) { + throw new ISE(e, "Unable to run"); + } + finally { + if (controller != null && queryListener != null) { + reports.put(controller.queryId(), queryListener.getReportMap()); + } + } + } + + @Override + public ListenableFuture cancelTask(String taskId) + { + inMemoryControllers.get(taskId).stop(); + return Futures.immediateFuture(null); + } + + @Override + public ListenableFuture taskReportAsMap(String taskId) + { + return Futures.immediateFuture(getReportForTask(taskId)); + } + + @Override + public ListenableFuture taskPayload(String taskId) + { + SettableFuture future = SettableFuture.create(); + future.set(new TaskPayloadResponse(taskId, getMSQControllerTask(taskId))); + return future; + } + + @Override + public ListenableFuture taskStatus(String taskId) + { + SettableFuture future = SettableFuture.create(); + TaskStatus taskStatus = inMemoryTaskStatus.get(taskId); + future.set(new TaskStatusResponse(taskId, new TaskStatusPlus( + taskId, + null, + MSQControllerTask.TYPE, + CREATED_TIME, + QUEUE_INSERTION_TIME, + taskStatus.getStatusCode(), + null, + DURATION, + taskStatus.getLocation(), + null, + taskStatus.getErrorMsg() + ))); + + return future; + } + + // hooks to pull stuff out for testing + @Nullable + public TaskReport.ReportMap getReportForTask(String id) + { + return reports.get(id); + } + + @Nullable + MSQControllerTask getMSQControllerTask(String id) + { + return inMemoryControllerTask.get(id); + } + + /** + * Listener that captures a report and makes it available through {@link #getReportMap()}. + */ + static class TestQueryListener implements QueryListener + { + private final String taskId; + private final MSQDestination destination; + private final List results = new ArrayList<>(); + + private List signature; + private List sqlTypeNames; + private boolean resultsTruncated = true; + private TaskReport.ReportMap reportMap; + + public TestQueryListener(final String taskId, final MSQDestination destination) + { + this.taskId = taskId; + this.destination = destination; + } + + @Override + public boolean readResults() + { + return destination.getRowsInTaskReport() == MSQDestination.UNLIMITED || destination.getRowsInTaskReport() > 0; + } + + @Override + public void onResultsStart(List signature, @Nullable List sqlTypeNames) + { + this.signature = signature; + this.sqlTypeNames = sqlTypeNames; + } + + @Override + public boolean onResultRow(Object[] row) + { + if (destination.getRowsInTaskReport() == MSQDestination.UNLIMITED + || results.size() < destination.getRowsInTaskReport()) { + results.add(row); + return true; + } else { + return false; + } + } + + @Override + public void onResultsComplete() + { + resultsTruncated = false; + } + + @Override + public void onQueryComplete(MSQTaskReportPayload report) + { + final MSQResultsReport resultsReport; + + if (signature != null) { + resultsReport = new MSQResultsReport( + signature, + sqlTypeNames, + results, + resultsTruncated + ); + } else { + resultsReport = null; + } + + final MSQTaskReport taskReport = new MSQTaskReport( + taskId, + new MSQTaskReportPayload( + report.getStatus(), + report.getStages(), + report.getCounters(), + resultsReport + ) + ); + + reportMap = TaskReport.buildTaskReports(taskReport); + } + + public TaskReport.ReportMap getReportMap() + { + return Preconditions.checkNotNull(reportMap, "reportMap"); + } + + public MSQStatusReport getStatusReport() + { + final MSQTaskReport taskReport = (MSQTaskReport) Iterables.getOnlyElement(getReportMap().values()); + return taskReport.getPayload().getStatus(); + } + } +} From 71e14d91805f0879b03877ce757528ae126662d2 Mon Sep 17 00:00:00 2001 From: Zoltan Haindrich Date: Mon, 3 Mar 2025 07:39:40 +0000 Subject: [PATCH 03/50] some-leftover --- .../dart/controller/sql/DartQueryMaker.java | 2 +- .../msq/test/CalciteSelectQueryMSQTest.java | 5 +- .../druid/msq/test/DartComponentSupplier.java | 46 ++++++++++++++++++- .../druid/quidem/DruidAvaticaTestDriver.java | 2 + .../sql/calcite/util/SqlTestFramework.java | 31 +++++++++++++ 5 files changed, 82 insertions(+), 4 deletions(-) diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartQueryMaker.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartQueryMaker.java index 4526754cf573..dee475ff5491 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartQueryMaker.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartQueryMaker.java @@ -141,7 +141,7 @@ public QueryResponse runQuery(DruidQuery druidQuery) final List> types = MSQTaskQueryMaker.getTypes(druidQuery, fieldMapping, plannerContext); - final String dartQueryId = druidQuery.getQuery().context().getString(DartSqlEngine.CTX_DART_QUERY_ID); + final String dartQueryId = druidQuery.getQuery().context().getString(DartSqlEngine.CTX_DART_QUERY_ID) +"asd"; final ControllerContext controllerContext = controllerContextFactory.newContext(dartQueryId); final ControllerImpl controller = new ControllerImpl( dartQueryId, diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/CalciteSelectQueryMSQTest.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/CalciteSelectQueryMSQTest.java index 0d6683496cf5..f19090ad6c0b 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/CalciteSelectQueryMSQTest.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/CalciteSelectQueryMSQTest.java @@ -22,7 +22,6 @@ import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import org.apache.druid.java.util.common.ISE; -import org.apache.druid.msq.sql.MSQTaskSqlEngine; import org.apache.druid.sql.calcite.CalciteQueryTest; import org.apache.druid.sql.calcite.QueryTestBuilder; import org.apache.druid.sql.calcite.SqlTestFrameworkConfig; @@ -31,6 +30,7 @@ import org.junit.jupiter.api.Test; import org.junit.jupiter.api.Timeout; +import java.util.UUID; import java.util.concurrent.TimeUnit; /** @@ -43,7 +43,8 @@ public class CalciteSelectQueryMSQTest extends CalciteQueryTest protected QueryTestBuilder testBuilder() { return new QueryTestBuilder(new CalciteTestConfig(true)) - .addCustomRunner(new ExtractResultsFactory(() -> (MSQTestOverlordServiceClient) ((MSQTaskSqlEngine) queryFramework().engine()).overlordClient())) + .addCustomRunner(new ExtractResultsFactory(() -> null))//(MSQTestOverlordServiceClient) ((DartSqlEngine) queryFramework().engine()).overlordClient())) + .queryContext(ImmutableMap.builder().put("asd",UUID.randomUUID().toString()).build()) .skipVectorize(true) .verifyNativeQueries(new VerifyMSQSupportedNativeQueriesPredicate()); } diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/DartComponentSupplier.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/DartComponentSupplier.java index 49649e9f98ce..dff0c8bd7cd3 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/DartComponentSupplier.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/DartComponentSupplier.java @@ -24,18 +24,30 @@ import com.google.inject.Binder; import com.google.inject.Injector; import com.google.inject.Provides; +import com.google.inject.TypeLiteral; +import org.apache.druid.collections.NonBlockingPool; +import org.apache.druid.discovery.DruidNodeDiscoveryProvider; import org.apache.druid.guice.LazySingleton; +import org.apache.druid.guice.annotations.EscalatedGlobal; +import org.apache.druid.guice.annotations.Merging; import org.apache.druid.initialization.DruidModule; +import org.apache.druid.java.util.http.client.HttpClient; import org.apache.druid.msq.dart.controller.sql.DartSqlEngine; import org.apache.druid.msq.dart.guice.DartControllerModule; +import org.apache.druid.msq.dart.guice.DartWorkerMemoryManagementModule; import org.apache.druid.msq.dart.guice.DartWorkerModule; import org.apache.druid.msq.exec.WorkerMemoryParameters; +import org.apache.druid.query.TestBufferPool; +import org.apache.druid.rpc.ServiceClientFactory; +import org.apache.druid.rpc.guice.ServiceClientModule; import org.apache.druid.server.QueryLifecycleFactory; import org.apache.druid.sql.calcite.TempDirProducer; import org.apache.druid.sql.calcite.run.SqlEngine; import org.apache.druid.sql.calcite.util.DruidModuleCollection; import org.apache.druid.sql.calcite.util.SqlTestFramework.StandardComponentSupplier; +import java.nio.ByteBuffer; + public class DartComponentSupplier extends AbstractMSQComponentSupplierDelegate { public DartComponentSupplier(TempDirProducer tempFolderProducer) @@ -56,7 +68,9 @@ public DruidModule getCoreModule() { return DruidModuleCollection.of(super.getCoreModule(), new DartControllerModule(), - new DartWorkerModule() + new DartWorkerModule(), + new DartWorkerMemoryManagementModule(), + new L1() ); } @@ -70,11 +84,41 @@ public SqlEngine createEngine( } + + static class L1 implements DruidModule{ + + @Provides + @EscalatedGlobal + final ServiceClientFactory getServiceClientFactory(HttpClient ht) { + return ServiceClientModule.makeServiceClientFactory(ht); + + } + + @Provides + final DruidNodeDiscoveryProvider getDiscoveryProvider() { + return null; + } + + + @Override + public void configure(Binder binder) + { +// binder.bind(new TypeLiteral>(){}) +// .annotatedWith(Merging.class) +// .to(TestBufferPool.class); + + } + + } static class LocalDartModule implements DruidModule{ @Override public void configure(Binder binder) { + binder.bind(new TypeLiteral>(){}) + .annotatedWith(Merging.class) + .to(TestBufferPool.class); + } // @Provides diff --git a/sql/src/test/java/org/apache/druid/quidem/DruidAvaticaTestDriver.java b/sql/src/test/java/org/apache/druid/quidem/DruidAvaticaTestDriver.java index ce575b52356f..16f8c87671d2 100644 --- a/sql/src/test/java/org/apache/druid/quidem/DruidAvaticaTestDriver.java +++ b/sql/src/test/java/org/apache/druid/quidem/DruidAvaticaTestDriver.java @@ -183,6 +183,7 @@ protected AbstractAvaticaHandler getAvaticaHandler(final DruidMeta druidMeta) { return new DruidAvaticaJsonHandler( druidMeta, + //FIXME new DruidNode("dummy", "dummy", false, 1, null, true, false), new AvaticaMonitor() ); @@ -205,6 +206,7 @@ public DruidModule getCoreModule() return DruidModuleCollection.of( super.getCoreModule(), binder -> { + // FIXME binder.bindConstant().annotatedWith(Names.named("serviceName")).to("test"); binder.bindConstant().annotatedWith(Names.named("servicePort")).to(0); binder.bindConstant().annotatedWith(Names.named("tlsServicePort")).to(-1); diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/util/SqlTestFramework.java b/sql/src/test/java/org/apache/druid/sql/calcite/util/SqlTestFramework.java index 90cc53d1b9a0..1fb05690f8b6 100644 --- a/sql/src/test/java/org/apache/druid/sql/calcite/util/SqlTestFramework.java +++ b/sql/src/test/java/org/apache/druid/sql/calcite/util/SqlTestFramework.java @@ -29,8 +29,12 @@ import com.google.inject.Module; import com.google.inject.Provides; import com.google.inject.TypeLiteral; +import org.apache.druid.client.SimpleServerView; +import org.apache.druid.client.TestHttpClient; +import org.apache.druid.client.TimelineServerView; import org.apache.druid.client.cache.Cache; import org.apache.druid.client.cache.CacheConfig; +import org.apache.druid.collections.BlockingPool; import org.apache.druid.collections.NonBlockingPool; import org.apache.druid.guice.BuiltInTypesModule; import org.apache.druid.guice.DruidInjectorBuilder; @@ -44,12 +48,14 @@ import org.apache.druid.guice.StartupInjectorBuilder; import org.apache.druid.guice.annotations.Global; import org.apache.druid.guice.annotations.Merging; +import org.apache.druid.guice.annotations.Self; import org.apache.druid.initialization.CoreInjectorBuilder; import org.apache.druid.initialization.DruidModule; import org.apache.druid.initialization.ServiceInjectorBuilder; import org.apache.druid.java.util.common.RE; import org.apache.druid.java.util.common.io.Closer; import org.apache.druid.java.util.emitter.service.ServiceEmitter; +import org.apache.druid.java.util.http.client.HttpClient; import org.apache.druid.math.expr.ExprMacroTable; import org.apache.druid.query.DefaultGenericQueryMetricsFactory; import org.apache.druid.query.DruidProcessingConfig; @@ -78,6 +84,7 @@ import org.apache.druid.segment.realtime.ChatHandlerProvider; import org.apache.druid.segment.realtime.NoopChatHandlerProvider; import org.apache.druid.server.ClientQuerySegmentWalker; +import org.apache.druid.server.DruidNode; import org.apache.druid.server.LocalQuerySegmentWalker; import org.apache.druid.server.QueryLifecycle; import org.apache.druid.server.QueryLifecycleFactory; @@ -771,6 +778,14 @@ public void configure(Binder binder) .annotatedWith(Global.class) .to(TestBufferPool.class); + binder.bind(new TypeLiteral>(){}) + .annotatedWith(Merging.class) + .to(TestBufferPool.class); + + // FIXME + DruidNode self = new DruidNode("test", "test-host", true, 80, 443, false, true); + binder.bind(DruidNode.class).annotatedWith(Self.class).toInstance(self); + TestRequestLogger testRequestLogger = new TestRequestLogger(); binder.bind(RequestLogger.class).toInstance(testRequestLogger); } @@ -1039,6 +1054,22 @@ public TestSegmentsBroker makeTimelines() return new TestSegmentsBroker(); } + @Provides + @LazySingleton + public TimelineServerView makeTimelines1(ObjectMapper objectMapper, QueryRunnerFactoryConglomerate conglomerate, + HttpClient httpClient) + { + //FIXME connect TSB with this + return new SimpleServerView(conglomerate, objectMapper, httpClient); + } + + @Provides + @LazySingleton + private HttpClient extracted(ObjectMapper objectMapper) + { + return new TestHttpClient(objectMapper); + } + @Provides @Named("empty") @LazySingleton From 064e362b29f2e4d42e8fcdbddeac7ec4325320ce Mon Sep 17 00:00:00 2001 From: Zoltan Haindrich Date: Mon, 3 Mar 2025 14:59:36 +0000 Subject: [PATCH 04/50] fixes --- .../controller/DartControllerContext.java | 3 ++ .../apache/druid/msq/exec/ControllerImpl.java | 1 + .../druid/sql/calcite/util/CalciteTests.java | 1 + .../sql/calcite/util/QueryFrameworkUtils.java | 1 + .../sql/calcite/util/SqlTestFramework.java | 40 +++++++++++++++++-- .../calcite/util/TestTimelineServerView.java | 9 +++-- 6 files changed, 48 insertions(+), 7 deletions(-) diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartControllerContext.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartControllerContext.java index c2257d160ed3..258ff6d7fa70 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartControllerContext.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartControllerContext.java @@ -121,6 +121,9 @@ public ControllerQueryKernelConfig queryKernelConfig( // allowed to float. If a segment moves to a new server that isn't part of our list after the WorkerManager is // created, we won't be able to find a valid server for certain segments. This isn't expected to be a problem, // since the serverView is referenced shortly after the worker list is created. + if(true) { +// throw DruidException.defensive("asf|"); + } final List workerIds = new ArrayList<>(servers.size()); for (final DruidServerMetadata server : servers) { if (server.getType() == ServerType.HISTORICAL) { diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/ControllerImpl.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/ControllerImpl.java index b370ab586433..dd4e924fe09e 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/ControllerImpl.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/ControllerImpl.java @@ -368,6 +368,7 @@ private MSQTaskReportPayload runInternal(final QueryListener queryListener, fina handleQueryResults(queryDef, queryKernel); } catch (Throwable e) { + log.info(e, "Execution failed"); exceptionEncountered = e; } diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/util/CalciteTests.java b/sql/src/test/java/org/apache/druid/sql/calcite/util/CalciteTests.java index abf32c9bf321..0ce02142e96f 100644 --- a/sql/src/test/java/org/apache/druid/sql/calcite/util/CalciteTests.java +++ b/sql/src/test/java/org/apache/druid/sql/calcite/util/CalciteTests.java @@ -474,6 +474,7 @@ private TaskStatusPlus createTaskStatus(String id, String datasource, Long durat new BrokerSegmentWatcherConfig(), BrokerSegmentMetadataCacheConfig.create() ), + //FIXME use the same new TestTimelineServerView(walker.getSegments()), new FakeServerInventoryView(), authorizerMapper, diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/util/QueryFrameworkUtils.java b/sql/src/test/java/org/apache/druid/sql/calcite/util/QueryFrameworkUtils.java index f28461e1711c..6a997fca9437 100644 --- a/sql/src/test/java/org/apache/druid/sql/calcite/util/QueryFrameworkUtils.java +++ b/sql/src/test/java/org/apache/druid/sql/calcite/util/QueryFrameworkUtils.java @@ -226,6 +226,7 @@ static DruidSchema createMockSchema( { final BrokerSegmentMetadataCache cache = new BrokerSegmentMetadataCache( createMockQueryLifecycleFactory(walker, conglomerate), + //FIXME use the same new TestTimelineServerView(walker.getSegments()), BrokerSegmentMetadataCacheConfig.create(), CalciteTests.TEST_AUTHENTICATOR_ESCALATOR, diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/util/SqlTestFramework.java b/sql/src/test/java/org/apache/druid/sql/calcite/util/SqlTestFramework.java index 1fb05690f8b6..021f0b394b80 100644 --- a/sql/src/test/java/org/apache/druid/sql/calcite/util/SqlTestFramework.java +++ b/sql/src/test/java/org/apache/druid/sql/calcite/util/SqlTestFramework.java @@ -29,6 +29,13 @@ import com.google.inject.Module; import com.google.inject.Provides; import com.google.inject.TypeLiteral; +import org.apache.druid.client.BrokerServerView; +import org.apache.druid.client.DirectDruidClientFactory; +import org.apache.druid.client.FilteredServerInventoryView; +import org.apache.druid.client.FilteredServerInventoryViewProvider; +import org.apache.druid.client.QueryableDruidServer; +import org.apache.druid.client.ServerInventoryView; +import org.apache.druid.client.ServerInventoryViewProvider; import org.apache.druid.client.SimpleServerView; import org.apache.druid.client.TestHttpClient; import org.apache.druid.client.TimelineServerView; @@ -41,6 +48,7 @@ import org.apache.druid.guice.ExpressionModule; import org.apache.druid.guice.LazySingleton; import org.apache.druid.guice.LifecycleModule; +import org.apache.druid.guice.ManageLifecycle; import org.apache.druid.guice.QueryRunnerFactoryModule; import org.apache.druid.guice.QueryableModule; import org.apache.druid.guice.SegmentWranglerModule; @@ -127,7 +135,6 @@ import org.apache.druid.sql.hook.DruidHookDispatcher; import org.apache.druid.timeline.DataSegment; import org.apache.druid.utils.JvmUtils; - import javax.inject.Named; import java.io.Closeable; @@ -1054,15 +1061,27 @@ public TestSegmentsBroker makeTimelines() return new TestSegmentsBroker(); } - @Provides @LazySingleton public TimelineServerView makeTimelines1(ObjectMapper objectMapper, QueryRunnerFactoryConglomerate conglomerate, HttpClient httpClient) { //FIXME connect TSB with this - return new SimpleServerView(conglomerate, objectMapper, httpClient); +///. return new SimpleServerView(conglomerate, objectMapper, httpClient); + + SimpleServerView simpleServerView = new SimpleServerView(conglomerate, objectMapper, httpClient); + return simpleServerView; + } + @Provides + @LazySingleton + public TimelineServerView makeTimelines2(SpecificSegmentsQuerySegmentWalker walker, QueryRunnerFactoryConglomerate conglomerate, + HttpClient httpClient) + { + List realtimeSegments = Collections.emptyList(); + return new TestTimelineServerView(walker.getSegments(), realtimeSegments); } + + @Provides @LazySingleton private HttpClient extracted(ObjectMapper objectMapper) @@ -1122,6 +1141,21 @@ public QueryScheduler makeQueryScheduler() @Override public void configure(Binder binder) { + if(true) { + if(false) { + binder.bind(TimelineServerView.class).to(SimpleServerView.class).in(LazySingleton.class); + } + + } else { + binder.bind(TimelineServerView.class).to(BrokerServerView.class).in(LazySingleton.class); + binder.bind(QueryableDruidServer.Maker.class).to(DirectDruidClientFactory.class).in(LazySingleton.class); + binder.bind(ServerInventoryView.class).toProvider(ServerInventoryViewProvider.class).in(ManageLifecycle.class); + binder.bind(FilteredServerInventoryView.class) + .toProvider(FilteredServerInventoryViewProvider.class) + .in(ManageLifecycle.class); + + // return new SimpleServerView(conglomerate, objectMapper, httpClient); + } } } diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/util/TestTimelineServerView.java b/sql/src/test/java/org/apache/druid/sql/calcite/util/TestTimelineServerView.java index 42aa4e491cb4..4e4c677de86e 100644 --- a/sql/src/test/java/org/apache/druid/sql/calcite/util/TestTimelineServerView.java +++ b/sql/src/test/java/org/apache/druid/sql/calcite/util/TestTimelineServerView.java @@ -47,7 +47,7 @@ public class TestTimelineServerView implements TimelineServerView { private static final DruidServerMetadata DUMMY_SERVER = new DruidServerMetadata( "dummy", - "dummy", + "dummy:15723", null, 0, ServerType.HISTORICAL, @@ -56,7 +56,7 @@ public class TestTimelineServerView implements TimelineServerView ); private static final DruidServerMetadata DUMMY_SERVER_REALTIME = new DruidServerMetadata( "dummy2", - "dummy2", + "dummy2:15723", null, 0, ServerType.REALTIME, @@ -65,7 +65,7 @@ public class TestTimelineServerView implements TimelineServerView ); private static final DruidServerMetadata DUMMY_BROKER = new DruidServerMetadata( "dummy3", - "dummy3", + "dummy3:15723", null, 0, ServerType.BROKER, @@ -93,7 +93,8 @@ public TestTimelineServerView(List segments, List real @Override public Optional> getTimeline(TableDataSource table) { - throw new UnsupportedOperationException(); + return Optional.empty(); +// throw new UnsupportedOperationException(); } @Override From 5aa1836a0d5b79440afb00f3f46b8dac7b0c9bcc Mon Sep 17 00:00:00 2001 From: Zoltan Haindrich Date: Tue, 4 Mar 2025 08:22:07 +0000 Subject: [PATCH 05/50] a --- .../DartControllerContextFactoryImpl.java | 8 ++- .../druid/msq/test/DartComponentSupplier.java | 4 ++ .../TestDartControllerContextFactoryImpl.java | 59 +++++++++++++++++++ 3 files changed, 69 insertions(+), 2 deletions(-) create mode 100644 extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/TestDartControllerContextFactoryImpl.java diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartControllerContextFactoryImpl.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartControllerContextFactoryImpl.java index 0a69a72fba23..a7b1e33abc70 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartControllerContextFactoryImpl.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartControllerContextFactoryImpl.java @@ -22,7 +22,6 @@ import com.fasterxml.jackson.databind.ObjectMapper; import com.google.inject.Inject; import com.google.inject.Injector; -import org.apache.druid.client.BrokerServerView; import org.apache.druid.client.TimelineServerView; import org.apache.druid.guice.annotations.EscalatedGlobal; import org.apache.druid.guice.annotations.Json; @@ -75,10 +74,15 @@ public ControllerContext newContext(final String queryId) injector, jsonMapper, selfNode, - new DartWorkerClient(queryId, serviceClientFactory, smileMapper, selfNode.getHostAndPortToUse()), + makeWorkerClient(queryId), memoryIntrospector, serverView, emitter ); } + + protected DartWorkerClient makeWorkerClient(final String queryId) + { + return new DartWorkerClient(queryId, serviceClientFactory, smileMapper, selfNode.getHostAndPortToUse()); + } } diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/DartComponentSupplier.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/DartComponentSupplier.java index dff0c8bd7cd3..9e8c03fef0be 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/DartComponentSupplier.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/DartComponentSupplier.java @@ -32,6 +32,7 @@ import org.apache.druid.guice.annotations.Merging; import org.apache.druid.initialization.DruidModule; import org.apache.druid.java.util.http.client.HttpClient; +import org.apache.druid.msq.dart.controller.DartControllerContextFactory; import org.apache.druid.msq.dart.controller.sql.DartSqlEngine; import org.apache.druid.msq.dart.guice.DartControllerModule; import org.apache.druid.msq.dart.guice.DartWorkerMemoryManagementModule; @@ -119,6 +120,9 @@ public void configure(Binder binder) .annotatedWith(Merging.class) .to(TestBufferPool.class); + binder.bind(DartControllerContextFactory.class) + .to(TestDartControllerContextFactoryImpl.class) + .in(LazySingleton.class); } // @Provides diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/TestDartControllerContextFactoryImpl.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/TestDartControllerContextFactoryImpl.java new file mode 100644 index 000000000000..0acc0898285e --- /dev/null +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/TestDartControllerContextFactoryImpl.java @@ -0,0 +1,59 @@ +/* + * 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.test; + +import com.fasterxml.jackson.databind.ObjectMapper; +import com.google.inject.Inject; +import com.google.inject.Injector; +import org.apache.druid.client.TimelineServerView; +import org.apache.druid.guice.annotations.EscalatedGlobal; +import org.apache.druid.guice.annotations.Json; +import org.apache.druid.guice.annotations.Self; +import org.apache.druid.guice.annotations.Smile; +import org.apache.druid.java.util.emitter.service.ServiceEmitter; +import org.apache.druid.msq.dart.controller.DartControllerContextFactoryImpl; +import org.apache.druid.msq.dart.worker.DartWorkerClient; +import org.apache.druid.msq.exec.MemoryIntrospector; +import org.apache.druid.rpc.ServiceClientFactory; +import org.apache.druid.server.DruidNode; + +public class TestDartControllerContextFactoryImpl extends DartControllerContextFactoryImpl +{ + @Inject + public TestDartControllerContextFactoryImpl( + final Injector injector, + @Json final ObjectMapper jsonMapper, + @Smile final ObjectMapper smileMapper, + @Self final DruidNode selfNode, + @EscalatedGlobal final ServiceClientFactory serviceClientFactory, + final MemoryIntrospector memoryIntrospector, + final TimelineServerView serverView, + final ServiceEmitter emitter) + { + super(injector, jsonMapper, smileMapper, selfNode, serviceClientFactory, memoryIntrospector, serverView, emitter); + } + + @Override + protected DartWorkerClient makeWorkerClient(String queryId) + { + return super.makeWorkerClient(queryId); + } + +} From c899e7ab68573a949288c667008f57cba0048f50 Mon Sep 17 00:00:00 2001 From: Zoltan Haindrich Date: Tue, 4 Mar 2025 08:24:06 +0000 Subject: [PATCH 06/50] move stuff --- .../apache/druid/msq/exec/ControllerImpl.java | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/ControllerImpl.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/ControllerImpl.java index dd4e924fe09e..2394f367b01f 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/ControllerImpl.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/ControllerImpl.java @@ -350,6 +350,15 @@ private MSQTaskReportPayload runInternal(final QueryListener queryListener, fina context.registerController(this, closer); queryDef = initializeQueryDefAndState(closer); + this.netClient = closer.register(new ExceptionWrappingWorkerClient(context.newWorkerClient())); + this.workerSketchFetcher = new WorkerSketchFetcher( + netClient, + workerManager, + queryKernelConfig.isFaultTolerant(), + MultiStageQueryContext.getSketchEncoding(querySpec.getContext()) + ); + closer.register(workerSketchFetcher::close); + // Execution-related: run the multi-stage QueryDefinition. final InputSpecSlicerFactory inputSpecSlicerFactory = makeInputSpecSlicerFactory(context.newTableInputSpecSlicer(workerManager)); @@ -590,7 +599,6 @@ public void addToKernelManipulationQueue(Consumer kernelC private QueryDefinition initializeQueryDefAndState(final Closer closer) { this.selfDruidNode = context.selfNode(); - this.netClient = closer.register(new ExceptionWrappingWorkerClient(context.newWorkerClient())); this.queryKernelConfig = context.queryKernelConfig(queryId, querySpec); final QueryContext queryContext = querySpec.getContext(); @@ -654,13 +662,7 @@ private QueryDefinition initializeQueryDefAndState(final Closer closer) ) ) ); - this.workerSketchFetcher = new WorkerSketchFetcher( - netClient, - workerManager, - queryKernelConfig.isFaultTolerant(), - MultiStageQueryContext.getSketchEncoding(queryContext) - ); - closer.register(workerSketchFetcher::close); + return queryDef; } From 0543df696612c837e2da943a49e346cacdc1f0a0 Mon Sep 17 00:00:00 2001 From: Zoltan Haindrich Date: Tue, 4 Mar 2025 11:40:45 +0000 Subject: [PATCH 07/50] my-wood --- .../DartControllerContextFactoryImpl.java | 18 +- .../msq/dart/worker/DartWorkerClient.java | 2 +- .../druid/msq/test/DartComponentSupplier.java | 7 + .../msq/test/MSQTestControllerContext.java | 9 +- .../TestDartControllerContextFactoryImpl.java | 156 +++++++++++++++++- 5 files changed, 180 insertions(+), 12 deletions(-) diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartControllerContextFactoryImpl.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartControllerContextFactoryImpl.java index a7b1e33abc70..b0d2b66d0907 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartControllerContextFactoryImpl.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartControllerContextFactoryImpl.java @@ -36,14 +36,14 @@ public class DartControllerContextFactoryImpl implements DartControllerContextFactory { - private final Injector injector; - private final ObjectMapper jsonMapper; - private final ObjectMapper smileMapper; - private final DruidNode selfNode; - private final ServiceClientFactory serviceClientFactory; - private final TimelineServerView serverView; - private final MemoryIntrospector memoryIntrospector; - private final ServiceEmitter emitter; + protected final Injector injector; + protected final ObjectMapper jsonMapper; + protected final ObjectMapper smileMapper; + protected final DruidNode selfNode; + protected final ServiceClientFactory serviceClientFactory; + protected final TimelineServerView serverView; + protected final MemoryIntrospector memoryIntrospector; + protected final ServiceEmitter emitter; @Inject public DartControllerContextFactoryImpl( @@ -83,6 +83,6 @@ public ControllerContext newContext(final String queryId) protected DartWorkerClient makeWorkerClient(final String queryId) { - return new DartWorkerClient(queryId, serviceClientFactory, smileMapper, selfNode.getHostAndPortToUse()); + return new DartWorkerClient(queryId, serviceClientFactory, smileMapper, selfNode.getHostAndPortToUse()); } } diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartWorkerClient.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartWorkerClient.java index 932300de217f..3bd21e970a0e 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartWorkerClient.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartWorkerClient.java @@ -158,7 +158,7 @@ public ListenableFuture stopWorker(String workerId) /** * Create a new client. Called by {@link #getClient(String)} if a new one is needed. */ - private Pair makeNewClient(final WorkerId workerId) + protected Pair makeNewClient(final WorkerId workerId) { final URI uri = workerId.toUri(); final FixedServiceLocator locator = new FixedServiceLocator(ServiceLocation.fromUri(uri)); diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/DartComponentSupplier.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/DartComponentSupplier.java index 9e8c03fef0be..1c7965b613a9 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/DartComponentSupplier.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/DartComponentSupplier.java @@ -120,9 +120,16 @@ public void configure(Binder binder) .annotatedWith(Merging.class) .to(TestBufferPool.class); + if(true) { binder.bind(DartControllerContextFactory.class) .to(TestDartControllerContextFactoryImpl.class) .in(LazySingleton.class); + }else { + binder.bind(DartControllerContextFactory.class) + .to(MSQTestControllerContext.class) + .in(LazySingleton.class); + + } } // @Provides diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestControllerContext.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestControllerContext.java index 6e632a4869ca..e1755757d04f 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestControllerContext.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestControllerContext.java @@ -44,6 +44,7 @@ import org.apache.druid.java.util.common.concurrent.Execs; import org.apache.druid.java.util.common.io.Closer; import org.apache.druid.java.util.common.logger.Logger; +import org.apache.druid.msq.dart.controller.DartControllerContextFactory; import org.apache.druid.msq.exec.Controller; import org.apache.druid.msq.exec.ControllerContext; import org.apache.druid.msq.exec.ControllerMemoryParameters; @@ -84,7 +85,7 @@ import java.util.concurrent.ConcurrentMap; import java.util.stream.Collectors; -public class MSQTestControllerContext implements ControllerContext +public class MSQTestControllerContext implements ControllerContext, DartControllerContextFactory { private static final Logger log = new Logger(MSQTestControllerContext.class); private static final int NUM_WORKERS = 4; @@ -387,4 +388,10 @@ public WorkerClient newWorkerClient() { return new MSQTestWorkerClient(inMemoryWorkers); } + + @Override + public ControllerContext newContext(String queryId) + { + return this; + } } diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/TestDartControllerContextFactoryImpl.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/TestDartControllerContextFactoryImpl.java index 0acc0898285e..44d0acdf1d89 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/TestDartControllerContextFactoryImpl.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/TestDartControllerContextFactoryImpl.java @@ -20,20 +20,38 @@ package org.apache.druid.msq.test; import com.fasterxml.jackson.databind.ObjectMapper; +import com.google.common.util.concurrent.ListenableFuture; import com.google.inject.Inject; import com.google.inject.Injector; +import it.unimi.dsi.fastutil.Pair; import org.apache.druid.client.TimelineServerView; +import org.apache.druid.frame.channel.ReadableByteChunksFrameChannel; +import org.apache.druid.frame.key.ClusterByPartitions; import org.apache.druid.guice.annotations.EscalatedGlobal; import org.apache.druid.guice.annotations.Json; import org.apache.druid.guice.annotations.Self; import org.apache.druid.guice.annotations.Smile; import org.apache.druid.java.util.emitter.service.ServiceEmitter; +import org.apache.druid.java.util.http.client.response.HttpResponseHandler; +import org.apache.druid.msq.counters.CounterSnapshotsTree; import org.apache.druid.msq.dart.controller.DartControllerContextFactoryImpl; import org.apache.druid.msq.dart.worker.DartWorkerClient; +import org.apache.druid.msq.dart.worker.WorkerId; import org.apache.druid.msq.exec.MemoryIntrospector; +import org.apache.druid.msq.exec.Worker; +import org.apache.druid.msq.kernel.StageId; +import org.apache.druid.msq.kernel.WorkOrder; +import org.apache.druid.msq.rpc.SketchEncoding; +import org.apache.druid.msq.statistics.ClusterByStatisticsSnapshot; +import org.apache.druid.rpc.RequestBuilder; +import org.apache.druid.rpc.ServiceClient; import org.apache.druid.rpc.ServiceClientFactory; +import org.apache.druid.rpc.ServiceRetryPolicy; import org.apache.druid.server.DruidNode; +import java.io.Closeable; +import java.util.Map; + public class TestDartControllerContextFactoryImpl extends DartControllerContextFactoryImpl { @Inject @@ -53,7 +71,143 @@ public TestDartControllerContextFactoryImpl( @Override protected DartWorkerClient makeWorkerClient(String queryId) { - return super.makeWorkerClient(queryId); + return + // super.makeWorkerClient(queryId); + + new DartTestWorkerClient(queryId, serviceClientFactory, smileMapper, selfNode.getHostAndPortToUse()); + // new HashMap()); + } + + static class DartTestWorkerClient extends DartWorkerClient + { + private Map m; + private MSQTestWorkerClient wc; + + public DartTestWorkerClient(String queryId, ServiceClientFactory clientFactory, ObjectMapper smileMapper, + String controllerHost) + { + super(queryId, clientFactory, smileMapper, controllerHost); + wc = new MSQTestWorkerClient(m); + } + + @Override + protected ServiceClient getClient(String workerIdString) + { + return super.getClient(workerIdString); + } + + @Override + public void closeClient(String workerHost) + { + super.closeClient(workerHost); + } + @Override + protected Object clone() throws CloneNotSupportedException + { + return super.clone(); + } + @Override + public ListenableFuture stopWorker(String workerId) + { + return super.stopWorker(workerId); + } + @Override + public ListenableFuture fetchChannelData(String workerId, StageId stageId, int partitionNumber, + long offset, ReadableByteChunksFrameChannel channel) + { + return super.fetchChannelData(workerId, stageId, partitionNumber, offset, channel); + } + + @Override + public ListenableFuture fetchClusterByStatisticsSnapshot(String workerId, + StageId stageId, SketchEncoding sketchEncoding) + { + return super.fetchClusterByStatisticsSnapshot(workerId, stageId, sketchEncoding); + + } + @Override + public ListenableFuture fetchClusterByStatisticsSnapshotForTimeChunk(String workerId, + StageId stageId, long timeChunk, SketchEncoding sketchEncoding) + { + return super.fetchClusterByStatisticsSnapshotForTimeChunk(workerId, stageId, timeChunk, sketchEncoding); + + } + @Override + public ListenableFuture getCounters(String workerId) + { + return super.getCounters(workerId); + + } + @Override + public ListenableFuture postCleanupStage(String workerId, StageId stageId) + { + return super.postCleanupStage(workerId, stageId); + + } + @Override + public ListenableFuture postFinish(String workerId) + { + return super.postFinish(workerId); + + } + + @Override + public ListenableFuture postResultPartitionBoundaries(String workerId, StageId stageId, + ClusterByPartitions partitionBoundaries) + { + return super.postResultPartitionBoundaries(workerId, stageId, partitionBoundaries); + + } + @Override + public ListenableFuture postWorkOrder(String workerId, WorkOrder workOrder) + { + return super.postWorkOrder(workerId, workOrder); + + } + + @Override + protected Pair makeNewClient(WorkerId workerId) + { + ServiceClient client =new TestDartServiceClient(workerId); + return super.makeNewClient(workerId); + } + + } + static class TestDartServiceClient implements ServiceClient { + + public TestDartServiceClient(WorkerId workerId) + { + if(true) + { + throw new RuntimeException("FIXME: Unimplemented!"); + } + + } + + @Override + public ListenableFuture asyncRequest(RequestBuilder requestBuilder, + HttpResponseHandler handler) + { + if(true) + { + throw new RuntimeException("FIXME: Unimplemented!"); + } + return null; + + } + + @Override + public ServiceClient withRetryPolicy(ServiceRetryPolicy retryPolicy) + { + if(true) + { + throw new RuntimeException("FIXME: Unimplemented!"); + } + return null; + + } + + } } From 99ccba03898a1d2fb22c0814adc97780da4bea73 Mon Sep 17 00:00:00 2001 From: Zoltan Haindrich Date: Tue, 4 Mar 2025 12:44:12 +0000 Subject: [PATCH 08/50] rename/etc --- .../DartControllerContextFactoryImpl.java | 6 +- .../msq/dart/worker/DartWorkerClient.java | 183 +------------ .../msq/dart/worker/DartWorkerClientImpl.java | 210 ++++++++++++++ .../dart/worker/DartWorkerFactoryImpl.java | 2 +- .../dart/worker/DartWorkerRetryPolicy.java | 2 +- .../controller/DartWorkerManagerTest.java | 4 +- .../TestDartControllerContextFactoryImpl.java | 259 ++++++++++-------- 7 files changed, 359 insertions(+), 307 deletions(-) create mode 100644 extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartWorkerClientImpl.java diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartControllerContextFactoryImpl.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartControllerContextFactoryImpl.java index b0d2b66d0907..40934f73e122 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartControllerContextFactoryImpl.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartControllerContextFactoryImpl.java @@ -29,6 +29,7 @@ import org.apache.druid.guice.annotations.Smile; import org.apache.druid.java.util.emitter.service.ServiceEmitter; import org.apache.druid.msq.dart.worker.DartWorkerClient; +import org.apache.druid.msq.dart.worker.DartWorkerClientImpl; import org.apache.druid.msq.exec.ControllerContext; import org.apache.druid.msq.exec.MemoryIntrospector; import org.apache.druid.rpc.ServiceClientFactory; @@ -54,8 +55,7 @@ public DartControllerContextFactoryImpl( @EscalatedGlobal final ServiceClientFactory serviceClientFactory, final MemoryIntrospector memoryIntrospector, final TimelineServerView serverView, - final ServiceEmitter emitter - ) + final ServiceEmitter emitter) { this.injector = injector; this.jsonMapper = jsonMapper; @@ -83,6 +83,6 @@ public ControllerContext newContext(final String queryId) protected DartWorkerClient makeWorkerClient(final String queryId) { - return new DartWorkerClient(queryId, serviceClientFactory, smileMapper, selfNode.getHostAndPortToUse()); + return new DartWorkerClientImpl(queryId, serviceClientFactory, smileMapper, selfNode.getHostAndPortToUse()); } } diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartWorkerClient.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartWorkerClient.java index 3bd21e970a0e..5c2f623755b2 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartWorkerClient.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartWorkerClient.java @@ -19,192 +19,15 @@ package org.apache.druid.msq.dart.worker; -import com.fasterxml.jackson.databind.ObjectMapper; -import com.fasterxml.jackson.jaxrs.smile.SmileMediaTypes; import com.google.common.util.concurrent.ListenableFuture; -import com.google.errorprone.annotations.concurrent.GuardedBy; -import it.unimi.dsi.fastutil.Pair; -import org.apache.druid.error.DruidException; -import org.apache.druid.java.util.common.logger.Logger; -import org.apache.druid.java.util.http.client.response.HttpResponseHandler; -import org.apache.druid.msq.dart.controller.DartWorkerManager; -import org.apache.druid.msq.dart.controller.sql.DartSqlEngine; -import org.apache.druid.msq.dart.worker.http.DartWorkerResource; import org.apache.druid.msq.exec.WorkerClient; -import org.apache.druid.msq.rpc.BaseWorkerClientImpl; -import org.apache.druid.rpc.FixedServiceLocator; -import org.apache.druid.rpc.IgnoreHttpResponseHandler; -import org.apache.druid.rpc.RequestBuilder; -import org.apache.druid.rpc.ServiceClient; -import org.apache.druid.rpc.ServiceClientFactory; -import org.apache.druid.rpc.ServiceLocation; -import org.apache.druid.rpc.ServiceRetryPolicy; -import org.apache.druid.utils.CloseableUtils; -import org.jboss.netty.handler.codec.http.HttpMethod; -import javax.annotation.Nullable; -import java.io.Closeable; -import java.net.URI; -import java.util.HashMap; -import java.util.Map; - -/** - * Dart implementation of {@link WorkerClient}. Uses the same {@link BaseWorkerClientImpl} as the task-based engine. - * Each instance of this class is scoped to a single query. - */ -public class DartWorkerClient extends BaseWorkerClientImpl +public interface DartWorkerClient extends WorkerClient { - private static final Logger log = new Logger(DartWorkerClient.class); - - private final String queryId; - private final ServiceClientFactory clientFactory; - private final ServiceRetryPolicy retryPolicy; - - @Nullable - private final String controllerHost; - - @GuardedBy("clientMap") - private final Map> clientMap = new HashMap<>(); - - /** - * Create a worker client. - * - * @param queryId dart query ID. see {@link DartSqlEngine#CTX_DART_QUERY_ID} - * @param clientFactory service client factor - * @param smileMapper Smile object mapper - * @param controllerHost Controller host (see {@link DartWorkerResource#HEADER_CONTROLLER_HOST}) if this is a - * controller-to-worker client. Null if this is a worker-to-worker client. - */ - public DartWorkerClient( - final String queryId, - final ServiceClientFactory clientFactory, - final ObjectMapper smileMapper, - @Nullable final String controllerHost - ) - { - super(smileMapper, SmileMediaTypes.APPLICATION_JACKSON_SMILE); - this.queryId = queryId; - this.clientFactory = clientFactory; - this.controllerHost = controllerHost; - - if (controllerHost == null) { - // worker -> worker client. Retry HTTP 503 in case worker A starts up before worker B, and needs to - // contact it immediately. - this.retryPolicy = new DartWorkerRetryPolicy(true); - } else { - // controller -> worker client. Do not retry any HTTP error codes. If we retry HTTP 503 for controller -> worker, - // we can get stuck trying to contact workers that have exited. - this.retryPolicy = new DartWorkerRetryPolicy(false); - } - } - - @Override - protected ServiceClient getClient(final String workerIdString) - { - final WorkerId workerId = WorkerId.fromString(workerIdString); - if (!queryId.equals(workerId.getQueryId())) { - throw DruidException.defensive("Unexpected queryId[%s]. Expected queryId[%s]", workerId.getQueryId(), queryId); - } - - synchronized (clientMap) { - return clientMap.computeIfAbsent(workerId.getHostAndPort(), ignored -> makeNewClient(workerId)).left(); - } - } - - /** - * Close a single worker's clients. Used when that worker fails, so we stop trying to contact it. - * - * @param workerHost worker host:port - */ - public void closeClient(final String workerHost) - { - synchronized (clientMap) { - final Pair clientPair = clientMap.remove(workerHost); - if (clientPair != null) { - CloseableUtils.closeAndWrapExceptions(clientPair.right()); - } - } - } - - /** - * Close all outstanding clients. - */ - @Override - public void close() - { - synchronized (clientMap) { - for (Map.Entry> entry : clientMap.entrySet()) { - CloseableUtils.closeAndSuppressExceptions( - entry.getValue().right(), - e -> log.warn(e, "Failed to close client[%s]", entry.getKey()) - ); - } - - clientMap.clear(); - } - } - - /** - * Stops a worker. Dart-only API, used by the {@link DartWorkerManager}. - */ - public ListenableFuture stopWorker(String workerId) - { - return getClient(workerId).asyncRequest( - new RequestBuilder(HttpMethod.POST, "/stop"), - IgnoreHttpResponseHandler.INSTANCE - ); - } - - /** - * Create a new client. Called by {@link #getClient(String)} if a new one is needed. - */ - protected Pair makeNewClient(final WorkerId workerId) - { - final URI uri = workerId.toUri(); - final FixedServiceLocator locator = new FixedServiceLocator(ServiceLocation.fromUri(uri)); - final ServiceClient baseClient = - clientFactory.makeClient(workerId.toString(), locator, retryPolicy); - final ServiceClient client; - - if (controllerHost != null) { - client = new ControllerDecoratedClient(baseClient, controllerHost); - } else { - client = baseClient; - } - - return Pair.of(client, locator); - } - /** - * Service client that adds the {@link DartWorkerResource#HEADER_CONTROLLER_HOST} header. - */ - private static class ControllerDecoratedClient implements ServiceClient - { - private final ServiceClient delegate; - private final String controllerHost; + void closeClient(String hostAndPort); - ControllerDecoratedClient(final ServiceClient delegate, final String controllerHost) - { - this.delegate = delegate; - this.controllerHost = controllerHost; - } + ListenableFuture stopWorker(String workerId); - @Override - public ListenableFuture asyncRequest( - final RequestBuilder requestBuilder, - final HttpResponseHandler handler - ) - { - return delegate.asyncRequest( - requestBuilder.header(DartWorkerResource.HEADER_CONTROLLER_HOST, controllerHost), - handler - ); - } - @Override - public ServiceClient withRetryPolicy(final ServiceRetryPolicy retryPolicy) - { - return new ControllerDecoratedClient(delegate.withRetryPolicy(retryPolicy), controllerHost); - } - } } diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartWorkerClientImpl.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartWorkerClientImpl.java new file mode 100644 index 000000000000..4814eb6d5bb3 --- /dev/null +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartWorkerClientImpl.java @@ -0,0 +1,210 @@ +/* + * 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.dart.worker; + +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.jaxrs.smile.SmileMediaTypes; +import com.google.common.util.concurrent.ListenableFuture; +import com.google.errorprone.annotations.concurrent.GuardedBy; +import it.unimi.dsi.fastutil.Pair; +import org.apache.druid.error.DruidException; +import org.apache.druid.java.util.common.logger.Logger; +import org.apache.druid.java.util.http.client.response.HttpResponseHandler; +import org.apache.druid.msq.dart.controller.DartWorkerManager; +import org.apache.druid.msq.dart.controller.sql.DartSqlEngine; +import org.apache.druid.msq.dart.worker.http.DartWorkerResource; +import org.apache.druid.msq.exec.WorkerClient; +import org.apache.druid.msq.rpc.BaseWorkerClientImpl; +import org.apache.druid.rpc.FixedServiceLocator; +import org.apache.druid.rpc.IgnoreHttpResponseHandler; +import org.apache.druid.rpc.RequestBuilder; +import org.apache.druid.rpc.ServiceClient; +import org.apache.druid.rpc.ServiceClientFactory; +import org.apache.druid.rpc.ServiceLocation; +import org.apache.druid.rpc.ServiceRetryPolicy; +import org.apache.druid.utils.CloseableUtils; +import org.jboss.netty.handler.codec.http.HttpMethod; + +import javax.annotation.Nullable; +import java.io.Closeable; +import java.net.URI; +import java.util.HashMap; +import java.util.Map; + +/** + * Dart implementation of {@link WorkerClient}. Uses the same {@link BaseWorkerClientImpl} as the task-based engine. + * Each instance of this class is scoped to a single query. + */ +public class DartWorkerClientImpl extends BaseWorkerClientImpl implements DartWorkerClient +{ + private static final Logger log = new Logger(DartWorkerClientImpl.class); + + private final String queryId; + private final ServiceClientFactory clientFactory; + private final ServiceRetryPolicy retryPolicy; + + @Nullable + private final String controllerHost; + + @GuardedBy("clientMap") + private final Map> clientMap = new HashMap<>(); + + /** + * Create a worker client. + * + * @param queryId dart query ID. see {@link DartSqlEngine#CTX_DART_QUERY_ID} + * @param clientFactory service client factor + * @param smileMapper Smile object mapper + * @param controllerHost Controller host (see {@link DartWorkerResource#HEADER_CONTROLLER_HOST}) if this is a + * controller-to-worker client. Null if this is a worker-to-worker client. + */ + public DartWorkerClientImpl( + final String queryId, + final ServiceClientFactory clientFactory, + final ObjectMapper smileMapper, + @Nullable final String controllerHost + ) + { + super(smileMapper, SmileMediaTypes.APPLICATION_JACKSON_SMILE); + this.queryId = queryId; + this.clientFactory = clientFactory; + this.controllerHost = controllerHost; + + if (controllerHost == null) { + // worker -> worker client. Retry HTTP 503 in case worker A starts up before worker B, and needs to + // contact it immediately. + this.retryPolicy = new DartWorkerRetryPolicy(true); + } else { + // controller -> worker client. Do not retry any HTTP error codes. If we retry HTTP 503 for controller -> worker, + // we can get stuck trying to contact workers that have exited. + this.retryPolicy = new DartWorkerRetryPolicy(false); + } + } + + @Override + protected ServiceClient getClient(final String workerIdString) + { + final WorkerId workerId = WorkerId.fromString(workerIdString); + if (!queryId.equals(workerId.getQueryId())) { + throw DruidException.defensive("Unexpected queryId[%s]. Expected queryId[%s]", workerId.getQueryId(), queryId); + } + + synchronized (clientMap) { + return clientMap.computeIfAbsent(workerId.getHostAndPort(), ignored -> makeNewClient(workerId)).left(); + } + } + + /** + * Close a single worker's clients. Used when that worker fails, so we stop trying to contact it. + * + * @param workerHost worker host:port + */ + public void closeClient(final String workerHost) + { + synchronized (clientMap) { + final Pair clientPair = clientMap.remove(workerHost); + if (clientPair != null) { + CloseableUtils.closeAndWrapExceptions(clientPair.right()); + } + } + } + + /** + * Close all outstanding clients. + */ + @Override + public void close() + { + synchronized (clientMap) { + for (Map.Entry> entry : clientMap.entrySet()) { + CloseableUtils.closeAndSuppressExceptions( + entry.getValue().right(), + e -> log.warn(e, "Failed to close client[%s]", entry.getKey()) + ); + } + + clientMap.clear(); + } + } + + /** + * Stops a worker. Dart-only API, used by the {@link DartWorkerManager}. + */ + public ListenableFuture stopWorker(String workerId) + { + return getClient(workerId).asyncRequest( + new RequestBuilder(HttpMethod.POST, "/stop"), + IgnoreHttpResponseHandler.INSTANCE + ); + } + + /** + * Create a new client. Called by {@link #getClient(String)} if a new one is needed. + */ + protected Pair makeNewClient(final WorkerId workerId) + { + final URI uri = workerId.toUri(); + final FixedServiceLocator locator = new FixedServiceLocator(ServiceLocation.fromUri(uri)); + final ServiceClient baseClient = + clientFactory.makeClient(workerId.toString(), locator, retryPolicy); + final ServiceClient client; + + if (controllerHost != null) { + client = new ControllerDecoratedClient(baseClient, controllerHost); + } else { + client = baseClient; + } + + return Pair.of(client, locator); + } + + /** + * Service client that adds the {@link DartWorkerResource#HEADER_CONTROLLER_HOST} header. + */ + private static class ControllerDecoratedClient implements ServiceClient + { + private final ServiceClient delegate; + private final String controllerHost; + + ControllerDecoratedClient(final ServiceClient delegate, final String controllerHost) + { + this.delegate = delegate; + this.controllerHost = controllerHost; + } + + @Override + public ListenableFuture asyncRequest( + final RequestBuilder requestBuilder, + final HttpResponseHandler handler + ) + { + return delegate.asyncRequest( + requestBuilder.header(DartWorkerResource.HEADER_CONTROLLER_HOST, controllerHost), + handler + ); + } + + @Override + public ServiceClient withRetryPolicy(final ServiceRetryPolicy retryPolicy) + { + return new ControllerDecoratedClient(delegate.withRetryPolicy(retryPolicy), controllerHost); + } + } +} diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartWorkerFactoryImpl.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartWorkerFactoryImpl.java index 34da929060d8..1960924b4b67 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartWorkerFactoryImpl.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartWorkerFactoryImpl.java @@ -101,7 +101,7 @@ public Worker build(String queryId, String controllerHost, File tempDir, QueryCo selfNode, jsonMapper, injector, - new DartWorkerClient(queryId, serviceClientFactory, smileMapper, null), + new DartWorkerClientImpl(queryId, serviceClientFactory, smileMapper, null), processingConfig, segmentWrangler, groupingEngine, diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartWorkerRetryPolicy.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartWorkerRetryPolicy.java index 5dbfe98ef0c5..9e4d2f668039 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartWorkerRetryPolicy.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartWorkerRetryPolicy.java @@ -25,7 +25,7 @@ import org.jboss.netty.handler.codec.http.HttpResponseStatus; /** - * Retry policy for {@link DartWorkerClient}. This is a {@link StandardRetryPolicy#unlimited()} with + * Retry policy for {@link DartWorkerClientImpl}. This is a {@link StandardRetryPolicy#unlimited()} with * {@link #retryHttpResponse(HttpResponse)} customized to retry fewer HTTP error codes. */ public class DartWorkerRetryPolicy implements ServiceRetryPolicy diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/DartWorkerManagerTest.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/DartWorkerManagerTest.java index f4441c984e70..ea781e2e9920 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/DartWorkerManagerTest.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/DartWorkerManagerTest.java @@ -28,7 +28,7 @@ import org.apache.druid.error.DruidException; import org.apache.druid.indexer.TaskState; import org.apache.druid.java.util.common.ISE; -import org.apache.druid.msq.dart.worker.DartWorkerClient; +import org.apache.druid.msq.dart.worker.DartWorkerClientImpl; import org.apache.druid.msq.dart.worker.WorkerId; import org.apache.druid.msq.exec.WorkerManager; import org.apache.druid.msq.exec.WorkerStats; @@ -56,7 +56,7 @@ public class DartWorkerManagerTest private AutoCloseable mockCloser; @Mock - private DartWorkerClient workerClient; + private DartWorkerClientImpl workerClient; @BeforeEach public void setUp() diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/TestDartControllerContextFactoryImpl.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/TestDartControllerContextFactoryImpl.java index 44d0acdf1d89..92b5a6312b86 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/TestDartControllerContextFactoryImpl.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/TestDartControllerContextFactoryImpl.java @@ -23,33 +23,20 @@ import com.google.common.util.concurrent.ListenableFuture; import com.google.inject.Inject; import com.google.inject.Injector; -import it.unimi.dsi.fastutil.Pair; import org.apache.druid.client.TimelineServerView; -import org.apache.druid.frame.channel.ReadableByteChunksFrameChannel; -import org.apache.druid.frame.key.ClusterByPartitions; import org.apache.druid.guice.annotations.EscalatedGlobal; import org.apache.druid.guice.annotations.Json; import org.apache.druid.guice.annotations.Self; import org.apache.druid.guice.annotations.Smile; import org.apache.druid.java.util.emitter.service.ServiceEmitter; -import org.apache.druid.java.util.http.client.response.HttpResponseHandler; -import org.apache.druid.msq.counters.CounterSnapshotsTree; import org.apache.druid.msq.dart.controller.DartControllerContextFactoryImpl; import org.apache.druid.msq.dart.worker.DartWorkerClient; -import org.apache.druid.msq.dart.worker.WorkerId; import org.apache.druid.msq.exec.MemoryIntrospector; import org.apache.druid.msq.exec.Worker; -import org.apache.druid.msq.kernel.StageId; -import org.apache.druid.msq.kernel.WorkOrder; -import org.apache.druid.msq.rpc.SketchEncoding; -import org.apache.druid.msq.statistics.ClusterByStatisticsSnapshot; -import org.apache.druid.rpc.RequestBuilder; -import org.apache.druid.rpc.ServiceClient; import org.apache.druid.rpc.ServiceClientFactory; -import org.apache.druid.rpc.ServiceRetryPolicy; import org.apache.druid.server.DruidNode; -import java.io.Closeable; +import java.util.HashMap; import java.util.Map; public class TestDartControllerContextFactoryImpl extends DartControllerContextFactoryImpl @@ -71,14 +58,13 @@ public TestDartControllerContextFactoryImpl( @Override protected DartWorkerClient makeWorkerClient(String queryId) { - return - // super.makeWorkerClient(queryId); + return new DartTestWorkerClient(queryId, serviceClientFactory, smileMapper, selfNode.getHostAndPortToUse()); // new HashMap()); } - static class DartTestWorkerClient extends DartWorkerClient + static class DartTestWorkerClient extends MSQTestWorkerClient implements DartWorkerClient { private Map m; private MSQTestWorkerClient wc; @@ -86,98 +72,11 @@ static class DartTestWorkerClient extends DartWorkerClient public DartTestWorkerClient(String queryId, ServiceClientFactory clientFactory, ObjectMapper smileMapper, String controllerHost) { - super(queryId, clientFactory, smileMapper, controllerHost); - wc = new MSQTestWorkerClient(m); - } - - @Override - protected ServiceClient getClient(String workerIdString) - { - return super.getClient(workerIdString); - } - - @Override - public void closeClient(String workerHost) - { - super.closeClient(workerHost); - } - @Override - protected Object clone() throws CloneNotSupportedException - { - return super.clone(); - } - @Override - public ListenableFuture stopWorker(String workerId) - { - return super.stopWorker(workerId); - } - @Override - public ListenableFuture fetchChannelData(String workerId, StageId stageId, int partitionNumber, - long offset, ReadableByteChunksFrameChannel channel) - { - return super.fetchChannelData(workerId, stageId, partitionNumber, offset, channel); - } - - @Override - public ListenableFuture fetchClusterByStatisticsSnapshot(String workerId, - StageId stageId, SketchEncoding sketchEncoding) - { - return super.fetchClusterByStatisticsSnapshot(workerId, stageId, sketchEncoding); - - } - @Override - public ListenableFuture fetchClusterByStatisticsSnapshotForTimeChunk(String workerId, - StageId stageId, long timeChunk, SketchEncoding sketchEncoding) - { - return super.fetchClusterByStatisticsSnapshotForTimeChunk(workerId, stageId, timeChunk, sketchEncoding); - - } - @Override - public ListenableFuture getCounters(String workerId) - { - return super.getCounters(workerId); - - } - @Override - public ListenableFuture postCleanupStage(String workerId, StageId stageId) - { - return super.postCleanupStage(workerId, stageId); - - } - @Override - public ListenableFuture postFinish(String workerId) - { - return super.postFinish(workerId); - - } - - @Override - public ListenableFuture postResultPartitionBoundaries(String workerId, StageId stageId, - ClusterByPartitions partitionBoundaries) - { - return super.postResultPartitionBoundaries(workerId, stageId, partitionBoundaries); - - } - @Override - public ListenableFuture postWorkOrder(String workerId, WorkOrder workOrder) - { - return super.postWorkOrder(workerId, workOrder); - + super(new HashMap<>()); } @Override - protected Pair makeNewClient(WorkerId workerId) - { - ServiceClient client =new TestDartServiceClient(workerId); - return super.makeNewClient(workerId); - } - - - } - - static class TestDartServiceClient implements ServiceClient { - - public TestDartServiceClient(WorkerId workerId) + public void closeClient(String hostAndPort) { if(true) { @@ -187,19 +86,7 @@ public TestDartServiceClient(WorkerId workerId) } @Override - public ListenableFuture asyncRequest(RequestBuilder requestBuilder, - HttpResponseHandler handler) - { - if(true) - { - throw new RuntimeException("FIXME: Unimplemented!"); - } - return null; - - } - - @Override - public ServiceClient withRetryPolicy(ServiceRetryPolicy retryPolicy) + public ListenableFuture stopWorker(String workerId) { if(true) { @@ -208,6 +95,138 @@ public ServiceClient withRetryPolicy(ServiceRetryPolicy retryPolicy) return null; } - } + + // @Override + // protected ServiceClient getClient(String workerIdString) + // { + // return super.getClient(workerIdString); + // } + // + // @Override + // public void closeClient(String workerHost) + // { + // super.closeClient(workerHost); + // } + // @Override + // protected Object clone() throws CloneNotSupportedException + // { + // return super.clone(); + // } + // @Override + // public ListenableFuture stopWorker(String workerId) + // { + // return super.stopWorker(workerId); + // } + // @Override + // public ListenableFuture fetchChannelData(String workerId, StageId + // stageId, int partitionNumber, + // long offset, ReadableByteChunksFrameChannel channel) + // { + // return super.fetchChannelData(workerId, stageId, partitionNumber, offset, + // channel); + // } + // + // @Override + // public ListenableFuture + // fetchClusterByStatisticsSnapshot(String workerId, + // StageId stageId, SketchEncoding sketchEncoding) + // { + // return super.fetchClusterByStatisticsSnapshot(workerId, stageId, + // sketchEncoding); + // + // } + // @Override + // public ListenableFuture + // fetchClusterByStatisticsSnapshotForTimeChunk(String workerId, + // StageId stageId, long timeChunk, SketchEncoding sketchEncoding) + // { + // return super.fetchClusterByStatisticsSnapshotForTimeChunk(workerId, + // stageId, timeChunk, sketchEncoding); + // + // } + // @Override + // public ListenableFuture getCounters(String workerId) + // { + // return super.getCounters(workerId); + // + // } + // @Override + // public ListenableFuture postCleanupStage(String workerId, StageId + // stageId) + // { + // return super.postCleanupStage(workerId, stageId); + // + // } + // @Override + // public ListenableFuture postFinish(String workerId) + // { + // return super.postFinish(workerId); + // + // } + // + // @Override + // public ListenableFuture postResultPartitionBoundaries(String + // workerId, StageId stageId, + // ClusterByPartitions partitionBoundaries) + // { + // return super.postResultPartitionBoundaries(workerId, stageId, + // partitionBoundaries); + // + // } + // @Override + // public ListenableFuture postWorkOrder(String workerId, WorkOrder + // workOrder) + // { + // return super.postWorkOrder(workerId, workOrder); + // + // } + // + // @Override + // protected Pair makeNewClient(WorkerId workerId) + // { + // ServiceClient client =new TestDartServiceClient(workerId); + // Closeable c=null; + // Pair p = Pair.of(client, c); + // return p; + // } + // + // + // } + // + // static class TestDartServiceClient implements ServiceClient { + // + // + // private MSQTestWorkerClient a; + // + // public TestDartServiceClient(WorkerId workerId) + // { + // a = new MSQTestWorkerClient(null); + // } + // + // @Override + // public ListenableFuture + // asyncRequest(RequestBuilder requestBuilder, + // HttpResponseHandler handler) + // { + // if(true) + // { + // throw new RuntimeException("FIXME: Unimplemented!"); + // } + // return null; + // + // } + // + // @Override + // public ServiceClient withRetryPolicy(ServiceRetryPolicy retryPolicy) + // { + // if(true) + // { + // throw new RuntimeException("FIXME: Unimplemented!"); + // } + // return null; + // + // } + // + // } } From 5de20dc9d2a18abe1f7e5b40cd74b4227a4fd925 Mon Sep 17 00:00:00 2001 From: Zoltan Haindrich Date: Tue, 4 Mar 2025 13:04:31 +0000 Subject: [PATCH 09/50] s --- .../druid/msq/test/DartComponentSupplier.java | 12 +++++ .../TestDartControllerContextFactoryImpl.java | 44 ++++++++++++++----- 2 files changed, 46 insertions(+), 10 deletions(-) diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/DartComponentSupplier.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/DartComponentSupplier.java index 1c7965b613a9..e204c2cd87c3 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/DartComponentSupplier.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/DartComponentSupplier.java @@ -32,11 +32,13 @@ import org.apache.druid.guice.annotations.Merging; import org.apache.druid.initialization.DruidModule; import org.apache.druid.java.util.http.client.HttpClient; +import org.apache.druid.msq.dart.Dart; import org.apache.druid.msq.dart.controller.DartControllerContextFactory; import org.apache.druid.msq.dart.controller.sql.DartSqlEngine; import org.apache.druid.msq.dart.guice.DartControllerModule; import org.apache.druid.msq.dart.guice.DartWorkerMemoryManagementModule; import org.apache.druid.msq.dart.guice.DartWorkerModule; +import org.apache.druid.msq.exec.Worker; import org.apache.druid.msq.exec.WorkerMemoryParameters; import org.apache.druid.query.TestBufferPool; import org.apache.druid.rpc.ServiceClientFactory; @@ -48,6 +50,8 @@ import org.apache.druid.sql.calcite.util.SqlTestFramework.StandardComponentSupplier; import java.nio.ByteBuffer; +import java.util.HashMap; +import java.util.Map; public class DartComponentSupplier extends AbstractMSQComponentSupplierDelegate { @@ -132,6 +136,14 @@ public void configure(Binder binder) } } + @Provides + @LazySingleton + @Dart + Map workerMap() + { + return new HashMap(); + } + // @Provides // @LazySingleton // public MSQTaskSqlEngine createEngine( diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/TestDartControllerContextFactoryImpl.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/TestDartControllerContextFactoryImpl.java index 92b5a6312b86..0c503f633b86 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/TestDartControllerContextFactoryImpl.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/TestDartControllerContextFactoryImpl.java @@ -29,18 +29,23 @@ import org.apache.druid.guice.annotations.Self; import org.apache.druid.guice.annotations.Smile; import org.apache.druid.java.util.emitter.service.ServiceEmitter; +import org.apache.druid.msq.dart.Dart; import org.apache.druid.msq.dart.controller.DartControllerContextFactoryImpl; import org.apache.druid.msq.dart.worker.DartWorkerClient; +import org.apache.druid.msq.exec.Controller; import org.apache.druid.msq.exec.MemoryIntrospector; import org.apache.druid.msq.exec.Worker; +import org.apache.druid.msq.exec.WorkerImpl; +import org.apache.druid.msq.exec.WorkerStorageParameters; import org.apache.druid.rpc.ServiceClientFactory; import org.apache.druid.server.DruidNode; -import java.util.HashMap; import java.util.Map; public class TestDartControllerContextFactoryImpl extends DartControllerContextFactoryImpl { + private Map workerMap; + @Inject public TestDartControllerContextFactoryImpl( final Injector injector, @@ -50,29 +55,48 @@ public TestDartControllerContextFactoryImpl( @EscalatedGlobal final ServiceClientFactory serviceClientFactory, final MemoryIntrospector memoryIntrospector, final TimelineServerView serverView, - final ServiceEmitter emitter) + final ServiceEmitter emitter, + @Dart Map workerMap + ) { super(injector, jsonMapper, smileMapper, selfNode, serviceClientFactory, memoryIntrospector, serverView, emitter); + this.workerMap = workerMap; } @Override protected DartWorkerClient makeWorkerClient(String queryId) { + Controller controller = null; + Worker worker = new WorkerImpl( + null, + new MSQTestWorkerContext( + queryId, + workerMap, + controller, + jsonMapper, + injector, + MSQTestBase.makeTestWorkerMemoryParameters(), + WorkerStorageParameters.createInstanceForTests(Long.MAX_VALUE) + ) + ); - return - new DartTestWorkerClient(queryId, serviceClientFactory, smileMapper, selfNode.getHostAndPortToUse()); - // new HashMap()); + return new DartTestWorkerClient( + queryId, serviceClientFactory, smileMapper, selfNode.getHostAndPortToUse(), workerMap + ); } static class DartTestWorkerClient extends MSQTestWorkerClient implements DartWorkerClient { - private Map m; - private MSQTestWorkerClient wc; + private Map workerMap; - public DartTestWorkerClient(String queryId, ServiceClientFactory clientFactory, ObjectMapper smileMapper, - String controllerHost) + public DartTestWorkerClient( + String queryId, + ServiceClientFactory clientFactory, + ObjectMapper smileMapper, + String controllerHost, Map workerMap + ) { - super(new HashMap<>()); + super(workerMap); } @Override From c0ed1b4687f1faf3646a549cc559f687f6239f1e Mon Sep 17 00:00:00 2001 From: Zoltan Haindrich Date: Tue, 4 Mar 2025 15:10:36 +0000 Subject: [PATCH 10/50] use TimelineServerView instead BrokerServerView --- .../druid/msq/dart/controller/DartControllerContext.java | 6 +++--- .../dart/controller/DartControllerContextFactoryImpl.java | 6 +++--- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartControllerContext.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartControllerContext.java index b7cf087a3f4d..c2257d160ed3 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartControllerContext.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartControllerContext.java @@ -21,7 +21,7 @@ import com.fasterxml.jackson.databind.ObjectMapper; import com.google.inject.Injector; -import org.apache.druid.client.BrokerServerView; +import org.apache.druid.client.TimelineServerView; import org.apache.druid.error.DruidException; import org.apache.druid.indexing.common.TaskLockType; import org.apache.druid.indexing.common.actions.TaskActionClient; @@ -84,7 +84,7 @@ public class DartControllerContext implements ControllerContext private final ObjectMapper jsonMapper; private final DruidNode selfNode; private final DartWorkerClient workerClient; - private final BrokerServerView serverView; + private final TimelineServerView serverView; private final MemoryIntrospector memoryIntrospector; private final ServiceMetricEvent.Builder metricBuilder; private final ServiceEmitter emitter; @@ -95,7 +95,7 @@ public DartControllerContext( final DruidNode selfNode, final DartWorkerClient workerClient, final MemoryIntrospector memoryIntrospector, - final BrokerServerView serverView, + final TimelineServerView serverView, final ServiceEmitter emitter ) { diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartControllerContextFactoryImpl.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartControllerContextFactoryImpl.java index 8cefb6af7ece..fe26554d3712 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartControllerContextFactoryImpl.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartControllerContextFactoryImpl.java @@ -22,7 +22,7 @@ import com.fasterxml.jackson.databind.ObjectMapper; import com.google.inject.Inject; import com.google.inject.Injector; -import org.apache.druid.client.BrokerServerView; +import org.apache.druid.client.TimelineServerView; import org.apache.druid.guice.annotations.EscalatedGlobal; import org.apache.druid.guice.annotations.Json; import org.apache.druid.guice.annotations.Self; @@ -41,7 +41,7 @@ public class DartControllerContextFactoryImpl implements DartControllerContextFa private final ObjectMapper smileMapper; private final DruidNode selfNode; private final ServiceClientFactory serviceClientFactory; - private final BrokerServerView serverView; + private final TimelineServerView serverView; private final MemoryIntrospector memoryIntrospector; private final ServiceEmitter emitter; @@ -53,7 +53,7 @@ public DartControllerContextFactoryImpl( @Self final DruidNode selfNode, @EscalatedGlobal final ServiceClientFactory serviceClientFactory, final MemoryIntrospector memoryIntrospector, - final BrokerServerView serverView, + final TimelineServerView serverView, final ServiceEmitter emitter ) { From fe641fddf659711ce57d4cf55d0da55d1507c835 Mon Sep 17 00:00:00 2001 From: Zoltan Haindrich Date: Wed, 5 Mar 2025 07:46:24 +0000 Subject: [PATCH 11/50] make TestTimelineServerView --- .../druid/sql/calcite/util/CalciteTests.java | 5 ++- .../sql/calcite/util/QueryFrameworkUtils.java | 38 ++++++++++++++++--- .../sql/calcite/util/SqlTestFramework.java | 26 +++++++++---- .../calcite/util/TestTimelineServerView.java | 2 +- 4 files changed, 56 insertions(+), 15 deletions(-) diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/util/CalciteTests.java b/sql/src/test/java/org/apache/druid/sql/calcite/util/CalciteTests.java index abf32c9bf321..cb3ceaf442d7 100644 --- a/sql/src/test/java/org/apache/druid/sql/calcite/util/CalciteTests.java +++ b/sql/src/test/java/org/apache/druid/sql/calcite/util/CalciteTests.java @@ -32,6 +32,7 @@ import org.apache.druid.client.FilteredServerInventoryView; import org.apache.druid.client.ServerInventoryView; import org.apache.druid.client.ServerView; +import org.apache.druid.client.TimelineServerView; import org.apache.druid.discovery.DiscoveryDruidNode; import org.apache.druid.discovery.DruidLeaderClient; import org.apache.druid.discovery.DruidNodeDiscovery; @@ -397,7 +398,7 @@ NodeRole.COORDINATOR, new FakeDruidNodeDiscovery(ImmutableMap.of(NodeRole.COORDI public static SystemSchema createMockSystemSchema( final DruidSchema druidSchema, - final SpecificSegmentsQuerySegmentWalker walker, + final TimelineServerView timelineServerView, final AuthorizerMapper authorizerMapper ) { @@ -474,7 +475,7 @@ private TaskStatusPlus createTaskStatus(String id, String datasource, Long durat new BrokerSegmentWatcherConfig(), BrokerSegmentMetadataCacheConfig.create() ), - new TestTimelineServerView(walker.getSegments()), + timelineServerView, new FakeServerInventoryView(), authorizerMapper, druidLeaderClient, diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/util/QueryFrameworkUtils.java b/sql/src/test/java/org/apache/druid/sql/calcite/util/QueryFrameworkUtils.java index f28461e1711c..f61f52343987 100644 --- a/sql/src/test/java/org/apache/druid/sql/calcite/util/QueryFrameworkUtils.java +++ b/sql/src/test/java/org/apache/druid/sql/calcite/util/QueryFrameworkUtils.java @@ -26,6 +26,7 @@ import org.apache.calcite.jdbc.CalciteSchema; import org.apache.calcite.schema.SchemaPlus; import org.apache.druid.client.InternalQueryConfig; +import org.apache.druid.client.TimelineServerView; import org.apache.druid.java.util.emitter.core.NoopEmitter; import org.apache.druid.java.util.emitter.service.ServiceEmitter; import org.apache.druid.query.DefaultGenericQueryMetricsFactory; @@ -125,7 +126,32 @@ public static DruidSchemaCatalog createMockRootSchema( @Nullable final ViewManager viewManager, final DruidSchemaManager druidSchemaManager, final AuthorizerMapper authorizerMapper, - final CatalogResolver catalogResolver + final CatalogResolver catalogResolver) + { + TimelineServerView timelineServerView = new TestTimelineServerView(walker.getSegments()); + return createMockRootSchema( + injector, + conglomerate, + walker, + plannerConfig, + viewManager, + druidSchemaManager, + authorizerMapper, + catalogResolver, + timelineServerView + ); + } + + public static DruidSchemaCatalog createMockRootSchema( + final Injector injector, + final QueryRunnerFactoryConglomerate conglomerate, + final SpecificSegmentsQuerySegmentWalker walker, + final PlannerConfig plannerConfig, + @Nullable final ViewManager viewManager, + final DruidSchemaManager druidSchemaManager, + final AuthorizerMapper authorizerMapper, + final CatalogResolver catalogResolver, + final TimelineServerView timelineServerView ) { DruidSchema druidSchema = createMockSchema( @@ -133,10 +159,11 @@ public static DruidSchemaCatalog createMockRootSchema( conglomerate, walker, druidSchemaManager, - catalogResolver + catalogResolver, + timelineServerView ); SystemSchema systemSchema = - CalciteTests.createMockSystemSchema(druidSchema, walker, authorizerMapper); + CalciteTests.createMockSystemSchema(druidSchema, timelineServerView, authorizerMapper); LookupSchema lookupSchema = createMockLookupSchema(injector); DruidOperatorTable createOperatorTable = createOperatorTable(injector); @@ -221,12 +248,13 @@ static DruidSchema createMockSchema( final QueryRunnerFactoryConglomerate conglomerate, final SpecificSegmentsQuerySegmentWalker walker, final DruidSchemaManager druidSchemaManager, - final CatalogResolver catalog + final CatalogResolver catalog, + final TimelineServerView timelineServerView ) { final BrokerSegmentMetadataCache cache = new BrokerSegmentMetadataCache( createMockQueryLifecycleFactory(walker, conglomerate), - new TestTimelineServerView(walker.getSegments()), + timelineServerView, BrokerSegmentMetadataCacheConfig.create(), CalciteTests.TEST_AUTHENTICATOR_ESCALATOR, new InternalQueryConfig(), diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/util/SqlTestFramework.java b/sql/src/test/java/org/apache/druid/sql/calcite/util/SqlTestFramework.java index 90cc53d1b9a0..656c5fe003f1 100644 --- a/sql/src/test/java/org/apache/druid/sql/calcite/util/SqlTestFramework.java +++ b/sql/src/test/java/org/apache/druid/sql/calcite/util/SqlTestFramework.java @@ -29,6 +29,7 @@ import com.google.inject.Module; import com.google.inject.Provides; import com.google.inject.TypeLiteral; +import org.apache.druid.client.TimelineServerView; import org.apache.druid.client.cache.Cache; import org.apache.druid.client.cache.CacheConfig; import org.apache.druid.collections.NonBlockingPool; @@ -672,7 +673,8 @@ public PlannerFixture( viewManager, componentSupplier.createSchemaManager(), framework.authorizerMapper, - framework.builder.catalogResolver + framework.builder.catalogResolver, + framework.injector.getInstance(TimelineServerView.class) ); this.plannerFactory = new PlannerFactory( @@ -923,26 +925,36 @@ private DruidSchema makeDruidSchema( final Injector injector, QueryRunnerFactoryConglomerate conglomerate, QuerySegmentWalker walker, - Builder builder) + Builder builder, + TimelineServerView timelineServerView) { return QueryFrameworkUtils.createMockSchema( injector, conglomerate, (SpecificSegmentsQuerySegmentWalker) walker, builder.componentSupplier.getPlannerComponentSupplier().createSchemaManager(), - builder.catalogResolver + builder.catalogResolver, + timelineServerView ); } @Provides @LazySingleton - private SystemSchema makeSystemSchema(QuerySegmentWalker walker, AuthorizerMapper authorizerMapper, - DruidSchema druidSchema) + private SystemSchema makeSystemSchema(AuthorizerMapper authorizerMapper, DruidSchema druidSchema, + TimelineServerView timelineServerView) + { + return CalciteTests.createMockSystemSchema(druidSchema, timelineServerView, authorizerMapper); + } + + + @Provides + @LazySingleton + private TimelineServerView makeTimelineServerView(SpecificSegmentsQuerySegmentWalker walker) { - return CalciteTests - .createMockSystemSchema(druidSchema, (SpecificSegmentsQuerySegmentWalker) walker, authorizerMapper); + return new TestTimelineServerView(walker.getSegments()); } + @Provides @LazySingleton private ColumnConfig getColumnConfig() diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/util/TestTimelineServerView.java b/sql/src/test/java/org/apache/druid/sql/calcite/util/TestTimelineServerView.java index 42aa4e491cb4..e5de23b0abbe 100644 --- a/sql/src/test/java/org/apache/druid/sql/calcite/util/TestTimelineServerView.java +++ b/sql/src/test/java/org/apache/druid/sql/calcite/util/TestTimelineServerView.java @@ -81,7 +81,7 @@ public class TestTimelineServerView implements TimelineServerView public TestTimelineServerView(List segments) { - this.segments.addAll(segments); + this(segments, Collections.emptyList()); } public TestTimelineServerView(List segments, List realtimeSegments) From d60fab3bdc3943e85551b0831eddaece8d69454d Mon Sep 17 00:00:00 2001 From: Zoltan Haindrich Date: Wed, 5 Mar 2025 11:00:36 +0000 Subject: [PATCH 12/50] stuff --- .../controller/DartControllerContext.java | 7 + .../msq/dart/worker/DartWorkerClient.java | 1 - .../msq/test/CalciteSelectQueryMSQTest.java | 1 + .../druid/msq/test/MSQTestWorkerClient.java | 30 +++-- .../TestDartControllerContextFactoryImpl.java | 120 ++++++++++++++---- .../dart.iq | 41 +----- 6 files changed, 128 insertions(+), 72 deletions(-) diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartControllerContext.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartControllerContext.java index 258ff6d7fa70..1d7056a99c7d 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartControllerContext.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartControllerContext.java @@ -89,6 +89,8 @@ public class DartControllerContext implements ControllerContext private final ServiceMetricEvent.Builder metricBuilder; private final ServiceEmitter emitter; + public Controller theController; + public DartControllerContext( final Injector injector, final ObjectMapper jsonMapper, @@ -218,6 +220,11 @@ public DartWorkerClient newWorkerClient() public void registerController(Controller controller, Closer closer) { closer.register(workerClient); + // FIXME: + //workerClient.setController(); +// controller + theController=controller; + } @Override diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartWorkerClient.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartWorkerClient.java index 5c2f623755b2..1ad7db873470 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartWorkerClient.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartWorkerClient.java @@ -24,7 +24,6 @@ public interface DartWorkerClient extends WorkerClient { - void closeClient(String hostAndPort); ListenableFuture stopWorker(String workerId); diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/CalciteSelectQueryMSQTest.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/CalciteSelectQueryMSQTest.java index f19090ad6c0b..dda6b8f38067 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/CalciteSelectQueryMSQTest.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/CalciteSelectQueryMSQTest.java @@ -43,6 +43,7 @@ public class CalciteSelectQueryMSQTest extends CalciteQueryTest protected QueryTestBuilder testBuilder() { return new QueryTestBuilder(new CalciteTestConfig(true)) +// .addCustomRunner(new ExtractResultsFactory(() -> null))//(MSQTestOverlordServiceClient) ((DartSqlEngine) queryFramework().engine()).overlordClient())) .addCustomRunner(new ExtractResultsFactory(() -> null))//(MSQTestOverlordServiceClient) ((DartSqlEngine) queryFramework().engine()).overlordClient())) .queryContext(ImmutableMap.builder().put("asd",UUID.randomUUID().toString()).build()) .skipVectorize(true) diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestWorkerClient.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestWorkerClient.java index 4c7ccd72efd0..d535084e6802 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestWorkerClient.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestWorkerClient.java @@ -39,7 +39,7 @@ public class MSQTestWorkerClient implements WorkerClient { - private final Map inMemoryWorkers; + protected final Map inMemoryWorkers; private final AtomicBoolean closed = new AtomicBoolean(); public MSQTestWorkerClient(Map inMemoryWorkers) @@ -50,10 +50,22 @@ public MSQTestWorkerClient(Map inMemoryWorkers) @Override public ListenableFuture postWorkOrder(String workerTaskId, WorkOrder workOrder) { - inMemoryWorkers.get(workerTaskId).postWorkOrder(workOrder); + getWorkerFor(workerTaskId).postWorkOrder(workOrder); + return Futures.immediateFuture(null); } + protected Worker getWorkerFor(String workerTaskId) + { + + return inMemoryWorkers.computeIfAbsent(workerTaskId, this::newWorker); + } + + protected Worker newWorker(String workerId) + { + throw new RuntimeException("FIXME: Unimplemented!"); + } + @Override public ListenableFuture fetchClusterByStatisticsSnapshot( String workerTaskId, @@ -61,7 +73,7 @@ public ListenableFuture fetchClusterByStatisticsSna SketchEncoding sketchEncoding ) { - return Futures.immediateFuture(inMemoryWorkers.get(workerTaskId).fetchStatisticsSnapshot(stageId)); + return Futures.immediateFuture(getWorkerFor(workerTaskId).fetchStatisticsSnapshot(stageId)); } @Override @@ -73,7 +85,7 @@ public ListenableFuture fetchClusterByStatisticsSna ) { return Futures.immediateFuture( - inMemoryWorkers.get(workerTaskId).fetchStatisticsSnapshotForTimeChunk(stageId, timeChunk) + getWorkerFor(workerTaskId).fetchStatisticsSnapshotForTimeChunk(stageId, timeChunk) ); } @@ -85,7 +97,7 @@ public ListenableFuture postResultPartitionBoundaries( ) { try { - inMemoryWorkers.get(workerTaskId).postResultPartitionBoundaries(stageId, partitionBoundaries); + getWorkerFor(workerTaskId).postResultPartitionBoundaries(stageId, partitionBoundaries); return Futures.immediateFuture(null); } catch (Exception e) { @@ -96,21 +108,21 @@ public ListenableFuture postResultPartitionBoundaries( @Override public ListenableFuture postCleanupStage(String workerTaskId, StageId stageId) { - inMemoryWorkers.get(workerTaskId).postCleanupStage(stageId); + getWorkerFor(workerTaskId).postCleanupStage(stageId); return Futures.immediateFuture(null); } @Override public ListenableFuture postFinish(String taskId) { - inMemoryWorkers.get(taskId).postFinish(); + getWorkerFor(taskId).postFinish(); return Futures.immediateFuture(null); } @Override public ListenableFuture getCounters(String taskId) { - return Futures.immediateFuture(inMemoryWorkers.get(taskId).getCounters()); + return Futures.immediateFuture(getWorkerFor(taskId).getCounters()); } @Override @@ -123,7 +135,7 @@ public ListenableFuture fetchChannelData( ) { try (InputStream inputStream = - inMemoryWorkers.get(workerTaskId).readStageOutput(stageId, partitionNumber, offset).get()) { + getWorkerFor(workerTaskId).readStageOutput(stageId, partitionNumber, offset).get()) { byte[] buffer = new byte[8 * 1024]; boolean didRead = false; int bytesRead; diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/TestDartControllerContextFactoryImpl.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/TestDartControllerContextFactoryImpl.java index 0c503f633b86..5bcb162dcab0 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/TestDartControllerContextFactoryImpl.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/TestDartControllerContextFactoryImpl.java @@ -30,17 +30,24 @@ import org.apache.druid.guice.annotations.Smile; import org.apache.druid.java.util.emitter.service.ServiceEmitter; import org.apache.druid.msq.dart.Dart; +import org.apache.druid.msq.dart.controller.DartControllerContext; import org.apache.druid.msq.dart.controller.DartControllerContextFactoryImpl; import org.apache.druid.msq.dart.worker.DartWorkerClient; import org.apache.druid.msq.exec.Controller; +import org.apache.druid.msq.exec.ControllerContext; import org.apache.druid.msq.exec.MemoryIntrospector; import org.apache.druid.msq.exec.Worker; import org.apache.druid.msq.exec.WorkerImpl; import org.apache.druid.msq.exec.WorkerStorageParameters; +import org.apache.druid.msq.kernel.StageId; +import org.apache.druid.msq.kernel.WorkOrder; import org.apache.druid.rpc.ServiceClientFactory; import org.apache.druid.server.DruidNode; import java.util.Map; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; public class TestDartControllerContextFactoryImpl extends DartControllerContextFactoryImpl { @@ -64,30 +71,50 @@ public TestDartControllerContextFactoryImpl( } @Override - protected DartWorkerClient makeWorkerClient(String queryId) + public ControllerContext newContext(String queryId) { - Controller controller = null; - Worker worker = new WorkerImpl( - null, - new MSQTestWorkerContext( - queryId, - workerMap, - controller, - jsonMapper, - injector, - MSQTestBase.makeTestWorkerMemoryParameters(), - WorkerStorageParameters.createInstanceForTests(Long.MAX_VALUE) - ) + DartControllerContext ctx; + DartTestWorkerClient wc; + ctx=new DartControllerContext( + injector, + jsonMapper, + selfNode, + wc=(DartTestWorkerClient) makeWorkerClient(queryId), + memoryIntrospector, + serverView, + emitter ); + wc.controllerCtx=ctx; + return ctx; + } + + @Override + protected DartWorkerClient makeWorkerClient(String queryId) + { +// Controller controller = null; +// Worker worker = new WorkerImpl( +// null, +// new MSQTestWorkerContext( +// queryId, +// workerMap, +// controller, +// jsonMapper, +// injector, +// MSQTestBase.makeTestWorkerMemoryParameters(), +// WorkerStorageParameters.createInstanceForTests(Long.MAX_VALUE) +// ) +// ); +// injector. return new DartTestWorkerClient( queryId, serviceClientFactory, smileMapper, selfNode.getHostAndPortToUse(), workerMap ); } - static class DartTestWorkerClient extends MSQTestWorkerClient implements DartWorkerClient + public class DartTestWorkerClient extends MSQTestWorkerClient implements DartWorkerClient { - private Map workerMap; + private final ExecutorService EXECUTOR = Executors.newCachedThreadPool(); + public DartControllerContext controllerCtx; public DartTestWorkerClient( String queryId, @@ -99,23 +126,70 @@ public DartTestWorkerClient( super(workerMap); } + @Override + protected Worker newWorker(String workerId) + { + String queryId=workerId; + Controller controller=controllerCtx.theController; + Worker worker = new WorkerImpl( + null, + new MSQTestWorkerContext( + queryId, + inMemoryWorkers, + controller, + jsonMapper, + injector, + MSQTestBase.makeTestWorkerMemoryParameters(), + WorkerStorageParameters.createInstanceForTests(Long.MAX_VALUE) + ) + ); + + Future future = EXECUTOR.submit(() -> { + try { + worker.run(); + } + catch (Exception e) { + throw new RuntimeException(e); + } + }); + + return worker; + } + + @Override + public ListenableFuture postWorkOrder(String workerTaskId, WorkOrder workOrder) + { + + return super.postWorkOrder(workerTaskId, workOrder); + + } + @Override + public ListenableFuture postCleanupStage(String workerTaskId, StageId stageId) + { +// if(true) +// { +// throw new RuntimeException("FIXME: Unimplemented!"); +// } + return super.postCleanupStage(workerTaskId, stageId); + + } @Override public void closeClient(String hostAndPort) { - if(true) - { - throw new RuntimeException("FIXME: Unimplemented!"); - } +// if(true) +// { +// throw new RuntimeException("FIXME: Unimplemented!"); +// } } @Override public ListenableFuture stopWorker(String workerId) { - if(true) - { - throw new RuntimeException("FIXME: Unimplemented!"); - } +// if(true) +// { +// throw new RuntimeException("FIXME: Unimplemented!"); +// } return null; } diff --git a/extensions-core/multi-stage-query/src/test/quidem/org.apache.druid.msq.quidem.MSQQuidemTest/dart.iq b/extensions-core/multi-stage-query/src/test/quidem/org.apache.druid.msq.quidem.MSQQuidemTest/dart.iq index a5ef2481f085..6a3bdc1754e8 100644 --- a/extensions-core/multi-stage-query/src/test/quidem/org.apache.druid.msq.quidem.MSQQuidemTest/dart.iq +++ b/extensions-core/multi-stage-query/src/test/quidem/org.apache.druid.msq.quidem.MSQQuidemTest/dart.iq @@ -2,14 +2,8 @@ !use druidtest://?componentSupplier=DartComponentSupplier !set outputformat mysql -# This test validates that all output rows are flushed to the output channel even if frame writer's capacity is reached. -select count(*) as actualNumRows -from ( - select countryName, cityName, channel, added, delta, row_number() over() as rowNumber - from wikipedia - group by countryName, cityName, channel, added, delta -); +select * from foo; +---------------+ | actualNumRows | +---------------+ @@ -19,36 +13,5 @@ from ( !ok -# Validate that all rows are outputted by window WindowOperatorQueryFrameProcessor layer for empty over() clause scenario. - -select count(*) as numRows, max(rowNumber) as maxRowNumber -from ( - select countryName, cityName, channel, added, delta, row_number() over() as rowNumber - from wikipedia - group by countryName, cityName, channel, added, delta -); -+---------+--------------+ -| numRows | maxRowNumber | -+---------+--------------+ -| 11631 | 11631 | -+---------+--------------+ -(1 row) - -!ok - -# Validate that all rows are outputted by window WindowOperatorQueryFrameProcessor layer for non-empty over() clause scenario. - -select rowNumber, count(rowNumber) as numRows -from ( - select countryName, cityName, channel, added, delta, row_number() over(partition by countryName, cityName, channel, added, delta) as rowNumber - from wikipedia - group by countryName, cityName, channel, added, delta -) group by rowNumber; -+-----------+---------+ -| rowNumber | numRows | -+-----------+---------+ -| 1 | 11631 | -+-----------+---------+ -(1 row) - +select * from foo group by dim1; !ok From 9b85514db677a8aead2d10b8a25170ba18fb2316 Mon Sep 17 00:00:00 2001 From: Zoltan Haindrich Date: Wed, 5 Mar 2025 13:53:48 +0000 Subject: [PATCH 13/50] aa --- .../msq/test/CalciteSelectQueryDartTest.java | 186 ++++++++++++++++++ .../msq/test/CalciteSelectQueryMSQTest.java | 2 +- .../druid/sql/calcite/CalciteQueryTest.java | 53 ++++- 3 files changed, 234 insertions(+), 7 deletions(-) create mode 100644 extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/CalciteSelectQueryDartTest.java diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/CalciteSelectQueryDartTest.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/CalciteSelectQueryDartTest.java new file mode 100644 index 000000000000..54eef7678569 --- /dev/null +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/CalciteSelectQueryDartTest.java @@ -0,0 +1,186 @@ +/* + * 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.test; + +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import org.apache.druid.java.util.common.ISE; +import org.apache.druid.sql.calcite.CalciteQueryTest; +import org.apache.druid.sql.calcite.QueryTestBuilder; +import org.apache.druid.sql.calcite.SqlTestFrameworkConfig; +import org.junit.Assert; +import org.junit.jupiter.api.Disabled; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; + +import java.util.UUID; +import java.util.concurrent.TimeUnit; + +/** + * Runs {@link CalciteQueryTest} but with MSQ engine + */ +@SqlTestFrameworkConfig.ComponentSupplier(DartComponentSupplier.class) +public class CalciteSelectQueryDartTest extends CalciteQueryTest +{ + @Override + protected QueryTestBuilder testBuilder() + { + return new QueryTestBuilder(new CalciteTestConfig(true)) +// .addCustomRunner(new ExtractResultsFactory(() -> null))//(MSQTestOverlordServiceClient) ((DartSqlEngine) queryFramework().engine()).overlordClient())) +// .addCustomRunner(new ExtractResultsFactory(() -> null))//(MSQTestOverlordServiceClient) ((DartSqlEngine) queryFramework().engine()).overlordClient())) + .queryContext(ImmutableMap.builder().put("asd",UUID.randomUUID().toString()).build()) + .skipVectorize(true) + .verifyNativeQueries(new VerifyMSQSupportedNativeQueriesPredicate()); + } + + @Disabled + @Override + @Test + public void testCannotInsertWithNativeEngine() + { + + } + + @Disabled + @Override + @Test + public void testCannotReplaceWithNativeEngine() + { + + } + + @Disabled + @Override + @Test + public void testRequireTimeConditionSimpleQueryNegative() + { + + } + + @Disabled + @Override + @Test + public void testRequireTimeConditionSubQueryNegative() + { + + } + + @Disabled + @Override + @Test + public void testRequireTimeConditionSemiJoinNegative() + { + + } + + @Disabled + @Override + @Test + public void testExactCountDistinctWithFilter() + { + + } + + @Disabled + @Override + @Test + public void testExactCountDistinctWithFilter2() + { + + } + + @Disabled + @Override + @Test + public void testUnplannableScanOrderByNonTime() + { + + } + + @Disabled + @Override + @Test + public void testUnSupportedNullsFirst() + { + } + + @Disabled + @Override + @Test + public void testUnSupportedNullsLast() + { + } + + /** + * Same query as {@link CalciteQueryTest#testArrayAggQueryOnComplexDatatypes}. ARRAY_AGG is not supported in MSQ currently. + * Once support is added, this test can be removed and msqCompatible() can be added to the one in CalciteQueryTest. + */ + @Test + @Override + public void testArrayAggQueryOnComplexDatatypes() + { + try { + testQuery("SELECT ARRAY_AGG(unique_dim1) FROM druid.foo", ImmutableList.of(), ImmutableList.of()); + Assert.fail("query execution should fail"); + } + catch (ISE e) { + Assert.assertTrue( + e.getMessage().contains("Cannot handle column [a0] with type [ARRAY>]") + ); + } + } + + @Test + @Timeout(value = 40000, unit = TimeUnit.MILLISECONDS) + public void testJoinMultipleTablesWithWhereCondition() + { + testBuilder() + .queryContext( + ImmutableMap.of( + "sqlJoinAlgorithm", "sortMerge" + ) + ) + .sql( + "SELECT f2.dim3,sum(f6.m1 * (1- f6.m2)) FROM" + + " druid.foo as f5, " + + " druid.foo as f6, " + + " druid.numfoo as f7, " + + " druid.foo2 as f2, " + + " druid.numfoo as f3, " + + " druid.foo as f4, " + + " druid.numfoo as f1, " + + " druid.foo2 as f8 " + + "where true" + + " and f1.dim1 = f2.dim2 " + + " and f3.dim1 = f4.dim2 " + + " and f5.dim1 = f6.dim2 " + + " and f7.dim2 = f8.dim3 " + + " and f2.dim1 = f4.dim2 " + + " and f6.dim1 = f8.dim2 " + + " and f1.dim1 = f7.dim2 " + + " and f8.dim2 = 'x' " + + " and f3.__time >= date '2011-11-11' " + + " and f3.__time < date '2013-11-11' " + + "group by 1 " + + "order by 2 desc limit 1001" + ) + .run(); + } +} diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/CalciteSelectQueryMSQTest.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/CalciteSelectQueryMSQTest.java index dda6b8f38067..ba0aca17a114 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/CalciteSelectQueryMSQTest.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/CalciteSelectQueryMSQTest.java @@ -44,7 +44,7 @@ protected QueryTestBuilder testBuilder() { return new QueryTestBuilder(new CalciteTestConfig(true)) // .addCustomRunner(new ExtractResultsFactory(() -> null))//(MSQTestOverlordServiceClient) ((DartSqlEngine) queryFramework().engine()).overlordClient())) - .addCustomRunner(new ExtractResultsFactory(() -> null))//(MSQTestOverlordServiceClient) ((DartSqlEngine) queryFramework().engine()).overlordClient())) +// .addCustomRunner(new ExtractResultsFactory(() -> null))//(MSQTestOverlordServiceClient) ((DartSqlEngine) queryFramework().engine()).overlordClient())) .queryContext(ImmutableMap.builder().put("asd",UUID.randomUUID().toString()).build()) .skipVectorize(true) .verifyNativeQueries(new VerifyMSQSupportedNativeQueriesPredicate()); diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/CalciteQueryTest.java b/sql/src/test/java/org/apache/druid/sql/calcite/CalciteQueryTest.java index 90a64d5fc703..c9606cdea8de 100644 --- a/sql/src/test/java/org/apache/druid/sql/calcite/CalciteQueryTest.java +++ b/sql/src/test/java/org/apache/druid/sql/calcite/CalciteQueryTest.java @@ -5256,8 +5256,8 @@ public void testInIsNotTrueFilter() @Test public void testNotInOrIsNullFilter() { - testQuery( - "SELECT dim1, COUNT(*) FROM druid.foo WHERE dim1 NOT IN ('ghi', 'abc', 'def') OR dim1 IS NULL GROUP BY dim1", + testQuery1( + "SELECT 1", ImmutableList.of( GroupByQuery.builder() .setDataSource(CalciteTests.DATASOURCE1) @@ -5277,14 +5277,55 @@ public void testNotInOrIsNullFilter() .build() ), ImmutableList.of( - new Object[]{"", 1L}, - new Object[]{"1", 1L}, - new Object[]{"10.1", 1L}, - new Object[]{"2", 1L} + new Object[]{1} + ) + ); + } + + + public void testQuery1( + final String sql, + final List> expectedQueries, + final List expectedResults + ) + { + testBuilder() + .sql(sql) +// .expectedQueries(expectedQueries) + .expectedResults(expectedResults) + .run(); + } + + @Test + public void testSelect1() + { + testQuery1( + "SELECT 1", + ImmutableList.of( + GroupByQuery.builder() + .setDataSource(CalciteTests.DATASOURCE1) + .setInterval(querySegmentSpec(Filtration.eternity())) + .setGranularity(Granularities.ALL) + .setDimensions(dimensions(new DefaultDimensionSpec("dim1", "d0"))) + .setDimFilter(or( + isNull("dim1"), + not(in("dim1", ColumnType.STRING, ImmutableList.of("abc", "def", "ghi"))) + )) + .setAggregatorSpecs( + aggregators( + new CountAggregatorFactory("a0") + ) + ) + .setContext(QUERY_CONTEXT_DEFAULT) + .build() + ), + ImmutableList.of( + new Object[]{1L} ) ); } + @Test public void testNotInAndIsNotNullFilter() { From 4e2fd3c74e10152b36c8263cf41ad0089ebbde02 Mon Sep 17 00:00:00 2001 From: Zoltan Haindrich Date: Wed, 5 Mar 2025 13:56:33 +0000 Subject: [PATCH 14/50] m --- .../msq/test/CalciteSelectQueryDartTest.java | 149 +----------------- 1 file changed, 8 insertions(+), 141 deletions(-) diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/CalciteSelectQueryDartTest.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/CalciteSelectQueryDartTest.java index 54eef7678569..8b8022242321 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/CalciteSelectQueryDartTest.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/CalciteSelectQueryDartTest.java @@ -21,166 +21,33 @@ import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; -import org.apache.druid.java.util.common.ISE; -import org.apache.druid.sql.calcite.CalciteQueryTest; +import org.apache.druid.sql.calcite.BaseCalciteQueryTest; import org.apache.druid.sql.calcite.QueryTestBuilder; import org.apache.druid.sql.calcite.SqlTestFrameworkConfig; -import org.junit.Assert; -import org.junit.jupiter.api.Disabled; import org.junit.jupiter.api.Test; -import org.junit.jupiter.api.Timeout; - import java.util.UUID; -import java.util.concurrent.TimeUnit; -/** - * Runs {@link CalciteQueryTest} but with MSQ engine - */ @SqlTestFrameworkConfig.ComponentSupplier(DartComponentSupplier.class) -public class CalciteSelectQueryDartTest extends CalciteQueryTest +public class CalciteSelectQueryDartTest extends BaseCalciteQueryTest { @Override protected QueryTestBuilder testBuilder() { return new QueryTestBuilder(new CalciteTestConfig(true)) -// .addCustomRunner(new ExtractResultsFactory(() -> null))//(MSQTestOverlordServiceClient) ((DartSqlEngine) queryFramework().engine()).overlordClient())) -// .addCustomRunner(new ExtractResultsFactory(() -> null))//(MSQTestOverlordServiceClient) ((DartSqlEngine) queryFramework().engine()).overlordClient())) - .queryContext(ImmutableMap.builder().put("asd",UUID.randomUUID().toString()).build()) + .queryContext(ImmutableMap.builder().put("asd", UUID.randomUUID().toString()).build()) .skipVectorize(true) .verifyNativeQueries(new VerifyMSQSupportedNativeQueriesPredicate()); } - @Disabled - @Override - @Test - public void testCannotInsertWithNativeEngine() - { - - } - - @Disabled - @Override - @Test - public void testCannotReplaceWithNativeEngine() - { - - } - - @Disabled - @Override - @Test - public void testRequireTimeConditionSimpleQueryNegative() - { - - } - - @Disabled - @Override @Test - public void testRequireTimeConditionSubQueryNegative() - { - - } - - @Disabled - @Override - @Test - public void testRequireTimeConditionSemiJoinNegative() - { - - } - - @Disabled - @Override - @Test - public void testExactCountDistinctWithFilter() - { - - } - - @Disabled - @Override - @Test - public void testExactCountDistinctWithFilter2() - { - - } - - @Disabled - @Override - @Test - public void testUnplannableScanOrderByNonTime() - { - - } - - @Disabled - @Override - @Test - public void testUnSupportedNullsFirst() - { - } - - @Disabled - @Override - @Test - public void testUnSupportedNullsLast() - { - } - - /** - * Same query as {@link CalciteQueryTest#testArrayAggQueryOnComplexDatatypes}. ARRAY_AGG is not supported in MSQ currently. - * Once support is added, this test can be removed and msqCompatible() can be added to the one in CalciteQueryTest. - */ - @Test - @Override - public void testArrayAggQueryOnComplexDatatypes() - { - try { - testQuery("SELECT ARRAY_AGG(unique_dim1) FROM druid.foo", ImmutableList.of(), ImmutableList.of()); - Assert.fail("query execution should fail"); - } - catch (ISE e) { - Assert.assertTrue( - e.getMessage().contains("Cannot handle column [a0] with type [ARRAY>]") - ); - } - } - - @Test - @Timeout(value = 40000, unit = TimeUnit.MILLISECONDS) - public void testJoinMultipleTablesWithWhereCondition() + public void testSelect1() { testBuilder() - .queryContext( - ImmutableMap.of( - "sqlJoinAlgorithm", "sortMerge" - ) - ) - .sql( - "SELECT f2.dim3,sum(f6.m1 * (1- f6.m2)) FROM" - + " druid.foo as f5, " - + " druid.foo as f6, " - + " druid.numfoo as f7, " - + " druid.foo2 as f2, " - + " druid.numfoo as f3, " - + " druid.foo as f4, " - + " druid.numfoo as f1, " - + " druid.foo2 as f8 " - + "where true" - + " and f1.dim1 = f2.dim2 " - + " and f3.dim1 = f4.dim2 " - + " and f5.dim1 = f6.dim2 " - + " and f7.dim2 = f8.dim3 " - + " and f2.dim1 = f4.dim2 " - + " and f6.dim1 = f8.dim2 " - + " and f1.dim1 = f7.dim2 " - + " and f8.dim2 = 'x' " - + " and f3.__time >= date '2011-11-11' " - + " and f3.__time < date '2013-11-11' " - + "group by 1 " - + "order by 2 desc limit 1001" + .sql("SELECT 1") + .expectedResults( + ImmutableList.of(new Object[] {1}) ) .run(); } + } From 061a7415e11347a96ac9e6225a3e7ac80d85c2bc Mon Sep 17 00:00:00 2001 From: Zoltan Haindrich Date: Wed, 5 Mar 2025 13:57:58 +0000 Subject: [PATCH 15/50] up --- .../druid/msq/test/CalciteSelectQueryDartTest.java | 12 ++++++++++++ .../apache/druid/sql/calcite/QueryTestBuilder.java | 10 +++++++--- 2 files changed, 19 insertions(+), 3 deletions(-) diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/CalciteSelectQueryDartTest.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/CalciteSelectQueryDartTest.java index 8b8022242321..8c6b3a3b9b38 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/CalciteSelectQueryDartTest.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/CalciteSelectQueryDartTest.java @@ -50,4 +50,16 @@ public void testSelect1() .run(); } + @Test + public void testSelectFromFoo() + { + testBuilder() + .sql("SELECT 2 from foo limit 2") + .expectedResults( + new Object[] {2}, + new Object[] {2} + ) + .run(); + } + } diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/QueryTestBuilder.java b/sql/src/test/java/org/apache/druid/sql/calcite/QueryTestBuilder.java index f53b2f37b1f5..91e275b6d8dd 100644 --- a/sql/src/test/java/org/apache/druid/sql/calcite/QueryTestBuilder.java +++ b/sql/src/test/java/org/apache/druid/sql/calcite/QueryTestBuilder.java @@ -21,6 +21,7 @@ import com.fasterxml.jackson.databind.ObjectMapper; import com.google.common.base.Preconditions; +import org.apache.curator.shaded.com.google.common.collect.ImmutableList; import org.apache.druid.query.Query; import org.apache.druid.query.QueryContexts; import org.apache.druid.quidem.DruidQTestInfo; @@ -163,13 +164,16 @@ public QueryTestBuilder expectedQueries(List> expectedQueries) return this; } - public QueryTestBuilder expectedResults( - final List expectedResults - ) + public QueryTestBuilder expectedResults(final List expectedResults) { return expectedResults(ResultMatchMode.EQUALS, expectedResults); } + public QueryTestBuilder expectedResults(Object[]... expectedResults) + { + return expectedResults(ResultMatchMode.EQUALS, ImmutableList.of(expectedResults)); + } + public QueryTestBuilder expectedResults( ResultMatchMode expecteMatchMode, final List expectedResults From fb4ec06e0cd62e753648f455062b4fa56f00cc16 Mon Sep 17 00:00:00 2001 From: Zoltan Haindrich Date: Wed, 5 Mar 2025 16:08:36 +0000 Subject: [PATCH 16/50] sx --- .../org/apache/druid/msq/exec/WorkerImpl.java | 2 +- .../msq/shuffle/output/StageOutputHolder.java | 33 +++++++++++++++++++ ...QueryDartTest.java => SimpleDartTest.java} | 8 +++-- .../druid/sql/calcite/QueryTestBuilder.java | 6 ---- .../sql/calcite/util/SqlTestFramework.java | 1 + .../calcite/util/TestTimelineServerView.java | 22 ++++++++++++- 6 files changed, 61 insertions(+), 11 deletions(-) rename extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/{CalciteSelectQueryDartTest.java => SimpleDartTest.java} (92%) diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/WorkerImpl.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/WorkerImpl.java index 44654810ac61..21b5e80f29be 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/WorkerImpl.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/WorkerImpl.java @@ -556,7 +556,7 @@ public ListenableFuture readStageOutput( final long offset ) { - return getOrCreateStageOutputHolder(stageId, partitionNumber).readRemotelyFrom(offset); + return getOrCreateStageOutputHolder(stageId, partitionNumber).readLocally12(offset); } /** diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/shuffle/output/StageOutputHolder.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/shuffle/output/StageOutputHolder.java index c19519dfb7bb..84f17d6cd4da 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/shuffle/output/StageOutputHolder.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/shuffle/output/StageOutputHolder.java @@ -22,6 +22,7 @@ import com.google.common.util.concurrent.ListenableFuture; import com.google.common.util.concurrent.SettableFuture; import org.apache.druid.common.guava.FutureUtils; +import org.apache.druid.frame.Frame; import org.apache.druid.frame.channel.ReadableFileFrameChannel; import org.apache.druid.frame.channel.ReadableFrameChannel; import org.apache.druid.frame.channel.ReadableNilFrameChannel; @@ -76,6 +77,38 @@ public ReadableFrameChannel readLocally() return new FutureReadableFrameChannel(FutureUtils.transform(readerFuture, StageOutputReader::readLocally)); } + public ReadableFrameChannel readLocally1() + { + return new FutureReadableFrameChannel(FutureUtils.transform(readerFuture, StageOutputReader::readLocally)); + } + + + public ListenableFuture readLocally12(long offset) + { + + ReadableFrameChannel l = readLocally(); + + ListenableFuture afr = l.readabilityFuture(); + Frame a = l.read(); + long b = a.numBytes(); +// WritableMemory wm = a.writableMemory(); + + return (ListenableFuture) afr; +//// byte[] buffer = new byte[8 * 1024]; +//// boolean didRead = false; +//// int bytesRead; +//// while ((bytesRead = inputStream.read(buffer)) != -1) { +//// channel.addChunk(Arrays.copyOf(buffer, bytesRead)); +//// didRead = true; +//// } +//// inputStream.close(); +// +// +// setChannel(readLocally()); +// ListenableFuture aa = FutureUtils.transform(channelFuture, StageOutputHolder::createReader); +// return FutureUtils.transformAsync(aa, reader -> reader.readRemotelyFrom(offset)); + } + /** * Sets the channel that backs {@link #readLocally()} and {@link #readRemotelyFrom(long)}. */ diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/CalciteSelectQueryDartTest.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/SimpleDartTest.java similarity index 92% rename from extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/CalciteSelectQueryDartTest.java rename to extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/SimpleDartTest.java index 8c6b3a3b9b38..018ea7830ea0 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/CalciteSelectQueryDartTest.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/SimpleDartTest.java @@ -28,7 +28,7 @@ import java.util.UUID; @SqlTestFrameworkConfig.ComponentSupplier(DartComponentSupplier.class) -public class CalciteSelectQueryDartTest extends BaseCalciteQueryTest +public class SimpleDartTest extends BaseCalciteQueryTest { @Override protected QueryTestBuilder testBuilder() @@ -56,8 +56,10 @@ public void testSelectFromFoo() testBuilder() .sql("SELECT 2 from foo limit 2") .expectedResults( - new Object[] {2}, - new Object[] {2} + ImmutableList.of( + new Object[] {2}, + new Object[] {2} + ) ) .run(); } diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/QueryTestBuilder.java b/sql/src/test/java/org/apache/druid/sql/calcite/QueryTestBuilder.java index 91e275b6d8dd..c35003e0e767 100644 --- a/sql/src/test/java/org/apache/druid/sql/calcite/QueryTestBuilder.java +++ b/sql/src/test/java/org/apache/druid/sql/calcite/QueryTestBuilder.java @@ -21,7 +21,6 @@ import com.fasterxml.jackson.databind.ObjectMapper; import com.google.common.base.Preconditions; -import org.apache.curator.shaded.com.google.common.collect.ImmutableList; import org.apache.druid.query.Query; import org.apache.druid.query.QueryContexts; import org.apache.druid.quidem.DruidQTestInfo; @@ -169,11 +168,6 @@ public QueryTestBuilder expectedResults(final List expectedResults) return expectedResults(ResultMatchMode.EQUALS, expectedResults); } - public QueryTestBuilder expectedResults(Object[]... expectedResults) - { - return expectedResults(ResultMatchMode.EQUALS, ImmutableList.of(expectedResults)); - } - public QueryTestBuilder expectedResults( ResultMatchMode expecteMatchMode, final List expectedResults diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/util/SqlTestFramework.java b/sql/src/test/java/org/apache/druid/sql/calcite/util/SqlTestFramework.java index 021f0b394b80..9467fef3433b 100644 --- a/sql/src/test/java/org/apache/druid/sql/calcite/util/SqlTestFramework.java +++ b/sql/src/test/java/org/apache/druid/sql/calcite/util/SqlTestFramework.java @@ -1068,6 +1068,7 @@ public TimelineServerView makeTimelines1(ObjectMapper objectMapper, QueryRunnerF //FIXME connect TSB with this ///. return new SimpleServerView(conglomerate, objectMapper, httpClient); + SimpleServerView simpleServerView = new SimpleServerView(conglomerate, objectMapper, httpClient); return simpleServerView; } diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/util/TestTimelineServerView.java b/sql/src/test/java/org/apache/druid/sql/calcite/util/TestTimelineServerView.java index 4e4c677de86e..c8b8a03185c1 100644 --- a/sql/src/test/java/org/apache/druid/sql/calcite/util/TestTimelineServerView.java +++ b/sql/src/test/java/org/apache/druid/sql/calcite/util/TestTimelineServerView.java @@ -25,7 +25,10 @@ import org.apache.druid.client.ImmutableDruidDataSource; import org.apache.druid.client.ImmutableDruidServer; import org.apache.druid.client.TimelineServerView; +import org.apache.druid.client.selector.HighestPriorityTierSelectorStrategy; +import org.apache.druid.client.selector.RandomServerSelectorStrategy; import org.apache.druid.client.selector.ServerSelector; +import org.apache.druid.client.selector.TierSelectorStrategy; import org.apache.druid.java.util.common.Pair; import org.apache.druid.query.QueryRunner; import org.apache.druid.query.TableDataSource; @@ -33,9 +36,12 @@ import org.apache.druid.server.coordination.ServerType; import org.apache.druid.timeline.DataSegment; import org.apache.druid.timeline.TimelineLookup; - +import org.apache.druid.timeline.VersionedIntervalTimeline; +import org.apache.druid.timeline.partition.PartitionChunk; +import org.apache.druid.timeline.partition.SingleElementPartitionChunk; import java.util.ArrayList; import java.util.Collections; +import java.util.Comparator; import java.util.List; import java.util.Optional; import java.util.concurrent.Executor; @@ -93,6 +99,20 @@ public TestTimelineServerView(List segments, List real @Override public Optional> getTimeline(TableDataSource table) { + for (DataSegment s : segments) { + if (!s.getDataSource().equals(table.getName())) { + continue; + } + VersionedIntervalTimeline a = new VersionedIntervalTimeline(Comparator.naturalOrder()); + TierSelectorStrategy st = new HighestPriorityTierSelectorStrategy(new RandomServerSelectorStrategy()); + ServerSelector sss = new ServerSelector(s, st ); + + PartitionChunk aa = new SingleElementPartitionChunk(sss); + a.add( s.getInterval(), s.getVersion(), aa); + return Optional.of(a); +// return Optional.of(); + + } return Optional.empty(); // throw new UnsupportedOperationException(); } From 1d806d67eab99fef27a853efda7b5273eade3d53 Mon Sep 17 00:00:00 2001 From: Zoltan Haindrich Date: Thu, 6 Mar 2025 08:21:44 +0000 Subject: [PATCH 17/50] crap --- .../dart/worker/DartWorkerFactoryImpl.java | 9 ++++- .../org/apache/druid/msq/exec/WorkerImpl.java | 9 +++-- .../msq/shuffle/output/StageOutputHolder.java | 33 ------------------- 3 files changed, 15 insertions(+), 36 deletions(-) diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartWorkerFactoryImpl.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartWorkerFactoryImpl.java index 1960924b4b67..f5ef8b4b9585 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartWorkerFactoryImpl.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartWorkerFactoryImpl.java @@ -35,6 +35,7 @@ import org.apache.druid.msq.exec.WorkerContext; import org.apache.druid.msq.exec.WorkerImpl; import org.apache.druid.msq.querykit.DataSegmentProvider; +import org.apache.druid.msq.shuffle.output.StageOutputHolder; import org.apache.druid.query.DruidProcessingConfig; import org.apache.druid.query.QueryContext; import org.apache.druid.query.groupby.GroupingEngine; @@ -113,6 +114,12 @@ public Worker build(String queryId, String controllerHost, File tempDir, QueryCo queryContext ); - return new WorkerImpl(null, workerContext); + return new WorkerImpl(null, workerContext) + { + protected StageOutputHolder makeStageOutputHolder() + { + return new StageOutputHolder(); + }; + }; } } diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/WorkerImpl.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/WorkerImpl.java index 21b5e80f29be..ec057e72fea6 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/WorkerImpl.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/WorkerImpl.java @@ -556,7 +556,7 @@ public ListenableFuture readStageOutput( final long offset ) { - return getOrCreateStageOutputHolder(stageId, partitionNumber).readLocally12(offset); + return getOrCreateStageOutputHolder(stageId, partitionNumber).readRemotelyFrom(offset); } /** @@ -964,7 +964,12 @@ private void removeStageDurableStorageOutput(final StageId stageId) private StageOutputHolder getOrCreateStageOutputHolder(final StageId stageId, final int partitionNumber) { return stageOutputs.computeIfAbsent(stageId, ignored1 -> new ConcurrentHashMap<>()) - .computeIfAbsent(partitionNumber, ignored -> new StageOutputHolder()); + .computeIfAbsent(partitionNumber, ignored -> makeStageOutputHolder()); + } + + protected StageOutputHolder makeStageOutputHolder() + { + return new StageOutputHolder(); } /** diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/shuffle/output/StageOutputHolder.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/shuffle/output/StageOutputHolder.java index 84f17d6cd4da..c19519dfb7bb 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/shuffle/output/StageOutputHolder.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/shuffle/output/StageOutputHolder.java @@ -22,7 +22,6 @@ import com.google.common.util.concurrent.ListenableFuture; import com.google.common.util.concurrent.SettableFuture; import org.apache.druid.common.guava.FutureUtils; -import org.apache.druid.frame.Frame; import org.apache.druid.frame.channel.ReadableFileFrameChannel; import org.apache.druid.frame.channel.ReadableFrameChannel; import org.apache.druid.frame.channel.ReadableNilFrameChannel; @@ -77,38 +76,6 @@ public ReadableFrameChannel readLocally() return new FutureReadableFrameChannel(FutureUtils.transform(readerFuture, StageOutputReader::readLocally)); } - public ReadableFrameChannel readLocally1() - { - return new FutureReadableFrameChannel(FutureUtils.transform(readerFuture, StageOutputReader::readLocally)); - } - - - public ListenableFuture readLocally12(long offset) - { - - ReadableFrameChannel l = readLocally(); - - ListenableFuture afr = l.readabilityFuture(); - Frame a = l.read(); - long b = a.numBytes(); -// WritableMemory wm = a.writableMemory(); - - return (ListenableFuture) afr; -//// byte[] buffer = new byte[8 * 1024]; -//// boolean didRead = false; -//// int bytesRead; -//// while ((bytesRead = inputStream.read(buffer)) != -1) { -//// channel.addChunk(Arrays.copyOf(buffer, bytesRead)); -//// didRead = true; -//// } -//// inputStream.close(); -// -// -// setChannel(readLocally()); -// ListenableFuture aa = FutureUtils.transform(channelFuture, StageOutputHolder::createReader); -// return FutureUtils.transformAsync(aa, reader -> reader.readRemotelyFrom(offset)); - } - /** * Sets the channel that backs {@link #readLocally()} and {@link #readRemotelyFrom(long)}. */ From b574226122e9e4fe3ee14ab0a37106a6d44a54b1 Mon Sep 17 00:00:00 2001 From: Zoltan Haindrich Date: Thu, 6 Mar 2025 08:36:26 +0000 Subject: [PATCH 18/50] refact --- .../org/apache/druid/msq/exec/RunWorkOrder.java | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/RunWorkOrder.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/RunWorkOrder.java index 01d2e925aa70..b4f1bbd5eaae 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/RunWorkOrder.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/RunWorkOrder.java @@ -370,13 +370,7 @@ private void makeWorkOutputChannelFactory() final OutputChannelFactory baseOutputChannelFactory; - if (workOrder.getStageDefinition().doesShuffle()) { - // Writing to a consumer in the same JVM (which will be set up later on in this method). - baseOutputChannelFactory = new BlockingQueueOutputChannelFactory(frameContext.memoryParameters().getFrameSize()); - } else { - // Writing stage output. - baseOutputChannelFactory = makeStageOutputChannelFactory(); - } + baseOutputChannelFactory = makeStageOutputChannelFactory(); workOutputChannelFactory = new CountingOutputChannelFactory( baseOutputChannelFactory, @@ -617,6 +611,11 @@ private void writeDurableStorageSuccessFile() private OutputChannelFactory makeStageOutputChannelFactory() { + if (workOrder.getStageDefinition().doesShuffle()) { + // Writing to a consumer in the same JVM (which will be set up later on in this method). + return new BlockingQueueOutputChannelFactory(frameContext.memoryParameters().getFrameSize()); + } + // Use the standard frame size, since we assume this size when computing how much is needed to merge output // files from different workers. final int frameSize = frameContext.memoryParameters().getFrameSize(); From 823d01cd4305d25d93ea44de46166cadc93318f4 Mon Sep 17 00:00:00 2001 From: Zoltan Haindrich Date: Thu, 6 Mar 2025 09:46:55 +0000 Subject: [PATCH 19/50] x --- .../dart/worker/DartWorkerFactoryImpl.java | 8 ++++++- .../org/apache/druid/msq/exec/WorkerImpl.java | 24 +++++++++++++++---- .../TestDartControllerContextFactoryImpl.java | 14 ++++++++--- 3 files changed, 37 insertions(+), 9 deletions(-) diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartWorkerFactoryImpl.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartWorkerFactoryImpl.java index f5ef8b4b9585..70abbb8bcb54 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartWorkerFactoryImpl.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartWorkerFactoryImpl.java @@ -34,6 +34,7 @@ import org.apache.druid.msq.exec.Worker; import org.apache.druid.msq.exec.WorkerContext; import org.apache.druid.msq.exec.WorkerImpl; +import org.apache.druid.msq.kernel.StageId; import org.apache.druid.msq.querykit.DataSegmentProvider; import org.apache.druid.msq.shuffle.output.StageOutputHolder; import org.apache.druid.query.DruidProcessingConfig; @@ -116,10 +117,15 @@ public Worker build(String queryId, String controllerHost, File tempDir, QueryCo return new WorkerImpl(null, workerContext) { - protected StageOutputHolder makeStageOutputHolder() + @Override + protected StageOutputHolder makeStageOutputHolder(StageId stageId, int partitionNumber) { return new StageOutputHolder(); }; }; } + + static class OutputKey { + StageId stageId; int partitionNumber; + } } diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/WorkerImpl.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/WorkerImpl.java index ec057e72fea6..9fface5b958f 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/WorkerImpl.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/WorkerImpl.java @@ -556,7 +556,21 @@ public ListenableFuture readStageOutput( final long offset ) { - return getOrCreateStageOutputHolder(stageId, partitionNumber).readRemotelyFrom(offset); + StageOutputHolder orCreateStageOutputHolder = getOrCreateStageOutputHolder(stageId, partitionNumber); + + return orCreateStageOutputHolder.readRemotelyFrom(offset); + } + + public ReadableFrameChannel readStageOutput1( + final StageId stageId, + final int partitionNumber, + final long offset + ) + { + StageOutputHolder orCreateStageOutputHolder = getOrCreateStageOutputHolder(stageId, partitionNumber); + return orCreateStageOutputHolder.readLocally(); + +// return orCreateStageOutputHolder.readRemotelyFrom(offset); } /** @@ -693,7 +707,7 @@ public CounterSnapshotsTree getCounters() * Create a {@link RunWorkOrderListener} for {@link RunWorkOrder} that hooks back into the {@link KernelHolders} * in the main loop. */ - private RunWorkOrderListener makeRunWorkOrderListener( + protected RunWorkOrderListener makeRunWorkOrderListener( final WorkOrder workOrder, final ControllerClient controllerClient, final Set criticalWarningCodes, @@ -795,7 +809,7 @@ public void onFailure(Throwable t) }; } - private InputChannelFactory makeBaseInputChannelFactory( + protected InputChannelFactory makeBaseInputChannelFactory( final WorkOrder workOrder, final ControllerClient controllerClient, final Closer closer @@ -964,10 +978,10 @@ private void removeStageDurableStorageOutput(final StageId stageId) private StageOutputHolder getOrCreateStageOutputHolder(final StageId stageId, final int partitionNumber) { return stageOutputs.computeIfAbsent(stageId, ignored1 -> new ConcurrentHashMap<>()) - .computeIfAbsent(partitionNumber, ignored -> makeStageOutputHolder()); + .computeIfAbsent(partitionNumber, ignored -> makeStageOutputHolder(stageId, partitionNumber)); } - protected StageOutputHolder makeStageOutputHolder() + protected StageOutputHolder makeStageOutputHolder(StageId stageId, int partitionNumber) { return new StageOutputHolder(); } diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/TestDartControllerContextFactoryImpl.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/TestDartControllerContextFactoryImpl.java index 5bcb162dcab0..7281af94bd13 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/TestDartControllerContextFactoryImpl.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/TestDartControllerContextFactoryImpl.java @@ -33,6 +33,7 @@ import org.apache.druid.msq.dart.controller.DartControllerContext; import org.apache.druid.msq.dart.controller.DartControllerContextFactoryImpl; import org.apache.druid.msq.dart.worker.DartWorkerClient; +import org.apache.druid.msq.dart.worker.DartWorkerClientImpl; import org.apache.druid.msq.exec.Controller; import org.apache.druid.msq.exec.ControllerContext; import org.apache.druid.msq.exec.MemoryIntrospector; @@ -41,6 +42,7 @@ import org.apache.druid.msq.exec.WorkerStorageParameters; import org.apache.druid.msq.kernel.StageId; import org.apache.druid.msq.kernel.WorkOrder; +import org.apache.druid.msq.shuffle.output.StageOutputHolder; import org.apache.druid.rpc.ServiceClientFactory; import org.apache.druid.server.DruidNode; @@ -111,7 +113,7 @@ protected DartWorkerClient makeWorkerClient(String queryId) ); } - public class DartTestWorkerClient extends MSQTestWorkerClient implements DartWorkerClient + public class DartTestWorkerClient extends DartWorkerClientImpl { private final ExecutorService EXECUTOR = Executors.newCachedThreadPool(); public DartControllerContext controllerCtx; @@ -123,7 +125,7 @@ public DartTestWorkerClient( String controllerHost, Map workerMap ) { - super(workerMap); + super(queryId, clientFactory, smileMapper, controllerHost); } @Override @@ -142,7 +144,13 @@ protected Worker newWorker(String workerId) MSQTestBase.makeTestWorkerMemoryParameters(), WorkerStorageParameters.createInstanceForTests(Long.MAX_VALUE) ) - ); + ) { + @Override + protected StageOutputHolder makeStageOutputHolder(StageId stageId, int partitionNumber) + { + return new StageOutputHolder(); + }; + }; Future future = EXECUTOR.submit(() -> { try { From 0a1059e7fcbae4ad0838fa073e687ae36d9754b1 Mon Sep 17 00:00:00 2001 From: Zoltan Haindrich Date: Thu, 6 Mar 2025 10:00:43 +0000 Subject: [PATCH 20/50] messedup? --- .../druid/msq/test/MSQTestWorkerClient.java | 38 ++++++++++--------- .../apache/druid/msq/test/SimpleDartTest.java | 3 +- .../TestDartControllerContextFactoryImpl.java | 5 +-- 3 files changed, 25 insertions(+), 21 deletions(-) diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestWorkerClient.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestWorkerClient.java index d535084e6802..13706292400c 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestWorkerClient.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestWorkerClient.java @@ -21,6 +21,7 @@ import com.google.common.util.concurrent.Futures; import com.google.common.util.concurrent.ListenableFuture; +import org.apache.druid.common.guava.FutureUtils; import org.apache.druid.frame.channel.ReadableByteChunksFrameChannel; import org.apache.druid.frame.key.ClusterByPartitions; import org.apache.druid.java.util.common.ISE; @@ -32,7 +33,6 @@ import org.apache.druid.msq.rpc.SketchEncoding; import org.apache.druid.msq.statistics.ClusterByStatisticsSnapshot; -import java.io.InputStream; import java.util.Arrays; import java.util.Map; import java.util.concurrent.atomic.AtomicBoolean; @@ -134,22 +134,26 @@ public ListenableFuture fetchChannelData( final ReadableByteChunksFrameChannel channel ) { - try (InputStream inputStream = - getWorkerFor(workerTaskId).readStageOutput(stageId, partitionNumber, offset).get()) { - byte[] buffer = new byte[8 * 1024]; - boolean didRead = false; - int bytesRead; - while ((bytesRead = inputStream.read(buffer)) != -1) { - channel.addChunk(Arrays.copyOf(buffer, bytesRead)); - didRead = true; - } - inputStream.close(); - - return Futures.immediateFuture(!didRead); - } - catch (Exception e) { - throw new ISE(e, "Error reading frame file channel"); - } + return FutureUtils.transform( + getWorkerFor(workerTaskId).readStageOutput(stageId, partitionNumber, offset), + inputStream -> { + try { + byte[] buffer = new byte[8 * 1024]; + boolean didRead = false; + int bytesRead; + while ((bytesRead = inputStream.read(buffer)) != -1) { + channel.addChunk(Arrays.copyOf(buffer, bytesRead)); + didRead = true; + } + inputStream.close(); + + return !didRead; + } + catch (Exception e) { + throw new ISE(e, "Error reading frame file channel"); + } + } + ); } @Override diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/SimpleDartTest.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/SimpleDartTest.java index 018ea7830ea0..3a6343d6967d 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/SimpleDartTest.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/SimpleDartTest.java @@ -54,7 +54,8 @@ public void testSelect1() public void testSelectFromFoo() { testBuilder() - .sql("SELECT 2 from foo limit 2") +// .sql("SELECT 2 from foo limit 2") + .sql("SELECT 2 from foo order by dim1") .expectedResults( ImmutableList.of( new Object[] {2}, diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/TestDartControllerContextFactoryImpl.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/TestDartControllerContextFactoryImpl.java index 7281af94bd13..99827e247711 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/TestDartControllerContextFactoryImpl.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/TestDartControllerContextFactoryImpl.java @@ -33,7 +33,6 @@ import org.apache.druid.msq.dart.controller.DartControllerContext; import org.apache.druid.msq.dart.controller.DartControllerContextFactoryImpl; import org.apache.druid.msq.dart.worker.DartWorkerClient; -import org.apache.druid.msq.dart.worker.DartWorkerClientImpl; import org.apache.druid.msq.exec.Controller; import org.apache.druid.msq.exec.ControllerContext; import org.apache.druid.msq.exec.MemoryIntrospector; @@ -113,7 +112,7 @@ protected DartWorkerClient makeWorkerClient(String queryId) ); } - public class DartTestWorkerClient extends DartWorkerClientImpl + public class DartTestWorkerClient extends MSQTestWorkerClient implements DartWorkerClient { private final ExecutorService EXECUTOR = Executors.newCachedThreadPool(); public DartControllerContext controllerCtx; @@ -125,7 +124,7 @@ public DartTestWorkerClient( String controllerHost, Map workerMap ) { - super(queryId, clientFactory, smileMapper, controllerHost); + super(workerMap); } @Override From d788f9fe506052c55a9446e22eb3e8107b71a34d Mon Sep 17 00:00:00 2001 From: Zoltan Haindrich Date: Thu, 6 Mar 2025 10:11:20 +0000 Subject: [PATCH 21/50] unlucky-for-some --- .../dart/worker/DartWorkerFactoryImpl.java | 11 ++-- .../org/apache/druid/msq/exec/WorkerImpl.java | 27 ++------- .../msq/shuffle/output/StageOutputHolder.java | 33 +++++++++++ .../druid/msq/test/MSQTestWorkerClient.java | 58 +++++++++++++------ .../apache/druid/msq/test/SimpleDartTest.java | 1 - .../TestDartControllerContextFactoryImpl.java | 9 +-- 6 files changed, 83 insertions(+), 56 deletions(-) diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartWorkerFactoryImpl.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartWorkerFactoryImpl.java index 70abbb8bcb54..04737b3225e0 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartWorkerFactoryImpl.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartWorkerFactoryImpl.java @@ -36,7 +36,6 @@ import org.apache.druid.msq.exec.WorkerImpl; import org.apache.druid.msq.kernel.StageId; import org.apache.druid.msq.querykit.DataSegmentProvider; -import org.apache.druid.msq.shuffle.output.StageOutputHolder; import org.apache.druid.query.DruidProcessingConfig; import org.apache.druid.query.QueryContext; import org.apache.druid.query.groupby.GroupingEngine; @@ -117,11 +116,11 @@ public Worker build(String queryId, String controllerHost, File tempDir, QueryCo return new WorkerImpl(null, workerContext) { - @Override - protected StageOutputHolder makeStageOutputHolder(StageId stageId, int partitionNumber) - { - return new StageOutputHolder(); - }; +// @Override +// protected StageOutputHolder makeStageOutputHolder(StageId stageId, int partitionNumber) +// { +// return new StageOutputHolder(); +// }; }; } diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/WorkerImpl.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/WorkerImpl.java index 9fface5b958f..44654810ac61 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/WorkerImpl.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/WorkerImpl.java @@ -556,21 +556,7 @@ public ListenableFuture readStageOutput( final long offset ) { - StageOutputHolder orCreateStageOutputHolder = getOrCreateStageOutputHolder(stageId, partitionNumber); - - return orCreateStageOutputHolder.readRemotelyFrom(offset); - } - - public ReadableFrameChannel readStageOutput1( - final StageId stageId, - final int partitionNumber, - final long offset - ) - { - StageOutputHolder orCreateStageOutputHolder = getOrCreateStageOutputHolder(stageId, partitionNumber); - return orCreateStageOutputHolder.readLocally(); - -// return orCreateStageOutputHolder.readRemotelyFrom(offset); + return getOrCreateStageOutputHolder(stageId, partitionNumber).readRemotelyFrom(offset); } /** @@ -707,7 +693,7 @@ public CounterSnapshotsTree getCounters() * Create a {@link RunWorkOrderListener} for {@link RunWorkOrder} that hooks back into the {@link KernelHolders} * in the main loop. */ - protected RunWorkOrderListener makeRunWorkOrderListener( + private RunWorkOrderListener makeRunWorkOrderListener( final WorkOrder workOrder, final ControllerClient controllerClient, final Set criticalWarningCodes, @@ -809,7 +795,7 @@ public void onFailure(Throwable t) }; } - protected InputChannelFactory makeBaseInputChannelFactory( + private InputChannelFactory makeBaseInputChannelFactory( final WorkOrder workOrder, final ControllerClient controllerClient, final Closer closer @@ -978,12 +964,7 @@ private void removeStageDurableStorageOutput(final StageId stageId) private StageOutputHolder getOrCreateStageOutputHolder(final StageId stageId, final int partitionNumber) { return stageOutputs.computeIfAbsent(stageId, ignored1 -> new ConcurrentHashMap<>()) - .computeIfAbsent(partitionNumber, ignored -> makeStageOutputHolder(stageId, partitionNumber)); - } - - protected StageOutputHolder makeStageOutputHolder(StageId stageId, int partitionNumber) - { - return new StageOutputHolder(); + .computeIfAbsent(partitionNumber, ignored -> new StageOutputHolder()); } /** diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/shuffle/output/StageOutputHolder.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/shuffle/output/StageOutputHolder.java index c19519dfb7bb..84f17d6cd4da 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/shuffle/output/StageOutputHolder.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/shuffle/output/StageOutputHolder.java @@ -22,6 +22,7 @@ import com.google.common.util.concurrent.ListenableFuture; import com.google.common.util.concurrent.SettableFuture; import org.apache.druid.common.guava.FutureUtils; +import org.apache.druid.frame.Frame; import org.apache.druid.frame.channel.ReadableFileFrameChannel; import org.apache.druid.frame.channel.ReadableFrameChannel; import org.apache.druid.frame.channel.ReadableNilFrameChannel; @@ -76,6 +77,38 @@ public ReadableFrameChannel readLocally() return new FutureReadableFrameChannel(FutureUtils.transform(readerFuture, StageOutputReader::readLocally)); } + public ReadableFrameChannel readLocally1() + { + return new FutureReadableFrameChannel(FutureUtils.transform(readerFuture, StageOutputReader::readLocally)); + } + + + public ListenableFuture readLocally12(long offset) + { + + ReadableFrameChannel l = readLocally(); + + ListenableFuture afr = l.readabilityFuture(); + Frame a = l.read(); + long b = a.numBytes(); +// WritableMemory wm = a.writableMemory(); + + return (ListenableFuture) afr; +//// byte[] buffer = new byte[8 * 1024]; +//// boolean didRead = false; +//// int bytesRead; +//// while ((bytesRead = inputStream.read(buffer)) != -1) { +//// channel.addChunk(Arrays.copyOf(buffer, bytesRead)); +//// didRead = true; +//// } +//// inputStream.close(); +// +// +// setChannel(readLocally()); +// ListenableFuture aa = FutureUtils.transform(channelFuture, StageOutputHolder::createReader); +// return FutureUtils.transformAsync(aa, reader -> reader.readRemotelyFrom(offset)); + } + /** * Sets the channel that backs {@link #readLocally()} and {@link #readRemotelyFrom(long)}. */ diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestWorkerClient.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestWorkerClient.java index 13706292400c..f084e3ab04f4 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestWorkerClient.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestWorkerClient.java @@ -33,6 +33,7 @@ import org.apache.druid.msq.rpc.SketchEncoding; import org.apache.druid.msq.statistics.ClusterByStatisticsSnapshot; +import java.io.InputStream; import java.util.Arrays; import java.util.Map; import java.util.concurrent.atomic.AtomicBoolean; @@ -134,26 +135,47 @@ public ListenableFuture fetchChannelData( final ReadableByteChunksFrameChannel channel ) { - return FutureUtils.transform( - getWorkerFor(workerTaskId).readStageOutput(stageId, partitionNumber, offset), - inputStream -> { - try { - byte[] buffer = new byte[8 * 1024]; - boolean didRead = false; - int bytesRead; - while ((bytesRead = inputStream.read(buffer)) != -1) { - channel.addChunk(Arrays.copyOf(buffer, bytesRead)); - didRead = true; + if(false) { + return FutureUtils.transform( + getWorkerFor(workerTaskId).readStageOutput(stageId, partitionNumber, offset), + inputStream -> { + try { + byte[] buffer = new byte[8 * 1024]; + boolean didRead = false; + int bytesRead; + while ((bytesRead = inputStream.read(buffer)) != -1) { + channel.addChunk(Arrays.copyOf(buffer, bytesRead)); + didRead = true; + } + inputStream.close(); + + return !didRead; + } + catch (Exception e) { + throw new ISE(e, "Error reading frame file channel"); } - inputStream.close(); - - return !didRead; - } - catch (Exception e) { - throw new ISE(e, "Error reading frame file channel"); } - } - ); + ); + } else + { + try (InputStream inputStream = + inMemoryWorkers.get(workerTaskId).readStageOutput(stageId, partitionNumber, offset).get()) { + byte[] buffer = new byte[8 * 1024]; + boolean didRead = false; + int bytesRead; + while ((bytesRead = inputStream.read(buffer)) != -1) { + channel.addChunk(Arrays.copyOf(buffer, bytesRead)); + didRead = true; + } + inputStream.close(); + + return Futures.immediateFuture(!didRead); + } + catch (Exception e) { + throw new ISE(e, "Error reading frame file channel"); + } + + } } @Override diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/SimpleDartTest.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/SimpleDartTest.java index 3a6343d6967d..09b675b07545 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/SimpleDartTest.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/SimpleDartTest.java @@ -54,7 +54,6 @@ public void testSelect1() public void testSelectFromFoo() { testBuilder() -// .sql("SELECT 2 from foo limit 2") .sql("SELECT 2 from foo order by dim1") .expectedResults( ImmutableList.of( diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/TestDartControllerContextFactoryImpl.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/TestDartControllerContextFactoryImpl.java index 99827e247711..5bcb162dcab0 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/TestDartControllerContextFactoryImpl.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/TestDartControllerContextFactoryImpl.java @@ -41,7 +41,6 @@ import org.apache.druid.msq.exec.WorkerStorageParameters; import org.apache.druid.msq.kernel.StageId; import org.apache.druid.msq.kernel.WorkOrder; -import org.apache.druid.msq.shuffle.output.StageOutputHolder; import org.apache.druid.rpc.ServiceClientFactory; import org.apache.druid.server.DruidNode; @@ -143,13 +142,7 @@ protected Worker newWorker(String workerId) MSQTestBase.makeTestWorkerMemoryParameters(), WorkerStorageParameters.createInstanceForTests(Long.MAX_VALUE) ) - ) { - @Override - protected StageOutputHolder makeStageOutputHolder(StageId stageId, int partitionNumber) - { - return new StageOutputHolder(); - }; - }; + ); Future future = EXECUTOR.submit(() -> { try { From 8f053d10eada6a60978af22305a36bf076b69041 Mon Sep 17 00:00:00 2001 From: Zoltan Haindrich Date: Thu, 6 Mar 2025 10:15:29 +0000 Subject: [PATCH 22/50] some changes which have broke it --- .../dart/worker/DartWorkerFactoryImpl.java | 15 ++--- .../apache/druid/msq/exec/RunWorkOrder.java | 13 +++-- .../org/apache/druid/msq/exec/WorkerImpl.java | 7 ++- .../msq/shuffle/output/StageOutputHolder.java | 33 ----------- .../druid/msq/test/MSQTestWorkerClient.java | 56 +++++-------------- 5 files changed, 33 insertions(+), 91 deletions(-) diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartWorkerFactoryImpl.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartWorkerFactoryImpl.java index 04737b3225e0..f5ef8b4b9585 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartWorkerFactoryImpl.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartWorkerFactoryImpl.java @@ -34,8 +34,8 @@ import org.apache.druid.msq.exec.Worker; import org.apache.druid.msq.exec.WorkerContext; import org.apache.druid.msq.exec.WorkerImpl; -import org.apache.druid.msq.kernel.StageId; import org.apache.druid.msq.querykit.DataSegmentProvider; +import org.apache.druid.msq.shuffle.output.StageOutputHolder; import org.apache.druid.query.DruidProcessingConfig; import org.apache.druid.query.QueryContext; import org.apache.druid.query.groupby.GroupingEngine; @@ -116,15 +116,10 @@ public Worker build(String queryId, String controllerHost, File tempDir, QueryCo return new WorkerImpl(null, workerContext) { -// @Override -// protected StageOutputHolder makeStageOutputHolder(StageId stageId, int partitionNumber) -// { -// return new StageOutputHolder(); -// }; + protected StageOutputHolder makeStageOutputHolder() + { + return new StageOutputHolder(); + }; }; } - - static class OutputKey { - StageId stageId; int partitionNumber; - } } diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/RunWorkOrder.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/RunWorkOrder.java index b4f1bbd5eaae..01d2e925aa70 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/RunWorkOrder.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/RunWorkOrder.java @@ -370,7 +370,13 @@ private void makeWorkOutputChannelFactory() final OutputChannelFactory baseOutputChannelFactory; - baseOutputChannelFactory = makeStageOutputChannelFactory(); + if (workOrder.getStageDefinition().doesShuffle()) { + // Writing to a consumer in the same JVM (which will be set up later on in this method). + baseOutputChannelFactory = new BlockingQueueOutputChannelFactory(frameContext.memoryParameters().getFrameSize()); + } else { + // Writing stage output. + baseOutputChannelFactory = makeStageOutputChannelFactory(); + } workOutputChannelFactory = new CountingOutputChannelFactory( baseOutputChannelFactory, @@ -611,11 +617,6 @@ private void writeDurableStorageSuccessFile() private OutputChannelFactory makeStageOutputChannelFactory() { - if (workOrder.getStageDefinition().doesShuffle()) { - // Writing to a consumer in the same JVM (which will be set up later on in this method). - return new BlockingQueueOutputChannelFactory(frameContext.memoryParameters().getFrameSize()); - } - // Use the standard frame size, since we assume this size when computing how much is needed to merge output // files from different workers. final int frameSize = frameContext.memoryParameters().getFrameSize(); diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/WorkerImpl.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/WorkerImpl.java index 44654810ac61..ec057e72fea6 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/WorkerImpl.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/WorkerImpl.java @@ -964,7 +964,12 @@ private void removeStageDurableStorageOutput(final StageId stageId) private StageOutputHolder getOrCreateStageOutputHolder(final StageId stageId, final int partitionNumber) { return stageOutputs.computeIfAbsent(stageId, ignored1 -> new ConcurrentHashMap<>()) - .computeIfAbsent(partitionNumber, ignored -> new StageOutputHolder()); + .computeIfAbsent(partitionNumber, ignored -> makeStageOutputHolder()); + } + + protected StageOutputHolder makeStageOutputHolder() + { + return new StageOutputHolder(); } /** diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/shuffle/output/StageOutputHolder.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/shuffle/output/StageOutputHolder.java index 84f17d6cd4da..c19519dfb7bb 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/shuffle/output/StageOutputHolder.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/shuffle/output/StageOutputHolder.java @@ -22,7 +22,6 @@ import com.google.common.util.concurrent.ListenableFuture; import com.google.common.util.concurrent.SettableFuture; import org.apache.druid.common.guava.FutureUtils; -import org.apache.druid.frame.Frame; import org.apache.druid.frame.channel.ReadableFileFrameChannel; import org.apache.druid.frame.channel.ReadableFrameChannel; import org.apache.druid.frame.channel.ReadableNilFrameChannel; @@ -77,38 +76,6 @@ public ReadableFrameChannel readLocally() return new FutureReadableFrameChannel(FutureUtils.transform(readerFuture, StageOutputReader::readLocally)); } - public ReadableFrameChannel readLocally1() - { - return new FutureReadableFrameChannel(FutureUtils.transform(readerFuture, StageOutputReader::readLocally)); - } - - - public ListenableFuture readLocally12(long offset) - { - - ReadableFrameChannel l = readLocally(); - - ListenableFuture afr = l.readabilityFuture(); - Frame a = l.read(); - long b = a.numBytes(); -// WritableMemory wm = a.writableMemory(); - - return (ListenableFuture) afr; -//// byte[] buffer = new byte[8 * 1024]; -//// boolean didRead = false; -//// int bytesRead; -//// while ((bytesRead = inputStream.read(buffer)) != -1) { -//// channel.addChunk(Arrays.copyOf(buffer, bytesRead)); -//// didRead = true; -//// } -//// inputStream.close(); -// -// -// setChannel(readLocally()); -// ListenableFuture aa = FutureUtils.transform(channelFuture, StageOutputHolder::createReader); -// return FutureUtils.transformAsync(aa, reader -> reader.readRemotelyFrom(offset)); - } - /** * Sets the channel that backs {@link #readLocally()} and {@link #readRemotelyFrom(long)}. */ diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestWorkerClient.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestWorkerClient.java index f084e3ab04f4..d535084e6802 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestWorkerClient.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestWorkerClient.java @@ -21,7 +21,6 @@ import com.google.common.util.concurrent.Futures; import com.google.common.util.concurrent.ListenableFuture; -import org.apache.druid.common.guava.FutureUtils; import org.apache.druid.frame.channel.ReadableByteChunksFrameChannel; import org.apache.druid.frame.key.ClusterByPartitions; import org.apache.druid.java.util.common.ISE; @@ -135,46 +134,21 @@ public ListenableFuture fetchChannelData( final ReadableByteChunksFrameChannel channel ) { - if(false) { - return FutureUtils.transform( - getWorkerFor(workerTaskId).readStageOutput(stageId, partitionNumber, offset), - inputStream -> { - try { - byte[] buffer = new byte[8 * 1024]; - boolean didRead = false; - int bytesRead; - while ((bytesRead = inputStream.read(buffer)) != -1) { - channel.addChunk(Arrays.copyOf(buffer, bytesRead)); - didRead = true; - } - inputStream.close(); - - return !didRead; - } - catch (Exception e) { - throw new ISE(e, "Error reading frame file channel"); - } - } - ); - } else - { - try (InputStream inputStream = - inMemoryWorkers.get(workerTaskId).readStageOutput(stageId, partitionNumber, offset).get()) { - byte[] buffer = new byte[8 * 1024]; - boolean didRead = false; - int bytesRead; - while ((bytesRead = inputStream.read(buffer)) != -1) { - channel.addChunk(Arrays.copyOf(buffer, bytesRead)); - didRead = true; - } - inputStream.close(); - - return Futures.immediateFuture(!didRead); - } - catch (Exception e) { - throw new ISE(e, "Error reading frame file channel"); - } - + try (InputStream inputStream = + getWorkerFor(workerTaskId).readStageOutput(stageId, partitionNumber, offset).get()) { + byte[] buffer = new byte[8 * 1024]; + boolean didRead = false; + int bytesRead; + while ((bytesRead = inputStream.read(buffer)) != -1) { + channel.addChunk(Arrays.copyOf(buffer, bytesRead)); + didRead = true; + } + inputStream.close(); + + return Futures.immediateFuture(!didRead); + } + catch (Exception e) { + throw new ISE(e, "Error reading frame file channel"); } } From 7cec2d273d8c7218259493d4eae432bed4b3f5b2 Mon Sep 17 00:00:00 2001 From: Zoltan Haindrich Date: Thu, 6 Mar 2025 12:08:26 +0000 Subject: [PATCH 23/50] t --- .../apache/druid/msq/test/SimpleDartTest.java | 30 +++++++++++++++++++ 1 file changed, 30 insertions(+) diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/SimpleDartTest.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/SimpleDartTest.java index 09b675b07545..3902d9a7edb2 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/SimpleDartTest.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/SimpleDartTest.java @@ -64,4 +64,34 @@ public void testSelectFromFoo() .run(); } + @Test + public void testSelectFromFoo11() + { + testBuilder() + .sql("SELECT dim1 from foo ") + .expectedResults( + ImmutableList.of( + new Object[] {2}, + new Object[] {2} + ) + ) + .run(); + } + @Test + public void testSelectFromFoo1() + { + String sql = "SELECT dim1, COUNT(*) FROM druid.foo " + + "WHERE dim1 NOT IN ('ghi', 'abc', 'def') AND dim1 IS NOT NULL " + + "GROUP BY dim1"; + testBuilder() + .sql(sql) + .expectedResults( + ImmutableList.of( + new Object[] {2}, + new Object[] {2} + ) + ) + .run(); + } + } From 35d2294f336e2a2ed7f4a4b2f5c9cb2577f660ab Mon Sep 17 00:00:00 2001 From: Zoltan Haindrich Date: Thu, 6 Mar 2025 12:15:43 +0000 Subject: [PATCH 24/50] update expecttions --- .../druid/msq/test/MSQTestWorkerClient.java | 2 +- .../apache/druid/msq/test/SimpleDartTest.java | 20 ++++++++++++++----- 2 files changed, 16 insertions(+), 6 deletions(-) diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestWorkerClient.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestWorkerClient.java index d535084e6802..5e37133d77dd 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestWorkerClient.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestWorkerClient.java @@ -57,7 +57,6 @@ public ListenableFuture postWorkOrder(String workerTaskId, WorkOrder workO protected Worker getWorkerFor(String workerTaskId) { - return inMemoryWorkers.computeIfAbsent(workerTaskId, this::newWorker); } @@ -157,6 +156,7 @@ public void close() { if (closed.compareAndSet(false, true)) { inMemoryWorkers.forEach((k, v) -> v.stop()); + inMemoryWorkers.clear(); } } } diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/SimpleDartTest.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/SimpleDartTest.java index 3902d9a7edb2..4e81e2104e5e 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/SimpleDartTest.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/SimpleDartTest.java @@ -57,6 +57,10 @@ public void testSelectFromFoo() .sql("SELECT 2 from foo order by dim1") .expectedResults( ImmutableList.of( + new Object[] {2}, + new Object[] {2}, + new Object[] {2}, + new Object[] {2}, new Object[] {2}, new Object[] {2} ) @@ -68,11 +72,15 @@ public void testSelectFromFoo() public void testSelectFromFoo11() { testBuilder() - .sql("SELECT dim1 from foo ") + .sql("SELECT dim1 from foo") .expectedResults( ImmutableList.of( - new Object[] {2}, - new Object[] {2} + new Object[]{""}, + new Object[]{"10.1"}, + new Object[]{"2"}, + new Object[]{"1"}, + new Object[]{"def"}, + new Object[]{"abc"} ) ) .run(); @@ -87,8 +95,10 @@ public void testSelectFromFoo1() .sql(sql) .expectedResults( ImmutableList.of( - new Object[] {2}, - new Object[] {2} + new Object[]{"", 1L}, + new Object[]{"1", 1L}, + new Object[]{"10.1", 1L}, + new Object[]{"2", 1L} ) ) .run(); From 008a15547fdc4cf5123da395b5e4a2fca4d10f45 Mon Sep 17 00:00:00 2001 From: Zoltan Haindrich Date: Thu, 6 Mar 2025 14:11:33 +0000 Subject: [PATCH 25/50] u --- .../controller/DartControllerContext.java | 9 +- .../TestDartControllerContextFactoryImpl.java | 184 +----------------- .../sql/calcite/util/SqlTestFramework.java | 12 +- 3 files changed, 21 insertions(+), 184 deletions(-) diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartControllerContext.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartControllerContext.java index 1d7056a99c7d..0947bcce180c 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartControllerContext.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartControllerContext.java @@ -123,9 +123,6 @@ public ControllerQueryKernelConfig queryKernelConfig( // allowed to float. If a segment moves to a new server that isn't part of our list after the WorkerManager is // created, we won't be able to find a valid server for certain segments. This isn't expected to be a problem, // since the serverView is referenced shortly after the worker list is created. - if(true) { -// throw DruidException.defensive("asf|"); - } final List workerIds = new ArrayList<>(servers.size()); for (final DruidServerMetadata server : servers) { if (server.getType() == ServerType.HISTORICAL) { @@ -221,9 +218,9 @@ public void registerController(Controller controller, Closer closer) { closer.register(workerClient); // FIXME: - //workerClient.setController(); -// controller - theController=controller; + // workerClient.setController(); + // controller + theController = controller; } diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/TestDartControllerContextFactoryImpl.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/TestDartControllerContextFactoryImpl.java index 5bcb162dcab0..fce5a6a24042 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/TestDartControllerContextFactoryImpl.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/TestDartControllerContextFactoryImpl.java @@ -47,7 +47,6 @@ import java.util.Map; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; -import java.util.concurrent.Future; public class TestDartControllerContextFactoryImpl extends DartControllerContextFactoryImpl { @@ -63,8 +62,7 @@ public TestDartControllerContextFactoryImpl( final MemoryIntrospector memoryIntrospector, final TimelineServerView serverView, final ServiceEmitter emitter, - @Dart Map workerMap - ) + @Dart Map workerMap) { super(injector, jsonMapper, smileMapper, selfNode, serviceClientFactory, memoryIntrospector, serverView, emitter); this.workerMap = workerMap; @@ -75,37 +73,22 @@ public ControllerContext newContext(String queryId) { DartControllerContext ctx; DartTestWorkerClient wc; - ctx=new DartControllerContext( + ctx = new DartControllerContext( injector, jsonMapper, selfNode, - wc=(DartTestWorkerClient) makeWorkerClient(queryId), + wc = (DartTestWorkerClient) makeWorkerClient(queryId), memoryIntrospector, serverView, emitter ); - wc.controllerCtx=ctx; + wc.controllerCtx = ctx; return ctx; } @Override protected DartWorkerClient makeWorkerClient(String queryId) { -// Controller controller = null; -// Worker worker = new WorkerImpl( -// null, -// new MSQTestWorkerContext( -// queryId, -// workerMap, -// controller, -// jsonMapper, -// injector, -// MSQTestBase.makeTestWorkerMemoryParameters(), -// WorkerStorageParameters.createInstanceForTests(Long.MAX_VALUE) -// ) -// ); - -// injector. return new DartTestWorkerClient( queryId, serviceClientFactory, smileMapper, selfNode.getHostAndPortToUse(), workerMap ); @@ -120,8 +103,7 @@ public DartTestWorkerClient( String queryId, ServiceClientFactory clientFactory, ObjectMapper smileMapper, - String controllerHost, Map workerMap - ) + String controllerHost, Map workerMap) { super(workerMap); } @@ -129,8 +111,8 @@ public DartTestWorkerClient( @Override protected Worker newWorker(String workerId) { - String queryId=workerId; - Controller controller=controllerCtx.theController; + String queryId = workerId; + Controller controller = controllerCtx.theController; Worker worker = new WorkerImpl( null, new MSQTestWorkerContext( @@ -144,7 +126,7 @@ protected Worker newWorker(String workerId) ) ); - Future future = EXECUTOR.submit(() -> { + EXECUTOR.submit(() -> { try { worker.run(); } @@ -159,172 +141,26 @@ protected Worker newWorker(String workerId) @Override public ListenableFuture postWorkOrder(String workerTaskId, WorkOrder workOrder) { - return super.postWorkOrder(workerTaskId, workOrder); - } + @Override public ListenableFuture postCleanupStage(String workerTaskId, StageId stageId) { -// if(true) -// { -// throw new RuntimeException("FIXME: Unimplemented!"); -// } return super.postCleanupStage(workerTaskId, stageId); } + @Override public void closeClient(String hostAndPort) { -// if(true) -// { -// throw new RuntimeException("FIXME: Unimplemented!"); -// } - } @Override public ListenableFuture stopWorker(String workerId) { -// if(true) -// { -// throw new RuntimeException("FIXME: Unimplemented!"); -// } return null; } } - - // @Override - // protected ServiceClient getClient(String workerIdString) - // { - // return super.getClient(workerIdString); - // } - // - // @Override - // public void closeClient(String workerHost) - // { - // super.closeClient(workerHost); - // } - // @Override - // protected Object clone() throws CloneNotSupportedException - // { - // return super.clone(); - // } - // @Override - // public ListenableFuture stopWorker(String workerId) - // { - // return super.stopWorker(workerId); - // } - // @Override - // public ListenableFuture fetchChannelData(String workerId, StageId - // stageId, int partitionNumber, - // long offset, ReadableByteChunksFrameChannel channel) - // { - // return super.fetchChannelData(workerId, stageId, partitionNumber, offset, - // channel); - // } - // - // @Override - // public ListenableFuture - // fetchClusterByStatisticsSnapshot(String workerId, - // StageId stageId, SketchEncoding sketchEncoding) - // { - // return super.fetchClusterByStatisticsSnapshot(workerId, stageId, - // sketchEncoding); - // - // } - // @Override - // public ListenableFuture - // fetchClusterByStatisticsSnapshotForTimeChunk(String workerId, - // StageId stageId, long timeChunk, SketchEncoding sketchEncoding) - // { - // return super.fetchClusterByStatisticsSnapshotForTimeChunk(workerId, - // stageId, timeChunk, sketchEncoding); - // - // } - // @Override - // public ListenableFuture getCounters(String workerId) - // { - // return super.getCounters(workerId); - // - // } - // @Override - // public ListenableFuture postCleanupStage(String workerId, StageId - // stageId) - // { - // return super.postCleanupStage(workerId, stageId); - // - // } - // @Override - // public ListenableFuture postFinish(String workerId) - // { - // return super.postFinish(workerId); - // - // } - // - // @Override - // public ListenableFuture postResultPartitionBoundaries(String - // workerId, StageId stageId, - // ClusterByPartitions partitionBoundaries) - // { - // return super.postResultPartitionBoundaries(workerId, stageId, - // partitionBoundaries); - // - // } - // @Override - // public ListenableFuture postWorkOrder(String workerId, WorkOrder - // workOrder) - // { - // return super.postWorkOrder(workerId, workOrder); - // - // } - // - // @Override - // protected Pair makeNewClient(WorkerId workerId) - // { - // ServiceClient client =new TestDartServiceClient(workerId); - // Closeable c=null; - // Pair p = Pair.of(client, c); - // return p; - // } - // - // - // } - // - // static class TestDartServiceClient implements ServiceClient { - // - // - // private MSQTestWorkerClient a; - // - // public TestDartServiceClient(WorkerId workerId) - // { - // a = new MSQTestWorkerClient(null); - // } - // - // @Override - // public ListenableFuture - // asyncRequest(RequestBuilder requestBuilder, - // HttpResponseHandler handler) - // { - // if(true) - // { - // throw new RuntimeException("FIXME: Unimplemented!"); - // } - // return null; - // - // } - // - // @Override - // public ServiceClient withRetryPolicy(ServiceRetryPolicy retryPolicy) - // { - // if(true) - // { - // throw new RuntimeException("FIXME: Unimplemented!"); - // } - // return null; - // - // } - // - // } } diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/util/SqlTestFramework.java b/sql/src/test/java/org/apache/druid/sql/calcite/util/SqlTestFramework.java index 685f0cd0a380..cff7434bca27 100644 --- a/sql/src/test/java/org/apache/druid/sql/calcite/util/SqlTestFramework.java +++ b/sql/src/test/java/org/apache/druid/sql/calcite/util/SqlTestFramework.java @@ -791,14 +791,18 @@ public void configure(Binder binder) .annotatedWith(Merging.class) .to(TestBufferPool.class); - // FIXME - DruidNode self = new DruidNode("test", "test-host", true, 80, 443, false, true); - binder.bind(DruidNode.class).annotatedWith(Self.class).toInstance(self); - TestRequestLogger testRequestLogger = new TestRequestLogger(); binder.bind(RequestLogger.class).toInstance(testRequestLogger); } + @Provides + @Self + @LazySingleton + public DruidNode makeSelfDruidNode() + { + return new DruidNode("test", "test-host", true, 80, 443, false, true); + } + @Provides AuthorizerMapper getAuthorizerMapper() { From 637ae373768503926ead778ab4ed58544bcbb455 Mon Sep 17 00:00:00 2001 From: Zoltan Haindrich Date: Thu, 6 Mar 2025 14:12:47 +0000 Subject: [PATCH 26/50] same --- .../druid/msq/test/MSQTestWorkerClient.java | 26 +++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestWorkerClient.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestWorkerClient.java index 5e37133d77dd..9b0e4c79da95 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestWorkerClient.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestWorkerClient.java @@ -21,6 +21,7 @@ import com.google.common.util.concurrent.Futures; import com.google.common.util.concurrent.ListenableFuture; +import org.apache.druid.common.guava.FutureUtils; import org.apache.druid.frame.channel.ReadableByteChunksFrameChannel; import org.apache.druid.frame.key.ClusterByPartitions; import org.apache.druid.java.util.common.ISE; @@ -133,6 +134,30 @@ public ListenableFuture fetchChannelData( final ReadableByteChunksFrameChannel channel ) { + if(true) { + return FutureUtils.transform( + getWorkerFor(workerTaskId).readStageOutput(stageId, partitionNumber, offset), + inputStream -> { + try { + byte[] buffer = new byte[8 * 1024]; + boolean didRead = false; + int bytesRead; + while ((bytesRead = inputStream.read(buffer)) != -1) { + channel.addChunk(Arrays.copyOf(buffer, bytesRead)); + didRead = true; + } + inputStream.close(); + + return !didRead; + } + catch (Exception e) { + throw new ISE(e, "Error reading frame file channel"); + } + } + ); + } else + { + try (InputStream inputStream = getWorkerFor(workerTaskId).readStageOutput(stageId, partitionNumber, offset).get()) { byte[] buffer = new byte[8 * 1024]; @@ -149,6 +174,7 @@ public ListenableFuture fetchChannelData( catch (Exception e) { throw new ISE(e, "Error reading frame file channel"); } + } } @Override From e1c46939451c76755bb721bcce99729c1dca64d5 Mon Sep 17 00:00:00 2001 From: Zoltan Haindrich Date: Thu, 6 Mar 2025 14:13:27 +0000 Subject: [PATCH 27/50] cleanup/etc --- .../druid/msq/test/MSQTestWorkerClient.java | 62 ++++++------------- 1 file changed, 20 insertions(+), 42 deletions(-) diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestWorkerClient.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestWorkerClient.java index 9b0e4c79da95..2dcf1b0dfb56 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestWorkerClient.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestWorkerClient.java @@ -33,7 +33,6 @@ import org.apache.druid.msq.rpc.SketchEncoding; import org.apache.druid.msq.statistics.ClusterByStatisticsSnapshot; -import java.io.InputStream; import java.util.Arrays; import java.util.Map; import java.util.concurrent.atomic.AtomicBoolean; @@ -134,47 +133,26 @@ public ListenableFuture fetchChannelData( final ReadableByteChunksFrameChannel channel ) { - if(true) { - return FutureUtils.transform( - getWorkerFor(workerTaskId).readStageOutput(stageId, partitionNumber, offset), - inputStream -> { - try { - byte[] buffer = new byte[8 * 1024]; - boolean didRead = false; - int bytesRead; - while ((bytesRead = inputStream.read(buffer)) != -1) { - channel.addChunk(Arrays.copyOf(buffer, bytesRead)); - didRead = true; - } - inputStream.close(); - - return !didRead; - } - catch (Exception e) { - throw new ISE(e, "Error reading frame file channel"); - } - } - ); - } else - { - - try (InputStream inputStream = - getWorkerFor(workerTaskId).readStageOutput(stageId, partitionNumber, offset).get()) { - byte[] buffer = new byte[8 * 1024]; - boolean didRead = false; - int bytesRead; - while ((bytesRead = inputStream.read(buffer)) != -1) { - channel.addChunk(Arrays.copyOf(buffer, bytesRead)); - didRead = true; - } - inputStream.close(); - - return Futures.immediateFuture(!didRead); - } - catch (Exception e) { - throw new ISE(e, "Error reading frame file channel"); - } - } + return FutureUtils.transform( + getWorkerFor(workerTaskId).readStageOutput(stageId, partitionNumber, offset), + inputStream -> { + try { + byte[] buffer = new byte[8 * 1024]; + boolean didRead = false; + int bytesRead; + while ((bytesRead = inputStream.read(buffer)) != -1) { + channel.addChunk(Arrays.copyOf(buffer, bytesRead)); + didRead = true; + } + inputStream.close(); + + return !didRead; + } + catch (Exception e) { + throw new ISE(e, "Error reading frame file channel"); + } + } + ); } @Override From 96b2a9c13ef5986ec4a88fad9121e55c0331cc5a Mon Sep 17 00:00:00 2001 From: Zoltan Haindrich Date: Thu, 6 Mar 2025 14:19:53 +0000 Subject: [PATCH 28/50] run dart w/o old msqtestmodule --- .../java/org/apache/druid/msq/exec/TestMSQSqlModule.java | 4 ++-- .../msq/test/AbstractMSQComponentSupplierDelegate.java | 4 ++++ .../org/apache/druid/rpc/guice/ServiceClientModule.java | 2 +- .../apache/druid/sql/calcite/util/SqlTestFramework.java | 9 +++++++++ 4 files changed, 16 insertions(+), 3 deletions(-) diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/TestMSQSqlModule.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/TestMSQSqlModule.java index 0b48d2904dd3..ba52e599eaf0 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/TestMSQSqlModule.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/TestMSQSqlModule.java @@ -38,7 +38,7 @@ public class TestMSQSqlModule extends TestDruidModule { - @Provides +// @Provides @MultiStageQuery @LazySingleton public SqlStatementFactory makeMSQSqlStatementFactory( @@ -48,7 +48,7 @@ public SqlStatementFactory makeMSQSqlStatementFactory( return new SqlStatementFactory(toolbox.withEngine(sqlEngine)); } - @Provides +// @Provides @LazySingleton public MSQTaskSqlEngine createEngine( ObjectMapper queryJsonMapper, diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/AbstractMSQComponentSupplierDelegate.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/AbstractMSQComponentSupplierDelegate.java index 42264ac46ba3..3a44d4c59139 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/AbstractMSQComponentSupplierDelegate.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/AbstractMSQComponentSupplierDelegate.java @@ -27,8 +27,10 @@ import org.apache.druid.msq.exec.TestMSQSqlModule; import org.apache.druid.msq.guice.MSQExternalDataSourceModule; import org.apache.druid.msq.guice.MSQIndexingModule; +import org.apache.druid.msq.guice.MSQSqlModule; import org.apache.druid.msq.sql.MSQTaskSqlEngine; import org.apache.druid.msq.test.CalciteMSQTestsHelper.MSQTestModule; +import org.apache.druid.rpc.guice.ServiceClientModule; import org.apache.druid.server.QueryLifecycleFactory; import org.apache.druid.sql.calcite.run.SqlEngine; import org.apache.druid.sql.calcite.util.DruidModuleCollection; @@ -56,6 +58,8 @@ public DruidModule getCoreModule() new JoinableFactoryModule(), new MSQExternalDataSourceModule(), new MSQIndexingModule(), + new MSQSqlModule(), + new ServiceClientModule(), new TestMSQSqlModule() ); } diff --git a/server/src/main/java/org/apache/druid/rpc/guice/ServiceClientModule.java b/server/src/main/java/org/apache/druid/rpc/guice/ServiceClientModule.java index a0d9a5e3725c..7df5647bfe39 100644 --- a/server/src/main/java/org/apache/druid/rpc/guice/ServiceClientModule.java +++ b/server/src/main/java/org/apache/druid/rpc/guice/ServiceClientModule.java @@ -59,7 +59,7 @@ public void configure(Binder binder) // Nothing to do. } - @Provides +// @Provides @LazySingleton @EscalatedGlobal public ServiceClientFactory getServiceClientFactory(@EscalatedGlobal final HttpClient httpClient) diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/util/SqlTestFramework.java b/sql/src/test/java/org/apache/druid/sql/calcite/util/SqlTestFramework.java index cff7434bca27..c06da3f29bf0 100644 --- a/sql/src/test/java/org/apache/druid/sql/calcite/util/SqlTestFramework.java +++ b/sql/src/test/java/org/apache/druid/sql/calcite/util/SqlTestFramework.java @@ -54,6 +54,7 @@ import org.apache.druid.guice.SegmentWranglerModule; import org.apache.druid.guice.ServerModule; import org.apache.druid.guice.StartupInjectorBuilder; +import org.apache.druid.guice.annotations.EscalatedGlobal; import org.apache.druid.guice.annotations.Global; import org.apache.druid.guice.annotations.Merging; import org.apache.druid.guice.annotations.Self; @@ -1096,6 +1097,14 @@ private HttpClient extracted(ObjectMapper objectMapper) return new TestHttpClient(objectMapper); } + @Provides + @EscalatedGlobal + @LazySingleton + private HttpClient extracted1(ObjectMapper objectMapper) + { + return new TestHttpClient(objectMapper); + } + @Provides @Named("empty") @LazySingleton From cff97bf30a6d1160eed87cb5ebde7de46b1d7eea Mon Sep 17 00:00:00 2001 From: Zoltan Haindrich Date: Thu, 6 Mar 2025 14:28:45 +0000 Subject: [PATCH 29/50] add-trial --- .../msq/test/DartComponentSupplier2.java | 190 ++++++++++++++++++ .../druid/msq/test/SimpleDartTest2.java | 107 ++++++++++ 2 files changed, 297 insertions(+) create mode 100644 extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/DartComponentSupplier2.java create mode 100644 extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/SimpleDartTest2.java diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/DartComponentSupplier2.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/DartComponentSupplier2.java new file mode 100644 index 000000000000..6ab0f8766946 --- /dev/null +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/DartComponentSupplier2.java @@ -0,0 +1,190 @@ +/* + * 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.test; + +import com.fasterxml.jackson.databind.ObjectMapper; +import com.google.common.collect.ImmutableList; +import com.google.inject.Binder; +import com.google.inject.Injector; +import com.google.inject.Provides; +import com.google.inject.TypeLiteral; +import org.apache.druid.collections.NonBlockingPool; +import org.apache.druid.discovery.DruidNodeDiscoveryProvider; +import org.apache.druid.guice.LazySingleton; +import org.apache.druid.guice.annotations.EscalatedGlobal; +import org.apache.druid.guice.annotations.Merging; +import org.apache.druid.initialization.DruidModule; +import org.apache.druid.java.util.http.client.HttpClient; +import org.apache.druid.msq.dart.Dart; +import org.apache.druid.msq.dart.controller.DartControllerContextFactory; +import org.apache.druid.msq.dart.guice.DartControllerModule; +import org.apache.druid.msq.dart.guice.DartWorkerMemoryManagementModule; +import org.apache.druid.msq.dart.guice.DartWorkerModule; +import org.apache.druid.msq.exec.Worker; +import org.apache.druid.msq.exec.WorkerMemoryParameters; +import org.apache.druid.msq.sql.MSQTaskSqlEngine; +import org.apache.druid.query.TestBufferPool; +import org.apache.druid.rpc.ServiceClientFactory; +import org.apache.druid.rpc.guice.ServiceClientModule; +import org.apache.druid.server.QueryLifecycleFactory; +import org.apache.druid.sql.calcite.TempDirProducer; +import org.apache.druid.sql.calcite.run.SqlEngine; +import org.apache.druid.sql.calcite.util.DruidModuleCollection; +import org.apache.druid.sql.calcite.util.SqlTestFramework.StandardComponentSupplier; + +import java.nio.ByteBuffer; +import java.util.HashMap; +import java.util.Map; + +public class DartComponentSupplier2 extends AbstractMSQComponentSupplierDelegate +{ + public DartComponentSupplier2(TempDirProducer tempFolderProducer) + { + super(new StandardComponentSupplier(tempFolderProducer)); + } + + @Override + public DruidModule getOverrideModule() + { + return DruidModuleCollection.of(super.getOverrideModule(), + new LocalDartModule() + ); + } + + @Override + public DruidModule getCoreModule() + { + return DruidModuleCollection.of(super.getCoreModule(), + new DartControllerModule(), + new DartWorkerModule(), + new DartWorkerMemoryManagementModule(), + new L1() + ); + } + + @Override + public SqlEngine createEngine( + QueryLifecycleFactory qlf, + ObjectMapper queryJsonMapper, + Injector injector) + { + return injector.getInstance(MSQTaskSqlEngine.class); + } + + + + static class L1 implements DruidModule{ + + @Provides + @EscalatedGlobal + final ServiceClientFactory getServiceClientFactory(HttpClient ht) { + return ServiceClientModule.makeServiceClientFactory(ht); + + } + + @Provides + final DruidNodeDiscoveryProvider getDiscoveryProvider() { + return null; + } + + + @Override + public void configure(Binder binder) + { +// binder.bind(new TypeLiteral>(){}) +// .annotatedWith(Merging.class) +// .to(TestBufferPool.class); + + } + + } + static class LocalDartModule implements DruidModule{ + + @Override + public void configure(Binder binder) + { + binder.bind(new TypeLiteral>(){}) + .annotatedWith(Merging.class) + .to(TestBufferPool.class); + + if(true) { + binder.bind(DartControllerContextFactory.class) + .to(TestDartControllerContextFactoryImpl.class) + .in(LazySingleton.class); + }else { + binder.bind(DartControllerContextFactory.class) + .to(MSQTestControllerContext.class) + .in(LazySingleton.class); + + } + } + + @Provides + @LazySingleton + @Dart + Map workerMap() + { + return new HashMap(); + } + +// @Provides +// @LazySingleton +// public MSQTaskSqlEngine createEngine( +// ObjectMapper queryJsonMapper, +// MSQTestOverlordServiceClient2 indexingServiceClient) +// { +// return new DartSqlEngine(indexingServiceClient, queryJsonMapper, new SegmentGenerationTerminalStageSpecFactory()); +// } +// +// @Provides +// @LazySingleton +// public DartSqlEngine makeSqlEngine( +// DartControllerContextFactory controllerContextFactory, +// DartControllerRegistry controllerRegistry, +// DartControllerConfig controllerConfig +// ) +// { +// return new DartSqlEngine( +// controllerContextFactory, +// controllerRegistry, +// controllerConfig, +// Execs.multiThreaded(controllerConfig.getConcurrentQueries(), "dart-controller-%s") +// ); +// } + @Provides + @LazySingleton + private MSQTestOverlordServiceClient2 makeOverlordServiceClient( + ObjectMapper queryJsonMapper, + Injector injector, + WorkerMemoryParameters workerMemoryParameters) + { + final MSQTestOverlordServiceClient2 indexingServiceClient = new MSQTestOverlordServiceClient2( + queryJsonMapper, + injector, + new MSQTestTaskActionClient(queryJsonMapper, injector), + workerMemoryParameters, + ImmutableList.of() + ); + return indexingServiceClient; + } + + + } +} diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/SimpleDartTest2.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/SimpleDartTest2.java new file mode 100644 index 000000000000..14f5363b89c9 --- /dev/null +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/SimpleDartTest2.java @@ -0,0 +1,107 @@ +/* + * 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.test; + +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import org.apache.druid.sql.calcite.BaseCalciteQueryTest; +import org.apache.druid.sql.calcite.QueryTestBuilder; +import org.apache.druid.sql.calcite.SqlTestFrameworkConfig; +import org.junit.jupiter.api.Test; +import java.util.UUID; + +@SqlTestFrameworkConfig.ComponentSupplier(DartComponentSupplier2.class) +public class SimpleDartTest2 extends BaseCalciteQueryTest +{ + @Override + protected QueryTestBuilder testBuilder() + { + return new QueryTestBuilder(new CalciteTestConfig(true)) + .queryContext(ImmutableMap.builder().put("asd", UUID.randomUUID().toString()).build()) + .skipVectorize(true) + .verifyNativeQueries(new VerifyMSQSupportedNativeQueriesPredicate()); + } + + @Test + public void testSelect1() + { + testBuilder() + .sql("SELECT 1") + .expectedResults( + ImmutableList.of(new Object[] {1}) + ) + .run(); + } + + @Test + public void testSelectFromFoo() + { + testBuilder() + .sql("SELECT 2 from foo order by dim1") + .expectedResults( + ImmutableList.of( + new Object[] {2}, + new Object[] {2}, + new Object[] {2}, + new Object[] {2}, + new Object[] {2}, + new Object[] {2} + ) + ) + .run(); + } + + @Test + public void testSelectFromFoo11() + { + testBuilder() + .sql("SELECT dim1 from foo") + .expectedResults( + ImmutableList.of( + new Object[]{""}, + new Object[]{"10.1"}, + new Object[]{"2"}, + new Object[]{"1"}, + new Object[]{"def"}, + new Object[]{"abc"} + ) + ) + .run(); + } + @Test + public void testSelectFromFoo1() + { + String sql = "SELECT dim1, COUNT(*) FROM druid.foo " + + "WHERE dim1 NOT IN ('ghi', 'abc', 'def') AND dim1 IS NOT NULL " + + "GROUP BY dim1"; + testBuilder() + .sql(sql) + .expectedResults( + ImmutableList.of( + new Object[]{"", 1L}, + new Object[]{"1", 1L}, + new Object[]{"10.1", 1L}, + new Object[]{"2", 1L} + ) + ) + .run(); + } + +} From 11b212aff8be079965af55ca1c226a3cd2b35edf Mon Sep 17 00:00:00 2001 From: Zoltan Haindrich Date: Thu, 6 Mar 2025 14:29:19 +0000 Subject: [PATCH 30/50] retain running dart with msqtestsqlmodule for now --- .../druid/msq/exec/TestMSQSqlModule.java | 4 +- .../AbstractMSQComponentSupplierDelegate.java | 4 - .../msq/test/DartComponentSupplier2.java | 190 ------------------ .../druid/msq/test/SimpleDartTest2.java | 107 ---------- .../druid/rpc/guice/ServiceClientModule.java | 2 +- .../sql/calcite/util/SqlTestFramework.java | 9 - 6 files changed, 3 insertions(+), 313 deletions(-) delete mode 100644 extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/DartComponentSupplier2.java delete mode 100644 extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/SimpleDartTest2.java diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/TestMSQSqlModule.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/TestMSQSqlModule.java index ba52e599eaf0..0b48d2904dd3 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/TestMSQSqlModule.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/TestMSQSqlModule.java @@ -38,7 +38,7 @@ public class TestMSQSqlModule extends TestDruidModule { -// @Provides + @Provides @MultiStageQuery @LazySingleton public SqlStatementFactory makeMSQSqlStatementFactory( @@ -48,7 +48,7 @@ public SqlStatementFactory makeMSQSqlStatementFactory( return new SqlStatementFactory(toolbox.withEngine(sqlEngine)); } -// @Provides + @Provides @LazySingleton public MSQTaskSqlEngine createEngine( ObjectMapper queryJsonMapper, diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/AbstractMSQComponentSupplierDelegate.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/AbstractMSQComponentSupplierDelegate.java index 3a44d4c59139..42264ac46ba3 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/AbstractMSQComponentSupplierDelegate.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/AbstractMSQComponentSupplierDelegate.java @@ -27,10 +27,8 @@ import org.apache.druid.msq.exec.TestMSQSqlModule; import org.apache.druid.msq.guice.MSQExternalDataSourceModule; import org.apache.druid.msq.guice.MSQIndexingModule; -import org.apache.druid.msq.guice.MSQSqlModule; import org.apache.druid.msq.sql.MSQTaskSqlEngine; import org.apache.druid.msq.test.CalciteMSQTestsHelper.MSQTestModule; -import org.apache.druid.rpc.guice.ServiceClientModule; import org.apache.druid.server.QueryLifecycleFactory; import org.apache.druid.sql.calcite.run.SqlEngine; import org.apache.druid.sql.calcite.util.DruidModuleCollection; @@ -58,8 +56,6 @@ public DruidModule getCoreModule() new JoinableFactoryModule(), new MSQExternalDataSourceModule(), new MSQIndexingModule(), - new MSQSqlModule(), - new ServiceClientModule(), new TestMSQSqlModule() ); } diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/DartComponentSupplier2.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/DartComponentSupplier2.java deleted file mode 100644 index 6ab0f8766946..000000000000 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/DartComponentSupplier2.java +++ /dev/null @@ -1,190 +0,0 @@ -/* - * 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.test; - -import com.fasterxml.jackson.databind.ObjectMapper; -import com.google.common.collect.ImmutableList; -import com.google.inject.Binder; -import com.google.inject.Injector; -import com.google.inject.Provides; -import com.google.inject.TypeLiteral; -import org.apache.druid.collections.NonBlockingPool; -import org.apache.druid.discovery.DruidNodeDiscoveryProvider; -import org.apache.druid.guice.LazySingleton; -import org.apache.druid.guice.annotations.EscalatedGlobal; -import org.apache.druid.guice.annotations.Merging; -import org.apache.druid.initialization.DruidModule; -import org.apache.druid.java.util.http.client.HttpClient; -import org.apache.druid.msq.dart.Dart; -import org.apache.druid.msq.dart.controller.DartControllerContextFactory; -import org.apache.druid.msq.dart.guice.DartControllerModule; -import org.apache.druid.msq.dart.guice.DartWorkerMemoryManagementModule; -import org.apache.druid.msq.dart.guice.DartWorkerModule; -import org.apache.druid.msq.exec.Worker; -import org.apache.druid.msq.exec.WorkerMemoryParameters; -import org.apache.druid.msq.sql.MSQTaskSqlEngine; -import org.apache.druid.query.TestBufferPool; -import org.apache.druid.rpc.ServiceClientFactory; -import org.apache.druid.rpc.guice.ServiceClientModule; -import org.apache.druid.server.QueryLifecycleFactory; -import org.apache.druid.sql.calcite.TempDirProducer; -import org.apache.druid.sql.calcite.run.SqlEngine; -import org.apache.druid.sql.calcite.util.DruidModuleCollection; -import org.apache.druid.sql.calcite.util.SqlTestFramework.StandardComponentSupplier; - -import java.nio.ByteBuffer; -import java.util.HashMap; -import java.util.Map; - -public class DartComponentSupplier2 extends AbstractMSQComponentSupplierDelegate -{ - public DartComponentSupplier2(TempDirProducer tempFolderProducer) - { - super(new StandardComponentSupplier(tempFolderProducer)); - } - - @Override - public DruidModule getOverrideModule() - { - return DruidModuleCollection.of(super.getOverrideModule(), - new LocalDartModule() - ); - } - - @Override - public DruidModule getCoreModule() - { - return DruidModuleCollection.of(super.getCoreModule(), - new DartControllerModule(), - new DartWorkerModule(), - new DartWorkerMemoryManagementModule(), - new L1() - ); - } - - @Override - public SqlEngine createEngine( - QueryLifecycleFactory qlf, - ObjectMapper queryJsonMapper, - Injector injector) - { - return injector.getInstance(MSQTaskSqlEngine.class); - } - - - - static class L1 implements DruidModule{ - - @Provides - @EscalatedGlobal - final ServiceClientFactory getServiceClientFactory(HttpClient ht) { - return ServiceClientModule.makeServiceClientFactory(ht); - - } - - @Provides - final DruidNodeDiscoveryProvider getDiscoveryProvider() { - return null; - } - - - @Override - public void configure(Binder binder) - { -// binder.bind(new TypeLiteral>(){}) -// .annotatedWith(Merging.class) -// .to(TestBufferPool.class); - - } - - } - static class LocalDartModule implements DruidModule{ - - @Override - public void configure(Binder binder) - { - binder.bind(new TypeLiteral>(){}) - .annotatedWith(Merging.class) - .to(TestBufferPool.class); - - if(true) { - binder.bind(DartControllerContextFactory.class) - .to(TestDartControllerContextFactoryImpl.class) - .in(LazySingleton.class); - }else { - binder.bind(DartControllerContextFactory.class) - .to(MSQTestControllerContext.class) - .in(LazySingleton.class); - - } - } - - @Provides - @LazySingleton - @Dart - Map workerMap() - { - return new HashMap(); - } - -// @Provides -// @LazySingleton -// public MSQTaskSqlEngine createEngine( -// ObjectMapper queryJsonMapper, -// MSQTestOverlordServiceClient2 indexingServiceClient) -// { -// return new DartSqlEngine(indexingServiceClient, queryJsonMapper, new SegmentGenerationTerminalStageSpecFactory()); -// } -// -// @Provides -// @LazySingleton -// public DartSqlEngine makeSqlEngine( -// DartControllerContextFactory controllerContextFactory, -// DartControllerRegistry controllerRegistry, -// DartControllerConfig controllerConfig -// ) -// { -// return new DartSqlEngine( -// controllerContextFactory, -// controllerRegistry, -// controllerConfig, -// Execs.multiThreaded(controllerConfig.getConcurrentQueries(), "dart-controller-%s") -// ); -// } - @Provides - @LazySingleton - private MSQTestOverlordServiceClient2 makeOverlordServiceClient( - ObjectMapper queryJsonMapper, - Injector injector, - WorkerMemoryParameters workerMemoryParameters) - { - final MSQTestOverlordServiceClient2 indexingServiceClient = new MSQTestOverlordServiceClient2( - queryJsonMapper, - injector, - new MSQTestTaskActionClient(queryJsonMapper, injector), - workerMemoryParameters, - ImmutableList.of() - ); - return indexingServiceClient; - } - - - } -} diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/SimpleDartTest2.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/SimpleDartTest2.java deleted file mode 100644 index 14f5363b89c9..000000000000 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/SimpleDartTest2.java +++ /dev/null @@ -1,107 +0,0 @@ -/* - * 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.test; - -import com.google.common.collect.ImmutableList; -import com.google.common.collect.ImmutableMap; -import org.apache.druid.sql.calcite.BaseCalciteQueryTest; -import org.apache.druid.sql.calcite.QueryTestBuilder; -import org.apache.druid.sql.calcite.SqlTestFrameworkConfig; -import org.junit.jupiter.api.Test; -import java.util.UUID; - -@SqlTestFrameworkConfig.ComponentSupplier(DartComponentSupplier2.class) -public class SimpleDartTest2 extends BaseCalciteQueryTest -{ - @Override - protected QueryTestBuilder testBuilder() - { - return new QueryTestBuilder(new CalciteTestConfig(true)) - .queryContext(ImmutableMap.builder().put("asd", UUID.randomUUID().toString()).build()) - .skipVectorize(true) - .verifyNativeQueries(new VerifyMSQSupportedNativeQueriesPredicate()); - } - - @Test - public void testSelect1() - { - testBuilder() - .sql("SELECT 1") - .expectedResults( - ImmutableList.of(new Object[] {1}) - ) - .run(); - } - - @Test - public void testSelectFromFoo() - { - testBuilder() - .sql("SELECT 2 from foo order by dim1") - .expectedResults( - ImmutableList.of( - new Object[] {2}, - new Object[] {2}, - new Object[] {2}, - new Object[] {2}, - new Object[] {2}, - new Object[] {2} - ) - ) - .run(); - } - - @Test - public void testSelectFromFoo11() - { - testBuilder() - .sql("SELECT dim1 from foo") - .expectedResults( - ImmutableList.of( - new Object[]{""}, - new Object[]{"10.1"}, - new Object[]{"2"}, - new Object[]{"1"}, - new Object[]{"def"}, - new Object[]{"abc"} - ) - ) - .run(); - } - @Test - public void testSelectFromFoo1() - { - String sql = "SELECT dim1, COUNT(*) FROM druid.foo " - + "WHERE dim1 NOT IN ('ghi', 'abc', 'def') AND dim1 IS NOT NULL " - + "GROUP BY dim1"; - testBuilder() - .sql(sql) - .expectedResults( - ImmutableList.of( - new Object[]{"", 1L}, - new Object[]{"1", 1L}, - new Object[]{"10.1", 1L}, - new Object[]{"2", 1L} - ) - ) - .run(); - } - -} diff --git a/server/src/main/java/org/apache/druid/rpc/guice/ServiceClientModule.java b/server/src/main/java/org/apache/druid/rpc/guice/ServiceClientModule.java index 7df5647bfe39..a0d9a5e3725c 100644 --- a/server/src/main/java/org/apache/druid/rpc/guice/ServiceClientModule.java +++ b/server/src/main/java/org/apache/druid/rpc/guice/ServiceClientModule.java @@ -59,7 +59,7 @@ public void configure(Binder binder) // Nothing to do. } -// @Provides + @Provides @LazySingleton @EscalatedGlobal public ServiceClientFactory getServiceClientFactory(@EscalatedGlobal final HttpClient httpClient) diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/util/SqlTestFramework.java b/sql/src/test/java/org/apache/druid/sql/calcite/util/SqlTestFramework.java index c06da3f29bf0..cff7434bca27 100644 --- a/sql/src/test/java/org/apache/druid/sql/calcite/util/SqlTestFramework.java +++ b/sql/src/test/java/org/apache/druid/sql/calcite/util/SqlTestFramework.java @@ -54,7 +54,6 @@ import org.apache.druid.guice.SegmentWranglerModule; import org.apache.druid.guice.ServerModule; import org.apache.druid.guice.StartupInjectorBuilder; -import org.apache.druid.guice.annotations.EscalatedGlobal; import org.apache.druid.guice.annotations.Global; import org.apache.druid.guice.annotations.Merging; import org.apache.druid.guice.annotations.Self; @@ -1097,14 +1096,6 @@ private HttpClient extracted(ObjectMapper objectMapper) return new TestHttpClient(objectMapper); } - @Provides - @EscalatedGlobal - @LazySingleton - private HttpClient extracted1(ObjectMapper objectMapper) - { - return new TestHttpClient(objectMapper); - } - @Provides @Named("empty") @LazySingleton From 83e36a3985c7b737f2d1bdc35e49eaf2da186bb4 Mon Sep 17 00:00:00 2001 From: Zoltan Haindrich Date: Fri, 7 Mar 2025 08:17:18 +0000 Subject: [PATCH 31/50] middle --- .../controller/DartControllerContext.java | 7 -- .../dart/controller/sql/DartQueryMaker.java | 1 + .../TestDartControllerContextFactoryImpl.java | 79 +++++++++++-------- 3 files changed, 49 insertions(+), 38 deletions(-) diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartControllerContext.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartControllerContext.java index 0947bcce180c..c2257d160ed3 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartControllerContext.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartControllerContext.java @@ -89,8 +89,6 @@ public class DartControllerContext implements ControllerContext private final ServiceMetricEvent.Builder metricBuilder; private final ServiceEmitter emitter; - public Controller theController; - public DartControllerContext( final Injector injector, final ObjectMapper jsonMapper, @@ -217,11 +215,6 @@ public DartWorkerClient newWorkerClient() public void registerController(Controller controller, Closer closer) { closer.register(workerClient); - // FIXME: - // workerClient.setController(); - // controller - theController = controller; - } @Override diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartQueryMaker.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartQueryMaker.java index dee475ff5491..f6a93e3605f8 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartQueryMaker.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartQueryMaker.java @@ -152,6 +152,7 @@ public QueryResponse runQuery(DruidQuery druidQuery) ), controllerContext ); + controllerContext.registerController(controller, null); final ControllerHolder controllerHolder = new ControllerHolder( controller, diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/TestDartControllerContextFactoryImpl.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/TestDartControllerContextFactoryImpl.java index fce5a6a24042..55f55de430f7 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/TestDartControllerContextFactoryImpl.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/TestDartControllerContextFactoryImpl.java @@ -28,6 +28,7 @@ import org.apache.druid.guice.annotations.Json; import org.apache.druid.guice.annotations.Self; import org.apache.druid.guice.annotations.Smile; +import org.apache.druid.java.util.common.io.Closer; import org.apache.druid.java.util.emitter.service.ServiceEmitter; import org.apache.druid.msq.dart.Dart; import org.apache.druid.msq.dart.controller.DartControllerContext; @@ -44,9 +45,11 @@ import org.apache.druid.rpc.ServiceClientFactory; import org.apache.druid.server.DruidNode; +import java.util.HashMap; import java.util.Map; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; +import java.util.function.Function; public class TestDartControllerContextFactoryImpl extends DartControllerContextFactoryImpl { @@ -71,52 +74,39 @@ public TestDartControllerContextFactoryImpl( @Override public ControllerContext newContext(String queryId) { - DartControllerContext ctx; - DartTestWorkerClient wc; - ctx = new DartControllerContext( - injector, - jsonMapper, - selfNode, - wc = (DartTestWorkerClient) makeWorkerClient(queryId), - memoryIntrospector, - serverView, - emitter - ); - wc.controllerCtx = ctx; - return ctx; + return new DartTestControllerContext(queryId); } - @Override + protected DartWorkerClient makeWorkerClient(String queryId) { - return new DartTestWorkerClient( - queryId, serviceClientFactory, smileMapper, selfNode.getHostAndPortToUse(), workerMap - ); + return new DartTestWorkerClient(null); } - public class DartTestWorkerClient extends MSQTestWorkerClient implements DartWorkerClient - { - private final ExecutorService EXECUTOR = Executors.newCachedThreadPool(); - public DartControllerContext controllerCtx; + public class DartTestControllerContext extends DartControllerContext { + + public ExecutorService executor = Executors.newCachedThreadPool(); + public Controller controller; + public Map inMemoryWorkers = new HashMap<>(); - public DartTestWorkerClient( - String queryId, - ServiceClientFactory clientFactory, - ObjectMapper smileMapper, - String controllerHost, Map workerMap) + public DartTestControllerContext(String queryId) { - super(workerMap); + super(injector, jsonMapper, selfNode, makeWorkerClient(queryId), memoryIntrospector, serverView, emitter); } @Override + public void registerController(Controller controller, Closer closer) + { + super.registerController(controller, closer); + this.controller= controller; + } + protected Worker newWorker(String workerId) { - String queryId = workerId; - Controller controller = controllerCtx.theController; Worker worker = new WorkerImpl( null, new MSQTestWorkerContext( - queryId, + workerId, inMemoryWorkers, controller, jsonMapper, @@ -126,7 +116,7 @@ protected Worker newWorker(String workerId) ) ); - EXECUTOR.submit(() -> { + executor.submit(() -> { try { worker.run(); } @@ -137,6 +127,33 @@ protected Worker newWorker(String workerId) return worker; } + } + + + public class DartTestWorkerClient extends MSQTestWorkerClient implements DartWorkerClient + { + public DartControllerContext controllerCtx; + + public DartTestWorkerClient(Function workerFactory) + { + super(workerMap); + } + + @Override + protected Worker newWorker(String workerId) + { + return makeWorker(workerId); + } + + private Worker makeWorker(String workerId) + { + if(true) + { + throw new RuntimeException("FIXME: Unimplemented!"); + } + return null; + + } @Override public ListenableFuture postWorkOrder(String workerTaskId, WorkOrder workOrder) From 9a9eef4f0adb46d2b0e2e33553bfdadb76ce7e71 Mon Sep 17 00:00:00 2001 From: Zoltan Haindrich Date: Fri, 7 Mar 2025 11:53:25 +0000 Subject: [PATCH 32/50] crap --- .../dart/controller/sql/DartQueryMaker.java | 1 - .../TestDartControllerContextFactoryImpl.java | 30 ++++++++++--------- 2 files changed, 16 insertions(+), 15 deletions(-) diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartQueryMaker.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartQueryMaker.java index f6a93e3605f8..dee475ff5491 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartQueryMaker.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartQueryMaker.java @@ -152,7 +152,6 @@ public QueryResponse runQuery(DruidQuery druidQuery) ), controllerContext ); - controllerContext.registerController(controller, null); final ControllerHolder controllerHolder = new ControllerHolder( controller, diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/TestDartControllerContextFactoryImpl.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/TestDartControllerContextFactoryImpl.java index 55f55de430f7..d28263c71444 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/TestDartControllerContextFactoryImpl.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/TestDartControllerContextFactoryImpl.java @@ -45,15 +45,19 @@ import org.apache.druid.rpc.ServiceClientFactory; import org.apache.druid.server.DruidNode; +import javax.management.RuntimeErrorException; + import java.util.HashMap; import java.util.Map; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; +import java.util.concurrent.atomic.AtomicReference; import java.util.function.Function; public class TestDartControllerContextFactoryImpl extends DartControllerContextFactoryImpl { private Map workerMap; + private DartTestControllerContext controllerContext; @Inject public TestDartControllerContextFactoryImpl( @@ -74,31 +78,31 @@ public TestDartControllerContextFactoryImpl( @Override public ControllerContext newContext(String queryId) { - return new DartTestControllerContext(queryId); + if (controllerContext != null) { + throw new RuntimeException(); + } + controllerContext = new DartTestControllerContext(queryId); + return controllerContext; } + public DartTestWorkerClient dartTestWorkerClient = new DartTestWorkerClient(); - protected DartWorkerClient makeWorkerClient(String queryId) + public class DartTestControllerContext extends DartControllerContext { - return new DartTestWorkerClient(null); - } - - public class DartTestControllerContext extends DartControllerContext { public ExecutorService executor = Executors.newCachedThreadPool(); - public Controller controller; public Map inMemoryWorkers = new HashMap<>(); public DartTestControllerContext(String queryId) { - super(injector, jsonMapper, selfNode, makeWorkerClient(queryId), memoryIntrospector, serverView, emitter); + super(injector, jsonMapper, selfNode, dartTestWorkerClient, memoryIntrospector, serverView, emitter); } @Override public void registerController(Controller controller, Closer closer) { super.registerController(controller, closer); - this.controller= controller; + dartTestWorkerClient.controller = controller; } protected Worker newWorker(String workerId) @@ -129,14 +133,13 @@ protected Worker newWorker(String workerId) } } - public class DartTestWorkerClient extends MSQTestWorkerClient implements DartWorkerClient { public DartControllerContext controllerCtx; - public DartTestWorkerClient(Function workerFactory) + public DartTestWorkerClient() { - super(workerMap); + super(new HashMap<>()); } @Override @@ -147,8 +150,7 @@ protected Worker newWorker(String workerId) private Worker makeWorker(String workerId) { - if(true) - { + if (true) { throw new RuntimeException("FIXME: Unimplemented!"); } return null; From 5e51e1ac089f63d24ddee729574c1a44b7b67ae2 Mon Sep 17 00:00:00 2001 From: Zoltan Haindrich Date: Fri, 7 Mar 2025 11:53:27 +0000 Subject: [PATCH 33/50] Revert "crap" This reverts commit 9a9eef4f0adb46d2b0e2e33553bfdadb76ce7e71. --- .../controller/DartControllerContext.java | 7 ++ .../TestDartControllerContextFactoryImpl.java | 85 +++++++------------ 2 files changed, 40 insertions(+), 52 deletions(-) diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartControllerContext.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartControllerContext.java index c2257d160ed3..0947bcce180c 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartControllerContext.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartControllerContext.java @@ -89,6 +89,8 @@ public class DartControllerContext implements ControllerContext private final ServiceMetricEvent.Builder metricBuilder; private final ServiceEmitter emitter; + public Controller theController; + public DartControllerContext( final Injector injector, final ObjectMapper jsonMapper, @@ -215,6 +217,11 @@ public DartWorkerClient newWorkerClient() public void registerController(Controller controller, Closer closer) { closer.register(workerClient); + // FIXME: + // workerClient.setController(); + // controller + theController = controller; + } @Override diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/TestDartControllerContextFactoryImpl.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/TestDartControllerContextFactoryImpl.java index d28263c71444..fce5a6a24042 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/TestDartControllerContextFactoryImpl.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/TestDartControllerContextFactoryImpl.java @@ -28,7 +28,6 @@ import org.apache.druid.guice.annotations.Json; import org.apache.druid.guice.annotations.Self; import org.apache.druid.guice.annotations.Smile; -import org.apache.druid.java.util.common.io.Closer; import org.apache.druid.java.util.emitter.service.ServiceEmitter; import org.apache.druid.msq.dart.Dart; import org.apache.druid.msq.dart.controller.DartControllerContext; @@ -45,19 +44,13 @@ import org.apache.druid.rpc.ServiceClientFactory; import org.apache.druid.server.DruidNode; -import javax.management.RuntimeErrorException; - -import java.util.HashMap; import java.util.Map; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; -import java.util.concurrent.atomic.AtomicReference; -import java.util.function.Function; public class TestDartControllerContextFactoryImpl extends DartControllerContextFactoryImpl { private Map workerMap; - private DartTestControllerContext controllerContext; @Inject public TestDartControllerContextFactoryImpl( @@ -78,39 +71,52 @@ public TestDartControllerContextFactoryImpl( @Override public ControllerContext newContext(String queryId) { - if (controllerContext != null) { - throw new RuntimeException(); - } - controllerContext = new DartTestControllerContext(queryId); - return controllerContext; + DartControllerContext ctx; + DartTestWorkerClient wc; + ctx = new DartControllerContext( + injector, + jsonMapper, + selfNode, + wc = (DartTestWorkerClient) makeWorkerClient(queryId), + memoryIntrospector, + serverView, + emitter + ); + wc.controllerCtx = ctx; + return ctx; } - public DartTestWorkerClient dartTestWorkerClient = new DartTestWorkerClient(); - - public class DartTestControllerContext extends DartControllerContext + @Override + protected DartWorkerClient makeWorkerClient(String queryId) { + return new DartTestWorkerClient( + queryId, serviceClientFactory, smileMapper, selfNode.getHostAndPortToUse(), workerMap + ); + } - public ExecutorService executor = Executors.newCachedThreadPool(); - public Map inMemoryWorkers = new HashMap<>(); + public class DartTestWorkerClient extends MSQTestWorkerClient implements DartWorkerClient + { + private final ExecutorService EXECUTOR = Executors.newCachedThreadPool(); + public DartControllerContext controllerCtx; - public DartTestControllerContext(String queryId) + public DartTestWorkerClient( + String queryId, + ServiceClientFactory clientFactory, + ObjectMapper smileMapper, + String controllerHost, Map workerMap) { - super(injector, jsonMapper, selfNode, dartTestWorkerClient, memoryIntrospector, serverView, emitter); + super(workerMap); } @Override - public void registerController(Controller controller, Closer closer) - { - super.registerController(controller, closer); - dartTestWorkerClient.controller = controller; - } - protected Worker newWorker(String workerId) { + String queryId = workerId; + Controller controller = controllerCtx.theController; Worker worker = new WorkerImpl( null, new MSQTestWorkerContext( - workerId, + queryId, inMemoryWorkers, controller, jsonMapper, @@ -120,7 +126,7 @@ protected Worker newWorker(String workerId) ) ); - executor.submit(() -> { + EXECUTOR.submit(() -> { try { worker.run(); } @@ -131,31 +137,6 @@ protected Worker newWorker(String workerId) return worker; } - } - - public class DartTestWorkerClient extends MSQTestWorkerClient implements DartWorkerClient - { - public DartControllerContext controllerCtx; - - public DartTestWorkerClient() - { - super(new HashMap<>()); - } - - @Override - protected Worker newWorker(String workerId) - { - return makeWorker(workerId); - } - - private Worker makeWorker(String workerId) - { - if (true) { - throw new RuntimeException("FIXME: Unimplemented!"); - } - return null; - - } @Override public ListenableFuture postWorkOrder(String workerTaskId, WorkOrder workOrder) From 5115dd78260491cdba8af4b111db89189ca39b4c Mon Sep 17 00:00:00 2001 From: Zoltan Haindrich Date: Fri, 7 Mar 2025 12:04:52 +0000 Subject: [PATCH 34/50] fixme --- .../dart/controller/DartControllerContext.java | 7 ------- .../msq/dart/guice/DartControllerModule.java | 3 +-- .../druid/msq/test/DartComponentSupplier.java | 9 +++++++++ .../TestDartControllerContextFactoryImpl.java | 15 ++++++++++----- 4 files changed, 20 insertions(+), 14 deletions(-) diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartControllerContext.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartControllerContext.java index 0947bcce180c..c2257d160ed3 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartControllerContext.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartControllerContext.java @@ -89,8 +89,6 @@ public class DartControllerContext implements ControllerContext private final ServiceMetricEvent.Builder metricBuilder; private final ServiceEmitter emitter; - public Controller theController; - public DartControllerContext( final Injector injector, final ObjectMapper jsonMapper, @@ -217,11 +215,6 @@ public DartWorkerClient newWorkerClient() public void registerController(Controller controller, Closer closer) { closer.register(workerClient); - // FIXME: - // workerClient.setController(); - // controller - theController = controller; - } @Override diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/guice/DartControllerModule.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/guice/DartControllerModule.java index 1a1f387a0eae..8a4b73bc9b0f 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/guice/DartControllerModule.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/guice/DartControllerModule.java @@ -65,8 +65,7 @@ public class DartControllerModule implements DruidModule @Override public void configure(Binder binder) { - // FIXME - if (true || DartModules.isDartEnabled(properties)) { + if (DartModules.isDartEnabled(properties)) { binder.install(new ActualModule()); } } diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/DartComponentSupplier.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/DartComponentSupplier.java index e204c2cd87c3..f91f7b2344bf 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/DartComponentSupplier.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/DartComponentSupplier.java @@ -36,6 +36,7 @@ import org.apache.druid.msq.dart.controller.DartControllerContextFactory; import org.apache.druid.msq.dart.controller.sql.DartSqlEngine; import org.apache.druid.msq.dart.guice.DartControllerModule; +import org.apache.druid.msq.dart.guice.DartModules; import org.apache.druid.msq.dart.guice.DartWorkerMemoryManagementModule; import org.apache.druid.msq.dart.guice.DartWorkerModule; import org.apache.druid.msq.exec.Worker; @@ -52,6 +53,7 @@ import java.nio.ByteBuffer; import java.util.HashMap; import java.util.Map; +import java.util.Properties; public class DartComponentSupplier extends AbstractMSQComponentSupplierDelegate { @@ -60,6 +62,13 @@ public DartComponentSupplier(TempDirProducer tempFolderProducer) super(new StandardComponentSupplier(tempFolderProducer)); } + @Override + public void gatherProperties(Properties properties) + { + super.gatherProperties(properties); + properties.put(DartModules.DART_ENABLED_PROPERTY, "true"); + } + @Override public DruidModule getOverrideModule() { diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/TestDartControllerContextFactoryImpl.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/TestDartControllerContextFactoryImpl.java index fce5a6a24042..45843f42272f 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/TestDartControllerContextFactoryImpl.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/TestDartControllerContextFactoryImpl.java @@ -51,6 +51,7 @@ public class TestDartControllerContextFactoryImpl extends DartControllerContextFactoryImpl { private Map workerMap; + public Controller theController1; @Inject public TestDartControllerContextFactoryImpl( @@ -81,11 +82,17 @@ public ControllerContext newContext(String queryId) memoryIntrospector, serverView, emitter - ); - wc.controllerCtx = ctx; + ) { + public void registerController(Controller controller, org.apache.druid.java.util.common.io.Closer closer) { + super.registerController(controller, closer); + theController1=controller; + + }; + }; return ctx; } + @Override protected DartWorkerClient makeWorkerClient(String queryId) { @@ -97,7 +104,6 @@ protected DartWorkerClient makeWorkerClient(String queryId) public class DartTestWorkerClient extends MSQTestWorkerClient implements DartWorkerClient { private final ExecutorService EXECUTOR = Executors.newCachedThreadPool(); - public DartControllerContext controllerCtx; public DartTestWorkerClient( String queryId, @@ -112,13 +118,12 @@ public DartTestWorkerClient( protected Worker newWorker(String workerId) { String queryId = workerId; - Controller controller = controllerCtx.theController; Worker worker = new WorkerImpl( null, new MSQTestWorkerContext( queryId, inMemoryWorkers, - controller, + theController1, jsonMapper, injector, MSQTestBase.makeTestWorkerMemoryParameters(), From a569e9551f0a0e0f455f81c77f1ae250dc62a112 Mon Sep 17 00:00:00 2001 From: Zoltan Haindrich Date: Fri, 7 Mar 2025 12:08:42 +0000 Subject: [PATCH 35/50] cleanup --- .../druid/msq/test/MSQTestWorkerClient.java | 3 +-- .../druid/quidem/DruidAvaticaTestDriver.java | 16 ---------------- .../druid/sql/calcite/util/SqlTestFramework.java | 14 +------------- 3 files changed, 2 insertions(+), 31 deletions(-) diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestWorkerClient.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestWorkerClient.java index 2dcf1b0dfb56..90ef5d4da78c 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestWorkerClient.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestWorkerClient.java @@ -51,7 +51,6 @@ public MSQTestWorkerClient(Map inMemoryWorkers) public ListenableFuture postWorkOrder(String workerTaskId, WorkOrder workOrder) { getWorkerFor(workerTaskId).postWorkOrder(workOrder); - return Futures.immediateFuture(null); } @@ -62,7 +61,7 @@ protected Worker getWorkerFor(String workerTaskId) protected Worker newWorker(String workerId) { - throw new RuntimeException("FIXME: Unimplemented!"); + throw new RuntimeException("Not implemented!"); } @Override diff --git a/sql/src/test/java/org/apache/druid/quidem/DruidAvaticaTestDriver.java b/sql/src/test/java/org/apache/druid/quidem/DruidAvaticaTestDriver.java index 16f8c87671d2..21e0cc7bd71a 100644 --- a/sql/src/test/java/org/apache/druid/quidem/DruidAvaticaTestDriver.java +++ b/sql/src/test/java/org/apache/druid/quidem/DruidAvaticaTestDriver.java @@ -24,7 +24,6 @@ import com.google.inject.Injector; import com.google.inject.Provides; import com.google.inject.name.Named; -import com.google.inject.name.Names; import org.apache.calcite.avatica.server.AbstractAvaticaHandler; import org.apache.druid.guice.LazySingleton; import org.apache.druid.initialization.DruidModule; @@ -183,7 +182,6 @@ protected AbstractAvaticaHandler getAvaticaHandler(final DruidMeta druidMeta) { return new DruidAvaticaJsonHandler( druidMeta, - //FIXME new DruidNode("dummy", "dummy", false, 1, null, true, false), new AvaticaMonitor() ); @@ -200,20 +198,6 @@ public AvaticaBasedTestConnectionSupplier(QueryComponentSupplier delegate) this.connectionModule = new AvaticaBasedConnectionModule(); } - @Override - public DruidModule getCoreModule() - { - return DruidModuleCollection.of( - super.getCoreModule(), - binder -> { - // FIXME - binder.bindConstant().annotatedWith(Names.named("serviceName")).to("test"); - binder.bindConstant().annotatedWith(Names.named("servicePort")).to(0); - binder.bindConstant().annotatedWith(Names.named("tlsServicePort")).to(-1); - } - ); - } - @Override public DruidModule getOverrideModule() { diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/util/SqlTestFramework.java b/sql/src/test/java/org/apache/druid/sql/calcite/util/SqlTestFramework.java index cff7434bca27..4f8a002cf445 100644 --- a/sql/src/test/java/org/apache/druid/sql/calcite/util/SqlTestFramework.java +++ b/sql/src/test/java/org/apache/druid/sql/calcite/util/SqlTestFramework.java @@ -1077,21 +1077,9 @@ public TestSegmentsBroker makeTimelines() return new TestSegmentsBroker(); } - @LazySingleton - public TimelineServerView makeTimelines1(ObjectMapper objectMapper, QueryRunnerFactoryConglomerate conglomerate, - HttpClient httpClient) - { - //FIXME connect TSB with this -///. return new SimpleServerView(conglomerate, objectMapper, httpClient); - - - SimpleServerView simpleServerView = new SimpleServerView(conglomerate, objectMapper, httpClient); - return simpleServerView; - } - @Provides @LazySingleton - private HttpClient extracted(ObjectMapper objectMapper) + private HttpClient makeHttpClient(ObjectMapper objectMapper) { return new TestHttpClient(objectMapper); } From 73872ba5010bf9ef45987075aad30969d03f93e7 Mon Sep 17 00:00:00 2001 From: Zoltan Haindrich Date: Fri, 7 Mar 2025 12:14:11 +0000 Subject: [PATCH 36/50] halfway --- .../dart.iq | 34 +++++++++++++------ 1 file changed, 24 insertions(+), 10 deletions(-) diff --git a/extensions-core/multi-stage-query/src/test/quidem/org.apache.druid.msq.quidem.MSQQuidemTest/dart.iq b/extensions-core/multi-stage-query/src/test/quidem/org.apache.druid.msq.quidem.MSQQuidemTest/dart.iq index 6a3bdc1754e8..a55baf16d6c4 100644 --- a/extensions-core/multi-stage-query/src/test/quidem/org.apache.druid.msq.quidem.MSQQuidemTest/dart.iq +++ b/extensions-core/multi-stage-query/src/test/quidem/org.apache.druid.msq.quidem.MSQQuidemTest/dart.iq @@ -1,17 +1,31 @@ -!set plannerStrategy DECOUPLED !use druidtest://?componentSupplier=DartComponentSupplier !set outputformat mysql - - select * from foo; -+---------------+ -| actualNumRows | -+---------------+ -| 11631 | -+---------------+ -(1 row) ++-------------------------+------+------+-----------+-----+-----+-----+--------------------+ +| __time | dim1 | dim2 | dim3 | cnt | m1 | m2 | unique_dim1 | ++-------------------------+------+------+-----------+-----+-----+-----+--------------------+ +| 2000-01-01 00:00:00.000 | | a | ["a","b"] | 1 | 1.0 | 1.0 | "AQAAAEAAAA==" | +| 2000-01-02 00:00:00.000 | 10.1 | | ["b","c"] | 1 | 2.0 | 2.0 | "AQAAAQAAAAHNBA==" | +| 2000-01-03 00:00:00.000 | 2 | | d | 1 | 3.0 | 3.0 | "AQAAAQAAAAOzAg==" | +| 2001-01-01 00:00:00.000 | 1 | a | | 1 | 4.0 | 4.0 | "AQAAAQAAAAFREA==" | +| 2001-01-02 00:00:00.000 | def | abc | | 1 | 5.0 | 5.0 | "AQAAAQAAAACyEA==" | +| 2001-01-03 00:00:00.000 | abc | | | 1 | 6.0 | 6.0 | "AQAAAQAAAAEkAQ==" | ++-------------------------+------+------+-----------+-----+-----+-----+--------------------+ +(6 rows) !ok -select * from foo group by dim1; +select 2 from foo group by dim1; ++--------+ +| EXPR$0 | ++--------+ +| 2 | +| 2 | +| 2 | +| 2 | +| 2 | +| 2 | ++--------+ +(6 rows) + !ok From 808a4c8e52765b8a519e19b73f273e2db446a34b Mon Sep 17 00:00:00 2001 From: Zoltan Haindrich Date: Fri, 7 Mar 2025 12:32:43 +0000 Subject: [PATCH 37/50] cleanup/etc --- .../druid/msq/exec/TestMSQSqlModule.java | 4 +- .../apache/druid/msq/test/SimpleDartTest.java | 46 ++++++++++++------- .../dart.iq | 37 ++++++--------- 3 files changed, 45 insertions(+), 42 deletions(-) diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/TestMSQSqlModule.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/TestMSQSqlModule.java index 0b48d2904dd3..9ca1e69fbebe 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/TestMSQSqlModule.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/TestMSQSqlModule.java @@ -33,8 +33,8 @@ import org.apache.druid.msq.test.MSQTestTaskActionClient; import org.apache.druid.sql.SqlStatementFactory; import org.apache.druid.sql.SqlToolbox; +import org.apache.druid.sql.avatica.DartDruidMeta; import org.apache.druid.sql.avatica.DruidMeta; -import org.apache.druid.sql.avatica.MSQDruidMeta; public class TestMSQSqlModule extends TestDruidModule { @@ -83,7 +83,7 @@ private WorkerMemoryParameters makeWorkerMemoryParameters() @Provides @LazySingleton - public DruidMeta createMeta(MSQDruidMeta druidMeta) + public DruidMeta createMeta(DartDruidMeta druidMeta) { return druidMeta; } diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/SimpleDartTest.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/SimpleDartTest.java index 4e81e2104e5e..66a8628fcee7 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/SimpleDartTest.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/SimpleDartTest.java @@ -44,9 +44,7 @@ public void testSelect1() { testBuilder() .sql("SELECT 1") - .expectedResults( - ImmutableList.of(new Object[] {1}) - ) + .expectedResults(ImmutableList.of(new Object[] {1})) .run(); } @@ -69,24 +67,41 @@ public void testSelectFromFoo() } @Test - public void testSelectFromFoo11() + public void testSelectDim1FromFoo11() { testBuilder() .sql("SELECT dim1 from foo") .expectedResults( ImmutableList.of( - new Object[]{""}, - new Object[]{"10.1"}, - new Object[]{"2"}, - new Object[]{"1"}, - new Object[]{"def"}, - new Object[]{"abc"} + new Object[] {""}, + new Object[] {"10.1"}, + new Object[] {"2"}, + new Object[] {"1"}, + new Object[] {"def"}, + new Object[] {"abc"} + ) + ) + .run(); + } + + @Test + public void testGby() + { + testBuilder() + .sql("SELECT 3 from foo group by dim2") + .expectedResults( + ImmutableList.of( + new Object[] {3}, + new Object[] {3}, + new Object[] {3}, + new Object[] {3} ) ) .run(); } + @Test - public void testSelectFromFoo1() + public void testComplexFromFoo() { String sql = "SELECT dim1, COUNT(*) FROM druid.foo " + "WHERE dim1 NOT IN ('ghi', 'abc', 'def') AND dim1 IS NOT NULL " @@ -95,13 +110,12 @@ public void testSelectFromFoo1() .sql(sql) .expectedResults( ImmutableList.of( - new Object[]{"", 1L}, - new Object[]{"1", 1L}, - new Object[]{"10.1", 1L}, - new Object[]{"2", 1L} + new Object[] {"", 1L}, + new Object[] {"1", 1L}, + new Object[] {"10.1", 1L}, + new Object[] {"2", 1L} ) ) .run(); } - } diff --git a/extensions-core/multi-stage-query/src/test/quidem/org.apache.druid.msq.quidem.MSQQuidemTest/dart.iq b/extensions-core/multi-stage-query/src/test/quidem/org.apache.druid.msq.quidem.MSQQuidemTest/dart.iq index a55baf16d6c4..646f601a230f 100644 --- a/extensions-core/multi-stage-query/src/test/quidem/org.apache.druid.msq.quidem.MSQQuidemTest/dart.iq +++ b/extensions-core/multi-stage-query/src/test/quidem/org.apache.druid.msq.quidem.MSQQuidemTest/dart.iq @@ -1,31 +1,20 @@ !use druidtest://?componentSupplier=DartComponentSupplier !set outputformat mysql -select * from foo; -+-------------------------+------+------+-----------+-----+-----+-----+--------------------+ -| __time | dim1 | dim2 | dim3 | cnt | m1 | m2 | unique_dim1 | -+-------------------------+------+------+-----------+-----+-----+-----+--------------------+ -| 2000-01-01 00:00:00.000 | | a | ["a","b"] | 1 | 1.0 | 1.0 | "AQAAAEAAAA==" | -| 2000-01-02 00:00:00.000 | 10.1 | | ["b","c"] | 1 | 2.0 | 2.0 | "AQAAAQAAAAHNBA==" | -| 2000-01-03 00:00:00.000 | 2 | | d | 1 | 3.0 | 3.0 | "AQAAAQAAAAOzAg==" | -| 2001-01-01 00:00:00.000 | 1 | a | | 1 | 4.0 | 4.0 | "AQAAAQAAAAFREA==" | -| 2001-01-02 00:00:00.000 | def | abc | | 1 | 5.0 | 5.0 | "AQAAAQAAAACyEA==" | -| 2001-01-03 00:00:00.000 | abc | | | 1 | 6.0 | 6.0 | "AQAAAQAAAAEkAQ==" | -+-------------------------+------+------+-----------+-----+-----+-----+--------------------+ + +--select dim1 from foo ; +select count(1) from foo ; ++------+ +| dim1 | ++------+ +| | +| 1 | +| 10.1 | +| 2 | +| abc | +| def | ++------+ (6 rows) !ok -select 2 from foo group by dim1; -+--------+ -| EXPR$0 | -+--------+ -| 2 | -| 2 | -| 2 | -| 2 | -| 2 | -| 2 | -+--------+ -(6 rows) -!ok From ae737bedfdbf6ab28a10ce01ec72704eb7eab566 Mon Sep 17 00:00:00 2001 From: Zoltan Haindrich Date: Fri, 7 Mar 2025 12:38:30 +0000 Subject: [PATCH 38/50] make quidem run --- .../druid/msq/exec/TestMSQSqlModule.java | 4 +- .../druid/msq/test/DartComponentSupplier.java | 9 +++++ .../druid/sql/avatica/DartDruidMeta.java | 38 +++++++++++++++++++ .../dart.iq | 1 - 4 files changed, 49 insertions(+), 3 deletions(-) create mode 100644 extensions-core/multi-stage-query/src/test/java/org/apache/druid/sql/avatica/DartDruidMeta.java diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/TestMSQSqlModule.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/TestMSQSqlModule.java index 9ca1e69fbebe..0b48d2904dd3 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/TestMSQSqlModule.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/TestMSQSqlModule.java @@ -33,8 +33,8 @@ import org.apache.druid.msq.test.MSQTestTaskActionClient; import org.apache.druid.sql.SqlStatementFactory; import org.apache.druid.sql.SqlToolbox; -import org.apache.druid.sql.avatica.DartDruidMeta; import org.apache.druid.sql.avatica.DruidMeta; +import org.apache.druid.sql.avatica.MSQDruidMeta; public class TestMSQSqlModule extends TestDruidModule { @@ -83,7 +83,7 @@ private WorkerMemoryParameters makeWorkerMemoryParameters() @Provides @LazySingleton - public DruidMeta createMeta(DartDruidMeta druidMeta) + public DruidMeta createMeta(MSQDruidMeta druidMeta) { return druidMeta; } diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/DartComponentSupplier.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/DartComponentSupplier.java index f91f7b2344bf..2ecc7b449794 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/DartComponentSupplier.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/DartComponentSupplier.java @@ -45,6 +45,8 @@ import org.apache.druid.rpc.ServiceClientFactory; import org.apache.druid.rpc.guice.ServiceClientModule; import org.apache.druid.server.QueryLifecycleFactory; +import org.apache.druid.sql.avatica.DartDruidMeta; +import org.apache.druid.sql.avatica.DruidMeta; import org.apache.druid.sql.calcite.TempDirProducer; import org.apache.druid.sql.calcite.run.SqlEngine; import org.apache.druid.sql.calcite.util.DruidModuleCollection; @@ -126,6 +128,13 @@ public void configure(Binder binder) } static class LocalDartModule implements DruidModule{ + @Provides + @LazySingleton + public DruidMeta createMeta(DartDruidMeta druidMeta) + { + return druidMeta; + } + @Override public void configure(Binder binder) { diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/sql/avatica/DartDruidMeta.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/sql/avatica/DartDruidMeta.java new file mode 100644 index 000000000000..48d2f831c2d8 --- /dev/null +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/sql/avatica/DartDruidMeta.java @@ -0,0 +1,38 @@ +/* + * 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.sql.avatica; + +import com.google.inject.Inject; +import org.apache.druid.msq.dart.Dart; +import org.apache.druid.server.security.AuthenticatorMapper; +import org.apache.druid.sql.SqlStatementFactory; + +public class DartDruidMeta extends DruidMeta +{ + @Inject + public DartDruidMeta( + final @Dart SqlStatementFactory sqlStatementFactory, + final AvaticaServerConfig config, + final ErrorHandler errorHandler, + final AuthenticatorMapper authMapper) + { + super(sqlStatementFactory, config, errorHandler, authMapper); + } +} diff --git a/extensions-core/multi-stage-query/src/test/quidem/org.apache.druid.msq.quidem.MSQQuidemTest/dart.iq b/extensions-core/multi-stage-query/src/test/quidem/org.apache.druid.msq.quidem.MSQQuidemTest/dart.iq index 646f601a230f..45ab2abf744e 100644 --- a/extensions-core/multi-stage-query/src/test/quidem/org.apache.druid.msq.quidem.MSQQuidemTest/dart.iq +++ b/extensions-core/multi-stage-query/src/test/quidem/org.apache.druid.msq.quidem.MSQQuidemTest/dart.iq @@ -17,4 +17,3 @@ select count(1) from foo ; !ok - From 0d8089dd77b8affa6cdb28c0ade0dbe7f6cb46e9 Mon Sep 17 00:00:00 2001 From: Zoltan Haindrich Date: Fri, 7 Mar 2025 12:40:22 +0000 Subject: [PATCH 39/50] moretest --- .../dart.iq | 30 +++++++++++++++++-- 1 file changed, 28 insertions(+), 2 deletions(-) diff --git a/extensions-core/multi-stage-query/src/test/quidem/org.apache.druid.msq.quidem.MSQQuidemTest/dart.iq b/extensions-core/multi-stage-query/src/test/quidem/org.apache.druid.msq.quidem.MSQQuidemTest/dart.iq index 45ab2abf744e..6d4c5706e8f5 100644 --- a/extensions-core/multi-stage-query/src/test/quidem/org.apache.druid.msq.quidem.MSQQuidemTest/dart.iq +++ b/extensions-core/multi-stage-query/src/test/quidem/org.apache.druid.msq.quidem.MSQQuidemTest/dart.iq @@ -1,8 +1,7 @@ !use druidtest://?componentSupplier=DartComponentSupplier !set outputformat mysql ---select dim1 from foo ; -select count(1) from foo ; +select dim1 from foo ; +------+ | dim1 | +------+ @@ -17,3 +16,30 @@ select count(1) from foo ; !ok +select dim1 from foo ; ++------+ +| dim1 | ++------+ +| | +| 1 | +| 10.1 | +| 2 | +| abc | +| def | ++------+ +(6 rows) + +!ok + +select dim2,count(dim3) from foo group by dim2; ++------+--------+ +| dim2 | EXPR$1 | ++------+--------+ +| | 1 | +| a | 2 | +| abc | 0 | +| | 1 | ++------+--------+ +(4 rows) + +!ok From 07eb9754c23faf0c846fca96a9e2943b859c9bad Mon Sep 17 00:00:00 2001 From: Zoltan Haindrich Date: Fri, 7 Mar 2025 12:52:28 +0000 Subject: [PATCH 40/50] cleanup --- .../DartControllerContextFactoryImpl.java | 8 +- .../druid/msq/test/DartComponentSupplier.java | 20 -- .../test/MSQTestOverlordServiceClient2.java | 293 ------------------ .../TestDartControllerContextFactoryImpl.java | 26 +- 4 files changed, 9 insertions(+), 338 deletions(-) delete mode 100644 extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestOverlordServiceClient2.java diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartControllerContextFactoryImpl.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartControllerContextFactoryImpl.java index 64cf5ee1ab69..113714aa9b92 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartControllerContextFactoryImpl.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartControllerContextFactoryImpl.java @@ -28,7 +28,6 @@ import org.apache.druid.guice.annotations.Self; import org.apache.druid.guice.annotations.Smile; import org.apache.druid.java.util.emitter.service.ServiceEmitter; -import org.apache.druid.msq.dart.worker.DartWorkerClient; import org.apache.druid.msq.dart.worker.DartWorkerClientImpl; import org.apache.druid.msq.exec.ControllerContext; import org.apache.druid.msq.exec.MemoryIntrospector; @@ -75,15 +74,10 @@ public ControllerContext newContext(final String queryId) injector, jsonMapper, selfNode, - makeWorkerClient(queryId), + new DartWorkerClientImpl(queryId, serviceClientFactory, smileMapper, selfNode.getHostAndPortToUse()), memoryIntrospector, serverView, emitter ); } - - protected DartWorkerClient makeWorkerClient(final String queryId) - { - return new DartWorkerClientImpl(queryId, serviceClientFactory, smileMapper, selfNode.getHostAndPortToUse()); - } } diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/DartComponentSupplier.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/DartComponentSupplier.java index 2ecc7b449794..c47a29b82e4b 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/DartComponentSupplier.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/DartComponentSupplier.java @@ -20,7 +20,6 @@ package org.apache.druid.msq.test; import com.fasterxml.jackson.databind.ObjectMapper; -import com.google.common.collect.ImmutableList; import com.google.inject.Binder; import com.google.inject.Injector; import com.google.inject.Provides; @@ -40,7 +39,6 @@ import org.apache.druid.msq.dart.guice.DartWorkerMemoryManagementModule; import org.apache.druid.msq.dart.guice.DartWorkerModule; import org.apache.druid.msq.exec.Worker; -import org.apache.druid.msq.exec.WorkerMemoryParameters; import org.apache.druid.query.TestBufferPool; import org.apache.druid.rpc.ServiceClientFactory; import org.apache.druid.rpc.guice.ServiceClientModule; @@ -186,23 +184,5 @@ Map workerMap() // Execs.multiThreaded(controllerConfig.getConcurrentQueries(), "dart-controller-%s") // ); // } - @Provides - @LazySingleton - private MSQTestOverlordServiceClient2 makeOverlordServiceClient( - ObjectMapper queryJsonMapper, - Injector injector, - WorkerMemoryParameters workerMemoryParameters) - { - final MSQTestOverlordServiceClient2 indexingServiceClient = new MSQTestOverlordServiceClient2( - queryJsonMapper, - injector, - new MSQTestTaskActionClient(queryJsonMapper, injector), - workerMemoryParameters, - ImmutableList.of() - ); - return indexingServiceClient; - } - - } } diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestOverlordServiceClient2.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestOverlordServiceClient2.java deleted file mode 100644 index cc05e0ea146e..000000000000 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestOverlordServiceClient2.java +++ /dev/null @@ -1,293 +0,0 @@ -/* - * 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.test; - -import com.fasterxml.jackson.databind.ObjectMapper; -import com.google.common.base.Preconditions; -import com.google.common.collect.Iterables; -import com.google.common.util.concurrent.Futures; -import com.google.common.util.concurrent.ListenableFuture; -import com.google.common.util.concurrent.SettableFuture; -import com.google.inject.Injector; -import org.apache.calcite.sql.type.SqlTypeName; -import org.apache.druid.client.ImmutableSegmentLoadInfo; -import org.apache.druid.client.indexing.TaskPayloadResponse; -import org.apache.druid.client.indexing.TaskStatusResponse; -import org.apache.druid.indexer.TaskStatus; -import org.apache.druid.indexer.TaskStatusPlus; -import org.apache.druid.indexer.report.TaskReport; -import org.apache.druid.indexing.common.actions.TaskActionClient; -import org.apache.druid.java.util.common.DateTimes; -import org.apache.druid.java.util.common.ISE; -import org.apache.druid.msq.exec.Controller; -import org.apache.druid.msq.exec.ControllerImpl; -import org.apache.druid.msq.exec.QueryListener; -import org.apache.druid.msq.exec.ResultsContext; -import org.apache.druid.msq.exec.WorkerMemoryParameters; -import org.apache.druid.msq.indexing.MSQControllerTask; -import org.apache.druid.msq.indexing.destination.MSQDestination; -import org.apache.druid.msq.indexing.report.MSQResultsReport; -import org.apache.druid.msq.indexing.report.MSQStatusReport; -import org.apache.druid.msq.indexing.report.MSQTaskReport; -import org.apache.druid.msq.indexing.report.MSQTaskReportPayload; -import org.apache.druid.rpc.indexing.NoopOverlordClient; -import org.joda.time.DateTime; - -import javax.annotation.Nullable; -import java.util.ArrayList; -import java.util.HashMap; -import java.util.List; -import java.util.Map; - -public class MSQTestOverlordServiceClient2 extends NoopOverlordClient -{ - private final Injector injector; - private final ObjectMapper objectMapper; - private final TaskActionClient taskActionClient; - private final WorkerMemoryParameters workerMemoryParameters; - private final List loadedSegmentMetadata; - private final Map inMemoryControllers = new HashMap<>(); - private final Map reports = new HashMap<>(); - private final Map inMemoryControllerTask = new HashMap<>(); - private final Map inMemoryTaskStatus = new HashMap<>(); - - public static final DateTime CREATED_TIME = DateTimes.of("2023-05-31T12:00Z"); - public static final DateTime QUEUE_INSERTION_TIME = DateTimes.of("2023-05-31T12:01Z"); - - public static final long DURATION = 100L; - - public MSQTestOverlordServiceClient2( - ObjectMapper objectMapper, - Injector injector, - TaskActionClient taskActionClient, - WorkerMemoryParameters workerMemoryParameters, - List loadedSegmentMetadata - ) - { - this.objectMapper = objectMapper; - this.injector = injector; - this.taskActionClient = taskActionClient; - this.workerMemoryParameters = workerMemoryParameters; - this.loadedSegmentMetadata = loadedSegmentMetadata; - } - - @Override - public ListenableFuture runTask(String taskId, Object taskObject) - { - TestQueryListener queryListener = null; - ControllerImpl controller = null; - MSQTestControllerContext msqTestControllerContext; - try { - MSQControllerTask cTask = objectMapper.convertValue(taskObject, MSQControllerTask.class); - - msqTestControllerContext = new MSQTestControllerContext( - objectMapper, - injector, - taskActionClient, - workerMemoryParameters, - loadedSegmentMetadata, - cTask.getTaskLockType(), - cTask.getQuerySpec().getContext() - ); - - inMemoryControllerTask.put(cTask.getId(), cTask); - - controller = new ControllerImpl( - cTask.getId(), - cTask.getQuerySpec(), - new ResultsContext(cTask.getSqlTypeNames(), cTask.getSqlResultsContext()), - msqTestControllerContext - ); - - inMemoryControllers.put(controller.queryId(), controller); - - queryListener = - new TestQueryListener( - cTask.getId(), - cTask.getQuerySpec().getDestination() - ); - - try { - controller.run(queryListener); - inMemoryTaskStatus.put(taskId, queryListener.getStatusReport().toTaskStatus(cTask.getId())); - } - catch (Exception e) { - inMemoryTaskStatus.put(taskId, TaskStatus.failure(cTask.getId(), e.toString())); - } - return Futures.immediateFuture(null); - } - catch (Exception e) { - throw new ISE(e, "Unable to run"); - } - finally { - if (controller != null && queryListener != null) { - reports.put(controller.queryId(), queryListener.getReportMap()); - } - } - } - - @Override - public ListenableFuture cancelTask(String taskId) - { - inMemoryControllers.get(taskId).stop(); - return Futures.immediateFuture(null); - } - - @Override - public ListenableFuture taskReportAsMap(String taskId) - { - return Futures.immediateFuture(getReportForTask(taskId)); - } - - @Override - public ListenableFuture taskPayload(String taskId) - { - SettableFuture future = SettableFuture.create(); - future.set(new TaskPayloadResponse(taskId, getMSQControllerTask(taskId))); - return future; - } - - @Override - public ListenableFuture taskStatus(String taskId) - { - SettableFuture future = SettableFuture.create(); - TaskStatus taskStatus = inMemoryTaskStatus.get(taskId); - future.set(new TaskStatusResponse(taskId, new TaskStatusPlus( - taskId, - null, - MSQControllerTask.TYPE, - CREATED_TIME, - QUEUE_INSERTION_TIME, - taskStatus.getStatusCode(), - null, - DURATION, - taskStatus.getLocation(), - null, - taskStatus.getErrorMsg() - ))); - - return future; - } - - // hooks to pull stuff out for testing - @Nullable - public TaskReport.ReportMap getReportForTask(String id) - { - return reports.get(id); - } - - @Nullable - MSQControllerTask getMSQControllerTask(String id) - { - return inMemoryControllerTask.get(id); - } - - /** - * Listener that captures a report and makes it available through {@link #getReportMap()}. - */ - static class TestQueryListener implements QueryListener - { - private final String taskId; - private final MSQDestination destination; - private final List results = new ArrayList<>(); - - private List signature; - private List sqlTypeNames; - private boolean resultsTruncated = true; - private TaskReport.ReportMap reportMap; - - public TestQueryListener(final String taskId, final MSQDestination destination) - { - this.taskId = taskId; - this.destination = destination; - } - - @Override - public boolean readResults() - { - return destination.getRowsInTaskReport() == MSQDestination.UNLIMITED || destination.getRowsInTaskReport() > 0; - } - - @Override - public void onResultsStart(List signature, @Nullable List sqlTypeNames) - { - this.signature = signature; - this.sqlTypeNames = sqlTypeNames; - } - - @Override - public boolean onResultRow(Object[] row) - { - if (destination.getRowsInTaskReport() == MSQDestination.UNLIMITED - || results.size() < destination.getRowsInTaskReport()) { - results.add(row); - return true; - } else { - return false; - } - } - - @Override - public void onResultsComplete() - { - resultsTruncated = false; - } - - @Override - public void onQueryComplete(MSQTaskReportPayload report) - { - final MSQResultsReport resultsReport; - - if (signature != null) { - resultsReport = new MSQResultsReport( - signature, - sqlTypeNames, - results, - resultsTruncated - ); - } else { - resultsReport = null; - } - - final MSQTaskReport taskReport = new MSQTaskReport( - taskId, - new MSQTaskReportPayload( - report.getStatus(), - report.getStages(), - report.getCounters(), - resultsReport - ) - ); - - reportMap = TaskReport.buildTaskReports(taskReport); - } - - public TaskReport.ReportMap getReportMap() - { - return Preconditions.checkNotNull(reportMap, "reportMap"); - } - - public MSQStatusReport getStatusReport() - { - final MSQTaskReport taskReport = (MSQTaskReport) Iterables.getOnlyElement(getReportMap().values()); - return taskReport.getPayload().getStatus(); - } - } -} diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/TestDartControllerContextFactoryImpl.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/TestDartControllerContextFactoryImpl.java index 45843f42272f..267e4dc412ef 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/TestDartControllerContextFactoryImpl.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/TestDartControllerContextFactoryImpl.java @@ -28,6 +28,7 @@ import org.apache.druid.guice.annotations.Json; import org.apache.druid.guice.annotations.Self; import org.apache.druid.guice.annotations.Smile; +import org.apache.druid.java.util.common.io.Closer; import org.apache.druid.java.util.emitter.service.ServiceEmitter; import org.apache.druid.msq.dart.Dart; import org.apache.druid.msq.dart.controller.DartControllerContext; @@ -72,33 +73,22 @@ public TestDartControllerContextFactoryImpl( @Override public ControllerContext newContext(String queryId) { - DartControllerContext ctx; - DartTestWorkerClient wc; - ctx = new DartControllerContext( + return new DartControllerContext( injector, jsonMapper, selfNode, - wc = (DartTestWorkerClient) makeWorkerClient(queryId), + new DartTestWorkerClient(queryId, serviceClientFactory, smileMapper, selfNode.getHostAndPortToUse(), workerMap), memoryIntrospector, serverView, emitter - ) { - public void registerController(Controller controller, org.apache.druid.java.util.common.io.Closer closer) { + ) + { + public void registerController(Controller controller, Closer closer) + { super.registerController(controller, closer); - theController1=controller; - + theController1 = controller; }; }; - return ctx; - } - - - @Override - protected DartWorkerClient makeWorkerClient(String queryId) - { - return new DartTestWorkerClient( - queryId, serviceClientFactory, smileMapper, selfNode.getHostAndPortToUse(), workerMap - ); } public class DartTestWorkerClient extends MSQTestWorkerClient implements DartWorkerClient From 0230874b9a9696e35bb9d0547c6e0e9701eedbb8 Mon Sep 17 00:00:00 2001 From: Zoltan Haindrich Date: Mon, 10 Mar 2025 09:47:59 +0000 Subject: [PATCH 41/50] cleanup tests/etc --- .../msq/test/CalciteSelectQueryMSQTest.java | 8 ++--- .../druid/msq/test/MSQTestWorkerClient.java | 34 +++++++++++++++---- .../apache/druid/msq/test/SimpleDartTest.java | 18 ++++++++++ 3 files changed, 48 insertions(+), 12 deletions(-) diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/CalciteSelectQueryMSQTest.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/CalciteSelectQueryMSQTest.java index ba0aca17a114..33aa39e21a53 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/CalciteSelectQueryMSQTest.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/CalciteSelectQueryMSQTest.java @@ -22,6 +22,7 @@ import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import org.apache.druid.java.util.common.ISE; +import org.apache.druid.msq.sql.MSQTaskSqlEngine; import org.apache.druid.sql.calcite.CalciteQueryTest; import org.apache.druid.sql.calcite.QueryTestBuilder; import org.apache.druid.sql.calcite.SqlTestFrameworkConfig; @@ -30,22 +31,19 @@ import org.junit.jupiter.api.Test; import org.junit.jupiter.api.Timeout; -import java.util.UUID; import java.util.concurrent.TimeUnit; /** * Runs {@link CalciteQueryTest} but with MSQ engine */ -@SqlTestFrameworkConfig.ComponentSupplier(DartComponentSupplier.class) +@SqlTestFrameworkConfig.ComponentSupplier(StandardMSQComponentSupplier.class) public class CalciteSelectQueryMSQTest extends CalciteQueryTest { @Override protected QueryTestBuilder testBuilder() { return new QueryTestBuilder(new CalciteTestConfig(true)) -// .addCustomRunner(new ExtractResultsFactory(() -> null))//(MSQTestOverlordServiceClient) ((DartSqlEngine) queryFramework().engine()).overlordClient())) -// .addCustomRunner(new ExtractResultsFactory(() -> null))//(MSQTestOverlordServiceClient) ((DartSqlEngine) queryFramework().engine()).overlordClient())) - .queryContext(ImmutableMap.builder().put("asd",UUID.randomUUID().toString()).build()) + .addCustomRunner(new ExtractResultsFactory(() -> (MSQTestOverlordServiceClient) ((MSQTaskSqlEngine) queryFramework().engine()).overlordClient())) .skipVectorize(true) .verifyNativeQueries(new VerifyMSQSupportedNativeQueriesPredicate()); } diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestWorkerClient.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestWorkerClient.java index 90ef5d4da78c..f19abb20fd37 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestWorkerClient.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestWorkerClient.java @@ -33,6 +33,7 @@ import org.apache.druid.msq.rpc.SketchEncoding; import org.apache.druid.msq.statistics.ClusterByStatisticsSnapshot; +import java.io.InputStream; import java.util.Arrays; import java.util.Map; import java.util.concurrent.atomic.AtomicBoolean; @@ -132,10 +133,9 @@ public ListenableFuture fetchChannelData( final ReadableByteChunksFrameChannel channel ) { - return FutureUtils.transform( - getWorkerFor(workerTaskId).readStageOutput(stageId, partitionNumber, offset), - inputStream -> { - try { + if(true) { + try (InputStream inputStream = + inMemoryWorkers.get(workerTaskId).readStageOutput(stageId, partitionNumber, offset).get()) { byte[] buffer = new byte[8 * 1024]; boolean didRead = false; int bytesRead; @@ -145,13 +145,33 @@ public ListenableFuture fetchChannelData( } inputStream.close(); - return !didRead; + return Futures.immediateFuture(!didRead); } catch (Exception e) { throw new ISE(e, "Error reading frame file channel"); } - } - ); + } else { + return FutureUtils.transform( + getWorkerFor(workerTaskId).readStageOutput(stageId, partitionNumber, offset), + inputStream -> { + try { + byte[] buffer = new byte[8 * 1024]; + boolean didRead = false; + int bytesRead; + while ((bytesRead = inputStream.read(buffer)) != -1) { + channel.addChunk(Arrays.copyOf(buffer, bytesRead)); + didRead = true; + } + inputStream.close(); + + return !didRead; + } + catch (Exception e) { + throw new ISE(e, "Error reading frame file channel"); + } + } + ); + } } @Override diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/SimpleDartTest.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/SimpleDartTest.java index 66a8628fcee7..6b8401903015 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/SimpleDartTest.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/SimpleDartTest.java @@ -24,7 +24,9 @@ import org.apache.druid.sql.calcite.BaseCalciteQueryTest; import org.apache.druid.sql.calcite.QueryTestBuilder; import org.apache.druid.sql.calcite.SqlTestFrameworkConfig; +import org.junit.jupiter.api.Disabled; import org.junit.jupiter.api.Test; + import java.util.UUID; @SqlTestFrameworkConfig.ComponentSupplier(DartComponentSupplier.class) @@ -66,6 +68,22 @@ public void testSelectFromFoo() .run(); } + @Test + @Disabled("this case currently stalls") + public void testSelectFromFooLimit2() + { + testBuilder() + .sql("SELECT 2 from foo limit 2") + .expectedResults( + ImmutableList.of( + new Object[] {2}, + new Object[] {2} + ) + ) + .run(); + } + + @Test public void testSelectDim1FromFoo11() { From 96fa16aee7c3acecb73b32e02b15504d819a4171 Mon Sep 17 00:00:00 2001 From: Zoltan Haindrich Date: Mon, 10 Mar 2025 09:48:22 +0000 Subject: [PATCH 42/50] remove --- .../druid/msq/test/MSQTestWorkerClient.java | 35 ++++--------------- 1 file changed, 7 insertions(+), 28 deletions(-) diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestWorkerClient.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestWorkerClient.java index f19abb20fd37..9f482912f315 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestWorkerClient.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestWorkerClient.java @@ -33,7 +33,6 @@ import org.apache.druid.msq.rpc.SketchEncoding; import org.apache.druid.msq.statistics.ClusterByStatisticsSnapshot; -import java.io.InputStream; import java.util.Arrays; import java.util.Map; import java.util.concurrent.atomic.AtomicBoolean; @@ -133,9 +132,10 @@ public ListenableFuture fetchChannelData( final ReadableByteChunksFrameChannel channel ) { - if(true) { - try (InputStream inputStream = - inMemoryWorkers.get(workerTaskId).readStageOutput(stageId, partitionNumber, offset).get()) { + return FutureUtils.transform( + getWorkerFor(workerTaskId).readStageOutput(stageId, partitionNumber, offset), + inputStream -> { + try { byte[] buffer = new byte[8 * 1024]; boolean didRead = false; int bytesRead; @@ -144,34 +144,13 @@ public ListenableFuture fetchChannelData( didRead = true; } inputStream.close(); - - return Futures.immediateFuture(!didRead); + return !didRead; } catch (Exception e) { throw new ISE(e, "Error reading frame file channel"); } - } else { - return FutureUtils.transform( - getWorkerFor(workerTaskId).readStageOutput(stageId, partitionNumber, offset), - inputStream -> { - try { - byte[] buffer = new byte[8 * 1024]; - boolean didRead = false; - int bytesRead; - while ((bytesRead = inputStream.read(buffer)) != -1) { - channel.addChunk(Arrays.copyOf(buffer, bytesRead)); - didRead = true; - } - inputStream.close(); - - return !didRead; - } - catch (Exception e) { - throw new ISE(e, "Error reading frame file channel"); - } - } - ); - } + } + ); } @Override From 39c4f42b9b85da9e8b528ba5803b5a9a0b6250a3 Mon Sep 17 00:00:00 2001 From: Zoltan Haindrich Date: Mon, 10 Mar 2025 10:07:51 +0000 Subject: [PATCH 43/50] cleanup --- .../dart/controller/sql/DartQueryMaker.java | 2 +- .../dart/worker/DartWorkerFactoryImpl.java | 9 +--- .../apache/druid/msq/exec/ControllerImpl.java | 1 - .../org/apache/druid/msq/exec/WorkerImpl.java | 7 +-- .../druid/msq/test/DartComponentSupplier.java | 22 +++----- .../apache/druid/msq/test/SimpleDartTest.java | 7 ++- .../druid/sql/calcite/CalciteQueryTest.java | 53 +++---------------- .../sql/calcite/util/SqlTestFramework.java | 24 --------- .../calcite/util/TestTimelineServerView.java | 22 ++++---- 9 files changed, 32 insertions(+), 115 deletions(-) diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartQueryMaker.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartQueryMaker.java index dee475ff5491..4526754cf573 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartQueryMaker.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartQueryMaker.java @@ -141,7 +141,7 @@ public QueryResponse runQuery(DruidQuery druidQuery) final List> types = MSQTaskQueryMaker.getTypes(druidQuery, fieldMapping, plannerContext); - final String dartQueryId = druidQuery.getQuery().context().getString(DartSqlEngine.CTX_DART_QUERY_ID) +"asd"; + final String dartQueryId = druidQuery.getQuery().context().getString(DartSqlEngine.CTX_DART_QUERY_ID); final ControllerContext controllerContext = controllerContextFactory.newContext(dartQueryId); final ControllerImpl controller = new ControllerImpl( dartQueryId, diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartWorkerFactoryImpl.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartWorkerFactoryImpl.java index f5ef8b4b9585..1960924b4b67 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartWorkerFactoryImpl.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartWorkerFactoryImpl.java @@ -35,7 +35,6 @@ import org.apache.druid.msq.exec.WorkerContext; import org.apache.druid.msq.exec.WorkerImpl; import org.apache.druid.msq.querykit.DataSegmentProvider; -import org.apache.druid.msq.shuffle.output.StageOutputHolder; import org.apache.druid.query.DruidProcessingConfig; import org.apache.druid.query.QueryContext; import org.apache.druid.query.groupby.GroupingEngine; @@ -114,12 +113,6 @@ public Worker build(String queryId, String controllerHost, File tempDir, QueryCo queryContext ); - return new WorkerImpl(null, workerContext) - { - protected StageOutputHolder makeStageOutputHolder() - { - return new StageOutputHolder(); - }; - }; + return new WorkerImpl(null, workerContext); } } diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/ControllerImpl.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/ControllerImpl.java index 2394f367b01f..cc73ea42c1a4 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/ControllerImpl.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/ControllerImpl.java @@ -377,7 +377,6 @@ private MSQTaskReportPayload runInternal(final QueryListener queryListener, fina handleQueryResults(queryDef, queryKernel); } catch (Throwable e) { - log.info(e, "Execution failed"); exceptionEncountered = e; } diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/WorkerImpl.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/WorkerImpl.java index ec057e72fea6..44654810ac61 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/WorkerImpl.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/WorkerImpl.java @@ -964,12 +964,7 @@ private void removeStageDurableStorageOutput(final StageId stageId) private StageOutputHolder getOrCreateStageOutputHolder(final StageId stageId, final int partitionNumber) { return stageOutputs.computeIfAbsent(stageId, ignored1 -> new ConcurrentHashMap<>()) - .computeIfAbsent(partitionNumber, ignored -> makeStageOutputHolder()); - } - - protected StageOutputHolder makeStageOutputHolder() - { - return new StageOutputHolder(); + .computeIfAbsent(partitionNumber, ignored -> new StageOutputHolder()); } /** diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/DartComponentSupplier.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/DartComponentSupplier.java index c47a29b82e4b..2393404168d1 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/DartComponentSupplier.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/DartComponentSupplier.java @@ -23,7 +23,6 @@ import com.google.inject.Binder; import com.google.inject.Injector; import com.google.inject.Provides; -import com.google.inject.TypeLiteral; import org.apache.druid.collections.NonBlockingPool; import org.apache.druid.discovery.DruidNodeDiscoveryProvider; import org.apache.druid.guice.LazySingleton; @@ -117,12 +116,7 @@ final DruidNodeDiscoveryProvider getDiscoveryProvider() { @Override public void configure(Binder binder) { -// binder.bind(new TypeLiteral>(){}) -// .annotatedWith(Merging.class) -// .to(TestBufferPool.class); - } - } static class LocalDartModule implements DruidModule{ @@ -136,20 +130,16 @@ public DruidMeta createMeta(DartDruidMeta druidMeta) @Override public void configure(Binder binder) { - binder.bind(new TypeLiteral>(){}) - .annotatedWith(Merging.class) - .to(TestBufferPool.class); - - if(true) { binder.bind(DartControllerContextFactory.class) .to(TestDartControllerContextFactoryImpl.class) .in(LazySingleton.class); - }else { - binder.bind(DartControllerContextFactory.class) - .to(MSQTestControllerContext.class) - .in(LazySingleton.class); + } - } + @Provides + @Merging + NonBlockingPool makeMergingBuffer(TestBufferPool bufferPool) + { + return bufferPool; } @Provides diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/SimpleDartTest.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/SimpleDartTest.java index 6b8401903015..a81798c3bee6 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/SimpleDartTest.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/SimpleDartTest.java @@ -21,6 +21,7 @@ import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; +import org.apache.druid.msq.dart.controller.sql.DartSqlEngine; import org.apache.druid.sql.calcite.BaseCalciteQueryTest; import org.apache.druid.sql.calcite.QueryTestBuilder; import org.apache.druid.sql.calcite.SqlTestFrameworkConfig; @@ -36,7 +37,11 @@ public class SimpleDartTest extends BaseCalciteQueryTest protected QueryTestBuilder testBuilder() { return new QueryTestBuilder(new CalciteTestConfig(true)) - .queryContext(ImmutableMap.builder().put("asd", UUID.randomUUID().toString()).build()) + .queryContext( + ImmutableMap.builder() + .put(DartSqlEngine.CTX_DART_QUERY_ID, UUID.randomUUID().toString()) + .build() + ) .skipVectorize(true) .verifyNativeQueries(new VerifyMSQSupportedNativeQueriesPredicate()); } diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/CalciteQueryTest.java b/sql/src/test/java/org/apache/druid/sql/calcite/CalciteQueryTest.java index c9606cdea8de..90a64d5fc703 100644 --- a/sql/src/test/java/org/apache/druid/sql/calcite/CalciteQueryTest.java +++ b/sql/src/test/java/org/apache/druid/sql/calcite/CalciteQueryTest.java @@ -5256,51 +5256,8 @@ public void testInIsNotTrueFilter() @Test public void testNotInOrIsNullFilter() { - testQuery1( - "SELECT 1", - ImmutableList.of( - GroupByQuery.builder() - .setDataSource(CalciteTests.DATASOURCE1) - .setInterval(querySegmentSpec(Filtration.eternity())) - .setGranularity(Granularities.ALL) - .setDimensions(dimensions(new DefaultDimensionSpec("dim1", "d0"))) - .setDimFilter(or( - isNull("dim1"), - not(in("dim1", ColumnType.STRING, ImmutableList.of("abc", "def", "ghi"))) - )) - .setAggregatorSpecs( - aggregators( - new CountAggregatorFactory("a0") - ) - ) - .setContext(QUERY_CONTEXT_DEFAULT) - .build() - ), - ImmutableList.of( - new Object[]{1} - ) - ); - } - - - public void testQuery1( - final String sql, - final List> expectedQueries, - final List expectedResults - ) - { - testBuilder() - .sql(sql) -// .expectedQueries(expectedQueries) - .expectedResults(expectedResults) - .run(); - } - - @Test - public void testSelect1() - { - testQuery1( - "SELECT 1", + testQuery( + "SELECT dim1, COUNT(*) FROM druid.foo WHERE dim1 NOT IN ('ghi', 'abc', 'def') OR dim1 IS NULL GROUP BY dim1", ImmutableList.of( GroupByQuery.builder() .setDataSource(CalciteTests.DATASOURCE1) @@ -5320,12 +5277,14 @@ public void testSelect1() .build() ), ImmutableList.of( - new Object[]{1L} + new Object[]{"", 1L}, + new Object[]{"1", 1L}, + new Object[]{"10.1", 1L}, + new Object[]{"2", 1L} ) ); } - @Test public void testNotInAndIsNotNullFilter() { diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/util/SqlTestFramework.java b/sql/src/test/java/org/apache/druid/sql/calcite/util/SqlTestFramework.java index 4f8a002cf445..ebe2e63d1326 100644 --- a/sql/src/test/java/org/apache/druid/sql/calcite/util/SqlTestFramework.java +++ b/sql/src/test/java/org/apache/druid/sql/calcite/util/SqlTestFramework.java @@ -29,14 +29,6 @@ import com.google.inject.Module; import com.google.inject.Provides; import com.google.inject.TypeLiteral; -import org.apache.druid.client.BrokerServerView; -import org.apache.druid.client.DirectDruidClientFactory; -import org.apache.druid.client.FilteredServerInventoryView; -import org.apache.druid.client.FilteredServerInventoryViewProvider; -import org.apache.druid.client.QueryableDruidServer; -import org.apache.druid.client.ServerInventoryView; -import org.apache.druid.client.ServerInventoryViewProvider; -import org.apache.druid.client.SimpleServerView; import org.apache.druid.client.TestHttpClient; import org.apache.druid.client.TimelineServerView; import org.apache.druid.client.cache.Cache; @@ -48,7 +40,6 @@ import org.apache.druid.guice.ExpressionModule; import org.apache.druid.guice.LazySingleton; import org.apache.druid.guice.LifecycleModule; -import org.apache.druid.guice.ManageLifecycle; import org.apache.druid.guice.QueryRunnerFactoryModule; import org.apache.druid.guice.QueryableModule; import org.apache.druid.guice.SegmentWranglerModule; @@ -1136,21 +1127,6 @@ public QueryScheduler makeQueryScheduler() @Override public void configure(Binder binder) { - if(true) { - if(false) { - binder.bind(TimelineServerView.class).to(SimpleServerView.class).in(LazySingleton.class); - } - - } else { - binder.bind(TimelineServerView.class).to(BrokerServerView.class).in(LazySingleton.class); - binder.bind(QueryableDruidServer.Maker.class).to(DirectDruidClientFactory.class).in(LazySingleton.class); - binder.bind(ServerInventoryView.class).toProvider(ServerInventoryViewProvider.class).in(ManageLifecycle.class); - binder.bind(FilteredServerInventoryView.class) - .toProvider(FilteredServerInventoryViewProvider.class) - .in(ManageLifecycle.class); - - // return new SimpleServerView(conglomerate, objectMapper, httpClient); - } } } diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/util/TestTimelineServerView.java b/sql/src/test/java/org/apache/druid/sql/calcite/util/TestTimelineServerView.java index dcbea0d059ef..88df3b80a4e5 100644 --- a/sql/src/test/java/org/apache/druid/sql/calcite/util/TestTimelineServerView.java +++ b/sql/src/test/java/org/apache/druid/sql/calcite/util/TestTimelineServerView.java @@ -99,22 +99,22 @@ public TestTimelineServerView(List segments, List real @Override public Optional> getTimeline(TableDataSource table) { - for (DataSegment s : segments) { - if (!s.getDataSource().equals(table.getName())) { + for (DataSegment segment : segments) { + if (!segment.getDataSource().equals(table.getName())) { continue; } - VersionedIntervalTimeline a = new VersionedIntervalTimeline(Comparator.naturalOrder()); - TierSelectorStrategy st = new HighestPriorityTierSelectorStrategy(new RandomServerSelectorStrategy()); - ServerSelector sss = new ServerSelector(s, st ); - PartitionChunk aa = new SingleElementPartitionChunk(sss); - a.add( s.getInterval(), s.getVersion(), aa); - return Optional.of(a); -// return Optional.of(); + VersionedIntervalTimeline timelineLookup = new VersionedIntervalTimeline( + Comparator.naturalOrder() + ); + TierSelectorStrategy st = new HighestPriorityTierSelectorStrategy(new RandomServerSelectorStrategy()); + ServerSelector sss = new ServerSelector(segment, st); + PartitionChunk partitionChunk = new SingleElementPartitionChunk(sss); + timelineLookup.add(segment.getInterval(), segment.getVersion(), partitionChunk); + return Optional.of(timelineLookup); } - return Optional.empty(); -// throw new UnsupportedOperationException(); + return Optional.empty(); } @Override From 07b986c05da76053f1379301639fdcffa1607bda Mon Sep 17 00:00:00 2001 From: Zoltan Haindrich Date: Mon, 10 Mar 2025 10:43:57 +0000 Subject: [PATCH 44/50] fixes --- .../druid/msq/test/DartComponentSupplier.java | 66 +++++++------------ .../TestDartControllerContextFactoryImpl.java | 13 ++-- .../dart.iq | 1 + 3 files changed, 29 insertions(+), 51 deletions(-) diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/DartComponentSupplier.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/DartComponentSupplier.java index 2393404168d1..7cfac91a4ce5 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/DartComponentSupplier.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/DartComponentSupplier.java @@ -69,22 +69,24 @@ public void gatherProperties(Properties properties) } @Override - public DruidModule getOverrideModule() + public DruidModule getCoreModule() { - return DruidModuleCollection.of(super.getOverrideModule(), - new LocalDartModule() - ); + return DruidModuleCollection.of( + super.getCoreModule(), + new DartControllerModule(), + new DartWorkerModule(), + new DartWorkerMemoryManagementModule(), + new DartTestCoreModule() + ); } @Override - public DruidModule getCoreModule() + public DruidModule getOverrideModule() { - return DruidModuleCollection.of(super.getCoreModule(), - new DartControllerModule(), - new DartWorkerModule(), - new DartWorkerMemoryManagementModule(), - new L1() - ); + return DruidModuleCollection.of( + super.getOverrideModule(), + new DartTestOverrideModule() + ); } @Override @@ -96,29 +98,30 @@ public SqlEngine createEngine( return injector.getInstance(DartSqlEngine.class); } - - - static class L1 implements DruidModule{ - + static class DartTestCoreModule implements DruidModule + { @Provides @EscalatedGlobal - final ServiceClientFactory getServiceClientFactory(HttpClient ht) { + final ServiceClientFactory getServiceClientFactory(HttpClient ht) + { return ServiceClientModule.makeServiceClientFactory(ht); } @Provides - final DruidNodeDiscoveryProvider getDiscoveryProvider() { + final DruidNodeDiscoveryProvider getDiscoveryProvider() + { return null; } - @Override public void configure(Binder binder) { } } - static class LocalDartModule implements DruidModule{ + + static class DartTestOverrideModule implements DruidModule + { @Provides @LazySingleton @@ -149,30 +152,5 @@ Map workerMap() { return new HashMap(); } - -// @Provides -// @LazySingleton -// public MSQTaskSqlEngine createEngine( -// ObjectMapper queryJsonMapper, -// MSQTestOverlordServiceClient2 indexingServiceClient) -// { -// return new DartSqlEngine(indexingServiceClient, queryJsonMapper, new SegmentGenerationTerminalStageSpecFactory()); -// } -// -// @Provides -// @LazySingleton -// public DartSqlEngine makeSqlEngine( -// DartControllerContextFactory controllerContextFactory, -// DartControllerRegistry controllerRegistry, -// DartControllerConfig controllerConfig -// ) -// { -// return new DartSqlEngine( -// controllerContextFactory, -// controllerRegistry, -// controllerConfig, -// Execs.multiThreaded(controllerConfig.getConcurrentQueries(), "dart-controller-%s") -// ); -// } } } diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/TestDartControllerContextFactoryImpl.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/TestDartControllerContextFactoryImpl.java index 267e4dc412ef..1c25537f0de3 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/TestDartControllerContextFactoryImpl.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/TestDartControllerContextFactoryImpl.java @@ -52,7 +52,7 @@ public class TestDartControllerContextFactoryImpl extends DartControllerContextFactoryImpl { private Map workerMap; - public Controller theController1; + public Controller controller; @Inject public TestDartControllerContextFactoryImpl( @@ -83,11 +83,11 @@ public ControllerContext newContext(String queryId) emitter ) { - public void registerController(Controller controller, Closer closer) + public void registerController(Controller currentController, Closer closer) { - super.registerController(controller, closer); - theController1 = controller; - }; + super.registerController(currentController, closer); + controller = currentController; + } }; } @@ -113,7 +113,7 @@ protected Worker newWorker(String workerId) new MSQTestWorkerContext( queryId, inMemoryWorkers, - theController1, + controller, jsonMapper, injector, MSQTestBase.makeTestWorkerMemoryParameters(), @@ -155,7 +155,6 @@ public void closeClient(String hostAndPort) public ListenableFuture stopWorker(String workerId) { return null; - } } } diff --git a/extensions-core/multi-stage-query/src/test/quidem/org.apache.druid.msq.quidem.MSQQuidemTest/dart.iq b/extensions-core/multi-stage-query/src/test/quidem/org.apache.druid.msq.quidem.MSQQuidemTest/dart.iq index 6d4c5706e8f5..d232a52ac20f 100644 --- a/extensions-core/multi-stage-query/src/test/quidem/org.apache.druid.msq.quidem.MSQQuidemTest/dart.iq +++ b/extensions-core/multi-stage-query/src/test/quidem/org.apache.druid.msq.quidem.MSQQuidemTest/dart.iq @@ -1,3 +1,4 @@ +!set dartQueryId 00000000-0000-0000-0000-000000000000 !use druidtest://?componentSupplier=DartComponentSupplier !set outputformat mysql From af3beae5f1afdf212e5c253bcd11f3482737de33 Mon Sep 17 00:00:00 2001 From: Zoltan Haindrich Date: Mon, 10 Mar 2025 10:52:12 +0000 Subject: [PATCH 45/50] move doc --- .../druid/msq/dart/worker/DartWorkerClient.java | 11 +++++++++-- .../druid/msq/dart/worker/DartWorkerClientImpl.java | 7 ++----- 2 files changed, 11 insertions(+), 7 deletions(-) diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartWorkerClient.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartWorkerClient.java index 1ad7db873470..e048f9fd6d67 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartWorkerClient.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartWorkerClient.java @@ -20,13 +20,20 @@ package org.apache.druid.msq.dart.worker; import com.google.common.util.concurrent.ListenableFuture; +import org.apache.druid.msq.dart.controller.DartWorkerManager; import org.apache.druid.msq.exec.WorkerClient; public interface DartWorkerClient extends WorkerClient { + /** + * Close a single worker's clients. Used when that worker fails, so we stop trying to contact it. + * + * @param workerHost worker host:port + */ void closeClient(String hostAndPort); + /** + * Stops a worker. Dart-only API, used by the {@link DartWorkerManager}. + */ ListenableFuture stopWorker(String workerId); - - } diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartWorkerClientImpl.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartWorkerClientImpl.java index 4814eb6d5bb3..a75c15533eca 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartWorkerClientImpl.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartWorkerClientImpl.java @@ -27,7 +27,6 @@ import org.apache.druid.error.DruidException; import org.apache.druid.java.util.common.logger.Logger; import org.apache.druid.java.util.http.client.response.HttpResponseHandler; -import org.apache.druid.msq.dart.controller.DartWorkerManager; import org.apache.druid.msq.dart.controller.sql.DartSqlEngine; import org.apache.druid.msq.dart.worker.http.DartWorkerResource; import org.apache.druid.msq.exec.WorkerClient; @@ -112,9 +111,7 @@ protected ServiceClient getClient(final String workerIdString) } /** - * Close a single worker's clients. Used when that worker fails, so we stop trying to contact it. - * - * @param workerHost worker host:port + * {@inheritDoc} */ public void closeClient(final String workerHost) { @@ -145,7 +142,7 @@ public void close() } /** - * Stops a worker. Dart-only API, used by the {@link DartWorkerManager}. + * {@inheritDoc} */ public ListenableFuture stopWorker(String workerId) { From a21433e572df185b34053bc906904ce943b574b6 Mon Sep 17 00:00:00 2001 From: Zoltan Haindrich Date: Mon, 10 Mar 2025 14:43:06 +0000 Subject: [PATCH 46/50] fixes bug --- .../druid/sql/calcite/util/CalciteTests.java | 5 +- .../sql/calcite/util/QueryFrameworkUtils.java | 38 ++------------ .../sql/calcite/util/SqlTestFramework.java | 50 +++---------------- 3 files changed, 14 insertions(+), 79 deletions(-) diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/util/CalciteTests.java b/sql/src/test/java/org/apache/druid/sql/calcite/util/CalciteTests.java index cb3ceaf442d7..abf32c9bf321 100644 --- a/sql/src/test/java/org/apache/druid/sql/calcite/util/CalciteTests.java +++ b/sql/src/test/java/org/apache/druid/sql/calcite/util/CalciteTests.java @@ -32,7 +32,6 @@ import org.apache.druid.client.FilteredServerInventoryView; import org.apache.druid.client.ServerInventoryView; import org.apache.druid.client.ServerView; -import org.apache.druid.client.TimelineServerView; import org.apache.druid.discovery.DiscoveryDruidNode; import org.apache.druid.discovery.DruidLeaderClient; import org.apache.druid.discovery.DruidNodeDiscovery; @@ -398,7 +397,7 @@ NodeRole.COORDINATOR, new FakeDruidNodeDiscovery(ImmutableMap.of(NodeRole.COORDI public static SystemSchema createMockSystemSchema( final DruidSchema druidSchema, - final TimelineServerView timelineServerView, + final SpecificSegmentsQuerySegmentWalker walker, final AuthorizerMapper authorizerMapper ) { @@ -475,7 +474,7 @@ private TaskStatusPlus createTaskStatus(String id, String datasource, Long durat new BrokerSegmentWatcherConfig(), BrokerSegmentMetadataCacheConfig.create() ), - timelineServerView, + new TestTimelineServerView(walker.getSegments()), new FakeServerInventoryView(), authorizerMapper, druidLeaderClient, diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/util/QueryFrameworkUtils.java b/sql/src/test/java/org/apache/druid/sql/calcite/util/QueryFrameworkUtils.java index f61f52343987..f28461e1711c 100644 --- a/sql/src/test/java/org/apache/druid/sql/calcite/util/QueryFrameworkUtils.java +++ b/sql/src/test/java/org/apache/druid/sql/calcite/util/QueryFrameworkUtils.java @@ -26,7 +26,6 @@ import org.apache.calcite.jdbc.CalciteSchema; import org.apache.calcite.schema.SchemaPlus; import org.apache.druid.client.InternalQueryConfig; -import org.apache.druid.client.TimelineServerView; import org.apache.druid.java.util.emitter.core.NoopEmitter; import org.apache.druid.java.util.emitter.service.ServiceEmitter; import org.apache.druid.query.DefaultGenericQueryMetricsFactory; @@ -126,32 +125,7 @@ public static DruidSchemaCatalog createMockRootSchema( @Nullable final ViewManager viewManager, final DruidSchemaManager druidSchemaManager, final AuthorizerMapper authorizerMapper, - final CatalogResolver catalogResolver) - { - TimelineServerView timelineServerView = new TestTimelineServerView(walker.getSegments()); - return createMockRootSchema( - injector, - conglomerate, - walker, - plannerConfig, - viewManager, - druidSchemaManager, - authorizerMapper, - catalogResolver, - timelineServerView - ); - } - - public static DruidSchemaCatalog createMockRootSchema( - final Injector injector, - final QueryRunnerFactoryConglomerate conglomerate, - final SpecificSegmentsQuerySegmentWalker walker, - final PlannerConfig plannerConfig, - @Nullable final ViewManager viewManager, - final DruidSchemaManager druidSchemaManager, - final AuthorizerMapper authorizerMapper, - final CatalogResolver catalogResolver, - final TimelineServerView timelineServerView + final CatalogResolver catalogResolver ) { DruidSchema druidSchema = createMockSchema( @@ -159,11 +133,10 @@ public static DruidSchemaCatalog createMockRootSchema( conglomerate, walker, druidSchemaManager, - catalogResolver, - timelineServerView + catalogResolver ); SystemSchema systemSchema = - CalciteTests.createMockSystemSchema(druidSchema, timelineServerView, authorizerMapper); + CalciteTests.createMockSystemSchema(druidSchema, walker, authorizerMapper); LookupSchema lookupSchema = createMockLookupSchema(injector); DruidOperatorTable createOperatorTable = createOperatorTable(injector); @@ -248,13 +221,12 @@ static DruidSchema createMockSchema( final QueryRunnerFactoryConglomerate conglomerate, final SpecificSegmentsQuerySegmentWalker walker, final DruidSchemaManager druidSchemaManager, - final CatalogResolver catalog, - final TimelineServerView timelineServerView + final CatalogResolver catalog ) { final BrokerSegmentMetadataCache cache = new BrokerSegmentMetadataCache( createMockQueryLifecycleFactory(walker, conglomerate), - timelineServerView, + new TestTimelineServerView(walker.getSegments()), BrokerSegmentMetadataCacheConfig.create(), CalciteTests.TEST_AUTHENTICATOR_ESCALATOR, new InternalQueryConfig(), diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/util/SqlTestFramework.java b/sql/src/test/java/org/apache/druid/sql/calcite/util/SqlTestFramework.java index ebe2e63d1326..90cc53d1b9a0 100644 --- a/sql/src/test/java/org/apache/druid/sql/calcite/util/SqlTestFramework.java +++ b/sql/src/test/java/org/apache/druid/sql/calcite/util/SqlTestFramework.java @@ -29,11 +29,8 @@ import com.google.inject.Module; import com.google.inject.Provides; import com.google.inject.TypeLiteral; -import org.apache.druid.client.TestHttpClient; -import org.apache.druid.client.TimelineServerView; import org.apache.druid.client.cache.Cache; import org.apache.druid.client.cache.CacheConfig; -import org.apache.druid.collections.BlockingPool; import org.apache.druid.collections.NonBlockingPool; import org.apache.druid.guice.BuiltInTypesModule; import org.apache.druid.guice.DruidInjectorBuilder; @@ -47,14 +44,12 @@ import org.apache.druid.guice.StartupInjectorBuilder; import org.apache.druid.guice.annotations.Global; import org.apache.druid.guice.annotations.Merging; -import org.apache.druid.guice.annotations.Self; import org.apache.druid.initialization.CoreInjectorBuilder; import org.apache.druid.initialization.DruidModule; import org.apache.druid.initialization.ServiceInjectorBuilder; import org.apache.druid.java.util.common.RE; import org.apache.druid.java.util.common.io.Closer; import org.apache.druid.java.util.emitter.service.ServiceEmitter; -import org.apache.druid.java.util.http.client.HttpClient; import org.apache.druid.math.expr.ExprMacroTable; import org.apache.druid.query.DefaultGenericQueryMetricsFactory; import org.apache.druid.query.DruidProcessingConfig; @@ -83,7 +78,6 @@ import org.apache.druid.segment.realtime.ChatHandlerProvider; import org.apache.druid.segment.realtime.NoopChatHandlerProvider; import org.apache.druid.server.ClientQuerySegmentWalker; -import org.apache.druid.server.DruidNode; import org.apache.druid.server.LocalQuerySegmentWalker; import org.apache.druid.server.QueryLifecycle; import org.apache.druid.server.QueryLifecycleFactory; @@ -678,8 +672,7 @@ public PlannerFixture( viewManager, componentSupplier.createSchemaManager(), framework.authorizerMapper, - framework.builder.catalogResolver, - framework.injector.getInstance(TimelineServerView.class) + framework.builder.catalogResolver ); this.plannerFactory = new PlannerFactory( @@ -778,22 +771,10 @@ public void configure(Binder binder) .annotatedWith(Global.class) .to(TestBufferPool.class); - binder.bind(new TypeLiteral>(){}) - .annotatedWith(Merging.class) - .to(TestBufferPool.class); - TestRequestLogger testRequestLogger = new TestRequestLogger(); binder.bind(RequestLogger.class).toInstance(testRequestLogger); } - @Provides - @Self - @LazySingleton - public DruidNode makeSelfDruidNode() - { - return new DruidNode("test", "test-host", true, 80, 443, false, true); - } - @Provides AuthorizerMapper getAuthorizerMapper() { @@ -942,36 +923,26 @@ private DruidSchema makeDruidSchema( final Injector injector, QueryRunnerFactoryConglomerate conglomerate, QuerySegmentWalker walker, - Builder builder, - TimelineServerView timelineServerView) + Builder builder) { return QueryFrameworkUtils.createMockSchema( injector, conglomerate, (SpecificSegmentsQuerySegmentWalker) walker, builder.componentSupplier.getPlannerComponentSupplier().createSchemaManager(), - builder.catalogResolver, - timelineServerView + builder.catalogResolver ); } @Provides @LazySingleton - private SystemSchema makeSystemSchema(AuthorizerMapper authorizerMapper, DruidSchema druidSchema, - TimelineServerView timelineServerView) - { - return CalciteTests.createMockSystemSchema(druidSchema, timelineServerView, authorizerMapper); - } - - - @Provides - @LazySingleton - private TimelineServerView makeTimelineServerView(SpecificSegmentsQuerySegmentWalker walker) + private SystemSchema makeSystemSchema(QuerySegmentWalker walker, AuthorizerMapper authorizerMapper, + DruidSchema druidSchema) { - return new TestTimelineServerView(walker.getSegments()); + return CalciteTests + .createMockSystemSchema(druidSchema, (SpecificSegmentsQuerySegmentWalker) walker, authorizerMapper); } - @Provides @LazySingleton private ColumnConfig getColumnConfig() @@ -1068,13 +1039,6 @@ public TestSegmentsBroker makeTimelines() return new TestSegmentsBroker(); } - @Provides - @LazySingleton - private HttpClient makeHttpClient(ObjectMapper objectMapper) - { - return new TestHttpClient(objectMapper); - } - @Provides @Named("empty") @LazySingleton From e60eea5b83aca575aab5a47277cd99818dd4d111 Mon Sep 17 00:00:00 2001 From: Zoltan Haindrich Date: Mon, 10 Mar 2025 14:43:08 +0000 Subject: [PATCH 47/50] Revert "fixes bug" This reverts commit a21433e572df185b34053bc906904ce943b574b6. --- .../druid/sql/calcite/util/CalciteTests.java | 5 +- .../sql/calcite/util/QueryFrameworkUtils.java | 38 ++++++++++++-- .../sql/calcite/util/SqlTestFramework.java | 50 ++++++++++++++++--- 3 files changed, 79 insertions(+), 14 deletions(-) diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/util/CalciteTests.java b/sql/src/test/java/org/apache/druid/sql/calcite/util/CalciteTests.java index abf32c9bf321..cb3ceaf442d7 100644 --- a/sql/src/test/java/org/apache/druid/sql/calcite/util/CalciteTests.java +++ b/sql/src/test/java/org/apache/druid/sql/calcite/util/CalciteTests.java @@ -32,6 +32,7 @@ import org.apache.druid.client.FilteredServerInventoryView; import org.apache.druid.client.ServerInventoryView; import org.apache.druid.client.ServerView; +import org.apache.druid.client.TimelineServerView; import org.apache.druid.discovery.DiscoveryDruidNode; import org.apache.druid.discovery.DruidLeaderClient; import org.apache.druid.discovery.DruidNodeDiscovery; @@ -397,7 +398,7 @@ NodeRole.COORDINATOR, new FakeDruidNodeDiscovery(ImmutableMap.of(NodeRole.COORDI public static SystemSchema createMockSystemSchema( final DruidSchema druidSchema, - final SpecificSegmentsQuerySegmentWalker walker, + final TimelineServerView timelineServerView, final AuthorizerMapper authorizerMapper ) { @@ -474,7 +475,7 @@ private TaskStatusPlus createTaskStatus(String id, String datasource, Long durat new BrokerSegmentWatcherConfig(), BrokerSegmentMetadataCacheConfig.create() ), - new TestTimelineServerView(walker.getSegments()), + timelineServerView, new FakeServerInventoryView(), authorizerMapper, druidLeaderClient, diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/util/QueryFrameworkUtils.java b/sql/src/test/java/org/apache/druid/sql/calcite/util/QueryFrameworkUtils.java index f28461e1711c..f61f52343987 100644 --- a/sql/src/test/java/org/apache/druid/sql/calcite/util/QueryFrameworkUtils.java +++ b/sql/src/test/java/org/apache/druid/sql/calcite/util/QueryFrameworkUtils.java @@ -26,6 +26,7 @@ import org.apache.calcite.jdbc.CalciteSchema; import org.apache.calcite.schema.SchemaPlus; import org.apache.druid.client.InternalQueryConfig; +import org.apache.druid.client.TimelineServerView; import org.apache.druid.java.util.emitter.core.NoopEmitter; import org.apache.druid.java.util.emitter.service.ServiceEmitter; import org.apache.druid.query.DefaultGenericQueryMetricsFactory; @@ -125,7 +126,32 @@ public static DruidSchemaCatalog createMockRootSchema( @Nullable final ViewManager viewManager, final DruidSchemaManager druidSchemaManager, final AuthorizerMapper authorizerMapper, - final CatalogResolver catalogResolver + final CatalogResolver catalogResolver) + { + TimelineServerView timelineServerView = new TestTimelineServerView(walker.getSegments()); + return createMockRootSchema( + injector, + conglomerate, + walker, + plannerConfig, + viewManager, + druidSchemaManager, + authorizerMapper, + catalogResolver, + timelineServerView + ); + } + + public static DruidSchemaCatalog createMockRootSchema( + final Injector injector, + final QueryRunnerFactoryConglomerate conglomerate, + final SpecificSegmentsQuerySegmentWalker walker, + final PlannerConfig plannerConfig, + @Nullable final ViewManager viewManager, + final DruidSchemaManager druidSchemaManager, + final AuthorizerMapper authorizerMapper, + final CatalogResolver catalogResolver, + final TimelineServerView timelineServerView ) { DruidSchema druidSchema = createMockSchema( @@ -133,10 +159,11 @@ public static DruidSchemaCatalog createMockRootSchema( conglomerate, walker, druidSchemaManager, - catalogResolver + catalogResolver, + timelineServerView ); SystemSchema systemSchema = - CalciteTests.createMockSystemSchema(druidSchema, walker, authorizerMapper); + CalciteTests.createMockSystemSchema(druidSchema, timelineServerView, authorizerMapper); LookupSchema lookupSchema = createMockLookupSchema(injector); DruidOperatorTable createOperatorTable = createOperatorTable(injector); @@ -221,12 +248,13 @@ static DruidSchema createMockSchema( final QueryRunnerFactoryConglomerate conglomerate, final SpecificSegmentsQuerySegmentWalker walker, final DruidSchemaManager druidSchemaManager, - final CatalogResolver catalog + final CatalogResolver catalog, + final TimelineServerView timelineServerView ) { final BrokerSegmentMetadataCache cache = new BrokerSegmentMetadataCache( createMockQueryLifecycleFactory(walker, conglomerate), - new TestTimelineServerView(walker.getSegments()), + timelineServerView, BrokerSegmentMetadataCacheConfig.create(), CalciteTests.TEST_AUTHENTICATOR_ESCALATOR, new InternalQueryConfig(), diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/util/SqlTestFramework.java b/sql/src/test/java/org/apache/druid/sql/calcite/util/SqlTestFramework.java index 90cc53d1b9a0..ebe2e63d1326 100644 --- a/sql/src/test/java/org/apache/druid/sql/calcite/util/SqlTestFramework.java +++ b/sql/src/test/java/org/apache/druid/sql/calcite/util/SqlTestFramework.java @@ -29,8 +29,11 @@ import com.google.inject.Module; import com.google.inject.Provides; import com.google.inject.TypeLiteral; +import org.apache.druid.client.TestHttpClient; +import org.apache.druid.client.TimelineServerView; import org.apache.druid.client.cache.Cache; import org.apache.druid.client.cache.CacheConfig; +import org.apache.druid.collections.BlockingPool; import org.apache.druid.collections.NonBlockingPool; import org.apache.druid.guice.BuiltInTypesModule; import org.apache.druid.guice.DruidInjectorBuilder; @@ -44,12 +47,14 @@ import org.apache.druid.guice.StartupInjectorBuilder; import org.apache.druid.guice.annotations.Global; import org.apache.druid.guice.annotations.Merging; +import org.apache.druid.guice.annotations.Self; import org.apache.druid.initialization.CoreInjectorBuilder; import org.apache.druid.initialization.DruidModule; import org.apache.druid.initialization.ServiceInjectorBuilder; import org.apache.druid.java.util.common.RE; import org.apache.druid.java.util.common.io.Closer; import org.apache.druid.java.util.emitter.service.ServiceEmitter; +import org.apache.druid.java.util.http.client.HttpClient; import org.apache.druid.math.expr.ExprMacroTable; import org.apache.druid.query.DefaultGenericQueryMetricsFactory; import org.apache.druid.query.DruidProcessingConfig; @@ -78,6 +83,7 @@ import org.apache.druid.segment.realtime.ChatHandlerProvider; import org.apache.druid.segment.realtime.NoopChatHandlerProvider; import org.apache.druid.server.ClientQuerySegmentWalker; +import org.apache.druid.server.DruidNode; import org.apache.druid.server.LocalQuerySegmentWalker; import org.apache.druid.server.QueryLifecycle; import org.apache.druid.server.QueryLifecycleFactory; @@ -672,7 +678,8 @@ public PlannerFixture( viewManager, componentSupplier.createSchemaManager(), framework.authorizerMapper, - framework.builder.catalogResolver + framework.builder.catalogResolver, + framework.injector.getInstance(TimelineServerView.class) ); this.plannerFactory = new PlannerFactory( @@ -771,10 +778,22 @@ public void configure(Binder binder) .annotatedWith(Global.class) .to(TestBufferPool.class); + binder.bind(new TypeLiteral>(){}) + .annotatedWith(Merging.class) + .to(TestBufferPool.class); + TestRequestLogger testRequestLogger = new TestRequestLogger(); binder.bind(RequestLogger.class).toInstance(testRequestLogger); } + @Provides + @Self + @LazySingleton + public DruidNode makeSelfDruidNode() + { + return new DruidNode("test", "test-host", true, 80, 443, false, true); + } + @Provides AuthorizerMapper getAuthorizerMapper() { @@ -923,26 +942,36 @@ private DruidSchema makeDruidSchema( final Injector injector, QueryRunnerFactoryConglomerate conglomerate, QuerySegmentWalker walker, - Builder builder) + Builder builder, + TimelineServerView timelineServerView) { return QueryFrameworkUtils.createMockSchema( injector, conglomerate, (SpecificSegmentsQuerySegmentWalker) walker, builder.componentSupplier.getPlannerComponentSupplier().createSchemaManager(), - builder.catalogResolver + builder.catalogResolver, + timelineServerView ); } @Provides @LazySingleton - private SystemSchema makeSystemSchema(QuerySegmentWalker walker, AuthorizerMapper authorizerMapper, - DruidSchema druidSchema) + private SystemSchema makeSystemSchema(AuthorizerMapper authorizerMapper, DruidSchema druidSchema, + TimelineServerView timelineServerView) + { + return CalciteTests.createMockSystemSchema(druidSchema, timelineServerView, authorizerMapper); + } + + + @Provides + @LazySingleton + private TimelineServerView makeTimelineServerView(SpecificSegmentsQuerySegmentWalker walker) { - return CalciteTests - .createMockSystemSchema(druidSchema, (SpecificSegmentsQuerySegmentWalker) walker, authorizerMapper); + return new TestTimelineServerView(walker.getSegments()); } + @Provides @LazySingleton private ColumnConfig getColumnConfig() @@ -1039,6 +1068,13 @@ public TestSegmentsBroker makeTimelines() return new TestSegmentsBroker(); } + @Provides + @LazySingleton + private HttpClient makeHttpClient(ObjectMapper objectMapper) + { + return new TestHttpClient(objectMapper); + } + @Provides @Named("empty") @LazySingleton From f27c023463e4523377fb036fc2c1d0665bee4b45 Mon Sep 17 00:00:00 2001 From: Zoltan Haindrich Date: Mon, 10 Mar 2025 14:58:44 +0000 Subject: [PATCH 48/50] set host --- .../org/apache/druid/msq/dart/worker/DartWorkerClientImpl.java | 2 ++ .../msq/test/{SimpleDartTest.java => CalciteDartTest.java} | 2 +- .../org/apache/druid/sql/calcite/util/SqlTestFramework.java | 2 +- 3 files changed, 4 insertions(+), 2 deletions(-) rename extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/{SimpleDartTest.java => CalciteDartTest.java} (98%) diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartWorkerClientImpl.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartWorkerClientImpl.java index a75c15533eca..adfebee111c9 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartWorkerClientImpl.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartWorkerClientImpl.java @@ -113,6 +113,7 @@ protected ServiceClient getClient(final String workerIdString) /** * {@inheritDoc} */ + @Override public void closeClient(final String workerHost) { synchronized (clientMap) { @@ -144,6 +145,7 @@ public void close() /** * {@inheritDoc} */ + @Override public ListenableFuture stopWorker(String workerId) { return getClient(workerId).asyncRequest( diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/SimpleDartTest.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/CalciteDartTest.java similarity index 98% rename from extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/SimpleDartTest.java rename to extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/CalciteDartTest.java index a81798c3bee6..f37162eb7aa7 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/SimpleDartTest.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/CalciteDartTest.java @@ -31,7 +31,7 @@ import java.util.UUID; @SqlTestFrameworkConfig.ComponentSupplier(DartComponentSupplier.class) -public class SimpleDartTest extends BaseCalciteQueryTest +public class CalciteDartTest extends BaseCalciteQueryTest { @Override protected QueryTestBuilder testBuilder() diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/util/SqlTestFramework.java b/sql/src/test/java/org/apache/druid/sql/calcite/util/SqlTestFramework.java index ebe2e63d1326..d88af75df644 100644 --- a/sql/src/test/java/org/apache/druid/sql/calcite/util/SqlTestFramework.java +++ b/sql/src/test/java/org/apache/druid/sql/calcite/util/SqlTestFramework.java @@ -791,7 +791,7 @@ public void configure(Binder binder) @LazySingleton public DruidNode makeSelfDruidNode() { - return new DruidNode("test", "test-host", true, 80, 443, false, true); + return new DruidNode("druid/broker", "local-test-host", false, 12345, 443, true, false); } @Provides From 26857dca06649798752eb09b306c2f2de4f31690 Mon Sep 17 00:00:00 2001 From: Zoltan Haindrich Date: Tue, 11 Mar 2025 08:26:44 +0000 Subject: [PATCH 49/50] override --- .../druid/msq/test/TestDartControllerContextFactoryImpl.java | 1 + 1 file changed, 1 insertion(+) diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/TestDartControllerContextFactoryImpl.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/TestDartControllerContextFactoryImpl.java index 1c25537f0de3..84e34c2c77e7 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/TestDartControllerContextFactoryImpl.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/TestDartControllerContextFactoryImpl.java @@ -83,6 +83,7 @@ public ControllerContext newContext(String queryId) emitter ) { + @Override public void registerController(Controller currentController, Closer closer) { super.registerController(currentController, closer); From 6faff9961da8953d0d47c941f7382418e7670c94 Mon Sep 17 00:00:00 2001 From: Zoltan Haindrich Date: Thu, 13 Mar 2025 14:50:17 +0000 Subject: [PATCH 50/50] cleanup --- .../msq/test/TestDartControllerContextFactoryImpl.java | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/TestDartControllerContextFactoryImpl.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/TestDartControllerContextFactoryImpl.java index 84e34c2c77e7..89209f012529 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/TestDartControllerContextFactoryImpl.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/TestDartControllerContextFactoryImpl.java @@ -77,7 +77,7 @@ public ControllerContext newContext(String queryId) injector, jsonMapper, selfNode, - new DartTestWorkerClient(queryId, serviceClientFactory, smileMapper, selfNode.getHostAndPortToUse(), workerMap), + new DartTestWorkerClient(), memoryIntrospector, serverView, emitter @@ -96,11 +96,7 @@ public class DartTestWorkerClient extends MSQTestWorkerClient implements DartWor { private final ExecutorService EXECUTOR = Executors.newCachedThreadPool(); - public DartTestWorkerClient( - String queryId, - ServiceClientFactory clientFactory, - ObjectMapper smileMapper, - String controllerHost, Map workerMap) + public DartTestWorkerClient() { super(workerMap); }