Skip to content

Commit

Permalink
ARROW-16035: [Java] Handling empty JDBC ResultSet
Browse files Browse the repository at this point in the history
ArrowVectorIterator.hasNext() delegates to the underlying resultset.isAfterLast() method, but per JDBC specs, this [should return false in the case of empty ResultSets](https://docs.oracle.com/en/java/javase/11/docs/api/java.sql/java/sql/ResultSet.html#isAfterLast()).  This causes hasNext() to return true for empty ResultSets, triggering infinite loops.

Closes #13049 from toddfarmer/tofarmer/fix-hasnext-for-empty-resultset

Authored-by: Todd Farmer <todd@fivefarmers.com>
Signed-off-by: David Li <li.davidm96@gmail.com>
  • Loading branch information
toddfarmer authored and lidavidm committed May 3, 2022
1 parent cf2a35c commit 2025673
Show file tree
Hide file tree
Showing 2 changed files with 38 additions and 7 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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.
*/
Expand Down Expand Up @@ -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;
}
}
}

Expand Down Expand Up @@ -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);
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -629,6 +645,9 @@ public boolean isBeforeFirst() throws SQLException {

@Override
public boolean isAfterLast() throws SQLException {
if (empty) {
return false;
}
return numRows < 0;
}

Expand All @@ -639,7 +658,7 @@ public boolean isFirst() throws SQLException {

@Override
public boolean isLast() throws SQLException {
return false;
return numRows <= 0;
}

@Override
Expand Down

0 comments on commit 2025673

Please sign in to comment.