fix(query): QueryReader.setRowIndex() now properly works with composite result sets

Author: zinal

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

@@ -28,7 +28,7 @@ public class QueryReader implements Iterable<ResultSetReader> {
     private final List<Issue> isssues;
     private final List<ResultSetParts> results;
 
-    private QueryReader(QueryInfo info, List<Issue> issues, List<ResultSetParts> results) {
+    QueryReader(QueryInfo info, List<Issue> issues, List<ResultSetParts> results) {
         this.info = info;
         this.isssues = issues;
         this.results = results;
@@ -113,7 +113,7 @@ public void onNextPart(QueryResultPart part) {
         }
     }
 
-    private static class ResultSetParts {
+    static class ResultSetParts {
         private final long resultSetIndex;
         private final List<QueryResultPart> parts = new ArrayList<>();
 
@@ -211,12 +211,15 @@ public void setRowIndex(int index) {
                 int readerRows = parts.get(partIndex).getRowCount();
                 if (currentIdx < readerRows) {
                     parts.get(partIndex).setRowIndex(currentIdx);
-                    return;
+                    break;
                 }
                 parts.get(partIndex).setRowIndex(readerRows - 1);
                 currentIdx -= readerRows;
                 partIndex++;
             }
+            for (int partStep = partIndex + 1; partStep < parts.size(); partStep++) {
+                parts.get(partStep).setRowIndex(0);
+            }
         }
 
         @Override

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

@@ -0,0 +1,118 @@
+package tech.ydb.query.tools;
+
+import java.util.Arrays;
+import org.junit.Assert;
+import org.junit.Test;
+import tech.ydb.proto.ValueProtos;
+import tech.ydb.query.result.QueryInfo;
+import tech.ydb.query.result.QueryResultPart;
+import tech.ydb.table.result.ResultSetReader;
+
+/**
+ *
+ * @author zinal
+ */
+public class QueryReaderTest {
+
+    @Test
+    public void testNextAndSetIndex() {
+        QueryInfo qi = new QueryInfo(null);
+        QueryReader.ResultSetParts rsp0 = initRsp(0L);
+        QueryReader.ResultSetParts rsp1 = initRsp(1L);
+        QueryReader.ResultSetParts rsp2 = initRsp(2L);
+        QueryReader qr = new QueryReader(qi, null, Arrays.asList(rsp0, rsp1, rsp2));
+        checkRows(qr, 0);
+        checkRows(qr, 1);
+        checkRows(qr, 2);
+    }
+
+    private QueryReader.ResultSetParts initRsp(long ix) {
+        QueryReader.ResultSetParts rsp = new QueryReader.ResultSetParts(ix);
+        rsp.addPart(initQrp(ix));
+        rsp.addPart(initQrp(ix));
+        rsp.addPart(initQrp(ix));
+        rsp.addPart(initQrp(ix));
+        return rsp;
+    }
+
+    private QueryResultPart initQrp(long ix) {
+        ValueProtos.ResultSet.Builder rsBuilder = ValueProtos.ResultSet.newBuilder();
+        rsBuilder.setTruncated(false);
+        rsBuilder.addColumns(ValueProtos.Column.newBuilder()
+                .setName("a")
+                .setType(ValueProtos.Type.newBuilder().setTypeIdValue(ValueProtos.Type.PrimitiveTypeId.INT32_VALUE).build())
+                .build()
+        );
+        rsBuilder.addColumns(ValueProtos.Column.newBuilder()
+                .setName("b")
+                .setType(ValueProtos.Type.newBuilder().setTypeIdValue(ValueProtos.Type.PrimitiveTypeId.UTF8_VALUE).build())
+                .build()
+        );
+        rsBuilder.addColumns(ValueProtos.Column.newBuilder()
+                .setName("c")
+                .setType(ValueProtos.Type.newBuilder().setTypeIdValue(ValueProtos.Type.PrimitiveTypeId.INT64_VALUE).build())
+                .build()
+        );
+        rsBuilder.addRows(ValueProtos.Value.newBuilder()
+                .addItems(ValueProtos.Value.newBuilder()
+                        .setInt32Value(10)
+                        .build())
+                .addItems(ValueProtos.Value.newBuilder()
+                        .setTextValue("aaa")
+                        .build())
+                .addItems(ValueProtos.Value.newBuilder()
+                        .setInt64Value(100L)
+                        .build())
+                .build());
+        rsBuilder.addRows(ValueProtos.Value.newBuilder()
+                .addItems(ValueProtos.Value.newBuilder()
+                        .setInt32Value(20)
+                        .build())
+                .addItems(ValueProtos.Value.newBuilder()
+                        .setTextValue("bbb")
+                        .build())
+                .addItems(ValueProtos.Value.newBuilder()
+                        .setInt64Value(200L)
+                        .build())
+                .build());
+        rsBuilder.addRows(ValueProtos.Value.newBuilder()
+                .addItems(ValueProtos.Value.newBuilder()
+                        .setInt32Value(30)
+                        .build())
+                .addItems(ValueProtos.Value.newBuilder()
+                        .setTextValue("ccc")
+                        .build())
+                .addItems(ValueProtos.Value.newBuilder()
+                        .setInt64Value(300L)
+                        .build())
+                .build());
+        return new QueryResultPart(ix, rsBuilder.build());
+    }
+
+    private void checkRows(QueryReader qr, int rsIndex) {
+        ResultSetReader rsr = qr.getResultSet(rsIndex);
+        Assert.assertNotNull(rsr);
+        
+        Assert.assertEquals(3, rsr.getColumnCount());
+        Assert.assertEquals("a", rsr.getColumnName(0));
+        Assert.assertEquals("b", rsr.getColumnName(1));
+        Assert.assertEquals("c", rsr.getColumnName(2));
+
+        int rowCount = 0;
+        while (rsr.next()) {
+            Assert.assertNotNull(rsr.getColumn(0));
+            Assert.assertNotNull(rsr.getColumn(1));
+            Assert.assertNotNull(rsr.getColumn(2));
+            rowCount += 1;
+        }
+        Assert.assertEquals(12, rowCount);
+
+        rsr.setRowIndex(0);
+        rowCount = 0;
+        while (rsr.next()) {
+            rowCount += 1;
+        }
+        Assert.assertEquals(12, rowCount);
+    }
+    
+}