From 861ee5e986d0a5b8156542e484353c4fee223a15 Mon Sep 17 00:00:00 2001 From: Paul Rogers Date: Sun, 12 Jun 2022 22:52:11 -0700 Subject: [PATCH 01/11] Basics of the operator and fragment mechanisms --- .../org/apache/druid/query/QueryPlus.java | 45 ++- .../java/org/apache/druid/queryng/README.md | 229 +++++++++++++++ .../java/org/apache/druid/queryng/Timer.java | 69 +++++ .../druid/queryng/config/QueryNGConfig.java | 51 ++++ .../druid/queryng/fragment/DAGBuilder.java | 46 +++ .../queryng/fragment/FragmentBuilder.java | 50 ++++ .../fragment/FragmentBuilderFactory.java | 33 +++ .../fragment/FragmentBuilderFactoryImpl.java | 62 ++++ .../queryng/fragment/FragmentBuilderImpl.java | 114 ++++++++ .../queryng/fragment/FragmentContext.java | 68 +++++ .../queryng/fragment/FragmentContextImpl.java | 161 +++++++++++ .../queryng/fragment/FragmentHandle.java | 128 +++++++++ .../queryng/fragment/FragmentHandleImpl.java | 228 +++++++++++++++ .../druid/queryng/fragment/FragmentRun.java | 46 +++ .../queryng/fragment/FragmentRunImpl.java | 79 ++++++ .../fragment/NullFragmentBuilderFactory.java | 35 +++ .../druid/queryng/guice/QueryNGModule.java | 48 ++++ .../queryng/operators/ConcatOperator.java | 100 +++++++ .../queryng/operators/LimitOperator.java | 54 ++++ .../queryng/operators/MappingOperator.java | 71 +++++ .../druid/queryng/operators/NullOperator.java | 56 ++++ .../druid/queryng/operators/Operator.java | 148 ++++++++++ .../druid/queryng/operators/Operators.java | 194 +++++++++++++ .../operators/OrderedMergeOperator.java | 167 +++++++++++ .../queryng/operators/PushBackOperator.java | 98 +++++++ .../queryng/operators/SequenceIterator.java | 81 ++++++ .../queryng/operators/SequenceOperator.java | 105 +++++++ .../queryng/operators/TransformOperator.java | 62 ++++ .../queryng/operators/WrappingOperator.java | 76 +++++ .../general/QueryRunnerOperator.java | 70 +++++ .../queryng/operators/ConcatOperatorTest.java | 136 +++++++++ .../druid/queryng/operators/FragmentTest.java | 265 ++++++++++++++++++ .../druid/queryng/operators/MockCursor.java | 205 ++++++++++++++ .../queryng/operators/MockFilterOperator.java | 58 ++++ .../druid/queryng/operators/MockOperator.java | 81 ++++++ .../queryng/operators/MockOperatorTest.java | 105 +++++++ .../queryng/operators/MockStorageAdapter.java | 143 ++++++++++ .../operators/PushBackOperatorTest.java | 82 ++++++ .../general/OrderedMergeOperatorTest.java | 140 +++++++++ 39 files changed, 3980 insertions(+), 9 deletions(-) create mode 100644 processing/src/main/java/org/apache/druid/queryng/README.md create mode 100644 processing/src/main/java/org/apache/druid/queryng/Timer.java create mode 100644 processing/src/main/java/org/apache/druid/queryng/config/QueryNGConfig.java create mode 100644 processing/src/main/java/org/apache/druid/queryng/fragment/DAGBuilder.java create mode 100644 processing/src/main/java/org/apache/druid/queryng/fragment/FragmentBuilder.java create mode 100644 processing/src/main/java/org/apache/druid/queryng/fragment/FragmentBuilderFactory.java create mode 100644 processing/src/main/java/org/apache/druid/queryng/fragment/FragmentBuilderFactoryImpl.java create mode 100644 processing/src/main/java/org/apache/druid/queryng/fragment/FragmentBuilderImpl.java create mode 100644 processing/src/main/java/org/apache/druid/queryng/fragment/FragmentContext.java create mode 100644 processing/src/main/java/org/apache/druid/queryng/fragment/FragmentContextImpl.java create mode 100644 processing/src/main/java/org/apache/druid/queryng/fragment/FragmentHandle.java create mode 100644 processing/src/main/java/org/apache/druid/queryng/fragment/FragmentHandleImpl.java create mode 100644 processing/src/main/java/org/apache/druid/queryng/fragment/FragmentRun.java create mode 100644 processing/src/main/java/org/apache/druid/queryng/fragment/FragmentRunImpl.java create mode 100644 processing/src/main/java/org/apache/druid/queryng/fragment/NullFragmentBuilderFactory.java create mode 100644 processing/src/main/java/org/apache/druid/queryng/guice/QueryNGModule.java create mode 100644 processing/src/main/java/org/apache/druid/queryng/operators/ConcatOperator.java create mode 100644 processing/src/main/java/org/apache/druid/queryng/operators/LimitOperator.java create mode 100644 processing/src/main/java/org/apache/druid/queryng/operators/MappingOperator.java create mode 100644 processing/src/main/java/org/apache/druid/queryng/operators/NullOperator.java create mode 100644 processing/src/main/java/org/apache/druid/queryng/operators/Operator.java create mode 100644 processing/src/main/java/org/apache/druid/queryng/operators/Operators.java create mode 100644 processing/src/main/java/org/apache/druid/queryng/operators/OrderedMergeOperator.java create mode 100644 processing/src/main/java/org/apache/druid/queryng/operators/PushBackOperator.java create mode 100644 processing/src/main/java/org/apache/druid/queryng/operators/SequenceIterator.java create mode 100644 processing/src/main/java/org/apache/druid/queryng/operators/SequenceOperator.java create mode 100644 processing/src/main/java/org/apache/druid/queryng/operators/TransformOperator.java create mode 100644 processing/src/main/java/org/apache/druid/queryng/operators/WrappingOperator.java create mode 100644 processing/src/main/java/org/apache/druid/queryng/operators/general/QueryRunnerOperator.java create mode 100644 processing/src/test/java/org/apache/druid/queryng/operators/ConcatOperatorTest.java create mode 100644 processing/src/test/java/org/apache/druid/queryng/operators/FragmentTest.java create mode 100644 processing/src/test/java/org/apache/druid/queryng/operators/MockCursor.java create mode 100644 processing/src/test/java/org/apache/druid/queryng/operators/MockFilterOperator.java create mode 100644 processing/src/test/java/org/apache/druid/queryng/operators/MockOperator.java create mode 100644 processing/src/test/java/org/apache/druid/queryng/operators/MockOperatorTest.java create mode 100644 processing/src/test/java/org/apache/druid/queryng/operators/MockStorageAdapter.java create mode 100644 processing/src/test/java/org/apache/druid/queryng/operators/PushBackOperatorTest.java create mode 100644 processing/src/test/java/org/apache/druid/queryng/operators/general/OrderedMergeOperatorTest.java diff --git a/processing/src/main/java/org/apache/druid/query/QueryPlus.java b/processing/src/main/java/org/apache/druid/query/QueryPlus.java index 1b18e9439099..e9001a7f3600 100644 --- a/processing/src/main/java/org/apache/druid/query/QueryPlus.java +++ b/processing/src/main/java/org/apache/druid/query/QueryPlus.java @@ -24,6 +24,7 @@ import org.apache.druid.guice.annotations.PublicApi; import org.apache.druid.java.util.common.guava.Sequence; import org.apache.druid.query.context.ResponseContext; +import org.apache.druid.queryng.fragment.FragmentBuilder; import javax.annotation.Nullable; @@ -40,18 +41,24 @@ public final class QueryPlus public static QueryPlus wrap(Query query) { Preconditions.checkNotNull(query); - return new QueryPlus<>(query, null, null); + return new QueryPlus<>(query, null, null, null); } private final Query query; private final QueryMetrics queryMetrics; private final String identity; + private final FragmentBuilder fragmentBuilder; - private QueryPlus(Query query, QueryMetrics queryMetrics, String identity) + private QueryPlus( + Query query, + QueryMetrics queryMetrics, + String identity, + FragmentBuilder fragmentBuilder) { this.query = query; this.queryMetrics = queryMetrics; this.identity = identity; + this.fragmentBuilder = fragmentBuilder; } public Query getQuery() @@ -71,7 +78,7 @@ public QueryMetrics getQueryMetrics() */ public QueryPlus withIdentity(String identity) { - return new QueryPlus<>(query, queryMetrics, identity); + return new QueryPlus<>(query, queryMetrics, identity, fragmentBuilder); } /** @@ -89,13 +96,13 @@ public QueryPlus withQueryMetrics(QueryToolChest> query if (queryMetrics != null) { return this; } else { - final QueryMetrics metrics = ((QueryToolChest) queryToolChest).makeMetrics(query); + final QueryMetrics metrics = ((QueryToolChest) queryToolChest).makeMetrics(query); if (identity != null) { metrics.identity(identity); } - return new QueryPlus<>(query, metrics, identity); + return new QueryPlus<>(query, metrics, identity, fragmentBuilder); } } @@ -120,7 +127,7 @@ private QueryPlus withoutQueryMetrics() if (queryMetrics == null) { return this; } else { - return new QueryPlus<>(query, null, identity); + return new QueryPlus<>(query, null, identity, fragmentBuilder); } } @@ -132,7 +139,8 @@ public QueryPlus withMaxQueuedBytes(long maxQueuedBytes) return new QueryPlus<>( query.withOverriddenContext(ImmutableMap.of(QueryContexts.MAX_QUEUED_BYTES_KEY, maxQueuedBytes)), queryMetrics, - identity + identity, + fragmentBuilder ); } @@ -141,7 +149,7 @@ public QueryPlus withMaxQueuedBytes(long maxQueuedBytes) */ public QueryPlus withQuery(Query replacementQuery) { - return new QueryPlus<>(replacementQuery, queryMetrics, identity); + return new QueryPlus<>(replacementQuery, queryMetrics, identity, fragmentBuilder); } public Sequence run(QuerySegmentWalker walker, ResponseContext context) @@ -151,6 +159,25 @@ public Sequence run(QuerySegmentWalker walker, ResponseContext context) public QueryPlus optimizeForSegment(PerSegmentQueryOptimizationContext optimizationContext) { - return new QueryPlus<>(query.optimizeForSegment(optimizationContext), queryMetrics, identity); + return new QueryPlus<>( + query.optimizeForSegment(optimizationContext), + queryMetrics, + identity, + fragmentBuilder); + } + + /** + * Returns the same QueryPlus object with the fragment builder added. The fragment + * builder enables this query to use the "Next Gen" query engine. The builder + * may be null, which indicates to use the "classic" rather than "NG" engine. + */ + public QueryPlus withFragmentBuilder(FragmentBuilder fragmentBuilder) + { + return new QueryPlus<>(query, queryMetrics, identity, fragmentBuilder); + } + + public FragmentBuilder fragmentBuilder() + { + return fragmentBuilder; } } diff --git a/processing/src/main/java/org/apache/druid/queryng/README.md b/processing/src/main/java/org/apache/druid/queryng/README.md new file mode 100644 index 000000000000..cdb9428193e7 --- /dev/null +++ b/processing/src/main/java/org/apache/druid/queryng/README.md @@ -0,0 +1,229 @@ + + +# Next-Generation Query Engine + +This package, and its children, implements a "next generation" version of the +Druid Broker/Historical query engine designed around the classic "operator +DAG" structure. The implementation is part of the "multi-stage query engine" +(MSQE) project proposed for Druid, and focuses on the low-latency query path. + +Background information can be found in [Issue #11933](https://github.com/apache/druid/issues/11933). + +This version is experimental and is a "shim" layer that works with much of +the existing structure of the scatter/gather engine. It establishes a +beach-head from which functionality will grow to allow multiple stages +beyond just scatter and gather. + +## Operator Overview + +Druid uses a unique query engine that evolved to serve a specific use case: +native queries with a scatter/gather distribution model. The structure is optimized +for efficiency, but is heavily tied to the scatter-gather model. As we move toward +and MSQE design, we want to revisit the engine architecture. + +Most modern database implementations involve a set of *operators*. The planner works +with a set of *logical operators*, such as the `RelNode` objects in Druid's Calcite +planner. The planner then converts these to descriptions of *physical operators* which +comprise the *physical plan*. Distributed engines divide the physical plan into +*fragments* (also called *slices*) which are farmed out to the various execution +nodes. Each fragment consists of a set of one or more operators that execute on the +given node. The execution node converts the fragment descriptions into a set of +concrete operators (called by many names) which then run to perform the desired +operations. + +The resulting structure forms a *DAG* (directed acyclic graph), typically in the form +of a tree. A root fragment drives the query. The root fragment lives in the Druid +Broker. Other fragments live in execution nodes (Historicals for Druid.) The set +of fragments forms a tree. In classic Druid, it is a two-level tree: the root fragment +represent the "gather" step, a single set of distributed leaf nodes forms the +"scatter" step. In MSQE there can be additional, intermediate fragments. + +Each fragment is itself composed of a DAG of operators, where an operator is a +single relational operation: filter, merge, limit, etc. Each operator reads from +one or more "upstream" operators, does one task, and sends results to the "downstream" +operators. Query engines typically use the "Volcano" iterator-based structure: +each operator has a `next()` method which returns the next results. For performance, +`next()` typically returns a *batch* of rows rather than a single row. In an RDBMS, +the batch might correspond to a single buffer. In Druid, it corresponds to an array +of rows. + +The result is that each operator is an interator that returns batches of rows until +EOF. The `Operator` interface has a very simple protocol: + +```java +public interface Operator +{ + Iterator open(FragmentContext context); + void close(boolean cascade); +} +``` + +Opening an operator provides an iterator to read results. Closing an operator releases +resources. Notice that the operator returns an iterator, but isn't itself an interator. +This structure allows an operator to be a "pass-through": it does nothing per-row but +only does things at open and close time. + +See `Operator`, `NullOperator` and `MockOperator` for very simple examples of the +idea. See `MockOperatorTest` to see how operators work in practice. + +## Constast with Sequences and Query Runners + +Druid has historically used an engine architecture based on the `QueryRunner` and +`Sequence` classes. A `QueryRunner` takes a native query, does some portion of the +work, and delegates the rest to a "base" or "delegate" query runner. The result is a +structure that roughly mimics that of operator. However, query runners are more +complex than an operator: query runners combine planning and execution, and are +often tightly coupled to their implementation and to their upstream (delegate) query +runners. Operators, by contrast, handle only execution: planning is done by a separate +abstraction. Operators are agnostic about their inputs and outputs: operators can be +composed in any number of ways and can be easily tested in isolation. + +Query runners return a `Sequence` which returns results on demand. The query runner is +really a mechanism for building the required `Sequence`, often via a set of lambdas +and anonyous inner classes. The result is that the `Sequence` is tightly coupled with +its query runner. The `Sequence` is also similar to an operator in that it returns +results. An operator, however, is simpler. The `Sequence` class is designed to ensure +that the sequence is always closed and uses a somewhat complex protocol to ensure that. +Operators accomplish the same result by employing a "fragment runner" to create, manage +and close the operators. The result is that the operator version of a task is generally +far simpler than the `Sequence` version. + +The `QueryRunner`/`Sequence` implementation tries to be stateless as much as possible, +which just means that state occurs in closure variables within the query runner +implementation. Such closures are hard to debug and make testing complex. Operators, by +contrast, are shamelessly stateful: they manage, via member variables, whatever state +is required to perform their job. Since operators run within a single thread at any +given time, such state is safe and well-defined. (Generally, in a data pipeline, it +doesn't even make sense for two threads to call the same operator: what does it mean +to do two sorts of the same data in parallel, for example?) + +The general mapping, in this version of the new engine is: + +* `QueryPlanner` corresponds to `QueryRunner`: it figures out which operators are +needed. +* `Operator` corresponds to `Sequence`: it is the physical implementation of a data +transformation task. + +## Support Structure + +The above description touched on several of the support classes that support +operators: + +* `QueryPlanner`: figures out which operators to create given a native query. +* `FragmentRunner`: runs a set of operators, handles failures, and ensures clean-up. +* `FragmentContext`: fragment-wide information shared by all operators. + +## Temporary Shim Structure + +Druid is a mature, complex project — it would be unwise to attempt to upgrade +the entire query stack in one go. Instead, we evolve the code: keep what works well +and replace the bits we need to improve incrementally. + +In this step, we replace query runners and their sequences one-by-one. The result is a +bit tricky: + +* Druid's query stack still produces the set of query runners, and we still invoke `run()`. +on each. +* When the "next gen" engine is enaled, each query runner delegates to the +`QueryPlanner` to work out an operator equivalent for what the query runner previously +did. When the engine is disabled, the query runner does its normal thing. +* The `QueryPlanner` has a modified copy of the query runner code, but in a form that +works out an operator to use. Sometimes the planner determines that no operator is +needed for a particular query. +* The `QueryPlanner` then creates an operator. Since the outer query runner is tasked +with returning a `Sequence`, we then wrap the operator in a `Sequence` so it fits into +the existing protocol. +* Most query runners have an upstream (delegate, base) runner. In that case, the +`QueryPlanner` calls `run()` on the base, converts its `Sequence` to an `Operator` and +hands that to the newly created operator as its input. + +The result is that we've got two parallel structures going: query runners speak `Sequence` +while the `QueryPlanner` speaks `Operator`. To make this work, we use "adapter" to +convert an operator to a sequence and a sequence to an operator. If we literally did +this, every query runner would produce a sequence/operator pair and a chain of two +query runners would look like this: + +```text +|--- query runner 1 ----|----- query runner 2 ----| +Sequence <-- Operator <-- Sequence <-- Operator <-- ... +``` + +Of course, the above would be silly. So, we emply another trick: when we wrap a sequence +in an operator, we check if the sequence in question is one which wraps an operator. If +so, we just discard the intermediate sequence so we get: + +```text +Sequence <-- Operator <-- Operator +``` + +Only the final stage is presented as a sequence: all the internal operations are +structured as operators. + +The result is a structure that plays well with both query runners and operators. When +we have a string of operators: we have an operator stack. When we interface to parts of +Druid that want a sequence, we have a sequence. + +The operator-wrapping sequence ensures that the operator is closed when the sequence +is closed. The operator cascades the close call to its children. This ensures that the +operator is closed as early as possible. + +A fragment context also holds all operators and ensures that they are all closed at +query completion, if not closed previously. In the eventual sructure, this mechanism +will ensure operators close even if an error occurs. + +### Next Steps + +The purpose of the current, interim structure is simply to introduce operators in a +safe, seamless way, and to enable testing. Later, we want to expand the scope. One way +is to expand the operator structure up through the early native query planning states so +that we produce an operator-based structure that cuts out the (now superflorous) query +runnners. + +Later, we can introduce an revised network exchange layer so that we can send a fragment +description to the Historical rather than sending a native query. + +During all of this, we will keep the "classic" engine in place, adding the "next gen" +engine alongside. + +## Configuration + +The "next gen" engine is off by default. Two steps are necessary to enable it: + +1. Set the `druid.queryng.enable` property in a config file or the command line. +2. Set the `queryng` context variable to `true` in each query. + +These steps are temporary since, at this time, the "next gen" query engine is +experimental. + +## Usage + +Once enabled, the next gen engine should "just work." There is no user-visible indication +of the new path. Developers can see the difference in stack traces: if an exception is thrown +in the query engine, or you set a breakpoint, you will see "Operator" methods on the stack +rather than "Sequence" or "Yielder" methods. + +Other than that, results should be identical. At scale (millions of rows), the operator +structure is more efficent. But, in most queries, the bulk of the time is spent in low-level +segment tasks outside the set of operations which have thus far been converted to operators. + +## Tests + +One handy feature of operators is that they can be tested independently. A set of tests +exist for each of the currently-supported operators. diff --git a/processing/src/main/java/org/apache/druid/queryng/Timer.java b/processing/src/main/java/org/apache/druid/queryng/Timer.java new file mode 100644 index 000000000000..b55e1eeab58e --- /dev/null +++ b/processing/src/main/java/org/apache/druid/queryng/Timer.java @@ -0,0 +1,69 @@ +/* + * 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.queryng; + +/** + * Very simple nano-second timer with an on/off switch. + */ +public class Timer +{ + private long totalTime; + private long startTime; + + public static Timer create() + { + return new Timer(); + } + + public static Timer createStarted() + { + Timer timer = create(); + timer.start(); + return timer; + } + + public static Timer createAt(long timeNs) + { + Timer timer = create(); + timer.startTime = timeNs; + return timer; + } + + public void start() + { + if (startTime == 0) { + startTime = System.nanoTime(); + } + } + + public void stop() + { + if (startTime != 0) { + totalTime += System.nanoTime() - startTime; + startTime = 0; + } + } + + public long get() + { + stop(); + return totalTime; + } +} diff --git a/processing/src/main/java/org/apache/druid/queryng/config/QueryNGConfig.java b/processing/src/main/java/org/apache/druid/queryng/config/QueryNGConfig.java new file mode 100644 index 000000000000..3fb0c1121906 --- /dev/null +++ b/processing/src/main/java/org/apache/druid/queryng/config/QueryNGConfig.java @@ -0,0 +1,51 @@ +/* + * 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.queryng.config; + +import com.fasterxml.jackson.annotation.JsonProperty; + +/** + * Configuration for the "NG" query engine. + */ +public class QueryNGConfig +{ + public static final String CONFIG_ROOT = "druid.queryng"; + + /** + * Whether the engine is enabled. It is disabled by default. + */ + @JsonProperty("enabled") + private boolean enabled; + + /** + * Create an instance for testing. + */ + public static QueryNGConfig create(boolean enabled) + { + QueryNGConfig config = new QueryNGConfig(); + config.enabled = enabled; + return config; + } + + public boolean enabled() + { + return enabled; + } +} diff --git a/processing/src/main/java/org/apache/druid/queryng/fragment/DAGBuilder.java b/processing/src/main/java/org/apache/druid/queryng/fragment/DAGBuilder.java new file mode 100644 index 000000000000..aae017807e77 --- /dev/null +++ b/processing/src/main/java/org/apache/druid/queryng/fragment/DAGBuilder.java @@ -0,0 +1,46 @@ +/* + * 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.queryng.fragment; + +import org.apache.druid.queryng.operators.Operator; + +/** + * Interface passed around while building a DAG of operators: + * provides access to the {@link FragmentContext} which each operator + * may use, and a method to register operators with the fragment. + */ +public interface DAGBuilder +{ + FragmentContext context(); + + /** + * Register an operator for this fragment. The operator will be + * closed automatically upon fragment completion both for the success + * and error cases. An operator may be closed earlier, if a + * DAG branch detects it is done during a run. Thus, every operator + * must handle a call to {@code close()} when the operator is already + * closed. + * + * Operators may be registered during a run, which is useful in the + * conversion from query runners as sometimes the query runner decides + * late what child to create. + */ + void register(Operator op); +} diff --git a/processing/src/main/java/org/apache/druid/queryng/fragment/FragmentBuilder.java b/processing/src/main/java/org/apache/druid/queryng/fragment/FragmentBuilder.java new file mode 100644 index 000000000000..f32170eb6f9f --- /dev/null +++ b/processing/src/main/java/org/apache/druid/queryng/fragment/FragmentBuilder.java @@ -0,0 +1,50 @@ +/* + * 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.queryng.fragment; + +import org.apache.druid.java.util.common.guava.Sequence; +import org.apache.druid.query.context.ResponseContext; +import org.apache.druid.queryng.operators.Operator; + +/** + * Build a fragment dynamically. For use during the transition when query runners + * build operators. Once built, the fragment can be run as a root operator, a + * sequence, or materialized into a list. The fragment can be run only once. + */ +public interface FragmentBuilder extends DAGBuilder +{ + FragmentHandle emptyHandle(); + FragmentHandle handle(Operator rootOp); + FragmentHandle handle(Sequence rootOp); + + FragmentRun run(Operator rootOp); + Sequence runAsSequence(Operator rootOp); + + /** + * A simple fragment builder for testing. + */ + static FragmentBuilder defaultBuilder() + { + return new FragmentBuilderImpl( + "unknown", + FragmentContext.NO_TIMEOUT, + ResponseContext.createEmpty()); + } +} diff --git a/processing/src/main/java/org/apache/druid/queryng/fragment/FragmentBuilderFactory.java b/processing/src/main/java/org/apache/druid/queryng/fragment/FragmentBuilderFactory.java new file mode 100644 index 000000000000..02224300807e --- /dev/null +++ b/processing/src/main/java/org/apache/druid/queryng/fragment/FragmentBuilderFactory.java @@ -0,0 +1,33 @@ +/* + * 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.queryng.fragment; + +import org.apache.druid.query.Query; +import org.apache.druid.query.context.ResponseContext; + +import javax.annotation.Nullable; + +public interface FragmentBuilderFactory +{ + @Nullable + FragmentBuilder create( + Query query, + ResponseContext responseContext); +} diff --git a/processing/src/main/java/org/apache/druid/queryng/fragment/FragmentBuilderFactoryImpl.java b/processing/src/main/java/org/apache/druid/queryng/fragment/FragmentBuilderFactoryImpl.java new file mode 100644 index 000000000000..5368fcae66e4 --- /dev/null +++ b/processing/src/main/java/org/apache/druid/queryng/fragment/FragmentBuilderFactoryImpl.java @@ -0,0 +1,62 @@ +/* + * 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.queryng.fragment; + +import org.apache.druid.query.Query; +import org.apache.druid.query.context.ResponseContext; +import org.apache.druid.queryng.config.QueryNGConfig; +import org.apache.druid.queryng.operators.Operators; + +import javax.inject.Inject; + +/** + * Creates a fragment context for the "shim" implementation of the + * NG query engine, but only if the engine is enabled. Queries should + * take the existence of the fragment context as their indication to use + * the NG engine, else stick with the "classic" engine. + */ +public class FragmentBuilderFactoryImpl implements FragmentBuilderFactory +{ + private final QueryNGConfig config; + + @Inject + public FragmentBuilderFactoryImpl(QueryNGConfig config) + { + this.config = config; + } + + @Override + public FragmentBuilder create( + final Query query, + final ResponseContext responseContext) + { + // Engine has to be enabled + if (!config.enabled()) { + return null; + } + // Client must explicitly ask for the engine + if (!Operators.isEnabled(query)) { + return null; + } + // Only then do we create a fragment builder which, implicitly, + // enables the NG engine. + return new FragmentBuilderImpl(query.getId(), 0, responseContext); + } +} diff --git a/processing/src/main/java/org/apache/druid/queryng/fragment/FragmentBuilderImpl.java b/processing/src/main/java/org/apache/druid/queryng/fragment/FragmentBuilderImpl.java new file mode 100644 index 000000000000..936c2bbf5a1c --- /dev/null +++ b/processing/src/main/java/org/apache/druid/queryng/fragment/FragmentBuilderImpl.java @@ -0,0 +1,114 @@ +/* + * 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.queryng.fragment; + +import org.apache.druid.java.util.common.guava.BaseSequence; +import org.apache.druid.java.util.common.guava.Sequence; +import org.apache.druid.query.context.ResponseContext; +import org.apache.druid.queryng.fragment.FragmentHandleImpl.EmptyFragmentHandle; +import org.apache.druid.queryng.fragment.FragmentHandleImpl.FragmentOperatorHandle; +import org.apache.druid.queryng.fragment.FragmentHandleImpl.FragmentSequenceHandle; +import org.apache.druid.queryng.operators.Operator; + +import java.util.Iterator; + +/** + * Constructs a fragment by registering a set of operators. Used during the + * transition when query runners create operators dynamically. + * + * Provides a "handle" to further construct the DAG for code that does not + * have access to this class, and provides a way to run the DAG given a + * root node (which must be one of the registered operators.) + */ +public class FragmentBuilderImpl implements FragmentBuilder +{ + private final FragmentContextImpl context; + + public FragmentBuilderImpl( + final String queryId, + long timeoutMs, + final ResponseContext responseContext) + { + this.context = new FragmentContextImpl(queryId, timeoutMs, responseContext); + } + + /** + * Adds an operator to the list of operators to close. Assumes operators are + * added bottom-up (as is required so that operators are given their inputs) + * so that the last operator in the list is the root we want to execute. + */ + @Override + public void register(Operator op) + { + context.register(op); + } + + @Override + public FragmentContext context() + { + return context; + } + + @Override + public FragmentHandle emptyHandle() + { + return new EmptyFragmentHandle(this); + } + + @Override + public FragmentHandle handle(Operator rootOp) + { + return new FragmentOperatorHandle(this, rootOp); + } + + @Override + public FragmentHandle handle(Sequence rootOp) + { + return new FragmentSequenceHandle(this, rootOp); + } + + @Override + public FragmentRun run(Operator rootOp) + { + return new FragmentRunImpl(context, rootOp); + } + + @Override + public Sequence runAsSequence(Operator rootOp) + { + final FragmentRun run = run(rootOp); + return new BaseSequence>( + new BaseSequence.IteratorMaker>() + { + @Override + public Iterator make() + { + return run.iterator(); + } + + @Override + public void cleanup(Iterator iterFromMake) + { + run.close(); + } + } + ); + } +} diff --git a/processing/src/main/java/org/apache/druid/queryng/fragment/FragmentContext.java b/processing/src/main/java/org/apache/druid/queryng/fragment/FragmentContext.java new file mode 100644 index 000000000000..c9b52d807048 --- /dev/null +++ b/processing/src/main/java/org/apache/druid/queryng/fragment/FragmentContext.java @@ -0,0 +1,68 @@ +/* + * 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.queryng.fragment; + +import org.apache.druid.query.SegmentDescriptor; +import org.apache.druid.query.context.ResponseContext; + +/** + * Provides fragment-level context to operators within a single + * fragment. + */ +public interface FragmentContext extends DAGBuilder +{ + long NO_TIMEOUT = -1; + + enum State + { + START, RUN, FAILED, CLOSED + } + + State state(); + String queryId(); + ResponseContext responseContext(); + + /** + * Checks if a query timeout has occurred. If so, will throw + * an unchecked exception. The operator need not catch this + * exception: the fragment runner will unwind the stack and + * call each operator's {@code close()} method on timeout. + */ + void checkTimeout(); + + void missingSegment(SegmentDescriptor descriptor); + + /** + * Reports the exception, if any, that terminated the fragment. + * Should be non-null only if the state is {@code FAILED}. + */ + Exception exception(); + + /** + * A simple fragment context for testing. + */ + static FragmentContext defaultContext() + { + return new FragmentContextImpl( + "unknown", + NO_TIMEOUT, + ResponseContext.createEmpty()); + } +} diff --git a/processing/src/main/java/org/apache/druid/queryng/fragment/FragmentContextImpl.java b/processing/src/main/java/org/apache/druid/queryng/fragment/FragmentContextImpl.java new file mode 100644 index 000000000000..71313c0a82c2 --- /dev/null +++ b/processing/src/main/java/org/apache/druid/queryng/fragment/FragmentContextImpl.java @@ -0,0 +1,161 @@ +/* + * 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.queryng.fragment; + +import com.google.common.base.Preconditions; +import org.apache.druid.java.util.common.JodaUtils; +import org.apache.druid.java.util.common.StringUtils; +import org.apache.druid.query.QueryTimeoutException; +import org.apache.druid.query.SegmentDescriptor; +import org.apache.druid.query.context.ResponseContext; +import org.apache.druid.queryng.operators.Operator; + +import java.util.ArrayList; +import java.util.Deque; +import java.util.List; +import java.util.concurrent.ConcurrentLinkedDeque; + +public class FragmentContextImpl implements FragmentContext +{ + private final long timeoutMs; + protected final Deque> operators = new ConcurrentLinkedDeque<>(); + private final ResponseContext responseContext; + private final String queryId; + private final long startTimeMillis; + private final long timeoutAt; + protected State state = State.START; + private Exception exception; + + protected FragmentContextImpl( + final String queryId, + long timeoutMs, + final ResponseContext responseContext) + { + this.queryId = queryId; + this.responseContext = responseContext; + this.startTimeMillis = System.currentTimeMillis(); + this.timeoutMs = timeoutMs; + if (timeoutMs > 0) { + this.timeoutAt = startTimeMillis + timeoutMs; + } else { + this.timeoutAt = JodaUtils.MAX_INSTANT; + } + } + + @Override + public FragmentContext context() + { + return this; + } + + @Override + public State state() + { + return state; + } + + @Override + public void register(Operator op) + { + Preconditions.checkState(state == State.START || state == State.RUN); + operators.add(op); + } + + @Override + public Exception exception() + { + return exception; + } + + @Override + public String queryId() + { + return queryId; + } + + @Override + public ResponseContext responseContext() + { + return responseContext; + } + + public void failed(Exception exception) + { + this.exception = exception; + this.state = State.FAILED; + } + + @Override + public void checkTimeout() + { + if (timeoutAt > 0 && System.currentTimeMillis() >= timeoutAt) { + throw new QueryTimeoutException( + StringUtils.nonStrictFormat("Query [%s] timed out after [%d] ms", + queryId, timeoutMs)); + } + } + + protected void recordRunTime() + { + if (timeoutAt == 0) { + return; + } + // This is very likely wrong + responseContext.put( + ResponseContext.Keys.TIMEOUT_AT, + timeoutAt - (System.currentTimeMillis() - startTimeMillis) + ); + } + + @Override + public void missingSegment(SegmentDescriptor descriptor) + { + responseContext.add(ResponseContext.Keys.MISSING_SEGMENTS, descriptor); + } + + /** + * Closes all operators from the leaves to the root. + * As a result, operators must not call their children during + * the {@code close()} call. Errors are collected, but all operators are closed + * regardless of exceptions. + */ + protected void close() + { + if (state == State.START) { + state = State.CLOSED; + } + if (state == State.CLOSED) { + return; + } + List exceptions = new ArrayList<>(); + Operator op; + while ((op = operators.pollFirst()) != null) { + try { + op.close(false); + } + catch (Exception e) { + exceptions.add(e); + } + } + // TODO: Do something with the exceptions + recordRunTime(); + state = State.CLOSED; + } +} diff --git a/processing/src/main/java/org/apache/druid/queryng/fragment/FragmentHandle.java b/processing/src/main/java/org/apache/druid/queryng/fragment/FragmentHandle.java new file mode 100644 index 000000000000..d69311eec78f --- /dev/null +++ b/processing/src/main/java/org/apache/druid/queryng/fragment/FragmentHandle.java @@ -0,0 +1,128 @@ +/* + * 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.queryng.fragment; + +import org.apache.druid.java.util.common.guava.Sequence; +import org.apache.druid.queryng.operators.Operator; + +/** + * Provides a handle to a fragment and the root operator or sequence. + *

+ * Primarily for the "transition period" in which operators work alongside + * query runners, and so the query stack is built as an alternating set + * of sequences and operators. The sequences are mostly "operator wrappers" + * which melt way if the DAG has the structure of: 

+ * operator --> sequence --> operator
+ *
+ *

+ * The operator stack speaks operators, while the "traditional" query runner + * stack speaks sequences. This handle allows converting the DAG from one + * form to the other, and running the stack as either a root operator or + * root sequence. + *

+ * The stack can be built upwards by adding another operator or sequence. + * The {@link #toOperator()} and {@link #toSequence()} methods convert the + * root, as needed, to the form needed as input to the next layer. + *

+ * Again, all this complexity is needed only during the interim period in + * which operators try to "blend in" with sequences. In the longer term, + * when the stack is just operators, this interface will become much + * simpler. + */ +public interface FragmentHandle +{ + FragmentBuilder builder(); + FragmentContext context(); + + /** + * Reports whether the root of this DAG is an operator. + */ + boolean rootIsOperator(); + + /** + * Returns the root operator if this handle is for an operator, + * or if it is for a sequence which wraps an operator. Returns + * {@code null} otherwise. + *

+ * Note: the reliable way to get a root operator is to call + * {code toOperator().rootOperator()}. + */ + Operator rootOperator(); + + /** + * Reports whether the root of this DAG is a sequence. + */ + boolean rootIsSequence(); + + /** + * Returns the root sequence if this handle is for a sequence, else + * {@code null}. Note: the reliable way to get a root sequence is to + * call: {@code toSequence().rootSequence()}. + */ + Sequence rootSequence(); + + /** + * Converts this DAG to one with an operator root. Returns the same + * handle if this is already an operator, unwraps the root sequence + * if possible, else adds an operator on top of the unwrappable + * sequence. + */ + FragmentHandle toOperator(); + + /** + * Converts this DAG to one with a sequence as root. Returns the same + * handle if this is already a sequence, else wraps the root operator + * in a sequence. + */ + FragmentHandle toSequence(); + + /** + * Extend the DAG upward by adding a new operator, which becomes the + * new root operator, if the fragment were opened after this point. + * Can only be applied to a fragment before it is opened. + * + * The operator may change the type of the elements returned by the + * new DAG. + */ + FragmentHandle compose(Operator newRoot); + + /** + * Extend the DAG upward by adding a new sequence, which becomes the + * new root sequence, if the fragment were opened after this point. + * Can only be applied to a fragment before it is opened. + * + * The sequence may change the type of the elements returned by the + * new DAG. + */ + FragmentHandle compose(Sequence newRoot); + + /** + * Run the DAG starting from the current root. Handles the + * sequence-to-operator conversion if needed. + */ + FragmentRun run(); + + /** + * Provide the fragment results as a sequence. The sequence closes the + * fragment at sequence closure. Handles the operator-to-sequence + * conversion if needed. + */ + Sequence runAsSequence(); +} diff --git a/processing/src/main/java/org/apache/druid/queryng/fragment/FragmentHandleImpl.java b/processing/src/main/java/org/apache/druid/queryng/fragment/FragmentHandleImpl.java new file mode 100644 index 000000000000..3f4ba3b7e148 --- /dev/null +++ b/processing/src/main/java/org/apache/druid/queryng/fragment/FragmentHandleImpl.java @@ -0,0 +1,228 @@ +/* + * 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.queryng.fragment; + +import com.google.common.base.Preconditions; +import org.apache.druid.java.util.common.ISE; +import org.apache.druid.java.util.common.guava.Sequence; +import org.apache.druid.java.util.common.guava.Sequences; +import org.apache.druid.queryng.fragment.FragmentContext.State; +import org.apache.druid.queryng.operators.NullOperator; +import org.apache.druid.queryng.operators.Operator; +import org.apache.druid.queryng.operators.Operators; + +public abstract class FragmentHandleImpl implements FragmentHandle +{ + protected final FragmentBuilder builder; + + public FragmentHandleImpl(FragmentBuilder builder) + { + this.builder = builder; + } + + @Override + public FragmentBuilder builder() + { + return builder; + } + + @Override + public FragmentContext context() + { + return builder.context(); + } + + @Override + public boolean rootIsSequence() + { + return !rootIsOperator(); + } + + @Override + public Sequence rootSequence() + { + return null; + } + + @Override + public FragmentHandle compose(Operator newRoot) + { + Preconditions.checkState(context().state() == State.START); + return new FragmentOperatorHandle(builder, newRoot); + } + + @Override + public FragmentHandle compose(Sequence newRoot) + { + Preconditions.checkState(context().state() == State.START); + return new FragmentSequenceHandle(builder, newRoot); + } + + protected static class FragmentOperatorHandle extends FragmentHandleImpl + { + private final Operator root; + + public FragmentOperatorHandle(FragmentBuilder builder, Operator root) + { + super(builder); + this.root = root; + } + + @Override + public FragmentRun run() + { + return builder.run(root); + } + + @Override + public boolean rootIsOperator() + { + return true; + } + + @Override + public Operator rootOperator() + { + return root; + } + + @Override + public Sequence runAsSequence() + { + return builder.runAsSequence(root); + } + + @Override + public FragmentHandle toOperator() + { + return this; + } + + @Override + public FragmentHandle toSequence() + { + return new FragmentSequenceHandle(builder, Operators.toSequence(root)); + } + } + + protected static class FragmentSequenceHandle extends FragmentHandleImpl + { + private final Sequence root; + + public FragmentSequenceHandle(FragmentBuilder builder, Sequence root) + { + super(builder); + this.root = root; + } + + @Override + public boolean rootIsOperator() + { + return false; + } + + @Override + public Sequence rootSequence() + { + return root; + } + + @Override + public Operator rootOperator() + { + return Operators.unwrapOperator(root); + } + + @Override + public FragmentHandle toOperator() + { + return new FragmentOperatorHandle(builder, Operators.toOperator(builder, root)); + } + + @Override + public FragmentHandle toSequence() + { + return this; + } + + @Override + public FragmentRun run() + { + return builder.run( + Operators.toOperator(builder, root)); + } + + @Override + public Sequence runAsSequence() + { + return root; + } + } + + protected static class EmptyFragmentHandle extends FragmentHandleImpl + { + public EmptyFragmentHandle(FragmentBuilder builder) + { + super(builder); + } + + @Override + public FragmentRun run() + { + return builder.run(new NullOperator(context())); + } + + @Override + public boolean rootIsOperator() + { + return false; + } + + @Override + public boolean rootIsSequence() + { + return false; + } + + @Override + public Operator rootOperator() + { + return null; + } + + @Override + public Sequence runAsSequence() + { + return Sequences.empty(); + } + + @Override + public FragmentHandle toOperator() + { + return this; + } + + @Override + public FragmentHandle toSequence() + { + throw new ISE("Cannot convert an empty fragment to a sequence"); + } + } +} diff --git a/processing/src/main/java/org/apache/druid/queryng/fragment/FragmentRun.java b/processing/src/main/java/org/apache/druid/queryng/fragment/FragmentRun.java new file mode 100644 index 000000000000..e63508ee5253 --- /dev/null +++ b/processing/src/main/java/org/apache/druid/queryng/fragment/FragmentRun.java @@ -0,0 +1,46 @@ +/* + * 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.queryng.fragment; + +import java.util.List; + +/** + * Runs the DAG. The fragment is opened (open is called on the root + * operator) upon construction. Callers can then obtain the iterator, + * or convert the DAG to a sequence or list. Callers must close + * this object at the end of the run. + * + * Callers should generally use only one of the access methods: obtain + * the iterator, a sequence, or convert the results to a list. The + * fragment is not reentrant: results can be obtained only once. + */ +public interface FragmentRun extends AutoCloseable, Iterable +{ + FragmentContext context(); + + /** + * Materializes the entire result set as a list. Primarily for testing. + * Opens the fragment, reads results, and closes the fragment. + */ + List toList(); + + @Override + void close(); +} diff --git a/processing/src/main/java/org/apache/druid/queryng/fragment/FragmentRunImpl.java b/processing/src/main/java/org/apache/druid/queryng/fragment/FragmentRunImpl.java new file mode 100644 index 000000000000..1b0973866501 --- /dev/null +++ b/processing/src/main/java/org/apache/druid/queryng/fragment/FragmentRunImpl.java @@ -0,0 +1,79 @@ +/* + * 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.queryng.fragment; + +import com.google.common.base.Preconditions; +import com.google.common.collect.Lists; +import org.apache.druid.queryng.fragment.FragmentContext.State; +import org.apache.druid.queryng.operators.Operator; + +import java.util.Iterator; +import java.util.List; + +public class FragmentRunImpl implements FragmentRun +{ + private final FragmentContextImpl context; + private Iterator rootIter; + + public FragmentRunImpl(FragmentContextImpl context, Operator root) + { + Preconditions.checkState(context.state == State.START); + this.context = context; + try { + rootIter = root.open(); + context.state = State.RUN; + } + catch (Exception e) { + context.failed(e); + context.state = State.FAILED; + throw e; + } + } + + @Override + public Iterator iterator() + { + Preconditions.checkState(context.state == State.RUN); + return rootIter; + } + + @Override + public FragmentContext context() + { + return context; + } + + @Override + public List toList() + { + try { + return Lists.newArrayList(this); + } + finally { + close(); + } + } + + @Override + public void close() + { + context.close(); + } +} diff --git a/processing/src/main/java/org/apache/druid/queryng/fragment/NullFragmentBuilderFactory.java b/processing/src/main/java/org/apache/druid/queryng/fragment/NullFragmentBuilderFactory.java new file mode 100644 index 000000000000..638235a9cd2b --- /dev/null +++ b/processing/src/main/java/org/apache/druid/queryng/fragment/NullFragmentBuilderFactory.java @@ -0,0 +1,35 @@ +/* + * 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.queryng.fragment; + +import org.apache.druid.query.Query; +import org.apache.druid.query.context.ResponseContext; + +/** + * Do-nothing version for testing. Tests do not (yet) use the NG engine. + */ +public class NullFragmentBuilderFactory implements FragmentBuilderFactory +{ + @Override + public FragmentBuilder create(Query query, ResponseContext responseContext) + { + return null; + } +} diff --git a/processing/src/main/java/org/apache/druid/queryng/guice/QueryNGModule.java b/processing/src/main/java/org/apache/druid/queryng/guice/QueryNGModule.java new file mode 100644 index 000000000000..1d1c351da7e9 --- /dev/null +++ b/processing/src/main/java/org/apache/druid/queryng/guice/QueryNGModule.java @@ -0,0 +1,48 @@ +/* + * 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.queryng.guice; + +import com.google.inject.Binder; +import com.google.inject.Module; +import org.apache.druid.guice.JsonConfigProvider; +import org.apache.druid.guice.LazySingleton; +import org.apache.druid.queryng.config.QueryNGConfig; +import org.apache.druid.queryng.fragment.FragmentBuilderFactory; +import org.apache.druid.queryng.fragment.FragmentBuilderFactoryImpl; + +/** + * Configure the "shim" version of the NG query engine which entails + * creating a config (to enable or disable the engine) and to create + * a factory for the fragment context. In this early version, all + * other parts of the engine are distributed across various query + * runners. + */ +public class QueryNGModule implements Module +{ + @Override + public void configure(Binder binder) + { + JsonConfigProvider.bind(binder, QueryNGConfig.CONFIG_ROOT, QueryNGConfig.class); + binder + .bind(FragmentBuilderFactory.class) + .to(FragmentBuilderFactoryImpl.class) + .in(LazySingleton.class); + } +} diff --git a/processing/src/main/java/org/apache/druid/queryng/operators/ConcatOperator.java b/processing/src/main/java/org/apache/druid/queryng/operators/ConcatOperator.java new file mode 100644 index 000000000000..658836fd65b5 --- /dev/null +++ b/processing/src/main/java/org/apache/druid/queryng/operators/ConcatOperator.java @@ -0,0 +1,100 @@ +/* + * 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.queryng.operators; + +import com.google.common.base.Preconditions; +import org.apache.druid.queryng.fragment.FragmentContext; +import org.apache.druid.queryng.operators.Operator.IterableOperator; + +import java.util.Iterator; +import java.util.List; + +/** + * Concatenate a series of inputs. Simply returns the input values, + * does not limit or coalesce batches. Starts child operators as late as + * possible, and closes them as early as possible. + * + * @see {@link org.apache.druid.query.scan.ScanQueryRunnerFactory#mergeRunners} + */ +public class ConcatOperator implements IterableOperator +{ + public static Operator concatOrNot( + FragmentContext context, + List> children) + { + if (children.size() > 1) { + return new ConcatOperator<>(context, children); + } + return children.get(0); + } + + private final Iterator> childIter; + private Operator current; + private Iterator currentIter; + + public ConcatOperator(FragmentContext context, List> children) + { + childIter = children.iterator(); + context.register(this); + } + + @Override + public Iterator open() + { + return this; + } + + @Override + public boolean hasNext() + { + while (true) { + if (current != null) { + if (currentIter.hasNext()) { + return true; + } + current.close(true); + current = null; + currentIter = null; + } + if (!childIter.hasNext()) { + return false; + } + current = childIter.next(); + currentIter = current.open(); + } + } + + @Override + public T next() + { + Preconditions.checkState(currentIter != null, "Missing call to hasNext()?"); + return currentIter.next(); + } + + @Override + public void close(boolean cascade) + { + if (cascade && current != null) { + current.close(cascade); + } + current = null; + currentIter = null; + } +} diff --git a/processing/src/main/java/org/apache/druid/queryng/operators/LimitOperator.java b/processing/src/main/java/org/apache/druid/queryng/operators/LimitOperator.java new file mode 100644 index 000000000000..93864d118f1e --- /dev/null +++ b/processing/src/main/java/org/apache/druid/queryng/operators/LimitOperator.java @@ -0,0 +1,54 @@ +/* + * 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.queryng.operators; + +import org.apache.druid.queryng.fragment.FragmentContext; + +/** + * Limits the results from the input operator to the given number + * of rows. + */ +public abstract class LimitOperator extends MappingOperator +{ + public static final long UNLIMITED = Long.MAX_VALUE; + + protected final long limit; + protected long rowCount; + protected int batchCount; + + public LimitOperator(FragmentContext context, long limit, Operator input) + { + super(context, input); + this.limit = limit; + } + + @Override + public boolean hasNext() + { + return rowCount < limit && super.hasNext(); + } + + @Override + public T next() + { + rowCount++; + return inputIter.next(); + } +} diff --git a/processing/src/main/java/org/apache/druid/queryng/operators/MappingOperator.java b/processing/src/main/java/org/apache/druid/queryng/operators/MappingOperator.java new file mode 100644 index 000000000000..877532030507 --- /dev/null +++ b/processing/src/main/java/org/apache/druid/queryng/operators/MappingOperator.java @@ -0,0 +1,71 @@ +/* + * 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.queryng.operators; + +import com.google.common.base.Preconditions; +import org.apache.druid.queryng.fragment.FragmentContext; +import org.apache.druid.queryng.operators.Operator.IterableOperator; + +import java.util.Iterator; + +/** + * Base class for operators that do a simple mapping of their input + * to their output. Handles the busy-work of managing the (single) + * input operator. + */ +public abstract class MappingOperator implements IterableOperator +{ + protected final FragmentContext context; + private final Operator input; + protected Iterator inputIter; + protected State state = State.START; + + public MappingOperator(FragmentContext context, Operator input) + { + this.context = context; + this.input = input; + context.register(this); + } + + @Override + public Iterator open() + { + Preconditions.checkState(state == State.START); + inputIter = input.open(); + state = State.RUN; + return this; + } + + @Override + public void close(boolean cascade) + { + if (state == State.RUN && cascade) { + input.close(cascade); + } + inputIter = null; + state = State.CLOSED; + } + + @Override + public boolean hasNext() + { + return state == State.RUN && inputIter.hasNext(); + } +} diff --git a/processing/src/main/java/org/apache/druid/queryng/operators/NullOperator.java b/processing/src/main/java/org/apache/druid/queryng/operators/NullOperator.java new file mode 100644 index 000000000000..8165c999dc3b --- /dev/null +++ b/processing/src/main/java/org/apache/druid/queryng/operators/NullOperator.java @@ -0,0 +1,56 @@ +/* + * 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.queryng.operators; + +import com.google.common.base.Preconditions; +import org.apache.druid.queryng.fragment.FragmentContext; + +import java.util.Collections; +import java.util.Iterator; + +/** + * World's simplest operator: does absolutely nothing + * (other than check that the protocol is followed.) Used in + * tests when we want an empty input, and for a fragment that + * somehow ended up with no operators. + */ +public class NullOperator implements Operator +{ + public State state = State.START; + + public NullOperator(FragmentContext context) + { + context.register(this); + } + + @Override + public Iterator open() + { + Preconditions.checkState(state == State.START); + state = State.RUN; + return Collections.emptyIterator(); + } + + @Override + public void close(boolean cascade) + { + state = State.CLOSED; + } +} diff --git a/processing/src/main/java/org/apache/druid/queryng/operators/Operator.java b/processing/src/main/java/org/apache/druid/queryng/operators/Operator.java new file mode 100644 index 000000000000..4e19b913205b --- /dev/null +++ b/processing/src/main/java/org/apache/druid/queryng/operators/Operator.java @@ -0,0 +1,148 @@ +/* + * 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.queryng.operators; + +import java.util.Iterator; + +/** + * An operator is a data pipeline transform: something that changes a stream of + * results in some way. An operator has a very simple lifecycle: + *

+ *

    + *
  • Created.
    • + *
  • Opened (once, bottom up in the operator DAG), which provides an + * iterator over the results to this operator.
    • + *
  • Closed.
    • + *
+ *

+ * Leaf operators produce results, internal operators transform them, and root + * operators do something with the results. Operators know nothing about their + * children other than that they follow the operator protocol and will produce a + * result when asked. Operators must agree on the type of the shared results. + *

+ * Unlike traditional QueryRunners, an operator does not create its + * children: that is the job of the planner that created a tree of operator + * definitions. The operator simply accepts the previously-created children and + * does its thing. + *

+ * Operators are assumed to be stateful since most data transform operations + * involve some kind of state. The {@code start()} and {@code close()} + * methods provide a well-defined way + * to handle resource. Operators are created as a DAG. Unlike a simple iterator, + * operators should not obtain resources in their constructors. Instead, + * they should obtain resources, open files or otherwise start things whirring in + * the {#code start()} method. In some cases, resource may be obtained a bit later, + * in the first call to {@code hasNext()}. In either case, resources should be + * released in {@code close()}. + *

+ * Operators may appear in a branch of the query DAG that can be deferred or + * closed early. Typical example: a UNION operator that starts its first child, + * runs it, closes it, then moves to the next child. This form of operation ensures + * resources are held for the briefest possible period of time. + *

+ * To make this work, operators should cascade start() and + * close() operations to their children. The query runner will + * start the root operator: the root must start is children, and so on. Closing is + * a bit more complex. Operators should close their children by cascading the + * close operation: but only if that call came from a parent operator (as + * indicated by the {@code cascade} parameter set to {@code true}.) + *

+ * The {@link FragmentRunner} will ensure all operators are closed by calling: + * close, from the bottom up, at the completion (successful or otherwise) of the + * query. In this case, the {@code cascade} parameter set to {@code false}, + * and each operator should not cascade this call to its children: the + * fragment runner will do so. + *

+ * Implementations can assume that calls to next() and + * get() occur within a single thread, so state can be maintained + * in normal (not atomic) variables. + *

+ * Operators are one-pass: they are not re-entrant and cannot be restarted. We + * assume that they read ephemeral values: once returned, they are gone. We also + * assume the results are volatile: once read, we may not be able to read the + * same set of values a second time even if we started the whole process over. + *

+ * The type of the values returned by the operator is determined by context. + * Druid native queries use a variety of Java objects: there is no single + * "row" class or interface. + *

+ * Operators do not have a return type parameter. Operators are generally created + * dynamically: that code is far simpler without having to deal with unknown + * types. Even test code will often call {@code assertEquals()} and the like + * which don't need the type. + *

+ * Having {@code open()} return the iterator for results accomplishes two goals. + * First is the minor benefit of ensuring that an operator is opened before + * fetching results. More substantially, this approach allows "wrapper" operators + * which only perform work in the open or close method. For those, the open + * method returns the iterator of the child, avoiding the overhead of pass-through + * calls for each data batch. The wrapper operator will not sit on the data + * path, only on the control (open/close) path. + * + * @param the type of the object (row, batch) returned by {@link #next()}. + */ +public interface Operator +{ + /** + * Convenience interface for an operator which is its own iterator. + */ + public interface IterableOperator extends Operator, Iterator + { + } + + /** + * State used to track the lifecycle of an operator when we + * need to know. + */ + enum State + { + START, RUN, CLOSED + } + + /** + * Called to prepare for processing. Allows the operator to be created early in + * the run, but resources to be obtained as late as possible in the run. An operator + * calls{@code open()} on its children when it is ready to read its input: either + * in the {@code open()} call for simple operators,or later, on demand, for more + * complex operators such as in a merge or union. + */ + Iterator open(); + + /** + * Called at two distinct times. An operator may choose to close a child + * when it is clear that the child will no longer be needed. For example, + * a union might close its first child when it moves onto the second. + *

+ * Operators should handle at least two calls to @{code close()}: an optional + * early close, and a definite final close when the fragment runner shuts + * down. + *

+ * Because the fragment runner will ensure a final close, operators are + * not required to ensure {@code close()} is called on children for odd + * paths, such as errors. + *

+ * If an operator needs to know if a query failed, it can check the status + * in the fragment context for the state of the query (i.e. failed.) + * + * @param cascade {@code false} if this is the final call from the fragment + * runner, {@code true} if it is an "early close" from a parent. + */ + void close(boolean cascade); +} diff --git a/processing/src/main/java/org/apache/druid/queryng/operators/Operators.java b/processing/src/main/java/org/apache/druid/queryng/operators/Operators.java new file mode 100644 index 000000000000..602baf9afcf0 --- /dev/null +++ b/processing/src/main/java/org/apache/druid/queryng/operators/Operators.java @@ -0,0 +1,194 @@ +/* + * 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.queryng.operators; + +import com.google.common.collect.Lists; +import org.apache.druid.java.util.common.guava.BaseSequence; +import org.apache.druid.java.util.common.guava.Sequence; +import org.apache.druid.query.Query; +import org.apache.druid.query.QueryPlus; +import org.apache.druid.query.QueryRunner; +import org.apache.druid.query.scan.ScanQuery; +import org.apache.druid.queryng.fragment.DAGBuilder; +import org.apache.druid.queryng.operators.general.QueryRunnerOperator; + +import java.util.Iterator; +import java.util.List; + +/** + * Utility functions related to operators. + */ +public class Operators +{ + public static final String CONTEXT_VAR = "queryng"; + + /** + * Determine if the Query NG (operator-based) engine is enabled for the + * given query (given as a QueryPlus). Query NG is enabled if the QueryPlus + * includes the fragment context needed by the Query NG engine. + */ + public static boolean enabledFor(final QueryPlus queryPlus) + { + return queryPlus.fragmentBuilder() != null; + } + + /** + * Determine if Query NG should be enabled for the given query; + * that is, if the query should have a fragment context attached. + * At present, Query NG is enabled if the query is a scan query and + * the query has the "queryng" context variable set. The caller + * should already have checked if the Query NG engine is enabled + * globally. If Query NG is enabled for a query, then the caller + * will attach a fragment context to the query's QueryPlus. + */ + public static boolean isEnabled(Query query) + { + // Query has to be of the currently-supported type + if (!(query instanceof ScanQuery)) { + return false; + } + return query.getContextBoolean(CONTEXT_VAR, false); + } + + /** + * Convenience function to open the operator and return its + * iterator as an {@code Iterable}. + */ + public static Iterable toIterable(Operator op) + { + return new Iterable() { + @Override + public Iterator iterator() + { + return op.open(); + } + }; + } + + public static Iterable toIterable(Iterator iter) + { + return new Iterable() { + @Override + public Iterator iterator() + { + return iter; + } + }; + } + + /** + * Wraps an operator in a sequence using the standard base sequence + * iterator mechanism (since an operator looks like an iterator.) + * + * This is a named class so we can unwrap the operator in + * {@link #runToProducer()} below. + */ + public static class OperatorWrapperSequence extends BaseSequence> + { + private final Operator op; + + public OperatorWrapperSequence(Operator op) + { + super(new BaseSequence.IteratorMaker>() + { + @Override + public Iterator make() + { + return (Iterator) op.open(); + } + + @Override + public void cleanup(Iterator iterFromMake) + { + op.close(true); + } + }); + this.op = op; + } + + public Operator unwrap() + { + return op; + } + + /** + * This will materialize the entire sequence from the wrapped + * operator. Use at your own risk. + */ + @Override + public List toList() + { + return Operators.toList(op); + } + } + + /** + * Converts a stand-alone operator to a sequence outside the context of a fragment + * runner. The sequence starts and closes the operator. + */ + public static Sequence toSequence(Operator op) + { + return new OperatorWrapperSequence<>(op); + } + + /** + * Wrap a sequence in an operator. + *

+ * If the input sequence is a wrapper around an operator, then + * (clumsily) unwraps that operator and returns it directly. + */ + public static Operator toOperator(DAGBuilder builder, Sequence sequence) + { + if (sequence instanceof OperatorWrapperSequence) { + return ((OperatorWrapperSequence) sequence).unwrap(); + } + return new SequenceOperator(builder.context(), sequence); + } + + public static Operator unwrapOperator(Sequence sequence) + { + if (sequence instanceof OperatorWrapperSequence) { + return ((OperatorWrapperSequence) sequence).unwrap(); + } + return null; + } + + /** + * Create an operator which wraps a query runner which allows a query runner + * to be an input to an operator. The runner, and its sequence, will be optimized + * away at runtime if both the upstream and downstream items are both operators, + * but the shim is left in place if the upstream is actually a query runner. + */ + public static QueryRunnerOperator toOperator(QueryRunner runner, QueryPlus query) + { + return new QueryRunnerOperator(runner, query); + } + + /** + * This will materialize the entire sequence from the wrapped + * operator. Use at your own risk. + */ + public static List toList(Operator op) + { + List results = Lists.newArrayList(Operators.toIterable(op)); + op.close(true); + return results; + } +} diff --git a/processing/src/main/java/org/apache/druid/queryng/operators/OrderedMergeOperator.java b/processing/src/main/java/org/apache/druid/queryng/operators/OrderedMergeOperator.java new file mode 100644 index 000000000000..1e19ed171e8f --- /dev/null +++ b/processing/src/main/java/org/apache/druid/queryng/operators/OrderedMergeOperator.java @@ -0,0 +1,167 @@ +/* + * 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.queryng.operators; + +import com.google.common.collect.Ordering; +import org.apache.druid.queryng.fragment.FragmentContext; +import org.apache.druid.queryng.operators.Operator.IterableOperator; + +import java.util.Iterator; +import java.util.PriorityQueue; +import java.util.function.Supplier; + +/** + * Perform an in-memory, n-way merge on n ordered inputs. + * Ordering is given by an {@link Ordering} which defines + * the set of fields to order by, and their comparison + * functions. + * + * @see {@link org.apache.druid.query.RetryQueryRunner} + */ +public class OrderedMergeOperator implements IterableOperator +{ + /** + * Manages a single input characterized by the operator for that + * input, an iterator over that operator, and the current value + * for that iterator. This class caches the current value so that + * the priority queue can perform comparisons on it. + */ + public static class Input + { + private final Operator input; + private final Iterator iter; + private T currentValue; + + public Input(Operator input) + { + this.input = input; + this.iter = input.open(); + if (iter.hasNext()) { + currentValue = iter.next(); + } else { + currentValue = null; + input.close(true); + } + } + + public Operator toOperator(FragmentContext context) + { + if (currentValue == null) { + return new NullOperator(context); + } else { + return new PushBackOperator(context, input, iter, currentValue); + } + } + + public T get() + { + return currentValue; + } + + public boolean next() + { + if (iter.hasNext()) { + currentValue = iter.next(); + return true; + } else { + currentValue = null; + input.close(true); + return false; + } + } + + public boolean eof() + { + return currentValue == null; + } + + public void close(boolean cascade) + { + if (currentValue != null) { + currentValue = null; + input.close(cascade); + } + } + } + + /** + * Supplier of a collection (iterable) of inputs as defined by an + * operator, iterator and current item. To unpack this a bit, the input + * to the merge is the mechanism which distributes fragments across + * multiple threads, waits for the response, checks for missing fragments + * and retries. All this is opaque to this class which just wants the + * final collection that should be merged. That collection is wrapped + * in a supplier so that it is not started until {@link #open()} time. + */ + private final Supplier>> inputs; + private final PriorityQueue> pQueue; + + public OrderedMergeOperator( + FragmentContext context, + Ordering ordering, + int approxInputCount, + Supplier>> inputs + ) + { + this.inputs = inputs; + this.pQueue = new PriorityQueue<>( + approxInputCount == 0 ? 1 : approxInputCount, + ordering.onResultOf(input -> input.get()) + ); + context.register(this); + } + + @Override + public Iterator open() + { + for (Input input : inputs.get()) { + if (!input.eof()) { + pQueue.add(input); + } + } + return this; + } + + @Override + public boolean hasNext() + { + return !pQueue.isEmpty(); + } + + @Override + public T next() + { + Input input = pQueue.remove(); + T result = input.get(); + if (input.next()) { + pQueue.add(input); + } + return result; + } + + @Override + public void close(boolean cascade) + { + while (!pQueue.isEmpty()) { + Input input = pQueue.remove(); + input.close(cascade); + } + } +} diff --git a/processing/src/main/java/org/apache/druid/queryng/operators/PushBackOperator.java b/processing/src/main/java/org/apache/druid/queryng/operators/PushBackOperator.java new file mode 100644 index 000000000000..7b0993d9e437 --- /dev/null +++ b/processing/src/main/java/org/apache/druid/queryng/operators/PushBackOperator.java @@ -0,0 +1,98 @@ +/* + * 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.queryng.operators; + +import org.apache.druid.java.util.common.ISE; +import org.apache.druid.queryng.fragment.FragmentContext; +import org.apache.druid.queryng.operators.Operator.IterableOperator; + +import java.util.Iterator; + +/** + * Operator which allows pushing a row back onto the input. The "pushed" + * row can occur at construction time, or during execution. + */ +public class PushBackOperator implements IterableOperator +{ + private final Operator input; + private Iterator inputIter; + private T pushed; + + public PushBackOperator( + FragmentContext context, + Operator input, + Iterator inputIter, + T pushed) + { + this.input = input; + this.inputIter = inputIter; + this.pushed = pushed; + context.register(this); + } + + public PushBackOperator(FragmentContext context, Operator input) + { + this(context, input, null, null); + } + + @Override + public Iterator open() + { + if (inputIter == null) { + inputIter = input.open(); + } + return this; + } + + @Override + public boolean hasNext() + { + return pushed != null || inputIter != null && inputIter.hasNext(); + } + + @Override + public T next() + { + if (pushed != null) { + T ret = pushed; + pushed = null; + return ret; + } + return inputIter.next(); + } + + public void push(T item) + { + if (pushed != null) { + throw new ISE("Cannot push more than one items onto PushBackOperator"); + } + pushed = item; + } + + @Override + public void close(boolean cascade) + { + if (cascade) { + input.close(cascade); + } + inputIter = null; + pushed = null; + } +} diff --git a/processing/src/main/java/org/apache/druid/queryng/operators/SequenceIterator.java b/processing/src/main/java/org/apache/druid/queryng/operators/SequenceIterator.java new file mode 100644 index 000000000000..b267a47ccaf7 --- /dev/null +++ b/processing/src/main/java/org/apache/druid/queryng/operators/SequenceIterator.java @@ -0,0 +1,81 @@ +/* + * 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.queryng.operators; + +import com.google.common.base.Preconditions; +import org.apache.druid.java.util.common.guava.Sequence; +import org.apache.druid.java.util.common.guava.Yielder; +import org.apache.druid.java.util.common.guava.YieldingAccumulator; + +import java.io.IOException; +import java.util.Iterator; + +/** + * Iterator over a sequence. + */ +public class SequenceIterator implements Iterator, AutoCloseable +{ + private Yielder yielder; + + public static SequenceIterator of(Sequence sequence) + { + return new SequenceIterator(sequence); + } + + public SequenceIterator(Sequence sequence) + { + this.yielder = sequence.toYielder( + null, + new YieldingAccumulator() + { + @Override + public T accumulate(T accumulated, T in) + { + yield(); + return in; + } + } + ); + } + + @Override + public boolean hasNext() + { + return !yielder.isDone(); + } + + @Override + public T next() + { + Preconditions.checkState(!yielder.isDone()); + T value = yielder.get(); + yielder = yielder.next(null); + return value; + } + + @Override + public void close() throws IOException + { + if (yielder != null) { + yielder.close(); + yielder = null; + } + } +} diff --git a/processing/src/main/java/org/apache/druid/queryng/operators/SequenceOperator.java b/processing/src/main/java/org/apache/druid/queryng/operators/SequenceOperator.java new file mode 100644 index 000000000000..377d57747b71 --- /dev/null +++ b/processing/src/main/java/org/apache/druid/queryng/operators/SequenceOperator.java @@ -0,0 +1,105 @@ +/* + * 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.queryng.operators; + +import com.google.common.base.Preconditions; +import org.apache.druid.java.util.common.guava.Sequence; +import org.apache.druid.java.util.common.guava.Yielder; +import org.apache.druid.java.util.common.guava.YieldingAccumulator; +import org.apache.druid.queryng.fragment.FragmentContext; +import org.apache.druid.queryng.operators.Operator.IterableOperator; + +import java.io.IOException; +import java.util.Iterator; + +/** + * The SequenceOperator wraps a {@link Sequence} in the + * operator protocol. The operator will make (at most) one pass through + * the sequence. The sequence's yielder will be defined in start(), + * which may cause the sequence to start doing work and obtaining resources. + * Each call to next()/get() will yield one result + * from the sequence. The close() call will close the yielder + * for the sequence, which should release any resources held by the sequence. + * + * @param The type of the item (row, batch) returned by the sequence + * and thus returned by the operator. + */ +public class SequenceOperator implements IterableOperator +{ + private final Sequence sequence; + private Yielder yielder; + + public SequenceOperator(FragmentContext context, Sequence sequence) + { + this.sequence = sequence; + context.register(this); + } + + @Override + public Iterator open() + { + Preconditions.checkState(yielder == null); + yielder = sequence.toYielder( + null, + new YieldingAccumulator() + { + @Override + public T accumulate(T accumulated, T in) + { + yield(); + return in; + } + } + ); + return this; + } + + @Override + public boolean hasNext() + { + return yielder != null && !yielder.isDone(); + } + + @Override + public T next() + { + Preconditions.checkState(yielder != null); + T value = yielder.get(); + yielder = yielder.next(null); + return value; + } + + @Override + public void close(boolean cascade) + { + if (yielder == null) { + return; + } + try { + yielder.close(); + } + catch (IOException e) { + throw new RuntimeException(e.getMessage(), e); + } + finally { + yielder = null; + } + } +} diff --git a/processing/src/main/java/org/apache/druid/queryng/operators/TransformOperator.java b/processing/src/main/java/org/apache/druid/queryng/operators/TransformOperator.java new file mode 100644 index 000000000000..bdfac97ad8a3 --- /dev/null +++ b/processing/src/main/java/org/apache/druid/queryng/operators/TransformOperator.java @@ -0,0 +1,62 @@ +/* + * 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.queryng.operators; + +import com.google.common.base.Function; +import org.apache.druid.query.BySegmentQueryRunner; +import org.apache.druid.query.Query; +import org.apache.druid.query.QueryToolChest; +import org.apache.druid.query.aggregation.MetricManipulationFn; +import org.apache.druid.queryng.fragment.FragmentContext; + +/** + * Operator that applies a function to each input item to produce the output item. + * + * Generalization of {@link QueryToolChest#makePostComputeManipulatorFn(Query, MetricManipulationFn)} to the + * result stream. When used in this role, the operator is expected to be the operator in the pipeline, + * after results are fully merged. + *

+ * Note that, when used in the above role, despite the type parameter "T", this runner may not actually + * return sequences with type T. This most + * commonly happens when an upstream {@link BySegmentQueryRunner} changes the result stream to type + * {@code Result>}, in which case this class will retain the structure, but call the finalizer + * function on each result in the by-segment list (which may change their type from T to something else). + * + * @see {@link org.apache.druid.query.FinalizeResultsQueryRunner} + */ +public class TransformOperator extends MappingOperator +{ + private final Function transformFn; + + public TransformOperator( + FragmentContext context, + final Function transformFn, + final Operator input) + { + super(context, input); + this.transformFn = transformFn; + } + + @Override + public OUT next() + { + return transformFn.apply(inputIter.next()); + } +} diff --git a/processing/src/main/java/org/apache/druid/queryng/operators/WrappingOperator.java b/processing/src/main/java/org/apache/druid/queryng/operators/WrappingOperator.java new file mode 100644 index 000000000000..922f00419a7d --- /dev/null +++ b/processing/src/main/java/org/apache/druid/queryng/operators/WrappingOperator.java @@ -0,0 +1,76 @@ +/* + * 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.queryng.operators; + +import com.google.common.base.Preconditions; +import org.apache.druid.queryng.fragment.FragmentContext; + +import java.util.Iterator; + +/** + * Operator which "wraps" another operator where the only behavior of + * interest is at the start or end of the run. The iterator from the + * input is the iterator for the output. + */ +public abstract class WrappingOperator implements Operator +{ + protected final FragmentContext context; + private final Operator input; + protected State state = State.START; + + public WrappingOperator(FragmentContext context, Operator input) + { + this.context = context; + this.input = input; + context.register(this); + } + + @Override + public Iterator open() + { + Preconditions.checkState(state == State.START); + Iterator inputIter = input.open(); + state = State.RUN; + onOpen(); + return inputIter; + } + + @Override + public void close(boolean cascade) + { + try { + if (state == State.RUN && cascade) { + input.close(cascade); + } + } + finally { + onClose(); + state = State.CLOSED; + } + } + + protected void onOpen() + { + } + + protected void onClose() + { + } +} diff --git a/processing/src/main/java/org/apache/druid/queryng/operators/general/QueryRunnerOperator.java b/processing/src/main/java/org/apache/druid/queryng/operators/general/QueryRunnerOperator.java new file mode 100644 index 000000000000..10c5524ff11c --- /dev/null +++ b/processing/src/main/java/org/apache/druid/queryng/operators/general/QueryRunnerOperator.java @@ -0,0 +1,70 @@ +/* + * 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.queryng.operators.general; + +import org.apache.druid.java.util.common.guava.Sequence; +import org.apache.druid.query.QueryPlus; +import org.apache.druid.query.QueryRunner; +import org.apache.druid.queryng.fragment.FragmentContext; +import org.apache.druid.queryng.operators.Operator; +import org.apache.druid.queryng.operators.Operators; + +import java.util.Iterator; + +/** + * Operator which wraps a query runner. When used in the interim + * "shim" architecture, this operator allows a query runner to be + * the input (upstream) to some other (downstream) operator. The + * upstream query runner may give rise to its own operator. In that + * case, at runtime, the sequence wrapper for that operator is + * optimized away, leaving just the two operators. + */ +public class QueryRunnerOperator implements Operator +{ + protected final FragmentContext context; + private final QueryRunner runner; + private final QueryPlus query; + private Operator child; + + public QueryRunnerOperator(QueryRunner runner, QueryPlus query) + { + this.context = query.fragmentBuilder().context(); + this.runner = runner; + this.query = query; + context.register(this); + } + + @Override + public Iterator open() + { + Sequence seq = runner.run(query, context.responseContext()); + child = Operators.toOperator(context, seq); + return child.open(); + } + + @Override + public void close(boolean cascade) + { + if (child != null && cascade) { + child.close(cascade); + } + child = null; + } +} diff --git a/processing/src/test/java/org/apache/druid/queryng/operators/ConcatOperatorTest.java b/processing/src/test/java/org/apache/druid/queryng/operators/ConcatOperatorTest.java new file mode 100644 index 000000000000..ac3d32b5bbdd --- /dev/null +++ b/processing/src/test/java/org/apache/druid/queryng/operators/ConcatOperatorTest.java @@ -0,0 +1,136 @@ +/* + * 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.queryng.operators; + +import org.apache.druid.queryng.fragment.FragmentContext; +import org.apache.druid.queryng.operators.Operator.State; +import org.junit.Test; + +import java.util.Arrays; +import java.util.Collections; +import java.util.Iterator; +import java.util.List; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertTrue; + +public class ConcatOperatorTest +{ + @Test + public void testEmpty() + { + FragmentContext context = FragmentContext.defaultContext(); + Operator op = new ConcatOperator(context, Collections.emptyList()); + List results = Operators.toList(op); + assertTrue(results.isEmpty()); + } + + @Test + public void testOneEmptyInput() + { + FragmentContext context = FragmentContext.defaultContext(); + Operator input = new NullOperator(context); + Operator op = new ConcatOperator(context, Arrays.asList(input)); + List results = Operators.toList(op); + assertTrue(results.isEmpty()); + } + + @Test + public void testTwoEmptyInputs() + { + FragmentContext context = FragmentContext.defaultContext(); + Operator input1 = new NullOperator(context); + Operator input2 = new NullOperator(context); + Operator op = new ConcatOperator(context, Arrays.asList(input1, input2)); + List results = Operators.toList(op); + assertTrue(results.isEmpty()); + } + + @Test + public void testOneInput() + { + FragmentContext context = FragmentContext.defaultContext(); + Operator input = MockOperator.ints(context, 2); + Operator op = new ConcatOperator(context, Arrays.asList(input)); + List results = Operators.toList(op); + List expected = Arrays.asList(0, 1); + assertEquals(expected, results); + } + + @Test + public void testEmptyThenNonEmptyInputs() + { + FragmentContext context = FragmentContext.defaultContext(); + Operator input1 = new NullOperator(context); + Operator input2 = MockOperator.ints(context, 2); + Operator op = new ConcatOperator(context, Arrays.asList(input1, input2)); + List results = Operators.toList(op); + List expected = Arrays.asList(0, 1); + assertEquals(expected, results); + } + + @Test + public void testNonEmptyThenEmptyInputs() + { + FragmentContext context = FragmentContext.defaultContext(); + Operator input1 = MockOperator.ints(context, 2); + Operator input2 = new NullOperator(context); + Operator op = new ConcatOperator(context, Arrays.asList(input1, input2)); + List results = Operators.toList(op); + List expected = Arrays.asList(0, 1); + assertEquals(expected, results); + } + + @Test + public void testTwoInputs() + { + FragmentContext context = FragmentContext.defaultContext(); + Operator input1 = MockOperator.ints(context, 2); + Operator input2 = MockOperator.ints(context, 2); + Operator op = new ConcatOperator(context, Arrays.asList(input1, input2)); + List results = Operators.toList(op); + List expected = Arrays.asList(0, 1, 0, 1); + assertEquals(expected, results); + } + + @Test + public void testClose() + { + FragmentContext context = FragmentContext.defaultContext(); + MockOperator input1 = MockOperator.ints(context, 2); + MockOperator input2 = MockOperator.ints(context, 2); + Operator op = new ConcatOperator(context, Arrays.asList(input1, input2)); + Iterator iter = op.open(); + assertTrue(iter.hasNext()); + assertEquals(0, (int) iter.next()); + + // Only first input has been opened. + assertEquals(State.RUN, input1.state); + assertEquals(State.START, input2.state); + + // Cascade closes inputs + op.close(true); + assertEquals(State.CLOSED, input1.state); + assertEquals(State.START, input2.state); + + // Close again does nothing. + op.close(false); + } +} diff --git a/processing/src/test/java/org/apache/druid/queryng/operators/FragmentTest.java b/processing/src/test/java/org/apache/druid/queryng/operators/FragmentTest.java new file mode 100644 index 000000000000..40213b8eebf2 --- /dev/null +++ b/processing/src/test/java/org/apache/druid/queryng/operators/FragmentTest.java @@ -0,0 +1,265 @@ +/* + * 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.queryng.operators; + +import org.apache.druid.java.util.common.ISE; +import org.apache.druid.java.util.common.guava.Sequence; +import org.apache.druid.java.util.common.guava.SequenceTestHelper; +import org.apache.druid.queryng.fragment.FragmentBuilder; +import org.apache.druid.queryng.fragment.FragmentHandle; +import org.apache.druid.queryng.fragment.FragmentRun; +import org.junit.Test; + +import java.io.IOException; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertSame; +import static org.junit.Assert.assertTrue; +import static org.junit.Assert.fail; + +/** + * Test the various fragment-level classes: {@code FragmentBuilder}, + * {@link FragmentHandle} and {@link FragmentRun}. Since the second two + * are convenience wrappers around the first, each test uses the wrappers + * to indirectly test the builder. + */ +public class FragmentTest +{ + @Test + public void testOperatorBasics() + { + FragmentBuilder builder = FragmentBuilder.defaultBuilder(); + MockOperator op = MockOperator.ints(builder.context(), 2); + + // Handle for the leaf operator + FragmentHandle handle = builder.handle(op); + assertTrue(handle.rootIsOperator()); + assertFalse(handle.rootIsSequence()); + assertSame(op, handle.rootOperator()); + assertNull(handle.rootSequence()); + assertSame(handle, handle.toOperator()); + assertSame(builder.context(), handle.context()); + + // Add another operator to the DAG. + Operator op2 = new MockFilterOperator(builder.context(), op, x -> x < 3); + FragmentHandle handle2 = handle.compose(op2); + assertTrue(handle2.rootIsOperator()); + assertFalse(handle2.rootIsSequence()); + assertSame(op2, handle2.rootOperator()); + assertNull(handle2.rootSequence()); + assertSame(handle2, handle2.toOperator()); + assertSame(builder.context(), handle2.context()); + } + + /** + * The empty handle is a "bootstrap" mechanism for building the + * leaf-most operator in a DAG. + */ + @Test + public void testEmptyHandle() + { + FragmentBuilder builder = FragmentBuilder.defaultBuilder(); + FragmentHandle emptyHandle = builder.emptyHandle(); + + assertFalse(emptyHandle.rootIsOperator()); + assertFalse(emptyHandle.rootIsSequence()); + assertNull(emptyHandle.rootOperator()); + assertNull(emptyHandle.rootSequence()); + assertSame(emptyHandle, emptyHandle.toOperator()); + assertSame(builder.context(), emptyHandle.context()); + + // There is no useful case for converting an empty handle + // to a sequence. + try { + emptyHandle.toSequence(); + fail(); + } + catch (ISE e) { + // Expected + } + + MockOperator op = MockOperator.ints(builder.context(), 2); + FragmentHandle handle = emptyHandle.compose(op); + assertTrue(handle.rootIsOperator()); + assertSame(op, handle.rootOperator()); + assertSame(builder.context(), handle.context()); + } + + @Test + public void testSequenceBasics() + { + FragmentBuilder builder = FragmentBuilder.defaultBuilder(); + MockOperator op = MockOperator.ints(builder.context(), 2); + Sequence seq = Operators.toSequence(op); + + // Handle for the leaf sequence + FragmentHandle handle = builder.handle(seq); + assertFalse(handle.rootIsOperator()); + assertTrue(handle.rootIsSequence()); + // Note asymmetry: sequences can be unwrapped to get an operator, + // but not visa-versa. + assertSame(op, handle.rootOperator()); + assertSame(seq, handle.rootSequence()); + assertSame(handle, handle.toSequence()); + assertSame(builder.context(), handle.context()); + + // Add another operator to the DAG by unwrapping the above + // sequence. + Operator op2 = new MockFilterOperator( + builder.context(), + handle.toOperator().rootOperator(), + x -> x < 3); + FragmentHandle handle2o = handle.compose(op2); + + // Use composition to get the new root sequence. + FragmentHandle handle2 = handle2o.toSequence(); + assertFalse(handle2.rootIsOperator()); + assertTrue(handle2.rootIsSequence()); + assertSame(op2, handle2.rootOperator()); + assertNotNull(handle2.rootSequence()); + assertSame(handle2, handle2.toSequence()); + assertSame(builder.context(), handle2.context()); + } + + @Test + public void testRun() + { + FragmentBuilder builder = FragmentBuilder.defaultBuilder(); + MockOperator op = MockOperator.ints(builder.context(), 2); + FragmentHandle handle = builder.handle(op); + int i = 0; + try (FragmentRun run = handle.run()) { + for (Integer value : run) { + assertEquals(i++, (int) value); + } + } + assertEquals(Operator.State.CLOSED, op.state); + } + + @Test + public void testToListWithOperator() + { + FragmentBuilder builder = FragmentBuilder.defaultBuilder(); + MockOperator op = MockOperator.ints(builder.context(), 2); + FragmentHandle handle = builder.handle(op); + List results = handle.run().toList(); + assertEquals(0, (int) results.get(0)); + assertEquals(1, (int) results.get(1)); + assertEquals(Operator.State.CLOSED, op.state); + } + + /** + * Test running a DAG with a root sequence as a sequence. + */ + @Test + public void testToListWithSequence() + { + FragmentBuilder builder = FragmentBuilder.defaultBuilder(); + MockOperator op = MockOperator.ints(builder.context(), 2); + FragmentHandle handle = builder.handle(op).toSequence(); + List results = handle.runAsSequence().toList(); + assertEquals(0, (int) results.get(0)); + assertEquals(1, (int) results.get(1)); + assertEquals(Operator.State.CLOSED, op.state); + } + + /** + * Test running a fragment as a sequence when the root is an + * operator. The operator will be wrapped in a sequence internally. + */ + @Test + public void testRunAsSequenceWithOperator() + { + FragmentBuilder builder = FragmentBuilder.defaultBuilder(); + MockOperator op = MockOperator.ints(builder.context(), 2); + FragmentHandle handle = builder.handle(op); + List results = handle.runAsSequence().toList(); + assertEquals(0, (int) results.get(0)); + assertEquals(1, (int) results.get(1)); + assertEquals(Operator.State.CLOSED, op.state); + } + + /** + * Test that if an operator stack has a sequence as its root, + * that running the DAG as an operator will unwrap that root + * sequence to get an operator. + */ + @Test + public void testRunWithSequence() + { + FragmentBuilder builder = FragmentBuilder.defaultBuilder(); + MockOperator op = MockOperator.ints(builder.context(), 2); + FragmentHandle handle = builder.handle(op).toSequence(); + List results = handle.run().toList(); + assertEquals(0, (int) results.get(0)); + assertEquals(1, (int) results.get(1)); + assertEquals(Operator.State.CLOSED, op.state); + } + + /** + * An operator is a one-pass object, don't try sequence tests that assume + * the sequence is reentrant. + */ + @Test + public void testSequenceYielder() throws IOException + { + FragmentBuilder builder = FragmentBuilder.defaultBuilder(); + MockOperator op = MockOperator.ints(builder.context(), 5); + final List expected = Arrays.asList(0, 1, 2, 3, 4); + Sequence seq = builder.runAsSequence(op); + SequenceTestHelper.testYield("op", 5, seq, expected); + assertEquals(Operator.State.CLOSED, op.state); + } + + @Test + public void testSequenceAccum() throws IOException + { + FragmentBuilder builder = FragmentBuilder.defaultBuilder(); + MockOperator op = MockOperator.ints(builder.context(), 4); + final List vals = Arrays.asList(0, 1, 2, 3); + FragmentHandle handle = builder.handle(op); + Sequence seq = handle.runAsSequence(); + SequenceTestHelper.testAccumulation("op", seq, vals); + assertEquals(Operator.State.CLOSED, op.state); + } + + @Test + public void testRunEmptyHandle() + { + FragmentBuilder builder = FragmentBuilder.defaultBuilder(); + FragmentHandle emptyHandle = builder.emptyHandle(); + assertTrue(emptyHandle.run().toList().isEmpty()); + } + + @Test + public void testRunEmptyHandleAsSequence() + { + FragmentBuilder builder = FragmentBuilder.defaultBuilder(); + FragmentHandle emptyHandle = builder.emptyHandle(); + Sequence seq = emptyHandle.runAsSequence(); + SequenceTestHelper.testAccumulation("empty", seq, Collections.emptyList()); + } +} diff --git a/processing/src/test/java/org/apache/druid/queryng/operators/MockCursor.java b/processing/src/test/java/org/apache/druid/queryng/operators/MockCursor.java new file mode 100644 index 000000000000..84ff056945a8 --- /dev/null +++ b/processing/src/test/java/org/apache/druid/queryng/operators/MockCursor.java @@ -0,0 +1,205 @@ +/* + * 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.queryng.operators; + +import org.apache.druid.java.util.common.ISE; +import org.apache.druid.query.dimension.DimensionSpec; +import org.apache.druid.query.monomorphicprocessing.RuntimeShapeInspector; +import org.apache.druid.segment.ColumnSelectorFactory; +import org.apache.druid.segment.ColumnValueSelector; +import org.apache.druid.segment.Cursor; +import org.apache.druid.segment.DimensionSelector; +import org.apache.druid.segment.LongColumnSelector; +import org.apache.druid.segment.column.ColumnCapabilities; +import org.joda.time.DateTime; +import org.joda.time.Interval; + +public class MockCursor implements Cursor, ColumnSelectorFactory +{ + private class MockTimeColumn implements LongColumnSelector + { + @Override + public long getLong() + { + return segmentBase + posn / divideBy; + } + + @Override + public void inspectRuntimeShape(RuntimeShapeInspector inspector) + { + } + + @Override + public boolean isNull() + { + return false; + } + } + + private class MockLongColumn implements LongColumnSelector + { + @Override + public long getLong() + { + return posn; + } + + @Override + public void inspectRuntimeShape(RuntimeShapeInspector inspector) + { + } + + @Override + public boolean isNull() + { + return false; + } + } + + private class MockStringColumn implements ColumnValueSelector + { + @Override + public long getLong() + { + return 0; + } + + @Override + public void inspectRuntimeShape(RuntimeShapeInspector inspector) + { + } + + @Override + public boolean isNull() + { + return false; + } + + @Override + public double getDouble() + { + return 0; + } + + @Override + public float getFloat() + { + return 0; + } + + @Override + public String getObject() + { + return "string value"; + } + + @Override + public Class classOfObject() + { + return String.class; + } + } + + private final int targetRowCount = 5_000_000; + private final long segmentBase; + private final int divideBy; + private int posn; + + public MockCursor(Interval interval) + { + segmentBase = interval.getStartMillis(); + long span = interval.getEndMillis() - segmentBase; + if (span > targetRowCount) { + divideBy = 1; + } else { + divideBy = (int) (targetRowCount / span); + } + } + + @Override + public ColumnSelectorFactory getColumnSelectorFactory() + { + return this; + } + + @Override + public DimensionSelector makeDimensionSelector( + DimensionSpec dimensionSpec) + { + throw new ISE("Not supported"); + } + + @Override + public ColumnValueSelector makeColumnValueSelector(String columnName) + { + switch (columnName) { + case "__time": + return new MockTimeColumn(); + case "delta": + return new MockLongColumn(); + case "page": + return new MockStringColumn(); + default: + return null; + } + } + + @Override + public ColumnCapabilities getColumnCapabilities(String column) + { + throw new ISE("Not supported"); + } + + @Override + public DateTime getTime() + { + return new DateTime(segmentBase + posn / divideBy); + } + + @Override + public void advance() + { + posn++; + } + + @Override + public void advanceUninterruptibly() + { + advance(); + } + + @Override + public boolean isDone() + { + return posn >= targetRowCount; + } + + @Override + public boolean isDoneOrInterrupted() + { + return isDone(); + } + + @Override + public void reset() + { + posn = 0; + } +} diff --git a/processing/src/test/java/org/apache/druid/queryng/operators/MockFilterOperator.java b/processing/src/test/java/org/apache/druid/queryng/operators/MockFilterOperator.java new file mode 100644 index 000000000000..81e0514e24f4 --- /dev/null +++ b/processing/src/test/java/org/apache/druid/queryng/operators/MockFilterOperator.java @@ -0,0 +1,58 @@ +/* + * 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.queryng.operators; + +import org.apache.druid.queryng.fragment.FragmentContext; + +import java.util.function.Predicate; + +/** + * Super-simple filter operator which allows tests to build up a + * stack of operators starting with values generated by MockOperator. + */ +public class MockFilterOperator extends MappingOperator +{ + private final Predicate predicate; + private T nextValue; + + public MockFilterOperator(FragmentContext context, Operator input, Predicate pred) + { + super(context, input); + this.predicate = pred; + } + + @Override + public boolean hasNext() + { + while (super.hasNext()) { + nextValue = inputIter.next(); + if (predicate.test(nextValue)) { + return true; + } + } + return false; + } + + @Override + public T next() + { + return nextValue; + } +} diff --git a/processing/src/test/java/org/apache/druid/queryng/operators/MockOperator.java b/processing/src/test/java/org/apache/druid/queryng/operators/MockOperator.java new file mode 100644 index 000000000000..f97853e09290 --- /dev/null +++ b/processing/src/test/java/org/apache/druid/queryng/operators/MockOperator.java @@ -0,0 +1,81 @@ +/* + * 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.queryng.operators; + +import com.google.common.base.Preconditions; +import org.apache.druid.queryng.fragment.FragmentContext; +import org.apache.druid.queryng.operators.Operator.IterableOperator; + +import java.util.Iterator; +import java.util.function.Function; + +public class MockOperator implements IterableOperator +{ + public final int targetCount; + private final Function generator; + private int rowPosn; + public State state = State.START; + + public MockOperator( + FragmentContext context, + int rowCount, + Function generator) + { + this.targetCount = rowCount; + this.generator = generator; + context.register(this); + } + + public static MockOperator ints(FragmentContext context, int rowCount) + { + return new MockOperator(context, rowCount, rid -> rid); + } + + public static MockOperator strings(FragmentContext context, int rowCount) + { + return new MockOperator(context, rowCount, rid -> "Mock row " + Integer.toString(rid)); + } + + @Override + public Iterator open() + { + Preconditions.checkState(state == State.START); + state = State.RUN; + return this; + } + + @Override + public boolean hasNext() + { + return rowPosn < targetCount; + } + + @Override + public T next() + { + return generator.apply(rowPosn++); + } + + @Override + public void close(boolean cascade) + { + state = State.CLOSED; + } +} diff --git a/processing/src/test/java/org/apache/druid/queryng/operators/MockOperatorTest.java b/processing/src/test/java/org/apache/druid/queryng/operators/MockOperatorTest.java new file mode 100644 index 000000000000..4462be4a0321 --- /dev/null +++ b/processing/src/test/java/org/apache/druid/queryng/operators/MockOperatorTest.java @@ -0,0 +1,105 @@ +/* + * 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.queryng.operators; + +import org.apache.druid.java.util.common.guava.Sequence; +import org.apache.druid.queryng.fragment.FragmentBuilder; +import org.apache.druid.queryng.fragment.FragmentHandle; +import org.apache.druid.queryng.fragment.FragmentRun; +import org.junit.Test; + +import java.util.Iterator; +import java.util.List; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; + +/** + * Use a mock operator to test (and illustrate) the basic operator + * mechanisms. + */ +public class MockOperatorTest +{ + @Test + public void testMockStringOperator() + { + FragmentBuilder builder = FragmentBuilder.defaultBuilder(); + MockOperator op = MockOperator.strings(builder.context(), 2); + FragmentHandle handle = builder.handle(op); + try (FragmentRun run = handle.run()) { + Iterator iter = run.iterator(); + assertTrue(iter.hasNext()); + assertEquals("Mock row 0", iter.next()); + assertTrue(iter.hasNext()); + assertEquals("Mock row 1", iter.next()); + assertFalse(iter.hasNext()); + } + assertEquals(Operator.State.CLOSED, op.state); + } + + @Test + public void testMockIntOperator() + { + // Test using the toList feature. + FragmentBuilder builder = FragmentBuilder.defaultBuilder(); + MockOperator op = MockOperator.ints(builder.context(), 2); + List results = builder.run(op).toList(); + assertEquals(0, (int) results.get(0)); + assertEquals(1, (int) results.get(1)); + assertEquals(Operator.State.CLOSED, op.state); + } + + /** + * Getting weird: an operator that wraps a sequence that wraps an operator. + */ + @Test + public void testSequenceOperator() + { + FragmentBuilder builder = FragmentBuilder.defaultBuilder(); + MockOperator op = MockOperator.strings(builder.context(), 2); + Sequence seq = Operators.toSequence(op); + Operator outer = Operators.toOperator(builder, seq); + FragmentRun run = builder.run(outer); + Iterator iter = run.iterator(); + assertTrue(iter.hasNext()); + assertEquals("Mock row 0", iter.next()); + assertTrue(iter.hasNext()); + assertEquals("Mock row 1", iter.next()); + assertFalse(iter.hasNext()); + run.close(); + assertEquals(Operator.State.CLOSED, op.state); + } + + @Test + public void testMockFilter() + { + FragmentBuilder builder = FragmentBuilder.defaultBuilder(); + MockOperator op = MockOperator.ints(builder.context(), 4); + Operator op2 = new MockFilterOperator( + builder.context(), + op, + x -> x % 2 == 0); + List results = builder.run(op2).toList(); + assertEquals(0, (int) results.get(0)); + assertEquals(2, (int) results.get(1)); + assertEquals(Operator.State.CLOSED, op.state); + } +} diff --git a/processing/src/test/java/org/apache/druid/queryng/operators/MockStorageAdapter.java b/processing/src/test/java/org/apache/druid/queryng/operators/MockStorageAdapter.java new file mode 100644 index 000000000000..0148e02a2cd1 --- /dev/null +++ b/processing/src/test/java/org/apache/druid/queryng/operators/MockStorageAdapter.java @@ -0,0 +1,143 @@ +/* + * 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.queryng.operators; + +import org.apache.druid.java.util.common.granularity.Granularity; +import org.apache.druid.java.util.common.guava.Sequence; +import org.apache.druid.java.util.common.guava.Sequences; +import org.apache.druid.query.QueryMetrics; +import org.apache.druid.query.filter.Filter; +import org.apache.druid.segment.Cursor; +import org.apache.druid.segment.Metadata; +import org.apache.druid.segment.StorageAdapter; +import org.apache.druid.segment.VirtualColumns; +import org.apache.druid.segment.column.ColumnCapabilities; +import org.apache.druid.segment.data.Indexed; +import org.apache.druid.segment.data.ListIndexed; +import org.joda.time.DateTime; +import org.joda.time.Interval; + +import java.util.Arrays; +import java.util.Collections; + +public class MockStorageAdapter implements StorageAdapter +{ + @Override + public Sequence makeCursors( + Filter filter, + Interval interval, + VirtualColumns virtualColumns, + Granularity gran, + boolean descending, + QueryMetrics queryMetrics) + { + return Sequences.simple(Collections.singletonList(new MockCursor(interval))); + } + + @Override + public Interval getInterval() + { + return Interval.parse("2015-09-12T13:00:00.000Z/2015-09-12T14:00:00.000Z"); + } + + @Override + public Indexed getAvailableDimensions() + { + return new ListIndexed<>(Arrays.asList("__time", "page")); + } + + @Override + public Iterable getAvailableMetrics() + { + return Arrays.asList("delta"); + } + + @Override + public int getDimensionCardinality(String column) + { + return Integer.MAX_VALUE; + } + + @Override + public DateTime getMinTime() + { + return DateTime.parse("2015-09-12T13:00:00.000Z"); + } + + @Override + public DateTime getMaxTime() + { + return DateTime.parse("2015-09-12T13:59:59.999Z"); + } + + @Override + public Comparable getMinValue(String column) + { + switch (column) { + case "__time": + return getMinTime(); + case "delta": + return 0; + case "page": + return ""; + } + return null; + } + + @Override + public Comparable getMaxValue(String column) + { + switch (column) { + case "__time": + return getMaxTime(); + case "delta": + return 10_000; + case "page": + return "zzzzzzzzzz"; + } + return null; + } + + @Override + public ColumnCapabilities getColumnCapabilities(String column) + { + // TODO Auto-generated method stub + return null; + } + + @Override + public int getNumRows() + { + return 5_000_000; + } + + @Override + public DateTime getMaxIngestedEventTime() + { + return getMaxTime(); + } + + @Override + public Metadata getMetadata() + { + // TODO Auto-generated method stub + return null; + } +} diff --git a/processing/src/test/java/org/apache/druid/queryng/operators/PushBackOperatorTest.java b/processing/src/test/java/org/apache/druid/queryng/operators/PushBackOperatorTest.java new file mode 100644 index 000000000000..911f5e136cd1 --- /dev/null +++ b/processing/src/test/java/org/apache/druid/queryng/operators/PushBackOperatorTest.java @@ -0,0 +1,82 @@ +/* + * 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.queryng.operators; + +import org.apache.druid.queryng.fragment.FragmentContext; +import org.junit.Test; + +import java.util.Arrays; +import java.util.Iterator; +import java.util.List; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertTrue; + +public class PushBackOperatorTest +{ + @Test + public void testEmptyInput() + { + FragmentContext context = FragmentContext.defaultContext(); + Operator input = new NullOperator(context); + Operator op = new PushBackOperator(context, input); + assertTrue(Operators.toList(op).isEmpty()); + } + + @Test + public void testSimpleInput() + { + FragmentContext context = FragmentContext.defaultContext(); + Operator input = MockOperator.ints(context, 2); + Operator op = new PushBackOperator(context, input); + List results = Operators.toList(op); + List expected = Arrays.asList(0, 1); + assertEquals(expected, results); + } + + @Test + public void testPush() + { + FragmentContext context = FragmentContext.defaultContext(); + Operator input = MockOperator.ints(context, 2); + PushBackOperator op = new PushBackOperator(context, input); + Iterator iter = op.open(); + assertTrue(iter.hasNext()); + Integer item = iter.next(); + op.push(item); + List results = Operators.toList(op); + List expected = Arrays.asList(0, 1); + assertEquals(expected, results); + } + + @Test + public void testInitialPush() + { + FragmentContext context = FragmentContext.defaultContext(); + Operator input = MockOperator.ints(context, 2); + Iterator iter = input.open(); + assertTrue(iter.hasNext()); + Integer item = iter.next(); + PushBackOperator op = new PushBackOperator(context, input, iter, item); + List results = Operators.toList(op); + List expected = Arrays.asList(0, 1); + assertEquals(expected, results); + } +} diff --git a/processing/src/test/java/org/apache/druid/queryng/operators/general/OrderedMergeOperatorTest.java b/processing/src/test/java/org/apache/druid/queryng/operators/general/OrderedMergeOperatorTest.java new file mode 100644 index 000000000000..88dda1789e1e --- /dev/null +++ b/processing/src/test/java/org/apache/druid/queryng/operators/general/OrderedMergeOperatorTest.java @@ -0,0 +1,140 @@ +/* + * 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.queryng.operators.general; + +import com.google.common.collect.Lists; +import com.google.common.collect.Ordering; +import org.apache.druid.queryng.fragment.FragmentContext; +import org.apache.druid.queryng.operators.MockOperator; +import org.apache.druid.queryng.operators.NullOperator; +import org.apache.druid.queryng.operators.Operator; +import org.apache.druid.queryng.operators.Operator.State; +import org.apache.druid.queryng.operators.OrderedMergeOperator; +import org.apache.druid.queryng.operators.OrderedMergeOperator.Input; +import org.junit.Test; + +import java.util.Arrays; +import java.util.Collections; +import java.util.Iterator; +import java.util.List; +import java.util.function.Supplier; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; + +public class OrderedMergeOperatorTest +{ + @Test + public void testNoInputs() + { + FragmentContext context = FragmentContext.defaultContext(); + Supplier>> inputs = + () -> Collections.emptyList(); + Operator op = new OrderedMergeOperator<>( + context, + Ordering.natural(), + 0, + inputs); + Iterator iter = op.open(); + assertFalse(iter.hasNext()); + op.close(true); + } + + @Test + public void testEmptyInputs() + { + FragmentContext context = FragmentContext.defaultContext(); + Supplier>> inputs = + () -> Arrays.asList( + new Input(new NullOperator(context)), + new Input(new NullOperator(context))); + Operator op = new OrderedMergeOperator<>( + context, + Ordering.natural(), + 2, + inputs); + Iterator iter = op.open(); + assertFalse(iter.hasNext()); + op.close(true); + } + + @Test + public void testOneInput() + { + FragmentContext context = FragmentContext.defaultContext(); + Supplier>> inputs = + () -> Arrays.asList( + new Input(MockOperator.ints(context, 3))); + Operator op = new OrderedMergeOperator<>( + context, + Ordering.natural(), + 1, + inputs); + Iterator iter = op.open(); + List results = Lists.newArrayList(iter); + op.close(true); + assertEquals(Arrays.asList(0, 1, 2), results); + } + + @Test + public void testTwoInputs() + { + FragmentContext context = FragmentContext.defaultContext(); + Supplier>> inputs = + () -> Arrays.asList( + new Input(MockOperator.ints(context, 3)), + new Input(MockOperator.ints(context, 5))); + Operator op = new OrderedMergeOperator<>( + context, + Ordering.natural(), + 2, + inputs); + Iterator iter = op.open(); + List results = Lists.newArrayList(iter); + op.close(true); + assertEquals(Arrays.asList(0, 0, 1, 1, 2, 2, 3, 4), results); + } + + @Test + public void testClose() + { + FragmentContext context = FragmentContext.defaultContext(); + MockOperator input1 = MockOperator.ints(context, 2); + MockOperator input2 = MockOperator.ints(context, 2); + Supplier>> inputs = + () -> Arrays.asList( + new Input(input1), + new Input(input2)); + Operator op = new OrderedMergeOperator<>( + context, + Ordering.natural(), + 2, + inputs); + Iterator iter = op.open(); + List results = Lists.newArrayList(iter); + assertEquals(Arrays.asList(0, 0, 1, 1), results); + + // Inputs are closed as exhausted. + assertTrue(input1.state == State.CLOSED); + assertTrue(input2.state == State.CLOSED); + op.close(true); + } +} From 824cec5ab63921fde9a70b1720cb515c6f408e86 Mon Sep 17 00:00:00 2001 From: Paul Rogers Date: Tue, 14 Jun 2022 15:58:17 -0700 Subject: [PATCH 02/11] Intellij inspections fixes --- .../org/apache/druid/queryng/operators/MockCursor.java | 3 ++- .../druid/queryng/operators/MockStorageAdapter.java | 10 +++++----- 2 files changed, 7 insertions(+), 6 deletions(-) diff --git a/processing/src/test/java/org/apache/druid/queryng/operators/MockCursor.java b/processing/src/test/java/org/apache/druid/queryng/operators/MockCursor.java index 84ff056945a8..6b6bdf941c65 100644 --- a/processing/src/test/java/org/apache/druid/queryng/operators/MockCursor.java +++ b/processing/src/test/java/org/apache/druid/queryng/operators/MockCursor.java @@ -19,6 +19,7 @@ package org.apache.druid.queryng.operators; +import org.apache.druid.java.util.common.DateTimes; import org.apache.druid.java.util.common.ISE; import org.apache.druid.query.dimension.DimensionSpec; import org.apache.druid.query.monomorphicprocessing.RuntimeShapeInspector; @@ -170,7 +171,7 @@ public ColumnCapabilities getColumnCapabilities(String column) @Override public DateTime getTime() { - return new DateTime(segmentBase + posn / divideBy); + return DateTimes.utc(segmentBase + posn / divideBy); } @Override diff --git a/processing/src/test/java/org/apache/druid/queryng/operators/MockStorageAdapter.java b/processing/src/test/java/org/apache/druid/queryng/operators/MockStorageAdapter.java index 0148e02a2cd1..be9d6a0a9d88 100644 --- a/processing/src/test/java/org/apache/druid/queryng/operators/MockStorageAdapter.java +++ b/processing/src/test/java/org/apache/druid/queryng/operators/MockStorageAdapter.java @@ -19,6 +19,8 @@ package org.apache.druid.queryng.operators; +import org.apache.druid.java.util.common.DateTimes; +import org.apache.druid.java.util.common.Intervals; import org.apache.druid.java.util.common.granularity.Granularity; import org.apache.druid.java.util.common.guava.Sequence; import org.apache.druid.java.util.common.guava.Sequences; @@ -54,7 +56,7 @@ public Sequence makeCursors( @Override public Interval getInterval() { - return Interval.parse("2015-09-12T13:00:00.000Z/2015-09-12T14:00:00.000Z"); + return Intervals.of("2015-09-12T13:00:00.000Z/2015-09-12T14:00:00.000Z"); } @Override @@ -78,13 +80,13 @@ public int getDimensionCardinality(String column) @Override public DateTime getMinTime() { - return DateTime.parse("2015-09-12T13:00:00.000Z"); + return DateTimes.of("2015-09-12T13:00:00.000Z"); } @Override public DateTime getMaxTime() { - return DateTime.parse("2015-09-12T13:59:59.999Z"); + return DateTimes.of("2015-09-12T13:59:59.999Z"); } @Override @@ -118,7 +120,6 @@ public Comparable getMaxValue(String column) @Override public ColumnCapabilities getColumnCapabilities(String column) { - // TODO Auto-generated method stub return null; } @@ -137,7 +138,6 @@ public DateTime getMaxIngestedEventTime() @Override public Metadata getMetadata() { - // TODO Auto-generated method stub return null; } } From 9fa862cfcd0d126140af1eaa715a79649b8a9f63 Mon Sep 17 00:00:00 2001 From: Paul Rogers Date: Tue, 14 Jun 2022 18:12:06 -0700 Subject: [PATCH 03/11] Many revisions to make static checks happy --- .../java/org/apache/druid/queryng/Timer.java | 69 -------- .../druid/queryng/config/QueryNGConfig.java | 1 + .../queryng/fragment/FragmentContext.java | 3 - .../queryng/fragment/FragmentContextImpl.java | 30 +--- .../druid/queryng/guice/QueryNGModule.java | 48 ----- .../queryng/operators/FilterOperator.java} | 8 +- .../queryng/operators/LimitOperator.java | 6 +- .../queryng/operators/MappingOperator.java | 2 - .../druid/queryng/operators/Operator.java | 88 ++++++---- .../druid/queryng/operators/Operators.java | 13 -- .../queryng/operators/SequenceIterator.java | 81 --------- .../queryng/operators/TransformOperator.java | 4 +- .../queryng/operators/WrappingOperator.java | 5 +- .../general/QueryRunnerOperator.java | 70 -------- .../queryng/operators/BasicOperatorTest.java | 166 ++++++++++++++++++ .../queryng/operators/ConcatOperatorTest.java | 38 +++- .../druid/queryng/operators/FragmentTest.java | 98 ++++++++++- .../druid/queryng/operators/MockCursor.java | 3 +- .../druid/queryng/operators/MockOperator.java | 2 +- .../queryng/operators/MockOperatorTest.java | 37 ---- .../OrderedMergeOperatorTest.java | 44 ++++- 21 files changed, 407 insertions(+), 409 deletions(-) delete mode 100644 processing/src/main/java/org/apache/druid/queryng/Timer.java delete mode 100644 processing/src/main/java/org/apache/druid/queryng/guice/QueryNGModule.java rename processing/src/{test/java/org/apache/druid/queryng/operators/MockFilterOperator.java => main/java/org/apache/druid/queryng/operators/FilterOperator.java} (82%) delete mode 100644 processing/src/main/java/org/apache/druid/queryng/operators/SequenceIterator.java delete mode 100644 processing/src/main/java/org/apache/druid/queryng/operators/general/QueryRunnerOperator.java create mode 100644 processing/src/test/java/org/apache/druid/queryng/operators/BasicOperatorTest.java rename processing/src/test/java/org/apache/druid/queryng/operators/{general => }/OrderedMergeOperatorTest.java (77%) diff --git a/processing/src/main/java/org/apache/druid/queryng/Timer.java b/processing/src/main/java/org/apache/druid/queryng/Timer.java deleted file mode 100644 index b55e1eeab58e..000000000000 --- a/processing/src/main/java/org/apache/druid/queryng/Timer.java +++ /dev/null @@ -1,69 +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.queryng; - -/** - * Very simple nano-second timer with an on/off switch. - */ -public class Timer -{ - private long totalTime; - private long startTime; - - public static Timer create() - { - return new Timer(); - } - - public static Timer createStarted() - { - Timer timer = create(); - timer.start(); - return timer; - } - - public static Timer createAt(long timeNs) - { - Timer timer = create(); - timer.startTime = timeNs; - return timer; - } - - public void start() - { - if (startTime == 0) { - startTime = System.nanoTime(); - } - } - - public void stop() - { - if (startTime != 0) { - totalTime += System.nanoTime() - startTime; - startTime = 0; - } - } - - public long get() - { - stop(); - return totalTime; - } -} diff --git a/processing/src/main/java/org/apache/druid/queryng/config/QueryNGConfig.java b/processing/src/main/java/org/apache/druid/queryng/config/QueryNGConfig.java index 3fb0c1121906..7aa989d1f7d0 100644 --- a/processing/src/main/java/org/apache/druid/queryng/config/QueryNGConfig.java +++ b/processing/src/main/java/org/apache/druid/queryng/config/QueryNGConfig.java @@ -37,6 +37,7 @@ public class QueryNGConfig /** * Create an instance for testing. */ + @SuppressWarnings("unused") // To be used later public static QueryNGConfig create(boolean enabled) { QueryNGConfig config = new QueryNGConfig(); diff --git a/processing/src/main/java/org/apache/druid/queryng/fragment/FragmentContext.java b/processing/src/main/java/org/apache/druid/queryng/fragment/FragmentContext.java index c9b52d807048..595bddf2847a 100644 --- a/processing/src/main/java/org/apache/druid/queryng/fragment/FragmentContext.java +++ b/processing/src/main/java/org/apache/druid/queryng/fragment/FragmentContext.java @@ -19,7 +19,6 @@ package org.apache.druid.queryng.fragment; -import org.apache.druid.query.SegmentDescriptor; import org.apache.druid.query.context.ResponseContext; /** @@ -47,8 +46,6 @@ enum State */ void checkTimeout(); - void missingSegment(SegmentDescriptor descriptor); - /** * Reports the exception, if any, that terminated the fragment. * Should be non-null only if the state is {@code FAILED}. diff --git a/processing/src/main/java/org/apache/druid/queryng/fragment/FragmentContextImpl.java b/processing/src/main/java/org/apache/druid/queryng/fragment/FragmentContextImpl.java index 71313c0a82c2..fed4be62f58f 100644 --- a/processing/src/main/java/org/apache/druid/queryng/fragment/FragmentContextImpl.java +++ b/processing/src/main/java/org/apache/druid/queryng/fragment/FragmentContextImpl.java @@ -23,13 +23,10 @@ import org.apache.druid.java.util.common.JodaUtils; import org.apache.druid.java.util.common.StringUtils; import org.apache.druid.query.QueryTimeoutException; -import org.apache.druid.query.SegmentDescriptor; import org.apache.druid.query.context.ResponseContext; import org.apache.druid.queryng.operators.Operator; -import java.util.ArrayList; import java.util.Deque; -import java.util.List; import java.util.concurrent.ConcurrentLinkedDeque; public class FragmentContextImpl implements FragmentContext @@ -41,7 +38,7 @@ public class FragmentContextImpl implements FragmentContext private final long startTimeMillis; private final long timeoutAt; protected State state = State.START; - private Exception exception; + private Deque exceptions = new ConcurrentLinkedDeque<>(); protected FragmentContextImpl( final String queryId, @@ -81,7 +78,7 @@ public void register(Operator op) @Override public Exception exception() { - return exception; + return exceptions.peek(); } @Override @@ -98,7 +95,7 @@ public ResponseContext responseContext() public void failed(Exception exception) { - this.exception = exception; + this.exceptions.add(exception); this.state = State.FAILED; } @@ -112,24 +109,6 @@ public void checkTimeout() } } - protected void recordRunTime() - { - if (timeoutAt == 0) { - return; - } - // This is very likely wrong - responseContext.put( - ResponseContext.Keys.TIMEOUT_AT, - timeoutAt - (System.currentTimeMillis() - startTimeMillis) - ); - } - - @Override - public void missingSegment(SegmentDescriptor descriptor) - { - responseContext.add(ResponseContext.Keys.MISSING_SEGMENTS, descriptor); - } - /** * Closes all operators from the leaves to the root. * As a result, operators must not call their children during @@ -144,7 +123,6 @@ protected void close() if (state == State.CLOSED) { return; } - List exceptions = new ArrayList<>(); Operator op; while ((op = operators.pollFirst()) != null) { try { @@ -154,8 +132,6 @@ protected void close() exceptions.add(e); } } - // TODO: Do something with the exceptions - recordRunTime(); state = State.CLOSED; } } diff --git a/processing/src/main/java/org/apache/druid/queryng/guice/QueryNGModule.java b/processing/src/main/java/org/apache/druid/queryng/guice/QueryNGModule.java deleted file mode 100644 index 1d1c351da7e9..000000000000 --- a/processing/src/main/java/org/apache/druid/queryng/guice/QueryNGModule.java +++ /dev/null @@ -1,48 +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.queryng.guice; - -import com.google.inject.Binder; -import com.google.inject.Module; -import org.apache.druid.guice.JsonConfigProvider; -import org.apache.druid.guice.LazySingleton; -import org.apache.druid.queryng.config.QueryNGConfig; -import org.apache.druid.queryng.fragment.FragmentBuilderFactory; -import org.apache.druid.queryng.fragment.FragmentBuilderFactoryImpl; - -/** - * Configure the "shim" version of the NG query engine which entails - * creating a config (to enable or disable the engine) and to create - * a factory for the fragment context. In this early version, all - * other parts of the engine are distributed across various query - * runners. - */ -public class QueryNGModule implements Module -{ - @Override - public void configure(Binder binder) - { - JsonConfigProvider.bind(binder, QueryNGConfig.CONFIG_ROOT, QueryNGConfig.class); - binder - .bind(FragmentBuilderFactory.class) - .to(FragmentBuilderFactoryImpl.class) - .in(LazySingleton.class); - } -} diff --git a/processing/src/test/java/org/apache/druid/queryng/operators/MockFilterOperator.java b/processing/src/main/java/org/apache/druid/queryng/operators/FilterOperator.java similarity index 82% rename from processing/src/test/java/org/apache/druid/queryng/operators/MockFilterOperator.java rename to processing/src/main/java/org/apache/druid/queryng/operators/FilterOperator.java index 81e0514e24f4..3d31862c2be7 100644 --- a/processing/src/test/java/org/apache/druid/queryng/operators/MockFilterOperator.java +++ b/processing/src/main/java/org/apache/druid/queryng/operators/FilterOperator.java @@ -24,15 +24,15 @@ import java.util.function.Predicate; /** - * Super-simple filter operator which allows tests to build up a - * stack of operators starting with values generated by MockOperator. + * Super-simple filter operator with the filter provided by a + * predicate. */ -public class MockFilterOperator extends MappingOperator +public class FilterOperator extends MappingOperator { private final Predicate predicate; private T nextValue; - public MockFilterOperator(FragmentContext context, Operator input, Predicate pred) + public FilterOperator(FragmentContext context, Operator input, Predicate pred) { super(context, input); this.predicate = pred; diff --git a/processing/src/main/java/org/apache/druid/queryng/operators/LimitOperator.java b/processing/src/main/java/org/apache/druid/queryng/operators/LimitOperator.java index 93864d118f1e..0235774d18fb 100644 --- a/processing/src/main/java/org/apache/druid/queryng/operators/LimitOperator.java +++ b/processing/src/main/java/org/apache/druid/queryng/operators/LimitOperator.java @@ -25,18 +25,18 @@ * Limits the results from the input operator to the given number * of rows. */ -public abstract class LimitOperator extends MappingOperator +public class LimitOperator extends MappingOperator { public static final long UNLIMITED = Long.MAX_VALUE; protected final long limit; protected long rowCount; - protected int batchCount; - public LimitOperator(FragmentContext context, long limit, Operator input) + public LimitOperator(FragmentContext context, Operator input, long limit) { super(context, input); this.limit = limit; + context.register(this); } @Override diff --git a/processing/src/main/java/org/apache/druid/queryng/operators/MappingOperator.java b/processing/src/main/java/org/apache/druid/queryng/operators/MappingOperator.java index 877532030507..c2fa7e286cf2 100644 --- a/processing/src/main/java/org/apache/druid/queryng/operators/MappingOperator.java +++ b/processing/src/main/java/org/apache/druid/queryng/operators/MappingOperator.java @@ -32,14 +32,12 @@ */ public abstract class MappingOperator implements IterableOperator { - protected final FragmentContext context; private final Operator input; protected Iterator inputIter; protected State state = State.START; public MappingOperator(FragmentContext context, Operator input) { - this.context = context; this.input = input; context.register(this); } diff --git a/processing/src/main/java/org/apache/druid/queryng/operators/Operator.java b/processing/src/main/java/org/apache/druid/queryng/operators/Operator.java index 4e19b913205b..fda745d8fb65 100644 --- a/processing/src/main/java/org/apache/druid/queryng/operators/Operator.java +++ b/processing/src/main/java/org/apache/druid/queryng/operators/Operator.java @@ -22,8 +22,8 @@ import java.util.Iterator; /** - * An operator is a data pipeline transform: something that changes a stream of - * results in some way. An operator has a very simple lifecycle: + * An operator is a data pipeline transform: something that operates on + * a stream of results in some way. An operator has a very simple lifecycle: *

*

    *
  • Created.
    • @@ -33,22 +33,24 @@ *
*

* Leaf operators produce results, internal operators transform them, and root - * operators do something with the results. Operators know nothing about their + * operators deliver results to some consumer. Operators know nothing about their * children other than that they follow the operator protocol and will produce a * result when asked. Operators must agree on the type of the shared results. *

* Unlike traditional QueryRunners, an operator does not create its - * children: that is the job of the planner that created a tree of operator - * definitions. The operator simply accepts the previously-created children and - * does its thing. - *

+ * children: that is the job of the planner that created a tree of operators + * The operator simply accepts the previously-created children and does its thing. + * + *

State

+ * * Operators are assumed to be stateful since most data transform operations - * involve some kind of state. The {@code start()} and {@code close()} - * methods provide a well-defined way - * to handle resource. Operators are created as a DAG. Unlike a simple iterator, + * involve some kind of state. + *

+ * The {@code open()} and {@code close()} methods provide a well-defined way + * to handle resources. Operators are created as a DAG. Unlike a simple iterator, * operators should not obtain resources in their constructors. Instead, * they should obtain resources, open files or otherwise start things whirring in - * the {#code start()} method. In some cases, resource may be obtained a bit later, + * the {#code open()} method. In some cases, resource may be obtained a bit later, * in the first call to {@code hasNext()}. In either case, resources should be * released in {@code close()}. *

@@ -57,53 +59,79 @@ * runs it, closes it, then moves to the next child. This form of operation ensures * resources are held for the briefest possible period of time. *

- * To make this work, operators should cascade start() and + * To make this work, operators cascade open() and * close() operations to their children. The query runner will * start the root operator: the root must start is children, and so on. Closing is * a bit more complex. Operators should close their children by cascading the * close operation: but only if that call came from a parent operator (as - * indicated by the {@code cascade} parameter set to {@code true}.) + * indicated by the {@code cascade} parameter set to {@code true}.) In general, + * each operator should open is children as late as possible, and close them as + * early as possible. *

- * The {@link FragmentRunner} will ensure all operators are closed by calling: + * A fragment runner will ensure all operators are closed by calling * close, from the bottom up, at the completion (successful or otherwise) of the * query. In this case, the {@code cascade} parameter set to {@code false}, * and each operator should not cascade this call to its children: the * fragment runner will do so. *

- * Implementations can assume that calls to next() and - * get() occur within a single thread, so state can be maintained - * in normal (not atomic) variables. - *

* Operators are one-pass: they are not re-entrant and cannot be restarted. We * assume that they read ephemeral values: once returned, they are gone. We also * assume the results are volatile: once read, we may not be able to read the * same set of values a second time even if we started the whole process over. *

- * The type of the values returned by the operator is determined by context. - * Druid native queries use a variety of Java objects: there is no single - * "row" class or interface. - *

- * Operators do not have a return type parameter. Operators are generally created - * dynamically: that code is far simpler without having to deal with unknown - * types. Even test code will often call {@code assertEquals()} and the like - * which don't need the type. - *

* Having {@code open()} return the iterator for results accomplishes two goals. * First is the minor benefit of ensuring that an operator is opened before * fetching results. More substantially, this approach allows "wrapper" operators * which only perform work in the open or close method. For those, the open * method returns the iterator of the child, avoiding the overhead of pass-through * calls for each data batch. The wrapper operator will not sit on the data - * path, only on the control (open/close) path. + * path, only on the control (open/close) path. See {@link WrappingOperator} for + * and example. + * + *

Single-Threaded

+ * + * Implementations can assume that calls to the operator methods, or to the + * iterator hasNext() and next() methods, always + * occur within a single thread, so state can be maintained + * in normal (not atomic) variables. If data is to be shuffled across threads, + * then a thread- (or network-)aware shuffle mechanism is required. + * + *

Type

+ * + * The type of the values returned by the operator is determined by its type. + * Druid native queries use a variety of Java objects: there is no single + * "row" class or interface. + *

+ * Operator type specifies the type of values returned by its iterator. + * Most operators require that the input and output types are the same, + * though some may transform data from one type to another. See + * {@link TransformOperator} for an example. + * + *

Errors

+ * + * Any operator may fail for any number of reasons: I/O errors, resource + * exhaustion, invalid state (such as divide-by-zero), code bugs, etc. + * When an operator fails, it simply throws an exception. The exception + * unwinds the stack up to the fragment runner, which is then responsible + * for calling close on the fragment context. That then cascades close down + * to all operators as described above. It may be that one or more operators + * are not in a bad state, and throw an exception on close. The fragment + * context "absorbs" those errors and continues to close all operators. + *

+ * The fragment context has no knowledge of whether an operator was called + * early. So, if an operator is particular, and wants to be closed only once, + * then that operator should maintain state and ignore subsequent calls to + * close. See {@link WrappingOperator} for a simple example. * - * @param the type of the object (row, batch) returned by {@link #next()}. + * @param the type of the object (row, batch) returned by the iterator + * returned from {@link #open()}. */ public interface Operator { /** * Convenience interface for an operator which is its own iterator. */ - public interface IterableOperator extends Operator, Iterator + interface IterableOperator extends Operator, Iterator { } diff --git a/processing/src/main/java/org/apache/druid/queryng/operators/Operators.java b/processing/src/main/java/org/apache/druid/queryng/operators/Operators.java index 602baf9afcf0..5f513c69e873 100644 --- a/processing/src/main/java/org/apache/druid/queryng/operators/Operators.java +++ b/processing/src/main/java/org/apache/druid/queryng/operators/Operators.java @@ -24,10 +24,8 @@ import org.apache.druid.java.util.common.guava.Sequence; import org.apache.druid.query.Query; import org.apache.druid.query.QueryPlus; -import org.apache.druid.query.QueryRunner; import org.apache.druid.query.scan.ScanQuery; import org.apache.druid.queryng.fragment.DAGBuilder; -import org.apache.druid.queryng.operators.general.QueryRunnerOperator; import java.util.Iterator; import java.util.List; @@ -170,17 +168,6 @@ public static Operator unwrapOperator(Sequence sequence) return null; } - /** - * Create an operator which wraps a query runner which allows a query runner - * to be an input to an operator. The runner, and its sequence, will be optimized - * away at runtime if both the upstream and downstream items are both operators, - * but the shim is left in place if the upstream is actually a query runner. - */ - public static QueryRunnerOperator toOperator(QueryRunner runner, QueryPlus query) - { - return new QueryRunnerOperator(runner, query); - } - /** * This will materialize the entire sequence from the wrapped * operator. Use at your own risk. diff --git a/processing/src/main/java/org/apache/druid/queryng/operators/SequenceIterator.java b/processing/src/main/java/org/apache/druid/queryng/operators/SequenceIterator.java deleted file mode 100644 index b267a47ccaf7..000000000000 --- a/processing/src/main/java/org/apache/druid/queryng/operators/SequenceIterator.java +++ /dev/null @@ -1,81 +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.queryng.operators; - -import com.google.common.base.Preconditions; -import org.apache.druid.java.util.common.guava.Sequence; -import org.apache.druid.java.util.common.guava.Yielder; -import org.apache.druid.java.util.common.guava.YieldingAccumulator; - -import java.io.IOException; -import java.util.Iterator; - -/** - * Iterator over a sequence. - */ -public class SequenceIterator implements Iterator, AutoCloseable -{ - private Yielder yielder; - - public static SequenceIterator of(Sequence sequence) - { - return new SequenceIterator(sequence); - } - - public SequenceIterator(Sequence sequence) - { - this.yielder = sequence.toYielder( - null, - new YieldingAccumulator() - { - @Override - public T accumulate(T accumulated, T in) - { - yield(); - return in; - } - } - ); - } - - @Override - public boolean hasNext() - { - return !yielder.isDone(); - } - - @Override - public T next() - { - Preconditions.checkState(!yielder.isDone()); - T value = yielder.get(); - yielder = yielder.next(null); - return value; - } - - @Override - public void close() throws IOException - { - if (yielder != null) { - yielder.close(); - yielder = null; - } - } -} diff --git a/processing/src/main/java/org/apache/druid/queryng/operators/TransformOperator.java b/processing/src/main/java/org/apache/druid/queryng/operators/TransformOperator.java index bdfac97ad8a3..90414f7d6db0 100644 --- a/processing/src/main/java/org/apache/druid/queryng/operators/TransformOperator.java +++ b/processing/src/main/java/org/apache/druid/queryng/operators/TransformOperator.java @@ -47,8 +47,8 @@ public class TransformOperator extends MappingOperator public TransformOperator( FragmentContext context, - final Function transformFn, - final Operator input) + final Operator input, + final Function transformFn) { super(context, input); this.transformFn = transformFn; diff --git a/processing/src/main/java/org/apache/druid/queryng/operators/WrappingOperator.java b/processing/src/main/java/org/apache/druid/queryng/operators/WrappingOperator.java index 922f00419a7d..f9b497375616 100644 --- a/processing/src/main/java/org/apache/druid/queryng/operators/WrappingOperator.java +++ b/processing/src/main/java/org/apache/druid/queryng/operators/WrappingOperator.java @@ -55,8 +55,11 @@ public Iterator open() @Override public void close(boolean cascade) { + if (state != State.RUN) { + return; + } try { - if (state == State.RUN && cascade) { + if (cascade) { input.close(cascade); } } diff --git a/processing/src/main/java/org/apache/druid/queryng/operators/general/QueryRunnerOperator.java b/processing/src/main/java/org/apache/druid/queryng/operators/general/QueryRunnerOperator.java deleted file mode 100644 index 10c5524ff11c..000000000000 --- a/processing/src/main/java/org/apache/druid/queryng/operators/general/QueryRunnerOperator.java +++ /dev/null @@ -1,70 +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.queryng.operators.general; - -import org.apache.druid.java.util.common.guava.Sequence; -import org.apache.druid.query.QueryPlus; -import org.apache.druid.query.QueryRunner; -import org.apache.druid.queryng.fragment.FragmentContext; -import org.apache.druid.queryng.operators.Operator; -import org.apache.druid.queryng.operators.Operators; - -import java.util.Iterator; - -/** - * Operator which wraps a query runner. When used in the interim - * "shim" architecture, this operator allows a query runner to be - * the input (upstream) to some other (downstream) operator. The - * upstream query runner may give rise to its own operator. In that - * case, at runtime, the sequence wrapper for that operator is - * optimized away, leaving just the two operators. - */ -public class QueryRunnerOperator implements Operator -{ - protected final FragmentContext context; - private final QueryRunner runner; - private final QueryPlus query; - private Operator child; - - public QueryRunnerOperator(QueryRunner runner, QueryPlus query) - { - this.context = query.fragmentBuilder().context(); - this.runner = runner; - this.query = query; - context.register(this); - } - - @Override - public Iterator open() - { - Sequence seq = runner.run(query, context.responseContext()); - child = Operators.toOperator(context, seq); - return child.open(); - } - - @Override - public void close(boolean cascade) - { - if (child != null && cascade) { - child.close(cascade); - } - child = null; - } -} diff --git a/processing/src/test/java/org/apache/druid/queryng/operators/BasicOperatorTest.java b/processing/src/test/java/org/apache/druid/queryng/operators/BasicOperatorTest.java new file mode 100644 index 000000000000..9d4e09face93 --- /dev/null +++ b/processing/src/test/java/org/apache/druid/queryng/operators/BasicOperatorTest.java @@ -0,0 +1,166 @@ +/* + * 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.queryng.operators; + +import org.apache.druid.java.util.common.guava.Sequence; +import org.apache.druid.queryng.fragment.FragmentBuilder; +import org.apache.druid.queryng.fragment.FragmentContext; +import org.apache.druid.queryng.fragment.FragmentRun; +import org.junit.Test; + +import java.util.Arrays; +import java.util.Iterator; +import java.util.List; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertTrue; + +public class BasicOperatorTest +{ + @Test + public void testFilter() + { + FragmentBuilder builder = FragmentBuilder.defaultBuilder(); + MockOperator op = MockOperator.ints(builder.context(), 4); + Operator op2 = new FilterOperator( + builder.context(), + op, + x -> x % 2 == 0); + List results = builder.run(op2).toList(); + assertEquals(Arrays.asList(0, 2), results); + assertEquals(Operator.State.CLOSED, op.state); + } + + @Test + public void testTransform() + { + FragmentBuilder builder = FragmentBuilder.defaultBuilder(); + MockOperator op = MockOperator.ints(builder.context(), 4); + Operator op2 = new TransformOperator( + builder.context(), + op, + x -> x * 2); + List results = builder.run(op2).toList(); + assertEquals(Arrays.asList(0, 2, 4, 6), results); + assertEquals(Operator.State.CLOSED, op.state); + } + + @Test + public void testLimit0() + { + FragmentBuilder builder = FragmentBuilder.defaultBuilder(); + MockOperator op = MockOperator.ints(builder.context(), 4); + Operator op2 = new LimitOperator( + builder.context(), + op, + 0); + List results = builder.run(op2).toList(); + assertTrue(results.isEmpty()); + assertEquals(Operator.State.CLOSED, op.state); + } + + @Test + public void testLimitNotHit() + { + FragmentBuilder builder = FragmentBuilder.defaultBuilder(); + MockOperator op = MockOperator.ints(builder.context(), 4); + Operator op2 = new LimitOperator( + builder.context(), + op, + 8); + List results = builder.run(op2).toList(); + assertEquals(Arrays.asList(0, 1, 2, 3), results); + } + + @Test + public void testLimitHit() + { + FragmentBuilder builder = FragmentBuilder.defaultBuilder(); + MockOperator op = MockOperator.ints(builder.context(), 4); + Operator op2 = new LimitOperator( + builder.context(), + op, + 2); + List results = builder.run(op2).toList(); + assertEquals(Arrays.asList(0, 1), results); + } + + private static class MockWrappingOperator extends WrappingOperator + { + boolean openCalled; + boolean closeCalled; + + public MockWrappingOperator(FragmentContext context, Operator input) + { + super(context, input); + } + + @Override + protected void onOpen() + { + openCalled = true; + // Silly check just to create a reference to the context variable. + assertNull(context.exception()); + } + + @Override + protected void onClose() + { + closeCalled = true; + } + } + + @Test + public void testWrappingOperator() + { + FragmentBuilder builder = FragmentBuilder.defaultBuilder(); + MockOperator op = MockOperator.ints(builder.context(), 4); + MockWrappingOperator op2 = new MockWrappingOperator( + builder.context(), + op); + List results = builder.run(op2).toList(); + assertEquals(Arrays.asList(0, 1, 2, 3), results); + assertTrue(op2.openCalled); + assertTrue(op2.closeCalled); + } + + /** + * Getting weird: an operator that wraps a sequence that wraps an operator. + */ + @Test + public void testSequenceOperator() + { + FragmentBuilder builder = FragmentBuilder.defaultBuilder(); + MockOperator op = MockOperator.strings(builder.context(), 2); + Sequence seq = Operators.toSequence(op); + Operator outer = Operators.toOperator(builder, seq); + FragmentRun run = builder.run(outer); + Iterator iter = run.iterator(); + assertTrue(iter.hasNext()); + assertEquals("Mock row 0", iter.next()); + assertTrue(iter.hasNext()); + assertEquals("Mock row 1", iter.next()); + assertFalse(iter.hasNext()); + run.close(); + assertEquals(Operator.State.CLOSED, op.state); + } +} diff --git a/processing/src/test/java/org/apache/druid/queryng/operators/ConcatOperatorTest.java b/processing/src/test/java/org/apache/druid/queryng/operators/ConcatOperatorTest.java index ac3d32b5bbdd..a91d56b05721 100644 --- a/processing/src/test/java/org/apache/druid/queryng/operators/ConcatOperatorTest.java +++ b/processing/src/test/java/org/apache/druid/queryng/operators/ConcatOperatorTest.java @@ -29,6 +29,7 @@ import java.util.List; import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertSame; import static org.junit.Assert.assertTrue; public class ConcatOperatorTest @@ -47,18 +48,31 @@ public void testOneEmptyInput() { FragmentContext context = FragmentContext.defaultContext(); Operator input = new NullOperator(context); - Operator op = new ConcatOperator(context, Arrays.asList(input)); + Operator op = new ConcatOperator( + context, + Collections.singletonList(input)); List results = Operators.toList(op); assertTrue(results.isEmpty()); } + @Test + public void testHelperNoConcat() + { + FragmentContext context = FragmentContext.defaultContext(); + MockOperator input1 = MockOperator.ints(context, 2); + Operator op = ConcatOperator.concatOrNot(context, Collections.singletonList(input1)); + assertSame(input1, op); + } + @Test public void testTwoEmptyInputs() { FragmentContext context = FragmentContext.defaultContext(); Operator input1 = new NullOperator(context); Operator input2 = new NullOperator(context); - Operator op = new ConcatOperator(context, Arrays.asList(input1, input2)); + Operator op = new ConcatOperator( + context, + Arrays.asList(input1, input2)); List results = Operators.toList(op); assertTrue(results.isEmpty()); } @@ -68,7 +82,9 @@ public void testOneInput() { FragmentContext context = FragmentContext.defaultContext(); Operator input = MockOperator.ints(context, 2); - Operator op = new ConcatOperator(context, Arrays.asList(input)); + Operator op = new ConcatOperator( + context, + Collections.singletonList(input)); List results = Operators.toList(op); List expected = Arrays.asList(0, 1); assertEquals(expected, results); @@ -80,7 +96,9 @@ public void testEmptyThenNonEmptyInputs() FragmentContext context = FragmentContext.defaultContext(); Operator input1 = new NullOperator(context); Operator input2 = MockOperator.ints(context, 2); - Operator op = new ConcatOperator(context, Arrays.asList(input1, input2)); + Operator op = ConcatOperator.concatOrNot( + context, + Arrays.asList(input1, input2)); List results = Operators.toList(op); List expected = Arrays.asList(0, 1); assertEquals(expected, results); @@ -92,7 +110,9 @@ public void testNonEmptyThenEmptyInputs() FragmentContext context = FragmentContext.defaultContext(); Operator input1 = MockOperator.ints(context, 2); Operator input2 = new NullOperator(context); - Operator op = new ConcatOperator(context, Arrays.asList(input1, input2)); + Operator op = ConcatOperator.concatOrNot( + context, + Arrays.asList(input1, input2)); List results = Operators.toList(op); List expected = Arrays.asList(0, 1); assertEquals(expected, results); @@ -104,7 +124,9 @@ public void testTwoInputs() FragmentContext context = FragmentContext.defaultContext(); Operator input1 = MockOperator.ints(context, 2); Operator input2 = MockOperator.ints(context, 2); - Operator op = new ConcatOperator(context, Arrays.asList(input1, input2)); + Operator op = ConcatOperator.concatOrNot( + context, + Arrays.asList(input1, input2)); List results = Operators.toList(op); List expected = Arrays.asList(0, 1, 0, 1); assertEquals(expected, results); @@ -116,7 +138,9 @@ public void testClose() FragmentContext context = FragmentContext.defaultContext(); MockOperator input1 = MockOperator.ints(context, 2); MockOperator input2 = MockOperator.ints(context, 2); - Operator op = new ConcatOperator(context, Arrays.asList(input1, input2)); + Operator op = ConcatOperator.concatOrNot( + context, + Arrays.asList(input1, input2)); Iterator iter = op.open(); assertTrue(iter.hasNext()); assertEquals(0, (int) iter.next()); diff --git a/processing/src/test/java/org/apache/druid/queryng/operators/FragmentTest.java b/processing/src/test/java/org/apache/druid/queryng/operators/FragmentTest.java index 40213b8eebf2..3ed0c002dd6b 100644 --- a/processing/src/test/java/org/apache/druid/queryng/operators/FragmentTest.java +++ b/processing/src/test/java/org/apache/druid/queryng/operators/FragmentTest.java @@ -19,12 +19,23 @@ package org.apache.druid.queryng.operators; +import com.google.common.collect.ImmutableMap; import org.apache.druid.java.util.common.ISE; import org.apache.druid.java.util.common.guava.Sequence; import org.apache.druid.java.util.common.guava.SequenceTestHelper; +import org.apache.druid.query.Druids; +import org.apache.druid.query.Query; +import org.apache.druid.query.QueryPlus; +import org.apache.druid.query.context.ResponseContext; +import org.apache.druid.queryng.config.QueryNGConfig; import org.apache.druid.queryng.fragment.FragmentBuilder; +import org.apache.druid.queryng.fragment.FragmentBuilderFactory; +import org.apache.druid.queryng.fragment.FragmentBuilderFactoryImpl; +import org.apache.druid.queryng.fragment.FragmentContext; +import org.apache.druid.queryng.fragment.FragmentContextImpl; import org.apache.druid.queryng.fragment.FragmentHandle; import org.apache.druid.queryng.fragment.FragmentRun; +import org.apache.druid.queryng.fragment.NullFragmentBuilderFactory; import org.junit.Test; import java.io.IOException; @@ -59,12 +70,13 @@ public void testOperatorBasics() assertTrue(handle.rootIsOperator()); assertFalse(handle.rootIsSequence()); assertSame(op, handle.rootOperator()); + assertSame(builder, handle.builder()); assertNull(handle.rootSequence()); assertSame(handle, handle.toOperator()); assertSame(builder.context(), handle.context()); // Add another operator to the DAG. - Operator op2 = new MockFilterOperator(builder.context(), op, x -> x < 3); + Operator op2 = new FilterOperator(builder.context(), op, x -> x < 3); FragmentHandle handle2 = handle.compose(op2); assertTrue(handle2.rootIsOperator()); assertFalse(handle2.rootIsSequence()); @@ -112,7 +124,7 @@ public void testEmptyHandle() public void testSequenceBasics() { FragmentBuilder builder = FragmentBuilder.defaultBuilder(); - MockOperator op = MockOperator.ints(builder.context(), 2); + MockOperator op = MockOperator.ints(builder.context(), 4); Sequence seq = Operators.toSequence(op); // Handle for the leaf sequence @@ -128,7 +140,7 @@ public void testSequenceBasics() // Add another operator to the DAG by unwrapping the above // sequence. - Operator op2 = new MockFilterOperator( + Operator op2 = new FilterOperator( builder.context(), handle.toOperator().rootOperator(), x -> x < 3); @@ -142,6 +154,18 @@ public void testSequenceBasics() assertNotNull(handle2.rootSequence()); assertSame(handle2, handle2.toSequence()); assertSame(builder.context(), handle2.context()); + + // Add a sequence. + Operator op3 = new FilterOperator( + builder.context(), + handle2.toOperator().rootOperator(), + x -> x > 1); + FragmentHandle handle3 = handle2.compose(Operators.toSequence(op3)); + assertFalse(handle3.rootIsOperator()); + assertTrue(handle3.rootIsSequence()); + List results = handle3.runAsSequence().toList(); + assertEquals(1, results.size()); + assertEquals(2, (int) results.get(0)); } @Test @@ -149,13 +173,18 @@ public void testRun() { FragmentBuilder builder = FragmentBuilder.defaultBuilder(); MockOperator op = MockOperator.ints(builder.context(), 2); + FragmentContext context = builder.context(); + assertEquals(FragmentContext.State.START, context.state()); FragmentHandle handle = builder.handle(op); int i = 0; try (FragmentRun run = handle.run()) { + assertEquals(FragmentContext.State.RUN, context.state()); + assertSame(context, run.context()); for (Integer value : run) { assertEquals(i++, (int) value); } } + assertEquals(FragmentContext.State.CLOSED, context.state()); assertEquals(Operator.State.CLOSED, op.state); } @@ -235,7 +264,7 @@ public void testSequenceYielder() throws IOException } @Test - public void testSequenceAccum() throws IOException + public void testSequenceAccum() { FragmentBuilder builder = FragmentBuilder.defaultBuilder(); MockOperator op = MockOperator.ints(builder.context(), 4); @@ -262,4 +291,65 @@ public void testRunEmptyHandleAsSequence() Sequence seq = emptyHandle.runAsSequence(); SequenceTestHelper.testAccumulation("empty", seq, Collections.emptyList()); } + + @Test + public void testFactory() + { + Query query = new Druids.ScanQueryBuilder() + .dataSource("foo") + .eternityInterval() + .build(); + + // Operators blocked by query: no gating context variable + QueryNGConfig enableConfig = QueryNGConfig.create(true); + assertTrue(enableConfig.enabled()); + FragmentBuilderFactory enableFactory = new FragmentBuilderFactoryImpl(enableConfig); + assertNull(enableFactory.create(query, ResponseContext.createEmpty())); + FragmentBuilderFactory nullFactory = new NullFragmentBuilderFactory(); + + QueryNGConfig disableConfig = QueryNGConfig.create(false); + assertFalse(disableConfig.enabled()); + FragmentBuilderFactory disableFactory = new FragmentBuilderFactoryImpl(disableConfig); + assertNull(disableFactory.create(query, ResponseContext.createEmpty())); + assertNull(nullFactory.create(query, ResponseContext.createEmpty())); + + // Enable at query level. Use of operators gated by config. + query = query.withOverriddenContext( + ImmutableMap.of(Operators.CONTEXT_VAR, true)); + assertNotNull(enableFactory.create(query, ResponseContext.createEmpty())); + assertNull(disableFactory.create(query, ResponseContext.createEmpty())); + assertNull(nullFactory.create(query, ResponseContext.createEmpty())); + } + + @Test + public void testQueryPlus() + { + FragmentBuilder builder = FragmentBuilder.defaultBuilder(); + Query query = new Druids.ScanQueryBuilder() + .dataSource("foo") + .eternityInterval() + .build(); + QueryPlus queryPlus = QueryPlus.wrap(query).withFragmentBuilder(builder); + assertSame(builder, queryPlus.fragmentBuilder()); + } + + @Test + public void testFragmentContext() + { + FragmentBuilder builder = FragmentBuilder.defaultBuilder(); + FragmentContext context = builder.context(); + assertEquals(FragmentContext.State.START, context.state()); + assertEquals("unknown", context.queryId()); + assertNotNull(context.responseContext()); + context.checkTimeout(); // Useless here, just prevents a "not used" error + assertNull(context.exception()); + MockOperator op = MockOperator.ints(builder.context(), 4); + FragmentHandle handle = builder.handle(op); + handle.run().iterator(); + assertEquals(FragmentContext.State.RUN, context.state()); + ISE ex = new ISE("oops"); + ((FragmentContextImpl) context).failed(ex); + assertEquals(FragmentContext.State.FAILED, context.state()); + assertSame(ex, context.exception()); + } } diff --git a/processing/src/test/java/org/apache/druid/queryng/operators/MockCursor.java b/processing/src/test/java/org/apache/druid/queryng/operators/MockCursor.java index 6b6bdf941c65..9e30c850d957 100644 --- a/processing/src/test/java/org/apache/druid/queryng/operators/MockCursor.java +++ b/processing/src/test/java/org/apache/druid/queryng/operators/MockCursor.java @@ -74,7 +74,8 @@ public boolean isNull() } } - private class MockStringColumn implements ColumnValueSelector + // Must be static to avoid Java 11 compile errors. + private static class MockStringColumn implements ColumnValueSelector { @Override public long getLong() diff --git a/processing/src/test/java/org/apache/druid/queryng/operators/MockOperator.java b/processing/src/test/java/org/apache/druid/queryng/operators/MockOperator.java index f97853e09290..f1bd21c26319 100644 --- a/processing/src/test/java/org/apache/druid/queryng/operators/MockOperator.java +++ b/processing/src/test/java/org/apache/druid/queryng/operators/MockOperator.java @@ -50,7 +50,7 @@ public static MockOperator ints(FragmentContext context, int rowCount) public static MockOperator strings(FragmentContext context, int rowCount) { - return new MockOperator(context, rowCount, rid -> "Mock row " + Integer.toString(rid)); + return new MockOperator(context, rowCount, rid -> "Mock row " + rid); } @Override diff --git a/processing/src/test/java/org/apache/druid/queryng/operators/MockOperatorTest.java b/processing/src/test/java/org/apache/druid/queryng/operators/MockOperatorTest.java index 4462be4a0321..12876e6b1d1e 100644 --- a/processing/src/test/java/org/apache/druid/queryng/operators/MockOperatorTest.java +++ b/processing/src/test/java/org/apache/druid/queryng/operators/MockOperatorTest.java @@ -19,7 +19,6 @@ package org.apache.druid.queryng.operators; -import org.apache.druid.java.util.common.guava.Sequence; import org.apache.druid.queryng.fragment.FragmentBuilder; import org.apache.druid.queryng.fragment.FragmentHandle; import org.apache.druid.queryng.fragment.FragmentRun; @@ -66,40 +65,4 @@ public void testMockIntOperator() assertEquals(1, (int) results.get(1)); assertEquals(Operator.State.CLOSED, op.state); } - - /** - * Getting weird: an operator that wraps a sequence that wraps an operator. - */ - @Test - public void testSequenceOperator() - { - FragmentBuilder builder = FragmentBuilder.defaultBuilder(); - MockOperator op = MockOperator.strings(builder.context(), 2); - Sequence seq = Operators.toSequence(op); - Operator outer = Operators.toOperator(builder, seq); - FragmentRun run = builder.run(outer); - Iterator iter = run.iterator(); - assertTrue(iter.hasNext()); - assertEquals("Mock row 0", iter.next()); - assertTrue(iter.hasNext()); - assertEquals("Mock row 1", iter.next()); - assertFalse(iter.hasNext()); - run.close(); - assertEquals(Operator.State.CLOSED, op.state); - } - - @Test - public void testMockFilter() - { - FragmentBuilder builder = FragmentBuilder.defaultBuilder(); - MockOperator op = MockOperator.ints(builder.context(), 4); - Operator op2 = new MockFilterOperator( - builder.context(), - op, - x -> x % 2 == 0); - List results = builder.run(op2).toList(); - assertEquals(0, (int) results.get(0)); - assertEquals(2, (int) results.get(1)); - assertEquals(Operator.State.CLOSED, op.state); - } } diff --git a/processing/src/test/java/org/apache/druid/queryng/operators/general/OrderedMergeOperatorTest.java b/processing/src/test/java/org/apache/druid/queryng/operators/OrderedMergeOperatorTest.java similarity index 77% rename from processing/src/test/java/org/apache/druid/queryng/operators/general/OrderedMergeOperatorTest.java rename to processing/src/test/java/org/apache/druid/queryng/operators/OrderedMergeOperatorTest.java index 88dda1789e1e..be90beb1f0a7 100644 --- a/processing/src/test/java/org/apache/druid/queryng/operators/general/OrderedMergeOperatorTest.java +++ b/processing/src/test/java/org/apache/druid/queryng/operators/OrderedMergeOperatorTest.java @@ -17,16 +17,12 @@ * under the License. */ -package org.apache.druid.queryng.operators.general; +package org.apache.druid.queryng.operators; import com.google.common.collect.Lists; import com.google.common.collect.Ordering; import org.apache.druid.queryng.fragment.FragmentContext; -import org.apache.druid.queryng.operators.MockOperator; -import org.apache.druid.queryng.operators.NullOperator; -import org.apache.druid.queryng.operators.Operator; import org.apache.druid.queryng.operators.Operator.State; -import org.apache.druid.queryng.operators.OrderedMergeOperator; import org.apache.druid.queryng.operators.OrderedMergeOperator.Input; import org.junit.Test; @@ -38,6 +34,7 @@ import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNull; import static org.junit.Assert.assertTrue; public class OrderedMergeOperatorTest @@ -81,7 +78,7 @@ public void testOneInput() { FragmentContext context = FragmentContext.defaultContext(); Supplier>> inputs = - () -> Arrays.asList( + () -> Collections.singletonList( new Input(MockOperator.ints(context, 3))); Operator op = new OrderedMergeOperator<>( context, @@ -137,4 +134,39 @@ public void testClose() assertTrue(input2.state == State.CLOSED); op.close(true); } + + /** + * Test the case where the input to the ordered merge is, instead, + * used directly as an operator (because it turns out there is only + * one such input, so a merge is unnecessary). + */ + @Test + public void testInput() + { + FragmentContext context = FragmentContext.defaultContext(); + MockOperator input1 = MockOperator.ints(context, 2); + Input input = new Input<>(input1); + assertFalse(input.eof()); + assertEquals(0, (int) input.get()); + Operator op2 = input.toOperator(context); + List results = Operators.toList(op2); + assertEquals(Arrays.asList(0, 1), results); + assertTrue(input1.state == State.CLOSED); + } + + /** + * As above, but with a no-rows input. + */ + @Test + public void testEmptyInput() + { + FragmentContext context = FragmentContext.defaultContext(); + Operator input1 = new NullOperator<>(context); + Input input = new Input<>(input1); + assertTrue(input.eof()); + assertNull(input.get()); + Operator op2 = input.toOperator(context); + List results = Operators.toList(op2); + assertTrue(results.isEmpty()); + } } From 5bfb2a3a9851026740f508d4ce3403680b1b9f3c Mon Sep 17 00:00:00 2001 From: Paul Rogers Date: Wed, 15 Jun 2022 19:36:09 -0700 Subject: [PATCH 04/11] Revised operator iterator to be simpler --- .../druid/queryng/config/QueryNGConfig.java | 33 ++++++ .../fragment/FragmentBuilderFactoryImpl.java | 3 +- .../queryng/fragment/FragmentBuilderImpl.java | 3 +- .../druid/queryng/fragment/FragmentRun.java | 6 +- .../queryng/fragment/FragmentRunImpl.java | 10 +- .../queryng/operators/ConcatOperator.java | 28 ++--- .../queryng/operators/FilterOperator.java | 16 +-- .../druid/queryng/operators/Iterators.java | 100 ++++++++++++++++++ .../queryng/operators/LimitOperator.java | 16 ++- .../queryng/operators/MappingOperator.java | 12 +-- .../druid/queryng/operators/NullOperator.java | 16 +-- .../druid/queryng/operators/Operator.java | 34 +++++- .../druid/queryng/operators/Operators.java | 63 +++-------- .../operators/OrderedMergeOperator.java | 26 +++-- .../queryng/operators/PushBackOperator.java | 16 +-- .../queryng/operators/SequenceOperator.java | 26 +++-- .../queryng/operators/TransformOperator.java | 2 +- .../queryng/operators/WrappingOperator.java | 6 +- .../queryng/operators/BasicOperatorTest.java | 26 +++-- .../queryng/operators/ConcatOperatorTest.java | 8 +- .../druid/queryng/operators/FragmentTest.java | 17 ++- .../druid/queryng/operators/MockOperator.java | 14 +-- .../queryng/operators/MockOperatorTest.java | 13 +-- .../queryng/operators/MockStorageAdapter.java | 2 +- .../queryng/operators/OperatorTests.java | 39 +++++++ .../operators/OrderedMergeOperatorTest.java | 23 ++-- .../operators/PushBackOperatorTest.java | 13 ++- 27 files changed, 360 insertions(+), 211 deletions(-) create mode 100644 processing/src/main/java/org/apache/druid/queryng/operators/Iterators.java create mode 100644 processing/src/test/java/org/apache/druid/queryng/operators/OperatorTests.java diff --git a/processing/src/main/java/org/apache/druid/queryng/config/QueryNGConfig.java b/processing/src/main/java/org/apache/druid/queryng/config/QueryNGConfig.java index 7aa989d1f7d0..9a5766b5acb0 100644 --- a/processing/src/main/java/org/apache/druid/queryng/config/QueryNGConfig.java +++ b/processing/src/main/java/org/apache/druid/queryng/config/QueryNGConfig.java @@ -20,6 +20,9 @@ package org.apache.druid.queryng.config; import com.fasterxml.jackson.annotation.JsonProperty; +import org.apache.druid.query.Query; +import org.apache.druid.query.QueryPlus; +import org.apache.druid.query.scan.ScanQuery; /** * Configuration for the "NG" query engine. @@ -34,6 +37,8 @@ public class QueryNGConfig @JsonProperty("enabled") private boolean enabled; + public static final String CONTEXT_VAR = "queryng"; + /** * Create an instance for testing. */ @@ -49,4 +54,32 @@ public boolean enabled() { return enabled; } + + /** + * Determine if Query NG should be enabled for the given query; + * that is, if the query should have a fragment context attached. + * At present, Query NG is enabled if the query is a scan query and + * the query has the "queryng" context variable set. The caller + * should already have checked if the Query NG engine is enabled + * globally. If Query NG is enabled for a query, then the caller + * will attach a fragment context to the query's QueryPlus. + */ + public static boolean isEnabled(Query query) + { + // Query has to be of the currently-supported type + if (!(query instanceof ScanQuery)) { + return false; + } + return query.getContextBoolean(CONTEXT_VAR, false); + } + + /** + * Determine if the Query NG (operator-based) engine is enabled for the + * given query (given as a QueryPlus). Query NG is enabled if the QueryPlus + * includes the fragment context needed by the Query NG engine. + */ + public static boolean enabledFor(final QueryPlus queryPlus) + { + return queryPlus.fragmentBuilder() != null; + } } diff --git a/processing/src/main/java/org/apache/druid/queryng/fragment/FragmentBuilderFactoryImpl.java b/processing/src/main/java/org/apache/druid/queryng/fragment/FragmentBuilderFactoryImpl.java index 5368fcae66e4..ee7a33e792d5 100644 --- a/processing/src/main/java/org/apache/druid/queryng/fragment/FragmentBuilderFactoryImpl.java +++ b/processing/src/main/java/org/apache/druid/queryng/fragment/FragmentBuilderFactoryImpl.java @@ -22,7 +22,6 @@ import org.apache.druid.query.Query; import org.apache.druid.query.context.ResponseContext; import org.apache.druid.queryng.config.QueryNGConfig; -import org.apache.druid.queryng.operators.Operators; import javax.inject.Inject; @@ -52,7 +51,7 @@ public FragmentBuilder create( return null; } // Client must explicitly ask for the engine - if (!Operators.isEnabled(query)) { + if (!QueryNGConfig.isEnabled(query)) { return null; } // Only then do we create a fragment builder which, implicitly, diff --git a/processing/src/main/java/org/apache/druid/queryng/fragment/FragmentBuilderImpl.java b/processing/src/main/java/org/apache/druid/queryng/fragment/FragmentBuilderImpl.java index 936c2bbf5a1c..a9357a8e5d50 100644 --- a/processing/src/main/java/org/apache/druid/queryng/fragment/FragmentBuilderImpl.java +++ b/processing/src/main/java/org/apache/druid/queryng/fragment/FragmentBuilderImpl.java @@ -25,6 +25,7 @@ import org.apache.druid.queryng.fragment.FragmentHandleImpl.EmptyFragmentHandle; import org.apache.druid.queryng.fragment.FragmentHandleImpl.FragmentOperatorHandle; import org.apache.druid.queryng.fragment.FragmentHandleImpl.FragmentSequenceHandle; +import org.apache.druid.queryng.operators.Iterators; import org.apache.druid.queryng.operators.Operator; import java.util.Iterator; @@ -100,7 +101,7 @@ public Sequence runAsSequence(Operator rootOp) @Override public Iterator make() { - return run.iterator(); + return Iterators.toIterator(run.iterator()); } @Override diff --git a/processing/src/main/java/org/apache/druid/queryng/fragment/FragmentRun.java b/processing/src/main/java/org/apache/druid/queryng/fragment/FragmentRun.java index e63508ee5253..873dc9aa64c4 100644 --- a/processing/src/main/java/org/apache/druid/queryng/fragment/FragmentRun.java +++ b/processing/src/main/java/org/apache/druid/queryng/fragment/FragmentRun.java @@ -19,6 +19,8 @@ package org.apache.druid.queryng.fragment; +import org.apache.druid.queryng.operators.Operator.RowIterator; + import java.util.List; /** @@ -31,10 +33,12 @@ * the iterator, a sequence, or convert the results to a list. The * fragment is not reentrant: results can be obtained only once. */ -public interface FragmentRun extends AutoCloseable, Iterable +public interface FragmentRun extends AutoCloseable { FragmentContext context(); + RowIterator iterator(); + /** * Materializes the entire result set as a list. Primarily for testing. * Opens the fragment, reads results, and closes the fragment. diff --git a/processing/src/main/java/org/apache/druid/queryng/fragment/FragmentRunImpl.java b/processing/src/main/java/org/apache/druid/queryng/fragment/FragmentRunImpl.java index 1b0973866501..d63ccd9eca65 100644 --- a/processing/src/main/java/org/apache/druid/queryng/fragment/FragmentRunImpl.java +++ b/processing/src/main/java/org/apache/druid/queryng/fragment/FragmentRunImpl.java @@ -20,17 +20,17 @@ package org.apache.druid.queryng.fragment; import com.google.common.base.Preconditions; -import com.google.common.collect.Lists; import org.apache.druid.queryng.fragment.FragmentContext.State; +import org.apache.druid.queryng.operators.Iterators; import org.apache.druid.queryng.operators.Operator; +import org.apache.druid.queryng.operators.Operator.RowIterator; -import java.util.Iterator; import java.util.List; public class FragmentRunImpl implements FragmentRun { private final FragmentContextImpl context; - private Iterator rootIter; + private RowIterator rootIter; public FragmentRunImpl(FragmentContextImpl context, Operator root) { @@ -48,7 +48,7 @@ public FragmentRunImpl(FragmentContextImpl context, Operator root) } @Override - public Iterator iterator() + public RowIterator iterator() { Preconditions.checkState(context.state == State.RUN); return rootIter; @@ -64,7 +64,7 @@ public FragmentContext context() public List toList() { try { - return Lists.newArrayList(this); + return Iterators.toList(rootIter); } finally { close(); diff --git a/processing/src/main/java/org/apache/druid/queryng/operators/ConcatOperator.java b/processing/src/main/java/org/apache/druid/queryng/operators/ConcatOperator.java index 658836fd65b5..de413d3b2e4f 100644 --- a/processing/src/main/java/org/apache/druid/queryng/operators/ConcatOperator.java +++ b/processing/src/main/java/org/apache/druid/queryng/operators/ConcatOperator.java @@ -19,7 +19,6 @@ package org.apache.druid.queryng.operators; -import com.google.common.base.Preconditions; import org.apache.druid.queryng.fragment.FragmentContext; import org.apache.druid.queryng.operators.Operator.IterableOperator; @@ -47,7 +46,7 @@ public static Operator concatOrNot( private final Iterator> childIter; private Operator current; - private Iterator currentIter; + private RowIterator currentIter; public ConcatOperator(FragmentContext context, List> children) { @@ -56,38 +55,33 @@ public ConcatOperator(FragmentContext context, List> children) } @Override - public Iterator open() + public RowIterator open() { return this; } @Override - public boolean hasNext() + public T next() throws EofException { while (true) { if (current != null) { - if (currentIter.hasNext()) { - return true; + try { + return currentIter.next(); + } + catch (EofException e) { + current.close(true); + current = null; + currentIter = null; } - current.close(true); - current = null; - currentIter = null; } if (!childIter.hasNext()) { - return false; + throw Operators.eof(); } current = childIter.next(); currentIter = current.open(); } } - @Override - public T next() - { - Preconditions.checkState(currentIter != null, "Missing call to hasNext()?"); - return currentIter.next(); - } - @Override public void close(boolean cascade) { diff --git a/processing/src/main/java/org/apache/druid/queryng/operators/FilterOperator.java b/processing/src/main/java/org/apache/druid/queryng/operators/FilterOperator.java index 3d31862c2be7..f2f8cf9be705 100644 --- a/processing/src/main/java/org/apache/druid/queryng/operators/FilterOperator.java +++ b/processing/src/main/java/org/apache/druid/queryng/operators/FilterOperator.java @@ -30,7 +30,6 @@ public class FilterOperator extends MappingOperator { private final Predicate predicate; - private T nextValue; public FilterOperator(FragmentContext context, Operator input, Predicate pred) { @@ -39,20 +38,13 @@ public FilterOperator(FragmentContext context, Operator input, Predicate p } @Override - public boolean hasNext() + public T next() throws EofException { - while (super.hasNext()) { - nextValue = inputIter.next(); + while (true) { + T nextValue = inputIter.next(); if (predicate.test(nextValue)) { - return true; + return nextValue; } } - return false; - } - - @Override - public T next() - { - return nextValue; } } diff --git a/processing/src/main/java/org/apache/druid/queryng/operators/Iterators.java b/processing/src/main/java/org/apache/druid/queryng/operators/Iterators.java new file mode 100644 index 000000000000..2773c3dcdaf8 --- /dev/null +++ b/processing/src/main/java/org/apache/druid/queryng/operators/Iterators.java @@ -0,0 +1,100 @@ +/* + * 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.queryng.operators; + +import com.google.common.collect.Lists; +import org.apache.druid.queryng.operators.Operator.EofException; +import org.apache.druid.queryng.operators.Operator.RowIterator; + +import java.util.Iterator; +import java.util.List; +import java.util.NoSuchElementException; + +/** + * Utility methods on top of {@link Operator.RowIterator RowIterator}, + * including conversion to a Java iterator (primarily for testing.) + */ +public class Iterators +{ + public static class ShimIterator implements Iterator + { + private final RowIterator operIter; + private boolean eof; + private T lookAhead; + + public ShimIterator(RowIterator operIter) + { + this.operIter = operIter; + } + + @Override + public boolean hasNext() + { + if (eof) { + return false; + } + try { + lookAhead = operIter.next(); + return true; + } + catch (EofException e) { + eof = false; + return false; + } + } + + @Override + public T next() + { + if (eof || lookAhead == null) { + throw new NoSuchElementException(); + } + return lookAhead; + } + + } + + public static Iterable toIterable(RowIterator iter) + { + return Iterators.toIterable(Iterators.toIterator(iter)); + } + + public static Iterable toIterable(Iterator iter) + { + return new Iterable() { + @Override + public Iterator iterator() + { + return iter; + } + }; + } + + public static Iterator toIterator(RowIterator opIter) + { + return new Iterators.ShimIterator(opIter); + } + + public static List toList(RowIterator operIter) + { + return Lists.newArrayList(new Iterators.ShimIterator(operIter)); + } + +} diff --git a/processing/src/main/java/org/apache/druid/queryng/operators/LimitOperator.java b/processing/src/main/java/org/apache/druid/queryng/operators/LimitOperator.java index 0235774d18fb..f9e2f7833ead 100644 --- a/processing/src/main/java/org/apache/druid/queryng/operators/LimitOperator.java +++ b/processing/src/main/java/org/apache/druid/queryng/operators/LimitOperator.java @@ -27,8 +27,6 @@ */ public class LimitOperator extends MappingOperator { - public static final long UNLIMITED = Long.MAX_VALUE; - protected final long limit; protected long rowCount; @@ -40,15 +38,13 @@ public LimitOperator(FragmentContext context, Operator input, long limit) } @Override - public boolean hasNext() - { - return rowCount < limit && super.hasNext(); - } - - @Override - public T next() + public T next() throws EofException { + if (rowCount >= limit) { + throw Operators.eof(); + } + T item = inputIter.next(); rowCount++; - return inputIter.next(); + return item; } } diff --git a/processing/src/main/java/org/apache/druid/queryng/operators/MappingOperator.java b/processing/src/main/java/org/apache/druid/queryng/operators/MappingOperator.java index c2fa7e286cf2..dd978094742b 100644 --- a/processing/src/main/java/org/apache/druid/queryng/operators/MappingOperator.java +++ b/processing/src/main/java/org/apache/druid/queryng/operators/MappingOperator.java @@ -23,8 +23,6 @@ import org.apache.druid.queryng.fragment.FragmentContext; import org.apache.druid.queryng.operators.Operator.IterableOperator; -import java.util.Iterator; - /** * Base class for operators that do a simple mapping of their input * to their output. Handles the busy-work of managing the (single) @@ -33,7 +31,7 @@ public abstract class MappingOperator implements IterableOperator { private final Operator input; - protected Iterator inputIter; + protected RowIterator inputIter; protected State state = State.START; public MappingOperator(FragmentContext context, Operator input) @@ -43,7 +41,7 @@ public MappingOperator(FragmentContext context, Operator input) } @Override - public Iterator open() + public RowIterator open() { Preconditions.checkState(state == State.START); inputIter = input.open(); @@ -60,10 +58,4 @@ public void close(boolean cascade) inputIter = null; state = State.CLOSED; } - - @Override - public boolean hasNext() - { - return state == State.RUN && inputIter.hasNext(); - } } diff --git a/processing/src/main/java/org/apache/druid/queryng/operators/NullOperator.java b/processing/src/main/java/org/apache/druid/queryng/operators/NullOperator.java index 8165c999dc3b..b5b103f3891b 100644 --- a/processing/src/main/java/org/apache/druid/queryng/operators/NullOperator.java +++ b/processing/src/main/java/org/apache/druid/queryng/operators/NullOperator.java @@ -21,9 +21,7 @@ import com.google.common.base.Preconditions; import org.apache.druid.queryng.fragment.FragmentContext; - -import java.util.Collections; -import java.util.Iterator; +import org.apache.druid.queryng.operators.Operator.IterableOperator; /** * World's simplest operator: does absolutely nothing @@ -31,7 +29,7 @@ * tests when we want an empty input, and for a fragment that * somehow ended up with no operators. */ -public class NullOperator implements Operator +public class NullOperator implements IterableOperator { public State state = State.START; @@ -41,11 +39,17 @@ public NullOperator(FragmentContext context) } @Override - public Iterator open() + public RowIterator open() { Preconditions.checkState(state == State.START); state = State.RUN; - return Collections.emptyIterator(); + return this; + } + + @Override + public T next() throws EofException + { + throw Operators.eof(); } @Override diff --git a/processing/src/main/java/org/apache/druid/queryng/operators/Operator.java b/processing/src/main/java/org/apache/druid/queryng/operators/Operator.java index fda745d8fb65..d58d88465d08 100644 --- a/processing/src/main/java/org/apache/druid/queryng/operators/Operator.java +++ b/processing/src/main/java/org/apache/druid/queryng/operators/Operator.java @@ -19,8 +19,6 @@ package org.apache.druid.queryng.operators; -import java.util.Iterator; - /** * An operator is a data pipeline transform: something that operates on * a stream of results in some way. An operator has a very simple lifecycle: @@ -32,6 +30,15 @@ *

  • Closed.
    • * *

    + * Opening an operator returns a {@link Operator.RowIterator RowIterator} + * which returns rows. The Java {@code Iterator} class has extra overhead + * which we want to avoid on the per-row inner loop code path. A + * {@code RowIterator} has one method: {@link Operator.RowIterator#next() next()}, + * which either returns a row (however the operator defines it), or throws an + * {@link Operator.EofException EofException} when there are no more rows. + * Downstream operators need not do any conditional checking: they can just + * propagate the exception if they have nothing to add at EOF. + *

    * Leaf operators produce results, internal operators transform them, and root * operators deliver results to some consumer. Operators know nothing about their * children other than that they follow the operator protocol and will produce a @@ -128,10 +135,29 @@ */ public interface Operator { + /** + * Exception thrown at EOF. + */ + class EofException extends Exception + { + } + + /** + * Iterator over operator results. Operators do not use the Java + * {@code Iterator} class: the simpler implementation here + * minimizes per-row overhead. An {@code OperatorIterator} can + * be converted to a Java {@code Iterator} by calling + * {@link Operators#toIterator()}#, but that adds overhead. + */ + interface RowIterator + { + T next() throws EofException; + } + /** * Convenience interface for an operator which is its own iterator. */ - interface IterableOperator extends Operator, Iterator + interface IterableOperator extends Operator, RowIterator { } @@ -151,7 +177,7 @@ enum State * in the {@code open()} call for simple operators,or later, on demand, for more * complex operators such as in a merge or union. */ - Iterator open(); + RowIterator open(); /** * Called at two distinct times. An operator may choose to close a child diff --git a/processing/src/main/java/org/apache/druid/queryng/operators/Operators.java b/processing/src/main/java/org/apache/druid/queryng/operators/Operators.java index 5f513c69e873..9152cf9d8cd9 100644 --- a/processing/src/main/java/org/apache/druid/queryng/operators/Operators.java +++ b/processing/src/main/java/org/apache/druid/queryng/operators/Operators.java @@ -22,10 +22,8 @@ import com.google.common.collect.Lists; import org.apache.druid.java.util.common.guava.BaseSequence; import org.apache.druid.java.util.common.guava.Sequence; -import org.apache.druid.query.Query; -import org.apache.druid.query.QueryPlus; -import org.apache.druid.query.scan.ScanQuery; import org.apache.druid.queryng.fragment.DAGBuilder; +import org.apache.druid.queryng.operators.Operator.EofException; import java.util.Iterator; import java.util.List; @@ -35,36 +33,6 @@ */ public class Operators { - public static final String CONTEXT_VAR = "queryng"; - - /** - * Determine if the Query NG (operator-based) engine is enabled for the - * given query (given as a QueryPlus). Query NG is enabled if the QueryPlus - * includes the fragment context needed by the Query NG engine. - */ - public static boolean enabledFor(final QueryPlus queryPlus) - { - return queryPlus.fragmentBuilder() != null; - } - - /** - * Determine if Query NG should be enabled for the given query; - * that is, if the query should have a fragment context attached. - * At present, Query NG is enabled if the query is a scan query and - * the query has the "queryng" context variable set. The caller - * should already have checked if the Query NG engine is enabled - * globally. If Query NG is enabled for a query, then the caller - * will attach a fragment context to the query's QueryPlus. - */ - public static boolean isEnabled(Query query) - { - // Query has to be of the currently-supported type - if (!(query instanceof ScanQuery)) { - return false; - } - return query.getContextBoolean(CONTEXT_VAR, false); - } - /** * Convenience function to open the operator and return its * iterator as an {@code Iterable}. @@ -75,18 +43,7 @@ public static Iterable toIterable(Operator op) @Override public Iterator iterator() { - return op.open(); - } - }; - } - - public static Iterable toIterable(Iterator iter) - { - return new Iterable() { - @Override - public Iterator iterator() - { - return iter; + return new Iterators.ShimIterator(op.open()); } }; } @@ -96,7 +53,7 @@ public Iterator iterator() * iterator mechanism (since an operator looks like an iterator.) * * This is a named class so we can unwrap the operator in - * {@link #runToProducer()} below. + * {@link #toOperator()} below. */ public static class OperatorWrapperSequence extends BaseSequence> { @@ -109,7 +66,7 @@ public OperatorWrapperSequence(Operator op) @Override public Iterator make() { - return (Iterator) op.open(); + return toIterator(op); } @Override @@ -168,14 +125,24 @@ public static Operator unwrapOperator(Sequence sequence) return null; } + public static Iterator toIterator(Operator op) + { + return new Iterators.ShimIterator(op.open()); + } + /** * This will materialize the entire sequence from the wrapped * operator. Use at your own risk. */ public static List toList(Operator op) { - List results = Lists.newArrayList(Operators.toIterable(op)); + List results = Lists.newArrayList(toIterator(op)); op.close(true); return results; } + + public static EofException eof() + { + return new EofException(); + } } diff --git a/processing/src/main/java/org/apache/druid/queryng/operators/OrderedMergeOperator.java b/processing/src/main/java/org/apache/druid/queryng/operators/OrderedMergeOperator.java index 1e19ed171e8f..5912337ceb88 100644 --- a/processing/src/main/java/org/apache/druid/queryng/operators/OrderedMergeOperator.java +++ b/processing/src/main/java/org/apache/druid/queryng/operators/OrderedMergeOperator.java @@ -23,7 +23,6 @@ import org.apache.druid.queryng.fragment.FragmentContext; import org.apache.druid.queryng.operators.Operator.IterableOperator; -import java.util.Iterator; import java.util.PriorityQueue; import java.util.function.Supplier; @@ -46,16 +45,17 @@ public class OrderedMergeOperator implements IterableOperator public static class Input { private final Operator input; - private final Iterator iter; + private final RowIterator iter; private T currentValue; public Input(Operator input) { this.input = input; this.iter = input.open(); - if (iter.hasNext()) { + try { currentValue = iter.next(); - } else { + } + catch (EofException e) { currentValue = null; input.close(true); } @@ -77,10 +77,11 @@ public T get() public boolean next() { - if (iter.hasNext()) { + try { currentValue = iter.next(); return true; - } else { + } + catch (EofException e) { currentValue = null; input.close(true); return false; @@ -129,7 +130,7 @@ public OrderedMergeOperator( } @Override - public Iterator open() + public RowIterator open() { for (Input input : inputs.get()) { if (!input.eof()) { @@ -140,14 +141,11 @@ public Iterator open() } @Override - public boolean hasNext() - { - return !pQueue.isEmpty(); - } - - @Override - public T next() + public T next() throws EofException { + if (pQueue.isEmpty()) { + throw Operators.eof(); + } Input input = pQueue.remove(); T result = input.get(); if (input.next()) { diff --git a/processing/src/main/java/org/apache/druid/queryng/operators/PushBackOperator.java b/processing/src/main/java/org/apache/druid/queryng/operators/PushBackOperator.java index 7b0993d9e437..0b7d9294583e 100644 --- a/processing/src/main/java/org/apache/druid/queryng/operators/PushBackOperator.java +++ b/processing/src/main/java/org/apache/druid/queryng/operators/PushBackOperator.java @@ -23,8 +23,6 @@ import org.apache.druid.queryng.fragment.FragmentContext; import org.apache.druid.queryng.operators.Operator.IterableOperator; -import java.util.Iterator; - /** * Operator which allows pushing a row back onto the input. The "pushed" * row can occur at construction time, or during execution. @@ -32,13 +30,13 @@ public class PushBackOperator implements IterableOperator { private final Operator input; - private Iterator inputIter; + private RowIterator inputIter; private T pushed; public PushBackOperator( FragmentContext context, Operator input, - Iterator inputIter, + RowIterator inputIter, T pushed) { this.input = input; @@ -53,7 +51,7 @@ public PushBackOperator(FragmentContext context, Operator input) } @Override - public Iterator open() + public RowIterator open() { if (inputIter == null) { inputIter = input.open(); @@ -62,13 +60,7 @@ public Iterator open() } @Override - public boolean hasNext() - { - return pushed != null || inputIter != null && inputIter.hasNext(); - } - - @Override - public T next() + public T next() throws EofException { if (pushed != null) { T ret = pushed; diff --git a/processing/src/main/java/org/apache/druid/queryng/operators/SequenceOperator.java b/processing/src/main/java/org/apache/druid/queryng/operators/SequenceOperator.java index 377d57747b71..cb2c24361063 100644 --- a/processing/src/main/java/org/apache/druid/queryng/operators/SequenceOperator.java +++ b/processing/src/main/java/org/apache/druid/queryng/operators/SequenceOperator.java @@ -27,7 +27,6 @@ import org.apache.druid.queryng.operators.Operator.IterableOperator; import java.io.IOException; -import java.util.Iterator; /** * The SequenceOperator wraps a {@link Sequence} in the @@ -53,7 +52,7 @@ public SequenceOperator(FragmentContext context, Sequence sequence) } @Override - public Iterator open() + public RowIterator open() { Preconditions.checkState(yielder == null); yielder = sequence.toYielder( @@ -72,14 +71,15 @@ public T accumulate(T accumulated, T in) } @Override - public boolean hasNext() - { - return yielder != null && !yielder.isDone(); - } - - @Override - public T next() + public T next() throws EofException { + if (yielder == null && yielder.isDone()) { + throw Operators.eof(); + } + if (yielder.isDone()) { + closeYielder(); + throw Operators.eof(); + } Preconditions.checkState(yielder != null); T value = yielder.get(); yielder = yielder.next(null); @@ -89,9 +89,13 @@ public T next() @Override public void close(boolean cascade) { - if (yielder == null) { - return; + if (yielder != null) { + closeYielder(); } + } + + private void closeYielder() + { try { yielder.close(); } diff --git a/processing/src/main/java/org/apache/druid/queryng/operators/TransformOperator.java b/processing/src/main/java/org/apache/druid/queryng/operators/TransformOperator.java index 90414f7d6db0..e2f510ce4558 100644 --- a/processing/src/main/java/org/apache/druid/queryng/operators/TransformOperator.java +++ b/processing/src/main/java/org/apache/druid/queryng/operators/TransformOperator.java @@ -55,7 +55,7 @@ public TransformOperator( } @Override - public OUT next() + public OUT next() throws EofException { return transformFn.apply(inputIter.next()); } diff --git a/processing/src/main/java/org/apache/druid/queryng/operators/WrappingOperator.java b/processing/src/main/java/org/apache/druid/queryng/operators/WrappingOperator.java index f9b497375616..5cdf13e43caf 100644 --- a/processing/src/main/java/org/apache/druid/queryng/operators/WrappingOperator.java +++ b/processing/src/main/java/org/apache/druid/queryng/operators/WrappingOperator.java @@ -22,8 +22,6 @@ import com.google.common.base.Preconditions; import org.apache.druid.queryng.fragment.FragmentContext; -import java.util.Iterator; - /** * Operator which "wraps" another operator where the only behavior of * interest is at the start or end of the run. The iterator from the @@ -43,10 +41,10 @@ public WrappingOperator(FragmentContext context, Operator input) } @Override - public Iterator open() + public RowIterator open() { Preconditions.checkState(state == State.START); - Iterator inputIter = input.open(); + RowIterator inputIter = input.open(); state = State.RUN; onOpen(); return inputIter; diff --git a/processing/src/test/java/org/apache/druid/queryng/operators/BasicOperatorTest.java b/processing/src/test/java/org/apache/druid/queryng/operators/BasicOperatorTest.java index 9d4e09face93..06768991640f 100644 --- a/processing/src/test/java/org/apache/druid/queryng/operators/BasicOperatorTest.java +++ b/processing/src/test/java/org/apache/druid/queryng/operators/BasicOperatorTest.java @@ -23,19 +23,32 @@ import org.apache.druid.queryng.fragment.FragmentBuilder; import org.apache.druid.queryng.fragment.FragmentContext; import org.apache.druid.queryng.fragment.FragmentRun; +import org.apache.druid.queryng.operators.Operator.EofException; +import org.apache.druid.queryng.operators.Operator.RowIterator; import org.junit.Test; import java.util.Arrays; -import java.util.Iterator; import java.util.List; import static org.junit.Assert.assertEquals; -import static org.junit.Assert.assertFalse; import static org.junit.Assert.assertNull; import static org.junit.Assert.assertTrue; public class BasicOperatorTest { + @Test + public void testToIterable() + { + FragmentBuilder builder = FragmentBuilder.defaultBuilder(); + MockOperator op = MockOperator.ints(builder.context(), 2); + int count = 0; + for (Integer row : Operators.toIterable(op)) { + assertEquals(count++, (int) row); + } + assertEquals(2, count); + op.close(false); + } + @Test public void testFilter() { @@ -145,21 +158,20 @@ public void testWrappingOperator() /** * Getting weird: an operator that wraps a sequence that wraps an operator. + * @throws EofException */ @Test - public void testSequenceOperator() + public void testSequenceOperator() throws EofException { FragmentBuilder builder = FragmentBuilder.defaultBuilder(); MockOperator op = MockOperator.strings(builder.context(), 2); Sequence seq = Operators.toSequence(op); Operator outer = Operators.toOperator(builder, seq); FragmentRun run = builder.run(outer); - Iterator iter = run.iterator(); - assertTrue(iter.hasNext()); + RowIterator iter = run.iterator(); assertEquals("Mock row 0", iter.next()); - assertTrue(iter.hasNext()); assertEquals("Mock row 1", iter.next()); - assertFalse(iter.hasNext()); + OperatorTests.assertEof(iter); run.close(); assertEquals(Operator.State.CLOSED, op.state); } diff --git a/processing/src/test/java/org/apache/druid/queryng/operators/ConcatOperatorTest.java b/processing/src/test/java/org/apache/druid/queryng/operators/ConcatOperatorTest.java index a91d56b05721..b823b1703c05 100644 --- a/processing/src/test/java/org/apache/druid/queryng/operators/ConcatOperatorTest.java +++ b/processing/src/test/java/org/apache/druid/queryng/operators/ConcatOperatorTest.java @@ -20,12 +20,13 @@ package org.apache.druid.queryng.operators; import org.apache.druid.queryng.fragment.FragmentContext; +import org.apache.druid.queryng.operators.Operator.EofException; +import org.apache.druid.queryng.operators.Operator.RowIterator; import org.apache.druid.queryng.operators.Operator.State; import org.junit.Test; import java.util.Arrays; import java.util.Collections; -import java.util.Iterator; import java.util.List; import static org.junit.Assert.assertEquals; @@ -133,7 +134,7 @@ public void testTwoInputs() } @Test - public void testClose() + public void testClose() throws EofException { FragmentContext context = FragmentContext.defaultContext(); MockOperator input1 = MockOperator.ints(context, 2); @@ -141,8 +142,7 @@ public void testClose() Operator op = ConcatOperator.concatOrNot( context, Arrays.asList(input1, input2)); - Iterator iter = op.open(); - assertTrue(iter.hasNext()); + RowIterator iter = op.open(); assertEquals(0, (int) iter.next()); // Only first input has been opened. diff --git a/processing/src/test/java/org/apache/druid/queryng/operators/FragmentTest.java b/processing/src/test/java/org/apache/druid/queryng/operators/FragmentTest.java index 3ed0c002dd6b..aef9e29dcfdd 100644 --- a/processing/src/test/java/org/apache/druid/queryng/operators/FragmentTest.java +++ b/processing/src/test/java/org/apache/druid/queryng/operators/FragmentTest.java @@ -36,6 +36,8 @@ import org.apache.druid.queryng.fragment.FragmentHandle; import org.apache.druid.queryng.fragment.FragmentRun; import org.apache.druid.queryng.fragment.NullFragmentBuilderFactory; +import org.apache.druid.queryng.operators.Operator.EofException; +import org.apache.druid.queryng.operators.Operator.RowIterator; import org.junit.Test; import java.io.IOException; @@ -180,7 +182,7 @@ public void testRun() try (FragmentRun run = handle.run()) { assertEquals(FragmentContext.State.RUN, context.state()); assertSame(context, run.context()); - for (Integer value : run) { + for (Integer value : Iterators.toIterable(run.iterator())) { assertEquals(i++, (int) value); } } @@ -315,7 +317,7 @@ public void testFactory() // Enable at query level. Use of operators gated by config. query = query.withOverriddenContext( - ImmutableMap.of(Operators.CONTEXT_VAR, true)); + ImmutableMap.of(QueryNGConfig.CONTEXT_VAR, true)); assertNotNull(enableFactory.create(query, ResponseContext.createEmpty())); assertNull(disableFactory.create(query, ResponseContext.createEmpty())); assertNull(nullFactory.create(query, ResponseContext.createEmpty())); @@ -329,12 +331,15 @@ public void testQueryPlus() .dataSource("foo") .eternityInterval() .build(); - QueryPlus queryPlus = QueryPlus.wrap(query).withFragmentBuilder(builder); + QueryPlus queryPlus = QueryPlus.wrap(query); + assertFalse(QueryNGConfig.enabledFor(queryPlus)); + queryPlus = queryPlus.withFragmentBuilder(builder); + assertTrue(QueryNGConfig.enabledFor(queryPlus)); assertSame(builder, queryPlus.fragmentBuilder()); } @Test - public void testFragmentContext() + public void testFragmentContext() throws EofException { FragmentBuilder builder = FragmentBuilder.defaultBuilder(); FragmentContext context = builder.context(); @@ -345,8 +350,10 @@ public void testFragmentContext() assertNull(context.exception()); MockOperator op = MockOperator.ints(builder.context(), 4); FragmentHandle handle = builder.handle(op); - handle.run().iterator(); + RowIterator iter = handle.run().iterator(); assertEquals(FragmentContext.State.RUN, context.state()); + // Read from the iterator, just to keep Java 11 happy. + assertNotNull(iter.next()); ISE ex = new ISE("oops"); ((FragmentContextImpl) context).failed(ex); assertEquals(FragmentContext.State.FAILED, context.state()); diff --git a/processing/src/test/java/org/apache/druid/queryng/operators/MockOperator.java b/processing/src/test/java/org/apache/druid/queryng/operators/MockOperator.java index f1bd21c26319..653dde5e1002 100644 --- a/processing/src/test/java/org/apache/druid/queryng/operators/MockOperator.java +++ b/processing/src/test/java/org/apache/druid/queryng/operators/MockOperator.java @@ -23,7 +23,6 @@ import org.apache.druid.queryng.fragment.FragmentContext; import org.apache.druid.queryng.operators.Operator.IterableOperator; -import java.util.Iterator; import java.util.function.Function; public class MockOperator implements IterableOperator @@ -54,7 +53,7 @@ public static MockOperator strings(FragmentContext context, int rowCount } @Override - public Iterator open() + public RowIterator open() { Preconditions.checkState(state == State.START); state = State.RUN; @@ -62,14 +61,11 @@ public Iterator open() } @Override - public boolean hasNext() - { - return rowPosn < targetCount; - } - - @Override - public T next() + public T next() throws EofException { + if (rowPosn >= targetCount) { + throw Operators.eof(); + } return generator.apply(rowPosn++); } diff --git a/processing/src/test/java/org/apache/druid/queryng/operators/MockOperatorTest.java b/processing/src/test/java/org/apache/druid/queryng/operators/MockOperatorTest.java index 12876e6b1d1e..04dbcbb07d2d 100644 --- a/processing/src/test/java/org/apache/druid/queryng/operators/MockOperatorTest.java +++ b/processing/src/test/java/org/apache/druid/queryng/operators/MockOperatorTest.java @@ -22,14 +22,13 @@ import org.apache.druid.queryng.fragment.FragmentBuilder; import org.apache.druid.queryng.fragment.FragmentHandle; import org.apache.druid.queryng.fragment.FragmentRun; +import org.apache.druid.queryng.operators.Operator.EofException; +import org.apache.druid.queryng.operators.Operator.RowIterator; import org.junit.Test; -import java.util.Iterator; import java.util.List; import static org.junit.Assert.assertEquals; -import static org.junit.Assert.assertFalse; -import static org.junit.Assert.assertTrue; /** * Use a mock operator to test (and illustrate) the basic operator @@ -38,18 +37,16 @@ public class MockOperatorTest { @Test - public void testMockStringOperator() + public void testMockStringOperator() throws EofException { FragmentBuilder builder = FragmentBuilder.defaultBuilder(); MockOperator op = MockOperator.strings(builder.context(), 2); FragmentHandle handle = builder.handle(op); try (FragmentRun run = handle.run()) { - Iterator iter = run.iterator(); - assertTrue(iter.hasNext()); + RowIterator iter = run.iterator(); assertEquals("Mock row 0", iter.next()); - assertTrue(iter.hasNext()); assertEquals("Mock row 1", iter.next()); - assertFalse(iter.hasNext()); + OperatorTests.assertEof(iter); } assertEquals(Operator.State.CLOSED, op.state); } diff --git a/processing/src/test/java/org/apache/druid/queryng/operators/MockStorageAdapter.java b/processing/src/test/java/org/apache/druid/queryng/operators/MockStorageAdapter.java index be9d6a0a9d88..c1bdd7ea7b5b 100644 --- a/processing/src/test/java/org/apache/druid/queryng/operators/MockStorageAdapter.java +++ b/processing/src/test/java/org/apache/druid/queryng/operators/MockStorageAdapter.java @@ -68,7 +68,7 @@ public Indexed getAvailableDimensions() @Override public Iterable getAvailableMetrics() { - return Arrays.asList("delta"); + return Collections.singletonList("delta"); } @Override diff --git a/processing/src/test/java/org/apache/druid/queryng/operators/OperatorTests.java b/processing/src/test/java/org/apache/druid/queryng/operators/OperatorTests.java new file mode 100644 index 000000000000..b2e92130e47e --- /dev/null +++ b/processing/src/test/java/org/apache/druid/queryng/operators/OperatorTests.java @@ -0,0 +1,39 @@ +/* + * 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.queryng.operators; + +import org.apache.druid.queryng.operators.Operator.EofException; +import org.apache.druid.queryng.operators.Operator.RowIterator; + +import static org.junit.Assert.fail; + +public class OperatorTests +{ + public static void assertEof(RowIterator operIter) + { + try { + operIter.next(); + fail(); + } + catch (EofException e) { + // Expected + } + } +} diff --git a/processing/src/test/java/org/apache/druid/queryng/operators/OrderedMergeOperatorTest.java b/processing/src/test/java/org/apache/druid/queryng/operators/OrderedMergeOperatorTest.java index be90beb1f0a7..d8a3142e718f 100644 --- a/processing/src/test/java/org/apache/druid/queryng/operators/OrderedMergeOperatorTest.java +++ b/processing/src/test/java/org/apache/druid/queryng/operators/OrderedMergeOperatorTest.java @@ -19,16 +19,15 @@ package org.apache.druid.queryng.operators; -import com.google.common.collect.Lists; import com.google.common.collect.Ordering; import org.apache.druid.queryng.fragment.FragmentContext; +import org.apache.druid.queryng.operators.Operator.RowIterator; import org.apache.druid.queryng.operators.Operator.State; import org.apache.druid.queryng.operators.OrderedMergeOperator.Input; import org.junit.Test; import java.util.Arrays; import java.util.Collections; -import java.util.Iterator; import java.util.List; import java.util.function.Supplier; @@ -50,8 +49,8 @@ public void testNoInputs() Ordering.natural(), 0, inputs); - Iterator iter = op.open(); - assertFalse(iter.hasNext()); + RowIterator iter = op.open(); + OperatorTests.assertEof(iter); op.close(true); } @@ -68,8 +67,8 @@ public void testEmptyInputs() Ordering.natural(), 2, inputs); - Iterator iter = op.open(); - assertFalse(iter.hasNext()); + RowIterator iter = op.open(); + OperatorTests.assertEof(iter); op.close(true); } @@ -85,8 +84,8 @@ public void testOneInput() Ordering.natural(), 1, inputs); - Iterator iter = op.open(); - List results = Lists.newArrayList(iter); + RowIterator iter = op.open(); + List results = Iterators.toList(iter); op.close(true); assertEquals(Arrays.asList(0, 1, 2), results); } @@ -104,8 +103,8 @@ public void testTwoInputs() Ordering.natural(), 2, inputs); - Iterator iter = op.open(); - List results = Lists.newArrayList(iter); + RowIterator iter = op.open(); + List results = Iterators.toList(iter); op.close(true); assertEquals(Arrays.asList(0, 0, 1, 1, 2, 2, 3, 4), results); } @@ -125,8 +124,8 @@ public void testClose() Ordering.natural(), 2, inputs); - Iterator iter = op.open(); - List results = Lists.newArrayList(iter); + RowIterator iter = op.open(); + List results = Iterators.toList(iter); assertEquals(Arrays.asList(0, 0, 1, 1), results); // Inputs are closed as exhausted. diff --git a/processing/src/test/java/org/apache/druid/queryng/operators/PushBackOperatorTest.java b/processing/src/test/java/org/apache/druid/queryng/operators/PushBackOperatorTest.java index 911f5e136cd1..2cc2a67b54b5 100644 --- a/processing/src/test/java/org/apache/druid/queryng/operators/PushBackOperatorTest.java +++ b/processing/src/test/java/org/apache/druid/queryng/operators/PushBackOperatorTest.java @@ -20,10 +20,11 @@ package org.apache.druid.queryng.operators; import org.apache.druid.queryng.fragment.FragmentContext; +import org.apache.druid.queryng.operators.Operator.EofException; +import org.apache.druid.queryng.operators.Operator.RowIterator; import org.junit.Test; import java.util.Arrays; -import java.util.Iterator; import java.util.List; import static org.junit.Assert.assertEquals; @@ -52,13 +53,12 @@ public void testSimpleInput() } @Test - public void testPush() + public void testPush() throws EofException { FragmentContext context = FragmentContext.defaultContext(); Operator input = MockOperator.ints(context, 2); PushBackOperator op = new PushBackOperator(context, input); - Iterator iter = op.open(); - assertTrue(iter.hasNext()); + RowIterator iter = op.open(); Integer item = iter.next(); op.push(item); List results = Operators.toList(op); @@ -67,12 +67,11 @@ public void testPush() } @Test - public void testInitialPush() + public void testInitialPush() throws EofException { FragmentContext context = FragmentContext.defaultContext(); Operator input = MockOperator.ints(context, 2); - Iterator iter = input.open(); - assertTrue(iter.hasNext()); + RowIterator iter = input.open(); Integer item = iter.next(); PushBackOperator op = new PushBackOperator(context, input, iter, item); List results = Operators.toList(op); From 63e3ed201b9cdd6c8377f8729f6a398e6134262a Mon Sep 17 00:00:00 2001 From: Paul Rogers Date: Thu, 16 Jun 2022 13:25:16 -0700 Subject: [PATCH 05/11] Build fixes --- .../java/org/apache/druid/queryng/config/QueryNGConfig.java | 1 + .../java/org/apache/druid/queryng/operators/Operator.java | 5 ++++- .../java/org/apache/druid/queryng/operators/Operators.java | 2 +- .../org/apache/druid/queryng/operators/SequenceOperator.java | 2 +- 4 files changed, 7 insertions(+), 3 deletions(-) diff --git a/processing/src/main/java/org/apache/druid/queryng/config/QueryNGConfig.java b/processing/src/main/java/org/apache/druid/queryng/config/QueryNGConfig.java index 9a5766b5acb0..5ca3c9ff119c 100644 --- a/processing/src/main/java/org/apache/druid/queryng/config/QueryNGConfig.java +++ b/processing/src/main/java/org/apache/druid/queryng/config/QueryNGConfig.java @@ -29,6 +29,7 @@ */ public class QueryNGConfig { + @SuppressWarnings("unused") // To be used later public static final String CONFIG_ROOT = "druid.queryng"; /** diff --git a/processing/src/main/java/org/apache/druid/queryng/operators/Operator.java b/processing/src/main/java/org/apache/druid/queryng/operators/Operator.java index d58d88465d08..b9d0fda512aa 100644 --- a/processing/src/main/java/org/apache/druid/queryng/operators/Operator.java +++ b/processing/src/main/java/org/apache/druid/queryng/operators/Operator.java @@ -19,6 +19,8 @@ package org.apache.druid.queryng.operators; +import java.util.Iterator; + /** * An operator is a data pipeline transform: something that operates on * a stream of results in some way. An operator has a very simple lifecycle: @@ -147,7 +149,8 @@ class EofException extends Exception * {@code Iterator} class: the simpler implementation here * minimizes per-row overhead. An {@code OperatorIterator} can * be converted to a Java {@code Iterator} by calling - * {@link Operators#toIterator()}#, but that adds overhead. + * {@link static Iterator Operators#toIterator(Operator)}, + * but that approach adds overhead. */ interface RowIterator { diff --git a/processing/src/main/java/org/apache/druid/queryng/operators/Operators.java b/processing/src/main/java/org/apache/druid/queryng/operators/Operators.java index 9152cf9d8cd9..4065edc37acd 100644 --- a/processing/src/main/java/org/apache/druid/queryng/operators/Operators.java +++ b/processing/src/main/java/org/apache/druid/queryng/operators/Operators.java @@ -53,7 +53,7 @@ public Iterator iterator() * iterator mechanism (since an operator looks like an iterator.) * * This is a named class so we can unwrap the operator in - * {@link #toOperator()} below. + * {@link static Operator #toOperator(DAGBuilder, Sequence)}. */ public static class OperatorWrapperSequence extends BaseSequence> { diff --git a/processing/src/main/java/org/apache/druid/queryng/operators/SequenceOperator.java b/processing/src/main/java/org/apache/druid/queryng/operators/SequenceOperator.java index cb2c24361063..60fb32b54f33 100644 --- a/processing/src/main/java/org/apache/druid/queryng/operators/SequenceOperator.java +++ b/processing/src/main/java/org/apache/druid/queryng/operators/SequenceOperator.java @@ -73,7 +73,7 @@ public T accumulate(T accumulated, T in) @Override public T next() throws EofException { - if (yielder == null && yielder.isDone()) { + if (yielder == null || yielder.isDone()) { throw Operators.eof(); } if (yielder.isDone()) { From b979649b1e017a5f041edecade5b091421159b26 Mon Sep 17 00:00:00 2001 From: Paul Rogers Date: Thu, 16 Jun 2022 16:19:01 -0700 Subject: [PATCH 06/11] Build fix --- .../main/java/org/apache/druid/queryng/operators/Operator.java | 2 -- 1 file changed, 2 deletions(-) diff --git a/processing/src/main/java/org/apache/druid/queryng/operators/Operator.java b/processing/src/main/java/org/apache/druid/queryng/operators/Operator.java index b9d0fda512aa..2b92d1235db5 100644 --- a/processing/src/main/java/org/apache/druid/queryng/operators/Operator.java +++ b/processing/src/main/java/org/apache/druid/queryng/operators/Operator.java @@ -19,8 +19,6 @@ package org.apache.druid.queryng.operators; -import java.util.Iterator; - /** * An operator is a data pipeline transform: something that operates on * a stream of results in some way. An operator has a very simple lifecycle: From 29d17028888fb19c28b0801869e6954a729d2d4d Mon Sep 17 00:00:00 2001 From: Paul Rogers Date: Fri, 17 Jun 2022 18:15:15 -0700 Subject: [PATCH 07/11] LGTM requested fix --- .../org/apache/druid/queryng/operators/SequenceOperator.java | 1 - 1 file changed, 1 deletion(-) diff --git a/processing/src/main/java/org/apache/druid/queryng/operators/SequenceOperator.java b/processing/src/main/java/org/apache/druid/queryng/operators/SequenceOperator.java index 60fb32b54f33..ec6a9327f117 100644 --- a/processing/src/main/java/org/apache/druid/queryng/operators/SequenceOperator.java +++ b/processing/src/main/java/org/apache/druid/queryng/operators/SequenceOperator.java @@ -80,7 +80,6 @@ public T next() throws EofException closeYielder(); throw Operators.eof(); } - Preconditions.checkState(yielder != null); T value = yielder.get(); yielder = yielder.next(null); return value; From a45f687b72d4eeec50ee2003e24227dbe30db27a Mon Sep 17 00:00:00 2001 From: Paul Rogers Date: Sat, 18 Jun 2022 15:53:48 -0700 Subject: [PATCH 08/11] Added scan query operators --- .../druid/query/scan/ScanResultValue.java | 22 +- .../druid/queryng/fragment/FragmentRun.java | 4 +- .../queryng/fragment/FragmentRunImpl.java | 6 +- .../druid/queryng/guice/QueryNGModule.java | 48 +++ .../queryng/operators/ConcatOperator.java | 4 +- .../druid/queryng/operators/Iterators.java | 14 +- .../queryng/operators/MappingOperator.java | 4 +- .../druid/queryng/operators/NullOperator.java | 2 +- .../druid/queryng/operators/Operator.java | 10 +- .../druid/queryng/operators/Operators.java | 14 + .../operators/OrderedMergeOperator.java | 4 +- .../queryng/operators/PushBackOperator.java | 6 +- .../queryng/operators/SequenceIterator.java | 81 +++++ .../queryng/operators/SequenceOperator.java | 2 +- .../queryng/operators/WrappingOperator.java | 4 +- .../general/QueryRunnerOperator.java | 68 ++++ .../queryng/operators/scan/CursorReader.java | 184 ++++++++++ .../scan/GroupedScanResultLimitOperator.java | 77 +++++ .../scan/ScanBatchToRowOperator.java | 66 ++++ .../scan/ScanCompactListToArrayOperator.java | 70 ++++ .../scan/ScanListToArrayOperator.java | 61 ++++ .../operators/scan/ScanQueryOperator.java | 321 ++++++++++++++++++ .../scan/ScanResultOffsetOperator.java | 76 +++++ .../scan/ScanRowToBatchOperator.java | 73 ++++ .../UngroupedScanResultLimitOperator.java | 85 +++++ .../druid/queryng/planner/ScanPlanner.java | 288 ++++++++++++++++ .../queryng/operators/BasicOperatorTest.java | 4 +- .../queryng/operators/ConcatOperatorTest.java | 4 +- .../druid/queryng/operators/FragmentTest.java | 4 +- .../druid/queryng/operators/MockOperator.java | 2 +- .../queryng/operators/MockOperatorTest.java | 4 +- .../queryng/operators/OperatorTests.java | 4 +- .../operators/OrderedMergeOperatorTest.java | 12 +- .../operators/PushBackOperatorTest.java | 6 +- .../operators/scan/MockScanResultReader.java | 206 +++++++++++ .../scan/ScanQueryOperatorsTest.java | 246 ++++++++++++++ 36 files changed, 2031 insertions(+), 55 deletions(-) create mode 100644 processing/src/main/java/org/apache/druid/queryng/guice/QueryNGModule.java create mode 100644 processing/src/main/java/org/apache/druid/queryng/operators/SequenceIterator.java create mode 100644 processing/src/main/java/org/apache/druid/queryng/operators/general/QueryRunnerOperator.java create mode 100644 processing/src/main/java/org/apache/druid/queryng/operators/scan/CursorReader.java create mode 100644 processing/src/main/java/org/apache/druid/queryng/operators/scan/GroupedScanResultLimitOperator.java create mode 100644 processing/src/main/java/org/apache/druid/queryng/operators/scan/ScanBatchToRowOperator.java create mode 100644 processing/src/main/java/org/apache/druid/queryng/operators/scan/ScanCompactListToArrayOperator.java create mode 100644 processing/src/main/java/org/apache/druid/queryng/operators/scan/ScanListToArrayOperator.java create mode 100644 processing/src/main/java/org/apache/druid/queryng/operators/scan/ScanQueryOperator.java create mode 100644 processing/src/main/java/org/apache/druid/queryng/operators/scan/ScanResultOffsetOperator.java create mode 100644 processing/src/main/java/org/apache/druid/queryng/operators/scan/ScanRowToBatchOperator.java create mode 100644 processing/src/main/java/org/apache/druid/queryng/operators/scan/UngroupedScanResultLimitOperator.java create mode 100644 processing/src/main/java/org/apache/druid/queryng/planner/ScanPlanner.java create mode 100644 processing/src/test/java/org/apache/druid/queryng/operators/scan/MockScanResultReader.java create mode 100644 processing/src/test/java/org/apache/druid/queryng/operators/scan/ScanQueryOperatorsTest.java diff --git a/processing/src/main/java/org/apache/druid/query/scan/ScanResultValue.java b/processing/src/main/java/org/apache/druid/query/scan/ScanResultValue.java index aa368102dd4f..4dc198354105 100644 --- a/processing/src/main/java/org/apache/druid/query/scan/ScanResultValue.java +++ b/processing/src/main/java/org/apache/druid/query/scan/ScanResultValue.java @@ -77,10 +77,17 @@ public Object getEvents() return events; } + @SuppressWarnings("unchecked") + public List getRows() + { + return (List) getEvents(); + } + public long getFirstEventTimestamp(ScanQuery.ResultFormat resultFormat) { if (resultFormat.equals(ScanQuery.ResultFormat.RESULT_FORMAT_LIST)) { - Object timestampObj = ((Map) ((List) this.getEvents()).get(0)).get(ColumnHolder.TIME_COLUMN_NAME); + final List> rows = getRows(); + Object timestampObj = rows.get(0).get(ColumnHolder.TIME_COLUMN_NAME); if (timestampObj == null) { throw new ISE("Unable to compare timestamp for rows without a time column"); } @@ -90,7 +97,8 @@ public long getFirstEventTimestamp(ScanQuery.ResultFormat resultFormat) if (timeColumnIndex == -1) { throw new ISE("Unable to compare timestamp for rows without a time column"); } - List firstEvent = (List) ((List) this.getEvents()).get(0); + final List> rows = getRows(); + final List firstEvent = rows.get(0); return DimensionHandlerUtils.convertObjectToLong(firstEvent.get(timeColumnIndex)); } throw new UOE("Unable to get first event timestamp using result format of [%s]", resultFormat.toString()); @@ -98,14 +106,18 @@ public long getFirstEventTimestamp(ScanQuery.ResultFormat resultFormat) public List toSingleEventScanResultValues() { - List singleEventScanResultValues = new ArrayList<>(); - List events = (List) this.getEvents(); - for (Object event : events) { + final List singleEventScanResultValues = new ArrayList<>(); + for (Object event : getRows()) { singleEventScanResultValues.add(new ScanResultValue(segmentId, columns, Collections.singletonList(event))); } return singleEventScanResultValues; } + public int rowCount() + { + return getRows().size(); + } + @Override public boolean equals(Object o) { diff --git a/processing/src/main/java/org/apache/druid/queryng/fragment/FragmentRun.java b/processing/src/main/java/org/apache/druid/queryng/fragment/FragmentRun.java index 873dc9aa64c4..ad7676e3cc09 100644 --- a/processing/src/main/java/org/apache/druid/queryng/fragment/FragmentRun.java +++ b/processing/src/main/java/org/apache/druid/queryng/fragment/FragmentRun.java @@ -19,7 +19,7 @@ package org.apache.druid.queryng.fragment; -import org.apache.druid.queryng.operators.Operator.RowIterator; +import org.apache.druid.queryng.operators.Operator.ResultIterator; import java.util.List; @@ -37,7 +37,7 @@ public interface FragmentRun extends AutoCloseable { FragmentContext context(); - RowIterator iterator(); + ResultIterator iterator(); /** * Materializes the entire result set as a list. Primarily for testing. diff --git a/processing/src/main/java/org/apache/druid/queryng/fragment/FragmentRunImpl.java b/processing/src/main/java/org/apache/druid/queryng/fragment/FragmentRunImpl.java index d63ccd9eca65..f3d1d7c3674c 100644 --- a/processing/src/main/java/org/apache/druid/queryng/fragment/FragmentRunImpl.java +++ b/processing/src/main/java/org/apache/druid/queryng/fragment/FragmentRunImpl.java @@ -23,14 +23,14 @@ import org.apache.druid.queryng.fragment.FragmentContext.State; import org.apache.druid.queryng.operators.Iterators; import org.apache.druid.queryng.operators.Operator; -import org.apache.druid.queryng.operators.Operator.RowIterator; +import org.apache.druid.queryng.operators.Operator.ResultIterator; import java.util.List; public class FragmentRunImpl implements FragmentRun { private final FragmentContextImpl context; - private RowIterator rootIter; + private ResultIterator rootIter; public FragmentRunImpl(FragmentContextImpl context, Operator root) { @@ -48,7 +48,7 @@ public FragmentRunImpl(FragmentContextImpl context, Operator root) } @Override - public RowIterator iterator() + public ResultIterator iterator() { Preconditions.checkState(context.state == State.RUN); return rootIter; diff --git a/processing/src/main/java/org/apache/druid/queryng/guice/QueryNGModule.java b/processing/src/main/java/org/apache/druid/queryng/guice/QueryNGModule.java new file mode 100644 index 000000000000..1d1c351da7e9 --- /dev/null +++ b/processing/src/main/java/org/apache/druid/queryng/guice/QueryNGModule.java @@ -0,0 +1,48 @@ +/* + * 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.queryng.guice; + +import com.google.inject.Binder; +import com.google.inject.Module; +import org.apache.druid.guice.JsonConfigProvider; +import org.apache.druid.guice.LazySingleton; +import org.apache.druid.queryng.config.QueryNGConfig; +import org.apache.druid.queryng.fragment.FragmentBuilderFactory; +import org.apache.druid.queryng.fragment.FragmentBuilderFactoryImpl; + +/** + * Configure the "shim" version of the NG query engine which entails + * creating a config (to enable or disable the engine) and to create + * a factory for the fragment context. In this early version, all + * other parts of the engine are distributed across various query + * runners. + */ +public class QueryNGModule implements Module +{ + @Override + public void configure(Binder binder) + { + JsonConfigProvider.bind(binder, QueryNGConfig.CONFIG_ROOT, QueryNGConfig.class); + binder + .bind(FragmentBuilderFactory.class) + .to(FragmentBuilderFactoryImpl.class) + .in(LazySingleton.class); + } +} diff --git a/processing/src/main/java/org/apache/druid/queryng/operators/ConcatOperator.java b/processing/src/main/java/org/apache/druid/queryng/operators/ConcatOperator.java index de413d3b2e4f..5d09a9bc7799 100644 --- a/processing/src/main/java/org/apache/druid/queryng/operators/ConcatOperator.java +++ b/processing/src/main/java/org/apache/druid/queryng/operators/ConcatOperator.java @@ -46,7 +46,7 @@ public static Operator concatOrNot( private final Iterator> childIter; private Operator current; - private RowIterator currentIter; + private ResultIterator currentIter; public ConcatOperator(FragmentContext context, List> children) { @@ -55,7 +55,7 @@ public ConcatOperator(FragmentContext context, List> children) } @Override - public RowIterator open() + public ResultIterator open() { return this; } diff --git a/processing/src/main/java/org/apache/druid/queryng/operators/Iterators.java b/processing/src/main/java/org/apache/druid/queryng/operators/Iterators.java index 2773c3dcdaf8..52bc987dd386 100644 --- a/processing/src/main/java/org/apache/druid/queryng/operators/Iterators.java +++ b/processing/src/main/java/org/apache/druid/queryng/operators/Iterators.java @@ -21,25 +21,25 @@ import com.google.common.collect.Lists; import org.apache.druid.queryng.operators.Operator.EofException; -import org.apache.druid.queryng.operators.Operator.RowIterator; +import org.apache.druid.queryng.operators.Operator.ResultIterator; import java.util.Iterator; import java.util.List; import java.util.NoSuchElementException; /** - * Utility methods on top of {@link Operator.RowIterator RowIterator}, + * Utility methods on top of {@link Operator.ResultIterator RowIterator}, * including conversion to a Java iterator (primarily for testing.) */ public class Iterators { public static class ShimIterator implements Iterator { - private final RowIterator operIter; + private final ResultIterator operIter; private boolean eof; private T lookAhead; - public ShimIterator(RowIterator operIter) + public ShimIterator(ResultIterator operIter) { this.operIter = operIter; } @@ -71,7 +71,7 @@ public T next() } - public static Iterable toIterable(RowIterator iter) + public static Iterable toIterable(ResultIterator iter) { return Iterators.toIterable(Iterators.toIterator(iter)); } @@ -87,12 +87,12 @@ public Iterator iterator() }; } - public static Iterator toIterator(RowIterator opIter) + public static Iterator toIterator(ResultIterator opIter) { return new Iterators.ShimIterator(opIter); } - public static List toList(RowIterator operIter) + public static List toList(ResultIterator operIter) { return Lists.newArrayList(new Iterators.ShimIterator(operIter)); } diff --git a/processing/src/main/java/org/apache/druid/queryng/operators/MappingOperator.java b/processing/src/main/java/org/apache/druid/queryng/operators/MappingOperator.java index dd978094742b..3416f02d2443 100644 --- a/processing/src/main/java/org/apache/druid/queryng/operators/MappingOperator.java +++ b/processing/src/main/java/org/apache/druid/queryng/operators/MappingOperator.java @@ -31,7 +31,7 @@ public abstract class MappingOperator implements IterableOperator { private final Operator input; - protected RowIterator inputIter; + protected ResultIterator inputIter; protected State state = State.START; public MappingOperator(FragmentContext context, Operator input) @@ -41,7 +41,7 @@ public MappingOperator(FragmentContext context, Operator input) } @Override - public RowIterator open() + public ResultIterator open() { Preconditions.checkState(state == State.START); inputIter = input.open(); diff --git a/processing/src/main/java/org/apache/druid/queryng/operators/NullOperator.java b/processing/src/main/java/org/apache/druid/queryng/operators/NullOperator.java index b5b103f3891b..31781a4ae30d 100644 --- a/processing/src/main/java/org/apache/druid/queryng/operators/NullOperator.java +++ b/processing/src/main/java/org/apache/druid/queryng/operators/NullOperator.java @@ -39,7 +39,7 @@ public NullOperator(FragmentContext context) } @Override - public RowIterator open() + public ResultIterator open() { Preconditions.checkState(state == State.START); state = State.RUN; diff --git a/processing/src/main/java/org/apache/druid/queryng/operators/Operator.java b/processing/src/main/java/org/apache/druid/queryng/operators/Operator.java index 2b92d1235db5..a393cd6c501e 100644 --- a/processing/src/main/java/org/apache/druid/queryng/operators/Operator.java +++ b/processing/src/main/java/org/apache/druid/queryng/operators/Operator.java @@ -30,10 +30,10 @@ *
  • Closed.
    • * *

    - * Opening an operator returns a {@link Operator.RowIterator RowIterator} + * Opening an operator returns a {@link Operator.ResultIterator RowIterator} * which returns rows. The Java {@code Iterator} class has extra overhead * which we want to avoid on the per-row inner loop code path. A - * {@code RowIterator} has one method: {@link Operator.RowIterator#next() next()}, + * {@code RowIterator} has one method: {@link Operator.ResultIterator#next() next()}, * which either returns a row (however the operator defines it), or throws an * {@link Operator.EofException EofException} when there are no more rows. * Downstream operators need not do any conditional checking: they can just @@ -150,7 +150,7 @@ class EofException extends Exception * {@link static Iterator Operators#toIterator(Operator)}, * but that approach adds overhead. */ - interface RowIterator + interface ResultIterator { T next() throws EofException; } @@ -158,7 +158,7 @@ interface RowIterator /** * Convenience interface for an operator which is its own iterator. */ - interface IterableOperator extends Operator, RowIterator + interface IterableOperator extends Operator, ResultIterator { } @@ -178,7 +178,7 @@ enum State * in the {@code open()} call for simple operators,or later, on demand, for more * complex operators such as in a merge or union. */ - RowIterator open(); + ResultIterator open(); /** * Called at two distinct times. An operator may choose to close a child diff --git a/processing/src/main/java/org/apache/druid/queryng/operators/Operators.java b/processing/src/main/java/org/apache/druid/queryng/operators/Operators.java index 4065edc37acd..533202edeb83 100644 --- a/processing/src/main/java/org/apache/druid/queryng/operators/Operators.java +++ b/processing/src/main/java/org/apache/druid/queryng/operators/Operators.java @@ -22,8 +22,11 @@ import com.google.common.collect.Lists; import org.apache.druid.java.util.common.guava.BaseSequence; import org.apache.druid.java.util.common.guava.Sequence; +import org.apache.druid.query.QueryPlus; +import org.apache.druid.query.QueryRunner; import org.apache.druid.queryng.fragment.DAGBuilder; import org.apache.druid.queryng.operators.Operator.EofException; +import org.apache.druid.queryng.operators.general.QueryRunnerOperator; import java.util.Iterator; import java.util.List; @@ -125,6 +128,17 @@ public static Operator unwrapOperator(Sequence sequence) return null; } + /** + * Create an operator which wraps a query runner which allows a query runner + * to be an input to an operator. The runner, and its sequence, will be optimized + * away at runtime if both the upstream and downstream items are both operators, + * but the shim is left in place if the upstream is actually a query runner. + */ + public static QueryRunnerOperator toOperator(QueryRunner runner, QueryPlus query) + { + return new QueryRunnerOperator(runner, query); + } + public static Iterator toIterator(Operator op) { return new Iterators.ShimIterator(op.open()); diff --git a/processing/src/main/java/org/apache/druid/queryng/operators/OrderedMergeOperator.java b/processing/src/main/java/org/apache/druid/queryng/operators/OrderedMergeOperator.java index 5912337ceb88..3d7dadb58afc 100644 --- a/processing/src/main/java/org/apache/druid/queryng/operators/OrderedMergeOperator.java +++ b/processing/src/main/java/org/apache/druid/queryng/operators/OrderedMergeOperator.java @@ -45,7 +45,7 @@ public class OrderedMergeOperator implements IterableOperator public static class Input { private final Operator input; - private final RowIterator iter; + private final ResultIterator iter; private T currentValue; public Input(Operator input) @@ -130,7 +130,7 @@ public OrderedMergeOperator( } @Override - public RowIterator open() + public ResultIterator open() { for (Input input : inputs.get()) { if (!input.eof()) { diff --git a/processing/src/main/java/org/apache/druid/queryng/operators/PushBackOperator.java b/processing/src/main/java/org/apache/druid/queryng/operators/PushBackOperator.java index 0b7d9294583e..f96f488b56d2 100644 --- a/processing/src/main/java/org/apache/druid/queryng/operators/PushBackOperator.java +++ b/processing/src/main/java/org/apache/druid/queryng/operators/PushBackOperator.java @@ -30,13 +30,13 @@ public class PushBackOperator implements IterableOperator { private final Operator input; - private RowIterator inputIter; + private ResultIterator inputIter; private T pushed; public PushBackOperator( FragmentContext context, Operator input, - RowIterator inputIter, + ResultIterator inputIter, T pushed) { this.input = input; @@ -51,7 +51,7 @@ public PushBackOperator(FragmentContext context, Operator input) } @Override - public RowIterator open() + public ResultIterator open() { if (inputIter == null) { inputIter = input.open(); diff --git a/processing/src/main/java/org/apache/druid/queryng/operators/SequenceIterator.java b/processing/src/main/java/org/apache/druid/queryng/operators/SequenceIterator.java new file mode 100644 index 000000000000..b267a47ccaf7 --- /dev/null +++ b/processing/src/main/java/org/apache/druid/queryng/operators/SequenceIterator.java @@ -0,0 +1,81 @@ +/* + * 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.queryng.operators; + +import com.google.common.base.Preconditions; +import org.apache.druid.java.util.common.guava.Sequence; +import org.apache.druid.java.util.common.guava.Yielder; +import org.apache.druid.java.util.common.guava.YieldingAccumulator; + +import java.io.IOException; +import java.util.Iterator; + +/** + * Iterator over a sequence. + */ +public class SequenceIterator implements Iterator, AutoCloseable +{ + private Yielder yielder; + + public static SequenceIterator of(Sequence sequence) + { + return new SequenceIterator(sequence); + } + + public SequenceIterator(Sequence sequence) + { + this.yielder = sequence.toYielder( + null, + new YieldingAccumulator() + { + @Override + public T accumulate(T accumulated, T in) + { + yield(); + return in; + } + } + ); + } + + @Override + public boolean hasNext() + { + return !yielder.isDone(); + } + + @Override + public T next() + { + Preconditions.checkState(!yielder.isDone()); + T value = yielder.get(); + yielder = yielder.next(null); + return value; + } + + @Override + public void close() throws IOException + { + if (yielder != null) { + yielder.close(); + yielder = null; + } + } +} diff --git a/processing/src/main/java/org/apache/druid/queryng/operators/SequenceOperator.java b/processing/src/main/java/org/apache/druid/queryng/operators/SequenceOperator.java index ec6a9327f117..3440c4ff1d33 100644 --- a/processing/src/main/java/org/apache/druid/queryng/operators/SequenceOperator.java +++ b/processing/src/main/java/org/apache/druid/queryng/operators/SequenceOperator.java @@ -52,7 +52,7 @@ public SequenceOperator(FragmentContext context, Sequence sequence) } @Override - public RowIterator open() + public ResultIterator open() { Preconditions.checkState(yielder == null); yielder = sequence.toYielder( diff --git a/processing/src/main/java/org/apache/druid/queryng/operators/WrappingOperator.java b/processing/src/main/java/org/apache/druid/queryng/operators/WrappingOperator.java index 5cdf13e43caf..466a883d2a1b 100644 --- a/processing/src/main/java/org/apache/druid/queryng/operators/WrappingOperator.java +++ b/processing/src/main/java/org/apache/druid/queryng/operators/WrappingOperator.java @@ -41,10 +41,10 @@ public WrappingOperator(FragmentContext context, Operator input) } @Override - public RowIterator open() + public ResultIterator open() { Preconditions.checkState(state == State.START); - RowIterator inputIter = input.open(); + ResultIterator inputIter = input.open(); state = State.RUN; onOpen(); return inputIter; diff --git a/processing/src/main/java/org/apache/druid/queryng/operators/general/QueryRunnerOperator.java b/processing/src/main/java/org/apache/druid/queryng/operators/general/QueryRunnerOperator.java new file mode 100644 index 000000000000..fb729e3dcb80 --- /dev/null +++ b/processing/src/main/java/org/apache/druid/queryng/operators/general/QueryRunnerOperator.java @@ -0,0 +1,68 @@ +/* + * 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.queryng.operators.general; + +import org.apache.druid.java.util.common.guava.Sequence; +import org.apache.druid.query.QueryPlus; +import org.apache.druid.query.QueryRunner; +import org.apache.druid.queryng.fragment.FragmentContext; +import org.apache.druid.queryng.operators.Operator; +import org.apache.druid.queryng.operators.Operators; + +/** + * Operator which wraps a query runner. When used in the interim + * "shim" architecture, this operator allows a query runner to be + * the input (upstream) to some other (downstream) operator. The + * upstream query runner may give rise to its own operator. In that + * case, at runtime, the sequence wrapper for that operator is + * optimized away, leaving just the two operators. + */ +public class QueryRunnerOperator implements Operator +{ + protected final FragmentContext context; + private final QueryRunner runner; + private final QueryPlus query; + private Operator child; + + public QueryRunnerOperator(QueryRunner runner, QueryPlus query) + { + this.context = query.fragmentBuilder().context(); + this.runner = runner; + this.query = query; + context.register(this); + } + + @Override + public ResultIterator open() + { + Sequence seq = runner.run(query, context.responseContext()); + child = Operators.toOperator(context, seq); + return child.open(); + } + + @Override + public void close(boolean cascade) + { + if (child != null && cascade) { + child.close(cascade); + } + child = null; + } +} diff --git a/processing/src/main/java/org/apache/druid/queryng/operators/scan/CursorReader.java b/processing/src/main/java/org/apache/druid/queryng/operators/scan/CursorReader.java new file mode 100644 index 000000000000..f9cee6d12de2 --- /dev/null +++ b/processing/src/main/java/org/apache/druid/queryng/operators/scan/CursorReader.java @@ -0,0 +1,184 @@ +/* + * 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.queryng.operators.scan; + +import com.google.common.base.Supplier; +import org.apache.druid.java.util.common.DateTimes; +import org.apache.druid.java.util.common.UOE; +import org.apache.druid.query.scan.ScanQuery.ResultFormat; +import org.apache.druid.segment.BaseObjectColumnValueSelector; +import org.apache.druid.segment.Cursor; +import org.apache.druid.segment.column.ColumnHolder; + +import java.util.ArrayList; +import java.util.Iterator; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; + +/** + * The cursor reader is a leaf operator which uses a cursor to access + * data, which is returned via the operator protocol. Converts cursor + * data into one of two supported Druid formats. Enforces a query row limit. + *

    + * Unlike most operators, this one is created on the fly by its parent + * to scan a specific query known only at runtime. A storage adapter may + * choose to create one or more queries: each is handled by an instance of this + * class. + * + * @see {@link org.apache.druid.query.scan.ScanQueryEngine} + */ +public class CursorReader implements Iterator +{ + private final Cursor cursor; + private final List selectedColumns; + private final long limit; + private final int batchSize; + private final ResultFormat resultFormat; + private final List> columnAccessors; + private long targetCount; + private long rowCount; + + public CursorReader( + final Cursor cursor, + final List selectedColumns, + final long limit, + final int batchSize, + final ResultFormat resultFormat, + final boolean isLegacy + ) + { + this.cursor = cursor; + this.selectedColumns = selectedColumns; + this.limit = limit; + this.batchSize = batchSize; + this.resultFormat = resultFormat; + this.columnAccessors = buildAccessors(isLegacy); + } + + private List> buildAccessors(final boolean isLegacy) + { + List> accessors = new ArrayList<>(selectedColumns.size()); + for (String column : selectedColumns) { + final Supplier accessor; + final BaseObjectColumnValueSelector selector; + if (isLegacy && ScanQueryOperator.LEGACY_TIMESTAMP_KEY.equals(column)) { + selector = cursor.getColumnSelectorFactory() + .makeColumnValueSelector(ColumnHolder.TIME_COLUMN_NAME); + accessor = new Supplier() { + @Override + public Object get() + { + return DateTimes.utc((Long) selector.getObject()); + } + }; + } else { + selector = cursor.getColumnSelectorFactory().makeColumnValueSelector(column); + if (selector == null) { + accessor = new Supplier() { + @Override + public Object get() + { + return null; + } + }; + } else { + accessor = new Supplier() { + @Override + public Object get() + { + return selector.getObject(); + } + }; + } + } + accessors.add(accessor); + } + return accessors; + } + + @Override + public boolean hasNext() + { + return !cursor.isDone() && rowCount < limit; + } + + @Override + public Object next() + { + targetCount = Math.min(limit - rowCount, rowCount + batchSize); + switch (resultFormat) { + case RESULT_FORMAT_LIST: + return nextAsListOfMaps(); + case RESULT_FORMAT_COMPACTED_LIST: + return nextAsCompactList(); + default: + throw new UOE("resultFormat[%s] is not supported", resultFormat.toString()); + } + } + + private void advance() + { + cursor.advance(); + rowCount++; + } + + private boolean hasNextRow() + { + return !cursor.isDone() && rowCount < targetCount; + } + + private Object getColumnValue(int i) + { + return columnAccessors.get(i).get(); + } + + /** + * Convert a cursor row into a simple list of maps, where each map + * represents a single event, and each map entry represents a column. + */ + public Object nextAsListOfMaps() + { + final List> events = new ArrayList<>(batchSize); + while (hasNextRow()) { + final Map theEvent = new LinkedHashMap<>(); + for (int j = 0; j < selectedColumns.size(); j++) { + theEvent.put(selectedColumns.get(j), getColumnValue(j)); + } + events.add(theEvent); + advance(); + } + return events; + } + + public Object nextAsCompactList() + { + final List> events = new ArrayList<>(batchSize); + while (hasNextRow()) { + final List theEvent = new ArrayList<>(selectedColumns.size()); + for (int j = 0; j < selectedColumns.size(); j++) { + theEvent.add(getColumnValue(j)); + } + events.add(theEvent); + advance(); + } + return events; + } +} diff --git a/processing/src/main/java/org/apache/druid/queryng/operators/scan/GroupedScanResultLimitOperator.java b/processing/src/main/java/org/apache/druid/queryng/operators/scan/GroupedScanResultLimitOperator.java new file mode 100644 index 000000000000..c318b289edbc --- /dev/null +++ b/processing/src/main/java/org/apache/druid/queryng/operators/scan/GroupedScanResultLimitOperator.java @@ -0,0 +1,77 @@ +/* + * 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.queryng.operators.scan; + +import com.google.common.annotations.VisibleForTesting; +import org.apache.druid.query.scan.ScanResultValue; +import org.apache.druid.queryng.fragment.FragmentContext; +import org.apache.druid.queryng.operators.MappingOperator; +import org.apache.druid.queryng.operators.Operator; +import org.apache.druid.queryng.operators.Operators; + +import java.util.List; + +/** + * Limit scan query results when each batch has multiple rows. + * + * @see {@link org.apache.druid.query.scan.ScanQueryLimitRowIterator} + */ +public class GroupedScanResultLimitOperator extends MappingOperator +{ + private final long limit; + private long rowCount; + + @VisibleForTesting + public GroupedScanResultLimitOperator( + FragmentContext context, + Operator child, + long limit) + { + super(context, child); + this.limit = limit; + } + + @Override + public ScanResultValue next() throws EofException + { + if (rowCount >= limit) { + // Already at limit. + throw Operators.eof(); + } + + // With throw EofException if no more input rows. + ScanResultValue batch = inputIter.next(); + List events = (List) batch.getEvents(); + if (events.size() <= limit - rowCount) { + // Entire batch is below limit. + rowCount += events.size(); + return batch; + } else { + // last batch + // single batch length is <= rowCount.MAX_VALUE, so this should not overflow + int numLeft = (int) (limit - rowCount); + rowCount = limit; + return new ScanResultValue( + batch.getSegmentId(), + batch.getColumns(), + events.subList(0, numLeft)); + } + } +} diff --git a/processing/src/main/java/org/apache/druid/queryng/operators/scan/ScanBatchToRowOperator.java b/processing/src/main/java/org/apache/druid/queryng/operators/scan/ScanBatchToRowOperator.java new file mode 100644 index 000000000000..dc70b1c415ca --- /dev/null +++ b/processing/src/main/java/org/apache/druid/queryng/operators/scan/ScanBatchToRowOperator.java @@ -0,0 +1,66 @@ +/* + * 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.queryng.operators.scan; + +import org.apache.druid.query.scan.ScanResultValue; +import org.apache.druid.queryng.fragment.FragmentContext; +import org.apache.druid.queryng.operators.MappingOperator; +import org.apache.druid.queryng.operators.Operator; + +import java.util.Iterator; +import java.util.List; + +/** + * Converts an input operator which returns scan query "batches" to individual map records. + * The record type is assumed to be one of the valid + * {@link org.apache.druid.query.scan.ScanQuery.ResultFormat + * ResultFormat} types. + */ +public class ScanBatchToRowOperator extends MappingOperator +{ + private Iterator batchIter; + + public ScanBatchToRowOperator(FragmentContext context, Operator input) + { + super(context, input); + } + + @Override + @SuppressWarnings("unchecked") + public T next() throws EofException + { + while (true) { + if (batchIter == null) { + batchIter = ((List) (inputIter.next().getRows())).iterator(); + } + if (batchIter.hasNext()) { + return batchIter.next(); + } + batchIter = null; + } + } + + @Override + public void close(boolean cascade) + { + batchIter = null; + super.close(cascade); + } +} diff --git a/processing/src/main/java/org/apache/druid/queryng/operators/scan/ScanCompactListToArrayOperator.java b/processing/src/main/java/org/apache/druid/queryng/operators/scan/ScanCompactListToArrayOperator.java new file mode 100644 index 000000000000..48b07690ee33 --- /dev/null +++ b/processing/src/main/java/org/apache/druid/queryng/operators/scan/ScanCompactListToArrayOperator.java @@ -0,0 +1,70 @@ +/* + * 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.queryng.operators.scan; + +import org.apache.druid.java.util.common.ISE; +import org.apache.druid.queryng.fragment.FragmentContext; +import org.apache.druid.queryng.operators.MappingOperator; +import org.apache.druid.queryng.operators.Operator; + +import java.util.List; + +/** + * Converts individual scan query rows with the + * {@link org.apache.druid.query.scan.ScanQuery.ResultFormat.RESULT_FORMAT_COMPACTED_LIST + * ResultFormat.RESULT_FORMAT_COMPACTED_LIST} format into an object array with fields + * in the order given by the output schema. + * + * @See {@link org.apache.druid.query.scan.ScanQueryQueryToolChest.resultsAsArrays + * ScanQueryQueryToolChest.resultsAsArrays} + */ +public class ScanCompactListToArrayOperator extends MappingOperator, Object[]> +{ + private final List fields; + + public ScanCompactListToArrayOperator( + FragmentContext context, + Operator> input, + List fields) + { + super(context, input); + this.fields = fields; + } + + @Override + public Object[] next() throws EofException + { + List row = inputIter.next(); + if (row.size() == fields.size()) { + return row.toArray(); + } else if (fields.isEmpty()) { + return new Object[0]; + } else { + // Uh oh... mismatch in expected and actual field count. I don't think + // this should happen, so let's throw an exception. If this really does + // happen, and there's a good reason for it, then we should remap + // the result row here. + throw new ISE( + "Mismatch in expected [%d] vs actual [%s] field count", + fields.size(), + row.size()); + } + } +} diff --git a/processing/src/main/java/org/apache/druid/queryng/operators/scan/ScanListToArrayOperator.java b/processing/src/main/java/org/apache/druid/queryng/operators/scan/ScanListToArrayOperator.java new file mode 100644 index 000000000000..cb53ed99902d --- /dev/null +++ b/processing/src/main/java/org/apache/druid/queryng/operators/scan/ScanListToArrayOperator.java @@ -0,0 +1,61 @@ +/* + * 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.queryng.operators.scan; + +import org.apache.druid.queryng.fragment.FragmentContext; +import org.apache.druid.queryng.operators.MappingOperator; +import org.apache.druid.queryng.operators.Operator; + +import java.util.List; +import java.util.Map; + +/** + * Converts individual scan query rows with the + * {@link org.apache.druid.query.scan.ScanQuery.ResultFormat.RESULT_FORMAT_LIST + * ResultFormat.RESULT_FORMAT_LIST} format into an object array with fields + * in the order given by the output schema. + * + * @See {@link org.apache.druid.query.scan.ScanQueryQueryToolChest.resultsAsArrays + * ScanQueryQueryToolChest.resultsAsArrays} + */ +public class ScanListToArrayOperator extends MappingOperator, Object[]> +{ + private final List fields; + + public ScanListToArrayOperator( + FragmentContext context, + Operator> input, + List fields) + { + super(context, input); + this.fields = fields; + } + + @Override + public Object[] next() throws EofException + { + Map row = inputIter.next(); + final Object[] rowArray = new Object[fields.size()]; + for (int i = 0; i < fields.size(); i++) { + rowArray[i] = row.get(fields.get(i)); + } + return rowArray; + } +} diff --git a/processing/src/main/java/org/apache/druid/queryng/operators/scan/ScanQueryOperator.java b/processing/src/main/java/org/apache/druid/queryng/operators/scan/ScanQueryOperator.java new file mode 100644 index 000000000000..b82f6209d4e5 --- /dev/null +++ b/processing/src/main/java/org/apache/druid/queryng/operators/scan/ScanQueryOperator.java @@ -0,0 +1,321 @@ +/* + * 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.queryng.operators.scan; + +import com.google.common.base.Preconditions; +import com.google.common.collect.Iterables; +import com.google.common.collect.Sets; +import org.apache.druid.java.util.common.ISE; +import org.apache.druid.java.util.common.granularity.Granularities; +import org.apache.druid.query.QueryContexts; +import org.apache.druid.query.context.ResponseContext; +import org.apache.druid.query.filter.Filter; +import org.apache.druid.query.scan.ScanQuery; +import org.apache.druid.query.scan.ScanResultValue; +import org.apache.druid.queryng.fragment.FragmentContext; +import org.apache.druid.queryng.operators.Operator; +import org.apache.druid.queryng.operators.Operators; +import org.apache.druid.queryng.operators.SequenceIterator; +import org.apache.druid.segment.Cursor; +import org.apache.druid.segment.Segment; +import org.apache.druid.segment.StorageAdapter; +import org.apache.druid.segment.VirtualColumn; +import org.apache.druid.segment.column.ColumnHolder; +import org.apache.druid.segment.filter.Filters; +import org.joda.time.Interval; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.Set; + +/** + * Implements a scan query against a fragment using a storage adapter which may + * return one or more cursors for the segment. Each cursor is processed using + * a {@link CursorReader}. The set of cursors is known only at run time. + * + * @see {@link org.apache.druid.query.scan.ScanQueryEngine} + */ +public class ScanQueryOperator implements Operator +{ + static final String LEGACY_TIMESTAMP_KEY = "timestamp"; + + /** + * Manages the run state for this operator. + */ + private class Impl implements ResultIterator + { + private final FragmentContext context; + private final SequenceIterator iter; + private final List selectedColumns; + private final long limit; + private CursorReader cursorReader; + private long rowCount; + @SuppressWarnings("unused") + private int batchCount; + + private Impl(FragmentContext context) + { + this.context = context; + ResponseContext responseContext = context.responseContext(); + responseContext.add(ResponseContext.Keys.NUM_SCANNED_ROWS, 0L); + long baseLimit = query.getScanRowsLimit(); + if (limitType() == Limit.GLOBAL) { + limit = baseLimit - (Long) responseContext.get(ResponseContext.Keys.NUM_SCANNED_ROWS); + } else { + limit = baseLimit; + } + final StorageAdapter adapter = segment.asStorageAdapter(); + //final StorageAdapter adapter = new MockStorageAdapter(); + if (adapter == null) { + throw new ISE( + "Null storage adapter found. Probably trying to issue a query against a segment being memory unmapped." + ); + } + if (isWildcard()) { + selectedColumns = inferColumns(adapter, isLegacy); + } else { + selectedColumns = columns; + } + iter = SequenceIterator.of(adapter.makeCursors( + filter, + interval(), + query.getVirtualColumns(), + Granularities.ALL, + isDescendingOrder(), + null + )); + } + + protected List inferColumns(StorageAdapter adapter, boolean isLegacy) + { + List cols = new ArrayList<>(); + final Set availableColumns = Sets.newLinkedHashSet( + Iterables.concat( + Collections.singleton(isLegacy ? LEGACY_TIMESTAMP_KEY : ColumnHolder.TIME_COLUMN_NAME), + Iterables.transform( + Arrays.asList(query.getVirtualColumns().getVirtualColumns()), + VirtualColumn::getOutputName + ), + adapter.getAvailableDimensions(), + adapter.getAvailableMetrics() + ) + ); + + cols.addAll(availableColumns); + + if (isLegacy) { + cols.remove(ColumnHolder.TIME_COLUMN_NAME); + } + return cols; + } + + /** + * Return the next batch of events from a cursor. Enforce the + * timeout limit. + *

    + * Checks if another batch of events is available. They are available if + * we have (or can get) a cursor which has rows, and we are not at the + * limit set for this operator. + * @throws EofException + */ + @Override + public ScanResultValue next() throws EofException + { + context.checkTimeout(); + while (true) { + if (cursorReader != null) { + if (cursorReader.hasNext()) { + // Happy path + List result = (List) cursorReader.next(); + batchCount++; + rowCount += result.size(); + return new ScanResultValue( + segmentId, + selectedColumns, + result); + } + // Cursor is done or was empty. + closeCursorReader(); + } + if (iter == null) { + // Done previously + throw Operators.eof(); + } + if (rowCount > limit) { + // Reached row limit + finish(); + throw Operators.eof(); + } + if (!iter.hasNext()) { + // No more cursors + finish(); + throw Operators.eof(); + } + // Read from the next cursor. + cursorReader = new CursorReader( + iter.next(), + selectedColumns, + limit - rowCount, + batchSize, + query.getResultFormat(), + isLegacy); + } + } + + private void closeCursorReader() + { + if (cursorReader != null) { + cursorReader = null; + } + } + + private void finish() + { + closeCursorReader(); + ResponseContext responseContext = context.responseContext(); + responseContext.add(ResponseContext.Keys.NUM_SCANNED_ROWS, rowCount); + try { + iter.close(); + } + catch (IOException e) { + // Ignore + } + } + } + + public enum Limit + { + NONE, + /** + * If we're performing time-ordering, we want to scan through the first `limit` rows in each + * segment ignoring the number of rows already counted on other segments. + */ + LOCAL, + GLOBAL + } + + protected final FragmentContext context; + private final ScanQuery query; + private final Segment segment; + private final String segmentId; + private final List columns; + private final Filter filter; + private final boolean isLegacy; + private final int batchSize; + private Impl impl; + + public ScanQueryOperator( + final FragmentContext context, + final ScanQuery query, + final Segment segment) + { + this.context = context; + this.query = query; + this.segment = segment; + this.segmentId = segment.getId().toString(); + this.columns = defineColumns(query); + List intervals = query.getQuerySegmentSpec().getIntervals(); + Preconditions.checkArgument(intervals.size() == 1, "Can only handle a single interval, got [%s]", intervals); + this.filter = Filters.convertToCNFFromQueryContext(query, Filters.toFilter(query.getFilter())); + this.isLegacy = Preconditions.checkNotNull(query.isLegacy(), "Expected non-null 'legacy' parameter"); + this.batchSize = query.getBatchSize(); + context.register(this); + } + + /** + * Define the query columns when the list is given by the query. + */ + private List defineColumns(ScanQuery query) + { + List queryCols = query.getColumns(); + + // Missing or empty list means wildcard + if (queryCols == null || queryCols.isEmpty()) { + return null; + } + final List planCols = new ArrayList<>(); + if (query.isLegacy() && !queryCols.contains(LEGACY_TIMESTAMP_KEY)) { + planCols.add(LEGACY_TIMESTAMP_KEY); + } + + // Unless we're in legacy mode, planCols equals query.getColumns() exactly. This is nice since it makes + // the compactedList form easier to use. + planCols.addAll(queryCols); + return planCols; + } + + public boolean isWildcard(ScanQuery query) + { + return (query.getColumns() == null || query.getColumns().isEmpty()); + } + + // TODO: Review against latest + public boolean isDescendingOrder() + { + return query.getTimeOrder().equals(ScanQuery.Order.DESCENDING) || + (query.getTimeOrder().equals(ScanQuery.Order.NONE) && query.isDescending()); + } + + public boolean hasTimeout() + { + return QueryContexts.hasTimeout(query); + } + + public boolean isWildcard() + { + return columns == null; + } + + // TODO: Review against latest + public Limit limitType() + { + if (!query.isLimited()) { + return Limit.NONE; + } else if (query.getTimeOrder().equals(ScanQuery.Order.NONE)) { + return Limit.LOCAL; + } else { + return Limit.GLOBAL; + } + } + + public Interval interval() + { + return query.getQuerySegmentSpec().getIntervals().get(0); + } + + @Override + public ResultIterator open() + { + impl = new Impl(context); + return impl; + } + + @Override + public void close(boolean cascade) + { + if (impl != null) { + impl.closeCursorReader(); + } + impl = null; + } +} diff --git a/processing/src/main/java/org/apache/druid/queryng/operators/scan/ScanResultOffsetOperator.java b/processing/src/main/java/org/apache/druid/queryng/operators/scan/ScanResultOffsetOperator.java new file mode 100644 index 000000000000..a2a8fc99124c --- /dev/null +++ b/processing/src/main/java/org/apache/druid/queryng/operators/scan/ScanResultOffsetOperator.java @@ -0,0 +1,76 @@ +/* + * 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.queryng.operators.scan; + +import org.apache.druid.query.scan.ScanResultValue; +import org.apache.druid.queryng.fragment.FragmentContext; +import org.apache.druid.queryng.operators.MappingOperator; +import org.apache.druid.queryng.operators.Operator; + +import java.util.List; + +/** + * Offset that skips a given number of rows on top of a skips ScanQuery. It is + * used to implement the "offset" feature. + * + * @see {@link org.apache.druid.query.scan.ScanQueryOffsetSequence} + */ +public class ScanResultOffsetOperator extends MappingOperator +{ + private final long offset; + private long rowCount; + + public ScanResultOffsetOperator( + FragmentContext context, + Operator input, + long offset) + { + super(context, input); + this.offset = offset; + } + + @Override + public ScanResultValue next() throws EofException + { + if (rowCount == 0) { + return skip(); + } + return inputIter.next(); + } + + private ScanResultValue skip() throws EofException + { + while (true) { + ScanResultValue batch = inputIter.next(); + final List rows = batch.getRows(); + final int eventCount = rows.size(); + final long toSkip = offset - rowCount; + if (toSkip >= eventCount) { + rowCount += eventCount; + continue; + } + rowCount += eventCount - toSkip; + return new ScanResultValue( + batch.getSegmentId(), + batch.getColumns(), + rows.subList((int) toSkip, eventCount)); + } + } +} diff --git a/processing/src/main/java/org/apache/druid/queryng/operators/scan/ScanRowToBatchOperator.java b/processing/src/main/java/org/apache/druid/queryng/operators/scan/ScanRowToBatchOperator.java new file mode 100644 index 000000000000..6816b8194fc3 --- /dev/null +++ b/processing/src/main/java/org/apache/druid/queryng/operators/scan/ScanRowToBatchOperator.java @@ -0,0 +1,73 @@ +/* + * 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.queryng.operators.scan; + +import org.apache.druid.query.scan.ScanResultValue; +import org.apache.druid.queryng.fragment.FragmentContext; +import org.apache.druid.queryng.operators.MappingOperator; +import org.apache.druid.queryng.operators.Operator; +import org.apache.druid.queryng.operators.Operators; + +import java.util.ArrayList; +import java.util.List; + +/** + * Pack a set of individual scan results into a batch up to the + * given size. + * + * @see {@link org.apache.druid.query.scan.ScanQueryLimitRowIterator} + */ +public class ScanRowToBatchOperator extends MappingOperator +{ + private final int batchSize; + + public ScanRowToBatchOperator( + FragmentContext context, + Operator child, + int batchSize) + { + super(context, child); + this.batchSize = batchSize; + } + + @Override + public ScanResultValue next() throws EofException + { + List eventsToAdd = new ArrayList<>(batchSize); + List columns = null; + while (eventsToAdd.size() < batchSize) { + try { + ScanResultValue srv = inputIter.next(); + if (columns == null) { + columns = srv.getColumns(); + } + eventsToAdd.add(srv.getRows().get(0)); + } + catch (EofException e) { + if (eventsToAdd.isEmpty()) { + throw Operators.eof(); + } + // We'll report EOF on the next call. + break; + } + } + return new ScanResultValue(null, columns, eventsToAdd); + } +} diff --git a/processing/src/main/java/org/apache/druid/queryng/operators/scan/UngroupedScanResultLimitOperator.java b/processing/src/main/java/org/apache/druid/queryng/operators/scan/UngroupedScanResultLimitOperator.java new file mode 100644 index 000000000000..aaec8ec43f10 --- /dev/null +++ b/processing/src/main/java/org/apache/druid/queryng/operators/scan/UngroupedScanResultLimitOperator.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.queryng.operators.scan; + +import com.google.common.annotations.VisibleForTesting; +import org.apache.druid.query.scan.ScanResultValue; +import org.apache.druid.queryng.fragment.FragmentContext; +import org.apache.druid.queryng.operators.MappingOperator; +import org.apache.druid.queryng.operators.Operator; +import org.apache.druid.queryng.operators.Operators; + +import java.util.ArrayList; +import java.util.List; + +/** + * Iterates over scan query results which each batch contains one row. + * + * @see {@link org.apache.druid.query.scan.ScanQueryLimitRowIterator} + */ +public class UngroupedScanResultLimitOperator extends MappingOperator +{ + private final long limit; + private final int batchSize; + private long rowCount; + + @VisibleForTesting + public UngroupedScanResultLimitOperator( + FragmentContext context, + Operator child, + long limit, + int batchSize) + { + super(context, child); + this.limit = limit; + this.batchSize = batchSize; + } + + @Override + public ScanResultValue next() throws EofException + { + if (rowCount >= limit) { + throw Operators.eof(); + } + // Perform single-event ScanResultValue batching at the outer level. Each + // scan result value from the input operator in this case will only have + // one event so there's no need to iterate through events. + List eventsToAdd = new ArrayList<>(batchSize); + List columns = null; + while (eventsToAdd.size() < batchSize && rowCount < limit) { + try { + ScanResultValue srv = inputIter.next(); + if (columns == null) { + columns = srv.getColumns(); + } + eventsToAdd.add(srv.getRows().get(0)); + rowCount++; + } + catch (EofException e) { + if (eventsToAdd.isEmpty()) { + throw Operators.eof(); + } + // We'll report EOF on the next call. + break; + } + } + return new ScanResultValue(null, columns, eventsToAdd); + } +} diff --git a/processing/src/main/java/org/apache/druid/queryng/planner/ScanPlanner.java b/processing/src/main/java/org/apache/druid/queryng/planner/ScanPlanner.java new file mode 100644 index 000000000000..494a4fdeb54e --- /dev/null +++ b/processing/src/main/java/org/apache/druid/queryng/planner/ScanPlanner.java @@ -0,0 +1,288 @@ +/* + * 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.queryng.planner; + +import com.google.common.collect.ImmutableMap; +import org.apache.druid.java.util.common.ISE; +import org.apache.druid.java.util.common.JodaUtils; +import org.apache.druid.java.util.common.UOE; +import org.apache.druid.java.util.common.guava.Sequence; +import org.apache.druid.query.Query; +import org.apache.druid.query.QueryContexts; +import org.apache.druid.query.QueryPlus; +import org.apache.druid.query.QueryRunner; +import org.apache.druid.query.context.ResponseContext; +import org.apache.druid.query.scan.ScanQuery; +import org.apache.druid.query.scan.ScanQueryConfig; +import org.apache.druid.query.scan.ScanResultValue; +import org.apache.druid.queryng.fragment.FragmentContext; +import org.apache.druid.queryng.operators.ConcatOperator; +import org.apache.druid.queryng.operators.Operator; +import org.apache.druid.queryng.operators.Operators; +import org.apache.druid.queryng.operators.scan.GroupedScanResultLimitOperator; +import org.apache.druid.queryng.operators.scan.ScanBatchToRowOperator; +import org.apache.druid.queryng.operators.scan.ScanCompactListToArrayOperator; +import org.apache.druid.queryng.operators.scan.ScanListToArrayOperator; +import org.apache.druid.queryng.operators.scan.ScanQueryOperator; +import org.apache.druid.queryng.operators.scan.ScanResultOffsetOperator; +import org.apache.druid.queryng.operators.scan.UngroupedScanResultLimitOperator; +import org.apache.druid.segment.Segment; + +import java.util.ArrayList; +import java.util.List; +import java.util.Map; + +/** + * Scan-specific parts of the hybrid query planner. + * + * @see {@link QueryPlanner} + */ +public class ScanPlanner +{ + /** + * Sets up an operator over a ScanResultValue operator. Its behaviour + * varies depending on whether the query is returning time-ordered values and whether the CTX_KEY_OUTERMOST + * flag is false. + *

    + * Behaviours: + *

      + *
    1. No time ordering: expects the child to produce ScanResultValues which each contain up to query.batchSize events. + * The operator will be "done" when the limit of events is reached. The final ScanResultValue might contain + * fewer than batchSize events so that the limit number of events is returned.
      2. + *
    2. Time Ordering, CTX_KEY_OUTERMOST false: Same behaviour as no time ordering.
      3. + *
    3. Time Ordering, CTX_KEY_OUTERMOST=true or null: The child operator in this case should produce ScanResultValues + * that contain only one event each for the CachingClusteredClient n-way merge. This operator will perform + * batching according to query batch size until the limit is reached.
      4. + *
    + * + * @see {@link org.apache.druid.query.scan.ScanQueryLimitRowIterator} + * @see {@link org.apache.druid.query.scan.ScanQueryQueryToolChest.mergeResults} + */ + public static Sequence runLimitAndOffset( + final QueryPlus queryPlus, + final QueryRunner input, + final ResponseContext responseContext, + final ScanQueryConfig scanQueryConfig) + { + // Remove "offset" and add it to the "limit" (we won't push the offset down, just apply it here, at the + // merge at the top of the stack). + final ScanQuery originalQuery = ((ScanQuery) (queryPlus.getQuery())); + ScanQuery.verifyOrderByForNativeExecution(originalQuery); + + final long newLimit; + if (!originalQuery.isLimited()) { + // Unlimited stays unlimited. + newLimit = Long.MAX_VALUE; + } else if (originalQuery.getScanRowsLimit() > Long.MAX_VALUE - originalQuery.getScanRowsOffset()) { + throw new ISE( + "Cannot apply limit[%d] with offset[%d] due to overflow", + originalQuery.getScanRowsLimit(), + originalQuery.getScanRowsOffset() + ); + } else { + newLimit = originalQuery.getScanRowsLimit() + originalQuery.getScanRowsOffset(); + } + + // Ensure "legacy" is a non-null value, such that all other nodes this query is forwarded to will treat it + // the same way, even if they have different default legacy values. + final ScanQuery queryToRun = originalQuery.withNonNullLegacy(scanQueryConfig) + .withOffset(0) + .withLimit(newLimit); + + final boolean hasLimit = queryToRun.isLimited(); + final boolean hasOffset = originalQuery.getScanRowsOffset() > 0; + + // Short-circuit if no limit or offset. + if (!hasLimit && !hasOffset) { + return input.run(queryPlus.withQuery(queryToRun), responseContext); + } + + Query historicalQuery = queryToRun; + if (hasLimit) { + ScanQuery.ResultFormat resultFormat = queryToRun.getResultFormat(); + if (ScanQuery.ResultFormat.RESULT_FORMAT_VALUE_VECTOR.equals(resultFormat)) { + throw new UOE(ScanQuery.ResultFormat.RESULT_FORMAT_VALUE_VECTOR + " is not supported yet"); + } + historicalQuery = + queryToRun.withOverriddenContext(ImmutableMap.of(ScanQuery.CTX_KEY_OUTERMOST, false)); + } + QueryPlus historicalQueryPlus = queryPlus.withQuery(historicalQuery); + Operator inputOp = Operators.toOperator( + input, + historicalQueryPlus); + if (hasLimit) { + final ScanQuery limitedQuery = (ScanQuery) historicalQuery; + if (isGrouped(queryToRun)) { + inputOp = new GroupedScanResultLimitOperator( + queryPlus.fragmentBuilder().context(), + inputOp, + limitedQuery.getScanRowsLimit() + ); + } else { + inputOp = new UngroupedScanResultLimitOperator( + queryPlus.fragmentBuilder().context(), + inputOp, + limitedQuery.getScanRowsLimit(), + limitedQuery.getBatchSize() + ); + } + } + if (hasOffset) { + ScanResultOffsetOperator op = new ScanResultOffsetOperator( + queryPlus.fragmentBuilder().context(), + inputOp, + queryToRun.getScanRowsOffset() + ); + inputOp = op; + } + return Operators.toSequence(inputOp); + } + + private static boolean isGrouped(ScanQuery query) + { + // TODO: Review + return query.getTimeOrder() == ScanQuery.Order.NONE || + !query.getContextBoolean(ScanQuery.CTX_KEY_OUTERMOST, true); + } + + /** + * @see {@link org.apache.druid.query.scan.ScanQueryRunnerFactory.mergeRunners} + */ + private static Sequence runConcatMerge( + final QueryPlus queryPlus, + final Iterable> queryRunners, + final ResponseContext responseContext) + { + List> inputs = new ArrayList<>(); + for (QueryRunner qr : queryRunners) { + inputs.add(Operators.toOperator(qr, queryPlus)); + } + Operator op = ConcatOperator.concatOrNot( + queryPlus.fragmentBuilder().context(), + inputs); + // TODO(paul): The original code applies a limit. Yet, when + // run, the stack shows two limits one top of one another, + // so the limit here seems unnecessary. + // That is, we are doing a concat operation. It does not matter + // if the limit is applied in the concat, or the next operator + // along: in either case, we'll stop reading from upstream when the + // limit is hit. + // + // ScanQuery query = (ScanQuery) queryPlus.getQuery(); + // if (query.isLimited()) { + // op = new ScanResultLimitOperator( + // query.getScanRowsLimit(), + // isGrouped(query), + // query.getBatchSize(), + // op + // ); + // } + return Operators.toSequence(op); + } + + /** + * @see {@link org.apache.druid.query.scan.ScanQueryRunnerFactory.mergeRunners} + * @see {@link org.apache.druid.query.scan.ScanQueryRunnerFactory.nWayMergeAndLimit} + */ + public static Sequence runMerge( + final QueryPlus queryPlus, + final Iterable> queryRunners, + final ResponseContext responseContext) + { + ScanQuery query = (ScanQuery) queryPlus.getQuery(); + ScanQuery.verifyOrderByForNativeExecution(query); + // Note: this variable is effective only when queryContext has a timeout. + // See the comment of ResponseContext.Key.TIMEOUT_AT. + final long timeoutAt = System.currentTimeMillis() + QueryContexts.getTimeout(queryPlus.getQuery()); + responseContext.putTimeoutTime(timeoutAt); + + // TODO: Review + if (query.getTimeOrder() == ScanQuery.Order.NONE) { + // Use normal strategy + return runConcatMerge( + queryPlus, + queryRunners, + responseContext); + } + return null; + } + + /** + * Convert the operator-based scan to that expected by the sequence-based + * query runner. + * + * @see {@link org.apache.druid.query.scan.ScanQueryRunnerFactory.ScanQueryRunner} + */ + public static Sequence runScan( + final QueryPlus queryPlus, + final Segment segment, + final ResponseContext responseContext) + { + if (!(queryPlus.getQuery() instanceof ScanQuery)) { + throw new ISE("Got a [%s] which isn't a %s", queryPlus.getQuery().getClass(), ScanQuery.class); + } + ScanQuery query = (ScanQuery) queryPlus.getQuery(); + ScanQuery.verifyOrderByForNativeExecution((ScanQuery) query); + final Long timeoutAt = responseContext.getTimeoutTime(); + if (timeoutAt == null || timeoutAt == 0L) { + responseContext.putTimeoutTime(JodaUtils.MAX_INSTANT); + } + // TODO (paul): Set the timeout at the overall fragment context level. + return Operators.toSequence( + new ScanQueryOperator( + queryPlus.fragmentBuilder().context(), + query, + segment)); + } + + public static Sequence resultsAsArrays( + QueryPlus queryPlus, + final List fields, + final Sequence resultSequence) + { + FragmentContext context = queryPlus.fragmentBuilder().context(); + Operator inputOp = Operators.toOperator( + context, + resultSequence); + Operator outputOp; + ScanQuery query = (ScanQuery) queryPlus.getQuery(); + switch (query.getResultFormat()) { + case RESULT_FORMAT_LIST: + outputOp = new ScanListToArrayOperator( + context, + new ScanBatchToRowOperator>( + context, + inputOp), + fields); + break; + case RESULT_FORMAT_COMPACTED_LIST: + outputOp = new ScanCompactListToArrayOperator( + context, + new ScanBatchToRowOperator>( + context, + inputOp), + fields); + break; + default: + throw new UOE("Unsupported resultFormat for array-based results: %s", query.getResultFormat()); + } + return Operators.toSequence(outputOp); + } +} diff --git a/processing/src/test/java/org/apache/druid/queryng/operators/BasicOperatorTest.java b/processing/src/test/java/org/apache/druid/queryng/operators/BasicOperatorTest.java index 06768991640f..be6e1045d287 100644 --- a/processing/src/test/java/org/apache/druid/queryng/operators/BasicOperatorTest.java +++ b/processing/src/test/java/org/apache/druid/queryng/operators/BasicOperatorTest.java @@ -24,7 +24,7 @@ import org.apache.druid.queryng.fragment.FragmentContext; import org.apache.druid.queryng.fragment.FragmentRun; import org.apache.druid.queryng.operators.Operator.EofException; -import org.apache.druid.queryng.operators.Operator.RowIterator; +import org.apache.druid.queryng.operators.Operator.ResultIterator; import org.junit.Test; import java.util.Arrays; @@ -168,7 +168,7 @@ public void testSequenceOperator() throws EofException Sequence seq = Operators.toSequence(op); Operator outer = Operators.toOperator(builder, seq); FragmentRun run = builder.run(outer); - RowIterator iter = run.iterator(); + ResultIterator iter = run.iterator(); assertEquals("Mock row 0", iter.next()); assertEquals("Mock row 1", iter.next()); OperatorTests.assertEof(iter); diff --git a/processing/src/test/java/org/apache/druid/queryng/operators/ConcatOperatorTest.java b/processing/src/test/java/org/apache/druid/queryng/operators/ConcatOperatorTest.java index b823b1703c05..724dd622fd61 100644 --- a/processing/src/test/java/org/apache/druid/queryng/operators/ConcatOperatorTest.java +++ b/processing/src/test/java/org/apache/druid/queryng/operators/ConcatOperatorTest.java @@ -21,7 +21,7 @@ import org.apache.druid.queryng.fragment.FragmentContext; import org.apache.druid.queryng.operators.Operator.EofException; -import org.apache.druid.queryng.operators.Operator.RowIterator; +import org.apache.druid.queryng.operators.Operator.ResultIterator; import org.apache.druid.queryng.operators.Operator.State; import org.junit.Test; @@ -142,7 +142,7 @@ public void testClose() throws EofException Operator op = ConcatOperator.concatOrNot( context, Arrays.asList(input1, input2)); - RowIterator iter = op.open(); + ResultIterator iter = op.open(); assertEquals(0, (int) iter.next()); // Only first input has been opened. diff --git a/processing/src/test/java/org/apache/druid/queryng/operators/FragmentTest.java b/processing/src/test/java/org/apache/druid/queryng/operators/FragmentTest.java index aef9e29dcfdd..febc05488e7d 100644 --- a/processing/src/test/java/org/apache/druid/queryng/operators/FragmentTest.java +++ b/processing/src/test/java/org/apache/druid/queryng/operators/FragmentTest.java @@ -37,7 +37,7 @@ import org.apache.druid.queryng.fragment.FragmentRun; import org.apache.druid.queryng.fragment.NullFragmentBuilderFactory; import org.apache.druid.queryng.operators.Operator.EofException; -import org.apache.druid.queryng.operators.Operator.RowIterator; +import org.apache.druid.queryng.operators.Operator.ResultIterator; import org.junit.Test; import java.io.IOException; @@ -350,7 +350,7 @@ public void testFragmentContext() throws EofException assertNull(context.exception()); MockOperator op = MockOperator.ints(builder.context(), 4); FragmentHandle handle = builder.handle(op); - RowIterator iter = handle.run().iterator(); + ResultIterator iter = handle.run().iterator(); assertEquals(FragmentContext.State.RUN, context.state()); // Read from the iterator, just to keep Java 11 happy. assertNotNull(iter.next()); diff --git a/processing/src/test/java/org/apache/druid/queryng/operators/MockOperator.java b/processing/src/test/java/org/apache/druid/queryng/operators/MockOperator.java index 653dde5e1002..c8e495a36abf 100644 --- a/processing/src/test/java/org/apache/druid/queryng/operators/MockOperator.java +++ b/processing/src/test/java/org/apache/druid/queryng/operators/MockOperator.java @@ -53,7 +53,7 @@ public static MockOperator strings(FragmentContext context, int rowCount } @Override - public RowIterator open() + public ResultIterator open() { Preconditions.checkState(state == State.START); state = State.RUN; diff --git a/processing/src/test/java/org/apache/druid/queryng/operators/MockOperatorTest.java b/processing/src/test/java/org/apache/druid/queryng/operators/MockOperatorTest.java index 04dbcbb07d2d..959df55d782a 100644 --- a/processing/src/test/java/org/apache/druid/queryng/operators/MockOperatorTest.java +++ b/processing/src/test/java/org/apache/druid/queryng/operators/MockOperatorTest.java @@ -23,7 +23,7 @@ import org.apache.druid.queryng.fragment.FragmentHandle; import org.apache.druid.queryng.fragment.FragmentRun; import org.apache.druid.queryng.operators.Operator.EofException; -import org.apache.druid.queryng.operators.Operator.RowIterator; +import org.apache.druid.queryng.operators.Operator.ResultIterator; import org.junit.Test; import java.util.List; @@ -43,7 +43,7 @@ public void testMockStringOperator() throws EofException MockOperator op = MockOperator.strings(builder.context(), 2); FragmentHandle handle = builder.handle(op); try (FragmentRun run = handle.run()) { - RowIterator iter = run.iterator(); + ResultIterator iter = run.iterator(); assertEquals("Mock row 0", iter.next()); assertEquals("Mock row 1", iter.next()); OperatorTests.assertEof(iter); diff --git a/processing/src/test/java/org/apache/druid/queryng/operators/OperatorTests.java b/processing/src/test/java/org/apache/druid/queryng/operators/OperatorTests.java index b2e92130e47e..0a7ad1f1f0b4 100644 --- a/processing/src/test/java/org/apache/druid/queryng/operators/OperatorTests.java +++ b/processing/src/test/java/org/apache/druid/queryng/operators/OperatorTests.java @@ -20,13 +20,13 @@ package org.apache.druid.queryng.operators; import org.apache.druid.queryng.operators.Operator.EofException; -import org.apache.druid.queryng.operators.Operator.RowIterator; +import org.apache.druid.queryng.operators.Operator.ResultIterator; import static org.junit.Assert.fail; public class OperatorTests { - public static void assertEof(RowIterator operIter) + public static void assertEof(ResultIterator operIter) { try { operIter.next(); diff --git a/processing/src/test/java/org/apache/druid/queryng/operators/OrderedMergeOperatorTest.java b/processing/src/test/java/org/apache/druid/queryng/operators/OrderedMergeOperatorTest.java index d8a3142e718f..64bf103a490f 100644 --- a/processing/src/test/java/org/apache/druid/queryng/operators/OrderedMergeOperatorTest.java +++ b/processing/src/test/java/org/apache/druid/queryng/operators/OrderedMergeOperatorTest.java @@ -21,7 +21,7 @@ import com.google.common.collect.Ordering; import org.apache.druid.queryng.fragment.FragmentContext; -import org.apache.druid.queryng.operators.Operator.RowIterator; +import org.apache.druid.queryng.operators.Operator.ResultIterator; import org.apache.druid.queryng.operators.Operator.State; import org.apache.druid.queryng.operators.OrderedMergeOperator.Input; import org.junit.Test; @@ -49,7 +49,7 @@ public void testNoInputs() Ordering.natural(), 0, inputs); - RowIterator iter = op.open(); + ResultIterator iter = op.open(); OperatorTests.assertEof(iter); op.close(true); } @@ -67,7 +67,7 @@ public void testEmptyInputs() Ordering.natural(), 2, inputs); - RowIterator iter = op.open(); + ResultIterator iter = op.open(); OperatorTests.assertEof(iter); op.close(true); } @@ -84,7 +84,7 @@ public void testOneInput() Ordering.natural(), 1, inputs); - RowIterator iter = op.open(); + ResultIterator iter = op.open(); List results = Iterators.toList(iter); op.close(true); assertEquals(Arrays.asList(0, 1, 2), results); @@ -103,7 +103,7 @@ public void testTwoInputs() Ordering.natural(), 2, inputs); - RowIterator iter = op.open(); + ResultIterator iter = op.open(); List results = Iterators.toList(iter); op.close(true); assertEquals(Arrays.asList(0, 0, 1, 1, 2, 2, 3, 4), results); @@ -124,7 +124,7 @@ public void testClose() Ordering.natural(), 2, inputs); - RowIterator iter = op.open(); + ResultIterator iter = op.open(); List results = Iterators.toList(iter); assertEquals(Arrays.asList(0, 0, 1, 1), results); diff --git a/processing/src/test/java/org/apache/druid/queryng/operators/PushBackOperatorTest.java b/processing/src/test/java/org/apache/druid/queryng/operators/PushBackOperatorTest.java index 2cc2a67b54b5..b25ad5328d00 100644 --- a/processing/src/test/java/org/apache/druid/queryng/operators/PushBackOperatorTest.java +++ b/processing/src/test/java/org/apache/druid/queryng/operators/PushBackOperatorTest.java @@ -21,7 +21,7 @@ import org.apache.druid.queryng.fragment.FragmentContext; import org.apache.druid.queryng.operators.Operator.EofException; -import org.apache.druid.queryng.operators.Operator.RowIterator; +import org.apache.druid.queryng.operators.Operator.ResultIterator; import org.junit.Test; import java.util.Arrays; @@ -58,7 +58,7 @@ public void testPush() throws EofException FragmentContext context = FragmentContext.defaultContext(); Operator input = MockOperator.ints(context, 2); PushBackOperator op = new PushBackOperator(context, input); - RowIterator iter = op.open(); + ResultIterator iter = op.open(); Integer item = iter.next(); op.push(item); List results = Operators.toList(op); @@ -71,7 +71,7 @@ public void testInitialPush() throws EofException { FragmentContext context = FragmentContext.defaultContext(); Operator input = MockOperator.ints(context, 2); - RowIterator iter = input.open(); + ResultIterator iter = input.open(); Integer item = iter.next(); PushBackOperator op = new PushBackOperator(context, input, iter, item); List results = Operators.toList(op); diff --git a/processing/src/test/java/org/apache/druid/queryng/operators/scan/MockScanResultReader.java b/processing/src/test/java/org/apache/druid/queryng/operators/scan/MockScanResultReader.java new file mode 100644 index 000000000000..584b7c3c368f --- /dev/null +++ b/processing/src/test/java/org/apache/druid/queryng/operators/scan/MockScanResultReader.java @@ -0,0 +1,206 @@ +/* + * 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.queryng.operators.scan; + +import com.google.common.annotations.VisibleForTesting; +import com.google.common.base.Preconditions; +import org.apache.druid.java.util.common.StringUtils; +import org.apache.druid.query.scan.ScanQuery.ResultFormat; +import org.apache.druid.query.scan.ScanResultValue; +import org.apache.druid.queryng.fragment.FragmentContext; +import org.apache.druid.queryng.operators.Operator.IterableOperator; +import org.apache.druid.queryng.operators.Operators; +import org.apache.druid.segment.column.ColumnHolder; +import org.joda.time.Interval; + +import java.time.Duration; +import java.time.Instant; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +public class MockScanResultReader implements IterableOperator +{ + public String segmentId = "mock-segment"; + protected final List columns; + public ResultFormat resultFormat; + private final int targetCount; + private final int batchSize; + private final long msPerRow; + private long nextTs; + private int rowCount; + @SuppressWarnings("unused") + private int batchCount; + + /** + * State allows tests to verify that the operator protocol + * was followed. Not really necessary here functionally, so this + * is just a test tool. + */ + public State state = State.START; + + public MockScanResultReader( + FragmentContext context, + int columnCount, + int targetCount, + int batchSize, + Interval interval) + { + this( + context, + columnCount, + targetCount, + batchSize, + interval, + ResultFormat.RESULT_FORMAT_COMPACTED_LIST); + } + + public MockScanResultReader( + FragmentContext context, + int columnCount, + int targetCount, + int batchSize, + Interval interval, + ResultFormat resultFormat) + { + this.columns = new ArrayList<>(columnCount); + if (columnCount > 0) { + columns.add(ColumnHolder.TIME_COLUMN_NAME); + } + for (int i = 1; i < columnCount; i++) { + columns.add("Column" + Integer.toString(i)); + } + this.targetCount = targetCount; + this.batchSize = batchSize; + this.resultFormat = resultFormat; + if (targetCount == 0) { + this.msPerRow = 0; + } else { + this.msPerRow = Math.toIntExact(interval.toDurationMillis() / targetCount); + } + this.nextTs = interval.getStartMillis(); + context.register(this); + } + + public static Interval interval(int offset) + { + Duration grain = Duration.ofMinutes(1); + Instant base = Instant.parse("2021-10-24T00:00:00Z"); + Duration grainOffset = grain.multipliedBy(offset); + Instant start = base.plus(grainOffset); + return new Interval(start.toEpochMilli(), start.plus(grain).toEpochMilli()); + } + + @Override + public ResultIterator open() + { + state = State.RUN; + return this; + } + + @Override + public ScanResultValue next() throws EofException + { + Preconditions.checkState(state == State.RUN); + if (rowCount >= targetCount) { + throw Operators.eof(); + } + int n = Math.min(targetCount - rowCount, batchSize); + Object batch; + if (resultFormat == ResultFormat.RESULT_FORMAT_COMPACTED_LIST) { + batch = compactBatch(n); + } else { + batch = listBatch(n); + } + batchCount++; + return new ScanResultValue(segmentId, columns, batch); + } + + private Object compactBatch(int n) + { + List> batch = new ArrayList<>(n); + for (int i = 0; i < n; i++) { + List values = new ArrayList<>(columns.size()); + if (!columns.isEmpty()) { + values.add(nextTs); + } + for (int j = 1; j < columns.size(); j++) { + values.add(StringUtils.format("Value %d.%d", rowCount, j)); + } + batch.add(values); + rowCount++; + nextTs += msPerRow; + } + return batch; + } + + private Object listBatch(int n) + { + List> batch = new ArrayList<>(n); + for (int i = 0; i < n; i++) { + Map values = new HashMap<>(columns.size()); + if (!columns.isEmpty()) { + values.put(ColumnHolder.TIME_COLUMN_NAME, nextTs); + } + for (int j = 1; j < columns.size(); j++) { + values.put(columns.get(j), StringUtils.format("Value %d.%d", rowCount, j)); + } + batch.add(values); + rowCount++; + nextTs += msPerRow; + } + return batch; + } + + @Override + public void close(boolean cascade) + { + state = State.CLOSED; + } + + @VisibleForTesting + public static long getTime(List row) + { + return (Long) row.get(0); + } + + @VisibleForTesting + public static long getFirstTime(Object row) + { + ScanResultValue value = (ScanResultValue) row; + List> events = value.getRows(); + if (events.isEmpty()) { + return 0; + } + return getTime(events.get(0)); + } + + @VisibleForTesting + public static long getLastTime(Object row) + { + ScanResultValue value = (ScanResultValue) row; + List> events = value.getRows(); + if (events.isEmpty()) { + return 0; + } + return getTime(events.get(events.size() - 1)); + } +} diff --git a/processing/src/test/java/org/apache/druid/queryng/operators/scan/ScanQueryOperatorsTest.java b/processing/src/test/java/org/apache/druid/queryng/operators/scan/ScanQueryOperatorsTest.java new file mode 100644 index 000000000000..cf97ddf1eae3 --- /dev/null +++ b/processing/src/test/java/org/apache/druid/queryng/operators/scan/ScanQueryOperatorsTest.java @@ -0,0 +1,246 @@ +/* + * 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.queryng.operators.scan; + +import com.google.common.base.Strings; +import org.apache.druid.java.util.common.StringUtils; +import org.apache.druid.query.scan.ScanQuery.ResultFormat; +import org.apache.druid.query.scan.ScanResultValue; +import org.apache.druid.queryng.fragment.FragmentBuilder; +import org.apache.druid.queryng.fragment.FragmentContext; +import org.apache.druid.queryng.fragment.FragmentRun; +import org.apache.druid.queryng.operators.Iterators; +import org.apache.druid.queryng.operators.Operator; +import org.apache.druid.queryng.operators.Operator.EofException; +import org.apache.druid.queryng.operators.Operator.ResultIterator; +import org.apache.druid.queryng.operators.OperatorTests; +import org.apache.druid.segment.column.ColumnHolder; +import org.junit.Test; + +import java.util.List; +import java.util.Map; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertTrue; + +public class ScanQueryOperatorsTest +{ + private MockScanResultReader scan( + FragmentContext context, + int columnCount, + int rowCount, + int batchSize) + { + return new MockScanResultReader( + context, + columnCount, + rowCount, + batchSize, + MockScanResultReader.interval(0)); + } + + private MockScanResultReader scan( + FragmentContext context, + int columnCount, + int rowCount, + int batchSize, + ResultFormat rowFormat) + { + return new MockScanResultReader( + context, + columnCount, + rowCount, + batchSize, + MockScanResultReader.interval(0), + rowFormat); + } + + // Tests for the mock reader used to power tests without the overhead + // of using an actual scan operator. The parent operators don't know + // the difference. + @Test + public void testMockReaderNull() + { + FragmentBuilder builder = FragmentBuilder.defaultBuilder(); + Operator op = scan(builder.context(), 0, 0, 3); + FragmentRun run = builder.run(op); + ResultIterator iter = run.iterator(); + OperatorTests.assertEof(iter); + run.close(); + } + + @Test + public void testMockReaderEmpty() throws EofException + { + FragmentBuilder builder = FragmentBuilder.defaultBuilder(); + MockScanResultReader scan = scan(builder.context(), 0, 1, 3); + assertFalse(Strings.isNullOrEmpty(scan.segmentId)); + FragmentRun run = builder.run(scan); + ResultIterator iter = run.iterator(); + ScanResultValue value = iter.next(); + assertTrue(value.getColumns().isEmpty()); + List> events = value.getRows(); + assertEquals(1, events.size()); + assertTrue(events.get(0).isEmpty()); + OperatorTests.assertEof(iter); + run.close(); + } + + @Test + public void testMockReader() + { + FragmentBuilder builder = FragmentBuilder.defaultBuilder(); + Operator scan = scan(builder.context(), 3, 10, 4); + FragmentRun run = builder.run(scan); + ResultIterator iter = run.iterator(); + int rowCount = 0; + for (ScanResultValue value : Iterators.toIterable(iter)) { + assertEquals(3, value.getColumns().size()); + assertEquals(ColumnHolder.TIME_COLUMN_NAME, value.getColumns().get(0)); + assertEquals("Column1", value.getColumns().get(1)); + assertEquals("Column2", value.getColumns().get(2)); + List> events = value.getRows(); + assertTrue(events.size() <= 4); + long prevTs = 0; + for (List row : events) { + for (int i = 0; i < 3; i++) { + Object colValue = row.get(i); + assertNotNull(row.get(i)); + if (i == 0) { + assertTrue(colValue instanceof Long); + long ts = (Long) colValue; + assertTrue(ts > prevTs); + prevTs = ts; + } else { + assertTrue(colValue instanceof String); + } + } + rowCount++; + } + } + assertEquals(10, rowCount); + run.close(); + } + + /** + * Test the offset operator for various numbers of input rows, spread across + * multiple batches. Tests for offsets that fall on a batch boundary + * and within a batch. + */ + @Test + public void testOffset() + { + final int totalRows = 10; + for (int offset = 1; offset < 2 * totalRows; offset++) { + FragmentBuilder builder = FragmentBuilder.defaultBuilder(); + Operator scan = scan(builder.context(), 3, totalRows, 4); + Operator root = + new ScanResultOffsetOperator(builder.context(), scan, offset); + int rowCount = 0; + final String firstVal = StringUtils.format("Value %d.1", offset); + FragmentRun run = builder.run(root); + ResultIterator iter = run.iterator(); + for (ScanResultValue row : Iterators.toIterable(iter)) { + ScanResultValue value = row; + List> events = value.getRows(); + if (rowCount == 0) { + assertEquals(firstVal, events.get(0).get(1)); + } + rowCount += events.size(); + } + assertEquals(Math.max(0, totalRows - offset), rowCount); + run.close(); + } + } + + /** + * Test the limit operator for various numbers of input rows, spread across + * multiple batches. Tests for limits that fall on a batch boundary + * and within a batch. + */ + @Test + public void testGroupedLimit() + { + final int totalRows = 10; + for (int limit = 0; limit < totalRows + 1; limit++) { + FragmentBuilder builder = FragmentBuilder.defaultBuilder(); + Operator scan = scan(builder.context(), 3, totalRows, 4); + Operator root = + new GroupedScanResultLimitOperator(builder.context(), scan, limit); + FragmentRun run = builder.run(root); + ResultIterator iter = run.iterator(); + int rowCount = 0; + for (ScanResultValue row : Iterators.toIterable(iter)) { + ScanResultValue value = row; + rowCount += value.rowCount(); + } + assertEquals(Math.min(totalRows, limit), rowCount); + run.close(); + } + } + + @Test + public void testUngroupedLimit() + { + final int totalRows = 10; + for (int limit = 0; limit < totalRows + 1; limit++) { + FragmentBuilder builder = FragmentBuilder.defaultBuilder(); + Operator scan = scan(builder.context(), 3, totalRows, 1); + Operator root = + new UngroupedScanResultLimitOperator(builder.context(), scan, limit, 4); + FragmentRun run = builder.run(root); + ResultIterator iter = run.iterator(); + int rowCount = 0; + for (ScanResultValue row : Iterators.toIterable(iter)) { + ScanResultValue value = row; + rowCount += value.rowCount(); + } + assertEquals(Math.min(totalRows, limit), rowCount); + run.close(); + } + } + + @Test + public void testBatchToRow() + { + { + FragmentBuilder builder = FragmentBuilder.defaultBuilder(); + MockScanResultReader scan = scan(builder.context(), 3, 25, 4, ResultFormat.RESULT_FORMAT_COMPACTED_LIST); + Operator> op = + new ScanBatchToRowOperator>(builder.context(), scan); + Operator root = + new ScanCompactListToArrayOperator(builder.context(), op, scan.columns); + List results = builder.run(root).toList(); + assertEquals(25, results.size()); + } + { + FragmentBuilder builder = FragmentBuilder.defaultBuilder(); + MockScanResultReader scan = scan(builder.context(), 3, 25, 4, ResultFormat.RESULT_FORMAT_LIST); + Operator> op = + new ScanBatchToRowOperator>(builder.context(), scan); + Operator root = + new ScanListToArrayOperator(builder.context(), op, scan.columns); + List results = builder.run(root).toList(); + assertEquals(25, results.size()); + } + } +} From 2eb249018ab84ebd486c3c7585c7e933aa58dc41 Mon Sep 17 00:00:00 2001 From: Paul Rogers Date: Sat, 18 Jun 2022 21:07:33 -0700 Subject: [PATCH 09/11] Starter set of scan query operators --- .../apache/druid/query/QueryToolChest.java | 6 +++ .../query/scan/ScanQueryLimitRowIterator.java | 2 +- .../query/scan/ScanQueryQueryToolChest.java | 21 +++++++- .../query/scan/ScanQueryRunnerFactory.java | 27 ++++++++--- .../druid/queryng/config/QueryNGConfig.java | 30 +++++++----- .../fragment/FragmentBuilderFactoryImpl.java | 8 +--- .../fragment/TestFragmentBuilderFactory.java | 48 +++++++++++++++++++ .../druid/queryng/guice/QueryNGModule.java | 5 +- .../scan/ScanCompactListToArrayOperator.java | 2 +- .../druid/queryng/operators/FragmentTest.java | 44 ++++++++++++++++- .../apache/druid/server/QueryLifecycle.java | 14 ++++-- .../druid/server/QueryLifecycleFactory.java | 7 ++- .../druid/guice/QueryableModuleTest.java | 2 + .../druid/server/QueryLifecycleTest.java | 37 +------------- .../druid/server/QueryResourceTest.java | 25 ++++++---- .../java/org/apache/druid/cli/CliBroker.java | 2 + .../org/apache/druid/cli/CliHistorical.java | 2 + .../java/org/apache/druid/cli/CliIndexer.java | 2 + .../java/org/apache/druid/cli/CliPeon.java | 2 + .../org/apache/druid/cli/CreateTables.java | 2 + .../org/apache/druid/cli/DumpSegment.java | 2 + .../org/apache/druid/cli/ExportMetadata.java | 4 +- .../org/apache/druid/cli/ResetCluster.java | 2 + .../apache/druid/cli/ValidateSegments.java | 2 + .../cli/validate/DruidJsonValidator.java | 2 + .../druid/sql/calcite/util/CalciteTests.java | 5 +- 26 files changed, 225 insertions(+), 80 deletions(-) create mode 100644 processing/src/main/java/org/apache/druid/queryng/fragment/TestFragmentBuilderFactory.java diff --git a/processing/src/main/java/org/apache/druid/query/QueryToolChest.java b/processing/src/main/java/org/apache/druid/query/QueryToolChest.java index ee400f814fb8..27635066543f 100644 --- a/processing/src/main/java/org/apache/druid/query/QueryToolChest.java +++ b/processing/src/main/java/org/apache/druid/query/QueryToolChest.java @@ -330,4 +330,10 @@ public Sequence resultsAsArrays(QueryType query, Sequence { throw new UOE("Query type '%s' does not support returning results as arrays", query.getType()); } + + @SuppressWarnings("unchecked") + public Sequence resultsAsArrays(QueryPlus query, Sequence resultSequence) + { + return resultsAsArrays((QueryType) query.getQuery(), resultSequence); + } } diff --git a/processing/src/main/java/org/apache/druid/query/scan/ScanQueryLimitRowIterator.java b/processing/src/main/java/org/apache/druid/query/scan/ScanQueryLimitRowIterator.java index 6a081b654865..1e81b36bc29d 100644 --- a/processing/src/main/java/org/apache/druid/query/scan/ScanQueryLimitRowIterator.java +++ b/processing/src/main/java/org/apache/druid/query/scan/ScanQueryLimitRowIterator.java @@ -68,7 +68,7 @@ public class ScanQueryLimitRowIterator implements CloseableIterator historicalQuery = queryPlus.getQuery().withOverriddenContext(ImmutableMap.of(ScanQuery.CTX_KEY_OUTERMOST, false)); - Sequence baseSequence = baseRunner.run(QueryPlus.wrap(historicalQuery), responseContext); + Sequence baseSequence = baseRunner.run(queryPlus.withQuery(historicalQuery), responseContext); this.yielder = baseSequence.toYielder( null, new YieldingAccumulator() diff --git a/processing/src/main/java/org/apache/druid/query/scan/ScanQueryQueryToolChest.java b/processing/src/main/java/org/apache/druid/query/scan/ScanQueryQueryToolChest.java index 70bba3a77714..930a6abf0b9a 100644 --- a/processing/src/main/java/org/apache/druid/query/scan/ScanQueryQueryToolChest.java +++ b/processing/src/main/java/org/apache/druid/query/scan/ScanQueryQueryToolChest.java @@ -32,9 +32,12 @@ import org.apache.druid.query.GenericQueryMetricsFactory; import org.apache.druid.query.Query; import org.apache.druid.query.QueryMetrics; +import org.apache.druid.query.QueryPlus; import org.apache.druid.query.QueryRunner; import org.apache.druid.query.QueryToolChest; import org.apache.druid.query.aggregation.MetricManipulationFn; +import org.apache.druid.queryng.config.QueryNGConfig; +import org.apache.druid.queryng.planner.ScanPlanner; import org.apache.druid.segment.VirtualColumn; import org.apache.druid.segment.column.ColumnType; import org.apache.druid.segment.column.RowSignature; @@ -66,6 +69,9 @@ public ScanQueryQueryToolChest( public QueryRunner mergeResults(final QueryRunner runner) { return (queryPlus, responseContext) -> { + if (QueryNGConfig.enabledFor(queryPlus)) { + return ScanPlanner.runLimitAndOffset(queryPlus, runner, responseContext, scanQueryConfig); + } final ScanQuery originalQuery = ((ScanQuery) (queryPlus.getQuery())); ScanQuery.verifyOrderByForNativeExecution(originalQuery); @@ -213,7 +219,7 @@ public Sequence resultsAsArrays(final ScanQuery query, final Sequence< // Uh oh... mismatch in expected and actual field count. I don't think this should happen, so let's // throw an exception. If this really does happen, and there's a good reason for it, then we should remap // the result row here. - throw new ISE("Mismatch in expected[%d] vs actual[%s] field count", fields.size(), row.size()); + throw new ISE("Mismatch in expected [%d] vs actual [%s] field count", fields.size(), row.size()); } }; break; @@ -230,4 +236,17 @@ public Sequence resultsAsArrays(final ScanQuery query, final Sequence< } ); } + + @Override + @SuppressWarnings("unchecked") + public Sequence resultsAsArrays(QueryPlus queryPlus, Sequence resultSequence) + { + ScanQuery query = (ScanQuery) queryPlus.getQuery(); + if (QueryNGConfig.enabledFor(queryPlus)) { + final List fields = resultArraySignature(query).getColumnNames(); + return ScanPlanner.resultsAsArrays(queryPlus, fields, resultSequence); + } else { + return resultsAsArrays(query, resultSequence); + } + } } diff --git a/processing/src/main/java/org/apache/druid/query/scan/ScanQueryRunnerFactory.java b/processing/src/main/java/org/apache/druid/query/scan/ScanQueryRunnerFactory.java index 8aec07679b3a..8a60844583ff 100644 --- a/processing/src/main/java/org/apache/druid/query/scan/ScanQueryRunnerFactory.java +++ b/processing/src/main/java/org/apache/druid/query/scan/ScanQueryRunnerFactory.java @@ -46,6 +46,8 @@ import org.apache.druid.query.spec.MultipleSpecificSegmentSpec; import org.apache.druid.query.spec.QuerySegmentSpec; import org.apache.druid.query.spec.SpecificSegmentSpec; +import org.apache.druid.queryng.config.QueryNGConfig; +import org.apache.druid.queryng.planner.ScanPlanner; import org.apache.druid.segment.Segment; import org.joda.time.Interval; @@ -87,8 +89,18 @@ public QueryRunner mergeRunners( final Iterable> queryRunners ) { - // in single thread and in Jetty thread instead of processing thread + // In single thread and in Jetty thread instead of processing thread return (queryPlus, responseContext) -> { + if (QueryNGConfig.enabledFor(queryPlus)) { + Sequence results = ScanPlanner.runMerge( + queryPlus, + queryRunners, + responseContext); + if (results != null) { + return results; + } + } + ScanQuery query = (ScanQuery) queryPlus.getQuery(); ScanQuery.verifyOrderByForNativeExecution(query); @@ -147,7 +159,7 @@ public QueryRunner mergeRunners( ((SinkQueryRunners) queryRunners).runnerIntervalMappingIterator() .forEachRemaining(intervalsAndRunnersOrdered::add); } else { - throw new ISE("Number of segment descriptors does not equal number of " + throw new ISE("Number of segment descriptors does not equal the number of " + "query runners...something went wrong!"); } @@ -295,7 +307,7 @@ List getIntervalsFromSpecificQuerySpec(QuerySegmentSpec spec) } else { throw new UOE( "Time-ordering on scan queries is only supported for queries with segment specs " - + "of type MultipleSpecificSegmentSpec or SpecificSegmentSpec...a [%s] was received instead.", + + "of type MultipleSpecificSegmentSpec or SpecificSegmentSpec. A [%s] was received instead.", spec.getClass().getSimpleName() ); } @@ -310,8 +322,8 @@ Sequence nWayMergeAndLimit( ) { // Starting from the innermost Sequences.map: - // (1) Deaggregate each ScanResultValue returned by the query runners - // (2) Combine the deaggregated ScanResultValues into a single sequence + // (1) Disaggregate each ScanResultValue returned by the query runners + // (2) Combine the disaggregated ScanResultValues into a single sequence // (3) Create a sequence of results from each runner in the group and flatmerge based on timestamp // (4) Create a sequence of results from each runner group // (5) Join all the results into a single sequence @@ -361,11 +373,14 @@ public ScanQueryRunner(ScanQueryEngine engine, Segment segment) @Override public Sequence run(QueryPlus queryPlus, ResponseContext responseContext) { + if (QueryNGConfig.enabledFor(queryPlus)) { + return ScanPlanner.runScan(queryPlus, segment, responseContext); + } + Query query = queryPlus.getQuery(); if (!(query instanceof ScanQuery)) { throw new ISE("Got a [%s] which isn't a %s", query.getClass(), ScanQuery.class); } - ScanQuery.verifyOrderByForNativeExecution((ScanQuery) query); // it happens in unit tests diff --git a/processing/src/main/java/org/apache/druid/queryng/config/QueryNGConfig.java b/processing/src/main/java/org/apache/druid/queryng/config/QueryNGConfig.java index 5ca3c9ff119c..1696bbfa88d5 100644 --- a/processing/src/main/java/org/apache/druid/queryng/config/QueryNGConfig.java +++ b/processing/src/main/java/org/apache/druid/queryng/config/QueryNGConfig.java @@ -32,22 +32,25 @@ public class QueryNGConfig @SuppressWarnings("unused") // To be used later public static final String CONFIG_ROOT = "druid.queryng"; + public static final String CONTEXT_VAR = "queryng"; + /** * Whether the engine is enabled. It is disabled by default. */ @JsonProperty("enabled") private boolean enabled; - public static final String CONTEXT_VAR = "queryng"; + @JsonProperty("requireContext") + private boolean requireContext = true; /** * Create an instance for testing. */ - @SuppressWarnings("unused") // To be used later - public static QueryNGConfig create(boolean enabled) + public static QueryNGConfig create(boolean enabled, boolean requireContext) { QueryNGConfig config = new QueryNGConfig(); config.enabled = enabled; + config.requireContext = requireContext; return config; } @@ -57,26 +60,29 @@ public boolean enabled() } /** - * Determine if Query NG should be enabled for the given query; - * that is, if the query should have a fragment context attached. + * Determine if Query NG should be enabled for the given query. Only scan + * queries are currently supported. For safety, the default config also + * requires that a context variable be set to enable the operatore-based + * engine. However, the configuration can skip the context check. A present, + * the skip-context option is primarily for testing. + * that is, if the query should have a fragment context attached. * At present, Query NG is enabled if the query is a scan query and * the query has the "queryng" context variable set. The caller * should already have checked if the Query NG engine is enabled * globally. If Query NG is enabled for a query, then the caller * will attach a fragment context to the query's QueryPlus. */ - public static boolean isEnabled(Query query) + public boolean isEnabled(Query query) { // Query has to be of the currently-supported type - if (!(query instanceof ScanQuery)) { - return false; - } - return query.getContextBoolean(CONTEXT_VAR, false); + return enabled + && (query instanceof ScanQuery) + && (!requireContext || query.getContextBoolean(CONTEXT_VAR, false)); } /** - * Determine if the Query NG (operator-based) engine is enabled for the - * given query (given as a QueryPlus). Query NG is enabled if the QueryPlus + * Determine if the Query NG (operator-based) engine is enabled for the given + * query (given as a QueryPlus). Query NG is enabled if the QueryPlus * includes the fragment context needed by the Query NG engine. */ public static boolean enabledFor(final QueryPlus queryPlus) diff --git a/processing/src/main/java/org/apache/druid/queryng/fragment/FragmentBuilderFactoryImpl.java b/processing/src/main/java/org/apache/druid/queryng/fragment/FragmentBuilderFactoryImpl.java index ee7a33e792d5..396612fbff62 100644 --- a/processing/src/main/java/org/apache/druid/queryng/fragment/FragmentBuilderFactoryImpl.java +++ b/processing/src/main/java/org/apache/druid/queryng/fragment/FragmentBuilderFactoryImpl.java @@ -46,12 +46,8 @@ public FragmentBuilder create( final Query query, final ResponseContext responseContext) { - // Engine has to be enabled - if (!config.enabled()) { - return null; - } - // Client must explicitly ask for the engine - if (!QueryNGConfig.isEnabled(query)) { + // Config imposes a number of obstacles. + if (!config.isEnabled(query)) { return null; } // Only then do we create a fragment builder which, implicitly, diff --git a/processing/src/main/java/org/apache/druid/queryng/fragment/TestFragmentBuilderFactory.java b/processing/src/main/java/org/apache/druid/queryng/fragment/TestFragmentBuilderFactory.java new file mode 100644 index 000000000000..9fdfe34658f3 --- /dev/null +++ b/processing/src/main/java/org/apache/druid/queryng/fragment/TestFragmentBuilderFactory.java @@ -0,0 +1,48 @@ +/* + * 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.queryng.fragment; + +import org.apache.druid.query.Query; +import org.apache.druid.query.context.ResponseContext; +import org.apache.druid.query.scan.ScanQuery; +import org.apache.druid.queryng.config.QueryNGConfig; + +/** + * Test version of the fragment factory which enables Query NG only if + * the {@code druid.queryng.enable} system property is set, and then, + * only for scan queries. + */ +public class TestFragmentBuilderFactory implements FragmentBuilderFactory +{ + private static final String ENABLED_KEY = QueryNGConfig.CONFIG_ROOT + ".enabled"; + private static final boolean ENABLED = Boolean.parseBoolean(System.getProperty(ENABLED_KEY)); + + @Override + public FragmentBuilder create(Query query, ResponseContext responseContext) + { + //if (!ENABLED) { + // return null; + //} + if (!(query instanceof ScanQuery)) { + return null; + } + return new FragmentBuilderImpl(query.getId(), 0, responseContext); + } +} diff --git a/processing/src/main/java/org/apache/druid/queryng/guice/QueryNGModule.java b/processing/src/main/java/org/apache/druid/queryng/guice/QueryNGModule.java index 1d1c351da7e9..e44fcd806a5b 100644 --- a/processing/src/main/java/org/apache/druid/queryng/guice/QueryNGModule.java +++ b/processing/src/main/java/org/apache/druid/queryng/guice/QueryNGModule.java @@ -25,7 +25,7 @@ import org.apache.druid.guice.LazySingleton; import org.apache.druid.queryng.config.QueryNGConfig; import org.apache.druid.queryng.fragment.FragmentBuilderFactory; -import org.apache.druid.queryng.fragment.FragmentBuilderFactoryImpl; +import org.apache.druid.queryng.fragment.NullFragmentBuilderFactory; /** * Configure the "shim" version of the NG query engine which entails @@ -42,7 +42,8 @@ public void configure(Binder binder) JsonConfigProvider.bind(binder, QueryNGConfig.CONFIG_ROOT, QueryNGConfig.class); binder .bind(FragmentBuilderFactory.class) - .to(FragmentBuilderFactoryImpl.class) + // Query NG disabled in production nodes for now. + .to(NullFragmentBuilderFactory.class) .in(LazySingleton.class); } } diff --git a/processing/src/main/java/org/apache/druid/queryng/operators/scan/ScanCompactListToArrayOperator.java b/processing/src/main/java/org/apache/druid/queryng/operators/scan/ScanCompactListToArrayOperator.java index 48b07690ee33..d30a8b09d991 100644 --- a/processing/src/main/java/org/apache/druid/queryng/operators/scan/ScanCompactListToArrayOperator.java +++ b/processing/src/main/java/org/apache/druid/queryng/operators/scan/ScanCompactListToArrayOperator.java @@ -62,7 +62,7 @@ public Object[] next() throws EofException // happen, and there's a good reason for it, then we should remap // the result row here. throw new ISE( - "Mismatch in expected [%d] vs actual [%s] field count", + "Mismatch in expected [%d] vs. actual [%s] field count", fields.size(), row.size()); } diff --git a/processing/src/test/java/org/apache/druid/queryng/operators/FragmentTest.java b/processing/src/test/java/org/apache/druid/queryng/operators/FragmentTest.java index febc05488e7d..0b0a170d03c6 100644 --- a/processing/src/test/java/org/apache/druid/queryng/operators/FragmentTest.java +++ b/processing/src/test/java/org/apache/druid/queryng/operators/FragmentTest.java @@ -19,14 +19,17 @@ package org.apache.druid.queryng.operators; +import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import org.apache.druid.java.util.common.ISE; +import org.apache.druid.java.util.common.Intervals; import org.apache.druid.java.util.common.guava.Sequence; import org.apache.druid.java.util.common.guava.SequenceTestHelper; import org.apache.druid.query.Druids; import org.apache.druid.query.Query; import org.apache.druid.query.QueryPlus; import org.apache.druid.query.context.ResponseContext; +import org.apache.druid.query.spec.MultipleIntervalSegmentSpec; import org.apache.druid.queryng.config.QueryNGConfig; import org.apache.druid.queryng.fragment.FragmentBuilder; import org.apache.druid.queryng.fragment.FragmentBuilderFactory; @@ -294,6 +297,43 @@ public void testRunEmptyHandleAsSequence() SequenceTestHelper.testAccumulation("empty", seq, Collections.emptyList()); } + @Test + public void testConfig() + { + Query scanQuery = new Druids.ScanQueryBuilder() + .dataSource("foo") + .eternityInterval() + .build(); + Query scanQueryWithContext = scanQuery.withOverriddenContext( + ImmutableMap.of(QueryNGConfig.CONTEXT_VAR, true)); + Query otherQuery = Druids.newTimeseriesQueryBuilder() + .dataSource("foo") + .intervals(new MultipleIntervalSegmentSpec( + ImmutableList.of(Intervals.ETERNITY))) + .build(); + + // Completely diabled. + QueryNGConfig config = QueryNGConfig.create(false, false); + assertFalse(config.enabled()); + assertFalse(config.isEnabled(scanQuery)); + assertFalse(config.isEnabled(scanQueryWithContext)); + assertFalse(config.isEnabled(otherQuery)); + + // Enabled, but only for scan. + config = QueryNGConfig.create(true, false); + assertTrue(config.enabled()); + assertTrue(config.isEnabled(scanQuery)); + assertTrue(config.isEnabled(scanQueryWithContext)); + assertFalse(config.isEnabled(otherQuery)); + + // Enabled, but only for scan, and only if requested in context. + config = QueryNGConfig.create(true, true); + assertTrue(config.enabled()); + assertFalse(config.isEnabled(scanQuery)); + assertTrue(config.isEnabled(scanQueryWithContext)); + assertFalse(config.isEnabled(otherQuery)); + } + @Test public void testFactory() { @@ -303,13 +343,13 @@ public void testFactory() .build(); // Operators blocked by query: no gating context variable - QueryNGConfig enableConfig = QueryNGConfig.create(true); + QueryNGConfig enableConfig = QueryNGConfig.create(true, true); assertTrue(enableConfig.enabled()); FragmentBuilderFactory enableFactory = new FragmentBuilderFactoryImpl(enableConfig); assertNull(enableFactory.create(query, ResponseContext.createEmpty())); FragmentBuilderFactory nullFactory = new NullFragmentBuilderFactory(); - QueryNGConfig disableConfig = QueryNGConfig.create(false); + QueryNGConfig disableConfig = QueryNGConfig.create(false, false); assertFalse(disableConfig.enabled()); FragmentBuilderFactory disableFactory = new FragmentBuilderFactoryImpl(disableConfig); assertNull(disableFactory.create(query, ResponseContext.createEmpty())); diff --git a/server/src/main/java/org/apache/druid/server/QueryLifecycle.java b/server/src/main/java/org/apache/druid/server/QueryLifecycle.java index 1bdba4515193..700024d4f6b7 100644 --- a/server/src/main/java/org/apache/druid/server/QueryLifecycle.java +++ b/server/src/main/java/org/apache/druid/server/QueryLifecycle.java @@ -46,6 +46,8 @@ import org.apache.druid.query.QueryToolChest; import org.apache.druid.query.QueryToolChestWarehouse; import org.apache.druid.query.context.ResponseContext; +import org.apache.druid.queryng.fragment.FragmentBuilder; +import org.apache.druid.queryng.fragment.FragmentBuilderFactory; import org.apache.druid.server.QueryResource.ResourceIOReaderWriter; import org.apache.druid.server.log.RequestLogger; import org.apache.druid.server.security.Access; @@ -69,8 +71,9 @@ import java.util.concurrent.TimeUnit; /** - * Class that helps a Druid server (broker, historical, etc) manage the lifecycle of a query that it is handling. It - * ensures that a query goes through the following stages, in the proper order: + * Class that helps a Druid server (broker, historical, etc) manage the + * lifecycle of a query that it is handling. It ensures that a query goes + * through the following stages, in the proper order: * *
      *
    1. Initialization ({@link #initialize(Query)})
      2. @@ -93,6 +96,7 @@ public class QueryLifecycle private final AuthorizerMapper authorizerMapper; private final DefaultQueryConfig defaultQueryConfig; private final AuthConfig authConfig; + private final FragmentBuilderFactory fragmentContextFactory; private final long startMs; private final long startNs; @@ -112,6 +116,7 @@ public QueryLifecycle( final AuthorizerMapper authorizerMapper, final DefaultQueryConfig defaultQueryConfig, final AuthConfig authConfig, + final FragmentBuilderFactory fragmentContextFactory, final long startMs, final long startNs ) @@ -124,11 +129,11 @@ public QueryLifecycle( this.authorizerMapper = authorizerMapper; this.defaultQueryConfig = defaultQueryConfig; this.authConfig = authConfig; + this.fragmentContextFactory = fragmentContextFactory; this.startMs = startMs; this.startNs = startNs; } - /** * For callers who have already authorized their query, and where simplicity is desired over flexibility. This method * does it all in one call. Logs and metrics are emitted when the Sequence is either fully iterated or throws an @@ -278,8 +283,11 @@ public QueryResponse execute() final ResponseContext responseContext = DirectDruidClient.makeResponseContextForQuery(); + final FragmentBuilder fragmentBuilder = fragmentContextFactory.create(baseQuery, responseContext); + @SuppressWarnings("rawtypes") final Sequence res = QueryPlus.wrap(baseQuery) .withIdentity(authenticationResult.getIdentity()) + .withFragmentBuilder(fragmentBuilder) .run(texasRanger, responseContext); return new QueryResponse(res == null ? Sequences.empty() : res, responseContext); diff --git a/server/src/main/java/org/apache/druid/server/QueryLifecycleFactory.java b/server/src/main/java/org/apache/druid/server/QueryLifecycleFactory.java index 766307cef5c9..022b62f48c7d 100644 --- a/server/src/main/java/org/apache/druid/server/QueryLifecycleFactory.java +++ b/server/src/main/java/org/apache/druid/server/QueryLifecycleFactory.java @@ -27,6 +27,7 @@ import org.apache.druid.query.GenericQueryMetricsFactory; import org.apache.druid.query.QuerySegmentWalker; import org.apache.druid.query.QueryToolChestWarehouse; +import org.apache.druid.queryng.fragment.FragmentBuilderFactory; import org.apache.druid.server.log.RequestLogger; import org.apache.druid.server.security.AuthConfig; import org.apache.druid.server.security.AuthorizerMapper; @@ -42,6 +43,7 @@ public class QueryLifecycleFactory private final AuthorizerMapper authorizerMapper; private final DefaultQueryConfig defaultQueryConfig; private final AuthConfig authConfig; + private final FragmentBuilderFactory fragmentContextFactory; @Inject public QueryLifecycleFactory( @@ -52,7 +54,8 @@ public QueryLifecycleFactory( final RequestLogger requestLogger, final AuthConfig authConfig, final AuthorizerMapper authorizerMapper, - final Supplier queryConfigSupplier + final Supplier queryConfigSupplier, + final FragmentBuilderFactory fragmentContextFactory ) { this.warehouse = warehouse; @@ -63,6 +66,7 @@ public QueryLifecycleFactory( this.authorizerMapper = authorizerMapper; this.defaultQueryConfig = queryConfigSupplier.get(); this.authConfig = authConfig; + this.fragmentContextFactory = fragmentContextFactory; } public QueryLifecycle factorize() @@ -76,6 +80,7 @@ public QueryLifecycle factorize() authorizerMapper, defaultQueryConfig, authConfig, + fragmentContextFactory, System.currentTimeMillis(), System.nanoTime() ); diff --git a/server/src/test/java/org/apache/druid/guice/QueryableModuleTest.java b/server/src/test/java/org/apache/druid/guice/QueryableModuleTest.java index b51f2c14853d..39d596871216 100644 --- a/server/src/test/java/org/apache/druid/guice/QueryableModuleTest.java +++ b/server/src/test/java/org/apache/druid/guice/QueryableModuleTest.java @@ -29,6 +29,7 @@ import org.apache.druid.guice.annotations.Json; import org.apache.druid.jackson.JacksonModule; import org.apache.druid.java.util.emitter.service.ServiceEmitter; +import org.apache.druid.queryng.guice.QueryNGModule; import org.apache.druid.server.log.EmittingRequestLogger; import org.apache.druid.server.log.EmittingRequestLoggerProvider; import org.apache.druid.server.log.NoopRequestLogger; @@ -77,6 +78,7 @@ private Injector makeInjector(Properties properties) new JacksonModule(), new ConfigModule(), new QueryRunnerFactoryModule(), + new QueryNGModule(), new DruidProcessingConfigModule(), new BrokerProcessingModule(), new LifecycleModule(), diff --git a/server/src/test/java/org/apache/druid/server/QueryLifecycleTest.java b/server/src/test/java/org/apache/druid/server/QueryLifecycleTest.java index 1d1840b6c72e..98cb1e21cca9 100644 --- a/server/src/test/java/org/apache/druid/server/QueryLifecycleTest.java +++ b/server/src/test/java/org/apache/druid/server/QueryLifecycleTest.java @@ -28,7 +28,6 @@ import org.apache.druid.query.DefaultQueryConfig; import org.apache.druid.query.Druids; import org.apache.druid.query.GenericQueryMetricsFactory; -import org.apache.druid.query.QueryContextTest; import org.apache.druid.query.QueryMetrics; import org.apache.druid.query.QueryRunner; import org.apache.druid.query.QuerySegmentWalker; @@ -36,6 +35,7 @@ import org.apache.druid.query.QueryToolChestWarehouse; import org.apache.druid.query.aggregation.CountAggregatorFactory; import org.apache.druid.query.timeseries.TimeseriesQuery; +import org.apache.druid.queryng.fragment.NullFragmentBuilderFactory; import org.apache.druid.server.log.RequestLogger; import org.apache.druid.server.security.Access; import org.apache.druid.server.security.Action; @@ -110,6 +110,7 @@ public void setup() authzMapper, queryConfig, authConfig, + new NullFragmentBuilderFactory(), millis, nanos ); @@ -245,40 +246,6 @@ public void testAuthorizeQueryContext_notAuthorized() Assert.assertFalse(lifecycle.authorize(mockRequest()).isAllowed()); } - @Test - public void testAuthorizeLegacyQueryContext_authorized() - { - EasyMock.expect(queryConfig.getContext()).andReturn(ImmutableMap.of()).anyTimes(); - EasyMock.expect(authConfig.authorizeQueryContextParams()).andReturn(true).anyTimes(); - EasyMock.expect(authenticationResult.getIdentity()).andReturn(IDENTITY).anyTimes(); - EasyMock.expect(authenticationResult.getAuthorizerName()).andReturn(AUTHORIZER).anyTimes(); - EasyMock.expect(authorizer.authorize(authenticationResult, new Resource("fake", ResourceType.DATASOURCE), Action.READ)) - .andReturn(Access.OK); - EasyMock.expect(authorizer.authorize(authenticationResult, new Resource("foo", ResourceType.QUERY_CONTEXT), Action.WRITE)) - .andReturn(Access.OK); - EasyMock.expect(authorizer.authorize(authenticationResult, new Resource("baz", ResourceType.QUERY_CONTEXT), Action.WRITE)).andReturn(Access.OK); - // to use legacy query context with context authorization, even system generated things like queryId need to be explicitly added - EasyMock.expect(authorizer.authorize(authenticationResult, new Resource("queryId", ResourceType.QUERY_CONTEXT), Action.WRITE)) - .andReturn(Access.OK); - - EasyMock.expect(toolChestWarehouse.getToolChest(EasyMock.anyObject())) - .andReturn(toolChest) - .once(); - - replayAll(); - - final QueryContextTest.LegacyContextQuery query = new QueryContextTest.LegacyContextQuery(ImmutableMap.of("foo", "bar", "baz", "qux")); - - lifecycle.initialize(query); - - Assert.assertNull(lifecycle.getQuery().getQueryContext()); - Assert.assertTrue(lifecycle.getQuery().getContext().containsKey("foo")); - Assert.assertTrue(lifecycle.getQuery().getContext().containsKey("baz")); - Assert.assertTrue(lifecycle.getQuery().getContext().containsKey("queryId")); - - Assert.assertTrue(lifecycle.authorize(mockRequest()).isAllowed()); - } - private HttpServletRequest mockRequest() { HttpServletRequest request = EasyMock.createNiceMock(HttpServletRequest.class); diff --git a/server/src/test/java/org/apache/druid/server/QueryResourceTest.java b/server/src/test/java/org/apache/druid/server/QueryResourceTest.java index c6cae9e1e1bb..d086e1ddf6d5 100644 --- a/server/src/test/java/org/apache/druid/server/QueryResourceTest.java +++ b/server/src/test/java/org/apache/druid/server/QueryResourceTest.java @@ -56,6 +56,7 @@ import org.apache.druid.query.SegmentDescriptor; import org.apache.druid.query.TruncatedResponseContextException; import org.apache.druid.query.timeboundary.TimeBoundaryResultValue; +import org.apache.druid.queryng.fragment.NullFragmentBuilderFactory; import org.apache.druid.server.initialization.ServerConfig; import org.apache.druid.server.log.TestRequestLogger; import org.apache.druid.server.metrics.NoopServiceEmitter; @@ -220,7 +221,8 @@ private QueryResource createQueryResource(ResponseContextConfig responseContextC testRequestLogger, new AuthConfig(), AuthTestUtils.TEST_AUTHORIZER_MAPPER, - Suppliers.ofInstance(new DefaultQueryConfig(ImmutableMap.of())) + Suppliers.ofInstance(new DefaultQueryConfig(ImmutableMap.of())), + new NullFragmentBuilderFactory() ), jsonMapper, smileMapper, @@ -266,7 +268,8 @@ public void testGoodQueryWithQueryConfigOverrideDefault() throws IOException testRequestLogger, new AuthConfig(), AuthTestUtils.TEST_AUTHORIZER_MAPPER, - Suppliers.ofInstance(overrideConfig) + Suppliers.ofInstance(overrideConfig), + new NullFragmentBuilderFactory() ), jsonMapper, smileMapper, @@ -318,7 +321,8 @@ public void testGoodQueryWithQueryConfigDoesNotOverrideQueryContext() throws IOE testRequestLogger, new AuthConfig(), AuthTestUtils.TEST_AUTHORIZER_MAPPER, - Suppliers.ofInstance(overrideConfig) + Suppliers.ofInstance(overrideConfig), + new NullFragmentBuilderFactory() ), jsonMapper, smileMapper, @@ -705,7 +709,8 @@ public Access authorize(AuthenticationResult authenticationResult, Resource reso testRequestLogger, new AuthConfig(), authMapper, - Suppliers.ofInstance(new DefaultQueryConfig(ImmutableMap.of())) + Suppliers.ofInstance(new DefaultQueryConfig(ImmutableMap.of())), + new NullFragmentBuilderFactory() ), jsonMapper, smileMapper, @@ -781,7 +786,8 @@ public QueryRunner getQueryRunnerForSegments(Query query, Iterable QueryRunner getQueryRunnerForSegments(Query query, Iterable getModules() new BrokerProcessingModule(), new QueryableModule(), new QueryRunnerFactoryModule(), + new QueryNGModule(), new SegmentWranglerModule(), new JoinableFactoryModule(), new BrokerServiceModule(), diff --git a/services/src/main/java/org/apache/druid/cli/CliHistorical.java b/services/src/main/java/org/apache/druid/cli/CliHistorical.java index dc1acc41f873..268dc59b6696 100644 --- a/services/src/main/java/org/apache/druid/cli/CliHistorical.java +++ b/services/src/main/java/org/apache/druid/cli/CliHistorical.java @@ -46,6 +46,7 @@ import org.apache.druid.java.util.common.logger.Logger; import org.apache.druid.query.QuerySegmentWalker; import org.apache.druid.query.lookup.LookupModule; +import org.apache.druid.queryng.guice.QueryNGModule; import org.apache.druid.server.QueryResource; import org.apache.druid.server.ResponseContextConfig; import org.apache.druid.server.SegmentManager; @@ -98,6 +99,7 @@ protected List getModules() new DruidProcessingModule(), new QueryableModule(), new QueryRunnerFactoryModule(), + new QueryNGModule(), new JoinableFactoryModule(), new HistoricalServiceModule(), binder -> { diff --git a/services/src/main/java/org/apache/druid/cli/CliIndexer.java b/services/src/main/java/org/apache/druid/cli/CliIndexer.java index e573e32c8b34..797ea5f50154 100644 --- a/services/src/main/java/org/apache/druid/cli/CliIndexer.java +++ b/services/src/main/java/org/apache/druid/cli/CliIndexer.java @@ -65,6 +65,7 @@ import org.apache.druid.metadata.input.InputSourceModule; import org.apache.druid.query.QuerySegmentWalker; import org.apache.druid.query.lookup.LookupModule; +import org.apache.druid.queryng.guice.QueryNGModule; import org.apache.druid.segment.realtime.appenderator.AppenderatorsManager; import org.apache.druid.segment.realtime.appenderator.UnifiedIndexerAppenderatorsManager; import org.apache.druid.server.DruidNode; @@ -122,6 +123,7 @@ protected List getModules() new DruidProcessingModule(), new QueryableModule(), new QueryRunnerFactoryModule(), + new QueryNGModule(), new JoinableFactoryModule(), new IndexerServiceModule(), new Module() diff --git a/services/src/main/java/org/apache/druid/cli/CliPeon.java b/services/src/main/java/org/apache/druid/cli/CliPeon.java index 717d036cf63f..ff1488d476c9 100644 --- a/services/src/main/java/org/apache/druid/cli/CliPeon.java +++ b/services/src/main/java/org/apache/druid/cli/CliPeon.java @@ -100,6 +100,7 @@ import org.apache.druid.metadata.input.InputSourceModule; import org.apache.druid.query.QuerySegmentWalker; import org.apache.druid.query.lookup.LookupModule; +import org.apache.druid.queryng.guice.QueryNGModule; import org.apache.druid.segment.handoff.CoordinatorBasedSegmentHandoffNotifierConfig; import org.apache.druid.segment.handoff.CoordinatorBasedSegmentHandoffNotifierFactory; import org.apache.druid.segment.handoff.SegmentHandoffNotifierFactory; @@ -196,6 +197,7 @@ protected List getModules() new DruidProcessingModule(), new QueryableModule(), new QueryRunnerFactoryModule(), + new QueryNGModule(), new JoinableFactoryModule(), new Module() { diff --git a/services/src/main/java/org/apache/druid/cli/CreateTables.java b/services/src/main/java/org/apache/druid/cli/CreateTables.java index 46f824f02df9..3cceb7eca01c 100644 --- a/services/src/main/java/org/apache/druid/cli/CreateTables.java +++ b/services/src/main/java/org/apache/druid/cli/CreateTables.java @@ -35,6 +35,7 @@ import org.apache.druid.metadata.MetadataStorageConnector; import org.apache.druid.metadata.MetadataStorageConnectorConfig; import org.apache.druid.metadata.MetadataStorageTablesConfig; +import org.apache.druid.queryng.guice.QueryNGModule; import org.apache.druid.server.DruidNode; import java.util.List; @@ -77,6 +78,7 @@ protected List getModules() new DruidProcessingModule(), new QueryableModule(), new QueryRunnerFactoryModule(), + new QueryNGModule(), binder -> { JsonConfigProvider.bindInstance( binder, diff --git a/services/src/main/java/org/apache/druid/cli/DumpSegment.java b/services/src/main/java/org/apache/druid/cli/DumpSegment.java index aece45c22248..f339f5f75662 100644 --- a/services/src/main/java/org/apache/druid/cli/DumpSegment.java +++ b/services/src/main/java/org/apache/druid/cli/DumpSegment.java @@ -66,6 +66,7 @@ import org.apache.druid.query.metadata.metadata.SegmentAnalysis; import org.apache.druid.query.metadata.metadata.SegmentMetadataQuery; import org.apache.druid.query.spec.SpecificSegmentSpec; +import org.apache.druid.queryng.guice.QueryNGModule; import org.apache.druid.segment.BaseObjectColumnValueSelector; import org.apache.druid.segment.ColumnSelectorFactory; import org.apache.druid.segment.Cursor; @@ -443,6 +444,7 @@ protected List getModules() new DruidProcessingModule(), new QueryableModule(), new QueryRunnerFactoryModule(), + new QueryNGModule(), new Module() { @Override diff --git a/services/src/main/java/org/apache/druid/cli/ExportMetadata.java b/services/src/main/java/org/apache/druid/cli/ExportMetadata.java index ece46fd11ebe..fb486227b9f6 100644 --- a/services/src/main/java/org/apache/druid/cli/ExportMetadata.java +++ b/services/src/main/java/org/apache/druid/cli/ExportMetadata.java @@ -41,6 +41,7 @@ import org.apache.druid.metadata.MetadataStorageConnectorConfig; import org.apache.druid.metadata.MetadataStorageTablesConfig; import org.apache.druid.metadata.SQLMetadataConnector; +import org.apache.druid.queryng.guice.QueryNGModule; import org.apache.druid.segment.loading.DataSegmentPusher; import org.apache.druid.server.DruidNode; import org.apache.druid.timeline.DataSegment; @@ -142,6 +143,7 @@ protected List getModules() new DruidProcessingModule(), new QueryableModule(), new QueryRunnerFactoryModule(), + new QueryNGModule(), binder -> { JsonConfigProvider.bindInstance( binder, @@ -208,7 +210,7 @@ public void run() final Injector injector = makeInjector(); SQLMetadataConnector dbConnector = injector.getInstance(SQLMetadataConnector.class); MetadataStorageTablesConfig metadataStorageTablesConfig = injector.getInstance(MetadataStorageTablesConfig.class); - + // We export a raw CSV first, and then apply some conversions for easier imports: // Boolean strings are rewritten as 1 and 0 // hexadecimal BLOB columns are rewritten with rewriteHexPayloadAsEscapedJson() diff --git a/services/src/main/java/org/apache/druid/cli/ResetCluster.java b/services/src/main/java/org/apache/druid/cli/ResetCluster.java index 16958398a82b..4b78cbb7bb69 100644 --- a/services/src/main/java/org/apache/druid/cli/ResetCluster.java +++ b/services/src/main/java/org/apache/druid/cli/ResetCluster.java @@ -36,6 +36,7 @@ import org.apache.druid.java.util.common.logger.Logger; import org.apache.druid.metadata.MetadataStorageConnector; import org.apache.druid.metadata.MetadataStorageTablesConfig; +import org.apache.druid.queryng.guice.QueryNGModule; import org.apache.druid.segment.loading.DataSegmentKiller; import org.apache.druid.server.DruidNode; import org.apache.druid.tasklogs.TaskLogKiller; @@ -84,6 +85,7 @@ protected List getModules() new DruidProcessingModule(), new QueryableModule(), new QueryRunnerFactoryModule(), + new QueryNGModule(), binder -> { JsonConfigProvider.bindInstance( binder, diff --git a/services/src/main/java/org/apache/druid/cli/ValidateSegments.java b/services/src/main/java/org/apache/druid/cli/ValidateSegments.java index de2b2bb6cb17..fe5161799f26 100644 --- a/services/src/main/java/org/apache/druid/cli/ValidateSegments.java +++ b/services/src/main/java/org/apache/druid/cli/ValidateSegments.java @@ -33,6 +33,7 @@ import org.apache.druid.java.util.common.IAE; import org.apache.druid.java.util.common.logger.Logger; import org.apache.druid.query.DruidProcessingConfig; +import org.apache.druid.queryng.guice.QueryNGModule; import org.apache.druid.segment.IndexIO; import org.apache.druid.segment.column.ColumnConfig; @@ -85,6 +86,7 @@ protected List getModules() new DruidProcessingModule(), new QueryableModule(), new QueryRunnerFactoryModule(), + new QueryNGModule(), new Module() { @Override diff --git a/services/src/main/java/org/apache/druid/cli/validate/DruidJsonValidator.java b/services/src/main/java/org/apache/druid/cli/validate/DruidJsonValidator.java index 063b164711be..f85169a770fb 100644 --- a/services/src/main/java/org/apache/druid/cli/validate/DruidJsonValidator.java +++ b/services/src/main/java/org/apache/druid/cli/validate/DruidJsonValidator.java @@ -54,6 +54,7 @@ import org.apache.druid.java.util.common.UOE; import org.apache.druid.java.util.common.logger.Logger; import org.apache.druid.query.Query; +import org.apache.druid.queryng.guice.QueryNGModule; import java.io.File; import java.io.IOException; @@ -106,6 +107,7 @@ protected List getModules() new DruidProcessingModule(), new QueryableModule(), new QueryRunnerFactoryModule(), + new QueryNGModule(), binder -> { binder.bindConstant().annotatedWith(Names.named("serviceName")).to("druid/validator"); binder.bindConstant().annotatedWith(Names.named("servicePort")).to(0); diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/util/CalciteTests.java b/sql/src/test/java/org/apache/druid/sql/calcite/util/CalciteTests.java index 54d3aea6f565..c96a9fcc9f3b 100644 --- a/sql/src/test/java/org/apache/druid/sql/calcite/util/CalciteTests.java +++ b/sql/src/test/java/org/apache/druid/sql/calcite/util/CalciteTests.java @@ -86,6 +86,7 @@ import org.apache.druid.query.expression.TestExprMacroTable; import org.apache.druid.query.lookup.LookupExtractorFactoryContainerProvider; import org.apache.druid.query.lookup.LookupSerdeModule; +import org.apache.druid.queryng.fragment.TestFragmentBuilderFactory; import org.apache.druid.segment.IndexBuilder; import org.apache.druid.segment.QueryableIndex; import org.apache.druid.segment.TestHelper; @@ -151,6 +152,7 @@ import org.joda.time.chrono.ISOChronology; import javax.annotation.Nullable; + import java.io.File; import java.util.ArrayList; import java.util.Arrays; @@ -826,7 +828,8 @@ public > QueryToolChest getToolChest new NoopRequestLogger(), new AuthConfig(), TEST_AUTHORIZER_MAPPER, - Suppliers.ofInstance(new DefaultQueryConfig(ImmutableMap.of())) + Suppliers.ofInstance(new DefaultQueryConfig(ImmutableMap.of())), + new TestFragmentBuilderFactory() ); } From caf24f7ee418a61a207affcff0a64e8fa0471300 Mon Sep 17 00:00:00 2001 From: Paul Rogers Date: Wed, 22 Jun 2022 17:54:59 -0700 Subject: [PATCH 10/11] Added more unit tests Extended ScanQueryRunnerTest to run with operators. Fixed a few bugs and build issues. --- .../fragment/TestFragmentBuilderFactory.java | 19 ++- .../druid/queryng/operators/Iterators.java | 12 ++ .../queryng/operators/scan/CursorReader.java | 39 +++++-- .../scan/ScanCompactListToArrayOperator.java | 4 +- .../scan/ScanListToArrayOperator.java | 4 +- .../operators/scan/ScanQueryOperator.java | 103 +++++++++-------- .../scan/ScanResultOffsetOperator.java | 3 +- .../scan/ScanRowToBatchOperator.java | 73 ------------ .../druid/queryng/planner/ScanPlanner.java | 108 ++++++++++-------- .../druid/query/scan/ScanQueryRunnerTest.java | 77 ++++++++----- .../ScanQueryRunnerWithOperatorsTest.java | 54 +++++++++ .../druid/query/scan/ScanQueryTest.java | 6 +- .../operators/scan/MockScanResultReader.java | 9 +- .../apache/druid/server/QueryLifecycle.java | 2 - .../SetAndVerifyContextQueryRunner.java | 2 +- .../apache/druid/sql/guice/SqlModuleTest.java | 6 +- 16 files changed, 298 insertions(+), 223 deletions(-) delete mode 100644 processing/src/main/java/org/apache/druid/queryng/operators/scan/ScanRowToBatchOperator.java create mode 100644 processing/src/test/java/org/apache/druid/query/scan/ScanQueryRunnerWithOperatorsTest.java diff --git a/processing/src/main/java/org/apache/druid/queryng/fragment/TestFragmentBuilderFactory.java b/processing/src/main/java/org/apache/druid/queryng/fragment/TestFragmentBuilderFactory.java index 9fdfe34658f3..a762c93cbb23 100644 --- a/processing/src/main/java/org/apache/druid/queryng/fragment/TestFragmentBuilderFactory.java +++ b/processing/src/main/java/org/apache/druid/queryng/fragment/TestFragmentBuilderFactory.java @@ -32,14 +32,25 @@ public class TestFragmentBuilderFactory implements FragmentBuilderFactory { private static final String ENABLED_KEY = QueryNGConfig.CONFIG_ROOT + ".enabled"; - private static final boolean ENABLED = Boolean.parseBoolean(System.getProperty(ENABLED_KEY)); + + private boolean enabled; + + public TestFragmentBuilderFactory() + { + this.enabled = Boolean.parseBoolean(System.getProperty(ENABLED_KEY)); + } + + public TestFragmentBuilderFactory(boolean enabled) + { + this.enabled = enabled; + } @Override public FragmentBuilder create(Query query, ResponseContext responseContext) { - //if (!ENABLED) { - // return null; - //} + if (!enabled) { + return null; + } if (!(query instanceof ScanQuery)) { return null; } diff --git a/processing/src/main/java/org/apache/druid/queryng/operators/Iterators.java b/processing/src/main/java/org/apache/druid/queryng/operators/Iterators.java index 52bc987dd386..0f77989636b7 100644 --- a/processing/src/main/java/org/apache/druid/queryng/operators/Iterators.java +++ b/processing/src/main/java/org/apache/druid/queryng/operators/Iterators.java @@ -97,4 +97,16 @@ public static List toList(ResultIterator operIter) return Lists.newArrayList(new Iterators.ShimIterator(operIter)); } + public static ResultIterator emptyIterator() + { + return new ResultIterator() + { + @Override + public T next() throws EofException + { + throw Operators.eof(); + } + }; + } + } diff --git a/processing/src/main/java/org/apache/druid/queryng/operators/scan/CursorReader.java b/processing/src/main/java/org/apache/druid/queryng/operators/scan/CursorReader.java index f9cee6d12de2..0776cf515c3f 100644 --- a/processing/src/main/java/org/apache/druid/queryng/operators/scan/CursorReader.java +++ b/processing/src/main/java/org/apache/druid/queryng/operators/scan/CursorReader.java @@ -21,7 +21,9 @@ import com.google.common.base.Supplier; import org.apache.druid.java.util.common.DateTimes; +import org.apache.druid.java.util.common.StringUtils; import org.apache.druid.java.util.common.UOE; +import org.apache.druid.query.QueryTimeoutException; import org.apache.druid.query.scan.ScanQuery.ResultFormat; import org.apache.druid.segment.BaseObjectColumnValueSelector; import org.apache.druid.segment.Cursor; @@ -32,20 +34,25 @@ import java.util.LinkedHashMap; import java.util.List; import java.util.Map; +import java.util.NoSuchElementException; /** - * The cursor reader is a leaf operator which uses a cursor to access - * data, which is returned via the operator protocol. Converts cursor - * data into one of two supported Druid formats. Enforces a query row limit. + * The cursor reader is a leaf operator which uses a cursor to access data, + * which is returned via the operator protocol. Converts cursor data into one + * of two supported Druid formats. Enforces a query row limit. *

      * Unlike most operators, this one is created on the fly by its parent * to scan a specific query known only at runtime. A storage adapter may * choose to create one or more queries: each is handled by an instance of this * class. + *

      + * As a small performance boost, we create a layer of "accessors" on top of + * the column selectors. The accesors encode the column-specific logic to + * avoid the need for if-statements on every row. * * @see {@link org.apache.druid.query.scan.ScanQueryEngine} */ -public class CursorReader implements Iterator +public class CursorReader implements Iterator> { private final Cursor cursor; private final List selectedColumns; @@ -53,6 +60,9 @@ public class CursorReader implements Iterator private final int batchSize; private final ResultFormat resultFormat; private final List> columnAccessors; + private final boolean hasTimeout; + private final long timeoutAt; + private final String queryId; private long targetCount; private long rowCount; @@ -62,7 +72,9 @@ public CursorReader( final long limit, final int batchSize, final ResultFormat resultFormat, - final boolean isLegacy + final boolean isLegacy, + final long timeoutAt, + final String queryId ) { this.cursor = cursor; @@ -71,6 +83,9 @@ public CursorReader( this.batchSize = batchSize; this.resultFormat = resultFormat; this.columnAccessors = buildAccessors(isLegacy); + this.hasTimeout = timeoutAt < Long.MAX_VALUE; + this.timeoutAt = timeoutAt; + this.queryId = queryId; } private List> buildAccessors(final boolean isLegacy) @@ -121,8 +136,14 @@ public boolean hasNext() } @Override - public Object next() + public List next() { + if (!hasNext()) { + throw new NoSuchElementException(); + } + if (hasTimeout && System.currentTimeMillis() >= timeoutAt) { + throw new QueryTimeoutException(StringUtils.nonStrictFormat("Query [%s] timed out", queryId)); + } targetCount = Math.min(limit - rowCount, rowCount + batchSize); switch (resultFormat) { case RESULT_FORMAT_LIST: @@ -130,7 +151,7 @@ public Object next() case RESULT_FORMAT_COMPACTED_LIST: return nextAsCompactList(); default: - throw new UOE("resultFormat[%s] is not supported", resultFormat.toString()); + throw new UOE("resultFormat [%s] is not supported", resultFormat.toString()); } } @@ -154,7 +175,7 @@ private Object getColumnValue(int i) * Convert a cursor row into a simple list of maps, where each map * represents a single event, and each map entry represents a column. */ - public Object nextAsListOfMaps() + public List nextAsListOfMaps() { final List> events = new ArrayList<>(batchSize); while (hasNextRow()) { @@ -168,7 +189,7 @@ public Object nextAsListOfMaps() return events; } - public Object nextAsCompactList() + public List nextAsCompactList() { final List> events = new ArrayList<>(batchSize); while (hasNextRow()) { diff --git a/processing/src/main/java/org/apache/druid/queryng/operators/scan/ScanCompactListToArrayOperator.java b/processing/src/main/java/org/apache/druid/queryng/operators/scan/ScanCompactListToArrayOperator.java index d30a8b09d991..756e37675f30 100644 --- a/processing/src/main/java/org/apache/druid/queryng/operators/scan/ScanCompactListToArrayOperator.java +++ b/processing/src/main/java/org/apache/druid/queryng/operators/scan/ScanCompactListToArrayOperator.java @@ -28,11 +28,11 @@ /** * Converts individual scan query rows with the - * {@link org.apache.druid.query.scan.ScanQuery.ResultFormat.RESULT_FORMAT_COMPACTED_LIST + * {@link org.apache.druid.query.scan.ScanQuery.ResultFormat#RESULT_FORMAT_COMPACTED_LIST * ResultFormat.RESULT_FORMAT_COMPACTED_LIST} format into an object array with fields * in the order given by the output schema. * - * @See {@link org.apache.druid.query.scan.ScanQueryQueryToolChest.resultsAsArrays + * @See {@link org.apache.druid.query.scan.ScanQueryQueryToolChest#resultsAsArrays * ScanQueryQueryToolChest.resultsAsArrays} */ public class ScanCompactListToArrayOperator extends MappingOperator, Object[]> diff --git a/processing/src/main/java/org/apache/druid/queryng/operators/scan/ScanListToArrayOperator.java b/processing/src/main/java/org/apache/druid/queryng/operators/scan/ScanListToArrayOperator.java index cb53ed99902d..221363a48d10 100644 --- a/processing/src/main/java/org/apache/druid/queryng/operators/scan/ScanListToArrayOperator.java +++ b/processing/src/main/java/org/apache/druid/queryng/operators/scan/ScanListToArrayOperator.java @@ -28,11 +28,11 @@ /** * Converts individual scan query rows with the - * {@link org.apache.druid.query.scan.ScanQuery.ResultFormat.RESULT_FORMAT_LIST + * {@link org.apache.druid.query.scan.ScanQuery.ResultFormat#RESULT_FORMAT_LIST * ResultFormat.RESULT_FORMAT_LIST} format into an object array with fields * in the order given by the output schema. * - * @See {@link org.apache.druid.query.scan.ScanQueryQueryToolChest.resultsAsArrays + * @See {@link org.apache.druid.query.scan.ScanQueryQueryToolChest#resultsAsArrays * ScanQueryQueryToolChest.resultsAsArrays} */ public class ScanListToArrayOperator extends MappingOperator, Object[]> diff --git a/processing/src/main/java/org/apache/druid/queryng/operators/scan/ScanQueryOperator.java b/processing/src/main/java/org/apache/druid/queryng/operators/scan/ScanQueryOperator.java index b82f6209d4e5..a79264348aee 100644 --- a/processing/src/main/java/org/apache/druid/queryng/operators/scan/ScanQueryOperator.java +++ b/processing/src/main/java/org/apache/druid/queryng/operators/scan/ScanQueryOperator.java @@ -25,11 +25,13 @@ import org.apache.druid.java.util.common.ISE; import org.apache.druid.java.util.common.granularity.Granularities; import org.apache.druid.query.QueryContexts; +import org.apache.druid.query.QueryMetrics; import org.apache.druid.query.context.ResponseContext; import org.apache.druid.query.filter.Filter; import org.apache.druid.query.scan.ScanQuery; import org.apache.druid.query.scan.ScanResultValue; import org.apache.druid.queryng.fragment.FragmentContext; +import org.apache.druid.queryng.operators.Iterators; import org.apache.druid.queryng.operators.Operator; import org.apache.druid.queryng.operators.Operators; import org.apache.druid.queryng.operators.SequenceIterator; @@ -41,6 +43,8 @@ import org.apache.druid.segment.filter.Filters; import org.joda.time.Interval; +import javax.annotation.Nullable; + import java.io.IOException; import java.util.ArrayList; import java.util.Arrays; @@ -77,15 +81,10 @@ private Impl(FragmentContext context) { this.context = context; ResponseContext responseContext = context.responseContext(); - responseContext.add(ResponseContext.Keys.NUM_SCANNED_ROWS, 0L); - long baseLimit = query.getScanRowsLimit(); - if (limitType() == Limit.GLOBAL) { - limit = baseLimit - (Long) responseContext.get(ResponseContext.Keys.NUM_SCANNED_ROWS); - } else { - limit = baseLimit; - } + // If the row count is not set, set it to 0, else do nothing. + responseContext.addRowScanCount(0); + limit = calculateRemainingScanRowsLimit(query, responseContext); final StorageAdapter adapter = segment.asStorageAdapter(); - //final StorageAdapter adapter = new MockStorageAdapter(); if (adapter == null) { throw new ISE( "Null storage adapter found. Probably trying to issue a query against a segment being memory unmapped." @@ -102,13 +101,24 @@ private Impl(FragmentContext context) query.getVirtualColumns(), Granularities.ALL, isDescendingOrder(), - null + queryMetrics )); } + /** + * If we're performing time-ordering, we want to scan through the first `limit` rows in each segment ignoring the number + * of rows already counted on other segments. + */ + private long calculateRemainingScanRowsLimit(ScanQuery query, ResponseContext responseContext) + { + if (query.getTimeOrder().equals(ScanQuery.Order.NONE)) { + return query.getScanRowsLimit() - (Long) responseContext.getRowScanCount(); + } + return query.getScanRowsLimit(); + } + protected List inferColumns(StorageAdapter adapter, boolean isLegacy) { - List cols = new ArrayList<>(); final Set availableColumns = Sets.newLinkedHashSet( Iterables.concat( Collections.singleton(isLegacy ? LEGACY_TIMESTAMP_KEY : ColumnHolder.TIME_COLUMN_NAME), @@ -121,12 +131,10 @@ protected List inferColumns(StorageAdapter adapter, boolean isLegacy) ) ); - cols.addAll(availableColumns); - if (isLegacy) { - cols.remove(ColumnHolder.TIME_COLUMN_NAME); + availableColumns.remove(ColumnHolder.TIME_COLUMN_NAME); } - return cols; + return new ArrayList<>(availableColumns); } /** @@ -178,7 +186,9 @@ public ScanResultValue next() throws EofException limit - rowCount, batchSize, query.getResultFormat(), - isLegacy); + isLegacy, + timeoutAt, + query.getId()); } } @@ -192,8 +202,7 @@ private void closeCursorReader() private void finish() { closeCursorReader(); - ResponseContext responseContext = context.responseContext(); - responseContext.add(ResponseContext.Keys.NUM_SCANNED_ROWS, rowCount); + context.responseContext().add(ResponseContext.Keys.NUM_SCANNED_ROWS, rowCount); try { iter.close(); } @@ -222,23 +231,33 @@ public enum Limit private final Filter filter; private final boolean isLegacy; private final int batchSize; + private final long timeoutAt; + @Nullable final QueryMetrics queryMetrics; private Impl impl; public ScanQueryOperator( final FragmentContext context, final ScanQuery query, - final Segment segment) + final Segment segment, + @Nullable final QueryMetrics queryMetrics) { this.context = context; this.query = query; this.segment = segment; this.segmentId = segment.getId().toString(); - this.columns = defineColumns(query); List intervals = query.getQuerySegmentSpec().getIntervals(); Preconditions.checkArgument(intervals.size() == 1, "Can only handle a single interval, got [%s]", intervals); this.filter = Filters.convertToCNFFromQueryContext(query, Filters.toFilter(query.getFilter())); + // "legacy" should be non-null due to toolChest.mergeResults this.isLegacy = Preconditions.checkNotNull(query.isLegacy(), "Expected non-null 'legacy' parameter"); this.batchSize = query.getBatchSize(); + this.queryMetrics = queryMetrics; + this.columns = defineColumns(query); + if (hasTimeout()) { + timeoutAt = context.responseContext().getTimeoutTime(); + } else { + timeoutAt = Long.MAX_VALUE; + } context.register(this); } @@ -247,26 +266,27 @@ public ScanQueryOperator( */ private List defineColumns(ScanQuery query) { - List queryCols = query.getColumns(); - - // Missing or empty list means wildcard - if (queryCols == null || queryCols.isEmpty()) { + if (isWildcard(query)) { return null; } - final List planCols = new ArrayList<>(); - if (query.isLegacy() && !queryCols.contains(LEGACY_TIMESTAMP_KEY)) { - planCols.add(LEGACY_TIMESTAMP_KEY); - } - // Unless we're in legacy mode, planCols equals query.getColumns() exactly. This is nice since it makes // the compactedList form easier to use. - planCols.addAll(queryCols); - return planCols; + List queryCols = query.getColumns(); + if (isLegacy && !queryCols.contains(LEGACY_TIMESTAMP_KEY)) { + final List planCols = new ArrayList<>(); + planCols.add(LEGACY_TIMESTAMP_KEY); + planCols.addAll(queryCols); + return planCols; + } else { + return queryCols; + } } public boolean isWildcard(ScanQuery query) { - return (query.getColumns() == null || query.getColumns().isEmpty()); + // Missing or empty list means wildcard + List queryCols = query.getColumns(); + return (queryCols == null || queryCols.isEmpty()); } // TODO: Review against latest @@ -286,18 +306,6 @@ public boolean isWildcard() return columns == null; } - // TODO: Review against latest - public Limit limitType() - { - if (!query.isLimited()) { - return Limit.NONE; - } else if (query.getTimeOrder().equals(ScanQuery.Order.NONE)) { - return Limit.LOCAL; - } else { - return Limit.GLOBAL; - } - } - public Interval interval() { return query.getQuerySegmentSpec().getIntervals().get(0); @@ -306,8 +314,13 @@ public Interval interval() @Override public ResultIterator open() { - impl = new Impl(context); - return impl; + final Long numScannedRows = context.responseContext().getRowScanCount(); + if (numScannedRows != null && numScannedRows >= query.getScanRowsLimit() && query.getTimeOrder().equals(ScanQuery.Order.NONE)) { + return Iterators.emptyIterator(); + } else { + impl = new Impl(context); + return impl; + } } @Override diff --git a/processing/src/main/java/org/apache/druid/queryng/operators/scan/ScanResultOffsetOperator.java b/processing/src/main/java/org/apache/druid/queryng/operators/scan/ScanResultOffsetOperator.java index a2a8fc99124c..111e887fa8aa 100644 --- a/processing/src/main/java/org/apache/druid/queryng/operators/scan/ScanResultOffsetOperator.java +++ b/processing/src/main/java/org/apache/druid/queryng/operators/scan/ScanResultOffsetOperator.java @@ -51,8 +51,9 @@ public ScanResultValue next() throws EofException { if (rowCount == 0) { return skip(); + } else { + return inputIter.next(); } - return inputIter.next(); } private ScanResultValue skip() throws EofException diff --git a/processing/src/main/java/org/apache/druid/queryng/operators/scan/ScanRowToBatchOperator.java b/processing/src/main/java/org/apache/druid/queryng/operators/scan/ScanRowToBatchOperator.java deleted file mode 100644 index 6816b8194fc3..000000000000 --- a/processing/src/main/java/org/apache/druid/queryng/operators/scan/ScanRowToBatchOperator.java +++ /dev/null @@ -1,73 +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.queryng.operators.scan; - -import org.apache.druid.query.scan.ScanResultValue; -import org.apache.druid.queryng.fragment.FragmentContext; -import org.apache.druid.queryng.operators.MappingOperator; -import org.apache.druid.queryng.operators.Operator; -import org.apache.druid.queryng.operators.Operators; - -import java.util.ArrayList; -import java.util.List; - -/** - * Pack a set of individual scan results into a batch up to the - * given size. - * - * @see {@link org.apache.druid.query.scan.ScanQueryLimitRowIterator} - */ -public class ScanRowToBatchOperator extends MappingOperator -{ - private final int batchSize; - - public ScanRowToBatchOperator( - FragmentContext context, - Operator child, - int batchSize) - { - super(context, child); - this.batchSize = batchSize; - } - - @Override - public ScanResultValue next() throws EofException - { - List eventsToAdd = new ArrayList<>(batchSize); - List columns = null; - while (eventsToAdd.size() < batchSize) { - try { - ScanResultValue srv = inputIter.next(); - if (columns == null) { - columns = srv.getColumns(); - } - eventsToAdd.add(srv.getRows().get(0)); - } - catch (EofException e) { - if (eventsToAdd.isEmpty()) { - throw Operators.eof(); - } - // We'll report EOF on the next call. - break; - } - } - return new ScanResultValue(null, columns, eventsToAdd); - } -} diff --git a/processing/src/main/java/org/apache/druid/queryng/planner/ScanPlanner.java b/processing/src/main/java/org/apache/druid/queryng/planner/ScanPlanner.java index 494a4fdeb54e..c98ee454a7ec 100644 --- a/processing/src/main/java/org/apache/druid/queryng/planner/ScanPlanner.java +++ b/processing/src/main/java/org/apache/druid/queryng/planner/ScanPlanner.java @@ -34,6 +34,7 @@ import org.apache.druid.query.scan.ScanResultValue; import org.apache.druid.queryng.fragment.FragmentContext; import org.apache.druid.queryng.operators.ConcatOperator; +import org.apache.druid.queryng.operators.NullOperator; import org.apache.druid.queryng.operators.Operator; import org.apache.druid.queryng.operators.Operators; import org.apache.druid.queryng.operators.scan.GroupedScanResultLimitOperator; @@ -43,6 +44,7 @@ import org.apache.druid.queryng.operators.scan.ScanQueryOperator; import org.apache.druid.queryng.operators.scan.ScanResultOffsetOperator; import org.apache.druid.queryng.operators.scan.UngroupedScanResultLimitOperator; +import org.apache.druid.segment.QueryableIndex; import org.apache.druid.segment.Segment; import java.util.ArrayList; @@ -51,8 +53,6 @@ /** * Scan-specific parts of the hybrid query planner. - * - * @see {@link QueryPlanner} */ public class ScanPlanner { @@ -73,7 +73,7 @@ public class ScanPlanner * * * @see {@link org.apache.druid.query.scan.ScanQueryLimitRowIterator} - * @see {@link org.apache.druid.query.scan.ScanQueryQueryToolChest.mergeResults} + * @see {@link org.apache.druid.query.scan.ScanQueryQueryToolChest#mergeResults} */ public static Sequence runLimitAndOffset( final QueryPlus queryPlus, @@ -81,33 +81,37 @@ public static Sequence runLimitAndOffset( final ResponseContext responseContext, final ScanQueryConfig scanQueryConfig) { - // Remove "offset" and add it to the "limit" (we won't push the offset down, just apply it here, at the - // merge at the top of the stack). - final ScanQuery originalQuery = ((ScanQuery) (queryPlus.getQuery())); + // Remove "offset" and add it to the "limit" (we won't push the offset + // down, just apply it here, at the merge at the top of the stack). + final ScanQuery originalQuery = (ScanQuery) queryPlus.getQuery(); ScanQuery.verifyOrderByForNativeExecution(originalQuery); + final boolean hasLimit = originalQuery.isLimited(); + final long limit = hasLimit ? originalQuery.getScanRowsLimit() : Long.MAX_VALUE; + final long offset = originalQuery.getScanRowsOffset(); final long newLimit; - if (!originalQuery.isLimited()) { + if (!hasLimit) { // Unlimited stays unlimited. - newLimit = Long.MAX_VALUE; - } else if (originalQuery.getScanRowsLimit() > Long.MAX_VALUE - originalQuery.getScanRowsOffset()) { + newLimit = limit; + } else if (limit > Long.MAX_VALUE - offset) { throw new ISE( - "Cannot apply limit[%d] with offset[%d] due to overflow", - originalQuery.getScanRowsLimit(), - originalQuery.getScanRowsOffset() + "Cannot apply limit [%d] with offset [%d] due to overflow", + limit, + offset ); } else { - newLimit = originalQuery.getScanRowsLimit() + originalQuery.getScanRowsOffset(); + newLimit = limit + offset; } - // Ensure "legacy" is a non-null value, such that all other nodes this query is forwarded to will treat it - // the same way, even if they have different default legacy values. + // Ensure "legacy" is a non-null value, such that all other nodes this + // query is forwarded to will treat it the same way, even if they have + // different default legacy values. final ScanQuery queryToRun = originalQuery.withNonNullLegacy(scanQueryConfig) .withOffset(0) .withLimit(newLimit); - final boolean hasLimit = queryToRun.isLimited(); - final boolean hasOffset = originalQuery.getScanRowsOffset() > 0; + final boolean hasOffset = offset > 0; + final boolean isGrouped = isGrouped(queryToRun); // Short-circuit if no limit or offset. if (!hasLimit && !hasOffset) { @@ -124,35 +128,34 @@ public static Sequence runLimitAndOffset( queryToRun.withOverriddenContext(ImmutableMap.of(ScanQuery.CTX_KEY_OUTERMOST, false)); } QueryPlus historicalQueryPlus = queryPlus.withQuery(historicalQuery); - Operator inputOp = Operators.toOperator( + FragmentContext fragmentContext = queryPlus.fragmentBuilder().context(); + Operator oper = Operators.toOperator( input, historicalQueryPlus); + if (hasOffset) { + oper = new ScanResultOffsetOperator( + fragmentContext, + oper, + offset + ); + } if (hasLimit) { - final ScanQuery limitedQuery = (ScanQuery) historicalQuery; - if (isGrouped(queryToRun)) { - inputOp = new GroupedScanResultLimitOperator( - queryPlus.fragmentBuilder().context(), - inputOp, - limitedQuery.getScanRowsLimit() + if (isGrouped) { + oper = new GroupedScanResultLimitOperator( + fragmentContext, + oper, + limit ); } else { - inputOp = new UngroupedScanResultLimitOperator( - queryPlus.fragmentBuilder().context(), - inputOp, - limitedQuery.getScanRowsLimit(), - limitedQuery.getBatchSize() + oper = new UngroupedScanResultLimitOperator( + fragmentContext, + oper, + limit, + queryToRun.getBatchSize() ); } } - if (hasOffset) { - ScanResultOffsetOperator op = new ScanResultOffsetOperator( - queryPlus.fragmentBuilder().context(), - inputOp, - queryToRun.getScanRowsOffset() - ); - inputOp = op; - } - return Operators.toSequence(inputOp); + return Operators.toSequence(oper); } private static boolean isGrouped(ScanQuery query) @@ -163,12 +166,12 @@ private static boolean isGrouped(ScanQuery query) } /** - * @see {@link org.apache.druid.query.scan.ScanQueryRunnerFactory.mergeRunners} + * @see {@link org.apache.druid.query.scan.ScanQueryRunnerFactory#mergeRunners} */ private static Sequence runConcatMerge( final QueryPlus queryPlus, - final Iterable> queryRunners, - final ResponseContext responseContext) + final Iterable> queryRunners + ) { List> inputs = new ArrayList<>(); for (QueryRunner qr : queryRunners) { @@ -198,8 +201,8 @@ private static Sequence runConcatMerge( } /** - * @see {@link org.apache.druid.query.scan.ScanQueryRunnerFactory.mergeRunners} - * @see {@link org.apache.druid.query.scan.ScanQueryRunnerFactory.nWayMergeAndLimit} + * @see {@link org.apache.druid.query.scan.ScanQueryRunnerFactory#mergeRunners} + * @see {@link org.apache.druid.query.scan.ScanQueryRunnerFactory#nWayMergeAndLimit} */ public static Sequence runMerge( final QueryPlus queryPlus, @@ -218,8 +221,7 @@ public static Sequence runMerge( // Use normal strategy return runConcatMerge( queryPlus, - queryRunners, - responseContext); + queryRunners); } return null; } @@ -229,12 +231,17 @@ public static Sequence runMerge( * query runner. * * @see {@link org.apache.druid.query.scan.ScanQueryRunnerFactory.ScanQueryRunner} + * @see {@link org.apache.druid.query.scan.ScanQueryEngine} */ public static Sequence runScan( final QueryPlus queryPlus, final Segment segment, final ResponseContext responseContext) { + FragmentContext fragmentContext = queryPlus.fragmentBuilder().context(); + if (isTombstone(segment)) { + return Operators.toSequence(new NullOperator<>(fragmentContext)); + } if (!(queryPlus.getQuery() instanceof ScanQuery)) { throw new ISE("Got a [%s] which isn't a %s", queryPlus.getQuery().getClass(), ScanQuery.class); } @@ -247,9 +254,16 @@ public static Sequence runScan( // TODO (paul): Set the timeout at the overall fragment context level. return Operators.toSequence( new ScanQueryOperator( - queryPlus.fragmentBuilder().context(), + fragmentContext, query, - segment)); + segment, + queryPlus.getQueryMetrics())); + } + + private static boolean isTombstone(final Segment segment) + { + QueryableIndex queryableIndex = segment.asQueryableIndex(); + return queryableIndex != null && queryableIndex.isFromTombstone(); } public static Sequence resultsAsArrays( diff --git a/processing/src/test/java/org/apache/druid/query/scan/ScanQueryRunnerTest.java b/processing/src/test/java/org/apache/druid/query/scan/ScanQueryRunnerTest.java index 31a7a454f374..4c22b07337c9 100644 --- a/processing/src/test/java/org/apache/druid/query/scan/ScanQueryRunnerTest.java +++ b/processing/src/test/java/org/apache/druid/query/scan/ScanQueryRunnerTest.java @@ -35,6 +35,7 @@ import org.apache.druid.java.util.common.DateTimes; import org.apache.druid.java.util.common.ISE; import org.apache.druid.java.util.common.Intervals; +import org.apache.druid.java.util.common.Pair; import org.apache.druid.java.util.common.StringUtils; import org.apache.druid.java.util.metrics.StubServiceEmitter; import org.apache.druid.query.DefaultGenericQueryMetricsFactory; @@ -47,7 +48,6 @@ import org.apache.druid.query.QueryRunnerTestHelper; import org.apache.druid.query.QueryTimeoutException; import org.apache.druid.query.TableDataSource; -import org.apache.druid.query.context.DefaultResponseContext; import org.apache.druid.query.context.ResponseContext; import org.apache.druid.query.expression.TestExprMacroTable; import org.apache.druid.query.extraction.MapLookupExtractor; @@ -84,7 +84,6 @@ @RunWith(Parameterized.class) public class ScanQueryRunnerTest extends InitializedNullHandlingTest { - private static final VirtualColumn EXPR_COLUMN = new ExpressionVirtualColumn("expr", "index * 2", ColumnType.LONG, TestExprMacroTable.INSTANCE); @@ -151,7 +150,6 @@ public Object getResult() @Parameterized.Parameters(name = "{0}, legacy = {1}") public static Iterable constructorFeeder() { - return QueryRunnerTestHelper.cartesian( QueryRunnerTestHelper.makeQueryRunners( FACTORY @@ -160,11 +158,11 @@ public static Iterable constructorFeeder() ); } - private final QueryRunner runner; + private final QueryRunner runner; private final boolean legacy; private final List columns; - public ScanQueryRunnerTest(final QueryRunner runner, final boolean legacy) + public ScanQueryRunnerTest(final QueryRunner runner, final boolean legacy) { this.runner = runner; this.legacy = legacy; @@ -204,10 +202,17 @@ private Druids.ScanQueryBuilder newTestQuery() .legacy(legacy); } + protected Pair, ResponseContext> queryPlusPlus(ScanQuery query) + { + ResponseContext responseContext = ResponseContext.createEmpty(); + QueryPlus queryPlus = QueryPlus + .wrap(query); + return Pair.of(queryPlus, responseContext); + } + @Test public void testFullOnSelect() { - ScanQuery query = newTestQuery() .intervals(I_0112_0114) .virtualColumns(EXPR_COLUMN) @@ -222,7 +227,8 @@ public void testFullOnSelect() (obj, lng) -> {}, (metrics) -> {} ).withWaitMeasuredFromNow(); - Iterable results = metricsEmittingQueryRunner.run(QueryPlus.wrap(query)).toList(); + Pair, ResponseContext> qpp = queryPlusPlus(query); + Iterable results = metricsEmittingQueryRunner.run(qpp.lhs, qpp.rhs).toList(); List expectedResults = toExpected( toFullEvents(V_0112_0114), @@ -244,7 +250,8 @@ public void testFullOnSelectAsCompactedList() .resultFormat(ScanQuery.ResultFormat.RESULT_FORMAT_COMPACTED_LIST) .build(); - Iterable results = runner.run(QueryPlus.wrap(query)).toList(); + Pair, ResponseContext> qpp = queryPlusPlus(query); + Iterable results = runner.run(qpp.lhs, qpp.rhs).toList(); List expectedResults = toExpected( toFullEvents(V_0112_0114), @@ -267,7 +274,8 @@ public void testSelectWithUnderscoreUnderscoreTime() ) .build(); - Iterable results = runner.run(QueryPlus.wrap(query)).toList(); + Pair, ResponseContext> qpp = queryPlusPlus(query); + Iterable results = runner.run(qpp.lhs, qpp.rhs).toList(); final List>> expectedEvents = toEvents( new String[]{ @@ -316,7 +324,8 @@ public void testSelectWithDimsAndMets() .columns(QueryRunnerTestHelper.MARKET_DIMENSION, QueryRunnerTestHelper.INDEX_METRIC) .build(); - Iterable results = runner.run(QueryPlus.wrap(query)).toList(); + Pair, ResponseContext> qpp = queryPlusPlus(query); + Iterable results = runner.run(qpp.lhs, qpp.rhs).toList(); List expectedResults = toExpected( toEvents( @@ -353,7 +362,8 @@ public void testSelectWithDimsAndMetsAsCompactedList() .resultFormat(ScanQuery.ResultFormat.RESULT_FORMAT_COMPACTED_LIST) .build(); - Iterable results = runner.run(QueryPlus.wrap(query)).toList(); + Pair, ResponseContext> qpp = queryPlusPlus(query); + Iterable results = runner.run(qpp.lhs, qpp.rhs).toList(); List expectedResults = toExpected( toEvents( @@ -393,7 +403,8 @@ public void testFullOnSelectWithFilterAndLimit() .limit(limit) .build(); - Iterable results = runner.run(QueryPlus.wrap(query)).toList(); + Pair, ResponseContext> qpp = queryPlusPlus(query); + Iterable results = runner.run(qpp.lhs, qpp.rhs).toList(); final List>> events = toEvents( new String[]{ @@ -452,10 +463,12 @@ public void testSelectWithFilterLookupExtractionFn() .columns(QueryRunnerTestHelper.QUALITY_DIMENSION, QueryRunnerTestHelper.INDEX_METRIC) .build(); - Iterable results = runner.run(QueryPlus.wrap(query)).toList(); + Pair, ResponseContext> qpp = queryPlusPlus(query); + Iterable results = runner.run(qpp.lhs, qpp.rhs).toList(); + qpp = queryPlusPlus(query); Iterable resultsOptimize = TOOL_CHEST .postMergeQueryDecoration(TOOL_CHEST.mergeResults(TOOL_CHEST.preMergeQueryDecoration(runner))) - .run(QueryPlus.wrap(query)) + .run(qpp.lhs, qpp.rhs) .toList(); final List>> events = toEvents( @@ -511,7 +524,8 @@ public void testFullSelectNoResults() ) .build(); - Iterable results = runner.run(QueryPlus.wrap(query)).toList(); + Pair, ResponseContext> qpp = queryPlusPlus(query); + Iterable results = runner.run(qpp.lhs, qpp.rhs).toList(); List expectedResults = Collections.emptyList(); @@ -526,7 +540,8 @@ public void testFullSelectNoDimensionAndMetric() .columns("foo", "foo2") .build(); - Iterable results = runner.run(QueryPlus.wrap(query)).toList(); + Pair, ResponseContext> qpp = queryPlusPlus(query); + Iterable results = runner.run(qpp.lhs, qpp.rhs).toList(); final List>> events = toEvents( legacy ? new String[]{getTimestampName() + ":TIME"} : new String[0], @@ -561,7 +576,8 @@ public void testFullOnSelectWithFilterLimitAndAscendingTimeOrderingListFormat() .context(ImmutableMap.of(ScanQuery.CTX_KEY_OUTERMOST, false)) .build(); - Iterable results = runner.run(QueryPlus.wrap(query)).toList(); + Pair, ResponseContext> qpp = queryPlusPlus(query); + Iterable results = runner.run(qpp.lhs, qpp.rhs).toList(); String[] seg1Results = new String[]{ "2011-01-12T00:00:00.000Z\tspot\tautomotive\tpreferred\tapreferred\t100.000000", "2011-01-12T00:00:00.000Z\tspot\tbusiness\tpreferred\tbpreferred\t100.000000", @@ -648,7 +664,8 @@ public void testFullOnSelectWithFilterLimitAndDescendingTimeOrderingListFormat() .order(ScanQuery.Order.DESCENDING) .build(); - Iterable results = runner.run(QueryPlus.wrap(query)).toList(); + Pair, ResponseContext> qpp = queryPlusPlus(query); + Iterable results = runner.run(qpp.lhs, qpp.rhs).toList(); String[] seg1Results = new String[]{ "2011-01-12T00:00:00.000Z\tspot\tautomotive\tpreferred\tapreferred\t100.000000", "2011-01-12T00:00:00.000Z\tspot\tbusiness\tpreferred\tbpreferred\t100.000000", @@ -760,7 +777,8 @@ public void testFullOnSelectWithFilterLimitAndAscendingTimeOrderingCompactedList .limit(limit) .build(); - Iterable results = runner.run(QueryPlus.wrap(query)).toList(); + Pair, ResponseContext> qpp = queryPlusPlus(query); + Iterable results = runner.run(qpp.lhs, qpp.rhs).toList(); final List>> ascendingEvents = toEvents( new String[]{ legacy ? getTimestampName() + ":TIME" : ColumnHolder.TIME_COLUMN_NAME, @@ -850,7 +868,8 @@ public void testFullOnSelectWithFilterLimitAndDescendingTimeOrderingCompactedLis .limit(limit) .build(); - Iterable results = runner.run(QueryPlus.wrap(query)).toList(); + Pair, ResponseContext> qpp = queryPlusPlus(query); + Iterable results = runner.run(qpp.lhs, qpp.rhs).toList(); String[] expectedRet = (String[]) ArrayUtils.addAll(seg1Results, seg2Results); ArrayUtils.reverse(expectedRet); final List>> descendingEvents = toEvents( @@ -908,17 +927,17 @@ public void testScanQueryTimeout() .virtualColumns(EXPR_COLUMN) .context(ImmutableMap.of(QueryContexts.TIMEOUT_KEY, 1)) .build(); - ResponseContext responseContext = DefaultResponseContext.createEmpty(); final long timeoutAt = System.currentTimeMillis(); - responseContext.putTimeoutTime(timeoutAt); + Pair, ResponseContext> qpp = queryPlusPlus(query); + qpp.rhs.putTimeoutTime(timeoutAt); try { - runner.run(QueryPlus.wrap(query), responseContext).toList(); + runner.run(qpp.lhs, qpp.rhs).toList(); Assert.fail("didn't timeout"); } catch (RuntimeException e) { Assert.assertTrue(e instanceof QueryTimeoutException); Assert.assertEquals("Query timeout", ((QueryTimeoutException) e).getErrorCode()); - Assert.assertEquals(timeoutAt, responseContext.getTimeoutTime().longValue()); + Assert.assertEquals(timeoutAt, qpp.rhs.getTimeoutTime().longValue()); } } @@ -931,6 +950,7 @@ public void testScanQueryTimeoutMerge() .context(ImmutableMap.of(QueryContexts.TIMEOUT_KEY, 1)) .build(); try { + Pair, ResponseContext> qpp = queryPlusPlus(query); FACTORY.mergeRunners( DirectQueryProcessingPool.INSTANCE, ImmutableList.of( @@ -942,7 +962,7 @@ public void testScanQueryTimeoutMerge() } return runner.run(queryPlus, responseContext); }) - ).run(QueryPlus.wrap(query), DefaultResponseContext.createEmpty()).toList(); + ).run(qpp.lhs, qpp.rhs).toList(); Assert.fail("didn't timeout"); } @@ -961,6 +981,7 @@ public void testScanQueryTimeoutZeroDoesntTimeOut() .context(ImmutableMap.of(QueryContexts.TIMEOUT_KEY, 0)) .build(); + Pair, ResponseContext> qpp = queryPlusPlus(query); Iterable results = FACTORY.mergeRunners( DirectQueryProcessingPool.INSTANCE, ImmutableList.of( @@ -972,7 +993,7 @@ public void testScanQueryTimeoutZeroDoesntTimeOut() } return runner.run(queryPlus, responseContext); }) - ).run(QueryPlus.wrap(query), DefaultResponseContext.createEmpty()).toList(); + ).run(qpp.lhs, qpp.rhs).toList(); List expectedResults = toExpected( toFullEvents(V_0112_0114), @@ -1148,8 +1169,8 @@ public static void verify( Assert.assertEquals(expected.getSegmentId(), actual.getSegmentId()); - Set exColumns = Sets.newTreeSet(expected.getColumns()); - Set acColumns = Sets.newTreeSet(actual.getColumns()); + Set exColumns = Sets.newTreeSet(expected.getColumns()); + Set acColumns = Sets.newTreeSet(actual.getColumns()); Assert.assertEquals(exColumns, acColumns); Iterator> expectedEvts = ((List>) expected.getEvents()).iterator(); diff --git a/processing/src/test/java/org/apache/druid/query/scan/ScanQueryRunnerWithOperatorsTest.java b/processing/src/test/java/org/apache/druid/query/scan/ScanQueryRunnerWithOperatorsTest.java new file mode 100644 index 000000000000..06e4d3f8da3c --- /dev/null +++ b/processing/src/test/java/org/apache/druid/query/scan/ScanQueryRunnerWithOperatorsTest.java @@ -0,0 +1,54 @@ +/* + * 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.scan; + +import org.apache.druid.java.util.common.Pair; +import org.apache.druid.query.QueryPlus; +import org.apache.druid.query.QueryRunner; +import org.apache.druid.query.context.ResponseContext; +import org.apache.druid.queryng.fragment.FragmentBuilderFactory; +import org.apache.druid.queryng.fragment.TestFragmentBuilderFactory; +import org.junit.runner.RunWith; +import org.junit.runners.Parameterized; + +/** + * Version of {@link ScanQueryRunnerTest} that enables the "Query NG" operator-based + * scan query implementation. + */ +@RunWith(Parameterized.class) +public class ScanQueryRunnerWithOperatorsTest extends ScanQueryRunnerTest +{ + public ScanQueryRunnerWithOperatorsTest(QueryRunner runner, boolean legacy) + { + super(runner, legacy); + } + + private FragmentBuilderFactory fragmentBuilderFactory = new TestFragmentBuilderFactory(true); + + @Override + protected Pair, ResponseContext> queryPlusPlus(ScanQuery query) + { + ResponseContext responseContext = ResponseContext.createEmpty(); + QueryPlus queryPlus = QueryPlus + .wrap(query) + .withFragmentBuilder(fragmentBuilderFactory.create(query, responseContext)); + return Pair.of(queryPlus, responseContext); + } +} diff --git a/processing/src/test/java/org/apache/druid/query/scan/ScanQueryTest.java b/processing/src/test/java/org/apache/druid/query/scan/ScanQueryTest.java index ad292d2fcfa0..0ebd7f669918 100644 --- a/processing/src/test/java/org/apache/druid/query/scan/ScanQueryTest.java +++ b/processing/src/test/java/org/apache/druid/query/scan/ScanQueryTest.java @@ -59,7 +59,7 @@ public static void setup() ArrayList> events1 = new ArrayList<>(); HashMap event1 = new HashMap<>(); - event1.put(ColumnHolder.TIME_COLUMN_NAME, new Long(42)); + event1.put(ColumnHolder.TIME_COLUMN_NAME, Long.valueOf(42)); events1.add(event1); s1 = new ScanResultValue( @@ -70,7 +70,7 @@ public static void setup() ArrayList> events2 = new ArrayList<>(); HashMap event2 = new HashMap<>(); - event2.put(ColumnHolder.TIME_COLUMN_NAME, new Long(43)); + event2.put(ColumnHolder.TIME_COLUMN_NAME, Long.valueOf(43)); events2.add(event2); s2 = new ScanResultValue( @@ -304,7 +304,7 @@ public void testTimeOrderingWithoutTimeColumn() ).flatMerge(seq -> seq, descendingOrderScan.getResultOrdering()); // This should throw an ISE - List res = borkedSequence.toList(); + borkedSequence.toList(); } @Test diff --git a/processing/src/test/java/org/apache/druid/queryng/operators/scan/MockScanResultReader.java b/processing/src/test/java/org/apache/druid/queryng/operators/scan/MockScanResultReader.java index 584b7c3c368f..9bee83b72989 100644 --- a/processing/src/test/java/org/apache/druid/queryng/operators/scan/MockScanResultReader.java +++ b/processing/src/test/java/org/apache/druid/queryng/operators/scan/MockScanResultReader.java @@ -21,6 +21,8 @@ import com.google.common.annotations.VisibleForTesting; import com.google.common.base.Preconditions; +import com.google.common.collect.Maps; +import org.apache.druid.java.util.common.Intervals; import org.apache.druid.java.util.common.StringUtils; import org.apache.druid.query.scan.ScanQuery.ResultFormat; import org.apache.druid.query.scan.ScanResultValue; @@ -33,7 +35,6 @@ import java.time.Duration; import java.time.Instant; import java.util.ArrayList; -import java.util.HashMap; import java.util.List; import java.util.Map; @@ -86,7 +87,7 @@ public MockScanResultReader( columns.add(ColumnHolder.TIME_COLUMN_NAME); } for (int i = 1; i < columnCount; i++) { - columns.add("Column" + Integer.toString(i)); + columns.add("Column" + i); } this.targetCount = targetCount; this.batchSize = batchSize; @@ -106,7 +107,7 @@ public static Interval interval(int offset) Instant base = Instant.parse("2021-10-24T00:00:00Z"); Duration grainOffset = grain.multipliedBy(offset); Instant start = base.plus(grainOffset); - return new Interval(start.toEpochMilli(), start.plus(grain).toEpochMilli()); + return Intervals.utc(start.toEpochMilli(), start.plus(grain).toEpochMilli()); } @Override @@ -156,7 +157,7 @@ private Object listBatch(int n) { List> batch = new ArrayList<>(n); for (int i = 0; i < n; i++) { - Map values = new HashMap<>(columns.size()); + Map values = Maps.newHashMapWithExpectedSize(columns.size()); if (!columns.isEmpty()) { values.put(ColumnHolder.TIME_COLUMN_NAME, nextTs); } diff --git a/server/src/main/java/org/apache/druid/server/QueryLifecycle.java b/server/src/main/java/org/apache/druid/server/QueryLifecycle.java index 700024d4f6b7..9d8d5bb0ef50 100644 --- a/server/src/main/java/org/apache/druid/server/QueryLifecycle.java +++ b/server/src/main/java/org/apache/druid/server/QueryLifecycle.java @@ -282,9 +282,7 @@ public QueryResponse execute() transition(State.AUTHORIZED, State.EXECUTING); final ResponseContext responseContext = DirectDruidClient.makeResponseContextForQuery(); - final FragmentBuilder fragmentBuilder = fragmentContextFactory.create(baseQuery, responseContext); - @SuppressWarnings("rawtypes") final Sequence res = QueryPlus.wrap(baseQuery) .withIdentity(authenticationResult.getIdentity()) .withFragmentBuilder(fragmentBuilder) diff --git a/server/src/main/java/org/apache/druid/server/SetAndVerifyContextQueryRunner.java b/server/src/main/java/org/apache/druid/server/SetAndVerifyContextQueryRunner.java index 579484ddab65..16b163772192 100644 --- a/server/src/main/java/org/apache/druid/server/SetAndVerifyContextQueryRunner.java +++ b/server/src/main/java/org/apache/druid/server/SetAndVerifyContextQueryRunner.java @@ -49,7 +49,7 @@ public SetAndVerifyContextQueryRunner(ServerConfig serverConfig, QueryRunner public Sequence run(QueryPlus queryPlus, ResponseContext responseContext) { return baseRunner.run( - QueryPlus.wrap(withTimeoutAndMaxScatterGatherBytes(queryPlus.getQuery(), serverConfig)), + queryPlus.withQuery(withTimeoutAndMaxScatterGatherBytes(queryPlus.getQuery(), serverConfig)), responseContext ); } diff --git a/sql/src/test/java/org/apache/druid/sql/guice/SqlModuleTest.java b/sql/src/test/java/org/apache/druid/sql/guice/SqlModuleTest.java index 4da433e46e6a..df82acc0dd77 100644 --- a/sql/src/test/java/org/apache/druid/sql/guice/SqlModuleTest.java +++ b/sql/src/test/java/org/apache/druid/sql/guice/SqlModuleTest.java @@ -49,6 +49,7 @@ import org.apache.druid.query.QuerySegmentWalker; import org.apache.druid.query.QueryToolChestWarehouse; import org.apache.druid.query.lookup.LookupExtractorFactoryContainerProvider; +import org.apache.druid.queryng.guice.QueryNGModule; import org.apache.druid.segment.join.JoinableFactory; import org.apache.druid.segment.loading.SegmentLoader; import org.apache.druid.server.QueryScheduler; @@ -148,7 +149,7 @@ public void testDefaultViewManagerBind() Assert.assertNotNull(viewManager); Assert.assertTrue(viewManager instanceof NoopViewManager); } - + @Test public void testNonDefaultViewManagerBind() { @@ -200,7 +201,8 @@ private Injector makeInjectorWithProperties(final Properties props) .in(LazySingleton.class); }, new SqlModule(props), - new TestViewManagerModule() + new TestViewManagerModule(), + new QueryNGModule() ) ); } From 5cdc99304e7b80b66b9927d604b98b7e58bc3b8a Mon Sep 17 00:00:00 2001 From: Paul Rogers Date: Thu, 23 Jun 2022 11:16:43 -0700 Subject: [PATCH 11/11] Test fixes --- .../org/apache/druid/query/QueryPlus.java | 8 ++++++++ .../apache/druid/query/QueryToolChest.java | 2 +- .../query/scan/ScanQueryLimitRowIterator.java | 4 +++- .../query/scan/ScanQueryQueryToolChest.java | 7 ++++--- .../operators/scan/ScanQueryOperator.java | 11 ----------- .../druid/queryng/planner/ScanPlanner.java | 19 +++++++++++-------- .../apache/druid/server/QueryLifecycle.java | 2 +- .../SetAndVerifyContextQueryRunner.java | 5 ++++- .../druid/sql/avatica/DruidStatement.java | 2 -- 9 files changed, 32 insertions(+), 28 deletions(-) diff --git a/processing/src/main/java/org/apache/druid/query/QueryPlus.java b/processing/src/main/java/org/apache/druid/query/QueryPlus.java index e9001a7f3600..d2fe559df65f 100644 --- a/processing/src/main/java/org/apache/druid/query/QueryPlus.java +++ b/processing/src/main/java/org/apache/druid/query/QueryPlus.java @@ -106,6 +106,14 @@ public QueryPlus withQueryMetrics(QueryToolChest> query } } + public QueryPlus withoutMetrics() + { + if (queryMetrics == null) { + return this; + } + return new QueryPlus<>(query, null, identity, fragmentBuilder); + } + /** * Returns a QueryPlus object without the components which are unsafe for concurrent use from multiple threads, * therefore couldn't be passed down in concurrent or async {@link QueryRunner}s. diff --git a/processing/src/main/java/org/apache/druid/query/QueryToolChest.java b/processing/src/main/java/org/apache/druid/query/QueryToolChest.java index 27635066543f..38575fae8b96 100644 --- a/processing/src/main/java/org/apache/druid/query/QueryToolChest.java +++ b/processing/src/main/java/org/apache/druid/query/QueryToolChest.java @@ -331,7 +331,7 @@ public Sequence resultsAsArrays(QueryType query, Sequence throw new UOE("Query type '%s' does not support returning results as arrays", query.getType()); } - @SuppressWarnings("unchecked") + @SuppressWarnings({"unchecked", "unused"}) public Sequence resultsAsArrays(QueryPlus query, Sequence resultSequence) { return resultsAsArrays((QueryType) query.getQuery(), resultSequence); diff --git a/processing/src/main/java/org/apache/druid/query/scan/ScanQueryLimitRowIterator.java b/processing/src/main/java/org/apache/druid/query/scan/ScanQueryLimitRowIterator.java index 1e81b36bc29d..7537136467dd 100644 --- a/processing/src/main/java/org/apache/druid/query/scan/ScanQueryLimitRowIterator.java +++ b/processing/src/main/java/org/apache/druid/query/scan/ScanQueryLimitRowIterator.java @@ -68,7 +68,9 @@ public class ScanQueryLimitRowIterator implements CloseableIterator historicalQuery = queryPlus.getQuery().withOverriddenContext(ImmutableMap.of(ScanQuery.CTX_KEY_OUTERMOST, false)); - Sequence baseSequence = baseRunner.run(queryPlus.withQuery(historicalQuery), responseContext); + // No metrics past this point: metrics are not thread-safe. + QueryPlus wrapped = queryPlus.withQuery(historicalQuery).withoutMetrics(); + Sequence baseSequence = baseRunner.run(wrapped, responseContext); this.yielder = baseSequence.toYielder( null, new YieldingAccumulator() diff --git a/processing/src/main/java/org/apache/druid/query/scan/ScanQueryQueryToolChest.java b/processing/src/main/java/org/apache/druid/query/scan/ScanQueryQueryToolChest.java index 930a6abf0b9a..252e01290443 100644 --- a/processing/src/main/java/org/apache/druid/query/scan/ScanQueryQueryToolChest.java +++ b/processing/src/main/java/org/apache/druid/query/scan/ScanQueryQueryToolChest.java @@ -230,15 +230,16 @@ public Sequence resultsAsArrays(final ScanQuery query, final Sequence< return resultSequence.flatMap( result -> { // Generics? Where we're going, we don't need generics. - final List rows = (List) result.getEvents(); - final Iterable arrays = Iterables.transform(rows, (Function) mapper); + @SuppressWarnings("unchecked") + final List rows = (List) result.getEvents(); + @SuppressWarnings("unchecked") + final Iterable arrays = Iterables.transform(rows, (Function) mapper); return Sequences.simple(arrays); } ); } @Override - @SuppressWarnings("unchecked") public Sequence resultsAsArrays(QueryPlus queryPlus, Sequence resultSequence) { ScanQuery query = (ScanQuery) queryPlus.getQuery(); diff --git a/processing/src/main/java/org/apache/druid/queryng/operators/scan/ScanQueryOperator.java b/processing/src/main/java/org/apache/druid/queryng/operators/scan/ScanQueryOperator.java index a79264348aee..c2f01433adfa 100644 --- a/processing/src/main/java/org/apache/druid/queryng/operators/scan/ScanQueryOperator.java +++ b/processing/src/main/java/org/apache/druid/queryng/operators/scan/ScanQueryOperator.java @@ -212,17 +212,6 @@ private void finish() } } - public enum Limit - { - NONE, - /** - * If we're performing time-ordering, we want to scan through the first `limit` rows in each - * segment ignoring the number of rows already counted on other segments. - */ - LOCAL, - GLOBAL - } - protected final FragmentContext context; private final ScanQuery query; private final Segment segment; diff --git a/processing/src/main/java/org/apache/druid/queryng/planner/ScanPlanner.java b/processing/src/main/java/org/apache/druid/queryng/planner/ScanPlanner.java index c98ee454a7ec..de80aee572fe 100644 --- a/processing/src/main/java/org/apache/druid/queryng/planner/ScanPlanner.java +++ b/processing/src/main/java/org/apache/druid/queryng/planner/ScanPlanner.java @@ -58,18 +58,20 @@ public class ScanPlanner { /** * Sets up an operator over a ScanResultValue operator. Its behaviour - * varies depending on whether the query is returning time-ordered values and whether the CTX_KEY_OUTERMOST - * flag is false. + * varies depending on whether the query is returning time-ordered values + * and whether the CTX_KEY_OUTERMOST flag is false. *

      * Behaviours: *

        - *
      1. No time ordering: expects the child to produce ScanResultValues which each contain up to query.batchSize events. - * The operator will be "done" when the limit of events is reached. The final ScanResultValue might contain + *
      2. No time ordering: expects the child to produce ScanResultValues which + * each contain up to query.batchSize events. The operator will be "done" + * when the limit of events is reached. The final ScanResultValue might contain * fewer than batchSize events so that the limit number of events is returned.
        3. *
      3. Time Ordering, CTX_KEY_OUTERMOST false: Same behaviour as no time ordering.
        4. - *
      4. Time Ordering, CTX_KEY_OUTERMOST=true or null: The child operator in this case should produce ScanResultValues - * that contain only one event each for the CachingClusteredClient n-way merge. This operator will perform - * batching according to query batch size until the limit is reached.
        5. + *
      5. Time Ordering, CTX_KEY_OUTERMOST=true or null: The child operator in this + * case should produce ScanResultValues that contain only one event each for + * the CachingClusteredClient n-way merge. This operator will perform + * batching according to query batch size until the limit is reached.
        6. *
      * * @see {@link org.apache.druid.query.scan.ScanQueryLimitRowIterator} @@ -127,7 +129,8 @@ public static Sequence runLimitAndOffset( historicalQuery = queryToRun.withOverriddenContext(ImmutableMap.of(ScanQuery.CTX_KEY_OUTERMOST, false)); } - QueryPlus historicalQueryPlus = queryPlus.withQuery(historicalQuery); + // No metrics past this point: metrics are not thread-safe. + QueryPlus historicalQueryPlus = queryPlus.withQuery(historicalQuery).withoutMetrics(); FragmentContext fragmentContext = queryPlus.fragmentBuilder().context(); Operator oper = Operators.toOperator( input, diff --git a/server/src/main/java/org/apache/druid/server/QueryLifecycle.java b/server/src/main/java/org/apache/druid/server/QueryLifecycle.java index 9d8d5bb0ef50..92ea54ac9e6f 100644 --- a/server/src/main/java/org/apache/druid/server/QueryLifecycle.java +++ b/server/src/main/java/org/apache/druid/server/QueryLifecycle.java @@ -71,7 +71,7 @@ import java.util.concurrent.TimeUnit; /** - * Class that helps a Druid server (broker, historical, etc) manage the + * Helps a Druid server (broker, historical, etc) manage the * lifecycle of a query that it is handling. It ensures that a query goes * through the following stages, in the proper order: * diff --git a/server/src/main/java/org/apache/druid/server/SetAndVerifyContextQueryRunner.java b/server/src/main/java/org/apache/druid/server/SetAndVerifyContextQueryRunner.java index 16b163772192..3318042bda34 100644 --- a/server/src/main/java/org/apache/druid/server/SetAndVerifyContextQueryRunner.java +++ b/server/src/main/java/org/apache/druid/server/SetAndVerifyContextQueryRunner.java @@ -48,8 +48,11 @@ public SetAndVerifyContextQueryRunner(ServerConfig serverConfig, QueryRunner @Override public Sequence run(QueryPlus queryPlus, ResponseContext responseContext) { + // No metrics past this point: metrics are not thread-safe. return baseRunner.run( - queryPlus.withQuery(withTimeoutAndMaxScatterGatherBytes(queryPlus.getQuery(), serverConfig)), + queryPlus + .withQuery(withTimeoutAndMaxScatterGatherBytes(queryPlus.getQuery(), serverConfig)) + .withoutMetrics(), responseContext ); } diff --git a/sql/src/main/java/org/apache/druid/sql/avatica/DruidStatement.java b/sql/src/main/java/org/apache/druid/sql/avatica/DruidStatement.java index b3c7e41284e0..cd54bf2c8c91 100644 --- a/sql/src/main/java/org/apache/druid/sql/avatica/DruidStatement.java +++ b/sql/src/main/java/org/apache/druid/sql/avatica/DruidStatement.java @@ -207,7 +207,6 @@ public DruidStatement prepare( } } - public DruidStatement execute(List parameters) { synchronized (lock) { @@ -385,7 +384,6 @@ private AvaticaParameter createParameter(RelDataTypeField field, RelDataType typ ); } - private DruidStatement closeAndPropagateThrowable(Throwable t) { this.throwable = t;