Search in sources :

Example 1 with BinaryFn

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));
}
Also used : Closeable(java.io.Closeable) CountDownLatch(java.util.concurrent.CountDownLatch) BinaryFn(io.druid.java.util.common.guava.nary.BinaryFn) Pair(io.druid.java.util.common.Pair)

Aggregations

Pair (io.druid.java.util.common.Pair)1 BinaryFn (io.druid.java.util.common.guava.nary.BinaryFn)1 Closeable (java.io.Closeable)1 CountDownLatch (java.util.concurrent.CountDownLatch)1