use of io.druid.java.util.common.guava.nary.BinaryFn in project druid by druid-io.
the class CombiningSequenceTest method testCombining.
private void testCombining(List<Pair<Integer, Integer>> pairs, List<Pair<Integer, Integer>> expected, int limit) throws Exception {
// Test that closing works too
final CountDownLatch closed = new CountDownLatch(1);
final Closeable closeable = new Closeable() {
@Override
public void close() throws IOException {
closed.countDown();
}
};
Sequence<Pair<Integer, Integer>> seq = Sequences.limit(CombiningSequence.create(Sequences.withBaggage(Sequences.simple(pairs), closeable), Ordering.natural().onResultOf(Pair.<Integer, Integer>lhsFn()), new BinaryFn<Pair<Integer, Integer>, Pair<Integer, Integer>, Pair<Integer, Integer>>() {
@Override
public Pair<Integer, Integer> apply(Pair<Integer, Integer> lhs, Pair<Integer, Integer> rhs) {
if (lhs == null) {
return rhs;
}
if (rhs == null) {
return lhs;
}
return Pair.of(lhs.lhs, lhs.rhs + rhs.rhs);
}
}), limit);
List<Pair<Integer, Integer>> merged = Sequences.toList(seq, Lists.<Pair<Integer, Integer>>newArrayList());
Assert.assertEquals(expected, merged);
Yielder<Pair<Integer, Integer>> yielder = seq.toYielder(null, new YieldingAccumulator<Pair<Integer, Integer>, Pair<Integer, Integer>>() {
int count = 0;
@Override
public Pair<Integer, Integer> accumulate(Pair<Integer, Integer> lhs, Pair<Integer, Integer> rhs) {
count++;
if (count % yieldEvery == 0) {
yield();
}
return rhs;
}
});
Iterator<Pair<Integer, Integer>> expectedVals = Iterators.filter(expected.iterator(), new Predicate<Pair<Integer, Integer>>() {
int count = 0;
@Override
public boolean apply(@Nullable Pair<Integer, Integer> input) {
count++;
if (count % yieldEvery == 0) {
return true;
}
return false;
}
});
if (expectedVals.hasNext()) {
while (!yielder.isDone()) {
final Pair<Integer, Integer> expectedVal = expectedVals.next();
final Pair<Integer, Integer> actual = yielder.get();
Assert.assertEquals(expectedVal, actual);
yielder = yielder.next(actual);
}
}
Assert.assertTrue(yielder.isDone());
Assert.assertFalse(expectedVals.hasNext());
yielder.close();
Assert.assertTrue("resource closed", closed.await(10000, TimeUnit.MILLISECONDS));
}
Aggregations