-
Notifications
You must be signed in to change notification settings - Fork 12
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 f28a6e0
Showing
12 changed files
with
432 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; | ||
| } | ||
| } |
106 changes: 106 additions & 0 deletions
106
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,106 @@ | ||
| package tech.ydb.yoj.repository.ydb.spliterator; | ||
|
|
||
| import tech.ydb.yoj.ExperimentalApi; | ||
|
|
||
| import javax.annotation.Nullable; | ||
| 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 void close() { | ||
| if (closed) { | ||
| return; | ||
| } | ||
| closed = true; | ||
| queue.close(); | ||
| } | ||
|
|
||
| @Override | ||
| public boolean tryAdvance(Consumer<? super V> action) { | ||
| if (closed) { | ||
| return false; | ||
| } | ||
|
|
||
| // WARNING: At one point in time, this spliterator will store up to queue.size() + 2 blocks from YDB in memory. | ||
| // One block right here, one in the queue, one in the grpc thread, waiting for free space in the queue. | ||
| valueIterator = getValueIterator(valueIterator, queue); | ||
| if (valueIterator == null) { | ||
| close(); | ||
| return false; | ||
| } | ||
|
|
||
| V value = valueIterator.next(); | ||
|
|
||
| action.accept(value); | ||
|
|
||
| return true; | ||
| } | ||
|
|
||
| /* | ||
| * Returns not empty valueIterator, null in case of end of stream | ||
| */ | ||
| @Nullable | ||
| private static <V> Iterator<V> getValueIterator( | ||
| @Nullable Iterator<V> valueIterator, YdbSpliteratorQueue<Iterator<V>> queue | ||
| ) { | ||
| // valueIterator could be null only on first call of tryAdvance | ||
| if (valueIterator == null) { | ||
| valueIterator = queue.poll(); | ||
| if (valueIterator == null) { | ||
| return null; | ||
| } | ||
| } | ||
|
|
||
| // queue could return empty iterator, we have to select one with elements | ||
| while (!valueIterator.hasNext()) { | ||
| valueIterator = queue.poll(); | ||
| if (valueIterator == null) { | ||
| return null; | ||
| } | ||
| } | ||
|
|
||
| return valueIterator; | ||
| } | ||
|
|
||
| @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.