From 4ef900dff56def420fd4d892a984b68fde77c64e Mon Sep 17 00:00:00 2001 From: Soumyava Das Date: Tue, 14 Nov 2023 08:36:07 -0800 Subject: [PATCH 01/42] Initial code --- .../apache/druid/msq/exec/ControllerImpl.java | 3 + .../msq/querykit/BaseLeafFrameProcessor.java | 5 +- .../WindowOperatorQueryFrameProcessor.java | 107 ++++++++++++++++++ ...dowOperatorQueryFrameProcessorFactory.java | 54 +++++++++ .../msq/querykit/WindowOperatorQueryKit.java | 64 +++++++++++ .../GroupByPostShuffleFrameProcessor.java | 27 +++-- .../scan/ScanQueryFrameProcessor.java | 5 + .../druid/msq/sql/MSQTaskSqlEngine.java | 2 +- .../apache/druid/msq/exec/MSQSelectTest.java | 71 +++++++++++- .../apache/druid/msq/test/MSQTestBase.java | 1 + .../druid/sql/calcite/rule/DruidJoinRule.java | 1 + .../calcite/tests/window/no_grouping2.sqlTest | 14 +-- 12 files changed, 328 insertions(+), 26 deletions(-) create mode 100644 extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryFrameProcessor.java create mode 100644 extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryFrameProcessorFactory.java create mode 100644 extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryKit.java 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 1ae3cdc5a6e1..fae4bebe4bf5 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 @@ -164,6 +164,7 @@ import org.apache.druid.msq.querykit.QueryKitUtils; import org.apache.druid.msq.querykit.ShuffleSpecFactories; import org.apache.druid.msq.querykit.ShuffleSpecFactory; +import org.apache.druid.msq.querykit.WindowOperatorQueryKit; import org.apache.druid.msq.querykit.groupby.GroupByQueryKit; import org.apache.druid.msq.querykit.results.QueryResultFrameProcessorFactory; import org.apache.druid.msq.querykit.scan.ScanQueryKit; @@ -182,6 +183,7 @@ import org.apache.druid.query.aggregation.AggregatorFactory; import org.apache.druid.query.groupby.GroupByQuery; import org.apache.druid.query.groupby.GroupByQueryConfig; +import org.apache.druid.query.operator.WindowOperatorQuery; import org.apache.druid.query.scan.ScanQuery; import org.apache.druid.segment.DimensionHandlerUtils; import org.apache.druid.segment.column.ColumnHolder; @@ -1169,6 +1171,7 @@ private QueryKit makeQueryControllerToolKit() ImmutableMap., QueryKit>builder() .put(ScanQuery.class, new ScanQueryKit(context.jsonMapper())) .put(GroupByQuery.class, new GroupByQueryKit(context.jsonMapper())) + .put(WindowOperatorQuery.class, new WindowOperatorQueryKit(context.jsonMapper())) .build(); return new MultiQueryKit(kitMap); diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/BaseLeafFrameProcessor.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/BaseLeafFrameProcessor.java index f67f30d0c5c6..d3dc01d80f83 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/BaseLeafFrameProcessor.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/BaseLeafFrameProcessor.java @@ -40,7 +40,7 @@ import java.util.List; import java.util.function.Function; -public abstract class BaseLeafFrameProcessor implements FrameProcessor +public abstract class BaseLeafFrameProcessor implements FrameProcessor { private final ReadableInput baseInput; private final ResourceHolder outputChannelHolder; @@ -63,6 +63,7 @@ protected BaseLeafFrameProcessor( @Override public List inputChannels() { + // somu: need to clarify if the data is in broker only if (baseInput.hasSegment()) { return Collections.emptyList(); } else { @@ -82,6 +83,7 @@ public ReturnOrAwait runIncrementally(final IntSet readableInputs) throw final ReturnOrAwait retVal; if (baseInput.hasSegment()) { + // run with data in historicals SegmentWithDescriptor segment = baseInput.getSegment(); if (segment.getDescriptor().isLoadedOnServer()) { retVal = runWithLoadedSegment(baseInput.getSegment()); @@ -89,6 +91,7 @@ public ReturnOrAwait runIncrementally(final IntSet readableInputs) throw retVal = runWithSegment(baseInput.getSegment()); } } else { + // run with data in broker retVal = runWithInputChannel(baseInput.getChannel(), baseInput.getChannelFrameReader()); } diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryFrameProcessor.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryFrameProcessor.java new file mode 100644 index 000000000000..b1315ee34ab5 --- /dev/null +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryFrameProcessor.java @@ -0,0 +1,107 @@ +package org.apache.druid.msq.querykit; + +import com.fasterxml.jackson.databind.ObjectMapper; +import org.apache.druid.collections.ResourceHolder; +import org.apache.druid.frame.channel.ReadableFrameChannel; +import org.apache.druid.frame.channel.WritableFrameChannel; +import org.apache.druid.frame.processor.ReturnOrAwait; +import org.apache.druid.frame.read.FrameReader; +import org.apache.druid.frame.util.SettableLongVirtualColumn; +import org.apache.druid.frame.write.FrameWriter; +import org.apache.druid.frame.write.FrameWriterFactory; +import org.apache.druid.java.util.common.Unit; +import org.apache.druid.java.util.common.io.Closer; +import org.apache.druid.java.util.common.logger.Logger; +import org.apache.druid.msq.input.ReadableInput; +import org.apache.druid.msq.input.table.SegmentWithDescriptor; +import org.apache.druid.msq.querykit.scan.ScanQueryFrameProcessor; +import org.apache.druid.query.operator.WindowOperatorQuery; +import org.apache.druid.query.scan.ScanQuery; +import org.apache.druid.segment.Cursor; +import org.apache.druid.segment.Segment; +import org.apache.druid.segment.SegmentReference; +import org.apache.druid.segment.SimpleAscendingOffset; +import org.apache.druid.segment.SimpleSettableOffset; +import org.apache.druid.segment.VirtualColumn; +import org.apache.druid.segment.VirtualColumns; + +import javax.annotation.Nullable; +import java.io.IOException; +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.atomic.AtomicLong; +import java.util.function.Function; + +public class WindowOperatorQueryFrameProcessor extends BaseLeafFrameProcessor +{ + + private static final Logger log = new Logger(ScanQueryFrameProcessor.class); + private final WindowOperatorQuery query; + private final ObjectMapper jsonMapper; + private final SettableLongVirtualColumn partitionBoostVirtualColumn; + private final VirtualColumns frameWriterVirtualColumns; + private final Closer closer = Closer.create(); + + private Cursor cursor; + private Segment segment; + private final SimpleSettableOffset cursorOffset = new SimpleAscendingOffset(Integer.MAX_VALUE); + private FrameWriter frameWriter; + private long currentAllocatorCapacity; // Used for generating FrameRowTooLargeException if needed + + public WindowOperatorQueryFrameProcessor( + final WindowOperatorQuery query, + final ObjectMapper jsonMapper, + final ReadableInput baseInput, + final Function segmentMapFn, + final ResourceHolder outputChannelHolder, + final ResourceHolder frameWriterFactoryHolder + ) + { + super( + baseInput, + segmentMapFn, + outputChannelHolder, + frameWriterFactoryHolder + ); + this.query = query; + this.jsonMapper = jsonMapper; + this.partitionBoostVirtualColumn = new SettableLongVirtualColumn(QueryKitUtils.PARTITION_BOOST_COLUMN); + + final List frameWriterVirtualColumns = new ArrayList<>(); + frameWriterVirtualColumns.add(partitionBoostVirtualColumn); + + final VirtualColumn segmentGranularityVirtualColumn = + QueryKitUtils.makeSegmentGranularityVirtualColumn(jsonMapper, query); + + if (segmentGranularityVirtualColumn != null) { + frameWriterVirtualColumns.add(segmentGranularityVirtualColumn); + } + + this.frameWriterVirtualColumns = VirtualColumns.create(frameWriterVirtualColumns); + } + // deep storage + @Override + protected ReturnOrAwait runWithSegment(SegmentWithDescriptor segment) throws IOException + { + return null; + } + + // realtime + @Override + protected ReturnOrAwait runWithLoadedSegment(SegmentWithDescriptor segment) throws IOException + { + return null; + } + + // previous stage output + @Override + protected ReturnOrAwait runWithInputChannel(ReadableFrameChannel inputChannel, FrameReader inputFrameReader) + throws IOException + { + // Read the frames from the broker + // convert to FrameRowsAndColumns + // Find a way to call the operators on top of the frames + // Find a way to flush the output to the outputChannel + return null; + } +} diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryFrameProcessorFactory.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryFrameProcessorFactory.java new file mode 100644 index 000000000000..71267185e28d --- /dev/null +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryFrameProcessorFactory.java @@ -0,0 +1,54 @@ +package org.apache.druid.msq.querykit; + +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonProperty; +import com.fasterxml.jackson.annotation.JsonTypeName; +import com.google.common.base.Preconditions; +import org.apache.druid.collections.ResourceHolder; +import org.apache.druid.frame.channel.WritableFrameChannel; +import org.apache.druid.frame.processor.FrameProcessor; +import org.apache.druid.frame.write.FrameWriterFactory; +import org.apache.druid.msq.input.ReadableInput; +import org.apache.druid.msq.kernel.FrameContext; +import org.apache.druid.query.operator.WindowOperatorQuery; +import org.apache.druid.segment.SegmentReference; + +import java.util.function.Function; + +@JsonTypeName("window") +public class WindowOperatorQueryFrameProcessorFactory extends BaseLeafFrameProcessorFactory +{ + private final WindowOperatorQuery query; + + @JsonCreator + public WindowOperatorQueryFrameProcessorFactory(@JsonProperty("query") WindowOperatorQuery query) + { + super(query); + this.query = Preconditions.checkNotNull(query, "query"); + } + + @JsonProperty + public WindowOperatorQuery getQuery() + { + return query; + } + + @Override + protected FrameProcessor makeProcessor( + ReadableInput baseInput, + Function segmentMapFn, + ResourceHolder outputChannelHolder, + ResourceHolder frameWriterFactoryHolder, + FrameContext frameContext + ) + { + return new WindowOperatorQueryFrameProcessor( + query, + frameContext.jsonMapper(), + baseInput, + segmentMapFn, + outputChannelHolder, + frameWriterFactoryHolder + ); + } +} diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryKit.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryKit.java new file mode 100644 index 000000000000..347fb0a54dad --- /dev/null +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryKit.java @@ -0,0 +1,64 @@ +package org.apache.druid.msq.querykit; + +import com.fasterxml.jackson.databind.ObjectMapper; +import org.apache.druid.msq.input.stage.StageInputSpec; +import org.apache.druid.msq.kernel.QueryDefinition; +import org.apache.druid.msq.kernel.QueryDefinitionBuilder; +import org.apache.druid.msq.kernel.StageDefinition; +import org.apache.druid.query.Query; +import org.apache.druid.query.operator.WindowOperatorQuery; +import org.apache.druid.segment.column.RowSignature; + +public class WindowOperatorQueryKit implements QueryKit +{ + private final ObjectMapper jsonMapper; + + public WindowOperatorQueryKit(ObjectMapper jsonMapper) + { + this.jsonMapper = jsonMapper; + } + @Override + public QueryDefinition makeQueryDefinition( + String queryId, + WindowOperatorQuery originalQuery, + QueryKit> queryKit, + ShuffleSpecFactory resultShuffleSpecFactory, + int maxWorkerCount, + int minStageNumber + ) + { + // need to validate query first + + final QueryDefinitionBuilder queryDefBuilder = QueryDefinition.builder().queryId(queryId); + final DataSourcePlan dataSourcePlan = DataSourcePlan.forDataSource( + queryKit, + queryId, + originalQuery.context(), + originalQuery.getDataSource(), + originalQuery.getQuerySegmentSpec(), + originalQuery.getFilter(), + null, + maxWorkerCount, + minStageNumber, + false + ); + + + dataSourcePlan.getSubQueryDefBuilder().ifPresent(queryDefBuilder::addAll); + + final int firstStageNumber = Math.max(minStageNumber, queryDefBuilder.getNextStageNumber()); + final WindowOperatorQuery queryToRun = (WindowOperatorQuery) originalQuery.withDataSource(dataSourcePlan.getNewDataSource()); + RowSignature rowSignature = queryToRun.getRowSignature(); + + // Create a new stage which takes in the subquery as an input + queryDefBuilder.add( + StageDefinition.builder(firstStageNumber) + .inputs(new StageInputSpec(firstStageNumber-1)) + .signature(rowSignature) + .maxWorkerCount(maxWorkerCount) + .shuffleSpec(null) + .processorFactory(new WindowOperatorQueryFrameProcessorFactory(queryToRun)) + ); + return queryDefBuilder.queryId(queryId).build(); + } +} diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/groupby/GroupByPostShuffleFrameProcessor.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/groupby/GroupByPostShuffleFrameProcessor.java index ed14dd52739c..3f2abd5c767e 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/groupby/GroupByPostShuffleFrameProcessor.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/groupby/GroupByPostShuffleFrameProcessor.java @@ -149,18 +149,23 @@ public ReturnOrAwait runIncrementally(final IntSet readableInputs) throw } final int fullRowSize = query.getResultRowSignature().size(); - rowSupplierFromFrameCursor = () -> { - final ResultRow row = ResultRow.create(fullRowSize); - for (int i = 0; i < fieldSuppliers.length; i++) { - row.set(i, fieldSuppliers[i].get()); + rowSupplierFromFrameCursor = new Supplier() + { + @Override + public ResultRow get() + { + final ResultRow row = ResultRow.create(fullRowSize); + for (int i = 0; i < fieldSuppliers.length; i++) { + row.set(i, fieldSuppliers[i].get()); + } + + for (int i = fieldSuppliers.length; i < fullRowSize; i++) { + // Post-aggregators. + row.set(i, null); + } + + return row; } - - for (int i = fieldSuppliers.length; i < fullRowSize; i++) { - // Post-aggregators. - row.set(i, null); - } - - return row; }; } } diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/scan/ScanQueryFrameProcessor.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/scan/ScanQueryFrameProcessor.java index ff15e116df6f..ede4856bc3f2 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/scan/ScanQueryFrameProcessor.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/scan/ScanQueryFrameProcessor.java @@ -194,6 +194,9 @@ private static ScanQuery prepareScanQueryForDataServer(@NotNull ScanQuery scanQu @Override protected ReturnOrAwait runWithLoadedSegment(final SegmentWithDescriptor segment) throws IOException { + // Try to run with the segment on the server + // if segment was not found for some reason + // call runWithSegment which would load from deep storage if (cursor == null) { ScanQuery preparedQuery = prepareScanQueryForDataServer(query); final Pair> statusSequencePair = @@ -244,7 +247,9 @@ protected ReturnOrAwait runWithLoadedSegment(final SegmentWithDescriptor s @Override protected ReturnOrAwait runWithSegment(final SegmentWithDescriptor segment) throws IOException { + // why not remove order by and limit here ? if (cursor == null) { + // load from deep storage final ResourceHolder segmentHolder = closer.register(segment.getOrLoad()); final Yielder cursorYielder = Yielders.each( diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/MSQTaskSqlEngine.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/MSQTaskSqlEngine.java index cb331760ca34..4e7d9023e252 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/MSQTaskSqlEngine.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/MSQTaskSqlEngine.java @@ -113,9 +113,9 @@ public boolean featureAvailable(EngineFeature feature, PlannerContext plannerCon case TOPN_QUERY: case TIME_BOUNDARY_QUERY: case GROUPING_SETS: - case WINDOW_FUNCTIONS: case ALLOW_TOP_LEVEL_UNION_ALL: return false; + case WINDOW_FUNCTIONS: case UNNEST: case CAN_SELECT: case CAN_INSERT: diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQSelectTest.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQSelectTest.java index e1c671a86650..c587c7d63e52 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQSelectTest.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQSelectTest.java @@ -130,12 +130,12 @@ public class MSQSelectTest extends MSQTestBase public static Collection data() { Object[][] data = new Object[][]{ - {DEFAULT, DEFAULT_MSQ_CONTEXT}, - {DURABLE_STORAGE, DURABLE_STORAGE_MSQ_CONTEXT}, - {FAULT_TOLERANCE, FAULT_TOLERANCE_MSQ_CONTEXT}, - {PARALLEL_MERGE, PARALLEL_MERGE_MSQ_CONTEXT}, - {QUERY_RESULTS_WITH_DURABLE_STORAGE, QUERY_RESULTS_WITH_DURABLE_STORAGE_CONTEXT}, - {QUERY_RESULTS_WITH_DEFAULT, QUERY_RESULTS_WITH_DEFAULT_CONTEXT} + {DEFAULT, DEFAULT_MSQ_CONTEXT} +// {DURABLE_STORAGE, DURABLE_STORAGE_MSQ_CONTEXT}, +// {FAULT_TOLERANCE, FAULT_TOLERANCE_MSQ_CONTEXT}, +// {PARALLEL_MERGE, PARALLEL_MERGE_MSQ_CONTEXT}, +// {QUERY_RESULTS_WITH_DURABLE_STORAGE, QUERY_RESULTS_WITH_DURABLE_STORAGE_CONTEXT}, +// {QUERY_RESULTS_WITH_DEFAULT, QUERY_RESULTS_WITH_DEFAULT_CONTEXT} }; return Arrays.asList(data); @@ -613,6 +613,65 @@ public void testGroupByOrderByDimension() .verifyResults(); } + @Test + public void testWindowOnFoo() + { + RowSignature rowSignature = RowSignature.builder() + .add("cnt", ColumnType.LONG) + .add("cnt1", ColumnType.LONG) + .build(); + + testSelectQuery() + .setSql("select m1,SUM(m1) OVER() cc from foo group by m1") + .setExpectedMSQSpec(MSQSpec.builder() + .query(GroupByQuery.builder() + .setDataSource(CalciteTests.DATASOURCE1) + .setInterval(querySegmentSpec(Filtration + .eternity())) + .setGranularity(Granularities.ALL) + .setDimensions(dimensions( + new DefaultDimensionSpec( + "cnt", + "d0", + ColumnType.LONG + ) + )) + .setAggregatorSpecs(aggregators(new CountAggregatorFactory( + "a0"))) + .setContext(context) + .build()) + .columnMappings( + new ColumnMappings(ImmutableList.of( + new ColumnMapping("d0", "cnt"), + new ColumnMapping("a0", "cnt1") + ) + )) + .tuningConfig(MSQTuningConfig.defaultConfig()) + .destination(isDurableStorageDestination() + ? DurableStorageMSQDestination.INSTANCE + : TaskReportMSQDestination.INSTANCE) + .build()) + .setExpectedRowSignature(rowSignature) + .setExpectedResultRows(ImmutableList.of(new Object[]{1L, 6L})) + .setQueryContext(context) + .setExpectedCountersForStageWorkerChannel( + CounterSnapshotMatcher + .with().totalFiles(1), + 0, 0, "input0" + ) + .setExpectedCountersForStageWorkerChannel( + CounterSnapshotMatcher + .with().rows(1).frames(1), + 0, 0, "output" + ) + .setExpectedCountersForStageWorkerChannel( + CounterSnapshotMatcher + .with().rows(1).frames(1), + 0, 0, "shuffle" + ) + .verifyResults(); + } + @Test public void testSelectWithLimit() { diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestBase.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestBase.java index ba78e213ca47..cbb325fb99a9 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestBase.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestBase.java @@ -251,6 +251,7 @@ public class MSQTestBase extends BaseCalciteQueryTest .put(MultiStageQueryContext.CTX_MAX_NUM_TASKS, 2) .put(MSQWarnings.CTX_MAX_PARSE_EXCEPTIONS_ALLOWED, 0) .put(MSQTaskQueryMaker.USER_KEY, "allowAll") + .put("enableWindowing", true) .build(); public static final Map DURABLE_STORAGE_MSQ_CONTEXT = diff --git a/sql/src/main/java/org/apache/druid/sql/calcite/rule/DruidJoinRule.java b/sql/src/main/java/org/apache/druid/sql/calcite/rule/DruidJoinRule.java index 6dc8ff00531b..6ca49ea4b95d 100644 --- a/sql/src/main/java/org/apache/druid/sql/calcite/rule/DruidJoinRule.java +++ b/sql/src/main/java/org/apache/druid/sql/calcite/rule/DruidJoinRule.java @@ -156,6 +156,7 @@ public void onMatch(RelOptRuleCall call) final Project rightProject = right.getPartialDruidQuery().getSelectProject(); // Right-side projection expressions rewritten to be on top of the join. + for (final RexNode rexNode : RexUtil.shift(rightProject.getProjects(), newLeft.getRowType().getFieldCount())) { if (join.getJoinType().generatesNullsOnRight()) { newProjectExprs.add(makeNullableIfLiteral(rexNode, rexBuilder)); diff --git a/sql/src/test/resources/calcite/tests/window/no_grouping2.sqlTest b/sql/src/test/resources/calcite/tests/window/no_grouping2.sqlTest index 3b2bef3e965a..7a579c3fc331 100644 --- a/sql/src/test/resources/calcite/tests/window/no_grouping2.sqlTest +++ b/sql/src/test/resources/calcite/tests/window/no_grouping2.sqlTest @@ -2,7 +2,7 @@ type: "operatorValidation" sql: | SELECT - m1, + m2, m1, SUM(m1) OVER () cc FROM druid.foo @@ -18,9 +18,9 @@ expectedOperators: name: "w0" fieldName: "m1" expectedResults: - - [1.0,21.0] - - [2.0,21.0] - - [3.0,21.0] - - [4.0,21.0] - - [5.0,21.0] - - [6.0,21.0] + - [1.0, 1.0, 21.0] + - [2.0, 2.0, 21.0] + - [3.0, 3.0, 21.0] + - [4.0,4.0,21.0] + - [5.0,5.0,21.0] + - [6.0,6.0,21.0] From 9ea01fc0c83ebdaae00f988a1306f794a3315e87 Mon Sep 17 00:00:00 2001 From: Soumyava Das Date: Thu, 30 Nov 2023 14:35:23 -0800 Subject: [PATCH 02/42] Hacky way of atleast getting things to work --- .../msq/querykit/BaseLeafFrameProcessor.java | 2 - .../WindowOperatorQueryFrameProcessor.java | 67 ++++++++++++++++++- .../concrete/FrameRowsAndColumns.java | 1 + .../concrete/RowBasedFrameRowAndColumns.java | 57 ++++++++++++++++ 4 files changed, 122 insertions(+), 5 deletions(-) create mode 100644 processing/src/main/java/org/apache/druid/query/rowsandcols/concrete/RowBasedFrameRowAndColumns.java diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/BaseLeafFrameProcessor.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/BaseLeafFrameProcessor.java index d3dc01d80f83..1cb7af037529 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/BaseLeafFrameProcessor.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/BaseLeafFrameProcessor.java @@ -83,7 +83,6 @@ public ReturnOrAwait runIncrementally(final IntSet readableInputs) throw final ReturnOrAwait retVal; if (baseInput.hasSegment()) { - // run with data in historicals SegmentWithDescriptor segment = baseInput.getSegment(); if (segment.getDescriptor().isLoadedOnServer()) { retVal = runWithLoadedSegment(baseInput.getSegment()); @@ -91,7 +90,6 @@ public ReturnOrAwait runIncrementally(final IntSet readableInputs) throw retVal = runWithSegment(baseInput.getSegment()); } } else { - // run with data in broker retVal = runWithInputChannel(baseInput.getChannel(), baseInput.getChannelFrameReader()); } diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryFrameProcessor.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryFrameProcessor.java index b1315ee34ab5..a8f6bd6c1629 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryFrameProcessor.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryFrameProcessor.java @@ -2,10 +2,12 @@ import com.fasterxml.jackson.databind.ObjectMapper; import org.apache.druid.collections.ResourceHolder; +import org.apache.druid.frame.Frame; import org.apache.druid.frame.channel.ReadableFrameChannel; import org.apache.druid.frame.channel.WritableFrameChannel; import org.apache.druid.frame.processor.ReturnOrAwait; import org.apache.druid.frame.read.FrameReader; +import org.apache.druid.frame.segment.FrameSegment; import org.apache.druid.frame.util.SettableLongVirtualColumn; import org.apache.druid.frame.write.FrameWriter; import org.apache.druid.frame.write.FrameWriterFactory; @@ -15,17 +17,28 @@ import org.apache.druid.msq.input.ReadableInput; import org.apache.druid.msq.input.table.SegmentWithDescriptor; import org.apache.druid.msq.querykit.scan.ScanQueryFrameProcessor; +import org.apache.druid.query.operator.OffsetLimit; +import org.apache.druid.query.operator.Operator; +import org.apache.druid.query.operator.OperatorFactory; +import org.apache.druid.query.operator.SegmentToRowsAndColumnsOperator; import org.apache.druid.query.operator.WindowOperatorQuery; +import org.apache.druid.query.rowsandcols.LazilyDecoratedRowsAndColumns; +import org.apache.druid.query.rowsandcols.RowsAndColumns; +import org.apache.druid.query.rowsandcols.concrete.FrameRowsAndColumns; +import org.apache.druid.query.rowsandcols.concrete.RowBasedFrameRowAndColumns; import org.apache.druid.query.scan.ScanQuery; import org.apache.druid.segment.Cursor; import org.apache.druid.segment.Segment; import org.apache.druid.segment.SegmentReference; import org.apache.druid.segment.SimpleAscendingOffset; import org.apache.druid.segment.SimpleSettableOffset; +import org.apache.druid.segment.StorageAdapter; import org.apache.druid.segment.VirtualColumn; import org.apache.druid.segment.VirtualColumns; +import org.apache.druid.timeline.SegmentId; import javax.annotation.Nullable; +import java.io.Closeable; import java.io.IOException; import java.util.ArrayList; import java.util.List; @@ -98,10 +111,58 @@ protected ReturnOrAwait runWithLoadedSegment(SegmentWithDescriptor segment protected ReturnOrAwait runWithInputChannel(ReadableFrameChannel inputChannel, FrameReader inputFrameReader) throws IOException { - // Read the frames from the broker + // Read the frames from the channel // convert to FrameRowsAndColumns - // Find a way to call the operators on top of the frames - // Find a way to flush the output to the outputChannel + + if (inputChannel.canRead()) { + Frame f = inputChannel.read(); + final FrameSegment frameSegment = new FrameSegment(f, inputFrameReader, SegmentId.dummy("x")); + + // the frame here is row based + // frame rows and columns need columnar. discuss with Eric + // Action item: need to implement a new rows and columns that accept a row-based frame + RowBasedFrameRowAndColumns frameRowsAndColumns = new RowBasedFrameRowAndColumns(f, inputFrameReader.signature()); + LazilyDecoratedRowsAndColumns ldrc = new LazilyDecoratedRowsAndColumns(frameRowsAndColumns, null, null, null, OffsetLimit.limit(Integer.MAX_VALUE), null, null); + // Create an operator on top of the created rows and columns + Operator op = new Operator() + { + @Nullable + @Override + public Closeable goOrContinue(Closeable continuationObject, Receiver receiver) + { + receiver.push(ldrc); + receiver.completed(); + return continuationObject; + } + }; + // + //Operator op = new SegmentToRowsAndColumnsOperator(frameSegment); + // On the operator created above add the operators present in the query that we want to chain + + for ( OperatorFactory of : query.getOperators()) { + op = of.wrap(op); + } + + // Let the operator run + // the results that come in the receiver must be pushed to the outout channel + // need to transform the output rows and columns back to frame + Operator.go(op, new Operator.Receiver() + { + @Override + public Operator.Signal push(RowsAndColumns rac) + { + //outputFrameChannel.output(rac.toFrame()); + return null; + } + + @Override + public void completed() + { + + } + }); + + } return null; } } diff --git a/processing/src/main/java/org/apache/druid/query/rowsandcols/concrete/FrameRowsAndColumns.java b/processing/src/main/java/org/apache/druid/query/rowsandcols/concrete/FrameRowsAndColumns.java index 59a97f12bf51..e448acd8ded0 100644 --- a/processing/src/main/java/org/apache/druid/query/rowsandcols/concrete/FrameRowsAndColumns.java +++ b/processing/src/main/java/org/apache/druid/query/rowsandcols/concrete/FrameRowsAndColumns.java @@ -44,6 +44,7 @@ public class FrameRowsAndColumns implements RowsAndColumns public FrameRowsAndColumns(Frame frame, RowSignature signature) { + //this.frame = frame; this.frame = FrameType.COLUMNAR.ensureType(frame); this.signature = signature; } diff --git a/processing/src/main/java/org/apache/druid/query/rowsandcols/concrete/RowBasedFrameRowAndColumns.java b/processing/src/main/java/org/apache/druid/query/rowsandcols/concrete/RowBasedFrameRowAndColumns.java new file mode 100644 index 000000000000..f8f4e79082cb --- /dev/null +++ b/processing/src/main/java/org/apache/druid/query/rowsandcols/concrete/RowBasedFrameRowAndColumns.java @@ -0,0 +1,57 @@ +package org.apache.druid.query.rowsandcols.concrete; + +import org.apache.druid.frame.Frame; +import org.apache.druid.frame.FrameType; +import org.apache.druid.frame.read.FrameReader; +import org.apache.druid.frame.segment.FrameStorageAdapter; +import org.apache.druid.java.util.common.Intervals; +import org.apache.druid.query.rowsandcols.RowsAndColumns; +import org.apache.druid.query.rowsandcols.column.Column; +import org.apache.druid.segment.StorageAdapter; +import org.apache.druid.segment.column.RowSignature; + +import javax.annotation.Nullable; +import java.util.Collection; +import java.util.LinkedHashMap; + +public class RowBasedFrameRowAndColumns implements RowsAndColumns +{ + private final Frame frame; + private final RowSignature signature; + private final LinkedHashMap colCache = new LinkedHashMap<>(); + + public RowBasedFrameRowAndColumns(Frame frame, RowSignature signature) { + this.frame = FrameType.ROW_BASED.ensureType(frame);; + this.signature = signature; + } + + @Override + public Collection getColumnNames() + { + return signature.getColumnNames(); + } + + @Override + public int numRows() + { + return frame.numRows(); + } + + @Nullable + @Override + public Column findColumn(String name) + { + return null; + } + + @SuppressWarnings("unchecked") + @Nullable + @Override + public T as(Class clazz) + { + if (StorageAdapter.class.equals(clazz)) { + return (T) new FrameStorageAdapter(frame, FrameReader.create(signature), Intervals.ETERNITY); + } + return null; + } +} From 9c4ac749c287fef452937dbd29e63ab9f038526b Mon Sep 17 00:00:00 2001 From: Soumyava Das Date: Fri, 1 Dec 2023 12:05:41 -0800 Subject: [PATCH 03/42] Temp unfinished changes --- .../msq/querykit/WindowOperatorQueryFrameProcessor.java | 8 ++++++-- .../java/org/apache/druid/msq/exec/MSQSelectTest.java | 2 +- 2 files changed, 7 insertions(+), 3 deletions(-) diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryFrameProcessor.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryFrameProcessor.java index a8f6bd6c1629..b251995571a9 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryFrameProcessor.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryFrameProcessor.java @@ -152,7 +152,7 @@ public Closeable goOrContinue(Closeable continuationObject, Receiver receiver) public Operator.Signal push(RowsAndColumns rac) { //outputFrameChannel.output(rac.toFrame()); - return null; + return Operator.Signal.GO; } @Override @@ -162,7 +162,11 @@ public void completed() } }); + } else if (inputChannel.isFinished()) { + return ReturnOrAwait.returnObject(Unit.instance()); + } else { + return ReturnOrAwait.awaitAll(inputChannels().size()); } - return null; + return ReturnOrAwait.runAgain(); } } diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQSelectTest.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQSelectTest.java index c587c7d63e52..0a4bf977a98c 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQSelectTest.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQSelectTest.java @@ -622,7 +622,7 @@ public void testWindowOnFoo() .build(); testSelectQuery() - .setSql("select m1,SUM(m1) OVER() cc from foo group by m1") + .setSql("select m1,SUM(m1) OVER(PARTITION BY m1) cc from foo group by m1") .setExpectedMSQSpec(MSQSpec.builder() .query(GroupByQuery.builder() .setDataSource(CalciteTests.DATASOURCE1) From 54f9ac31724a4be3c644a85e581fd2ea072c3233 Mon Sep 17 00:00:00 2001 From: Soumyava Das Date: Wed, 6 Dec 2023 20:18:01 -0800 Subject: [PATCH 04/42] Converting rac back to frames --- .../WindowOperatorQueryFrameProcessor.java | 48 +++++--- .../apache/druid/msq/exec/MSQSelectTest.java | 19 ++- .../apache/druid/frame/write/FrameWriter.java | 2 +- .../LazilyDecoratedRowsAndColumns.java | 112 ++++++++++++++++++ .../concrete/RowBasedFrameRowAndColumns.java | 1 - 5 files changed, 160 insertions(+), 22 deletions(-) diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryFrameProcessor.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryFrameProcessor.java index b251995571a9..ae154b383824 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryFrameProcessor.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryFrameProcessor.java @@ -1,8 +1,10 @@ package org.apache.druid.msq.querykit; import com.fasterxml.jackson.databind.ObjectMapper; +import com.google.common.collect.Iterables; import org.apache.druid.collections.ResourceHolder; import org.apache.druid.frame.Frame; +import org.apache.druid.frame.channel.FrameWithPartition; import org.apache.druid.frame.channel.ReadableFrameChannel; import org.apache.druid.frame.channel.WritableFrameChannel; import org.apache.druid.frame.processor.ReturnOrAwait; @@ -11,6 +13,7 @@ import org.apache.druid.frame.util.SettableLongVirtualColumn; import org.apache.druid.frame.write.FrameWriter; import org.apache.druid.frame.write.FrameWriterFactory; +import org.apache.druid.java.util.common.Pair; import org.apache.druid.java.util.common.Unit; import org.apache.druid.java.util.common.io.Closer; import org.apache.druid.java.util.common.logger.Logger; @@ -35,6 +38,7 @@ import org.apache.druid.segment.StorageAdapter; import org.apache.druid.segment.VirtualColumn; import org.apache.druid.segment.VirtualColumns; +import org.apache.druid.segment.column.RowSignature; import org.apache.druid.timeline.SegmentId; import javax.annotation.Nullable; @@ -116,25 +120,14 @@ protected ReturnOrAwait runWithInputChannel(ReadableFrameChannel inputChan if (inputChannel.canRead()) { Frame f = inputChannel.read(); - final FrameSegment frameSegment = new FrameSegment(f, inputFrameReader, SegmentId.dummy("x")); // the frame here is row based // frame rows and columns need columnar. discuss with Eric // Action item: need to implement a new rows and columns that accept a row-based frame + + // Create a frame rows and columns what would RowBasedFrameRowAndColumns frameRowsAndColumns = new RowBasedFrameRowAndColumns(f, inputFrameReader.signature()); - LazilyDecoratedRowsAndColumns ldrc = new LazilyDecoratedRowsAndColumns(frameRowsAndColumns, null, null, null, OffsetLimit.limit(Integer.MAX_VALUE), null, null); - // Create an operator on top of the created rows and columns - Operator op = new Operator() - { - @Nullable - @Override - public Closeable goOrContinue(Closeable continuationObject, Receiver receiver) - { - receiver.push(ldrc); - receiver.completed(); - return continuationObject; - } - }; + Operator op = getOperator(frameRowsAndColumns); // //Operator op = new SegmentToRowsAndColumnsOperator(frameSegment); // On the operator created above add the operators present in the query that we want to chain @@ -152,6 +145,15 @@ public Closeable goOrContinue(Closeable continuationObject, Receiver receiver) public Operator.Signal push(RowsAndColumns rac) { //outputFrameChannel.output(rac.toFrame()); + LazilyDecoratedRowsAndColumns tmpldrc = new LazilyDecoratedRowsAndColumns(rac, null, null, null, OffsetLimit.limit(Integer.MAX_VALUE), null, null); + Pair pairFrames = tmpldrc.naiveMaterializeToRowFrames(rac); + Frame f = Frame.wrap(pairFrames.lhs); + try { + Iterables.getOnlyElement(outputChannels()).write(new FrameWithPartition(f, FrameWithPartition.NO_PARTITION)); + } + catch (IOException e) { + throw new RuntimeException(e); + } return Operator.Signal.GO; } @@ -169,4 +171,22 @@ public void completed() } return ReturnOrAwait.runAgain(); } + + private static Operator getOperator(RowBasedFrameRowAndColumns frameRowsAndColumns) + { + LazilyDecoratedRowsAndColumns ldrc = new LazilyDecoratedRowsAndColumns(frameRowsAndColumns, null, null, null, OffsetLimit.limit(Integer.MAX_VALUE), null, null); + // Create an operator on top of the created rows and columns + Operator op = new Operator() + { + @Nullable + @Override + public Closeable goOrContinue(Closeable continuationObject, Receiver receiver) + { + receiver.push(ldrc); + receiver.completed(); + return continuationObject; + } + }; + return op; + } } diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQSelectTest.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQSelectTest.java index 0a4bf977a98c..ba2a4b66c239 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQSelectTest.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQSelectTest.java @@ -617,12 +617,12 @@ public void testGroupByOrderByDimension() public void testWindowOnFoo() { RowSignature rowSignature = RowSignature.builder() - .add("cnt", ColumnType.LONG) - .add("cnt1", ColumnType.LONG) + .add("m1", ColumnType.FLOAT) + .add("cc", ColumnType.DOUBLE) .build(); testSelectQuery() - .setSql("select m1,SUM(m1) OVER(PARTITION BY m1) cc from foo group by m1") + .setSql("select m1,SUM(m1) OVER() cc from foo group by m1") .setExpectedMSQSpec(MSQSpec.builder() .query(GroupByQuery.builder() .setDataSource(CalciteTests.DATASOURCE1) @@ -652,7 +652,14 @@ public void testWindowOnFoo() : TaskReportMSQDestination.INSTANCE) .build()) .setExpectedRowSignature(rowSignature) - .setExpectedResultRows(ImmutableList.of(new Object[]{1L, 6L})) + .setExpectedResultRows(ImmutableList.of( + new Object[]{1.0f, 1.0}, + new Object[]{2.0f, 2.0}, + new Object[]{3.0f, 3.0}, + new Object[]{4.0f, 4.0}, + new Object[]{5.0f, 5.0}, + new Object[]{6.0f, 6.0} + )) .setQueryContext(context) .setExpectedCountersForStageWorkerChannel( CounterSnapshotMatcher @@ -661,12 +668,12 @@ public void testWindowOnFoo() ) .setExpectedCountersForStageWorkerChannel( CounterSnapshotMatcher - .with().rows(1).frames(1), + .with().rows(6).frames(1), 0, 0, "output" ) .setExpectedCountersForStageWorkerChannel( CounterSnapshotMatcher - .with().rows(1).frames(1), + .with().rows(6).frames(1), 0, 0, "shuffle" ) .verifyResults(); diff --git a/processing/src/main/java/org/apache/druid/frame/write/FrameWriter.java b/processing/src/main/java/org/apache/druid/frame/write/FrameWriter.java index f37ed8571ac1..050bb6513d21 100644 --- a/processing/src/main/java/org/apache/druid/frame/write/FrameWriter.java +++ b/processing/src/main/java/org/apache/druid/frame/write/FrameWriter.java @@ -28,7 +28,7 @@ /** * Writer for {@link org.apache.druid.frame.Frame}. See that class for format information. * - * Generally obtained through a {@link FrameWriters#makeFrameWriterFactory}. + * Generally obtained through a {@link FrameWriters#makeFrameWriterFactory} */ public interface FrameWriter extends Closeable { diff --git a/processing/src/main/java/org/apache/druid/query/rowsandcols/LazilyDecoratedRowsAndColumns.java b/processing/src/main/java/org/apache/druid/query/rowsandcols/LazilyDecoratedRowsAndColumns.java index 3e938eb2d9d8..264057232a6a 100644 --- a/processing/src/main/java/org/apache/druid/query/rowsandcols/LazilyDecoratedRowsAndColumns.java +++ b/processing/src/main/java/org/apache/druid/query/rowsandcols/LazilyDecoratedRowsAndColumns.java @@ -418,4 +418,116 @@ private Pair naiveMaterialize(RowsAndColumns rac) return Pair.of(frameWriter.toByteArray(), sigBob.build()); } + + + @Nullable + public Pair naiveMaterializeToRowFrames(RowsAndColumns rac) + { + final int numRows = rac.numRows(); + + BitSet rowsToSkip = null; + + if (interval != null) { + rowsToSkip = new BitSet(numRows); + + final Column racColumn = rac.findColumn("__time"); + if (racColumn == null) { + // The time column doesn't exist, this means that we have a null column. A null column when coerced into a + // long as is required by the time filter produces all 0s, so either 0 is included and matches all rows or + // it's not and we skip all rows. + if (!interval.contains(0)) { + return null; + } + } else { + final ColumnAccessor accessor = racColumn.toAccessor(); + for (int i = 0; i < accessor.numRows(); ++i) { + rowsToSkip.set(i, !interval.contains(accessor.getLong(i))); + } + } + } + + AtomicInteger rowId = new AtomicInteger(0); + final ColumnSelectorFactoryMaker csfm = ColumnSelectorFactoryMaker.fromRAC(rac); + final ColumnSelectorFactory selectorFactory = csfm.make(rowId); + + if (filter != null) { + if (rowsToSkip == null) { + rowsToSkip = new BitSet(numRows); + } + + final ValueMatcher matcher = filter.makeMatcher(selectorFactory); + + for (; rowId.get() < numRows; rowId.incrementAndGet()) { + final int theId = rowId.get(); + if (rowsToSkip.get(theId)) { + continue; + } + + if (!matcher.matches(false)) { + rowsToSkip.set(theId); + } + } + } + + if (virtualColumns != null) { + throw new UOE("Cannot apply virtual columns [%s] with naive apply.", virtualColumns); + } + + ArrayList columnsToGenerate = new ArrayList<>(); + if (viewableColumns != null) { + columnsToGenerate.addAll(viewableColumns); + } else { + columnsToGenerate.addAll(rac.getColumnNames()); + // When/if we support virtual columns from here, we should auto-add them to the list here as well as they expand + // the implicit project when no projection is defined + } + + // There is all sorts of sub-optimal things in this code, but we just ignore them for now as it is difficult to + // optimally build frames for incremental data processing. In order to build the frame object here, we must first + // materialize everything in memory so that we know how long things are, such that we can allocate a big byte[] + // so that the Frame.wrap() call can be given just a big byte[] to do its reading from. + // + // It would be a bit better if we could just build the per-column byte[] and somehow re-generate a Frame from + // that. But it's also possible that this impedence mis-match is a function of using column-oriented frames and + // row-oriented frames are a much more friendly format. Anyway, this long comment is here because this exploration + // is being left as an exercise for the future. + + final RowSignature.Builder sigBob = RowSignature.builder(); + final ArenaMemoryAllocatorFactory memFactory = new ArenaMemoryAllocatorFactory(200 << 20); + + + for (String column : columnsToGenerate) { + final Column racColumn = rac.findColumn(column); + if (racColumn == null) { + continue; + } + sigBob.add(column, racColumn.toAccessor().getType()); + } + + long remainingRowsToSkip = limit.getOffset(); + long remainingRowsToFetch = limit.getLimitOrMax(); + + final FrameWriter frameWriter = FrameWriters.makeFrameWriterFactory( + FrameType.ROW_BASED, + memFactory, + sigBob.build(), + Collections.emptyList() + ).newFrameWriter(selectorFactory); + + rowId.set(0); + for (; rowId.get() < numRows && remainingRowsToFetch > 0; rowId.incrementAndGet()) { + final int theId = rowId.get(); + if (rowsToSkip != null && rowsToSkip.get(theId)) { + continue; + } + if (remainingRowsToSkip > 0) { + remainingRowsToSkip--; + continue; + } + remainingRowsToFetch--; + frameWriter.addSelection(); + } + + return Pair.of(frameWriter.toByteArray(), sigBob.build()); + } } diff --git a/processing/src/main/java/org/apache/druid/query/rowsandcols/concrete/RowBasedFrameRowAndColumns.java b/processing/src/main/java/org/apache/druid/query/rowsandcols/concrete/RowBasedFrameRowAndColumns.java index f8f4e79082cb..06f386161d93 100644 --- a/processing/src/main/java/org/apache/druid/query/rowsandcols/concrete/RowBasedFrameRowAndColumns.java +++ b/processing/src/main/java/org/apache/druid/query/rowsandcols/concrete/RowBasedFrameRowAndColumns.java @@ -18,7 +18,6 @@ public class RowBasedFrameRowAndColumns implements RowsAndColumns { private final Frame frame; private final RowSignature signature; - private final LinkedHashMap colCache = new LinkedHashMap<>(); public RowBasedFrameRowAndColumns(Frame frame, RowSignature signature) { this.frame = FrameType.ROW_BASED.ensureType(frame);; From d6cef471fd78d5fccd1307814ccc1c5dbbd1e6fa Mon Sep 17 00:00:00 2001 From: Soumyava Das Date: Thu, 7 Dec 2023 10:41:11 -0800 Subject: [PATCH 05/42] Working UTs --- .../msq/querykit/BaseLeafFrameProcessor.java | 2 +- .../WindowOperatorQueryFrameProcessor.java | 148 +++++++++------- ...dowOperatorQueryFrameProcessorFactory.java | 19 +++ .../msq/querykit/WindowOperatorQueryKit.java | 22 ++- .../apache/druid/msq/exec/MSQSelectTest.java | 159 +++++++++++++++--- .../WindowFramedAggregateProcessor.java | 22 +++ .../concrete/RowBasedFrameRowAndColumns.java | 25 ++- 7 files changed, 308 insertions(+), 89 deletions(-) diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/BaseLeafFrameProcessor.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/BaseLeafFrameProcessor.java index 1cb7af037529..ac3ef234a75b 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/BaseLeafFrameProcessor.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/BaseLeafFrameProcessor.java @@ -40,7 +40,7 @@ import java.util.List; import java.util.function.Function; -public abstract class BaseLeafFrameProcessor implements FrameProcessor +public abstract class BaseLeafFrameProcessor implements FrameProcessor { private final ReadableInput baseInput; private final ResourceHolder outputChannelHolder; diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryFrameProcessor.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryFrameProcessor.java index ae154b383824..222b9d4a0bbb 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryFrameProcessor.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryFrameProcessor.java @@ -1,3 +1,22 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + package org.apache.druid.msq.querykit; import com.fasterxml.jackson.databind.ObjectMapper; @@ -9,7 +28,6 @@ import org.apache.druid.frame.channel.WritableFrameChannel; import org.apache.druid.frame.processor.ReturnOrAwait; import org.apache.druid.frame.read.FrameReader; -import org.apache.druid.frame.segment.FrameSegment; import org.apache.druid.frame.util.SettableLongVirtualColumn; import org.apache.druid.frame.write.FrameWriter; import org.apache.druid.frame.write.FrameWriterFactory; @@ -23,30 +41,24 @@ import org.apache.druid.query.operator.OffsetLimit; import org.apache.druid.query.operator.Operator; import org.apache.druid.query.operator.OperatorFactory; -import org.apache.druid.query.operator.SegmentToRowsAndColumnsOperator; import org.apache.druid.query.operator.WindowOperatorQuery; import org.apache.druid.query.rowsandcols.LazilyDecoratedRowsAndColumns; import org.apache.druid.query.rowsandcols.RowsAndColumns; -import org.apache.druid.query.rowsandcols.concrete.FrameRowsAndColumns; import org.apache.druid.query.rowsandcols.concrete.RowBasedFrameRowAndColumns; -import org.apache.druid.query.scan.ScanQuery; import org.apache.druid.segment.Cursor; import org.apache.druid.segment.Segment; import org.apache.druid.segment.SegmentReference; import org.apache.druid.segment.SimpleAscendingOffset; import org.apache.druid.segment.SimpleSettableOffset; -import org.apache.druid.segment.StorageAdapter; import org.apache.druid.segment.VirtualColumn; import org.apache.druid.segment.VirtualColumns; import org.apache.druid.segment.column.RowSignature; -import org.apache.druid.timeline.SegmentId; import javax.annotation.Nullable; import java.io.Closeable; import java.io.IOException; import java.util.ArrayList; import java.util.List; -import java.util.concurrent.atomic.AtomicLong; import java.util.function.Function; public class WindowOperatorQueryFrameProcessor extends BaseLeafFrameProcessor @@ -96,6 +108,7 @@ public WindowOperatorQueryFrameProcessor( this.frameWriterVirtualColumns = VirtualColumns.create(frameWriterVirtualColumns); } + // deep storage @Override protected ReturnOrAwait runWithSegment(SegmentWithDescriptor segment) throws IOException @@ -118,63 +131,80 @@ protected ReturnOrAwait runWithInputChannel(ReadableFrameChannel inputChan // Read the frames from the channel // convert to FrameRowsAndColumns - if (inputChannel.canRead()) { - Frame f = inputChannel.read(); - - // the frame here is row based - // frame rows and columns need columnar. discuss with Eric - // Action item: need to implement a new rows and columns that accept a row-based frame - - // Create a frame rows and columns what would - RowBasedFrameRowAndColumns frameRowsAndColumns = new RowBasedFrameRowAndColumns(f, inputFrameReader.signature()); - Operator op = getOperator(frameRowsAndColumns); - // - //Operator op = new SegmentToRowsAndColumnsOperator(frameSegment); - // On the operator created above add the operators present in the query that we want to chain - - for ( OperatorFactory of : query.getOperators()) { - op = of.wrap(op); - } - - // Let the operator run - // the results that come in the receiver must be pushed to the outout channel - // need to transform the output rows and columns back to frame - Operator.go(op, new Operator.Receiver() - { - @Override - public Operator.Signal push(RowsAndColumns rac) - { - //outputFrameChannel.output(rac.toFrame()); - LazilyDecoratedRowsAndColumns tmpldrc = new LazilyDecoratedRowsAndColumns(rac, null, null, null, OffsetLimit.limit(Integer.MAX_VALUE), null, null); - Pair pairFrames = tmpldrc.naiveMaterializeToRowFrames(rac); - Frame f = Frame.wrap(pairFrames.lhs); - try { - Iterables.getOnlyElement(outputChannels()).write(new FrameWithPartition(f, FrameWithPartition.NO_PARTITION)); - } - catch (IOException e) { - throw new RuntimeException(e); - } - return Operator.Signal.GO; - } - - @Override - public void completed() - { - - } - }); - - } else if (inputChannel.isFinished()) { - return ReturnOrAwait.returnObject(Unit.instance()); - } else { - return ReturnOrAwait.awaitAll(inputChannels().size()); - } + if (inputChannel.canRead()) { + Frame f = inputChannel.read(); + + // the frame here is row based + // frame rows and columns need columnar. discuss with Eric + // Action item: need to implement a new rows and columns that accept a row-based frame + + // Create a frame rows and columns what would + RowBasedFrameRowAndColumns frameRowsAndColumns = new RowBasedFrameRowAndColumns(f, inputFrameReader.signature()); + Operator op = getOperator(frameRowsAndColumns); + // + //Operator op = new SegmentToRowsAndColumnsOperator(frameSegment); + // On the operator created above add the operators present in the query that we want to chain + + for (OperatorFactory of : query.getOperators()) { + op = of.wrap(op); + } + + // Let the operator run + // the results that come in the receiver must be pushed to the outout channel + // need to transform the output rows and columns back to frame + Operator.go(op, new Operator.Receiver() + { + @Override + public Operator.Signal push(RowsAndColumns rac) + { + //outputFrameChannel.output(rac.toFrame()); + LazilyDecoratedRowsAndColumns tmpldrc = new LazilyDecoratedRowsAndColumns( + rac, + null, + null, + null, + OffsetLimit.limit(Integer.MAX_VALUE), + null, + null + ); + Pair pairFrames = tmpldrc.naiveMaterializeToRowFrames(rac); + Frame f = Frame.wrap(pairFrames.lhs); + try { + Iterables.getOnlyElement(outputChannels()) + .write(new FrameWithPartition(f, FrameWithPartition.NO_PARTITION)); + } + catch (IOException e) { + throw new RuntimeException(e); + } + return Operator.Signal.GO; + } + + @Override + public void completed() + { + + } + }); + + } else if (inputChannel.isFinished()) { + return ReturnOrAwait.returnObject(Unit.instance()); + } else { + return ReturnOrAwait.awaitAll(inputChannels().size()); + } return ReturnOrAwait.runAgain(); } private static Operator getOperator(RowBasedFrameRowAndColumns frameRowsAndColumns) { - LazilyDecoratedRowsAndColumns ldrc = new LazilyDecoratedRowsAndColumns(frameRowsAndColumns, null, null, null, OffsetLimit.limit(Integer.MAX_VALUE), null, null); + LazilyDecoratedRowsAndColumns ldrc = new LazilyDecoratedRowsAndColumns( + frameRowsAndColumns, + null, + null, + null, + OffsetLimit.limit(Integer.MAX_VALUE), + null, + null + ); // Create an operator on top of the created rows and columns Operator op = new Operator() { diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryFrameProcessorFactory.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryFrameProcessorFactory.java index 71267185e28d..0afd863460af 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryFrameProcessorFactory.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryFrameProcessorFactory.java @@ -1,3 +1,22 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + package org.apache.druid.msq.querykit; import com.fasterxml.jackson.annotation.JsonCreator; diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryKit.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryKit.java index 347fb0a54dad..5028ada05a77 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryKit.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryKit.java @@ -1,3 +1,22 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + package org.apache.druid.msq.querykit; import com.fasterxml.jackson.databind.ObjectMapper; @@ -17,6 +36,7 @@ public WindowOperatorQueryKit(ObjectMapper jsonMapper) { this.jsonMapper = jsonMapper; } + @Override public QueryDefinition makeQueryDefinition( String queryId, @@ -53,7 +73,7 @@ public QueryDefinition makeQueryDefinition( // Create a new stage which takes in the subquery as an input queryDefBuilder.add( StageDefinition.builder(firstStageNumber) - .inputs(new StageInputSpec(firstStageNumber-1)) + .inputs(new StageInputSpec(firstStageNumber - 1)) .signature(rowSignature) .maxWorkerCount(maxWorkerCount) .shuffleSpec(null) diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQSelectTest.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQSelectTest.java index ba2a4b66c239..b0013f00b2d8 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQSelectTest.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQSelectTest.java @@ -53,6 +53,7 @@ import org.apache.druid.query.TableDataSource; import org.apache.druid.query.UnionDataSource; import org.apache.druid.query.UnnestDataSource; +import org.apache.druid.query.aggregation.AggregatorFactory; import org.apache.druid.query.aggregation.CountAggregatorFactory; import org.apache.druid.query.aggregation.DoubleSumAggregatorFactory; import org.apache.druid.query.aggregation.FilteredAggregatorFactory; @@ -65,8 +66,14 @@ import org.apache.druid.query.groupby.GroupByQuery; import org.apache.druid.query.groupby.orderby.DefaultLimitSpec; import org.apache.druid.query.groupby.orderby.OrderByColumnSpec; +import org.apache.druid.query.operator.NaivePartitioningOperatorFactory; +import org.apache.druid.query.operator.WindowOperatorQuery; +import org.apache.druid.query.operator.window.WindowFrame; +import org.apache.druid.query.operator.window.WindowFramedAggregateProcessor; +import org.apache.druid.query.operator.window.WindowOperatorFactory; import org.apache.druid.query.ordering.StringComparators; import org.apache.druid.query.scan.ScanQuery; +import org.apache.druid.query.spec.LegacySegmentSpec; import org.apache.druid.segment.column.ColumnType; import org.apache.druid.segment.column.RowSignature; import org.apache.druid.segment.join.JoinType; @@ -130,12 +137,12 @@ public class MSQSelectTest extends MSQTestBase public static Collection data() { Object[][] data = new Object[][]{ - {DEFAULT, DEFAULT_MSQ_CONTEXT} -// {DURABLE_STORAGE, DURABLE_STORAGE_MSQ_CONTEXT}, -// {FAULT_TOLERANCE, FAULT_TOLERANCE_MSQ_CONTEXT}, -// {PARALLEL_MERGE, PARALLEL_MERGE_MSQ_CONTEXT}, -// {QUERY_RESULTS_WITH_DURABLE_STORAGE, QUERY_RESULTS_WITH_DURABLE_STORAGE_CONTEXT}, -// {QUERY_RESULTS_WITH_DEFAULT, QUERY_RESULTS_WITH_DEFAULT_CONTEXT} + {DEFAULT, DEFAULT_MSQ_CONTEXT}, + {DURABLE_STORAGE, DURABLE_STORAGE_MSQ_CONTEXT}, + {FAULT_TOLERANCE, FAULT_TOLERANCE_MSQ_CONTEXT}, + {PARALLEL_MERGE, PARALLEL_MERGE_MSQ_CONTEXT}, + {QUERY_RESULTS_WITH_DURABLE_STORAGE, QUERY_RESULTS_WITH_DURABLE_STORAGE_CONTEXT}, + {QUERY_RESULTS_WITH_DEFAULT, QUERY_RESULTS_WITH_DEFAULT_CONTEXT} }; return Arrays.asList(data); @@ -621,29 +628,47 @@ public void testWindowOnFoo() .add("cc", ColumnType.DOUBLE) .build(); + final Query groupByQuery = GroupByQuery.builder() + .setDataSource(CalciteTests.DATASOURCE1) + .setInterval(querySegmentSpec(Filtration + .eternity())) + .setGranularity(Granularities.ALL) + .setDimensions(dimensions( + new DefaultDimensionSpec( + "m1", + "d0", + ColumnType.FLOAT + ) + )) + .setContext(context) + .build(); + + + final WindowFrame theFrame = new WindowFrame(WindowFrame.PeerType.ROWS, true, 0, true, 0); + final AggregatorFactory[] theAggs = { + new DoubleSumAggregatorFactory("w0", "d0") + }; + WindowFramedAggregateProcessor proc = new WindowFramedAggregateProcessor(theFrame, theAggs); + + final WindowOperatorQuery query = new WindowOperatorQuery( + new QueryDataSource(groupByQuery), + new LegacySegmentSpec(Intervals.ETERNITY), + context, + RowSignature.builder().add("d0", ColumnType.FLOAT).add("w0", ColumnType.DOUBLE).build(), + ImmutableList.of( + new NaivePartitioningOperatorFactory(ImmutableList.of("d0")), + new WindowOperatorFactory(proc) + ), + null + ); testSelectQuery() - .setSql("select m1,SUM(m1) OVER() cc from foo group by m1") + .setSql("select m1,SUM(m1) OVER(PARTITION BY m1) cc from foo group by m1") .setExpectedMSQSpec(MSQSpec.builder() - .query(GroupByQuery.builder() - .setDataSource(CalciteTests.DATASOURCE1) - .setInterval(querySegmentSpec(Filtration - .eternity())) - .setGranularity(Granularities.ALL) - .setDimensions(dimensions( - new DefaultDimensionSpec( - "cnt", - "d0", - ColumnType.LONG - ) - )) - .setAggregatorSpecs(aggregators(new CountAggregatorFactory( - "a0"))) - .setContext(context) - .build()) + .query(query) .columnMappings( new ColumnMappings(ImmutableList.of( - new ColumnMapping("d0", "cnt"), - new ColumnMapping("a0", "cnt1") + new ColumnMapping("d0", "m1"), + new ColumnMapping("w0", "cc") ) )) .tuningConfig(MSQTuningConfig.defaultConfig()) @@ -679,6 +704,90 @@ public void testWindowOnFoo() .verifyResults(); } + @Test + public void testWindowOnFooWithEmptyOver() + { + RowSignature rowSignature = RowSignature.builder() + .add("m1", ColumnType.FLOAT) + .add("cc", ColumnType.DOUBLE) + .build(); + + final Query groupByQuery = GroupByQuery.builder() + .setDataSource(CalciteTests.DATASOURCE1) + .setInterval(querySegmentSpec(Filtration + .eternity())) + .setGranularity(Granularities.ALL) + .setDimensions(dimensions( + new DefaultDimensionSpec( + "m1", + "d0", + ColumnType.FLOAT + ) + )) + .setContext(context) + .build(); + + + final WindowFrame theFrame = new WindowFrame(WindowFrame.PeerType.ROWS, true, 0, true, 0); + final AggregatorFactory[] theAggs = { + new DoubleSumAggregatorFactory("w0", "d0") + }; + WindowFramedAggregateProcessor proc = new WindowFramedAggregateProcessor(theFrame, theAggs); + + final WindowOperatorQuery query = new WindowOperatorQuery( + new QueryDataSource(groupByQuery), + new LegacySegmentSpec(Intervals.ETERNITY), + context, + RowSignature.builder().add("d0", ColumnType.FLOAT).add("w0", ColumnType.DOUBLE).build(), + ImmutableList.of( + new NaivePartitioningOperatorFactory(ImmutableList.of()), + new WindowOperatorFactory(proc) + ), + null + ); + testSelectQuery() + .setSql("select m1,SUM(m1) OVER() cc from foo group by m1") + .setExpectedMSQSpec(MSQSpec.builder() + .query(query) + .columnMappings( + new ColumnMappings(ImmutableList.of( + new ColumnMapping("d0", "m1"), + new ColumnMapping("w0", "cc") + ) + )) + .tuningConfig(MSQTuningConfig.defaultConfig()) + .destination(isDurableStorageDestination() + ? DurableStorageMSQDestination.INSTANCE + : TaskReportMSQDestination.INSTANCE) + .build()) + .setExpectedRowSignature(rowSignature) + .setExpectedResultRows(ImmutableList.of( + new Object[]{1.0f, 21.0}, + new Object[]{2.0f, 21.0}, + new Object[]{3.0f, 21.0}, + new Object[]{4.0f, 21.0}, + new Object[]{5.0f, 21.0}, + new Object[]{6.0f, 21.0} + )) + .setQueryContext(context) + .setExpectedCountersForStageWorkerChannel( + CounterSnapshotMatcher + .with().totalFiles(1), + 0, 0, "input0" + ) + .setExpectedCountersForStageWorkerChannel( + CounterSnapshotMatcher + .with().rows(6).frames(1), + 0, 0, "output" + ) + .setExpectedCountersForStageWorkerChannel( + CounterSnapshotMatcher + .with().rows(6).frames(1), + 0, 0, "shuffle" + ) + .verifyResults(); + } + @Test public void testSelectWithLimit() { diff --git a/processing/src/main/java/org/apache/druid/query/operator/window/WindowFramedAggregateProcessor.java b/processing/src/main/java/org/apache/druid/query/operator/window/WindowFramedAggregateProcessor.java index a67cb519407f..22d61ca8f2a1 100644 --- a/processing/src/main/java/org/apache/druid/query/operator/window/WindowFramedAggregateProcessor.java +++ b/processing/src/main/java/org/apache/druid/query/operator/window/WindowFramedAggregateProcessor.java @@ -28,6 +28,7 @@ import javax.annotation.Nullable; import java.util.Arrays; +import java.util.Objects; public class WindowFramedAggregateProcessor implements Processor { @@ -96,4 +97,25 @@ public String toString() ", aggregations=" + Arrays.toString(aggregations) + '}'; } + + @Override + public boolean equals(Object o) + { + if (this == o) { + return true; + } + if (o == null || getClass() != o.getClass()) { + return false; + } + WindowFramedAggregateProcessor that = (WindowFramedAggregateProcessor) o; + return Objects.equals(frame, that.frame) && Arrays.equals(aggregations, that.aggregations); + } + + @Override + public int hashCode() + { + int result = Objects.hash(frame); + result = 31 * result + Arrays.hashCode(aggregations); + return result; + } } diff --git a/processing/src/main/java/org/apache/druid/query/rowsandcols/concrete/RowBasedFrameRowAndColumns.java b/processing/src/main/java/org/apache/druid/query/rowsandcols/concrete/RowBasedFrameRowAndColumns.java index 06f386161d93..7b7937ebb84b 100644 --- a/processing/src/main/java/org/apache/druid/query/rowsandcols/concrete/RowBasedFrameRowAndColumns.java +++ b/processing/src/main/java/org/apache/druid/query/rowsandcols/concrete/RowBasedFrameRowAndColumns.java @@ -1,3 +1,22 @@ +/* + * 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.query.rowsandcols.concrete; import org.apache.druid.frame.Frame; @@ -12,15 +31,15 @@ import javax.annotation.Nullable; import java.util.Collection; -import java.util.LinkedHashMap; public class RowBasedFrameRowAndColumns implements RowsAndColumns { private final Frame frame; private final RowSignature signature; - public RowBasedFrameRowAndColumns(Frame frame, RowSignature signature) { - this.frame = FrameType.ROW_BASED.ensureType(frame);; + public RowBasedFrameRowAndColumns(Frame frame, RowSignature signature) + { + this.frame = FrameType.ROW_BASED.ensureType(frame); this.signature = signature; } From 40a18f10039d3bb68c5464e05535808fa630c6a0 Mon Sep 17 00:00:00 2001 From: Soumyava Das Date: Fri, 8 Dec 2023 12:30:31 -0800 Subject: [PATCH 06/42] Fixing running window function in console --- .../java/org/apache/druid/msq/guice/MSQIndexingModule.java | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/guice/MSQIndexingModule.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/guice/MSQIndexingModule.java index 4af832705c3f..8ab890dcd7a0 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/guice/MSQIndexingModule.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/guice/MSQIndexingModule.java @@ -84,6 +84,7 @@ import org.apache.druid.msq.input.table.TableInputSpec; import org.apache.druid.msq.kernel.NilExtraInfoHolder; import org.apache.druid.msq.querykit.InputNumberDataSource; +import org.apache.druid.msq.querykit.WindowOperatorQueryFrameProcessorFactory; import org.apache.druid.msq.querykit.common.OffsetLimitFrameProcessorFactory; import org.apache.druid.msq.querykit.common.SortMergeJoinFrameProcessorFactory; import org.apache.druid.msq.querykit.groupby.GroupByPostShuffleFrameProcessorFactory; @@ -158,7 +159,7 @@ public List getJacksonModules() NilExtraInfoHolder.class, SortMergeJoinFrameProcessorFactory.class, QueryResultFrameProcessorFactory.class, - + WindowOperatorQueryFrameProcessorFactory.class, // DataSource classes (note: ExternalDataSource is in MSQSqlModule) InputNumberDataSource.class, From 7e34aa8eb2643d1b9aa0b1701fec474a8008517a Mon Sep 17 00:00:00 2001 From: Soumyava Das Date: Wed, 3 Jan 2024 10:04:22 +0530 Subject: [PATCH 07/42] Adding shuffle spec and separating out stages for each window --- .../WindowLeafOperatorProcessorFactory.java | 85 ++++++++++ .../WindowOperatorQueryFrameProcessor.java | 92 ++++++++-- ...dowOperatorQueryFrameProcessorFactory.java | 23 ++- .../msq/querykit/WindowOperatorQueryKit.java | 157 +++++++++++++++++- .../apache/druid/msq/exec/MSQSelectTest.java | 89 ++++++++++ .../LazilyDecoratedRowsAndColumns.java | 118 +------------ .../concrete/FrameRowsAndColumns.java | 4 + 7 files changed, 426 insertions(+), 142 deletions(-) create mode 100644 extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowLeafOperatorProcessorFactory.java diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowLeafOperatorProcessorFactory.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowLeafOperatorProcessorFactory.java new file mode 100644 index 000000000000..293ddd609a57 --- /dev/null +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowLeafOperatorProcessorFactory.java @@ -0,0 +1,85 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.druid.msq.querykit; + +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonTypeName; +import org.apache.druid.frame.processor.OutputChannelFactory; +import org.apache.druid.msq.counters.CounterTracker; +import org.apache.druid.msq.input.InputSlice; +import org.apache.druid.msq.input.InputSliceReader; +import org.apache.druid.msq.kernel.FrameContext; +import org.apache.druid.msq.kernel.ProcessorsAndChannels; +import org.apache.druid.msq.kernel.StageDefinition; +import org.apache.druid.query.operator.OperatorFactory; + +import javax.annotation.Nullable; +import java.io.IOException; +import java.util.List; +import java.util.Objects; +import java.util.function.Consumer; + +@JsonTypeName("leafWindow") +public class WindowLeafOperatorProcessorFactory extends BaseFrameProcessorFactory +{ + private final List operatorFactoryList; + + @JsonCreator + public WindowLeafOperatorProcessorFactory(List operatorFactoryList) + { + this.operatorFactoryList = operatorFactoryList; + } + + @Override + public ProcessorsAndChannels makeProcessors( + StageDefinition stageDefinition, + int workerNumber, + List inputSlices, + InputSliceReader inputSliceReader, + @Nullable Object extra, + OutputChannelFactory outputChannelFactory, + FrameContext frameContext, + int maxOutstandingProcessors, + CounterTracker counters, + Consumer warningPublisher + ) throws IOException + { + return null; + } + + @Override + public boolean equals(Object o) + { + if (this == o) { + return true; + } + if (o == null || getClass() != o.getClass()) { + return false; + } + WindowLeafOperatorProcessorFactory that = (WindowLeafOperatorProcessorFactory) o; + return Objects.equals(operatorFactoryList, that.operatorFactoryList); + } + + @Override + public int hashCode() + { + return Objects.hash(operatorFactoryList); + } +} diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryFrameProcessor.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryFrameProcessor.java index 222b9d4a0bbb..8c3504475952 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryFrameProcessor.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryFrameProcessor.java @@ -23,6 +23,8 @@ import com.google.common.collect.Iterables; import org.apache.druid.collections.ResourceHolder; import org.apache.druid.frame.Frame; +import org.apache.druid.frame.FrameType; +import org.apache.druid.frame.allocation.ArenaMemoryAllocatorFactory; import org.apache.druid.frame.channel.FrameWithPartition; import org.apache.druid.frame.channel.ReadableFrameChannel; import org.apache.druid.frame.channel.WritableFrameChannel; @@ -31,6 +33,7 @@ import org.apache.druid.frame.util.SettableLongVirtualColumn; import org.apache.druid.frame.write.FrameWriter; import org.apache.druid.frame.write.FrameWriterFactory; +import org.apache.druid.frame.write.FrameWriters; import org.apache.druid.java.util.common.Pair; import org.apache.druid.java.util.common.Unit; import org.apache.druid.java.util.common.io.Closer; @@ -44,7 +47,10 @@ import org.apache.druid.query.operator.WindowOperatorQuery; import org.apache.druid.query.rowsandcols.LazilyDecoratedRowsAndColumns; import org.apache.druid.query.rowsandcols.RowsAndColumns; +import org.apache.druid.query.rowsandcols.column.Column; import org.apache.druid.query.rowsandcols.concrete.RowBasedFrameRowAndColumns; +import org.apache.druid.query.rowsandcols.semantic.ColumnSelectorFactoryMaker; +import org.apache.druid.segment.ColumnSelectorFactory; import org.apache.druid.segment.Cursor; import org.apache.druid.segment.Segment; import org.apache.druid.segment.SegmentReference; @@ -58,7 +64,10 @@ import java.io.Closeable; import java.io.IOException; import java.util.ArrayList; +import java.util.BitSet; +import java.util.Collections; import java.util.List; +import java.util.concurrent.atomic.AtomicInteger; import java.util.function.Function; public class WindowOperatorQueryFrameProcessor extends BaseLeafFrameProcessor @@ -66,10 +75,13 @@ public class WindowOperatorQueryFrameProcessor extends BaseLeafFrameProcessor private static final Logger log = new Logger(ScanQueryFrameProcessor.class); private final WindowOperatorQuery query; + + private final List operatorFactoryList; private final ObjectMapper jsonMapper; private final SettableLongVirtualColumn partitionBoostVirtualColumn; private final VirtualColumns frameWriterVirtualColumns; private final Closer closer = Closer.create(); + private final RowSignature outputStageSignature; private Cursor cursor; private Segment segment; @@ -79,11 +91,13 @@ public class WindowOperatorQueryFrameProcessor extends BaseLeafFrameProcessor public WindowOperatorQueryFrameProcessor( final WindowOperatorQuery query, + final List operatorFactoryList, final ObjectMapper jsonMapper, final ReadableInput baseInput, final Function segmentMapFn, final ResourceHolder outputChannelHolder, - final ResourceHolder frameWriterFactoryHolder + final ResourceHolder frameWriterFactoryHolder, + RowSignature rearrangePointer ) { super( @@ -95,6 +109,8 @@ public WindowOperatorQueryFrameProcessor( this.query = query; this.jsonMapper = jsonMapper; this.partitionBoostVirtualColumn = new SettableLongVirtualColumn(QueryKitUtils.PARTITION_BOOST_COLUMN); + this.operatorFactoryList = operatorFactoryList; + this.outputStageSignature = rearrangePointer; final List frameWriterVirtualColumns = new ArrayList<>(); frameWriterVirtualColumns.add(partitionBoostVirtualColumn); @@ -111,14 +127,14 @@ public WindowOperatorQueryFrameProcessor( // deep storage @Override - protected ReturnOrAwait runWithSegment(SegmentWithDescriptor segment) throws IOException + protected ReturnOrAwait runWithSegment(SegmentWithDescriptor segment) { return null; } // realtime @Override - protected ReturnOrAwait runWithLoadedSegment(SegmentWithDescriptor segment) throws IOException + protected ReturnOrAwait runWithLoadedSegment(SegmentWithDescriptor segment) { return null; } @@ -126,7 +142,6 @@ protected ReturnOrAwait runWithLoadedSegment(SegmentWithDescriptor segment // previous stage output @Override protected ReturnOrAwait runWithInputChannel(ReadableFrameChannel inputChannel, FrameReader inputFrameReader) - throws IOException { // Read the frames from the channel // convert to FrameRowsAndColumns @@ -145,7 +160,7 @@ protected ReturnOrAwait runWithInputChannel(ReadableFrameChannel inputChan //Operator op = new SegmentToRowsAndColumnsOperator(frameSegment); // On the operator created above add the operators present in the query that we want to chain - for (OperatorFactory of : query.getOperators()) { + for (OperatorFactory of : operatorFactoryList) { op = of.wrap(op); } @@ -157,17 +172,8 @@ protected ReturnOrAwait runWithInputChannel(ReadableFrameChannel inputChan @Override public Operator.Signal push(RowsAndColumns rac) { - //outputFrameChannel.output(rac.toFrame()); - LazilyDecoratedRowsAndColumns tmpldrc = new LazilyDecoratedRowsAndColumns( - rac, - null, - null, - null, - OffsetLimit.limit(Integer.MAX_VALUE), - null, - null - ); - Pair pairFrames = tmpldrc.naiveMaterializeToRowFrames(rac); + // convert the rac to a row-based frame + Pair pairFrames = materializeRacToRowFrames(rac, outputStageSignature); Frame f = Frame.wrap(pairFrames.lhs); try { Iterables.getOnlyElement(outputChannels()) @@ -219,4 +225,58 @@ public Closeable goOrContinue(Closeable continuationObject, Receiver receiver) }; return op; } + + public Pair materializeRacToRowFrames(RowsAndColumns rac, RowSignature outputSignature) + { + final int numRows = rac.numRows(); + + BitSet rowsToSkip = null; + + AtomicInteger rowId = new AtomicInteger(0); + final ColumnSelectorFactoryMaker csfm = ColumnSelectorFactoryMaker.fromRAC(rac); + final ColumnSelectorFactory selectorFactory = csfm.make(rowId); + + //columnsToGenerate.addAll(rac.getColumnNames()); + ArrayList columnsToGenerate = new ArrayList<>(outputSignature.getColumnNames()); + + + final RowSignature.Builder sigBob = RowSignature.builder(); + final ArenaMemoryAllocatorFactory memFactory = new ArenaMemoryAllocatorFactory(200 << 20); + + + for (String column : columnsToGenerate) { + final Column racColumn = rac.findColumn(column); + if (racColumn == null) { + continue; + } + sigBob.add(column, racColumn.toAccessor().getType()); + } + + long remainingRowsToSkip = 0; + long remainingRowsToFetch = Integer.MAX_VALUE; + + final FrameWriter frameWriter; + frameWriter = FrameWriters.makeFrameWriterFactory( + FrameType.ROW_BASED, + memFactory, + sigBob.build(), + Collections.emptyList() + ).newFrameWriter(selectorFactory); + + rowId.set(0); + for (; rowId.get() < numRows && remainingRowsToFetch > 0; rowId.incrementAndGet()) { + final int theId = rowId.get(); + if (rowsToSkip != null && rowsToSkip.get(theId)) { + continue; + } + if (remainingRowsToSkip > 0) { + remainingRowsToSkip--; + continue; + } + remainingRowsToFetch--; + frameWriter.addSelection(); + } + + return Pair.of(frameWriter.toByteArray(), sigBob.build()); + } } diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryFrameProcessorFactory.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryFrameProcessorFactory.java index 0afd863460af..5321d6f968d3 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryFrameProcessorFactory.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryFrameProcessorFactory.java @@ -29,21 +29,32 @@ import org.apache.druid.frame.write.FrameWriterFactory; import org.apache.druid.msq.input.ReadableInput; import org.apache.druid.msq.kernel.FrameContext; +import org.apache.druid.query.operator.OperatorFactory; import org.apache.druid.query.operator.WindowOperatorQuery; import org.apache.druid.segment.SegmentReference; +import org.apache.druid.segment.column.RowSignature; +import java.util.List; import java.util.function.Function; @JsonTypeName("window") public class WindowOperatorQueryFrameProcessorFactory extends BaseLeafFrameProcessorFactory { private final WindowOperatorQuery query; + private final List operatorList; + RowSignature rearrangePoiter; @JsonCreator - public WindowOperatorQueryFrameProcessorFactory(@JsonProperty("query") WindowOperatorQuery query) + public WindowOperatorQueryFrameProcessorFactory( + @JsonProperty("query") WindowOperatorQuery query, + @JsonProperty("operators") List operatorFactoryList, + RowSignature pointers + ) { super(query); this.query = Preconditions.checkNotNull(query, "query"); + this.operatorList = operatorFactoryList; + this.rearrangePoiter = pointers; } @JsonProperty @@ -52,6 +63,12 @@ public WindowOperatorQuery getQuery() return query; } + @JsonProperty + public List getOperators() + { + return operatorList; + } + @Override protected FrameProcessor makeProcessor( ReadableInput baseInput, @@ -63,11 +80,13 @@ protected FrameProcessor makeProcessor( { return new WindowOperatorQueryFrameProcessor( query, + operatorList, frameContext.jsonMapper(), baseInput, segmentMapFn, outputChannelHolder, - frameWriterFactoryHolder + frameWriterFactoryHolder, + rearrangePoiter ); } } diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryKit.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryKit.java index 5028ada05a77..b30c4b169430 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryKit.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryKit.java @@ -20,14 +20,29 @@ package org.apache.druid.msq.querykit; import com.fasterxml.jackson.databind.ObjectMapper; +import org.apache.druid.frame.key.ClusterBy; +import org.apache.druid.frame.key.KeyColumn; +import org.apache.druid.frame.key.KeyOrder; import org.apache.druid.msq.input.stage.StageInputSpec; +import org.apache.druid.msq.kernel.HashShuffleSpec; import org.apache.druid.msq.kernel.QueryDefinition; import org.apache.druid.msq.kernel.QueryDefinitionBuilder; +import org.apache.druid.msq.kernel.ShuffleSpec; import org.apache.druid.msq.kernel.StageDefinition; import org.apache.druid.query.Query; +import org.apache.druid.query.operator.ColumnWithDirection; +import org.apache.druid.query.operator.NaivePartitioningOperatorFactory; +import org.apache.druid.query.operator.NaiveSortOperatorFactory; +import org.apache.druid.query.operator.OperatorFactory; import org.apache.druid.query.operator.WindowOperatorQuery; +import org.apache.druid.query.operator.window.WindowOperatorFactory; import org.apache.druid.segment.column.RowSignature; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + public class WindowOperatorQueryKit implements QueryKit { private final ObjectMapper jsonMapper; @@ -48,6 +63,16 @@ public QueryDefinition makeQueryDefinition( ) { // need to validate query first + // populate the group of operators to be processed as each stage + // the size of the operators is the number of serialized stages + // later we should also check if these can be parallelized + // check there is an empty over clause or not + List> operatorList = new ArrayList<>(); + boolean status = validateAndReturnOperatorList(originalQuery, operatorList); + + + ShuffleSpec nextShuffleSpec = findShuffleSpecForNextWindow(operatorList.get(0), maxWorkerCount); + // add this shuffle spec to the last stage of the inner query final QueryDefinitionBuilder queryDefBuilder = QueryDefinition.builder().queryId(queryId); final DataSourcePlan dataSourcePlan = DataSourcePlan.forDataSource( @@ -70,15 +95,129 @@ public QueryDefinition makeQueryDefinition( final WindowOperatorQuery queryToRun = (WindowOperatorQuery) originalQuery.withDataSource(dataSourcePlan.getNewDataSource()); RowSignature rowSignature = queryToRun.getRowSignature(); - // Create a new stage which takes in the subquery as an input - queryDefBuilder.add( - StageDefinition.builder(firstStageNumber) - .inputs(new StageInputSpec(firstStageNumber - 1)) - .signature(rowSignature) - .maxWorkerCount(maxWorkerCount) - .shuffleSpec(null) - .processorFactory(new WindowOperatorQueryFrameProcessorFactory(queryToRun)) - ); + if (status) { + // empty over clause found + // moving everything to a single partition + queryDefBuilder.add( + StageDefinition.builder(firstStageNumber) + .inputs(new StageInputSpec(firstStageNumber - 1)) + .signature(rowSignature) + .maxWorkerCount(maxWorkerCount) + .shuffleSpec(null) + .processorFactory(new WindowOperatorQueryFrameProcessorFactory( + queryToRun, + queryToRun.getOperators(), + rowSignature + )) + ); + } else { + // there are multiple windows present in the query + // Create stages for each window in the query + // These stages will be serialized + // the partition by clause of the next window will be the shuffle key for the previous window + RowSignature.Builder bob = RowSignature.builder(); + final int numberOfWindows = operatorList.size(); + final int baseSize = rowSignature.size() - numberOfWindows; + for (int i = 0; i < baseSize; i++) { + bob.add(rowSignature.getColumnName(i), rowSignature.getColumnType(i).get()); + } + + for (int i = 0; i < numberOfWindows; i++) { + bob.add(rowSignature.getColumnName(baseSize + i), rowSignature.getColumnType(baseSize + i).get()).build(); + // find the shuffle spec of the next stage + // if it is the last stage set the next shuffle spec to null + if (i + 1 == numberOfWindows) { + nextShuffleSpec = null; + } else { + nextShuffleSpec = findShuffleSpecForNextWindow(operatorList.get(i + 1), maxWorkerCount); + } + + final RowSignature intermediateSignature = bob.build(); + final RowSignature stageRowSignature; + if (nextShuffleSpec == null) { + stageRowSignature = intermediateSignature; + } else { + stageRowSignature = QueryKitUtils.sortableSignature( + intermediateSignature, + nextShuffleSpec.clusterBy().getColumns() + ); + } + + queryDefBuilder.add( + StageDefinition.builder(firstStageNumber + i) + .inputs(new StageInputSpec(firstStageNumber + i - 1)) + .signature(stageRowSignature) + .maxWorkerCount(maxWorkerCount) + .shuffleSpec(nextShuffleSpec) + .processorFactory(new WindowOperatorQueryFrameProcessorFactory( + queryToRun, + operatorList.get(i), + stageRowSignature + )) + ); + } + } return queryDefBuilder.queryId(queryId).build(); } + + private boolean validateAndReturnOperatorList( + WindowOperatorQuery originalQuery, + List> operatorList + ) + { + final List operators = originalQuery.getOperators(); + List operatorFactoryList = new ArrayList<>(); + for (OperatorFactory of : operators) { + operatorFactoryList.add(of); + if (of instanceof WindowOperatorFactory) { + operatorList.add(operatorFactoryList); + operatorFactoryList = new ArrayList<>(); + } else if (of instanceof NaivePartitioningOperatorFactory) { + if (((NaivePartitioningOperatorFactory) of).getPartitionColumns().isEmpty()) { + operatorList.clear(); + operatorList.add(originalQuery.getOperators()); + return true; + } + } + } + return false; + } + + private ShuffleSpec findShuffleSpecForNextWindow(List operatorFactories, int maxWorkerCount) + { + NaivePartitioningOperatorFactory partition = null; + NaiveSortOperatorFactory sort = null; + List keyColsOfWindow = new ArrayList<>(); + for (OperatorFactory of : operatorFactories) { + if (of instanceof NaivePartitioningOperatorFactory) { + partition = (NaivePartitioningOperatorFactory) of; + } else if (of instanceof NaiveSortOperatorFactory) { + sort = (NaiveSortOperatorFactory) of; + } + } + Map colMap = new HashMap<>(); + if (sort != null) { + for (ColumnWithDirection sortColumn : sort.getSortColumns()) { + colMap.put(sortColumn.getColumn(), sortColumn.getDirection()); + } + } + assert partition != null; + if (partition.getPartitionColumns().isEmpty()) { + return null; + } + for (String partitionColumn : partition.getPartitionColumns()) { + KeyColumn kc; + if (colMap.containsKey(partitionColumn)) { + if (colMap.get(partitionColumn) == ColumnWithDirection.Direction.ASC) { + kc = new KeyColumn(partitionColumn, KeyOrder.ASCENDING); + } else { + kc = new KeyColumn(partitionColumn, KeyOrder.DESCENDING); + } + } else { + kc = new KeyColumn(partitionColumn, KeyOrder.ASCENDING); + } + keyColsOfWindow.add(kc); + } + return new HashShuffleSpec(new ClusterBy(keyColsOfWindow, 0), maxWorkerCount); + } } diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQSelectTest.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQSelectTest.java index 92828c423c28..484a1828d9c7 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQSelectTest.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQSelectTest.java @@ -704,6 +704,95 @@ public void testWindowOnFoo() .verifyResults(); } + @Test + public void testWindowOnFooWith2Windows() + { + RowSignature rowSignature = RowSignature.builder() + .add("m1", ColumnType.FLOAT) + .add("cc", ColumnType.DOUBLE) + .build(); + + final Query groupByQuery = GroupByQuery.builder() + .setDataSource(CalciteTests.DATASOURCE1) + .setInterval(querySegmentSpec(Filtration + .eternity())) + .setGranularity(Granularities.ALL) + .setDimensions(dimensions( + new DefaultDimensionSpec( + "m1", + "d0", + ColumnType.FLOAT + ) + )) + .setContext(context) + .build(); + + + final WindowFrame theFrame = new WindowFrame(WindowFrame.PeerType.ROWS, true, 0, true, 0); + final AggregatorFactory[] theAggs = { + new DoubleSumAggregatorFactory("w0", "d0") + }; + WindowFramedAggregateProcessor proc = new WindowFramedAggregateProcessor(theFrame, theAggs); + + final WindowOperatorQuery query = new WindowOperatorQuery( + new QueryDataSource(groupByQuery), + new LegacySegmentSpec(Intervals.ETERNITY), + context, + RowSignature.builder().add("d0", ColumnType.FLOAT).add("w0", ColumnType.DOUBLE).build(), + ImmutableList.of( + new NaivePartitioningOperatorFactory(ImmutableList.of("d0")), + new WindowOperatorFactory(proc) + ), + null + ); + testSelectQuery() + .setSql("SELECT m1, m2,\n" + + "SUM(m2) OVER() as summ2\n" + + ",SUM(m1) OVER(PARTITION BY m2) as summ1\n" + + "from foo\n" + + "GROUP BY m1,m2") + .setExpectedMSQSpec(MSQSpec.builder() + .query(query) + .columnMappings( + new ColumnMappings(ImmutableList.of( + new ColumnMapping("d0", "m1"), + new ColumnMapping("w0", "cc") + ) + )) + .tuningConfig(MSQTuningConfig.defaultConfig()) + .destination(isDurableStorageDestination() + ? DurableStorageMSQDestination.INSTANCE + : TaskReportMSQDestination.INSTANCE) + .build()) + .setExpectedRowSignature(rowSignature) + .setExpectedResultRows(ImmutableList.of( + new Object[]{1.0f, 1.0}, + new Object[]{2.0f, 2.0}, + new Object[]{3.0f, 3.0}, + new Object[]{4.0f, 4.0}, + new Object[]{5.0f, 5.0}, + new Object[]{6.0f, 6.0} + )) + .setQueryContext(context) + .setExpectedCountersForStageWorkerChannel( + CounterSnapshotMatcher + .with().totalFiles(1), + 0, 0, "input0" + ) + .setExpectedCountersForStageWorkerChannel( + CounterSnapshotMatcher + .with().rows(6).frames(1), + 0, 0, "output" + ) + .setExpectedCountersForStageWorkerChannel( + CounterSnapshotMatcher + .with().rows(6).frames(1), + 0, 0, "shuffle" + ) + .verifyResults(); + } + + @Test public void testWindowOnFooWithEmptyOver() { diff --git a/processing/src/main/java/org/apache/druid/query/rowsandcols/LazilyDecoratedRowsAndColumns.java b/processing/src/main/java/org/apache/druid/query/rowsandcols/LazilyDecoratedRowsAndColumns.java index 264057232a6a..380b94acaafe 100644 --- a/processing/src/main/java/org/apache/druid/query/rowsandcols/LazilyDecoratedRowsAndColumns.java +++ b/processing/src/main/java/org/apache/druid/query/rowsandcols/LazilyDecoratedRowsAndColumns.java @@ -225,9 +225,9 @@ private Pair materializeStorageAdapter(StorageAdapter as) cols = base.getColumnNames(); } else { cols = ImmutableList.builder() - .addAll(base.getColumnNames()) - .addAll(virtualColumns.getColumnNames()) - .build(); + .addAll(base.getColumnNames()) + .addAll(virtualColumns.getColumnNames()) + .build(); } } AtomicReference siggy = new AtomicReference<>(null); @@ -418,116 +418,4 @@ private Pair naiveMaterialize(RowsAndColumns rac) return Pair.of(frameWriter.toByteArray(), sigBob.build()); } - - - @Nullable - public Pair naiveMaterializeToRowFrames(RowsAndColumns rac) - { - final int numRows = rac.numRows(); - - BitSet rowsToSkip = null; - - if (interval != null) { - rowsToSkip = new BitSet(numRows); - - final Column racColumn = rac.findColumn("__time"); - if (racColumn == null) { - // The time column doesn't exist, this means that we have a null column. A null column when coerced into a - // long as is required by the time filter produces all 0s, so either 0 is included and matches all rows or - // it's not and we skip all rows. - if (!interval.contains(0)) { - return null; - } - } else { - final ColumnAccessor accessor = racColumn.toAccessor(); - for (int i = 0; i < accessor.numRows(); ++i) { - rowsToSkip.set(i, !interval.contains(accessor.getLong(i))); - } - } - } - - AtomicInteger rowId = new AtomicInteger(0); - final ColumnSelectorFactoryMaker csfm = ColumnSelectorFactoryMaker.fromRAC(rac); - final ColumnSelectorFactory selectorFactory = csfm.make(rowId); - - if (filter != null) { - if (rowsToSkip == null) { - rowsToSkip = new BitSet(numRows); - } - - final ValueMatcher matcher = filter.makeMatcher(selectorFactory); - - for (; rowId.get() < numRows; rowId.incrementAndGet()) { - final int theId = rowId.get(); - if (rowsToSkip.get(theId)) { - continue; - } - - if (!matcher.matches(false)) { - rowsToSkip.set(theId); - } - } - } - - if (virtualColumns != null) { - throw new UOE("Cannot apply virtual columns [%s] with naive apply.", virtualColumns); - } - - ArrayList columnsToGenerate = new ArrayList<>(); - if (viewableColumns != null) { - columnsToGenerate.addAll(viewableColumns); - } else { - columnsToGenerate.addAll(rac.getColumnNames()); - // When/if we support virtual columns from here, we should auto-add them to the list here as well as they expand - // the implicit project when no projection is defined - } - - // There is all sorts of sub-optimal things in this code, but we just ignore them for now as it is difficult to - // optimally build frames for incremental data processing. In order to build the frame object here, we must first - // materialize everything in memory so that we know how long things are, such that we can allocate a big byte[] - // so that the Frame.wrap() call can be given just a big byte[] to do its reading from. - // - // It would be a bit better if we could just build the per-column byte[] and somehow re-generate a Frame from - // that. But it's also possible that this impedence mis-match is a function of using column-oriented frames and - // row-oriented frames are a much more friendly format. Anyway, this long comment is here because this exploration - // is being left as an exercise for the future. - - final RowSignature.Builder sigBob = RowSignature.builder(); - final ArenaMemoryAllocatorFactory memFactory = new ArenaMemoryAllocatorFactory(200 << 20); - - - for (String column : columnsToGenerate) { - final Column racColumn = rac.findColumn(column); - if (racColumn == null) { - continue; - } - sigBob.add(column, racColumn.toAccessor().getType()); - } - - long remainingRowsToSkip = limit.getOffset(); - long remainingRowsToFetch = limit.getLimitOrMax(); - - final FrameWriter frameWriter = FrameWriters.makeFrameWriterFactory( - FrameType.ROW_BASED, - memFactory, - sigBob.build(), - Collections.emptyList() - ).newFrameWriter(selectorFactory); - - rowId.set(0); - for (; rowId.get() < numRows && remainingRowsToFetch > 0; rowId.incrementAndGet()) { - final int theId = rowId.get(); - if (rowsToSkip != null && rowsToSkip.get(theId)) { - continue; - } - if (remainingRowsToSkip > 0) { - remainingRowsToSkip--; - continue; - } - remainingRowsToFetch--; - frameWriter.addSelection(); - } - - return Pair.of(frameWriter.toByteArray(), sigBob.build()); - } } diff --git a/processing/src/main/java/org/apache/druid/query/rowsandcols/concrete/FrameRowsAndColumns.java b/processing/src/main/java/org/apache/druid/query/rowsandcols/concrete/FrameRowsAndColumns.java index e448acd8ded0..c22d60e28da8 100644 --- a/processing/src/main/java/org/apache/druid/query/rowsandcols/concrete/FrameRowsAndColumns.java +++ b/processing/src/main/java/org/apache/druid/query/rowsandcols/concrete/FrameRowsAndColumns.java @@ -28,6 +28,7 @@ import org.apache.druid.java.util.common.Intervals; import org.apache.druid.query.rowsandcols.RowsAndColumns; import org.apache.druid.query.rowsandcols.column.Column; +import org.apache.druid.query.rowsandcols.semantic.WireTransferable; import org.apache.druid.segment.StorageAdapter; import org.apache.druid.segment.column.ColumnType; import org.apache.druid.segment.column.RowSignature; @@ -90,6 +91,9 @@ public T as(Class clazz) if (StorageAdapter.class.equals(clazz)) { return (T) new FrameStorageAdapter(frame, FrameReader.create(signature), Intervals.ETERNITY); } + if (WireTransferable.class.equals(clazz)) { + return (T) this; + } return null; } } From f5a1f5980857538647894542ed13701b63064de9 Mon Sep 17 00:00:00 2001 From: Soumyava Das Date: Mon, 8 Jan 2024 14:19:58 +0530 Subject: [PATCH 08/42] serde stuff by adding ops to proc factory --- .../WindowLeafOperatorProcessorFactory.java | 85 --------- .../WindowOperatorQueryFrameProcessor.java | 20 ++- ...dowOperatorQueryFrameProcessorFactory.java | 54 +++++- .../apache/druid/msq/exec/MSQInsertTest.java | 96 ++++++++++ .../apache/druid/msq/exec/MSQSelectTest.java | 166 ++++++++++++++++-- 5 files changed, 301 insertions(+), 120 deletions(-) delete mode 100644 extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowLeafOperatorProcessorFactory.java diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowLeafOperatorProcessorFactory.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowLeafOperatorProcessorFactory.java deleted file mode 100644 index 293ddd609a57..000000000000 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowLeafOperatorProcessorFactory.java +++ /dev/null @@ -1,85 +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.querykit; - -import com.fasterxml.jackson.annotation.JsonCreator; -import com.fasterxml.jackson.annotation.JsonTypeName; -import org.apache.druid.frame.processor.OutputChannelFactory; -import org.apache.druid.msq.counters.CounterTracker; -import org.apache.druid.msq.input.InputSlice; -import org.apache.druid.msq.input.InputSliceReader; -import org.apache.druid.msq.kernel.FrameContext; -import org.apache.druid.msq.kernel.ProcessorsAndChannels; -import org.apache.druid.msq.kernel.StageDefinition; -import org.apache.druid.query.operator.OperatorFactory; - -import javax.annotation.Nullable; -import java.io.IOException; -import java.util.List; -import java.util.Objects; -import java.util.function.Consumer; - -@JsonTypeName("leafWindow") -public class WindowLeafOperatorProcessorFactory extends BaseFrameProcessorFactory -{ - private final List operatorFactoryList; - - @JsonCreator - public WindowLeafOperatorProcessorFactory(List operatorFactoryList) - { - this.operatorFactoryList = operatorFactoryList; - } - - @Override - public ProcessorsAndChannels makeProcessors( - StageDefinition stageDefinition, - int workerNumber, - List inputSlices, - InputSliceReader inputSliceReader, - @Nullable Object extra, - OutputChannelFactory outputChannelFactory, - FrameContext frameContext, - int maxOutstandingProcessors, - CounterTracker counters, - Consumer warningPublisher - ) throws IOException - { - return null; - } - - @Override - public boolean equals(Object o) - { - if (this == o) { - return true; - } - if (o == null || getClass() != o.getClass()) { - return false; - } - WindowLeafOperatorProcessorFactory that = (WindowLeafOperatorProcessorFactory) o; - return Objects.equals(operatorFactoryList, that.operatorFactoryList); - } - - @Override - public int hashCode() - { - return Objects.hash(operatorFactoryList); - } -} diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryFrameProcessor.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryFrameProcessor.java index 8c3504475952..a28048029fdf 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryFrameProcessor.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryFrameProcessor.java @@ -40,7 +40,6 @@ import org.apache.druid.java.util.common.logger.Logger; import org.apache.druid.msq.input.ReadableInput; import org.apache.druid.msq.input.table.SegmentWithDescriptor; -import org.apache.druid.msq.querykit.scan.ScanQueryFrameProcessor; import org.apache.druid.query.operator.OffsetLimit; import org.apache.druid.query.operator.Operator; import org.apache.druid.query.operator.OperatorFactory; @@ -73,7 +72,7 @@ public class WindowOperatorQueryFrameProcessor extends BaseLeafFrameProcessor { - private static final Logger log = new Logger(ScanQueryFrameProcessor.class); + private static final Logger log = new Logger(WindowOperatorQueryFrameProcessor.class); private final WindowOperatorQuery query; private final List operatorFactoryList; @@ -97,7 +96,7 @@ public WindowOperatorQueryFrameProcessor( final Function segmentMapFn, final ResourceHolder outputChannelHolder, final ResourceHolder frameWriterFactoryHolder, - RowSignature rearrangePointer + final RowSignature rowSignature ) { super( @@ -110,7 +109,7 @@ public WindowOperatorQueryFrameProcessor( this.jsonMapper = jsonMapper; this.partitionBoostVirtualColumn = new SettableLongVirtualColumn(QueryKitUtils.PARTITION_BOOST_COLUMN); this.operatorFactoryList = operatorFactoryList; - this.outputStageSignature = rearrangePointer; + this.outputStageSignature = rowSignature; final List frameWriterVirtualColumns = new ArrayList<>(); frameWriterVirtualColumns.add(partitionBoostVirtualColumn); @@ -125,14 +124,12 @@ public WindowOperatorQueryFrameProcessor( this.frameWriterVirtualColumns = VirtualColumns.create(frameWriterVirtualColumns); } - // deep storage @Override protected ReturnOrAwait runWithSegment(SegmentWithDescriptor segment) { return null; } - // realtime @Override protected ReturnOrAwait runWithLoadedSegment(SegmentWithDescriptor segment) { @@ -154,12 +151,18 @@ protected ReturnOrAwait runWithInputChannel(ReadableFrameChannel inputChan // Action item: need to implement a new rows and columns that accept a row-based frame // Create a frame rows and columns what would + /** + * OVER(PARTITION BY m1) + */ RowBasedFrameRowAndColumns frameRowsAndColumns = new RowBasedFrameRowAndColumns(f, inputFrameReader.signature()); Operator op = getOperator(frameRowsAndColumns); // //Operator op = new SegmentToRowsAndColumnsOperator(frameSegment); // On the operator created above add the operators present in the query that we want to chain + // previous shuffle has partitioning + // need to remove partitioning here + for (OperatorFactory of : operatorFactoryList) { op = of.wrap(op); } @@ -173,6 +176,9 @@ protected ReturnOrAwait runWithInputChannel(ReadableFrameChannel inputChan public Operator.Signal push(RowsAndColumns rac) { // convert the rac to a row-based frame + // Note that there can be multiple racs here + // one per partition + // But one rac will be written to 1 frame Pair pairFrames = materializeRacToRowFrames(rac, outputStageSignature); Frame f = Frame.wrap(pairFrames.lhs); try { @@ -236,7 +242,6 @@ public Pair materializeRacToRowFrames(RowsAndColumns rac, final ColumnSelectorFactoryMaker csfm = ColumnSelectorFactoryMaker.fromRAC(rac); final ColumnSelectorFactory selectorFactory = csfm.make(rowId); - //columnsToGenerate.addAll(rac.getColumnNames()); ArrayList columnsToGenerate = new ArrayList<>(outputSignature.getColumnNames()); @@ -276,7 +281,6 @@ public Pair materializeRacToRowFrames(RowsAndColumns rac, remainingRowsToFetch--; frameWriter.addSelection(); } - return Pair.of(frameWriter.toByteArray(), sigBob.build()); } } diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryFrameProcessorFactory.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryFrameProcessorFactory.java index 5321d6f968d3..f729b5fa9277 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryFrameProcessorFactory.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryFrameProcessorFactory.java @@ -35,6 +35,7 @@ import org.apache.druid.segment.column.RowSignature; import java.util.List; +import java.util.Objects; import java.util.function.Function; @JsonTypeName("window") @@ -42,33 +43,60 @@ public class WindowOperatorQueryFrameProcessorFactory extends BaseLeafFrameProce { private final WindowOperatorQuery query; private final List operatorList; - RowSignature rearrangePoiter; + private final RowSignature stageRowSignature; + + @Override + public boolean equals(Object o) + { + if (this == o) { + return true; + } + if (o == null || getClass() != o.getClass()) { + return false; + } + WindowOperatorQueryFrameProcessorFactory that = (WindowOperatorQueryFrameProcessorFactory) o; + return Objects.equals(query, that.query) + && Objects.equals(operatorList, that.operatorList) + && Objects.equals(stageRowSignature, that.stageRowSignature); + } + + @Override + public int hashCode() + { + return Objects.hash(query, operatorList, stageRowSignature); + } @JsonCreator public WindowOperatorQueryFrameProcessorFactory( @JsonProperty("query") WindowOperatorQuery query, - @JsonProperty("operators") List operatorFactoryList, - RowSignature pointers + @JsonProperty("operatorList") List operatorFactoryList, + @JsonProperty("stageRowSignature") RowSignature stageRowSignature ) { super(query); this.query = Preconditions.checkNotNull(query, "query"); - this.operatorList = operatorFactoryList; - this.rearrangePoiter = pointers; + this.operatorList = Preconditions.checkNotNull(operatorFactoryList, "bad operator"); + this.stageRowSignature = Preconditions.checkNotNull(stageRowSignature, "stageSignature"); } - @JsonProperty + @JsonProperty("query") public WindowOperatorQuery getQuery() { return query; } - @JsonProperty + @JsonProperty("operatorList") public List getOperators() { return operatorList; } + @JsonProperty("stageRowSignature") + public RowSignature getSignature() + { + return stageRowSignature; + } + @Override protected FrameProcessor makeProcessor( ReadableInput baseInput, @@ -86,7 +114,17 @@ protected FrameProcessor makeProcessor( segmentMapFn, outputChannelHolder, frameWriterFactoryHolder, - rearrangePoiter + stageRowSignature ); } + + @Override + public String toString() + { + return "WindowOperatorQueryFrameProcessorFactory{" + + "query=" + query + + ", operatorList=" + operatorList + + ", stageRowSignature=" + stageRowSignature + + '}'; + } } diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQInsertTest.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQInsertTest.java index edbf7216dc01..f1415aa44af4 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQInsertTest.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQInsertTest.java @@ -282,6 +282,102 @@ public void testInsertWithUnnestWithVirtualColumns() } + @Test + public void testInsertWithWindow() + { + List expectedRows = ImmutableList.of( + new Object[]{946684800000L, 1.0f, 1.0}, + new Object[]{946771200000L, 2.0f, 2.0}, + new Object[]{946857600000L, 3.0f, 3.0}, + new Object[]{978307200000L, 4.0f, 4.0}, + new Object[]{978393600000L, 5.0f, 5.0}, + new Object[]{978480000000L, 6.0f, 6.0} + ); + + RowSignature rowSignature = RowSignature.builder() + .add("__time", ColumnType.LONG) + .add("m1", ColumnType.FLOAT) + .add("summ1", ColumnType.DOUBLE) + .build(); + + + testIngestQuery().setSql( + "insert into foo1 SELECT __time, m1,\n" + + "SUM(m1) OVER(PARTITION BY m1) as summ1\n" + + "from foo\n" + + "GROUP BY __time, m1 PARTITIONED BY ALL") + .setQueryContext(context) + .setExpectedResultRows(expectedRows) + .setExpectedDataSource("foo1") + .setExpectedRowSignature(rowSignature) + .verifyResults(); + + } + + @Test + public void testInsertWithWindowEmptyOver() + { + List expectedRows = ImmutableList.of( + new Object[]{946684800000L, 1.0f, 21.0}, + new Object[]{946771200000L, 2.0f, 21.0}, + new Object[]{946857600000L, 3.0f, 21.0}, + new Object[]{978307200000L, 4.0f, 21.0}, + new Object[]{978393600000L, 5.0f, 21.0}, + new Object[]{978480000000L, 6.0f, 21.0} + ); + + RowSignature rowSignature = RowSignature.builder() + .add("__time", ColumnType.LONG) + .add("m1", ColumnType.FLOAT) + .add("summ1", ColumnType.DOUBLE) + .build(); + + + testIngestQuery().setSql( + "insert into foo1 SELECT __time, m1,\n" + + "SUM(m1) OVER() as summ1\n" + + "from foo\n" + + "GROUP BY __time, m1 PARTITIONED BY ALL") + .setQueryContext(context) + .setExpectedResultRows(expectedRows) + .setExpectedDataSource("foo1") + .setExpectedRowSignature(rowSignature) + .verifyResults(); + + } + + @Test + public void testInsertWithWindowPartitionByOrderBy() + { + List expectedRows = ImmutableList.of( + new Object[]{946684800000L, 1.0f, 1.0}, + new Object[]{946771200000L, 2.0f, 2.0}, + new Object[]{946857600000L, 3.0f, 3.0}, + new Object[]{978307200000L, 4.0f, 4.0}, + new Object[]{978393600000L, 5.0f, 5.0}, + new Object[]{978480000000L, 6.0f, 6.0} + ); + + RowSignature rowSignature = RowSignature.builder() + .add("__time", ColumnType.LONG) + .add("m1", ColumnType.FLOAT) + .add("summ1", ColumnType.DOUBLE) + .build(); + + + testIngestQuery().setSql( + "insert into foo1 SELECT __time, m1,\n" + + "SUM(m1) OVER(PARTITION BY m1 ORDER BY m1 ASC) as summ1\n" + + "from foo\n" + + "GROUP BY __time, m1 PARTITIONED BY ALL") + .setQueryContext(context) + .setExpectedResultRows(expectedRows) + .setExpectedDataSource("foo1") + .setExpectedRowSignature(rowSignature) + .verifyResults(); + + } + @Test public void testInsertOnExternalDataSource() throws IOException { diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQSelectTest.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQSelectTest.java index 484a1828d9c7..28810f011670 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQSelectTest.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQSelectTest.java @@ -137,12 +137,12 @@ public class MSQSelectTest extends MSQTestBase public static Collection data() { Object[][] data = new Object[][]{ - {DEFAULT, DEFAULT_MSQ_CONTEXT}, - {DURABLE_STORAGE, DURABLE_STORAGE_MSQ_CONTEXT}, - {FAULT_TOLERANCE, FAULT_TOLERANCE_MSQ_CONTEXT}, - {PARALLEL_MERGE, PARALLEL_MERGE_MSQ_CONTEXT}, - {QUERY_RESULTS_WITH_DURABLE_STORAGE, QUERY_RESULTS_WITH_DURABLE_STORAGE_CONTEXT}, - {QUERY_RESULTS_WITH_DEFAULT, QUERY_RESULTS_WITH_DEFAULT_CONTEXT} + {DEFAULT, DEFAULT_MSQ_CONTEXT} +// {DURABLE_STORAGE, DURABLE_STORAGE_MSQ_CONTEXT}, +// {FAULT_TOLERANCE, FAULT_TOLERANCE_MSQ_CONTEXT}, +// {PARALLEL_MERGE, PARALLEL_MERGE_MSQ_CONTEXT}, +// {QUERY_RESULTS_WITH_DURABLE_STORAGE, QUERY_RESULTS_WITH_DURABLE_STORAGE_CONTEXT}, +// {QUERY_RESULTS_WITH_DEFAULT, QUERY_RESULTS_WITH_DEFAULT_CONTEXT} }; return Arrays.asList(data); @@ -705,11 +705,13 @@ public void testWindowOnFoo() } @Test - public void testWindowOnFooWith2Windows() + public void testWindowOnFooWith2WindowsV1() { RowSignature rowSignature = RowSignature.builder() .add("m1", ColumnType.FLOAT) - .add("cc", ColumnType.DOUBLE) + .add("m2", ColumnType.DOUBLE) + .add("summ2", ColumnType.DOUBLE) + .add("summ1", ColumnType.DOUBLE) .build(); final Query groupByQuery = GroupByQuery.builder() @@ -722,6 +724,11 @@ public void testWindowOnFooWith2Windows() "m1", "d0", ColumnType.FLOAT + ), + new DefaultDimensionSpec( + "m2", + "d1", + ColumnType.DOUBLE ) )) .setContext(context) @@ -730,24 +737,143 @@ public void testWindowOnFooWith2Windows() final WindowFrame theFrame = new WindowFrame(WindowFrame.PeerType.ROWS, true, 0, true, 0); final AggregatorFactory[] theAggs = { - new DoubleSumAggregatorFactory("w0", "d0") + new DoubleSumAggregatorFactory("w0", "d1") + }; + final AggregatorFactory[] nextAggs = { + new DoubleSumAggregatorFactory("w1", "d0") }; WindowFramedAggregateProcessor proc = new WindowFramedAggregateProcessor(theFrame, theAggs); + WindowFramedAggregateProcessor proc1 = new WindowFramedAggregateProcessor(theFrame, nextAggs); final WindowOperatorQuery query = new WindowOperatorQuery( new QueryDataSource(groupByQuery), new LegacySegmentSpec(Intervals.ETERNITY), context, - RowSignature.builder().add("d0", ColumnType.FLOAT).add("w0", ColumnType.DOUBLE).build(), + RowSignature.builder() + .add("d0", ColumnType.FLOAT) + .add("d1", ColumnType.DOUBLE) + .add("w0", ColumnType.DOUBLE) + .add("w1", ColumnType.DOUBLE) + .build(), ImmutableList.of( new NaivePartitioningOperatorFactory(ImmutableList.of("d0")), - new WindowOperatorFactory(proc) + new WindowOperatorFactory(proc), + new NaivePartitioningOperatorFactory(ImmutableList.of()), + new WindowOperatorFactory(proc1) + ), + null + ); + testSelectQuery() + .setSql("SELECT m1, m2,\n" + + "SUM(m2) OVER(PARTITION BY m1) as summ2\n" + + ",SUM(m1) OVER() as summ1\n" + + "from foo\n" + + "GROUP BY m1,m2") + .setExpectedMSQSpec(MSQSpec.builder() + .query(query) + .columnMappings( + new ColumnMappings(ImmutableList.of( + new ColumnMapping("d0", "m1"), + new ColumnMapping("d1", "m2"), + new ColumnMapping("w0", "summ2"), + new ColumnMapping("w1", "summ1") + ) + )) + .tuningConfig(MSQTuningConfig.defaultConfig()) + .destination(isDurableStorageDestination() + ? DurableStorageMSQDestination.INSTANCE + : TaskReportMSQDestination.INSTANCE) + .build()) + .setExpectedRowSignature(rowSignature) + .setExpectedResultRows(ImmutableList.of( + new Object[]{1.0f, 1.0, 1.0, 1.0}, + new Object[]{2.0f, 2.0, 2.0, 2.0}, + new Object[]{3.0f, 3.0, 3.0, 3.0}, + new Object[]{4.0f, 4.0, 4.0, 4.0}, + new Object[]{5.0f, 5.0, 5.0, 5.0}, + new Object[]{6.0f, 6.0, 6.0, 6.0} + )) + .setQueryContext(context) + .setExpectedCountersForStageWorkerChannel( + CounterSnapshotMatcher + .with().totalFiles(1), + 0, 0, "input0" + ) + .setExpectedCountersForStageWorkerChannel( + CounterSnapshotMatcher + .with().rows(6).frames(1), + 0, 0, "output" + ) + .setExpectedCountersForStageWorkerChannel( + CounterSnapshotMatcher + .with().rows(6).frames(1), + 0, 0, "shuffle" + ) + .verifyResults(); + } + + @Test + public void testWindowOnFooWith2WindowsV2() + { + RowSignature rowSignature = RowSignature.builder() + .add("m1", ColumnType.FLOAT) + .add("m2", ColumnType.DOUBLE) + .add("summ2", ColumnType.DOUBLE) + .add("summ1", ColumnType.DOUBLE) + .build(); + + final Query groupByQuery = GroupByQuery.builder() + .setDataSource(CalciteTests.DATASOURCE1) + .setInterval(querySegmentSpec(Filtration + .eternity())) + .setGranularity(Granularities.ALL) + .setDimensions(dimensions( + new DefaultDimensionSpec( + "m1", + "d0", + ColumnType.FLOAT + ), + new DefaultDimensionSpec( + "m2", + "d1", + ColumnType.DOUBLE + ) + )) + .setContext(context) + .build(); + + + final WindowFrame theFrame = new WindowFrame(WindowFrame.PeerType.ROWS, true, 0, true, 0); + final AggregatorFactory[] theAggs = { + new DoubleSumAggregatorFactory("w0", "d1") + }; + final AggregatorFactory[] nextAggs = { + new DoubleSumAggregatorFactory("w1", "d0") + }; + WindowFramedAggregateProcessor proc = new WindowFramedAggregateProcessor(theFrame, theAggs); + WindowFramedAggregateProcessor proc1 = new WindowFramedAggregateProcessor(theFrame, nextAggs); + + final WindowOperatorQuery query = new WindowOperatorQuery( + new QueryDataSource(groupByQuery), + new LegacySegmentSpec(Intervals.ETERNITY), + context, + RowSignature.builder() + .add("d0", ColumnType.FLOAT) + .add("d1", ColumnType.DOUBLE) + .add("w0", ColumnType.DOUBLE) + .add("w1", ColumnType.DOUBLE) + .build(), + ImmutableList.of( + new NaivePartitioningOperatorFactory(ImmutableList.of("d0")), + new WindowOperatorFactory(proc), + new NaivePartitioningOperatorFactory(ImmutableList.of("d1")), + new WindowOperatorFactory(proc1) ), null ); testSelectQuery() .setSql("SELECT m1, m2,\n" - + "SUM(m2) OVER() as summ2\n" + + "SUM(m2) OVER(PARTITION BY m1) as summ2\n" + ",SUM(m1) OVER(PARTITION BY m2) as summ1\n" + "from foo\n" + "GROUP BY m1,m2") @@ -756,7 +882,9 @@ public void testWindowOnFooWith2Windows() .columnMappings( new ColumnMappings(ImmutableList.of( new ColumnMapping("d0", "m1"), - new ColumnMapping("w0", "cc") + new ColumnMapping("d1", "m2"), + new ColumnMapping("w0", "summ2"), + new ColumnMapping("w1", "summ1") ) )) .tuningConfig(MSQTuningConfig.defaultConfig()) @@ -766,12 +894,12 @@ public void testWindowOnFooWith2Windows() .build()) .setExpectedRowSignature(rowSignature) .setExpectedResultRows(ImmutableList.of( - new Object[]{1.0f, 1.0}, - new Object[]{2.0f, 2.0}, - new Object[]{3.0f, 3.0}, - new Object[]{4.0f, 4.0}, - new Object[]{5.0f, 5.0}, - new Object[]{6.0f, 6.0} + new Object[]{1.0f, 1.0, 1.0, 1.0}, + new Object[]{2.0f, 2.0, 2.0, 2.0}, + new Object[]{3.0f, 3.0, 3.0, 3.0}, + new Object[]{4.0f, 4.0, 4.0, 4.0}, + new Object[]{5.0f, 5.0, 5.0, 5.0}, + new Object[]{6.0f, 6.0, 6.0, 6.0} )) .setQueryContext(context) .setExpectedCountersForStageWorkerChannel( From ab6e317e9a6e62fba79034be4b17554df898ac56 Mon Sep 17 00:00:00 2001 From: Soumyava Das Date: Tue, 9 Jan 2024 11:23:22 +0530 Subject: [PATCH 09/42] Updating for first set of reviews --- .../WindowOperatorQueryFrameProcessor.java | 128 +++++++++--------- .../apache/druid/msq/exec/MSQSelectTest.java | 12 +- 2 files changed, 67 insertions(+), 73 deletions(-) diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryFrameProcessor.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryFrameProcessor.java index a28048029fdf..897973810119 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryFrameProcessor.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryFrameProcessor.java @@ -23,8 +23,6 @@ import com.google.common.collect.Iterables; import org.apache.druid.collections.ResourceHolder; import org.apache.druid.frame.Frame; -import org.apache.druid.frame.FrameType; -import org.apache.druid.frame.allocation.ArenaMemoryAllocatorFactory; import org.apache.druid.frame.channel.FrameWithPartition; import org.apache.druid.frame.channel.ReadableFrameChannel; import org.apache.druid.frame.channel.WritableFrameChannel; @@ -33,8 +31,6 @@ import org.apache.druid.frame.util.SettableLongVirtualColumn; import org.apache.druid.frame.write.FrameWriter; import org.apache.druid.frame.write.FrameWriterFactory; -import org.apache.druid.frame.write.FrameWriters; -import org.apache.druid.java.util.common.Pair; import org.apache.druid.java.util.common.Unit; import org.apache.druid.java.util.common.io.Closer; import org.apache.druid.java.util.common.logger.Logger; @@ -46,7 +42,6 @@ import org.apache.druid.query.operator.WindowOperatorQuery; import org.apache.druid.query.rowsandcols.LazilyDecoratedRowsAndColumns; import org.apache.druid.query.rowsandcols.RowsAndColumns; -import org.apache.druid.query.rowsandcols.column.Column; import org.apache.druid.query.rowsandcols.concrete.RowBasedFrameRowAndColumns; import org.apache.druid.query.rowsandcols.semantic.ColumnSelectorFactoryMaker; import org.apache.druid.segment.ColumnSelectorFactory; @@ -64,7 +59,6 @@ import java.io.IOException; import java.util.ArrayList; import java.util.BitSet; -import java.util.Collections; import java.util.List; import java.util.concurrent.atomic.AtomicInteger; import java.util.function.Function; @@ -81,10 +75,9 @@ public class WindowOperatorQueryFrameProcessor extends BaseLeafFrameProcessor private final VirtualColumns frameWriterVirtualColumns; private final Closer closer = Closer.create(); private final RowSignature outputStageSignature; - + private final SimpleSettableOffset cursorOffset = new SimpleAscendingOffset(Integer.MAX_VALUE); private Cursor cursor; private Segment segment; - private final SimpleSettableOffset cursorOffset = new SimpleAscendingOffset(Integer.MAX_VALUE); private FrameWriter frameWriter; private long currentAllocatorCapacity; // Used for generating FrameRowTooLargeException if needed @@ -124,6 +117,32 @@ public WindowOperatorQueryFrameProcessor( this.frameWriterVirtualColumns = VirtualColumns.create(frameWriterVirtualColumns); } + private static Operator getOperator(RowBasedFrameRowAndColumns frameRowsAndColumns) + { + LazilyDecoratedRowsAndColumns ldrc = new LazilyDecoratedRowsAndColumns( + frameRowsAndColumns, + null, + null, + null, + OffsetLimit.limit(Integer.MAX_VALUE), + null, + null + ); + // Create an operator on top of the created rows and columns + Operator op = new Operator() + { + @Nullable + @Override + public Closeable goOrContinue(Closeable continuationObject, Receiver receiver) + { + receiver.push(ldrc); + receiver.completed(); + return continuationObject; + } + }; + return op; + } + @Override protected ReturnOrAwait runWithSegment(SegmentWithDescriptor segment) { @@ -179,11 +198,11 @@ public Operator.Signal push(RowsAndColumns rac) // Note that there can be multiple racs here // one per partition // But one rac will be written to 1 frame - Pair pairFrames = materializeRacToRowFrames(rac, outputStageSignature); - Frame f = Frame.wrap(pairFrames.lhs); + AtomicInteger rowId = new AtomicInteger(0); + createFrameWriterIfNeeded(rac, rowId); + materializeRacToRowFrames(rac, outputStageSignature, rowId); try { - Iterables.getOnlyElement(outputChannels()) - .write(new FrameWithPartition(f, FrameWithPartition.NO_PARTITION)); + flushFrameWriter(); } catch (IOException e) { throw new RuntimeException(e); @@ -206,68 +225,15 @@ public void completed() return ReturnOrAwait.runAgain(); } - private static Operator getOperator(RowBasedFrameRowAndColumns frameRowsAndColumns) - { - LazilyDecoratedRowsAndColumns ldrc = new LazilyDecoratedRowsAndColumns( - frameRowsAndColumns, - null, - null, - null, - OffsetLimit.limit(Integer.MAX_VALUE), - null, - null - ); - // Create an operator on top of the created rows and columns - Operator op = new Operator() - { - @Nullable - @Override - public Closeable goOrContinue(Closeable continuationObject, Receiver receiver) - { - receiver.push(ldrc); - receiver.completed(); - return continuationObject; - } - }; - return op; - } - - public Pair materializeRacToRowFrames(RowsAndColumns rac, RowSignature outputSignature) + public void materializeRacToRowFrames(RowsAndColumns rac, RowSignature outputSignature, AtomicInteger rowId) { final int numRows = rac.numRows(); BitSet rowsToSkip = null; - AtomicInteger rowId = new AtomicInteger(0); - final ColumnSelectorFactoryMaker csfm = ColumnSelectorFactoryMaker.fromRAC(rac); - final ColumnSelectorFactory selectorFactory = csfm.make(rowId); - - ArrayList columnsToGenerate = new ArrayList<>(outputSignature.getColumnNames()); - - - final RowSignature.Builder sigBob = RowSignature.builder(); - final ArenaMemoryAllocatorFactory memFactory = new ArenaMemoryAllocatorFactory(200 << 20); - - - for (String column : columnsToGenerate) { - final Column racColumn = rac.findColumn(column); - if (racColumn == null) { - continue; - } - sigBob.add(column, racColumn.toAccessor().getType()); - } - long remainingRowsToSkip = 0; long remainingRowsToFetch = Integer.MAX_VALUE; - final FrameWriter frameWriter; - frameWriter = FrameWriters.makeFrameWriterFactory( - FrameType.ROW_BASED, - memFactory, - sigBob.build(), - Collections.emptyList() - ).newFrameWriter(selectorFactory); - rowId.set(0); for (; rowId.get() < numRows && remainingRowsToFetch > 0; rowId.incrementAndGet()) { final int theId = rowId.get(); @@ -281,6 +247,34 @@ public Pair materializeRacToRowFrames(RowsAndColumns rac, remainingRowsToFetch--; frameWriter.addSelection(); } - return Pair.of(frameWriter.toByteArray(), sigBob.build()); + } + + private void createFrameWriterIfNeeded(RowsAndColumns rac, AtomicInteger rowId) + { + if (frameWriter == null) { + final ColumnSelectorFactoryMaker csfm = ColumnSelectorFactoryMaker.fromRAC(rac); + final ColumnSelectorFactory frameWriterColumnSelectorFactory = csfm.make(rowId); + final FrameWriterFactory frameWriterFactory = getFrameWriterFactory(); + frameWriter = frameWriterFactory.newFrameWriter(frameWriterColumnSelectorFactory); + currentAllocatorCapacity = frameWriterFactory.allocatorCapacity(); + } + } + + private long flushFrameWriter() throws IOException + { + if (frameWriter != null && frameWriter.getNumRows() > 0) { + final Frame frame = Frame.wrap(frameWriter.toByteArray()); + Iterables.getOnlyElement(outputChannels()).write(new FrameWithPartition(frame, FrameWithPartition.NO_PARTITION)); + frameWriter.close(); + frameWriter = null; + return frame.numRows(); + } else { + if (frameWriter != null) { + frameWriter.close(); + frameWriter = null; + } + + return 0; + } } } diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQSelectTest.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQSelectTest.java index 28810f011670..501b6122fc29 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQSelectTest.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQSelectTest.java @@ -137,12 +137,12 @@ public class MSQSelectTest extends MSQTestBase public static Collection data() { Object[][] data = new Object[][]{ - {DEFAULT, DEFAULT_MSQ_CONTEXT} -// {DURABLE_STORAGE, DURABLE_STORAGE_MSQ_CONTEXT}, -// {FAULT_TOLERANCE, FAULT_TOLERANCE_MSQ_CONTEXT}, -// {PARALLEL_MERGE, PARALLEL_MERGE_MSQ_CONTEXT}, -// {QUERY_RESULTS_WITH_DURABLE_STORAGE, QUERY_RESULTS_WITH_DURABLE_STORAGE_CONTEXT}, -// {QUERY_RESULTS_WITH_DEFAULT, QUERY_RESULTS_WITH_DEFAULT_CONTEXT} + {DEFAULT, DEFAULT_MSQ_CONTEXT}, + {DURABLE_STORAGE, DURABLE_STORAGE_MSQ_CONTEXT}, + {FAULT_TOLERANCE, FAULT_TOLERANCE_MSQ_CONTEXT}, + {PARALLEL_MERGE, PARALLEL_MERGE_MSQ_CONTEXT}, + {QUERY_RESULTS_WITH_DURABLE_STORAGE, QUERY_RESULTS_WITH_DURABLE_STORAGE_CONTEXT}, + {QUERY_RESULTS_WITH_DEFAULT, QUERY_RESULTS_WITH_DEFAULT_CONTEXT} }; return Arrays.asList(data); From f1efec32a6fd4eeed4ab424601330e8ad4380b5c Mon Sep 17 00:00:00 2001 From: Soumyava Das Date: Tue, 16 Jan 2024 15:41:48 +0530 Subject: [PATCH 10/42] Changes for partition boundary detection --- .idea/inspectionProfiles/Druid.xml | 21 +- .../druid/msq/querykit/DataSourcePlan.java | 28 +- .../WindowOperatorQueryFrameProcessor.java | 406 +++++++++++++++--- .../msq/querykit/WindowOperatorQueryKit.java | 8 + .../GroupByPostShuffleFrameProcessor.java | 1 - .../msq/querykit/groupby/GroupByQueryKit.java | 29 +- .../apache/druid/msq/exec/MSQSelectTest.java | 3 + .../MapOfColumnsRowsAndColumns.java | 31 ++ .../druid/sql/calcite/CalciteQueryTest.java | 14 + 9 files changed, 453 insertions(+), 88 deletions(-) diff --git a/.idea/inspectionProfiles/Druid.xml b/.idea/inspectionProfiles/Druid.xml index d1f72dd2763f..88ffd72e1098 100644 --- a/.idea/inspectionProfiles/Druid.xml +++ b/.idea/inspectionProfiles/Druid.xml @@ -39,9 +39,13 @@