fix returned ResultSet for Arrays

Author: xborder

flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcArray.java

@@ -26,6 +26,7 @@
 import org.apache.arrow.driver.jdbc.utils.SqlTypes;
 import org.apache.arrow.memory.util.LargeMemoryUtil;
 import org.apache.arrow.vector.FieldVector;
+import org.apache.arrow.vector.IntVector;
 import org.apache.arrow.vector.ValueVector;
 import org.apache.arrow.vector.VectorSchemaRoot;
 import org.apache.arrow.vector.types.pojo.ArrowType;
@@ -135,12 +136,22 @@ public ResultSet getResultSet(long index, int count) throws SQLException {
 
   private static ResultSet getResultSetNoBoundariesCheck(
       ValueVector dataVector, long start, long count) throws SQLException {
+    int intStart = LargeMemoryUtil.checkedCastToInt(start);
+    int intCount = LargeMemoryUtil.checkedCastToInt(count);
+
+    // Create an index vector with 1-based indices (per JDBC spec) to return with value vector
+    IntVector indexVector = new IntVector("INDEX", dataVector.getAllocator());
+    indexVector.allocateNew(intCount);
+    for (int i = 0; i < intCount; i++) {
+      indexVector.set(i, i + 1);
+    }
+    indexVector.setValueCount(intCount);
+
     TransferPair transferPair = dataVector.getTransferPair(dataVector.getAllocator());
-    transferPair.splitAndTransfer(
-        LargeMemoryUtil.checkedCastToInt(start), LargeMemoryUtil.checkedCastToInt(count));
-    FieldVector vectorSlice = (FieldVector) transferPair.getTo();
+    transferPair.splitAndTransfer(intStart, intCount);
+    FieldVector valueVector = (FieldVector) transferPair.getTo();
 
-    VectorSchemaRoot vectorSchemaRoot = VectorSchemaRoot.of(vectorSlice);
+    VectorSchemaRoot vectorSchemaRoot = VectorSchemaRoot.of(indexVector, valueVector);
     return ArrowFlightJdbcVectorSchemaRootResultSet.fromVectorSchemaRoot(vectorSchemaRoot);
   }
 

flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/impl/complex/AbstractArrowFlightJdbcListAccessorTest.java

@@ -21,6 +21,7 @@
 
 import java.sql.Array;
 import java.sql.ResultSet;
+import java.sql.ResultSetMetaData;
 import java.util.Arrays;
 import java.util.List;
 import java.util.function.Supplier;
@@ -191,7 +192,12 @@ public void testShouldGetArrayGetResultSetReturnValidResultSet(
           try (ResultSet rs = array.getResultSet()) {
             int count = 0;
             while (rs.next()) {
-              final int value = rs.getInt(1);
+              // Column 1: 1-based index (per JDBC spec)
+              final int index = rs.getInt(1);
+              assertThat(index, equalTo(count + 1));
+
+              // Column 2: actual value (per JDBC spec)
+              final int value = rs.getInt(2);
               assertThat(value, equalTo(currentRow * count));
               count++;
             }