BatchingShuffleEntryReader implements ShuffleEntryReader { @Override public Reiterator<ShuffleEntry> read( @Nullable ShufflePosition startPosition, @Nullable ShufflePosition endPosition) { return new ShuffleReadIterator(startPosition, endPosition); } BatchingShuffleEntryReader(ShuffleBatchReader batchReader); @Override Reiterator<ShuffleEntry> read(
      @Nullable ShufflePosition startPosition, @Nullable ShufflePosition endPosition); @Override void close();  }
@Test public void readerCanRead() throws Exception { ShuffleEntry e1 = new ShuffleEntry(KEY, SKEY, VALUE); ShuffleEntry e2 = new ShuffleEntry(KEY, SKEY, VALUE); ArrayList<ShuffleEntry> entries = new ArrayList<>(); entries.add(e1); entries.add(e2); when(batchReader.read(START_POSITION, END_POSITION)) .thenReturn(new ShuffleBatchReader.Batch(entries, null)); List<ShuffleEntry> results = newArrayList(reader.read(START_POSITION, END_POSITION)); assertThat(results, contains(e1, e2)); }
@Test public void readerIteratorCanBeCopied() throws Exception { ShuffleEntry e1 = new ShuffleEntry(KEY, SKEY, VALUE); ShuffleEntry e2 = new ShuffleEntry(KEY, SKEY, VALUE); ArrayList<ShuffleEntry> entries = new ArrayList<>(); entries.add(e1); entries.add(e2); when(batchReader.read(START_POSITION, END_POSITION)) .thenReturn(new ShuffleBatchReader.Batch(entries, null)); Reiterator<ShuffleEntry> it = reader.read(START_POSITION, END_POSITION); assertThat(it.hasNext(), equalTo(Boolean.TRUE)); assertThat(it.next(), equalTo(e1)); Reiterator<ShuffleEntry> copy = it.copy(); assertThat(it.hasNext(), equalTo(Boolean.TRUE)); assertThat(it.next(), equalTo(e2)); assertThat(it.hasNext(), equalTo(Boolean.FALSE)); assertThat(copy.hasNext(), equalTo(Boolean.TRUE)); assertThat(copy.next(), equalTo(e2)); assertThat(copy.hasNext(), equalTo(Boolean.FALSE)); }
@Test public void readerShouldMergeMultipleBatchResults() throws Exception { ShuffleEntry e1 = new ShuffleEntry(KEY, SKEY, VALUE); List<ShuffleEntry> e1s = Collections.singletonList(e1); ShuffleEntry e2 = new ShuffleEntry(KEY, SKEY, VALUE); List<ShuffleEntry> e2s = Collections.singletonList(e2); when(batchReader.read(START_POSITION, END_POSITION)) .thenReturn(new ShuffleBatchReader.Batch(e1s, NEXT_START_POSITION)); when(batchReader.read(NEXT_START_POSITION, END_POSITION)) .thenReturn(new ShuffleBatchReader.Batch(e2s, null)); List<ShuffleEntry> results = newArrayList(reader.read(START_POSITION, END_POSITION)); assertThat(results, contains(e1, e2)); verify(batchReader).read(START_POSITION, END_POSITION); verify(batchReader).read(NEXT_START_POSITION, END_POSITION); verifyNoMoreInteractions(batchReader); }
@Test public void readerShouldMergeMultipleBatchResultsIncludingEmptyShards() throws Exception { List<ShuffleEntry> e1s = new ArrayList<>(); List<ShuffleEntry> e2s = new ArrayList<>(); ShuffleEntry e3 = new ShuffleEntry(KEY, SKEY, VALUE); List<ShuffleEntry> e3s = Collections.singletonList(e3); when(batchReader.read(START_POSITION, END_POSITION)) .thenReturn(new ShuffleBatchReader.Batch(e1s, NEXT_START_POSITION)); when(batchReader.read(NEXT_START_POSITION, END_POSITION)) .thenReturn(new ShuffleBatchReader.Batch(e2s, SECOND_NEXT_START_POSITION)); when(batchReader.read(SECOND_NEXT_START_POSITION, END_POSITION)) .thenReturn(new ShuffleBatchReader.Batch(e3s, null)); List<ShuffleEntry> results = newArrayList(reader.read(START_POSITION, END_POSITION)); assertThat(results, contains(e3)); verify(batchReader).read(START_POSITION, END_POSITION); verify(batchReader).read(NEXT_START_POSITION, END_POSITION); verify(batchReader).read(SECOND_NEXT_START_POSITION, END_POSITION); verifyNoMoreInteractions(batchReader); }
@Test public void readerCanRead() throws Exception { ShuffleEntry e1 = new ShuffleEntry(KEY, SKEY, VALUE); ShuffleEntry e2 = new ShuffleEntry(KEY, SKEY, VALUE); ArrayList<ShuffleEntry> entries = new ArrayList<>(); entries.add(e1); entries.add(e2); long batchSize = (long) e1.length() + e2.length(); when(batchReader.read(START_POSITION, END_POSITION)) .thenReturn(new ShuffleBatchReader.Batch(entries, null, batchSize)); List<ShuffleEntry> results = newArrayList(reader.read(START_POSITION, END_POSITION)); assertThat(results, contains(e1, e2)); }
@Test public void readerIteratorCanBeCopied() throws Exception { ShuffleEntry e1 = new ShuffleEntry(KEY, SKEY, VALUE); ShuffleEntry e2 = new ShuffleEntry(KEY, SKEY, VALUE); ArrayList<ShuffleEntry> entries = new ArrayList<>(); entries.add(e1); entries.add(e2); long batchSize = (long) e1.length() + e2.length(); when(batchReader.read(START_POSITION, END_POSITION)) .thenReturn(new ShuffleBatchReader.Batch(entries, null, batchSize)); Reiterator<ShuffleEntry> it = reader.read(START_POSITION, END_POSITION); assertThat(it.hasNext(), equalTo(Boolean.TRUE)); assertThat(it.next(), equalTo(e1)); Reiterator<ShuffleEntry> copy = it.copy(); assertThat(it.hasNext(), equalTo(Boolean.TRUE)); assertThat(it.next(), equalTo(e2)); assertThat(it.hasNext(), equalTo(Boolean.FALSE)); assertThat(copy.hasNext(), equalTo(Boolean.TRUE)); assertThat(copy.next(), equalTo(e2)); assertThat(copy.hasNext(), equalTo(Boolean.FALSE)); }
@Test public void readerShouldMergeMultipleBatchResults() throws Exception { ShuffleEntry e1 = new ShuffleEntry(KEY, SKEY, VALUE); List<ShuffleEntry> e1s = Collections.singletonList(e1); ShuffleEntry e2 = new ShuffleEntry(KEY, SKEY, VALUE); List<ShuffleEntry> e2s = Collections.singletonList(e2); when(batchReader.read(START_POSITION, END_POSITION)) .thenReturn(new ShuffleBatchReader.Batch(e1s, NEXT_START_POSITION, e1.length())); when(batchReader.read(NEXT_START_POSITION, END_POSITION)) .thenReturn(new ShuffleBatchReader.Batch(e2s, null, e2.length())); List<ShuffleEntry> results = newArrayList(reader.read(START_POSITION, END_POSITION)); assertThat(results, contains(e1, e2)); verify(batchReader).read(START_POSITION, END_POSITION); verify(batchReader).read(NEXT_START_POSITION, END_POSITION); verifyNoMoreInteractions(batchReader); }
@Test public void readerShouldMergeMultipleBatchResultsIncludingEmptyShards() throws Exception { List<ShuffleEntry> e1s = new ArrayList<>(); List<ShuffleEntry> e2s = new ArrayList<>(); ShuffleEntry e3 = new ShuffleEntry(KEY, SKEY, VALUE); List<ShuffleEntry> e3s = Collections.singletonList(e3); when(batchReader.read(START_POSITION, END_POSITION)) .thenReturn(new ShuffleBatchReader.Batch(e1s, NEXT_START_POSITION, 0)); when(batchReader.read(NEXT_START_POSITION, END_POSITION)) .thenReturn(new ShuffleBatchReader.Batch(e2s, SECOND_NEXT_START_POSITION, 0)); when(batchReader.read(SECOND_NEXT_START_POSITION, END_POSITION)) .thenReturn(new ShuffleBatchReader.Batch(e3s, null, e3.length())); List<ShuffleEntry> results = newArrayList(reader.read(START_POSITION, END_POSITION)); assertThat(results, contains(e3)); verify(batchReader).read(START_POSITION, END_POSITION); verify(batchReader).read(NEXT_START_POSITION, END_POSITION); verify(batchReader).read(SECOND_NEXT_START_POSITION, END_POSITION); verifyNoMoreInteractions(batchReader); }