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

Example 16 with TableKey

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

the class KeyValueTableTestBase method checkGlobalIterator.

private void checkGlobalIterator(KeyValueTable keyValueTable, Versions versions, int iteration) {
    val itemsAtOnce = getSecondaryKeyCount() / 5;
    // All expected Entries, indexed by Keys.
    val allEntries = new HashMap<TableKey, TableEntry>();
    forEveryKey((pk, sk) -> {
        val k = new TableKey(pk, sk);
        val e = new TableEntry(k, versions.get(getUniqueKeyId(k.getPrimaryKey(), k.getSecondaryKey())), getValue(pk, sk, iteration));
        allEntries.put(k, e);
    });
    // All Keys, sorted.
    val allKeys = allEntries.keySet().stream().sorted(KEY_COMPARATOR).collect(Collectors.toList());
    // All unique Primary Keys.
    val allPrimaryKeys = allKeys.stream().map(TableKey::getPrimaryKey).distinct().sorted(KEY_COMPARATOR::compare).collect(Collectors.toList());
    // Issue various iterators (range) and verify against expected values.
    for (int startIndex = 0; startIndex < allPrimaryKeys.size() / 2; startIndex++) {
        val endIndex = allPrimaryKeys.size() - startIndex - 1;
        val firstPK = allPrimaryKeys.get(startIndex);
        val lastPK = allPrimaryKeys.get(endIndex);
        val expectedKeys = allKeys.stream().filter(k -> KEY_COMPARATOR.compare(k.getPrimaryKey(), firstPK) >= 0 && KEY_COMPARATOR.compare(k.getPrimaryKey(), lastPK) <= 0).collect(Collectors.toList());
        val iterator = keyValueTable.iterator().maxIterationSize(itemsAtOnce).forRange(firstPK, lastPK);
        val iteratorKeys = new ArrayList<TableKey>();
        iterator.keys().collectRemaining(ii -> iteratorKeys.addAll(ii.getItems())).join();
        AssertExtensions.assertListEquals("Unexpected keys returned from iterator.keys()", expectedKeys, iteratorKeys, this::areEqual);
        val expectedEntries = expectedKeys.stream().map(allEntries::get).collect(Collectors.toList());
        val iteratorEntries = new ArrayList<TableEntry>();
        iterator.entries().collectRemaining(ii -> iteratorEntries.addAll(ii.getItems())).join();
        AssertExtensions.assertListEquals("Unexpected entries returned from iterator.entries()", expectedEntries, iteratorEntries, (e1, e2) -> areEqual(e1, e2) && e1.getVersion().equals(e2.getVersion()));
    }
    // Check all() iterator.
    val allIterator = keyValueTable.iterator().maxIterationSize(itemsAtOnce).all();
    val expectedAllEntries = allKeys.stream().map(allEntries::get).collect(Collectors.toList());
    val allIteratorEntries = new ArrayList<TableEntry>();
    allIterator.entries().collectRemaining(ii -> allIteratorEntries.addAll(ii.getItems())).join();
    AssertExtensions.assertListEquals("Unexpected entries returned from allIterator.entries()", expectedAllEntries, allIteratorEntries, (e1, e2) -> areEqual(e1, e2) && e1.getVersion().equals(e2.getVersion()));
}

Example 17 with TableKey

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

the class KeyValueTableTestBase method testIterators.

@Test
public void testIterators() {
    @Cleanup val kvt = createKeyValueTable();
    val iteration = new AtomicInteger(0);
    // Populate everything.
    val versions = new Versions();
    forEveryPrimaryKey((pk, secondaryKeys) -> {
        List<TableModification> inserts = secondaryKeys.stream().map(sk -> new Insert(new TableKey(pk, sk), getValue(pk, sk, iteration.get()))).collect(Collectors.toList());
        val keyVersions = kvt.update(inserts).join();
        for (int i = 0; i < inserts.size(); i++) {
            versions.add(getUniqueKeyId(pk, inserts.get(i).getKey().getSecondaryKey()), keyVersions.get(i));
        }
    });
    // Check the Primary Key iterator.
    checkPrimaryKeyIterator(kvt, versions, iteration.get());
    // Check the Global iterators.
    checkGlobalIterator(kvt, versions, iteration.get());
}

Example 18 with TableKey

use of io.pravega.client.tables.TableKey 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);
}

Example 19 with TableKey

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

the class KeyValueTableIteratorImplTests method testGlobalIterators.

@Test
public void testGlobalIterators() {
    val pk1 = newBuffer(DEFAULT_CONFIG.getPrimaryKeyLength());
    val pk2 = newBuffer(DEFAULT_CONFIG.getPrimaryKeyLength());
    val maxIterationSize = 10;
    val c = new KeyValueTableIteratorImpl.TableKeyComparator();
    val allEntries = IntStream.range(0, maxIterationSize).mapToObj(i -> new AbstractMap.SimpleImmutableEntry<>(newBuffer(TOTAL_KEY_LENGTH), newBuffer(5))).sorted((e1, e2) -> c.compare(e1.getKey(), e2.getKey())).collect(Collectors.toList());
    val mockSegment = mock(TableSegment.class);
    when(mockSegment.keyIterator(any())).thenAnswer(arg -> {
        val iteratorArgs = (SegmentIteratorArgs) arg.getArgument(0);
        checkSegmentIteratorArgs(iteratorArgs, pk1, pk2, 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, pk1, pk2, 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());
    when(selector.getAllTableSegments()).thenAnswer(arg -> IntStream.range(0, DEFAULT_CONFIG.getPartitionCount()).mapToObj(i -> mockSegment).collect(Collectors.toList()));
    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.forRange(pk1, pk2);
    // 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() * DEFAULT_CONFIG.getPartitionCount(), iteratorKeys.size());
    Assert.assertEquals(allEntries.size() * DEFAULT_CONFIG.getPartitionCount(), iteratorEntries.size());
    for (int i = 0; i < allEntries.size(); i++) {
        val expected = allEntries.get(i);
        for (int p = 0; p < DEFAULT_CONFIG.getPartitionCount(); p++) {
            val actualKey = iteratorKeys.get(i * DEFAULT_CONFIG.getPartitionCount() + p);
            val actualEntry = iteratorEntries.get(i * DEFAULT_CONFIG.getPartitionCount() + p);
            Assert.assertEquals(Unpooled.wrappedBuffer(expected.getKey()), entryHelper.serializeKey(actualKey.getPrimaryKey(), actualKey.getSecondaryKey()));
            Assert.assertEquals(actualKey, actualEntry.getKey());
            Assert.assertEquals(expected.getValue(), actualEntry.getValue());
        }
    }
}