Examples with IteratorItem io.pravega.client.tables.IteratorItem used on opensource projects

Example 1 with IteratorItem

use of io.pravega.client.tables.IteratorItem in project pravega by pravega.

the class TableSegmentIteratorTests method testIterator.

@Test
public void testIterator() {
    val items = new TreeMap<ByteBuf, ByteBuf>();
    for (int i = 0; i < 256; i++) {
        val value = Unpooled.wrappedBuffer(new byte[] { (byte) i });
        items.put(value, value);
    }
    val initialArgs = new SegmentIteratorArgs(items.firstKey(), items.lastKey(), 2);
    val tsi = new TableSegmentIterator<>(args -> CompletableFuture.completedFuture(new IteratorItem<>(items.subMap(args.getFromKey(), true, args.getToKey(), true).keySet().stream().limit(args.getMaxItemsAtOnce()).collect(Collectors.toList()))), r -> r, initialArgs);
    val result = new ArrayList<ByteBuf>();
    tsi.collectRemaining(i -> {
        Assert.assertEquals(initialArgs.getMaxItemsAtOnce(), i.getItems().size());
        result.addAll(i.getItems());
        return true;
    }).join();
    val expectedResult = new ArrayList<>(items.keySet());
    AssertExtensions.assertListEquals("Unexpected result.", expectedResult, result, ByteBuf::equals);
    Assert.assertNull("Not expecting any more items.", tsi.getNext().join());
}

Example 2 with IteratorItem

use of io.pravega.client.tables.IteratorItem in project pravega by pravega.

the class TableSegmentIterator method getNext.

@Override
public CompletableFuture<IteratorItem<T>> getNext() {
    if (this.args.get() == null) {
        // We are done.
        return CompletableFuture.completedFuture(null);
    }
    CompletableFuture<IteratorItem<T>> result = this.fetchNext.apply(this.args.get()).thenApply(r -> {
        if (r == null) {
            // We are done.
            this.args.set(null);
            return null;
        }
        ByteBuf lastKey = r.getItems().isEmpty() ? null : this.getKey.apply(r.getItems().get(r.getItems().size() - 1));
        this.args.set(this.args.get().next(lastKey));
        return r;
    });
    Futures.exceptionListener(result, ObjectClosedException.class, ex -> this.args.set(null));
    return result;
}

Example 3 with IteratorItem

use of io.pravega.client.tables.IteratorItem in project pravega by pravega.

the class TableSegmentImplTest method testIterator.

private <T> void testIterator(Function<SegmentIteratorArgs, AsyncIterator<IteratorItem<T>>> newIterator, Supplier<ByteBuf> getLastRequestFromKey, Supplier<ByteBuf> getLastRequestToKey, Function<TableSegmentEntry, T> getItemFromEntry, Consumer<List<T>> sendReply, BiPredicate<T, T> checkItemEquality) throws Exception {
    val suggestedKeyCount = 3;
    // Generate 100 Entries and split them into batches.
    val allEntries = IntStream.range(0, 100).mapToObj(i -> versionedEntry(i * 10L, Integer.toString(i * 10), 1L)).collect(Collectors.toList());
    val inputEntries = splitIteratorInputs(allEntries);
    // Do an full iteration as well.
    inputEntries.add(allEntries);
    // Check regular iteration.
    for (int i = 0; i < inputEntries.size(); i++) {
        val entryList = inputEntries.get(i);
        SegmentIteratorArgs args = SegmentIteratorArgs.builder().fromKey(entryList.get(0).getKey().getKey()).toKey(entryList.get(entryList.size() - 1).getKey().getKey()).maxItemsAtOnce(suggestedKeyCount).build();
        // We collect iterated items in this list.
        val actualItems = new ArrayList<T>();
        val itemsToReturn = entryList.iterator();
        val tableIterator = newIterator.apply(args);
        while (itemsToReturn.hasNext()) {
            val iteratorFuture = tableIterator.getNext();
            // Verify the wire command got sent as expected.
            val requestFromKey = getLastRequestFromKey.get();
            Assert.assertEquals("Unexpected fromKey sent.", args.getFromKey(), requestFromKey);
            val requestToKey = getLastRequestToKey.get();
            Assert.assertEquals("Unexpected toKey sent.", args.getToKey(), requestToKey);
            // Send a reply.
            val expectedResult = new ArrayList<T>();
            int count = suggestedKeyCount;
            while (itemsToReturn.hasNext() && count > 0) {
                val next = itemsToReturn.next();
                expectedResult.add(getItemFromEntry.apply(next));
                args = args.next(next.getKey().getKey());
                count--;
            }
            sendReply.accept(expectedResult);
            // Check the partial result.
            val iteratorResult = iteratorFuture.get(SHORT_TIMEOUT, TimeUnit.MILLISECONDS);
            AssertExtensions.assertListEquals("Unexpected partial result.", expectedResult, iteratorResult.getItems(), checkItemEquality);
            actualItems.addAll(iteratorResult.getItems());
        }
        // Then the final result.
        val expected = entryList.stream().map(getItemFromEntry).collect(Collectors.toList());
        AssertExtensions.assertListEquals("Unexpected result.", expected, actualItems, checkItemEquality);
    }
}

Example 4 with IteratorItem

use of io.pravega.client.tables.IteratorItem in project pravega by pravega.

the class KeyValueTableIteratorImplTests method testIteratorSingleSegment.

/**
 * Tests the iterator when there is a single segment involved.
 */
@Test
public void testIteratorSingleSegment() {
    val pk = newBuffer(DEFAULT_CONFIG.getPrimaryKeyLength());
    val sk1 = newBuffer(DEFAULT_CONFIG.getSecondaryKeyLength());
    val sk2 = newBuffer(DEFAULT_CONFIG.getSecondaryKeyLength());
    val maxIterationSize = 10;
    val allEntries = IntStream.range(0, maxIterationSize).mapToObj(i -> new AbstractMap.SimpleImmutableEntry<>(newBuffer(TOTAL_KEY_LENGTH), newBuffer(5))).collect(Collectors.toList());
    val mockSegment = mock(TableSegment.class);
    when(mockSegment.keyIterator(any())).thenAnswer(arg -> {
        val iteratorArgs = (SegmentIteratorArgs) arg.getArgument(0);
        checkSegmentIteratorArgs(iteratorArgs, pk, sk1, sk2, maxIterationSize);
        val keys = allEntries.stream().map(e -> new TableSegmentKey(Unpooled.wrappedBuffer(e.getKey()), TableSegmentKeyVersion.NO_VERSION)).collect(Collectors.toList());
        return AsyncIterator.singleton(new IteratorItem<>(keys));
    });
    when(mockSegment.entryIterator(any())).thenAnswer(arg -> {
        val iteratorArgs = (SegmentIteratorArgs) arg.getArgument(0);
        checkSegmentIteratorArgs(iteratorArgs, pk, sk1, sk2, maxIterationSize);
        val entries = allEntries.stream().map(e -> new TableSegmentEntry(new TableSegmentKey(Unpooled.wrappedBuffer(e.getKey()), TableSegmentKeyVersion.NO_VERSION), Unpooled.wrappedBuffer(e.getValue()))).collect(Collectors.toList());
        return AsyncIterator.singleton(new IteratorItem<>(entries));
    });
    val selector = mock(SegmentSelector.class);
    when(selector.getKvt()).thenReturn(KVT);
    when(selector.getSegmentCount()).thenReturn(DEFAULT_CONFIG.getPartitionCount());
    val segmentRequests = new HashSet<ByteBuffer>();
    when(selector.getTableSegment(any())).thenAnswer(arg -> {
        val key = (ByteBuffer) arg.getArgument(0);
        segmentRequests.add(key.duplicate());
        return mockSegment;
    });
    val entryHelper = new TableEntryHelper(selector, DEFAULT_CONFIG);
    val b = new KeyValueTableIteratorImpl.Builder(DEFAULT_CONFIG, entryHelper, executorService()).maxIterationSize(maxIterationSize);
    // Everything up until here has been verified in the Builder unit tests (no need to do it again).
    // Create a Primary Key range iterator.
    val iterator = b.forPrimaryKey(pk, sk1, sk2);
    // Issue a Key iterator and then an Entry iterator, then collect all the results.
    val iteratorKeys = new ArrayList<TableKey>();
    iterator.keys().collectRemaining(ii -> iteratorKeys.addAll(ii.getItems())).join();
    val iteratorEntries = new ArrayList<TableEntry>();
    iterator.entries().collectRemaining(ii -> iteratorEntries.addAll(ii.getItems())).join();
    // Validate the results are as expected.
    Assert.assertEquals(allEntries.size(), iteratorKeys.size());
    Assert.assertEquals(allEntries.size(), iteratorEntries.size());
    for (int i = 0; i < allEntries.size(); i++) {
        val expected = allEntries.get(i);
        val actualKey = iteratorKeys.get(i);
        val actualEntry = iteratorEntries.get(i);
        Assert.assertEquals(Unpooled.wrappedBuffer(expected.getKey()), entryHelper.serializeKey(actualKey.getPrimaryKey(), actualKey.getSecondaryKey()));
        Assert.assertEquals(actualKey, actualEntry.getKey());
        Assert.assertEquals(expected.getValue(), actualEntry.getValue());
    }
    Assert.assertEquals("Only expecting 1 segment to be requested.", 1, segmentRequests.size());
}

Example 5 with IteratorItem

use of io.pravega.client.tables.IteratorItem in project pravega by pravega.

the class KeyValueTableIteratorImplTests method testMergeAsyncIterator.

/**
 * Tests the {@link KeyValueTableIteratorImpl.MergeAsyncIterator} class.
 */
@Test
public void testMergeAsyncIterator() {
    val segmentCount = 5;
    val minItemsPerSegment = 11;
    val maxItemsPerSegment = 101;
    val iterationSize = 3;
    // Generate test data.
    val c = new KeyValueTableIteratorImpl.TableKeyComparator();
    // Sorted.
    val segmentIterators = new ArrayList<AsyncIterator<IteratorItem<TableKey>>>();
    // Sorted.
    val expectedData = new ArrayList<TableKey>();
    for (int i = 0; i < segmentCount; i++) {
        val count = minItemsPerSegment + random.nextInt(maxItemsPerSegment - minItemsPerSegment);
        val segmentBatches = new ArrayList<List<TableKey>>();
        // Generate segment contents.
        val segmentSortedContents = IntStream.range(0, count).mapToObj(x -> new TableKey(newBuffer(DEFAULT_CONFIG.getPrimaryKeyLength()), newBuffer(DEFAULT_CONFIG.getSecondaryKeyLength()))).sorted(c).collect(Collectors.toList());
        // Break it down into batches and create a "segment iterator" from them.
        int index = 0;
        while (index < count) {
            int batchCount = Math.min(iterationSize, count - index);
            segmentBatches.add(segmentSortedContents.subList(index, index + batchCount));
            index += batchCount;
        }
        segmentIterators.add(createAsyncIterator(segmentBatches));
        expectedData.addAll(segmentSortedContents);
    }
    expectedData.sort(c);
    // Create a merge iterator and collect its contents.
    val mergeIterator = new KeyValueTableIteratorImpl.MergeAsyncIterator<>(segmentIterators.iterator(), k -> k, iterationSize, executorService());
    val actualData = new ArrayList<TableKey>();
    mergeIterator.collectRemaining(ii -> {
        val expected = Math.min(iterationSize, expectedData.size() - actualData.size());
        Assert.assertEquals(expected, ii.getItems().size());
        return actualData.addAll(ii.getItems());
    });
    // Verify it returns the correct items.
    AssertExtensions.assertListEquals("", expectedData, actualData, TableKey::equals);
}