Examples with SplitReader co.cask.cdap.api.data.batch.SplitReader used on opensource projects

Example 1 with SplitReader

use of co.cask.cdap.api.data.batch.SplitReader in project cdap by caskdata.

the class ObjectMappedTableDatasetTest method testGetSplits.

@Test
public void testGetSplits() throws Exception {
    dsFrameworkUtil.createInstance(ObjectMappedTable.class.getName(), RECORDS_ID, ObjectMappedTableProperties.builder().setType(Record.class).build());
    try {
        final ObjectMappedTableDataset<Record> records = dsFrameworkUtil.getInstance(RECORDS_ID);
        TransactionExecutor txnl = dsFrameworkUtil.newInMemoryTransactionExecutor((TransactionAware) records);
        final Record record = new Record(Integer.MAX_VALUE, Long.MAX_VALUE, Float.MAX_VALUE, Double.MAX_VALUE, "foobar", Bytes.toBytes("foobar"), ByteBuffer.wrap(Bytes.toBytes("foobar")), UUID.randomUUID());
        final byte[] rowkey = Bytes.toBytes("row1");
        txnl.execute(new TransactionExecutor.Subroutine() {

            @Override
            public void apply() throws Exception {
                records.write(rowkey, record);
            }
        });
        // should not include the record, since upper bound is not inclusive
        txnl.execute(new TransactionExecutor.Subroutine() {

            @Override
            public void apply() throws Exception {
                List<Split> splits = records.getSplits(1, null, rowkey);
                List<Record> recordsRead = new ArrayList<>();
                for (Split split : splits) {
                    SplitReader<byte[], Record> splitReader = records.createSplitReader(split);
                    try {
                        splitReader.initialize(split);
                        while (splitReader.nextKeyValue()) {
                            recordsRead.add(splitReader.getCurrentValue());
                        }
                    } finally {
                        splitReader.close();
                    }
                }
                Assert.assertEquals(0, recordsRead.size());
            }
        });
        // should include the record, since lower bound is inclusive
        txnl.execute(new TransactionExecutor.Subroutine() {

            @Override
            public void apply() throws Exception {
                List<Split> splits = records.getSplits(1, rowkey, null);
                List<Record> recordsRead = new ArrayList<>();
                for (Split split : splits) {
                    SplitReader<byte[], Record> splitReader = records.createSplitReader(split);
                    try {
                        splitReader.initialize(split);
                        while (splitReader.nextKeyValue()) {
                            recordsRead.add(splitReader.getCurrentValue());
                        }
                    } finally {
                        splitReader.close();
                    }
                }
                Assert.assertEquals(1, recordsRead.size());
                Assert.assertEquals(record, recordsRead.get(0));
            }
        });
    } finally {
        dsFrameworkUtil.deleteInstance(RECORDS_ID);
    }
}

Example 2 with SplitReader

use of co.cask.cdap.api.data.batch.SplitReader in project cdap by caskdata.

the class ObjectStoreDatasetTest method testBatchCustomList.

@Test
public void testBatchCustomList() throws Exception {
    DatasetId customlist = DatasetFrameworkTestUtil.NAMESPACE_ID.dataset("customlist");
    createObjectStoreInstance(customlist, new TypeToken<List<Custom>>() {
    }.getType());
    final ObjectStoreDataset<List<Custom>> customStore = dsFrameworkUtil.getInstance(customlist);
    TransactionExecutor txnl = dsFrameworkUtil.newInMemoryTransactionExecutor(customStore);
    final SortedSet<Long> keysWritten = Sets.newTreeSet();
    txnl.execute(new TransactionExecutor.Subroutine() {

        @Override
        public void apply() throws Exception {
            List<Custom> customList1 = Arrays.asList(new Custom(1, Lists.newArrayList("one", "ONE")), new Custom(2, Lists.newArrayList("two", "TWO")));
            Random rand = new Random(100);
            long key1 = rand.nextLong();
            keysWritten.add(key1);
            customStore.write(Bytes.toBytes(key1), customList1);
            List<Custom> customList2 = Arrays.asList(new Custom(3, Lists.newArrayList("three", "THREE")), new Custom(4, Lists.newArrayList("four", "FOUR")));
            long key2 = rand.nextLong();
            keysWritten.add(key2);
            customStore.write(Bytes.toBytes(key2), customList2);
        }
    });
    final SortedSet<Long> keysWrittenCopy = ImmutableSortedSet.copyOf(keysWritten);
    txnl.execute(new TransactionExecutor.Subroutine() {

        @Override
        public void apply() throws Exception {
            // get the splits for the table
            List<Split> splits = customStore.getSplits();
            for (Split split : splits) {
                SplitReader<byte[], List<Custom>> reader = customStore.createSplitReader(split);
                reader.initialize(split);
                while (reader.nextKeyValue()) {
                    byte[] key = reader.getCurrentKey();
                    Assert.assertTrue(keysWritten.remove(Bytes.toLong(key)));
                }
            }
            // verify all keys have been read
            if (!keysWritten.isEmpty()) {
                System.out.println("Remaining [" + keysWritten.size() + "]: " + keysWritten);
            }
            Assert.assertTrue(keysWritten.isEmpty());
        }
    });
    deleteAndVerifyInBatch(customStore, txnl, keysWrittenCopy);
    dsFrameworkUtil.deleteInstance(customlist);
}