diff --git a/query/src/main/java/tech/ydb/query/tools/QueryReader.java b/query/src/main/java/tech/ydb/query/tools/QueryReader.java index 4d5d76401..e4d277730 100644 --- a/query/src/main/java/tech/ydb/query/tools/QueryReader.java +++ b/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(); diff --git a/query/src/test/java/tech/ydb/query/tools/QueryReaderTest.java b/query/src/test/java/tech/ydb/query/tools/QueryReaderTest.java index 9dda8ef40..a568cbc46 100644 --- a/query/src/test/java/tech/ydb/query/tools/QueryReaderTest.java +++ b/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) { diff --git a/table/src/main/java/tech/ydb/table/result/ResultSetReader.java b/table/src/main/java/tech/ydb/table/result/ResultSetReader.java index b5c78805a..3781df403 100644 --- a/table/src/main/java/tech/ydb/table/result/ResultSetReader.java +++ b/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. + *

+ * Valid row indexes are in the range 0 <= index < getRowCount(), where + * {@code 0} is the first row and {@code getRowCount() - 1} is the last row. + * Passing a negative value positions the reader before 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 after 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. + *

+ * 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 + * (index < 0 || index >= getColumnCount()) */ 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 + * (index < 0 || index >= getColumnCount()) */ - 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 + * (index < 0 || index >= getColumnCount()) + * @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); } diff --git a/table/src/main/java/tech/ydb/table/result/impl/ProtoResultSetReader.java b/table/src/main/java/tech/ydb/table/result/impl/ProtoResultSetReader.java index 1d8cbd862..76ca61bdf 100644 --- a/table/src/main/java/tech/ydb/table/result/impl/ProtoResultSetReader.java +++ b/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 columnIndexes; // TODO: use better data structure - private final AbstractValueReader[] columnReaders; + private final ValueProtos.ResultSet rs; + private final Map 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; } } diff --git a/table/src/main/java/tech/ydb/table/result/impl/ProtoValueReaders.java b/table/src/main/java/tech/ydb/table/result/impl/ProtoValueReaders.java index addb7842f..ae0793a19 100644 --- a/table/src/main/java/tech/ydb/table/result/impl/ProtoValueReaders.java +++ b/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 resultSets) { // TODO: add lightweight implementation instead of proto joining Preconditions.checkArgument(!resultSets.isEmpty(), "Expect multiple result sets to join from"); diff --git a/table/src/test/java/tech/ydb/table/result/ResultSetReaderTest.java b/table/src/test/java/tech/ydb/table/result/ResultSetReaderTest.java index 3313cedfc..9552708e6 100644 --- a/table/src/test/java/tech/ydb/table/result/ResultSetReaderTest.java +++ b/table/src/test/java/tech/ydb/table/result/ResultSetReaderTest.java @@ -2,9 +2,11 @@ import org.junit.Assert; import org.junit.Test; +import org.junit.function.ThrowingRunnable; import tech.ydb.proto.ValueProtos; import tech.ydb.table.result.impl.ProtoValueReaders; +import tech.ydb.table.values.PrimitiveType; import tech.ydb.table.values.proto.ProtoType; import tech.ydb.table.values.proto.ProtoValue; @@ -14,20 +16,118 @@ */ public class ResultSetReaderTest { + private void assertNotPositioned(ThrowingRunnable runnable) { + IllegalStateException ex = Assert.assertThrows(IllegalStateException.class, runnable); + Assert.assertEquals("ResultSetReader not positioned properly, perhaps you need to call next.", ex.getMessage()); + } + + private void assertIllegalArgumentException(String msg, ThrowingRunnable runnable) { + IllegalArgumentException ex = Assert.assertThrows(IllegalArgumentException.class, runnable); + Assert.assertEquals(msg, ex.getMessage()); + } + + @Test + public void columnInfoTest() { + ValueProtos.ResultSet resultSet = ValueProtos.ResultSet.newBuilder() + .addColumns(newColumn("name", ProtoType.getText())) + .addColumns(newColumn("year", ProtoType.getUint32())) + .addRows(newRow(ProtoValue.fromText("Edgy"), ProtoValue.fromUint32(2006))) + .build(); + + ResultSetReader reader = ProtoValueReaders.forResultSet(resultSet); + + Assert.assertEquals("name", reader.getColumnName(0)); + Assert.assertEquals(PrimitiveType.Text, reader.getColumnType(0)); + Assert.assertEquals(0, reader.getColumnIndex("name")); + + Assert.assertEquals("year", reader.getColumnName(1)); + Assert.assertEquals(PrimitiveType.Uint32, reader.getColumnType(1)); + Assert.assertEquals(1, reader.getColumnIndex("year")); + + Assert.assertEquals(-1, reader.getColumnIndex("test")); // not found + assertIllegalArgumentException("Column index: -1, columns count: 2", () -> reader.getColumnName(-1)); + assertIllegalArgumentException("Column index: 2, columns count: 2", () -> reader.getColumnName(2)); + assertIllegalArgumentException("Column index: -1, columns count: 2", () -> reader.getColumnType(-1)); + assertIllegalArgumentException("Column index: 2, columns count: 2", () -> reader.getColumnType(2)); + + // before first row state + assertNotPositioned(() -> reader.getColumn("name")); + assertNotPositioned(() -> reader.getColumn(0)); + + Assert.assertTrue(reader.next()); // first row + + assertIllegalArgumentException("Column index: -1, columns count: 2", () -> reader.getColumn(-1)); + assertIllegalArgumentException("Column index: 2, columns count: 2", () -> reader.getColumn(2)); + + assertIllegalArgumentException("Unknown column 'test'", () -> reader.getColumn("test")); + } + + @Test + public void iterateReaderTest() { + ValueProtos.ResultSet resultSet = ValueProtos.ResultSet.newBuilder() + .addColumns(newColumn("name", ProtoType.getText())) + .addColumns(newColumn("year", ProtoType.getUint32())) + .addRows(newRow(ProtoValue.fromText("Edgy"), ProtoValue.fromUint32(2006))) + .addRows(newRow(ProtoValue.fromText("Feisty"), ProtoValue.fromUint32(2007))) + .addRows(newRow(ProtoValue.fromText("Gutsy"), ProtoValue.fromUint32(2007))) + .addRows(newRow(ProtoValue.fromText("Hardy"), ProtoValue.fromUint32(2008))) + .build(); + + ResultSetReader reader = ProtoValueReaders.forResultSet(resultSet); + + // before first row state + assertNotPositioned(() -> reader.getColumn("name")); + assertNotPositioned(() -> reader.getColumn(0)); + + Assert.assertTrue(reader.next()); // first row + Assert.assertEquals("Edgy", reader.getColumn("name").getText()); + Assert.assertEquals(2006, reader.getColumn("year").getUint32()); + + reader.setRowIndex(-10); // reset to before first + assertNotPositioned(() -> reader.getColumn("name")); + assertNotPositioned(() -> reader.getColumn(0)); + + Assert.assertTrue(reader.next()); // first row + Assert.assertEquals("Edgy", reader.getColumn("name").getText()); + Assert.assertEquals(2006, reader.getColumn("year").getUint32()); + + reader.setRowIndex(3); // to last row + Assert.assertEquals("Hardy", reader.getColumn("name").getText()); + Assert.assertEquals(2008, reader.getColumn("year").getUint32()); + + Assert.assertFalse(reader.next()); // after last row + assertNotPositioned(() -> reader.getColumn("name")); + assertNotPositioned(() -> reader.getColumn(0)); + + reader.setRowIndex(0); // reset to first + Assert.assertEquals("Edgy", reader.getColumn("name").getText()); + Assert.assertEquals(2006, reader.getColumn("year").getUint32()); + + Assert.assertTrue(reader.next()); // second row + Assert.assertEquals("Feisty", reader.getColumn("name").getText()); + Assert.assertEquals(2007, reader.getColumn("year").getUint32()); + + reader.setRowIndex(1000); // after last row + assertNotPositioned(() -> reader.getColumn("name")); + assertNotPositioned(() -> reader.getColumn(0)); + } + @Test public void readPrimitives() { ValueProtos.ResultSet resultSet = ValueProtos.ResultSet.newBuilder() - .addColumns(newColumn("name", ProtoType.getText())) - .addColumns(newColumn("year", ProtoType.getUint32())) - .addRows(newRow(ProtoValue.fromText("Edgy"), ProtoValue.fromUint32(2006))) - .addRows(newRow(ProtoValue.fromText("Feisty"), ProtoValue.fromUint32(2007))) - .addRows(newRow(ProtoValue.fromText("Gutsy"), ProtoValue.fromUint32(2007))) - .addRows(newRow(ProtoValue.fromText("Hardy"), ProtoValue.fromUint32(2008))) - .build(); + .addColumns(newColumn("name", ProtoType.getText())) + .addColumns(newColumn("year", ProtoType.getUint32())) + .addRows(newRow(ProtoValue.fromText("Edgy"), ProtoValue.fromUint32(2006))) + .addRows(newRow(ProtoValue.fromText("Feisty"), ProtoValue.fromUint32(2007))) + .addRows(newRow(ProtoValue.fromText("Gutsy"), ProtoValue.fromUint32(2007))) + .addRows(newRow(ProtoValue.fromText("Hardy"), ProtoValue.fromUint32(2008))) + .setTruncated(true) + .build(); ResultSetReader reader = ProtoValueReaders.forResultSet(resultSet); Assert.assertEquals(2, reader.getColumnCount()); Assert.assertEquals(4, reader.getRowCount()); + Assert.assertTrue(reader.isTruncated()); Assert.assertTrue(reader.next()); Assert.assertSame(reader.getColumn("name"), reader.getColumn(0)); @@ -54,24 +154,26 @@ public void readPrimitives() { @Test public void readUnsignedInts() { ValueProtos.ResultSet resultSet = ValueProtos.ResultSet.newBuilder() - .addColumns(newColumn("u8", ProtoType.getUint8())) - .addColumns(newColumn("u16", ProtoType.getUint16())) - .addColumns(newColumn("u32", ProtoType.getUint32())) - .addColumns(newColumn("u64", ProtoType.getUint64())) - .addRows(newRow( - ProtoValue.fromUint8(1), ProtoValue.fromUint16(1), - ProtoValue.fromUint32(1), ProtoValue.fromUint64(1))) - .addRows(newRow( - ProtoValue.fromUint8(-1), ProtoValue.fromUint16(-1), - ProtoValue.fromUint32(-1), ProtoValue.fromUint64(-1))) - .addRows(newRow( - ProtoValue.fromUint8(0xFF), ProtoValue.fromUint16(0xFFFF), - ProtoValue.fromUint32(0xFFFFFFFFl), ProtoValue.fromUint64(0xFFFFFFFFFFFFFFFFl))) - .build(); + .addColumns(newColumn("u8", ProtoType.getUint8())) + .addColumns(newColumn("u16", ProtoType.getUint16())) + .addColumns(newColumn("u32", ProtoType.getUint32())) + .addColumns(newColumn("u64", ProtoType.getUint64())) + .addRows(newRow( + ProtoValue.fromUint8(1), ProtoValue.fromUint16(1), + ProtoValue.fromUint32(1), ProtoValue.fromUint64(1))) + .addRows(newRow( + ProtoValue.fromUint8(-1), ProtoValue.fromUint16(-1), + ProtoValue.fromUint32(-1), ProtoValue.fromUint64(-1))) + .addRows(newRow( + ProtoValue.fromUint8(0xFF), ProtoValue.fromUint16(0xFFFF), + ProtoValue.fromUint32(0xFFFFFFFFl), ProtoValue.fromUint64(0xFFFFFFFFFFFFFFFFl))) + .setTruncated(false) + .build(); ResultSetReader reader = ProtoValueReaders.forResultSet(resultSet); Assert.assertEquals(4, reader.getColumnCount()); Assert.assertEquals(3, reader.getRowCount()); + Assert.assertFalse(reader.isTruncated()); Assert.assertTrue(reader.next()); Assert.assertSame(reader.getColumn("u8"), reader.getColumn(0));