Fixed ResultSetReader inconsistency

query/src/main/java/tech/ydb/query/tools/QueryReader.java

@@ -166,7 +166,7 @@ public ValueReader getColumn(String name) {
 
         @Override
         public Type getColumnType(int index) {
-            if (partIndex < 0) {
+            if (partIndex < 0 || partIndex >= parts.length) {
                 return null;
             }
             return parts[partIndex].getColumnType(index);
@@ -179,38 +179,37 @@ public int getRowCount() {
 
         @Override
         public void setRowIndex(int index) {
-            // TODO: Enable after JDBC fixing
-//            if (index < 0 || index >= rowsCount) {
-//                throw new IndexOutOfBoundsException(String.format("Index %s out of bounds for length %s",
-//                        index, rowsCount));
-//            }
-//            int currentIdx = index;
-            int currentIdx = Math.max(0, index);
+            if (index < 0) { // reset all
+                partIndex = parts.length == 0 ? -1 : 0;
+                for (ResultSetReader rs: parts) {
+                    rs.setRowIndex(-1);
+                }
+                return;
+            }
+
+            int currentIdx = index;
+
             partIndex = 0;
             while (partIndex < parts.length) {
                 int readerRows = parts[partIndex].getRowCount();
                 if (currentIdx < readerRows) {
                     parts[partIndex].setRowIndex(currentIdx);
-                    break;
+                    for (int partStep = partIndex + 1; partStep < parts.length; partStep++) {
+                        parts[partStep].setRowIndex(-1);
+                    }
+                    return;
                 }
                 parts[partIndex].setRowIndex(readerRows);
                 currentIdx -= readerRows;
                 partIndex++;
             }
 
-            // TODO: remove after JDBC fixing
-            if (partIndex >= parts.length) {
-                partIndex = parts.length - 1;
-            }
-
-            for (int partStep = partIndex + 1; partStep < parts.length; partStep++) {
-                parts[partStep].setRowIndex(0);
-            }
+            partIndex = parts.length - 1;
         }
 
         @Override
         public boolean next() {
-            if (partIndex < 0) {
+            if (partIndex < 0 || partIndex >= parts.length) {
                 return false;
             }
             boolean res = parts[partIndex].next();

query/src/test/java/tech/ydb/query/tools/QueryReaderTest.java

@@ -132,26 +132,32 @@ public void compositeResultSetTest() {
         Assert.assertEquals(6, readAll(rsr, 0));
         Assert.assertEquals(0, readAll(rsr, 0));
 
-        rsr.setRowIndex(0);
+        rsr.setRowIndex(-1);
+        Assert.assertEquals(6, readAll(rsr, 0));
+        Assert.assertEquals(0, readAll(rsr, 0));
+
+        rsr.setRowIndex(-100);
         Assert.assertEquals(6, readAll(rsr, 0));
         Assert.assertEquals(0, readAll(rsr, 0));
 
+        rsr.setRowIndex(0);
+        Assert.assertEquals(5, readAll(rsr, 1));
+        Assert.assertEquals(0, readAll(rsr, 0));
+
         rsr.setRowIndex(3);
-        Assert.assertEquals(3, readAll(rsr, 3));
+        Assert.assertEquals(2, readAll(rsr, 4));
 
         rsr.setRowIndex(5);
-        Assert.assertEquals(1, readAll(rsr, 5));
+        Assert.assertEquals(0, readAll(rsr, 0));
 
         rsr.setRowIndex(-1);
         Assert.assertEquals(6, readAll(rsr, 0));
 
         rsr.setRowIndex(6);
         Assert.assertEquals(0, readAll(rsr, 0));
-//        IndexOutOfBoundsException ex1 = Assert.assertThrows(IndexOutOfBoundsException.class, () -> rsr.setRowIndex(6));
-//        Assert.assertEquals("Index 6 out of bounds for length 6", ex1.getMessage());
-//
-//        IndexOutOfBoundsException ex2 = Assert.assertThrows(IndexOutOfBoundsException.class, () -> rsr.setRowIndex(-1));
-//        Assert.assertEquals("Index -1 out of bounds for length 6", ex2.getMessage());
+
+        rsr.setRowIndex(100);
+        Assert.assertEquals(0, readAll(rsr, 0));
     }
 
     private int readAll(ResultSetReader rsr, int startKey) {

table/src/main/java/tech/ydb/table/result/ResultSetReader.java

@@ -8,55 +8,96 @@
 public interface ResultSetReader {
 
     /**
-     * Returns {@code true} if the result was truncated, {@code false} otherwise.
+     * Gets whether this result set was truncated.
+     *
+     * @return {@code true} if the result was truncated, {@code false} otherwise.
      */
     boolean isTruncated();
 
     /**
-     * Returns number of columns.
+     * Gets number of this result set columns
+     *
+     * @return the result set columns count
      */
     int getColumnCount();
 
     /**
-     * Returns number of rows.
+     * Gets number of this result set rows
+     *
+     * @return the result set rows count
      */
     int getRowCount();
 
     /**
-     * Explicitly switch to a specific row.
+     * Explicitly sets the reader on a specific row position.
+     * <p>
+     * Valid row indexes are in the range <code>0 &lt;= index &lt; getRowCount()</code>, where
+     * {@code 0} is the first row and {@code getRowCount() - 1} is the last row.
+     * Passing a negative value positions the reader <em>before</em> the first row, so that
+     * a subsequent {@link #next()} call advances to the first row. Passing a value greater
+     * than or equal to {@link #getRowCount()} positions the reader <em>after</em> the last
+     * row, so that {@link #next()} will return {@code false}.
+     *
+     * @param index the desired row index, zero is the first row; see above for special values
      */
     void setRowIndex(int index);
 
     /**
-     * Set iterator to the next table row.
+     * Set reader to the next table row.
+     * <p>
+     * On success {@link #next()} will reset all column parsers to the values in next row. Column parsers
+     * are invalid before the first {@link #next()} call.
      *
-     * On success tryNextRow will reset all column parsers to the values in next row.
-     * Column parsers are invalid before the first TryNextRow call.
+     * @return returns {@code true} if the next row is available, {@code false} otherwise.
      */
     boolean next();
 
     /**
-     * Returns column name by index.
+     * Gets column name by index.
+     *
+     * @param index the column index, zero is the first column
+     * @return the column name
+     * @throws IllegalArgumentException if the index is out of range
+     *         (<code>index &lt; 0 || index &gt;= getColumnCount()</code>)
      */
     String getColumnName(int index);
 
     /**
-     * Returns column index by name or {@code -1} if column with given name is not present.
+     * Gets column type by index.
+     *
+     * @param index the column index, zero is the first column
+     * @return the column type
+     * @throws IllegalArgumentException if the index is out of range
+     *         (<code>index &lt; 0 || index &gt;= getColumnCount()</code>)
      */
-    int getColumnIndex(String name);
+    Type getColumnType(int index);
 
     /**
-     * Returns value reader for column by index.
+     * Gets column index by name or {@code -1} if column with given name is not present.
+     *
+     * @param name the column name
+     * @return the column index
      */
-    ValueReader getColumn(int index);
+    int getColumnIndex(String name);
 
     /**
-     * Returns value reader for column by name.
+     * Gets the current row value reader by the column index
+     *
+     * @param index the column index, zero is the first column
+     * @return the value reader
+     * @throws IllegalArgumentException if the column index is out of range
+     *         (<code>index &lt; 0 || index &gt;= getColumnCount()</code>)
+     * @throws IllegalStateException if the result set reading was not started or was already finished
      */
-    ValueReader getColumn(String name);
+    ValueReader getColumn(int index);
 
     /**
-     * Returns column type by index (always create a new type instance)
+     * Gets the current row value reader by the column name
+     *
+     * @param name the column name
+     * @return the value reader
+     * @throws IllegalArgumentException if column with given name is not present
+     * @throws IllegalStateException if the result set reading was not started or was already finished
      */
-    Type getColumnType(int index);
+    ValueReader getColumn(String name);
 }

table/src/main/java/tech/ydb/table/result/impl/ProtoResultSetReader.java

@@ -15,75 +15,78 @@
  */
 final class ProtoResultSetReader implements ResultSetReader {
 
-    private final ValueProtos.ResultSet resultSet;
-    private final Map<String, Integer> columnIndexes; // TODO: use better data structure
-    private final AbstractValueReader[] columnReaders;
+    private final ValueProtos.ResultSet rs;
+    private final Map<String, Integer> columnIndexes;
+    private final AbstractValueReader[] readers;
 
-    private int rowIndex;
-    private ValueProtos.Value currentRow;
+    private int rowIndex = -1; // before first
+    private ValueProtos.Value currentRow = null;
 
     ProtoResultSetReader(ValueProtos.ResultSet resultSet) {
-        this.resultSet = resultSet;
+        this.rs = resultSet;
         this.columnIndexes = Maps.newHashMapWithExpectedSize(resultSet.getColumnsCount());
-        this.columnReaders = new AbstractValueReader[resultSet.getColumnsCount()];
+        this.readers = new AbstractValueReader[resultSet.getColumnsCount()];
 
         for (int i = 0; i < resultSet.getColumnsCount(); i++) {
             ValueProtos.Column columnMeta = resultSet.getColumns(i);
             this.columnIndexes.put(columnMeta.getName(), i);
-            this.columnReaders[i] = ProtoValueReaders.forTypeImpl(columnMeta.getType());
+            this.readers[i] = ProtoValueReaders.forTypeImpl(columnMeta.getType());
         }
     }
 
     @Override
     public boolean isTruncated() {
-        return resultSet.getTruncated();
+        return rs.getTruncated();
     }
 
-    /**
-     * Returns number of columns.
-     */
     @Override
     public int getColumnCount() {
-        return resultSet.getColumnsCount();
+        return rs.getColumnsCount();
     }
 
-    /**
-     * Returns number of rows.
-     */
     @Override
     public int getRowCount() {
-        return resultSet.getRowsCount();
+        return rs.getRowsCount();
     }
 
     @Override
     public void setRowIndex(int index) {
-        if (index < 0 || index >= resultSet.getRowsCount()) {
+        if (index <= -1) {
+            rowIndex = -1; // before first
             currentRow = null;
-        } else {
-            rowIndex = index;
-            currentRow = resultSet.getRows(index);
+            return;
         }
+
+        if (index >= rs.getRowsCount()) {
+            rowIndex = rs.getRowsCount(); // after last
+            currentRow = null;
+            return;
+        }
+
+        rowIndex = index;
+        currentRow = rs.getRows(rowIndex);
     }
 
-    /**
-     * Set iterator to the next table row.
-     *
-     * On success tryNextRow will reset all column parsers to the values in next row.
-     * Column parsers are invalid before the first TryNextRow call.
-     */
     @Override
     public boolean next() {
-        if (rowIndex >= resultSet.getRowsCount()) {
+        rowIndex++;
+
+        if (rowIndex >= rs.getRowsCount()) {
+            rowIndex = rs.getRowsCount(); // after last
             currentRow = null;
             return false;
         }
-        currentRow = resultSet.getRows(rowIndex++);
+
+        currentRow = rs.getRows(rowIndex);
         return true;
     }
 
     @Override
     public String getColumnName(int index) {
-        return resultSet.getColumns(index).getName();
+        if (index < 0 || index >= rs.getColumnsCount()) {
+            throw new IllegalArgumentException("Column index: " + index + ", columns count: " + readers.length);
+        }
+        return rs.getColumns(index).getName();
     }
 
     @Override
@@ -95,34 +98,36 @@ public int getColumnIndex(String name) {
     @Override
     public ValueReader getColumn(int index) {
         if (currentRow == null) {
-            throw new IllegalStateException("empty result set or next() was never called");
+            throw new IllegalStateException("ResultSetReader not positioned properly, perhaps you need to call next.");
         }
-        AbstractValueReader reader = columnReaders[index];
+        if (index < 0 || index >= readers.length) {
+            throw new IllegalArgumentException("Column index: " + index + ", columns count: " + readers.length);
+        }
+        AbstractValueReader reader = readers[index];
         reader.setProtoValue(currentRow.getItems(index));
         return reader;
     }
 
     @Override
     public ValueReader getColumn(String name) {
-        int index = columnIndex(name);
+        Integer index = columnIndexes.get(name);
+        if (index == null) {
+            throw new IllegalArgumentException("Unknown column '" + name + "'");
+        }
         return getColumn(index);
     }
 
     @Override
     public Type getColumnType(int index) {
-        AbstractValueReader reader = columnReaders[index];
-        return reader.getType();
-    }
-
-    private int columnIndex(String name) {
-        Integer index = columnIndexes.get(name);
-        if (index == null) {
-            throw new IllegalArgumentException("unknown column '" + name + "'");
+        if (index < 0 || index >= readers.length) {
+            throw new IllegalArgumentException("Column index: " + index + ", columns count: " + readers.length);
         }
-        return index;
+        AbstractValueReader reader = readers[index];
+        return reader.getType();
     }
 
+    @Deprecated
     ValueProtos.ResultSet getResultSet() {
-        return resultSet;
+        return rs;
     }
 }

table/src/main/java/tech/ydb/table/result/impl/ProtoValueReaders.java

@@ -20,6 +20,7 @@ public static ResultSetReader forResultSet(ValueProtos.ResultSet resultSet) {
         return new ProtoResultSetReader(resultSet);
     }
 
+    @Deprecated
     public static ResultSetReader forResultSets(Collection<ResultSetReader> resultSets) {
         // TODO: add lightweight implementation instead of proto joining
         Preconditions.checkArgument(!resultSets.isEmpty(), "Expect multiple result sets to join from");