| /* |
| * Licensed to the Apache Software Foundation (ASF) under one |
| * or more contributor license agreements. See the NOTICE file |
| * distributed with this work for additional information |
| * regarding copyright ownership. The ASF licenses this file |
| * to you under the Apache License, Version 2.0 (the |
| * "License"); you may not use this file except in compliance |
| * with the License. You may obtain a copy of the License at |
| * |
| * http://www.apache.org/licenses/LICENSE-2.0 |
| * |
| * Unless required by applicable law or agreed to in writing, software |
| * distributed under the License is distributed on an "AS IS" BASIS, |
| * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| * See the License for the specific language governing permissions and |
| * limitations under the License. |
| */ |
| package org.apache.beam.runners.dataflow.worker; |
| |
| import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState; |
| import static org.junit.Assert.assertArrayEquals; |
| import static org.junit.Assert.assertEquals; |
| import static org.junit.Assert.assertFalse; |
| import static org.junit.Assert.assertNotNull; |
| import static org.junit.Assert.assertNull; |
| import static org.junit.Assert.assertTrue; |
| import static org.junit.Assert.fail; |
| |
| import java.io.Closeable; |
| import java.io.File; |
| import java.io.IOException; |
| import java.nio.ByteBuffer; |
| import java.nio.channels.NonWritableChannelException; |
| import java.nio.channels.SeekableByteChannel; |
| import java.util.ArrayList; |
| import java.util.Arrays; |
| import java.util.Collections; |
| import java.util.Comparator; |
| import java.util.HashMap; |
| import java.util.Iterator; |
| import java.util.List; |
| import java.util.Map; |
| import java.util.NavigableMap; |
| import java.util.NavigableSet; |
| import java.util.NoSuchElementException; |
| import java.util.Random; |
| import java.util.SortedMap; |
| import java.util.SortedSet; |
| import java.util.TreeMap; |
| import java.util.TreeSet; |
| import org.apache.beam.runners.dataflow.internal.IsmFormat; |
| import org.apache.beam.runners.dataflow.internal.IsmFormat.IsmRecord; |
| import org.apache.beam.runners.dataflow.internal.IsmFormat.IsmRecordCoder; |
| import org.apache.beam.runners.dataflow.internal.IsmFormat.MetadataKeyCoder; |
| import org.apache.beam.runners.dataflow.options.DataflowPipelineOptions; |
| import org.apache.beam.runners.dataflow.util.RandomAccessData; |
| import org.apache.beam.runners.dataflow.worker.DataflowOperationContext.DataflowExecutionState; |
| import org.apache.beam.runners.dataflow.worker.IsmReaderImpl.CachedTailSeekableByteChannel; |
| import org.apache.beam.runners.dataflow.worker.IsmReaderImpl.IsmShardKey; |
| import org.apache.beam.runners.dataflow.worker.profiler.ScopedProfiler.NoopProfileScope; |
| import org.apache.beam.runners.dataflow.worker.util.ValueInEmptyWindows; |
| import org.apache.beam.runners.dataflow.worker.util.common.worker.ExecutorTestUtils.TestReaderObserver; |
| import org.apache.beam.runners.dataflow.worker.util.common.worker.NativeReader; |
| import org.apache.beam.runners.dataflow.worker.util.common.worker.Sink.SinkWriter; |
| import org.apache.beam.sdk.coders.ByteArrayCoder; |
| import org.apache.beam.sdk.coders.Coder; |
| import org.apache.beam.sdk.coders.VoidCoder; |
| import org.apache.beam.sdk.io.FileSystems; |
| import org.apache.beam.sdk.options.PipelineOptionsFactory; |
| import org.apache.beam.sdk.testing.CoderPropertiesTest.NonDeterministicCoder; |
| import org.apache.beam.sdk.util.CoderUtils; |
| import org.apache.beam.sdk.util.WeightedValue; |
| import org.apache.beam.sdk.util.WindowedValue; |
| import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Predicate; |
| import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Predicates; |
| import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.cache.Cache; |
| import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.cache.CacheBuilder; |
| import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.FluentIterable; |
| import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList; |
| import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterables; |
| import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.primitives.Ints; |
| import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.primitives.UnsignedBytes; |
| import org.junit.After; |
| import org.junit.Before; |
| import org.junit.Rule; |
| import org.junit.Test; |
| import org.junit.rules.ExpectedException; |
| import org.junit.rules.TemporaryFolder; |
| import org.junit.runner.RunWith; |
| import org.junit.runners.JUnit4; |
| |
| /** Tests for {@link IsmReader}. */ |
| @RunWith(JUnit4.class) |
| public class IsmReaderTest { |
| private static final long BLOOM_FILTER_SIZE_LIMIT = 10_000; |
| private static final int TEST_BLOCK_SIZE = 1024; |
| private static final IsmRecordCoder<byte[]> CODER = |
| IsmRecordCoder.of( |
| 1, // number or shard key coders for value records |
| 1, // number of shard key coders for metadata records |
| ImmutableList.<Coder<?>>of(MetadataKeyCoder.of(ByteArrayCoder.of()), ByteArrayCoder.of()), |
| ByteArrayCoder.of()); |
| |
| private static final Coder<String> NON_DETERMINISTIC_CODER = new NonDeterministicCoder(); |
| private static final byte[] EMPTY = new byte[0]; |
| |
| @Rule public TemporaryFolder tmpFolder = new TemporaryFolder(); |
| @Rule public ExpectedException expectedException = ExpectedException.none(); |
| |
| private Cache< |
| IsmReaderImpl.IsmShardKey, |
| WeightedValue<NavigableMap<RandomAccessData, WindowedValue<IsmRecord<byte[]>>>>> |
| cache; |
| private DataflowExecutionContext executionContext; |
| private DataflowOperationContext operationContext; |
| private SideInputReadCounter sideInputReadCounter; |
| private Closeable stateCloseable; |
| |
| @Before |
| public void setUp() { |
| cache = |
| CacheBuilder.newBuilder() |
| .weigher(Weighers.fixedWeightKeys(1)) |
| .maximumWeight(10_000) |
| .build(); |
| |
| executionContext = |
| BatchModeExecutionContext.forTesting( |
| PipelineOptionsFactory.as(DataflowPipelineOptions.class), "testStage"); |
| |
| DataflowExecutionState state = |
| executionContext |
| .getExecutionStateRegistry() |
| .getState( |
| NameContextsForTests.nameContextForTest(), "test", null, NoopProfileScope.NOOP); |
| operationContext = |
| executionContext.createOperationContext(NameContextsForTests.nameContextForTest()); |
| stateCloseable = executionContext.getExecutionStateTracker().enterState(state); |
| sideInputReadCounter = new DataflowSideInputReadCounter(executionContext, operationContext, 1); |
| } |
| |
| @After |
| public void tearDown() throws IOException { |
| stateCloseable.close(); |
| } |
| |
| @Test |
| public void testReadEmpty() throws Exception { |
| writeElementsToFileAndReadInOrder(Collections.<IsmRecord<byte[]>>emptyList()); |
| } |
| |
| @Test |
| public void testUsingNonDeterministicShardKeyCoder() throws Exception { |
| expectedException.expect(IllegalArgumentException.class); |
| expectedException.expectMessage("is expected to be deterministic"); |
| new IsmReaderImpl<>( |
| FileSystems.matchSingleFileSpec(tmpFolder.newFile().getPath()).resourceId(), |
| IsmRecordCoder.of( |
| 1, // number or shard key coders for value records |
| 0, // number of shard key coders for metadata records |
| ImmutableList.<Coder<?>>of(NON_DETERMINISTIC_CODER, ByteArrayCoder.of()), |
| ByteArrayCoder.of()), |
| cache); |
| } |
| |
| @Test |
| public void testUsingNonDeterministicNonShardKeyCoder() throws Exception { |
| expectedException.expect(IllegalArgumentException.class); |
| expectedException.expectMessage("is expected to be deterministic"); |
| new IsmReaderImpl<>( |
| FileSystems.matchSingleFileSpec(tmpFolder.newFile().getPath()).resourceId(), |
| IsmRecordCoder.of( |
| 1, // number or shard key coders for value records |
| 0, // number of shard key coders for metadata records |
| ImmutableList.<Coder<?>>of(ByteArrayCoder.of(), NON_DETERMINISTIC_CODER), |
| ByteArrayCoder.of()), |
| cache); |
| } |
| |
| @Test |
| public void testIsEmpty() throws Exception { |
| File tmpFile = tmpFolder.newFile(); |
| List<IsmRecord<byte[]>> data = new ArrayList<>(); |
| writeElementsToFile(data, tmpFile); |
| |
| IsmReader<byte[]> reader = |
| new IsmReaderImpl<byte[]>( |
| FileSystems.matchSingleFileSpec(tmpFile.getAbsolutePath()).resourceId(), CODER, cache); |
| assertFalse(reader.isInitialized()); |
| assertTrue(reader.isEmpty()); |
| assertTrue(reader.isInitialized()); |
| } |
| |
| @Test |
| public void testRead() throws Exception { |
| Random random = new Random(23498321490L); |
| for (int i : Arrays.asList(1, 2, 3, 4, 5, 6, 7, 8, 9, 10)) { |
| int minElements = (int) Math.pow(2, i); |
| int valueSize = 128; |
| // Generates between 2^i and 2^(i + 1) elements. |
| writeElementsToFileAndReadInOrder( |
| dataGenerator( |
| 8 /* number of primary keys */, |
| minElements + random.nextInt(minElements) /* number of secondary keys */, |
| 8 /* max key size */, |
| valueSize)); |
| } |
| } |
| |
| @Test |
| public void testReadThatProducesIndexEntries() throws Exception { |
| Random random = new Random(23498323891L); |
| int minElements = (int) Math.pow(2, 6); |
| int valueSize = 128; |
| // Since we are generating more than 2 blocks worth of data, we are guaranteed that |
| // at least one index entry is generated per shard. |
| checkState(minElements * valueSize > 2 * TEST_BLOCK_SIZE); |
| writeElementsToFileAndReadInOrder( |
| dataGenerator( |
| 8 /* number of primary keys */, |
| minElements + random.nextInt(minElements) /* number of secondary keys */, |
| 8 /* max key size */, |
| valueSize /* max value size */)); |
| } |
| |
| @Test |
| public void testReadRandomOrder() throws Exception { |
| Random random = new Random(2348238943L); |
| for (int i : Arrays.asList(1, 2, 3, 4, 5, 6, 7, 8)) { |
| int minElements = (int) Math.pow(2, i); |
| int valueSize = 128; |
| // Generates between 2^i and 2^(i + 1) elements. |
| writeElementsToFileAndReadInRandomOrder( |
| dataGenerator( |
| 7 /* number of primary keys */, |
| minElements + random.nextInt(minElements) /* number of secondary keys */, |
| 8 /* max key size */, |
| valueSize /* max value size */)); |
| } |
| } |
| |
| @Test |
| public void testGetLastWithPrefix() throws Exception { |
| Random random = new Random(2348238943L); |
| for (int i : Arrays.asList(1, 2, 3, 4, 5, 6, 7, 8, 9, 10)) { |
| int minElements = (int) Math.pow(2, i); |
| int valueSize = 128; |
| // Generates between 2^i and 2^(i + 1) elements. |
| writeElementsToFileAndFindLastElementPerPrimaryKey( |
| dataGenerator( |
| 7, |
| minElements + random.nextInt(minElements), |
| 8 /* max key size */, |
| valueSize /* max value size */)); |
| } |
| } |
| |
| @Test |
| public void testReadMissingKeys() throws Exception { |
| File tmpFile = tmpFolder.newFile(); |
| List<IsmRecord<byte[]>> data = new ArrayList<>(); |
| data.add(IsmRecord.<byte[]>of(ImmutableList.of(EMPTY, new byte[] {0x04}), EMPTY)); |
| data.add(IsmRecord.<byte[]>of(ImmutableList.of(EMPTY, new byte[] {0x08}), EMPTY)); |
| writeElementsToFile(data, tmpFile); |
| |
| IsmReader<byte[]> reader = |
| new IsmReaderImpl<byte[]>( |
| FileSystems.matchSingleFileSpec(tmpFile.getAbsolutePath()).resourceId(), CODER, cache); |
| |
| // Check that we got false with a key before all keys contained in the file. |
| assertFalse(reader.overKeyComponents(ImmutableList.of(EMPTY, new byte[] {0x02})).start()); |
| // Check that we got false with a key between two other keys contained in the file. |
| assertFalse(reader.overKeyComponents(ImmutableList.of(EMPTY, new byte[] {0x06})).start()); |
| // Check that we got false with a key that is after all keys contained in the file. |
| assertFalse(reader.overKeyComponents(ImmutableList.of(EMPTY, new byte[] {0x10})).start()); |
| } |
| |
| @Test |
| public void testReadMissingKeysBypassingBloomFilter() throws Exception { |
| File tmpFile = tmpFolder.newFile(); |
| List<IsmRecord<byte[]>> data = new ArrayList<>(); |
| data.add(IsmRecord.<byte[]>of(ImmutableList.of(EMPTY, new byte[] {0x04}), EMPTY)); |
| data.add(IsmRecord.<byte[]>of(ImmutableList.of(EMPTY, new byte[] {0x08}), EMPTY)); |
| writeElementsToFile(data, tmpFile); |
| |
| IsmReader<byte[]> reader = |
| new IsmReaderImpl<byte[]>( |
| FileSystems.matchSingleFileSpec(tmpFile.getAbsolutePath()).resourceId(), CODER, cache) { |
| // We use this override to get around the Bloom filter saying that the key doesn't exist. |
| @Override |
| boolean bloomFilterMightContain(RandomAccessData keyBytes) { |
| return true; |
| } |
| }; |
| |
| // Check that we got false with a key before all keys contained in the file. |
| assertFalse(reader.overKeyComponents(ImmutableList.of(EMPTY, new byte[] {0x02})).start()); |
| // Check that we got false with a key between two other keys contained in the file. |
| assertFalse(reader.overKeyComponents(ImmutableList.of(EMPTY, new byte[] {0x06})).start()); |
| // Check that we got false with a key that is after all keys contained in the file. |
| assertFalse(reader.overKeyComponents(ImmutableList.of(EMPTY, new byte[] {0x10})).start()); |
| } |
| |
| @Test |
| public void testReadKeyThatEncodesToEmptyByteArray() throws Exception { |
| File tmpFile = tmpFolder.newFile(); |
| IsmRecordCoder<Void> coder = |
| IsmRecordCoder.of(1, 0, ImmutableList.<Coder<?>>of(VoidCoder.of()), VoidCoder.of()); |
| IsmSink<Void> sink = |
| new IsmSink<>( |
| FileSystems.matchNewResource(tmpFile.getPath(), false), coder, BLOOM_FILTER_SIZE_LIMIT); |
| IsmRecord<Void> element = IsmRecord.of(Arrays.asList((Void) null), (Void) null); |
| try (SinkWriter<WindowedValue<IsmRecord<Void>>> writer = sink.writer()) { |
| writer.add(new ValueInEmptyWindows<>(element)); |
| } |
| |
| Cache< |
| IsmShardKey, |
| WeightedValue<NavigableMap<RandomAccessData, WindowedValue<IsmRecord<Void>>>>> |
| cache = |
| CacheBuilder.newBuilder() |
| .weigher(Weighers.fixedWeightKeys(1)) |
| .maximumWeight(10_000) |
| .build(); |
| IsmReader<Void> reader = |
| new IsmReaderImpl<>( |
| FileSystems.matchSingleFileSpec(tmpFile.getAbsolutePath()).resourceId(), coder, cache); |
| IsmReader<Void>.IsmPrefixReaderIterator iterator = reader.iterator(); |
| assertTrue(iterator.start()); |
| assertEquals( |
| coder.structuralValue(element), coder.structuralValue(iterator.getCurrent().getValue())); |
| } |
| |
| @Test |
| public void testInitializationForSmallFilesIsCached() throws Exception { |
| File tmpFile = tmpFolder.newFile(); |
| IsmShardKey expectedShardKey = |
| new IsmShardKey(tmpFile.getAbsolutePath(), new RandomAccessData(0), 0, 13); |
| List<IsmRecord<byte[]>> data = new ArrayList<>(); |
| data.add(IsmRecord.<byte[]>of(ImmutableList.of(EMPTY, new byte[] {0x04}), new byte[] {0x04})); |
| data.add(IsmRecord.<byte[]>of(ImmutableList.of(EMPTY, new byte[] {0x08}), new byte[] {0x08})); |
| writeElementsToFile(data, tmpFile); |
| |
| IsmReader<byte[]> reader = |
| new IsmReaderImpl<byte[]>( |
| FileSystems.matchSingleFileSpec(tmpFile.getAbsolutePath()).resourceId(), CODER, cache); |
| // Validate that reader and cache are in initial state |
| assertFalse(reader.isInitialized()); |
| assertEquals(0, cache.size()); |
| // Force initialization |
| reader.overKeyComponents(ImmutableList.of()); |
| // Validate reader is initialized and expected entry is cached |
| assertTrue(reader.isInitialized()); |
| WeightedValue<NavigableMap<RandomAccessData, WindowedValue<IsmRecord<byte[]>>>> block = |
| cache.getIfPresent(expectedShardKey); |
| assertNotNull(block); |
| assertArrayEquals( |
| new byte[] {0x04}, block.getValue().firstEntry().getValue().getValue().getValue()); |
| assertArrayEquals( |
| new byte[] {0x08}, block.getValue().lastEntry().getValue().getValue().getValue()); |
| } |
| |
| @Test |
| public void testInitializationForLargeFilesIsNotCached() throws Exception { |
| File tmpFile = tmpFolder.newFile(); |
| List<IsmRecord<byte[]>> data = new ArrayList<>(); |
| // Use enough data records which are smaller than the cache limit to exceed the 1 MB |
| // footer read buffer used by IsmReader to optimize small file handling. |
| for (int i = 0; i < IsmReaderImpl.MAX_SHARD_INDEX_AND_FOOTER_SIZE / TEST_BLOCK_SIZE + 1; ++i) { |
| data.add( |
| IsmRecord.<byte[]>of( |
| ImmutableList.of(EMPTY, Ints.toByteArray(i)), new byte[TEST_BLOCK_SIZE])); |
| } |
| writeElementsToFile(data, tmpFile); |
| |
| IsmReader<byte[]> reader = |
| new IsmReaderImpl<byte[]>( |
| FileSystems.matchSingleFileSpec(tmpFile.getAbsolutePath()).resourceId(), CODER, cache); |
| assertFalse(reader.isInitialized()); |
| assertEquals(0, cache.size()); |
| reader.overKeyComponents(ImmutableList.of()); |
| assertTrue(reader.isInitialized()); |
| assertEquals(0, cache.size()); |
| } |
| |
| /** Write input elements to the specified file. */ |
| static void writeElementsToFile(Iterable<IsmRecord<byte[]>> elements, File tmpFile) |
| throws Exception { |
| IsmSink<byte[]> sink = |
| new IsmSink<byte[]>( |
| FileSystems.matchNewResource(tmpFile.getPath(), false), |
| CODER, |
| BLOOM_FILTER_SIZE_LIMIT) { |
| @Override |
| long getBlockSize() { |
| return TEST_BLOCK_SIZE; |
| } |
| }; |
| |
| try (SinkWriter<WindowedValue<IsmRecord<byte[]>>> writer = sink.writer()) { |
| for (IsmRecord<byte[]> element : elements) { |
| writer.add(new ValueInEmptyWindows<>(element)); |
| } |
| } |
| } |
| |
| /** |
| * Writes elements to an Ism file using an IsmSink. Then reads them back with an IsmReader, |
| * verifying the values read match those that were written. |
| */ |
| private void writeElementsToFileAndReadInOrder(Iterable<IsmRecord<byte[]>> elements) |
| throws Exception { |
| File tmpFile = tmpFolder.newFile(); |
| writeElementsToFile(elements, tmpFile); |
| IsmReader<byte[]> reader = |
| new IsmReaderImpl<>( |
| FileSystems.matchSingleFileSpec(tmpFile.getAbsolutePath()).resourceId(), CODER, cache); |
| assertFalse(reader.isInitialized()); |
| |
| TestReaderObserver observer = new TestReaderObserver(reader); |
| reader.addObserver(observer); |
| |
| Iterator<IsmRecord<byte[]>> elementsIterator = elements.iterator(); |
| try (NativeReader.NativeReaderIterator<WindowedValue<IsmRecord<byte[]>>> iterator = |
| reader.iterator()) { |
| |
| boolean more = iterator.start(); |
| assertTrue(reader.isInitialized()); |
| for (; more; more = iterator.advance()) { |
| |
| if (!elementsIterator.hasNext()) { |
| break; |
| } |
| IsmRecord<byte[]> expected = elementsIterator.next(); |
| IsmRecord<byte[]> actual = iterator.getCurrent().getValue(); |
| assertIsmEquals(actual, expected); |
| |
| final int expectedLength; |
| if (IsmFormat.isMetadataKey(expected.getKeyComponents())) { |
| expectedLength = expected.getMetadata().length; |
| } else { |
| expectedLength = expected.getValue().length; |
| } |
| // Verify that the observer saw at least as many bytes as the size of the value. |
| assertTrue( |
| expectedLength <= observer.getActualSizes().get(observer.getActualSizes().size() - 1)); |
| } |
| if (iterator.advance()) { |
| fail("Read more elements then expected, did not expect: " + iterator.getCurrent()); |
| } else if (elementsIterator.hasNext()) { |
| fail("Read less elements then expected, expected: " + elementsIterator.next()); |
| } |
| |
| // Verify that we see a {@link NoSuchElementException} if we attempt to go further. |
| try { |
| iterator.getCurrent(); |
| fail("Expected a NoSuchElementException to have been thrown."); |
| } catch (NoSuchElementException expected) { |
| } |
| } |
| } |
| |
| private static void assertIsmEquals(IsmRecord<byte[]> actual, IsmRecord<byte[]> expected) { |
| assertEquals(expected.getKeyComponents().size(), actual.getKeyComponents().size()); |
| for (int i = 0; i < expected.getKeyComponents().size(); ++i) { |
| if (actual.getKeyComponent(i) != expected.getKeyComponent(i)) { |
| assertArrayEquals((byte[]) actual.getKeyComponent(i), (byte[]) expected.getKeyComponent(i)); |
| } |
| } |
| |
| if (IsmFormat.isMetadataKey(expected.getKeyComponents())) { |
| assertArrayEquals(actual.getMetadata(), expected.getMetadata()); |
| } else { |
| assertArrayEquals(actual.getValue(), expected.getValue()); |
| } |
| } |
| |
| /** |
| * A predicate which filters elements on whether the second key's last byte is odd or even. Allows |
| * for a stable partitioning of generated data. |
| */ |
| private static class EvenFilter implements Predicate<IsmRecord<byte[]>> { |
| private static final EvenFilter INSTANCE = new EvenFilter(); |
| |
| @Override |
| public boolean apply(IsmRecord<byte[]> input) { |
| byte[] secondKey = (byte[]) input.getKeyComponent(1); |
| return secondKey[secondKey.length - 1] % 2 == 0; |
| } |
| } |
| |
| /** |
| * Writes elements to an Ism file using an IsmSink. Then reads them back with an IsmReader using a |
| * random order. |
| */ |
| private void writeElementsToFileAndReadInRandomOrder(Iterable<IsmRecord<byte[]>> elements) |
| throws Exception { |
| File tmpFile = tmpFolder.newFile(); |
| List<IsmRecord<byte[]>> oddSecondaryKeys = |
| new ArrayList<>( |
| ImmutableList.copyOf(Iterables.filter(elements, Predicates.not(EvenFilter.INSTANCE)))); |
| List<IsmRecord<byte[]>> evenSecondaryKeys = |
| new ArrayList<>(ImmutableList.copyOf(Iterables.filter(elements, EvenFilter.INSTANCE))); |
| |
| writeElementsToFile(oddSecondaryKeys, tmpFile); |
| IsmReader<byte[]> reader = |
| new IsmReaderImpl<>( |
| FileSystems.matchSingleFileSpec(tmpFile.getAbsolutePath()).resourceId(), CODER, cache); |
| |
| // Test using next() for a within shard Ism prefix reader iterator |
| Collections.shuffle(oddSecondaryKeys); |
| for (IsmRecord<byte[]> expectedNext : oddSecondaryKeys) { |
| IsmReader<byte[]>.IsmPrefixReaderIterator iterator = |
| reader.overKeyComponents(expectedNext.getKeyComponents()); |
| assertTrue(iterator.start()); |
| assertIsmEquals(iterator.getCurrent().getValue(), expectedNext); |
| } |
| |
| Collections.shuffle(oddSecondaryKeys); |
| // Test using get() for a shard aware Ism prefix reader |
| IsmReader<byte[]>.IsmPrefixReaderIterator readerIterator = |
| reader.overKeyComponents(ImmutableList.of()); |
| for (IsmRecord<byte[]> expectedNext : oddSecondaryKeys) { |
| assertIsmEquals(readerIterator.get(expectedNext.getKeyComponents()).getValue(), expectedNext); |
| } |
| |
| // Test using next() for a within shard Ism prefix reader iterator |
| Collections.shuffle(evenSecondaryKeys); |
| for (IsmRecord<byte[]> missingNext : evenSecondaryKeys) { |
| assertFalse(reader.overKeyComponents(missingNext.getKeyComponents()).start()); |
| } |
| |
| Collections.shuffle(evenSecondaryKeys); |
| // Test using get() for a shard aware Ism prefix reader |
| readerIterator = reader.overKeyComponents(ImmutableList.of()); |
| for (IsmRecord<byte[]> missingNext : evenSecondaryKeys) { |
| assertNull(readerIterator.get(missingNext.getKeyComponents())); |
| } |
| } |
| |
| private void writeElementsToFileAndFindLastElementPerPrimaryKey( |
| Iterable<IsmRecord<byte[]>> elements) throws Exception { |
| File tmpFile = tmpFolder.newFile(); |
| Iterable<IsmRecord<byte[]>> oddValues = |
| Iterables.filter(elements, Predicates.not(EvenFilter.INSTANCE)); |
| Iterable<IsmRecord<byte[]>> evenValues = Iterables.filter(elements, EvenFilter.INSTANCE); |
| writeElementsToFile(oddValues, tmpFile); |
| IsmReader<byte[]> reader = |
| new IsmReaderImpl<>( |
| FileSystems.matchSingleFileSpec(tmpFile.getAbsolutePath()).resourceId(), CODER, cache); |
| |
| SortedMap<byte[], NavigableSet<IsmRecord<byte[]>>> sortedBySecondKey = |
| new TreeMap<>(UnsignedBytes.lexicographicalComparator()); |
| for (IsmRecord<byte[]> element : oddValues) { |
| byte[] encodedPrimaryKey = |
| CoderUtils.encodeToByteArray(CODER.getKeyComponentCoder(0), element.getKeyComponent(0)); |
| if (!sortedBySecondKey.containsKey(encodedPrimaryKey)) { |
| sortedBySecondKey.put( |
| encodedPrimaryKey, new TreeSet<>(new IsmRecordKeyComparator<>(CODER))); |
| } |
| sortedBySecondKey.get(encodedPrimaryKey).add(element); |
| } |
| |
| // The returned value should have the element as a prefix of itself. |
| for (IsmRecord<byte[]> element : oddValues) { |
| byte[] encodedPrimaryKey = |
| CoderUtils.encodeToByteArray(CODER.getKeyComponentCoder(0), element.getKeyComponent(0)); |
| assertIsmEquals( |
| reader |
| .overKeyComponents(ImmutableList.of(element.getKeyComponent(0))) |
| .getLast() |
| .getValue(), |
| sortedBySecondKey.get(encodedPrimaryKey).last()); |
| } |
| |
| // The returned value should always have the element as a prefix of itself or not exist. |
| for (IsmRecord<byte[]> element : evenValues) { |
| byte[] encodedPrimaryKey = |
| CoderUtils.encodeToByteArray(CODER.getKeyComponentCoder(0), element.getKeyComponent(0)); |
| IsmReader<byte[]>.IsmPrefixReaderIterator readerIterator = |
| reader.overKeyComponents(ImmutableList.of(element.getKeyComponent(0))); |
| WindowedValue<IsmRecord<byte[]>> lastWindowedValue = readerIterator.getLast(); |
| if (lastWindowedValue != null) { |
| assertIsmEquals( |
| lastWindowedValue.getValue(), sortedBySecondKey.get(encodedPrimaryKey).last()); |
| } |
| } |
| } |
| |
| static class IsmRecordKeyComparator<V> implements Comparator<IsmRecord<V>> { |
| private final IsmRecordCoder<V> coder; |
| |
| IsmRecordKeyComparator(IsmRecordCoder<V> coder) { |
| this.coder = coder; |
| } |
| |
| @Override |
| public int compare(IsmRecord<V> first, IsmRecord<V> second) { |
| RandomAccessData firstKeyBytes = new RandomAccessData(); |
| coder.encodeAndHash(first.getKeyComponents(), firstKeyBytes); |
| RandomAccessData secondKeyBytes = new RandomAccessData(); |
| coder.encodeAndHash(second.getKeyComponents(), secondKeyBytes); |
| return RandomAccessData.UNSIGNED_LEXICOGRAPHICAL_COMPARATOR.compare( |
| firstKeyBytes, secondKeyBytes); |
| } |
| } |
| |
| /** |
| * Specifies the minimum key size so that we can produce a random byte array with enough of a |
| * prefix to be able to create successively larger secondary keys. |
| */ |
| private static final int MIN_KEY_SIZE = 4; |
| |
| /** Specifies the percentage of keys that are metadata records when using the data generator. */ |
| private static final double PERCENT_METADATA_RECORDS = 0.01; |
| |
| /** Creates a map from Ism shard to a sorted set of IsmRecords. */ |
| private Map<Integer, SortedSet<IsmRecord<byte[]>>> dataGeneratorPerShard( |
| final int numberOfPrimaryKeys, |
| final int minNumberOfSecondaryKeys, |
| final int maxKeySize, |
| final int maxValueSize) { |
| checkState(maxKeySize >= MIN_KEY_SIZE); |
| |
| final Random random = new Random(minNumberOfSecondaryKeys); |
| |
| Map<Integer, SortedSet<IsmRecord<byte[]>>> shardToRecordMap = new HashMap<>(); |
| while (shardToRecordMap.keySet().size() < numberOfPrimaryKeys) { |
| // Generate the next primary key |
| byte[] primaryKey = new byte[random.nextInt(maxKeySize - MIN_KEY_SIZE) + MIN_KEY_SIZE]; |
| random.nextBytes(primaryKey); |
| int shardId = CODER.hash(ImmutableList.of(primaryKey)); |
| // Add a sorted set for the shard id if this shard id has never been generated before. |
| if (!shardToRecordMap.containsKey(shardId)) { |
| shardToRecordMap.put( |
| shardId, new TreeSet<IsmRecord<byte[]>>(new IsmRecordKeyComparator<byte[]>(CODER))); |
| } |
| |
| // Generate the requested number of secondary keys using the newly generated primary key. |
| byte[] secondaryKey = new byte[maxKeySize]; |
| for (int j = 0; j < minNumberOfSecondaryKeys; ++j) { |
| secondaryKey = generateNextSecondaryKey(random, maxKeySize, secondaryKey); |
| |
| // Generate the value bytes. |
| byte[] value = new byte[random.nextInt(maxValueSize)]; |
| random.nextBytes(value); |
| |
| // 1% of keys are metadata records |
| if (random.nextFloat() < PERCENT_METADATA_RECORDS) { |
| IsmRecord<byte[]> ismRecord = |
| IsmRecord.meta(ImmutableList.of(IsmFormat.getMetadataKey(), secondaryKey), value); |
| int metadataShardId = CODER.hash(ismRecord.getKeyComponents()); |
| // Add a sorted set for the shard id if this shard id has never been generated before. |
| if (!shardToRecordMap.containsKey(metadataShardId)) { |
| shardToRecordMap.put( |
| metadataShardId, |
| new TreeSet<IsmRecord<byte[]>>(new IsmRecordKeyComparator<byte[]>(CODER))); |
| } |
| shardToRecordMap.get(metadataShardId).add(ismRecord); |
| } else { |
| IsmRecord<byte[]> ismRecord = |
| IsmRecord.<byte[]>of(ImmutableList.of(primaryKey, secondaryKey), value); |
| shardToRecordMap.get(shardId).add(ismRecord); |
| } |
| } |
| } |
| return shardToRecordMap; |
| } |
| |
| private byte[] generateNextSecondaryKey( |
| Random random, int maxKeySize, byte[] previousSecondaryKey) { |
| byte[] currentSecondaryKey = new byte[random.nextInt(maxKeySize - MIN_KEY_SIZE) + MIN_KEY_SIZE]; |
| int matchingPrefix = |
| Math.min( |
| currentSecondaryKey.length, random.nextInt(maxKeySize - MIN_KEY_SIZE) + MIN_KEY_SIZE); |
| byte[] randomSuffix = new byte[currentSecondaryKey.length - matchingPrefix]; |
| random.nextBytes(randomSuffix); |
| |
| System.arraycopy( |
| previousSecondaryKey, |
| 0, |
| currentSecondaryKey, |
| 0, |
| Math.min(currentSecondaryKey.length, previousSecondaryKey.length)); |
| System.arraycopy(randomSuffix, 0, currentSecondaryKey, matchingPrefix, randomSuffix.length); |
| |
| matchingPrefix -= 1; |
| // Find the first byte which is less than 255 at the end of the matching portion. |
| while ((currentSecondaryKey[matchingPrefix] & 0xFF) == 0xFF) { |
| currentSecondaryKey[matchingPrefix] = 0; |
| matchingPrefix -= 1; |
| } |
| // Increment the last byte of the matching prefix to make sure this key is |
| // larger than the previous key. |
| currentSecondaryKey[matchingPrefix] = (byte) ((currentSecondaryKey[matchingPrefix] & 0xFF) + 1); |
| return currentSecondaryKey; |
| } |
| |
| /** Creates an iterable of IsmRecords grouped by shard id, and in ascending order per shard. */ |
| private Iterable<IsmRecord<byte[]>> dataGenerator( |
| final int numberOfPrimaryKeys, |
| final int numberOfSecondaryKeys, |
| final int approximateKeySize, |
| final int maxValueSize) { |
| |
| FluentIterable<IsmRecord<byte[]>> records = |
| FluentIterable.from( |
| dataGeneratorPerShard( |
| numberOfPrimaryKeys, numberOfSecondaryKeys, |
| approximateKeySize, maxValueSize) |
| .entrySet()) |
| .transformAndConcat(Map.Entry::getValue); |
| return records; |
| } |
| |
| @Test |
| public void testCachedTailSeekableByteChannelThrowsOnTruncate() throws Exception { |
| try (SeekableByteChannel channel = new CachedTailSeekableByteChannel(0, new byte[0])) { |
| expectedException.expect(NonWritableChannelException.class); |
| channel.truncate(0); |
| } |
| } |
| |
| @Test |
| public void testCachedTailSeekableByteChannelThrowsOnWrite() throws Exception { |
| try (SeekableByteChannel channel = new CachedTailSeekableByteChannel(0, new byte[0])) { |
| expectedException.expect(NonWritableChannelException.class); |
| channel.write(ByteBuffer.wrap(new byte[0])); |
| } |
| } |
| |
| @Test |
| public void testCachedTailSeekableByteChannelRead() throws Exception { |
| final int offset = 10; |
| try (SeekableByteChannel channel = |
| new CachedTailSeekableByteChannel(offset, new byte[] {0, 1, 2})) { |
| ByteBuffer buffer = ByteBuffer.allocate(1); |
| |
| channel.position(offset); |
| assertEquals(1, channel.read(buffer)); |
| assertEquals(0, buffer.get(0)); |
| assertEquals(offset + 1, channel.position()); |
| buffer.clear(); |
| |
| assertEquals(1, channel.read(buffer)); |
| assertEquals(1, buffer.get(0)); |
| assertEquals(offset + 2, channel.position()); |
| buffer.clear(); |
| |
| assertEquals(1, channel.read(buffer)); |
| assertEquals(2, buffer.get(0)); |
| assertEquals(offset + 3, channel.position()); |
| buffer.clear(); |
| |
| // Reposition the stream and do a read |
| channel.position(offset + 1); |
| assertEquals(1, channel.read(buffer)); |
| assertEquals(1, buffer.get(0)); |
| assertEquals(offset + 2, channel.position()); |
| buffer.clear(); |
| |
| assertEquals(1, channel.read(buffer)); |
| assertEquals(2, buffer.get(0)); |
| assertEquals(offset + 3, channel.position()); |
| buffer.clear(); |
| |
| // This read is expected to return EOF |
| assertEquals(-1, channel.read(buffer)); |
| buffer.clear(); |
| |
| // Reposition the stream to EOF and do a read, expected to return EOF |
| channel.position(offset + 3); |
| assertEquals(-1, channel.read(buffer)); |
| buffer.clear(); |
| } |
| } |
| |
| @Test |
| public void testCachedTailSeekableByteChannelSeekBeforeBounds() throws Exception { |
| try (SeekableByteChannel channel = new CachedTailSeekableByteChannel(1, new byte[0])) { |
| // Seek to only valid position |
| channel.position(1); |
| expectedException.expect(IllegalArgumentException.class); |
| channel.position(0); |
| } |
| } |
| |
| @Test |
| public void testCachedTailSeekableByteChannelSeekBeyondBounds() throws Exception { |
| try (SeekableByteChannel channel = new CachedTailSeekableByteChannel(1, new byte[0])) { |
| // Seek to only valid position |
| channel.position(1); |
| expectedException.expect(IllegalArgumentException.class); |
| channel.position(2); |
| } |
| } |
| } |
