Line data Source code
1 : // Route classifier tests: subqueries, hand-built nodes, and edge cases.
2 :
3 : #include "backend/engine/coordinator/route_classifier_test_fixture.h"
4 : #include "backend/engine/disposition.h"
5 : #include "googlesql/public/id_string.h"
6 : #include "googlesql/public/value.h"
7 : #include "googlesql/resolved_ast/resolved_column.h"
8 :
9 : namespace bigquery_emulator {
10 : namespace backend {
11 : namespace engine {
12 : namespace coordinator {
13 : namespace {
14 1 : TEST_F(RouteClassifierTest, UncorrelatedSubqueryExprStaysOnFastPath) {
15 : // `WHERE id IN (SELECT n FROM UNNEST([1, 2, 3]) AS n)` is a
16 : // non-correlated IN subquery: the inner SELECT does not
17 : // reference any column from the outer scan. The analyzer marks
18 : // this by leaving `parameter_list()` empty. The transpiler's
19 : // `EmitSubqueryExpr` lowers it directly to DuckDB's
20 : // `(<lhs> IN (<sub>))` shape, so the classifier MUST keep the
21 : // query on `kDuckdbNative` -- a false promotion would force the
22 : // semantic executor to run a shape the fast path handles
23 : // correctly. See `docs/ENGINE_POLICY.md` Family 3.
24 1 : const auto* stmt = Analyze(
25 1 : "SELECT id FROM people "
26 1 : "WHERE id IN (SELECT n FROM UNNEST([1, 2, 3]) AS n)");
27 1 : ASSERT_NE(stmt, nullptr);
28 1 : RouteDecision d = classifier_.Classify(*stmt);
29 1 : EXPECT_EQ(d.disposition, Disposition::kDuckdbNative);
30 2 : EXPECT_TRUE(d.offending_node.empty()) << d.offending_node;
31 1 : }
32 :
33 : TEST_F(RouteClassifierTest,
34 1 : CorrelatedScalarSubqueryExprPromotesToSemanticExecutor) {
35 : // `(SELECT COUNT(*) FROM <inner> WHERE <inner>.k = outer.k)` is
36 : // a correlated scalar subquery: the inner WHERE clause
37 : // references the outer scan's column. The analyzer marks the
38 : // referenced outer columns in `ResolvedSubqueryExpr::parameter_list()`,
39 : // which the classifier inspects via `VisitResolvedSubqueryExpr`.
40 : // Promotion to `kSemanticExecutor` is mandatory: DuckDB's
41 : // correlated-subquery decorrelation does not guarantee BigQuery
42 : // per-outer-row evaluation order for every shape, and the only
43 : // way to avoid silent approximation is to evaluate the inner
44 : // subquery once per outer row in the local interpreter.
45 : //
46 : // The semantic executor's correlated-subquery evaluator is
47 : // `docs/ENGINE_POLICY.md` Family 4 (deferred to a
48 : // follow-up subagent); until it lands the gateway surfaces
49 : // UNIMPLEMENTED via the executor stub. That is the same
50 : // end-user-visible outcome the fast path's empty-string
51 : // contract would produce, but going through the classifier
52 : // means the disposition is a deliberate route choice, not a
53 : // surprise transpiler bailout.
54 1 : const auto* stmt = Analyze(
55 1 : "SELECT (SELECT COUNT(*) FROM people AS p WHERE p.id = people.id) AS c "
56 1 : "FROM people");
57 1 : ASSERT_NE(stmt, nullptr);
58 1 : RouteDecision d = classifier_.Classify(*stmt);
59 1 : EXPECT_EQ(d.disposition, Disposition::kSemanticExecutor);
60 1 : EXPECT_EQ(d.offending_node, "ResolvedSubqueryExpr(correlated)");
61 1 : }
62 :
63 : TEST_F(RouteClassifierTest,
64 1 : CorrelatedExistsSubqueryExprPromotesToSemanticExecutor) {
65 : // `EXISTS (SELECT 1 FROM <inner> WHERE <inner>.k = outer.k)` is
66 : // the most common correlated subquery shape (semi-join
67 : // expression). The classifier promotes it for the same reason
68 : // as the scalar case: per-outer-row evaluation order is
69 : // BigQuery-defined, not DuckDB's call. Self-join the only
70 : // available table (`people`) so the test does not depend on a
71 : // second catalog table. The semantic executor evaluator is
72 : // Family 4.
73 1 : const auto* stmt = Analyze(
74 1 : "SELECT id FROM people "
75 1 : "WHERE EXISTS (SELECT 1 FROM people AS p WHERE p.id = people.id)");
76 1 : ASSERT_NE(stmt, nullptr);
77 1 : RouteDecision d = classifier_.Classify(*stmt);
78 1 : EXPECT_EQ(d.disposition, Disposition::kSemanticExecutor);
79 1 : EXPECT_EQ(d.offending_node, "ResolvedSubqueryExpr(correlated)");
80 1 : }
81 :
82 1 : TEST_F(RouteClassifierTest, BarrierScanPromotesToSemanticExecutor) {
83 : // `docs/ENGINE_POLICY.md` Family 2. A
84 : // `ResolvedBarrierScan` is a pipe-operator optimizer marker that
85 : // blocks fusion across its boundary. DuckDB has no analog
86 : // contract, so the classifier MUST promote any query containing
87 : // one to `kSemanticExecutor` (the row's YAML disposition).
88 : //
89 : // We exercise the YAML row directly via a hand-built statement
90 : // because the analyzer does not emit `ResolvedBarrierScan` for
91 : // surface SQL today (pipe operators are an analyzer feature flag).
92 : // The hand-built shape mirrors what the analyzer emits for
93 : // `<expr> |> BARRIER` once the flag flips.
94 1 : auto single = ::googlesql::MakeResolvedSingleRowScan();
95 1 : auto barrier = ::googlesql::MakeResolvedBarrierScan(
96 1 : /*column_list=*/{}, std::move(single));
97 1 : ::googlesql::ResolvedColumn out_col(
98 1 : /*column_id=*/200,
99 1 : /*table_name=*/::googlesql::IdString::MakeGlobal("$query"),
100 1 : /*name=*/::googlesql::IdString::MakeGlobal("c"),
101 1 : type_factory_->get_int64());
102 1 : std::vector<std::unique_ptr<const ::googlesql::ResolvedComputedColumn>> exprs;
103 1 : exprs.push_back(::googlesql::MakeResolvedComputedColumn(
104 1 : out_col, ::googlesql::MakeResolvedLiteral(::googlesql::Value::Int64(7))));
105 1 : auto project = ::googlesql::MakeResolvedProjectScan(
106 1 : /*column_list=*/{out_col}, std::move(exprs), std::move(barrier));
107 1 : std::vector<std::unique_ptr<const ::googlesql::ResolvedOutputColumn>> outputs;
108 1 : outputs.push_back(::googlesql::MakeResolvedOutputColumn("c", out_col));
109 1 : auto query_stmt = ::googlesql::MakeResolvedQueryStmt(
110 1 : std::move(outputs), /*is_value_table=*/false, std::move(project));
111 :
112 1 : RouteDecision d = classifier_.Classify(*query_stmt);
113 1 : EXPECT_EQ(d.disposition, Disposition::kSemanticExecutor);
114 1 : EXPECT_EQ(d.offending_node, "ResolvedBarrierScan");
115 1 : }
116 :
117 1 : TEST_F(RouteClassifierTest, PivotScanRoutesToDuckdbRewrite) {
118 : // `docs/ENGINE_POLICY.md` Family 3. The engine
119 : // disables `REWRITE_PIVOT` so the analyzer hands us a raw
120 : // `ResolvedPivotScan`; the disposition table routes it through
121 : // `kDuckdbRewrite`, and the transpiler's `EmitPivotScan` lowers
122 : // it to DuckDB conditional aggregation (FILTER).
123 1 : const ::googlesql::ResolvedStatement* stmt =
124 1 : Analyze("SELECT * FROM people PIVOT(COUNT(*) FOR name IN ('a', 'b'))");
125 1 : ASSERT_NE(stmt, nullptr);
126 1 : RouteDecision d = classifier_.Classify(*stmt);
127 1 : EXPECT_EQ(d.disposition, Disposition::kDuckdbRewrite);
128 1 : }
129 :
130 1 : TEST_F(RouteClassifierTest, UnpivotScanRoutesToDuckdbRewrite) {
131 : // Same as the PIVOT test above but for `ResolvedUnpivotScan`.
132 : // The engine disables `REWRITE_UNPIVOT`; the disposition table
133 : // routes it through `kDuckdbRewrite`; the transpiler's
134 : // `EmitUnpivotScan` lowers it to UNION ALL.
135 1 : const ::googlesql::ResolvedStatement* stmt =
136 1 : Analyze("SELECT * FROM people UNPIVOT(value FOR label IN (id))");
137 1 : ASSERT_NE(stmt, nullptr);
138 1 : RouteDecision d = classifier_.Classify(*stmt);
139 1 : EXPECT_EQ(d.disposition, Disposition::kDuckdbRewrite);
140 1 : }
141 :
142 1 : TEST_F(RouteClassifierTest, RecursiveScanRoutesToDuckdbRewrite) {
143 : // `docs/ENGINE_POLICY.md` Family 4. The disposition
144 : // table routes `ResolvedRecursiveScan` (and its
145 : // `ResolvedRecursiveRefScan` reference) through `kDuckdbRewrite`;
146 : // the transpiler's `EmitRecursiveScan` lowers it to DuckDB's
147 : // `WITH RECURSIVE`.
148 1 : const ::googlesql::ResolvedStatement* stmt = Analyze(
149 1 : "WITH RECURSIVE r AS ("
150 1 : " SELECT 1 AS n"
151 1 : " UNION ALL"
152 1 : " SELECT n FROM r"
153 1 : ")"
154 1 : "SELECT n FROM r");
155 1 : ASSERT_NE(stmt, nullptr);
156 1 : RouteDecision d = classifier_.Classify(*stmt);
157 1 : EXPECT_EQ(d.disposition, Disposition::kDuckdbRewrite);
158 1 : }
159 :
160 1 : TEST_F(RouteClassifierTest, DeferredComputedColumnStaysOnDuckdbNative) {
161 : // R17 follow-up: ResolvedDeferredComputedColumn is duckdb_native.
162 : // Hand-build an AggregateScan whose aggregate_list holds a
163 : // DeferredComputedColumn so the classifier resolves via YAML lookup
164 : // without depending on analyzer side-effect rewriting.
165 1 : ::googlesql::ResolvedColumn out_col(
166 1 : /*column_id=*/300,
167 1 : /*table_name=*/::googlesql::IdString::MakeGlobal("$query"),
168 1 : /*name=*/::googlesql::IdString::MakeGlobal("v"),
169 1 : type_factory_->get_int64());
170 1 : ::googlesql::ResolvedColumn side_col(
171 1 : /*column_id=*/301,
172 1 : /*table_name=*/::googlesql::IdString::MakeGlobal("$query"),
173 1 : /*name=*/::googlesql::IdString::MakeGlobal("_se"),
174 1 : type_factory_->get_bytes());
175 1 : std::vector<std::unique_ptr<const ::googlesql::ResolvedComputedColumnBase>>
176 1 : aggregate_list;
177 1 : aggregate_list.push_back(::googlesql::MakeResolvedDeferredComputedColumn(
178 1 : out_col,
179 1 : ::googlesql::MakeResolvedLiteral(::googlesql::Value::Int64(0)),
180 1 : side_col));
181 1 : auto agg_scan = ::googlesql::MakeResolvedAggregateScan(
182 1 : /*column_list=*/{out_col},
183 1 : ::googlesql::MakeResolvedSingleRowScan(),
184 1 : /*group_by_list=*/{},
185 1 : std::move(aggregate_list),
186 1 : /*grouping_set_list=*/{},
187 1 : /*rollup_column_list=*/{});
188 1 : std::vector<std::unique_ptr<const ::googlesql::ResolvedOutputColumn>> outputs;
189 1 : outputs.push_back(::googlesql::MakeResolvedOutputColumn("v", out_col));
190 1 : auto query_stmt = ::googlesql::MakeResolvedQueryStmt(
191 1 : std::move(outputs), /*is_value_table=*/false, std::move(agg_scan));
192 :
193 1 : RouteDecision d = classifier_.Classify(*query_stmt);
194 2 : EXPECT_EQ(d.disposition, Disposition::kDuckdbNative) << d.offending_node;
195 2 : EXPECT_NE(d.disposition, Disposition::kSemanticExecutor) << d.offending_node;
196 1 : }
197 :
198 1 : TEST_F(RouteClassifierTest, DifferentialPrivacyAggregateScanRoutesToLocalStub) {
199 1 : const auto* stmt = Analyze(
200 1 : "SELECT WITH DIFFERENTIAL_PRIVACY "
201 1 : "OPTIONS(epsilon=10, delta=0.01, privacy_unit_column=id) "
202 1 : "name, COUNT(*) AS c FROM people GROUP BY name");
203 1 : ASSERT_NE(stmt, nullptr);
204 :
205 1 : RouteDecision d = classifier_.Classify(*stmt);
206 1 : EXPECT_EQ(d.disposition, Disposition::kLocalStub);
207 1 : EXPECT_EQ(d.offending_node, "ResolvedDifferentialPrivacyAggregateScan");
208 2 : EXPECT_NE(d.reason.find("local-stub"), std::string::npos)
209 2 : << "reason should mention the local-stub route; got: " << d.reason;
210 1 : }
211 :
212 : TEST_F(RouteClassifierTest,
213 1 : AggregationThresholdAggregateScanRoutesToLocalStub) {
214 1 : const auto* stmt = Analyze(
215 1 : "SELECT WITH AGGREGATION_THRESHOLD "
216 1 : "OPTIONS(threshold=1, privacy_unit_column=id) "
217 1 : "name, COUNT(*) AS c FROM people GROUP BY name");
218 1 : ASSERT_NE(stmt, nullptr);
219 :
220 1 : RouteDecision d = classifier_.Classify(*stmt);
221 1 : EXPECT_EQ(d.disposition, Disposition::kLocalStub);
222 1 : EXPECT_EQ(d.offending_node, "ResolvedAggregationThresholdAggregateScan");
223 1 : }
224 :
225 1 : TEST_F(RouteClassifierTest, AnonymizedAggregateScanRoutesToLocalStub) {
226 1 : const auto* stmt = Analyze(
227 1 : "SELECT WITH ANONYMIZATION "
228 1 : "OPTIONS(k_threshold=1, epsilon=10) "
229 1 : "name, COUNT(*) AS c FROM people GROUP BY name");
230 1 : ASSERT_NE(stmt, nullptr);
231 :
232 1 : RouteDecision d = classifier_.Classify(*stmt);
233 1 : EXPECT_EQ(d.disposition, Disposition::kLocalStub);
234 1 : EXPECT_EQ(d.offending_node, "ResolvedAnonymizedAggregateScan");
235 1 : }
236 :
237 1 : TEST_F(RouteClassifierTest, MlPredictRoutesToLocalStub) {
238 1 : const auto* stmt = Analyze(
239 1 : "SELECT * FROM ML.PREDICT(MODEL `ds.unregistered_model`, "
240 1 : "(SELECT 1.0 AS f1))");
241 1 : ASSERT_NE(stmt, nullptr);
242 :
243 1 : RouteDecision d = classifier_.Classify(*stmt);
244 1 : EXPECT_EQ(d.disposition, Disposition::kLocalStub);
245 1 : EXPECT_EQ(d.offending_node, "function:ml.predict");
246 2 : EXPECT_NE(d.reason.find("local-stub"), std::string::npos)
247 2 : << "reason should mention the local-stub route; got: " << d.reason;
248 1 : }
249 :
250 1 : TEST_F(RouteClassifierTest, KeysFunctionRoutesToLocalStub) {
251 : // `KEYS.NEW_KEYSET('AEAD_AES_GCM_256')` is a `local_stub` row in
252 : // `functions.yaml` per `docs/ENGINE_POLICY.md`. A
253 : // SELECT referencing it must promote the route to `kLocalStub`
254 : // (above `kSemanticExecutor`, below `kUnsupported`) so the
255 : // coordinator dispatches into the semantic executor's per-family
256 : // stub handler (`backend/engine/semantic/stubs/keys.cc`). The
257 : // offending node carries the function name so a future
258 : // gateway-side error envelope can attribute the stub to the
259 : // right family.
260 1 : const auto* stmt = Analyze("SELECT KEYS.NEW_KEYSET('AEAD_AES_GCM_256')");
261 1 : ASSERT_NE(stmt, nullptr);
262 :
263 1 : RouteDecision d = classifier_.Classify(*stmt);
264 1 : EXPECT_EQ(d.disposition, Disposition::kLocalStub);
265 1 : EXPECT_EQ(d.offending_node, "function:keys.new_keyset");
266 2 : EXPECT_NE(d.reason.find("local-stub"), std::string::npos)
267 2 : << "reason should mention the local-stub route; got: " << d.reason;
268 1 : }
269 :
270 1 : TEST_F(RouteClassifierTest, LocalStubOutranksSemanticExecutorInSameQuery) {
271 : // When a `local_stub` function (`KEYS.NEW_KEYSET`) and a
272 : // `semantic_executor` function (`APPROX_QUANTILES`) appear together,
273 : // the local-stub promotion wins (priority 5 > 4).
274 1 : const auto* stmt = Analyze(
275 1 : "SELECT APPROX_QUANTILES(id, 4) AS q, "
276 1 : "KEYS.NEW_KEYSET('AEAD_AES_GCM_256') AS k FROM people");
277 1 : ASSERT_NE(stmt, nullptr);
278 :
279 1 : RouteDecision d = classifier_.Classify(*stmt);
280 1 : EXPECT_EQ(d.disposition, Disposition::kLocalStub);
281 1 : EXPECT_EQ(d.offending_node, "function:keys.new_keyset");
282 1 : }
283 :
284 1 : TEST_F(RouteClassifierTest, MatchRecognizeScanPromotesToSemanticExecutor) {
285 1 : const auto* stmt = Analyze(
286 1 : "SELECT m FROM people MATCH_RECOGNIZE("
287 1 : "ORDER BY id MEASURES COUNT(*) AS m PATTERN (A) DEFINE A AS id > 0)");
288 1 : ASSERT_NE(stmt, nullptr);
289 :
290 1 : RouteDecision d = classifier_.Classify(*stmt);
291 1 : EXPECT_EQ(d.disposition, Disposition::kSemanticExecutor);
292 1 : EXPECT_EQ(d.offending_node, "ResolvedMatchRecognizeScan");
293 1 : }
294 :
295 1 : TEST_F(RouteClassifierTest, ExplainStatementRoutesToUnsupported) {
296 : // EXPLAIN is not a BigQuery statement (bq dry-run: "Statement not
297 : // supported: ExplainStatement"); the analyzer here still parses it
298 : // because the test fixture enables all statement kinds, but the
299 : // classifier must route it to the deliberate `unsupported` envelope.
300 1 : const auto* stmt = Analyze("EXPLAIN SELECT * FROM people");
301 1 : ASSERT_NE(stmt, nullptr);
302 :
303 1 : RouteDecision d = classifier_.Classify(*stmt);
304 1 : EXPECT_EQ(d.disposition, Disposition::kUnsupported);
305 1 : EXPECT_EQ(d.offending_node, "ResolvedExplainStmt");
306 1 : }
307 : } // namespace
308 : } // namespace coordinator
309 : } // namespace engine
310 : } // namespace backend
311 : } // namespace bigquery_emulator
|
