diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/LogicalResultSinkToShortCircuitPointQuery.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/LogicalResultSinkToShortCircuitPointQuery.java index f839b81b6011e2..51bdc44b66b9c6 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/LogicalResultSinkToShortCircuitPointQuery.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/LogicalResultSinkToShortCircuitPointQuery.java @@ -31,6 +31,7 @@ import org.apache.doris.nereids.trees.plans.logical.LogicalFilter; import org.apache.doris.nereids.trees.plans.logical.LogicalOlapScan; import org.apache.doris.qe.ConnectContext; +import org.apache.doris.qe.ConnectContext.ConnectType; import com.google.common.annotations.VisibleForTesting; import com.google.common.collect.ImmutableList; @@ -64,7 +65,20 @@ private boolean filterMatchShortCircuitCondition(LogicalFilter @VisibleForTesting boolean scanMatchShortCircuitCondition(LogicalOlapScan olapScan) { - if (!ConnectContext.get().getSessionVariable().isEnableShortCircuitQuery()) { + ConnectContext connectContext = ConnectContext.get(); + if (!connectContext.getSessionVariable().isEnableShortCircuitQuery()) { + return false; + } + // The short circuit produces no Arrow result at either end. PointQueryExecutor is not a + // Coordinator, and Coordinator/NereidsCoordinator are the only places that register a + // FlightSqlEndpointsLocation, so GetFlightInfo found none and failed the query with + // "no FlightSqlEndpointsLocations"; the BE side cannot be pointed at either, since the lookup rpc + // serializes with VMysqlResultWriter into PTabletKeyLookupResponse.row_batch and never creates the + // ArrowFlightResultBlockBuffer that fetch_arrow_flight_schema looks up. Keep Arrow Flight SQL on + // the normal execution path. This has to be decided here at plan time rather than when picking the + // executor: OlapScanNode.computeTabletInfo and several rewrite and property rules read + // StatementContext.isShortCircuitQuery() while building the plan. See #67368. + if (connectContext.getConnectType() == ConnectType.ARROW_FLIGHT_SQL) { return false; } // Lazy point-query pruning does not preserve explicit PARTITION/TABLET restrictions. diff --git a/fe/fe-core/src/main/java/org/apache/doris/qe/StmtExecutor.java b/fe/fe-core/src/main/java/org/apache/doris/qe/StmtExecutor.java index a9a38a884d5823..8a733fb47ab176 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/qe/StmtExecutor.java +++ b/fe/fe-core/src/main/java/org/apache/doris/qe/StmtExecutor.java @@ -1550,8 +1550,10 @@ public void executeAndSendResult(boolean isOutfileQuery, boolean isSendFields, // need deferral (the BE buffers their result independently) but are captured by the // same gate; the trade-off is their coordinator, query queue slot and query // registration stay held until the next query / teardown instead of being released - // at the end of GetFlightInfo. Point queries use a different coordBase (not - // deferred). See #62259. + // at the end of GetFlightInfo. A short-circuit point query is the one case with a + // different coordBase, and it can no longer reach here: it has no Arrow result on + // either side, so LogicalResultSinkToShortCircuitPointQuery keeps Arrow Flight SQL + // on the normal execution path. See #62259 and #67368. if (coordBase == coord) { deferredForArrowFlight = true; context.addFlightSqlDeferredExecutor(this); diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/ShortCircuitPointQueryTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/ShortCircuitPointQueryTest.java index 1b25c4728d76f0..a94f6ae46f2c66 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/ShortCircuitPointQueryTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/ShortCircuitPointQueryTest.java @@ -29,11 +29,15 @@ import org.apache.doris.nereids.trees.plans.logical.LogicalOlapScan; import org.apache.doris.nereids.util.MemoPatternMatchSupported; import org.apache.doris.nereids.util.PlanChecker; +import org.apache.doris.qe.ConnectContext; +import org.apache.doris.qe.ConnectContext.ConnectType; import org.apache.doris.utframe.TestWithFeService; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; +import java.lang.reflect.Field; + /** * Regression test: * For short-circuit point query, we should not rewrite LogicalOlapScan to LogicalEmptyRelation @@ -154,6 +158,34 @@ void testRemoteOlapTableDoesNotUseShortCircuit() throws Exception { .scanMatchShortCircuitCondition(scan)); } + @Test + void testArrowFlightSqlConnectionDoesNotUseShortCircuit() throws Exception { + // The short circuit hands its rows back through PointQueryExecutor, which registers no + // FlightSqlEndpointsLocation and leaves no Arrow result on the BE, so GetFlightInfo used to fail + // with "fetch arrow flight schema failed, no FlightSqlEndpointsLocations" and drop the row. + // An Arrow Flight SQL connection has to plan the normal execution path. See #67368. + Field connectTypeField = ConnectContext.class.getDeclaredField("connectType"); + connectTypeField.setAccessible(true); + ConnectType originConnectType = (ConnectType) connectTypeField.get(connectContext); + try { + connectTypeField.set(connectContext, ConnectType.ARROW_FLIGHT_SQL); + Plan plan = rewrite("select * from tbl_point_query where `key` = 1"); + + Assertions.assertFalse(connectContext.getStatementContext().isShortCircuitQuery()); + // And the plan really is the ordinary one: tbl_point_query is empty, so it prunes to a + // LogicalEmptyRelation, which is exactly what the short circuit suppresses in + // testShortCircuitPointQueryKeepOlapScanWhenTableEmpty above. + Assertions.assertTrue(plan.anyMatch(p -> p instanceof LogicalEmptyRelation)); + Assertions.assertFalse(plan.anyMatch(p -> p instanceof LogicalOlapScan)); + } finally { + connectTypeField.set(connectContext, originConnectType); + } + + // The very same statement still short circuits on a MySQL connection. + rewrite("select * from tbl_point_query where `key` = 1"); + Assertions.assertTrue(connectContext.getStatementContext().isShortCircuitQuery()); + } + private long getTabletId(String partitionName) throws Exception { Database database = Env.getCurrentInternalCatalog().getDbOrMetaException("test"); OlapTable table = (OlapTable) database.getTableOrMetaException("tbl_partitioned_point_query"); diff --git a/regression-test/suites/arrow_flight_sql_p0/test_point_query_over_arrow_flight.groovy b/regression-test/suites/arrow_flight_sql_p0/test_point_query_over_arrow_flight.groovy new file mode 100644 index 00000000000000..a20b3bc6d4565f --- /dev/null +++ b/regression-test/suites/arrow_flight_sql_p0/test_point_query_over_arrow_flight.groovy @@ -0,0 +1,131 @@ +// 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. + +import org.apache.doris.regression.util.JdbcUtils + +// Regression for https://github.com/apache/doris/issues/67368 +// +// A UNIQUE KEY point query that matches the short circuit is executed by PointQueryExecutor instead of +// a Coordinator, and neither end of that path can serve an Arrow Flight result: +// +// * Coordinator and NereidsCoordinator are the only places that register a FlightSqlEndpointsLocation, +// so GetFlightInfo found no endpoint and failed with +// "fetch arrow flight schema failed, no FlightSqlEndpointsLocations", dropping the row. +// * The BE could not be pointed at either. tablet_fetch_data serializes with VMysqlResultWriter into +// PTabletKeyLookupResponse.row_batch and runs no fragment, so the ArrowFlightResultBlockBuffer that +// fetch_arrow_flight_schema looks up by finst id never exists. +// +// The fix keeps Arrow Flight SQL connections on the normal execution path, decided at plan time in +// LogicalResultSinkToShortCircuitPointQuery. The table below is the one from the issue. +suite("test_point_query_over_arrow_flight") { + def mysqlConn = context.getConn() + def flightConn = context.getArrowFlightSqlConnection() + + def runOnMysql = { String stmt -> + def (result, meta) = JdbcUtils.executeToList(mysqlConn, stmt) + return result + } + def runOnFlight = { String stmt -> + def (result, meta) = JdbcUtils.executeToList(flightConn, stmt) + return result + } + // The suite level explain{} action always runs on the MySQL connection, but the whole point here is + // which protocol asked for the plan, so read the explain text off each connection explicitly. + def explainOn = { conn, String stmt -> + def (rows, meta) = JdbcUtils.executeToList(conn, "explain " + stmt) + return rows.collect { row -> row.get(0).toString() }.join("\n") + } + + def dbName = context.dbName + runOnMysql "USE `${dbName}`" + runOnFlight "USE `${dbName}`" + + def tblName = "test_point_query_over_arrow_flight_tbl" + runOnMysql "DROP TABLE IF EXISTS ${tblName}" + runOnMysql """ + CREATE TABLE ${tblName} ( + `col1` SMALLINT NOT NULL, + `col2` INT NOT NULL, + `loc3` CHAR(10) NOT NULL, + `value` CHAR(10) NOT NULL, + INDEX col3 (`loc3`) USING INVERTED, + INDEX col2_idx (`col2`) USING INVERTED + ) ENGINE=OLAP + UNIQUE KEY(`col1`, `col2`, `loc3`) + DISTRIBUTED BY HASH(`col1`, `col2`, `loc3`) BUCKETS 1 + PROPERTIES ( + "replication_allocation" = "tag.location.default: 1", + "disable_auto_compaction" = "true", + "bloom_filter_columns" = "col1", + "store_row_column" = "true", + "enable_mow_light_delete" = "false" + ) + """ + runOnMysql "INSERT INTO ${tblName} VALUES (10, 20, 'aabc', 'value')" + + def pointQuery = "SELECT * FROM ${tblName} WHERE col1 = 10 AND col2 = 20 AND loc3 = 'aabc'" + + // The short circuit is still taken on a MySQL connection: the fix is scoped to one protocol, it does + // not disable the optimization. Assert this first, so a table that stopped qualifying for the short + // circuit (schema or session variable drift) fails loudly here instead of making the flight + // assertions below pass for the wrong reason. + def mysqlExplain = explainOn(mysqlConn, pointQuery) + assertTrue(mysqlExplain.contains("SHORT-CIRCUIT"), + "the point query must still short circuit on a mysql connection, but got:\n" + mysqlExplain) + + // The same statement must be planned on the normal path over Arrow Flight SQL. + def flightExplain = explainOn(flightConn, pointQuery) + assertFalse(flightExplain.contains("SHORT-CIRCUIT"), + "the point query must not short circuit on an arrow flight connection, but got:\n" + flightExplain) + + // This is the call that used to fail with "no FlightSqlEndpointsLocations". + def (flightRows, flightMeta) = JdbcUtils.executeToList(flightConn, pointQuery) + assertEquals(1, flightRows.size()) + assertEquals(10, flightRows[0][0] as int) + assertEquals(20, flightRows[0][1] as int) + assertEquals("aabc", flightRows[0][2].toString()) + assertEquals("value", flightRows[0][3].toString()) + + // Both protocols must see the same row, one through the short circuit and one through the normal + // plan. + def mysqlRows = runOnMysql(pointQuery) + assertEquals(1, mysqlRows.size()) + assertEquals(mysqlRows[0].collect { it.toString() }, flightRows[0].collect { it.toString() }) + + // The BE produces the arrow batch, so the column types survive. Serving the point query result from + // the FE instead would hand every column back as a string, because FlightSqlChannel.addResult builds + // varchar vectors only. + assertTrue(flightRows[0][0] instanceof Number, + "col1 must stay numeric over arrow flight, but got: " + flightRows[0][0].getClass()) + assertTrue(flightRows[0][1] instanceof Number, + "col2 must stay numeric over arrow flight, but got: " + flightRows[0][1].getClass()) + assertEquals(4, flightMeta.getColumnCount()) + + // A key that matches no row is planned the same way and used to fail with the same error, so it is + // not enough for the statement above to be the only shape that works. + def emptyRows = runOnFlight("SELECT * FROM ${tblName} WHERE col1 = 11 AND col2 = 20 AND loc3 = 'aabc'") + assertEquals(0, emptyRows.size()) + + // The workaround reported in the issue keeps working, and a plain non point query on the same table + // is unaffected. + def hintRows = runOnFlight("SELECT /*+ SET_VAR(enable_short_circuit_query=false) */ * FROM ${tblName} " + + "WHERE col1 = 10 AND col2 = 20 AND loc3 = 'aabc'") + assertEquals(1, hintRows.size()) + assertEquals(1, runOnFlight("SELECT col1 FROM ${tblName} ORDER BY col1").size()) + + runOnMysql "DROP TABLE IF EXISTS ${tblName}" +}