Skip to content

Commit 5d68e0b

Browse files
mergify[bot]Wenjun7J
mergify[bot]
and
Wenjun7J
authored
[BugFix] fix arrow flight empty set column name is 'r' (backport #71534) (#71632)
Signed-off-by: SevenJ <wenjun7j@gmail.com> Co-authored-by: SevenJ <166966490+Wenjun7J@users.noreply.github.com>
1 parent cb6e15f commit 5d68e0b

2 files changed

Lines changed: 183 additions & 51 deletions

File tree

‎fe/fe-core/src/main/java/com/starrocks/service/arrow/flight/sql/ArrowFlightSqlServiceImpl.java‎

Lines changed: 63 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -27,11 +27,14 @@
2727
import com.starrocks.common.InvalidConfException;
2828
import com.starrocks.common.Pair;
2929
import com.starrocks.common.ThreadPoolManager;
30-
import com.starrocks.common.util.ArrowUtil;
3130
import com.starrocks.common.util.DebugUtil;
3231
import com.starrocks.qe.GlobalVariable;
3332
import com.starrocks.server.GlobalStateMgr;
3433
import com.starrocks.service.arrow.flight.sql.session.ArrowFlightSqlSessionManager;
34+
import com.starrocks.sql.analyzer.Analyzer;
35+
import com.starrocks.sql.ast.OriginStatement;
36+
import com.starrocks.sql.ast.QueryStatement;
37+
import com.starrocks.sql.ast.StatementBase;
3538
import com.starrocks.sql.ast.expression.Expr;
3639
import com.starrocks.sql.plan.ExecPlan;
3740
import com.starrocks.system.ComputeNode;
@@ -194,19 +197,19 @@ public void createPreparedStatement(FlightSql.ActionCreatePreparedStatementReque
194197

195198
String preparedStmtId = ctx.addPreparedStatement(request.getQuery());
196199

197-
// To prevent the client from mistakenly interpreting an empty Schema as an update statement (instead of a query statement),
198-
// we need to ensure that the Schema returned by createPreparedStatement includes the query metadata.
199-
// This means we need to correctly set the DatasetSchema and ParameterSchema in ActionCreatePreparedStatementResult.
200-
// We generate a minimal Schema. This minimal Schema can include an integer column to ensure the Schema is not empty.
201-
try (VectorSchemaRoot schemaRoot = ArrowUtil.createSingleSchemaRoot("r", "0")) {
202-
Schema schema = schemaRoot.getSchema();
203-
FlightSql.ActionCreatePreparedStatementResult result =
204-
FlightSql.ActionCreatePreparedStatementResult.newBuilder()
205-
.setPreparedStatementHandle(ByteString.copyFromUtf8(preparedStmtId))
206-
.setDatasetSchema(ByteString.copyFrom(serializeMetadata(schema)))
207-
.setParameterSchema(ByteString.copyFrom(serializeMetadata(schema))).build();
208-
listener.onNext(new Result(Any.pack(result).toByteArray()));
209-
}
200+
// Try to plan the query to get the real schema for the prepared statement.
201+
// This is important because clients (JDBC, ADBC) use the schema returned here
202+
// to determine column names and types. Without this, a placeholder schema with
203+
// a single column named "r" would be returned, which is incorrect.
204+
Schema schema = buildSchemaFromQuery(ctx, request.getQuery());
205+
206+
FlightSql.ActionCreatePreparedStatementResult result =
207+
FlightSql.ActionCreatePreparedStatementResult.newBuilder()
208+
.setPreparedStatementHandle(ByteString.copyFromUtf8(preparedStmtId))
209+
.setDatasetSchema(ByteString.copyFrom(serializeMetadata(schema)))
210+
.setParameterSchema(ByteString.copyFrom(serializeMetadata(new Schema(Collections.emptyList()))))
211+
.build();
212+
listener.onNext(new Result(Any.pack(result).toByteArray()));
210213
listener.onCompleted();
211214
} catch (Exception e) {
212215
listener.onError(
@@ -900,6 +903,52 @@ protected <T extends Message> FlightInfo buildFlightInfo(T request, FlightDescri
900903
return new FlightInfo(schema, descriptor, endpoints, -1, -1);
901904
}
902905

906+
/**
907+
* Analyze the query to obtain the real output schema (column names and types).
908+
* Only performs semantic analysis (not full planning) to avoid side effects that
909+
* could interfere with the subsequent query execution in getFlightInfoPreparedStatement.
910+
* For non-query statements or if analysis fails, a placeholder schema is returned.
911+
*/
912+
private Schema buildSchemaFromQuery(ArrowFlightSqlConnectContext ctx, String query) {
913+
try {
914+
try (var scope = ctx.bindScope()) {
915+
List<StatementBase> stmts = com.starrocks.sql.parser.SqlParser.parse(query, ctx.getSessionVariable());
916+
if (stmts.isEmpty()) {
917+
return buildPlaceholderSchema();
918+
}
919+
StatementBase stmt = stmts.get(0);
920+
if (!(stmt instanceof QueryStatement)) {
921+
return buildPlaceholderSchema();
922+
}
923+
stmt.setOrigStmt(new OriginStatement(query));
924+
925+
Analyzer.analyze(stmt, ctx);
926+
927+
QueryStatement queryStmt = (QueryStatement) stmt;
928+
List<String> colNames = queryStmt.getQueryRelation().getColumnOutputNames();
929+
List<Expr> outputExprs = queryStmt.getQueryRelation().getOutputExpression();
930+
931+
List<Field> arrowFields = Lists.newArrayList();
932+
for (int i = 0; i < colNames.size(); i++) {
933+
Expr expr = outputExprs.get(i);
934+
Field arrowField = ArrowUtils.convertToArrowType(
935+
expr.getOriginType(), colNames.get(i), expr.isNullable());
936+
arrowFields.add(arrowField);
937+
}
938+
return new Schema(arrowFields);
939+
}
940+
} catch (Exception e) {
941+
LOG.warn("[ARROW] Failed to analyze query for schema in createPreparedStatement, " +
942+
"falling back to placeholder schema. query={}", query, e);
943+
return buildPlaceholderSchema();
944+
}
945+
}
946+
947+
private static Schema buildPlaceholderSchema() {
948+
return new Schema(Lists.newArrayList(
949+
ArrowUtils.convertToArrowType(com.starrocks.type.IntegerType.INT, "result", true)));
950+
}
951+
903952
public static Schema buildSchema(ExecPlan execPlan) {
904953
List<Field> arrowFields = Lists.newArrayList();
905954

‎fe/fe-core/src/test/java/com/starrocks/service/arrow/flight/sql/ArrowFlightSqlServiceImplTest.java‎

Lines changed: 120 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -15,13 +15,14 @@
1515
package com.starrocks.service.arrow.flight.sql;
1616

1717
import com.google.common.cache.Cache;
18+
import com.google.protobuf.Any;
1819
import com.google.protobuf.ByteString;
1920
import com.starrocks.common.Config;
2021
import com.starrocks.common.DdlException;
21-
import com.starrocks.common.util.ArrowUtil;
2222
import com.starrocks.metric.LongCounterMetric;
2323
import com.starrocks.metric.MetricRepo;
2424
import com.starrocks.plugin.AuditEvent;
25+
import com.starrocks.qe.ConnectContext;
2526
import com.starrocks.qe.GlobalVariable;
2627
import com.starrocks.qe.QueryState;
2728
import com.starrocks.qe.SessionVariable;
@@ -56,24 +57,30 @@
5657
import org.apache.arrow.flight.sql.FlightSqlProducer;
5758
import org.apache.arrow.flight.sql.impl.FlightSql;
5859
import org.apache.arrow.vector.VectorSchemaRoot;
60+
import org.apache.arrow.vector.ipc.ReadChannel;
61+
import org.apache.arrow.vector.ipc.message.MessageSerializer;
5962
import org.apache.arrow.vector.types.pojo.Schema;
6063
import org.junit.jupiter.api.BeforeEach;
6164
import org.junit.jupiter.api.Test;
6265
import org.mockito.ArgumentCaptor;
6366
import org.mockito.MockedStatic;
6467

68+
import java.io.ByteArrayInputStream;
69+
import java.io.IOException;
6570
import java.lang.reflect.Field;
66-
import java.nio.ByteBuffer;
67-
import java.nio.charset.StandardCharsets;
71+
import java.lang.reflect.Method;
72+
import java.nio.channels.Channels;
6873
import java.util.HashMap;
6974
import java.util.Map;
75+
import java.util.concurrent.CountDownLatch;
76+
import java.util.concurrent.TimeUnit;
7077

7178
import static org.junit.jupiter.api.Assertions.assertEquals;
7279
import static org.junit.jupiter.api.Assertions.assertNotNull;
7380
import static org.junit.jupiter.api.Assertions.assertNull;
81+
import static org.junit.jupiter.api.Assertions.assertSame;
7482
import static org.junit.jupiter.api.Assertions.assertThrows;
7583
import static org.junit.jupiter.api.Assertions.assertTrue;
76-
import static org.mockito.Answers.CALLS_REAL_METHODS;
7784
import static org.mockito.ArgumentMatchers.any;
7885
import static org.mockito.ArgumentMatchers.anyBoolean;
7986
import static org.mockito.ArgumentMatchers.anyLong;
@@ -93,6 +100,7 @@ public class ArrowFlightSqlServiceImplTest {
93100
private ArrowFlightSqlSessionManager sessionManager;
94101
private ArrowFlightSqlConnectContext mockContext;
95102
private FlightProducer.CallContext mockCallContext;
103+
private SessionVariable mockSessionVariable;
96104
@Mocked
97105
AuditEncryptionChecker mockChecker;
98106

@@ -116,6 +124,34 @@ protected boolean isProxyEnabled() {
116124
}
117125
}
118126

127+
private static class CapturingResultListener implements FlightProducer.StreamListener<Result> {
128+
private final CountDownLatch finished = new CountDownLatch(1);
129+
private volatile Result result;
130+
private volatile Throwable error;
131+
private volatile boolean completed;
132+
133+
@Override
134+
public void onNext(Result val) {
135+
this.result = val;
136+
}
137+
138+
@Override
139+
public void onError(Throwable t) {
140+
this.error = t;
141+
finished.countDown();
142+
}
143+
144+
@Override
145+
public void onCompleted() {
146+
this.completed = true;
147+
finished.countDown();
148+
}
149+
150+
public boolean await() throws InterruptedException {
151+
return finished.await(5, TimeUnit.SECONDS);
152+
}
153+
}
154+
119155
@BeforeEach
120156
public void setUp() throws IllegalAccessException, NoSuchFieldException, InterruptedException {
121157
sessionManager = mock(ArrowFlightSqlSessionManager.class);
@@ -129,7 +165,7 @@ public void setUp() throws IllegalAccessException, NoSuchFieldException, Interru
129165
when(mockContext.getExecutionId()).thenReturn(new com.starrocks.thrift.TUniqueId(1, 1));
130166
when(mockContext.getArrowFlightSqlToken()).thenReturn("token123");
131167

132-
SessionVariable mockSessionVariable = mock(SessionVariable.class);
168+
mockSessionVariable = mock(SessionVariable.class);
133169
when(mockContext.getSessionVariable()).thenReturn(mockSessionVariable);
134170
when(mockContext.acquireRunningToken(anyLong())).thenReturn(true);
135171
when(mockSessionVariable.getQueryTimeoutS()).thenReturn(10);
@@ -519,42 +555,75 @@ public void testGetStreamSqlInfo() {
519555
}
520556

521557
@Test
522-
public void testCreatePreparedStatement() throws Exception {
523-
// mock request
558+
public void testCreatePreparedStatementUsesAnalyzedSchema() throws Exception {
559+
String query = "SELECT 1 AS c1, 'x' AS c2";
524560
FlightSql.ActionCreatePreparedStatementRequest request = mock(FlightSql.ActionCreatePreparedStatementRequest.class);
525-
when(request.getQuery()).thenReturn("SELECT 1");
561+
when(request.getQuery()).thenReturn(query);
562+
ArrowFlightSqlConnectContext realContext = new ArrowFlightSqlConnectContext("token123");
563+
when(sessionManager.validateAndGetConnectContext("token123")).thenReturn(realContext);
526564

527-
// mock context
528-
FlightProducer.CallContext callContext = mock(FlightProducer.CallContext.class);
529-
when(callContext.peerIdentity()).thenReturn("token123");
565+
CapturingResultListener listener = new CapturingResultListener();
566+
service.createPreparedStatement(request, mockCallContext, listener);
530567

531-
// mock sessionManager
532-
when(sessionManager.validateAndGetConnectContext("token123")).thenReturn(mockContext);
568+
FlightSql.ActionCreatePreparedStatementResult preparedStatementResult = awaitPreparedStatementResult(listener);
533569

534-
// mock listener
535-
FlightProducer.StreamListener<Result> listener = mock(FlightProducer.StreamListener.class);
536-
537-
// mock ArrowUtil.createSingleSchemaRoot
538-
try (
539-
MockedStatic<ArrowUtil> mockArrowUtil = mockStatic(ArrowUtil.class);
540-
MockedStatic<ArrowFlightSqlServiceImpl> mockStatic =
541-
mockStatic(ArrowFlightSqlServiceImpl.class, CALLS_REAL_METHODS)
542-
) {
543-
// mock createSingleSchemaRoot
544-
VectorSchemaRoot mockRoot = mock(VectorSchemaRoot.class);
545-
Schema mockSchema = mock(Schema.class);
546-
when(mockRoot.getSchema()).thenReturn(mockSchema);
547-
mockArrowUtil.when(() -> ArrowUtil.createSingleSchemaRoot(anyString(), anyString())).thenReturn(mockRoot);
548-
549-
// mock serializeMetadata
550-
mockStatic.when(() -> ArrowFlightSqlServiceImpl.serializeMetadata(mockSchema))
551-
.thenReturn(ByteBuffer.wrap("mock".getBytes(StandardCharsets.UTF_8)));
552-
553-
// call method
554-
service.createPreparedStatement(request, callContext, listener);
555-
556-
// wait executor
557-
Thread.sleep(500);
570+
String preparedStatementHandle = preparedStatementResult.getPreparedStatementHandle().toStringUtf8();
571+
assertNotNull(preparedStatementHandle);
572+
assertEquals(query, realContext.getPreparedStatement(preparedStatementHandle));
573+
574+
Schema datasetSchema = deserializeSchema(preparedStatementResult.getDatasetSchema());
575+
assertEquals(2, datasetSchema.getFields().size());
576+
assertEquals("c1", datasetSchema.getFields().get(0).getName());
577+
assertTrue(datasetSchema.getFields().get(0).getType() instanceof org.apache.arrow.vector.types.pojo.ArrowType.Int);
578+
assertEquals("c2", datasetSchema.getFields().get(1).getName());
579+
assertTrue(datasetSchema.getFields().get(1).getType() instanceof org.apache.arrow.vector.types.pojo.ArrowType.Utf8);
580+
581+
Schema parameterSchema = deserializeSchema(preparedStatementResult.getParameterSchema());
582+
assertTrue(parameterSchema.getFields().isEmpty());
583+
}
584+
585+
@Test
586+
public void testCreatePreparedStatementFallsBackToPlaceholderSchemaForNonQueryStatement() throws Exception {
587+
String query = "SHOW DATABASES";
588+
FlightSql.ActionCreatePreparedStatementRequest request = mock(FlightSql.ActionCreatePreparedStatementRequest.class);
589+
when(request.getQuery()).thenReturn(query);
590+
ArrowFlightSqlConnectContext realContext = new ArrowFlightSqlConnectContext("token123");
591+
when(sessionManager.validateAndGetConnectContext("token123")).thenReturn(realContext);
592+
593+
CapturingResultListener listener = new CapturingResultListener();
594+
service.createPreparedStatement(request, mockCallContext, listener);
595+
596+
FlightSql.ActionCreatePreparedStatementResult preparedStatementResult = awaitPreparedStatementResult(listener);
597+
598+
String preparedStatementHandle = preparedStatementResult.getPreparedStatementHandle().toStringUtf8();
599+
assertNotNull(preparedStatementHandle);
600+
assertEquals(query, realContext.getPreparedStatement(preparedStatementHandle));
601+
602+
Schema datasetSchema = deserializeSchema(preparedStatementResult.getDatasetSchema());
603+
assertEquals(1, datasetSchema.getFields().size());
604+
assertEquals("result", datasetSchema.getFields().get(0).getName());
605+
assertTrue(datasetSchema.getFields().get(0).getType() instanceof org.apache.arrow.vector.types.pojo.ArrowType.Int);
606+
assertTrue(datasetSchema.getFields().get(0).isNullable());
607+
608+
Schema parameterSchema = deserializeSchema(preparedStatementResult.getParameterSchema());
609+
assertTrue(parameterSchema.getFields().isEmpty());
610+
}
611+
612+
@Test
613+
public void testBuildSchemaFromQueryRestoresPreviousConnectContext() throws Exception {
614+
ConnectContext previous = new ConnectContext();
615+
previous.setThreadLocalInfo();
616+
ArrowFlightSqlConnectContext arrowContext = new ArrowFlightSqlConnectContext("token123");
617+
Method method = ArrowFlightSqlServiceImpl.class
618+
.getDeclaredMethod("buildSchemaFromQuery", ArrowFlightSqlConnectContext.class, String.class);
619+
method.setAccessible(true);
620+
621+
try {
622+
Schema schema = (Schema) method.invoke(service, arrowContext, "SELECT 1 AS c1");
623+
assertNotNull(schema);
624+
assertSame(previous, ConnectContext.get());
625+
} finally {
626+
ConnectContext.remove();
558627
}
559628
}
560629

@@ -611,6 +680,20 @@ private void setFinalField(Object target, String fieldName, Object value) throws
611680
field.set(target, value);
612681
}
613682

683+
private FlightSql.ActionCreatePreparedStatementResult awaitPreparedStatementResult(CapturingResultListener listener)
684+
throws Exception {
685+
assertTrue(listener.await(), "Timed out waiting for createPreparedStatement result");
686+
assertNull(listener.error, "Unexpected createPreparedStatement error: " + listener.error);
687+
assertTrue(listener.completed);
688+
assertNotNull(listener.result);
689+
return Any.parseFrom(listener.result.getBody()).unpack(FlightSql.ActionCreatePreparedStatementResult.class);
690+
}
691+
692+
private Schema deserializeSchema(ByteString schemaBytes) throws IOException {
693+
return MessageSerializer.deserializeSchema(
694+
new ReadChannel(Channels.newChannel(new ByteArrayInputStream(schemaBytes.toByteArray()))));
695+
}
696+
614697
@Test
615698
public void testCacheClose() {
616699
ArrowFlightSqlServiceImpl service = new ArrowFlightSqlServiceImpl(mock(ArrowFlightSqlSessionManager.class), null);

0 commit comments

Comments
 (0)