-
Notifications
You must be signed in to change notification settings - Fork 13
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
better-spliterator: Better spliterator
- Loading branch information
Alexander Lavrukov
committed
Jul 2, 2024
1 parent
58cba6e
commit 030388c
Showing
12 changed files
with
499 additions
and
31 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
7 changes: 7 additions & 0 deletions
7
...ory-ydb-v2/src/main/java/tech/ydb/yoj/repository/ydb/spliterator/ClosableSpliterator.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,7 @@ | ||
| package tech.ydb.yoj.repository.ydb.spliterator; | ||
|
|
||
| import java.util.Spliterator; | ||
|
|
||
| public interface ClosableSpliterator<V> extends Spliterator<V> { | ||
| void close(); | ||
| } |
10 changes: 10 additions & 0 deletions
10
repository-ydb-v2/src/main/java/tech/ydb/yoj/repository/ydb/spliterator/ResultConverter.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,10 @@ | ||
| package tech.ydb.yoj.repository.ydb.spliterator; | ||
|
|
||
| import tech.ydb.proto.ValueProtos; | ||
|
|
||
| import java.util.List; | ||
|
|
||
| @FunctionalInterface | ||
| public interface ResultConverter<V> { | ||
| V convert(List<ValueProtos.Column> columns, ValueProtos.Value value); | ||
| } |
68 changes: 68 additions & 0 deletions
68
...itory-ydb-v2/src/main/java/tech/ydb/yoj/repository/ydb/spliterator/ResultSetIterator.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,68 @@ | ||
| package tech.ydb.yoj.repository.ydb.spliterator; | ||
|
|
||
| import tech.ydb.proto.ValueProtos; | ||
| import tech.ydb.table.result.ResultSetReader; | ||
| import tech.ydb.yoj.repository.ydb.client.YdbConverter; | ||
|
|
||
| import java.util.ArrayList; | ||
| import java.util.Iterator; | ||
| import java.util.List; | ||
| import java.util.NoSuchElementException; | ||
|
|
||
| public final class ResultSetIterator<V> implements Iterator<V> { | ||
| private final ResultSetReader resultSet; | ||
| private final ResultConverter<V> converter; | ||
| private final List<ValueProtos.Column> columns; | ||
|
|
||
| private int position = 0; | ||
|
|
||
| public ResultSetIterator(ResultSetReader resultSet, ResultConverter<V> converter) { | ||
| List<ValueProtos.Column> columns; | ||
| if (resultSet.getRowCount() > 0) { | ||
| resultSet.setRowIndex(0); | ||
| columns = getColumns(resultSet); | ||
| } else { | ||
| columns = new ArrayList<>(); | ||
| } | ||
|
|
||
| this.resultSet = resultSet; | ||
| this.converter = converter; | ||
| this.columns = columns; | ||
| } | ||
|
|
||
| @Override | ||
| public boolean hasNext() { | ||
| return position < resultSet.getRowCount(); | ||
| } | ||
|
|
||
| @Override | ||
| public V next() { | ||
| if (!hasNext()) { | ||
| throw new NoSuchElementException(); | ||
| } | ||
|
|
||
| ValueProtos.Value value = buildValue(position++); | ||
|
|
||
| return converter.convert(columns, value); | ||
| } | ||
|
|
||
| private ValueProtos.Value buildValue(int rowIndex) { | ||
| resultSet.setRowIndex(rowIndex); | ||
| ValueProtos.Value.Builder value = ValueProtos.Value.newBuilder(); | ||
| for (int i = 0; i < columns.size(); i++) { | ||
| value.addItems(YdbConverter.convertValueToProto(resultSet.getColumn(i))); | ||
| } | ||
| return value.build(); | ||
| } | ||
|
|
||
| private static List<ValueProtos.Column> getColumns(ResultSetReader resultSet) { | ||
| List<ValueProtos.Column> columns = new ArrayList<>(); | ||
| for (int i = 0; i < resultSet.getColumnCount(); i++) { | ||
| columns.add(ValueProtos.Column.newBuilder() | ||
| .setName(resultSet.getColumnName(i)) | ||
| .build() | ||
| ); | ||
| } | ||
| return columns; | ||
| } | ||
| } |
77 changes: 77 additions & 0 deletions
77
repository-ydb-v2/src/main/java/tech/ydb/yoj/repository/ydb/spliterator/YdbSpliterator.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,77 @@ | ||
| package tech.ydb.yoj.repository.ydb.spliterator; | ||
|
|
||
| import tech.ydb.yoj.ExperimentalApi; | ||
|
|
||
| import java.util.Iterator; | ||
| import java.util.Spliterator; | ||
| import java.util.function.Consumer; | ||
| import java.util.stream.Stream; | ||
| import java.util.stream.StreamSupport; | ||
|
|
||
| @ExperimentalApi(issue = "https://github.com/ydb-platform/yoj-project/issues/42") | ||
| public final class YdbSpliterator<V> implements ClosableSpliterator<V> { | ||
| private final YdbSpliteratorQueue<Iterator<V>> queue; | ||
| private final int flags; | ||
|
|
||
| private Iterator<V> valueIterator; | ||
|
|
||
| private boolean closed = false; | ||
|
|
||
| public YdbSpliterator(YdbSpliteratorQueue<Iterator<V>> queue, boolean isOrdered) { | ||
| this.queue = queue; | ||
| this.flags = (isOrdered ? ORDERED : 0) | NONNULL; | ||
| } | ||
|
|
||
| // Correct way to create stream with YdbSpliterator. onClose call is important for avoid supplier thread leak. | ||
| public Stream<V> createStream() { | ||
| return StreamSupport.stream(this, false).onClose(this::close); | ||
| } | ||
|
|
||
| @Override | ||
| public boolean tryAdvance(Consumer<? super V> action) { | ||
| if (closed) { | ||
| return false; | ||
| } | ||
|
|
||
| if (valueIterator == null || !valueIterator.hasNext()) { | ||
| valueIterator = queue.poll(); | ||
| if (valueIterator == null || !valueIterator.hasNext()) { | ||
| closed = true; | ||
| return false; | ||
| } | ||
| } | ||
|
|
||
| V value = valueIterator.next(); | ||
|
|
||
| action.accept(value); | ||
|
|
||
| return true; | ||
| } | ||
|
|
||
| @Override | ||
| public void close() { | ||
| closed = true; | ||
| queue.close(); | ||
| } | ||
|
|
||
| @Override | ||
| public Spliterator<V> trySplit() { | ||
| return null; | ||
| } | ||
|
|
||
| @Override | ||
| public long estimateSize() { | ||
| return Long.MAX_VALUE; | ||
| } | ||
|
|
||
| @Override | ||
| public long getExactSizeIfKnown() { | ||
| return -1; | ||
| } | ||
|
|
||
| @Override | ||
| public int characteristics() { | ||
| return flags; | ||
| } | ||
|
|
||
| } |
Oops, something went wrong.