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);
           }
         }
       }

Reply via email to