This is an automated email from the ASF dual-hosted git repository. jackietien pushed a commit to branch ai-code/flight-sql in repository https://gitbox.apache.org/repos/asf/iotdb.git
commit c6808daf8c9bcce1b7d5af6fe8dfd5dc1662beae Author: JackieTien97 <[email protected]> AuthorDate: Tue Feb 24 20:22:33 2026 +0800 gemini 3.1 pro fix IT --- external-service-impl/flight-sql/pom.xml | 8 +++ .../iotdb/flight/TsBlockToArrowConverter.java | 13 +++- .../it/flightsql/IoTDBArrowFlightSqlIT.java | 76 +++++++++------------- 3 files changed, 51 insertions(+), 46 deletions(-) diff --git a/external-service-impl/flight-sql/pom.xml b/external-service-impl/flight-sql/pom.xml index be363c336b7..e1160323af6 100644 --- a/external-service-impl/flight-sql/pom.xml +++ b/external-service-impl/flight-sql/pom.xml @@ -113,6 +113,14 @@ <configuration> <ignoredDependencies> <ignoredDependency>org.apache.tsfile:common</ignoredDependency> + <ignoredDependency>org.apache.iotdb:iotdb-thrift-commons</ignoredDependency> + <ignoredDependency>com.google.protobuf:protobuf-java</ignoredDependency> + <ignoredDependency>org.apache.iotdb:iotdb-thrift</ignoredDependency> + <ignoredDependency>org.apache.arrow:flight-core</ignoredDependency> + <ignoredDependency>org.apache.arrow:arrow-vector</ignoredDependency> + <ignoredDependency>org.apache.arrow:arrow-memory-core</ignoredDependency> + <ignoredDependency>org.apache.iotdb:service-rpc</ignoredDependency> + <ignoredDependency>org.apache.arrow:arrow-memory-netty</ignoredDependency> </ignoredDependencies> </configuration> </plugin> diff --git a/external-service-impl/flight-sql/src/main/java/org/apache/iotdb/flight/TsBlockToArrowConverter.java b/external-service-impl/flight-sql/src/main/java/org/apache/iotdb/flight/TsBlockToArrowConverter.java index 396f230d0b9..fad93708e78 100644 --- a/external-service-impl/flight-sql/src/main/java/org/apache/iotdb/flight/TsBlockToArrowConverter.java +++ b/external-service-impl/flight-sql/src/main/java/org/apache/iotdb/flight/TsBlockToArrowConverter.java @@ -121,8 +121,17 @@ public class TsBlockToArrowConverter { root.allocateNew(); for (int colIdx = 0; colIdx < columnNames.size(); colIdx++) { - String colName = columnNames.get(colIdx); - Integer sourceIdx = headerMap.get(colName); + // Use headerMap for column index mapping when available; fall back to + // sequential index + // for queries like SHOW DATABASES where getColumnNameIndexMap() returns null. + int sourceIdx; + if (headerMap != null) { + String colName = columnNames.get(colIdx); + Integer idx = headerMap.get(colName); + sourceIdx = (idx != null) ? idx : colIdx; + } else { + sourceIdx = colIdx; + } Column column = tsBlock.getColumn(sourceIdx); TSDataType dataType = dataTypes.get(colIdx); FieldVector fieldVector = root.getVector(colIdx); diff --git a/integration-test/src/test/java/org/apache/iotdb/relational/it/flightsql/IoTDBArrowFlightSqlIT.java b/integration-test/src/test/java/org/apache/iotdb/relational/it/flightsql/IoTDBArrowFlightSqlIT.java index b683df9a1c4..1710c15adb2 100644 --- a/integration-test/src/test/java/org/apache/iotdb/relational/it/flightsql/IoTDBArrowFlightSqlIT.java +++ b/integration-test/src/test/java/org/apache/iotdb/relational/it/flightsql/IoTDBArrowFlightSqlIT.java @@ -157,62 +157,50 @@ public class IoTDBArrowFlightSqlIT { } @Test - public void testQueryWithAllDataTypes() throws Exception { - FlightInfo flightInfo = + public void testFlightSqlQueries() throws Exception { + FlightInfo flightInfo; + List<List<String>> rows; + Schema schema; + List<Field> fields; + // 1. Query with all data types + flightInfo = flightSqlClient.execute( "SELECT time, id1, s1, s2, s3, s4, s5, s6 FROM " + TABLE + " ORDER BY time", credentials); - - // Validate schema - Schema schema = flightInfo.getSchemaOptional().orElse(null); + schema = flightInfo.getSchemaOptional().orElse(null); assertNotNull("Schema should not be null", schema); - List<Field> fields = schema.getFields(); + fields = schema.getFields(); assertEquals("Should have 8 columns", 8, fields.size()); - - // Fetch all data - List<List<String>> rows = fetchAllRows(flightInfo); + rows = fetchAllRows(flightInfo); assertEquals("Should have 3 rows", 3, rows.size()); - } - - @Test - public void testQueryWithFilter() throws Exception { - FlightInfo flightInfo = + // 2. Query with filter + flightInfo = flightSqlClient.execute( "SELECT id1, s1 FROM " + TABLE + " WHERE id1 = 'device1' ORDER BY time", credentials); - - List<List<String>> rows = fetchAllRows(flightInfo); + rows = fetchAllRows(flightInfo); assertEquals("Should have 2 rows for device1", 2, rows.size()); - } - @Test - public void testQueryWithAggregation() throws Exception { - FlightInfo flightInfo = + // 3. Query with aggregation + flightInfo = flightSqlClient.execute( "SELECT id1, COUNT(*) as cnt, SUM(s1) as s1_sum " + "FROM " + TABLE + " GROUP BY id1 ORDER BY id1", credentials); - - List<List<String>> rows = fetchAllRows(flightInfo); + rows = fetchAllRows(flightInfo); assertEquals("Should have 2 groups", 2, rows.size()); - } - @Test - public void testEmptyResult() throws Exception { - FlightInfo flightInfo = + // 4. Empty result query + flightInfo = flightSqlClient.execute( "SELECT * FROM " + TABLE + " WHERE id1 = 'nonexistent'", credentials); - - List<List<String>> rows = fetchAllRows(flightInfo); + rows = fetchAllRows(flightInfo); assertEquals("Should have 0 rows", 0, rows.size()); - } - - @Test - public void testShowDatabases() throws Exception { - FlightInfo flightInfo = flightSqlClient.execute("SHOW DATABASES", credentials); - List<List<String>> rows = fetchAllRows(flightInfo); + // 5. Show databases + flightInfo = flightSqlClient.execute("SHOW DATABASES", credentials); + rows = fetchAllRows(flightInfo); assertTrue("Should have at least 1 database", rows.size() >= 1); boolean found = false; @@ -224,7 +212,7 @@ public class IoTDBArrowFlightSqlIT { } } } - assertTrue("Should find test database " + DATABASE, found); + assertTrue("Should find the created database", found); } /** @@ -234,17 +222,17 @@ public class IoTDBArrowFlightSqlIT { private List<List<String>> fetchAllRows(FlightInfo flightInfo) throws Exception { List<List<String>> rows = new ArrayList<>(); for (FlightEndpoint endpoint : flightInfo.getEndpoints()) { - try (FlightStream stream = flightSqlClient.getStream(endpoint.getTicket(), credentials)) { + try (FlightStream stream = flightSqlClient.getStream(endpoint.getTicket())) { while (stream.next()) { - VectorSchemaRoot root = stream.getRoot(); - int rowCount = root.getRowCount(); - for (int i = 0; i < rowCount; i++) { - List<String> row = new ArrayList<>(); - for (FieldVector vector : root.getFieldVectors()) { - Object value = vector.getObject(i); - row.add(value == null ? "null" : value.toString()); + try (VectorSchemaRoot root = stream.getRoot()) { + for (int i = 0; i < root.getRowCount(); i++) { + List<String> row = new ArrayList<>(); + for (FieldVector vector : root.getFieldVectors()) { + Object value = vector.getObject(i); + row.add(value == null ? "null" : value.toString()); + } + rows.add(row); } - rows.add(row); } } }
