CassandraIO { public static <T> Read<T> read() { return new AutoValue_CassandraIO_Read.Builder<T>().build(); } private  CassandraIO(); static Read<T> read(); static Write<T> write(); static Write<T> delete();  }
@Test public void testEstimatedSizeBytes() throws Exception { PipelineOptions pipelineOptions = PipelineOptionsFactory.create(); CassandraIO.Read<Scientist> read = CassandraIO.<Scientist>read() .withHosts(Collections.singletonList(CASSANDRA_HOST)) .withPort(cassandraPort) .withKeyspace(CASSANDRA_KEYSPACE) .withTable(CASSANDRA_TABLE); CassandraIO.CassandraSource<Scientist> source = new CassandraIO.CassandraSource<>(read, null); long estimatedSizeBytes = source.getEstimatedSizeBytes(pipelineOptions); assertThat( "wrong estimated size in " + CASSANDRA_KEYSPACE + "/" + CASSANDRA_TABLE, estimatedSizeBytes, greaterThan(0L)); }
@Test public void testRead() throws Exception { PCollection<Scientist> output = pipeline.apply( CassandraIO.<Scientist>read() .withHosts(Collections.singletonList(CASSANDRA_HOST)) .withPort(cassandraPort) .withKeyspace(CASSANDRA_KEYSPACE) .withTable(CASSANDRA_TABLE) .withCoder(SerializableCoder.of(Scientist.class)) .withEntity(Scientist.class)); PAssert.thatSingleton(output.apply("Count", Count.globally())).isEqualTo(NUM_ROWS); PCollection<KV<String, Integer>> mapped = output.apply( MapElements.via( new SimpleFunction<Scientist, KV<String, Integer>>() { @Override public KV<String, Integer> apply(Scientist scientist) { return KV.of(scientist.name, scientist.id); } })); PAssert.that(mapped.apply("Count occurrences per scientist", Count.perKey())) .satisfies( input -> { for (KV<String, Long> element : input) { assertEquals(element.getKey(), NUM_ROWS / 10, element.getValue().longValue()); } return null; }); pipeline.run(); }
@Test public void testReadWithQuery() throws Exception { PCollection<Scientist> output = pipeline.apply( CassandraIO.<Scientist>read() .withHosts(Collections.singletonList(CASSANDRA_HOST)) .withPort(cassandraPort) .withKeyspace(CASSANDRA_KEYSPACE) .withTable(CASSANDRA_TABLE) .withQuery( "select person_id, writetime(person_name) from beam_ks.scientist where person_id=10") .withCoder(SerializableCoder.of(Scientist.class)) .withEntity(Scientist.class)); PAssert.thatSingleton(output.apply("Count", Count.globally())).isEqualTo(1L); PAssert.that(output) .satisfies( input -> { for (Scientist sci : input) { assertNull(sci.name); assertTrue(sci.nameTs != null && sci.nameTs > 0); } return null; }); pipeline.run(); }
@Test public void testReadWithMapper() throws Exception { counter.set(0); SerializableFunction<Session, Mapper> factory = new NOOPMapperFactory(); pipeline.apply( CassandraIO.<String>read() .withHosts(Collections.singletonList(CASSANDRA_HOST)) .withPort(cassandraPort) .withKeyspace(CASSANDRA_KEYSPACE) .withTable(CASSANDRA_TABLE) .withCoder(SerializableCoder.of(String.class)) .withEntity(String.class) .withMapperFactoryFn(factory)); pipeline.run(); assertEquals(NUM_ROWS, counter.intValue()); }
@Test public void testSplit() throws Exception { PipelineOptions options = PipelineOptionsFactory.create(); CassandraIO.Read<Scientist> read = CassandraIO.<Scientist>read() .withHosts(Collections.singletonList(CASSANDRA_HOST)) .withPort(cassandraPort) .withKeyspace(CASSANDRA_KEYSPACE) .withTable(CASSANDRA_TABLE) .withEntity(Scientist.class) .withCoder(SerializableCoder.of(Scientist.class)); String splitQuery = QueryBuilder.select().from(CASSANDRA_KEYSPACE, CASSANDRA_TABLE).toString(); CassandraIO.CassandraSource<Scientist> initialSource = new CassandraIO.CassandraSource<>(read, Collections.singletonList(splitQuery)); int desiredBundleSizeBytes = 2048; long estimatedSize = initialSource.getEstimatedSizeBytes(options); List<BoundedSource<Scientist>> splits = initialSource.split(desiredBundleSizeBytes, options); SourceTestUtils.assertSourcesEqualReferenceSource(initialSource, splits, options); float expectedNumSplitsloat = (float) initialSource.getEstimatedSizeBytes(options) / desiredBundleSizeBytes; long sum = 0; for (BoundedSource<Scientist> subSource : splits) { sum += subSource.getEstimatedSizeBytes(options); } assertEquals((long) (estimatedSize / splits.size()) * splits.size(), sum); int expectedNumSplits = (int) Math.ceil(expectedNumSplitsloat); assertEquals("Wrong number of splits", expectedNumSplits, splits.size()); int emptySplits = 0; for (BoundedSource<Scientist> subSource : splits) { if (readFromSource(subSource, options).isEmpty()) { emptySplits += 1; } } assertThat( "There are too many empty splits, parallelism is sub-optimal", emptySplits, lessThan((int) (ACCEPTABLE_EMPTY_SPLITS_PERCENTAGE * splits.size()))); }
@Test public void testEstimatedSizeBytes() throws Exception { PipelineOptions pipelineOptions = PipelineOptionsFactory.create(); CassandraIO.Read<Scientist> read = CassandraIO.<Scientist>read() .withHosts(Collections.singletonList(CASSANDRA_HOST)) .withPort(cassandraPort) .withKeyspace(CASSANDRA_KEYSPACE) .withTable(CASSANDRA_TABLE); CassandraIO.CassandraSource<Scientist> source = new CassandraIO.CassandraSource<>(read, null); long estimatedSizeBytes = source.getEstimatedSizeBytes(pipelineOptions); assertTrue((estimatedSizeBytes >= 12960L * 0.9f) && (estimatedSizeBytes <= 12960L * 1.1f)); }