From 038c5b7e55cbc700792deb17c8eab787ab2fd034 Mon Sep 17 00:00:00 2001
From: Todd Farmer <todd@fivefarmers.com>
Date: Mon, 2 May 2022 09:13:49 -0600
Subject: [PATCH] ARROW-16035: Fixing ArrowVectorIterator.hasNext() in case of
 empty ResultSet

---
 .../adapter/jdbc/ArrowVectorIterator.java     | 20 ++++++++++++---
 .../adapter/jdbc/h2/JdbcToArrowTest.java      | 25 ++++++++++++++++---
 2 files changed, 38 insertions(+), 7 deletions(-)

diff --git a/java/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/ArrowVectorIterator.java b/java/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/ArrowVectorIterator.java
index 1dfc462b5ef88..ba0b700415cb7 100644
--- a/java/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/ArrowVectorIterator.java
+++ b/java/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/ArrowVectorIterator.java
@@ -53,6 +53,10 @@ public class ArrowVectorIterator implements Iterator<VectorSchemaRoot>, AutoClos
 
   private final int targetBatchSize;
 
+  // This is used to track whether the ResultSet has been fully read, and is needed spcifically for cases where there
+  // is a ResultSet having zero rows (empty):
+  private boolean readComplete = false;
+
   /**
    * Construct an instance.
    */
@@ -107,10 +111,15 @@ private void consumeData(VectorSchemaRoot root) {
           compositeConsumer.consume(resultSet);
           readRowCount++;
         }
+        readComplete = true;
       } else {
-        while (readRowCount < targetBatchSize && resultSet.next()) {
-          compositeConsumer.consume(resultSet);
-          readRowCount++;
+        while ((readRowCount < targetBatchSize) && !readComplete) {
+          if (resultSet.next()) {
+            compositeConsumer.consume(resultSet);
+            readRowCount++;
+          } else {
+            readComplete = true;
+          }
         }
       }
 
@@ -154,8 +163,11 @@ private void load(VectorSchemaRoot root) throws SQLException {
 
   @Override
   public boolean hasNext() {
+    if (readComplete) {
+      return false;
+    }
     try {
-      return !resultSet.isAfterLast();
+      return !resultSet.isLast();
     } catch (SQLException e) {
       throw new RuntimeException(e);
     }
diff --git a/java/adapter/jdbc/src/test/java/org/apache/arrow/adapter/jdbc/h2/JdbcToArrowTest.java b/java/adapter/jdbc/src/test/java/org/apache/arrow/adapter/jdbc/h2/JdbcToArrowTest.java
index ca1c0c00bf758..32db254ac3437 100644
--- a/java/adapter/jdbc/src/test/java/org/apache/arrow/adapter/jdbc/h2/JdbcToArrowTest.java
+++ b/java/adapter/jdbc/src/test/java/org/apache/arrow/adapter/jdbc/h2/JdbcToArrowTest.java
@@ -40,7 +40,7 @@
 import static org.apache.arrow.adapter.jdbc.JdbcToArrowTestHelper.getFloatValues;
 import static org.apache.arrow.adapter.jdbc.JdbcToArrowTestHelper.getIntValues;
 import static org.apache.arrow.adapter.jdbc.JdbcToArrowTestHelper.getLongValues;
-import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.*;
 
 import java.io.IOException;
 import java.io.InputStream;
@@ -251,15 +251,31 @@ public void runLargeNumberOfRows() throws IOException, SQLException {
       allocator.close();
     }
 
-    assertEquals(x, targetRows);
+    assertEquals(targetRows, x);
+  }
+
+  @Test
+  public void testZeroRowResultSet() throws Exception {
+    BufferAllocator allocator = new RootAllocator(Integer.MAX_VALUE);
+    int x = 0;
+    final int targetRows = 0;
+    ResultSet rs = new FakeResultSet(targetRows);
+    JdbcToArrowConfig config = new JdbcToArrowConfigBuilder(
+            allocator, JdbcToArrowUtils.getUtcCalendar(), /* include metadata */ false)
+            .setReuseVectorSchemaRoot(reuseVectorSchemaRoot).build();
+
+    ArrowVectorIterator iter = JdbcToArrow.sqlToArrowVectorIterator(rs, config);
+    assertFalse("Iterator on zero row ResultSet should not haveNext()", iter.hasNext());
   }
 
   private class FakeResultSet implements ResultSet {
 
     public int numRows;
+    public boolean empty;
 
     FakeResultSet(int numRows) {
       this.numRows = numRows;
+      this.empty = numRows <= 0;
     }
 
     @Override
@@ -629,6 +645,9 @@ public boolean isBeforeFirst() throws SQLException {
 
     @Override
     public boolean isAfterLast() throws SQLException {
+      if (empty) {
+        return false;
+      }
       return numRows < 0;
     }
 
@@ -639,7 +658,7 @@ public boolean isFirst() throws SQLException {
 
     @Override
     public boolean isLast() throws SQLException {
-      return false;
+      return numRows <= 0;
     }
 
     @Override