cassandra-four-zero-bridge/src/test/java/org/apache/cassandra/spark/reader/IndexDbTests.java - cassandra-analytics - Git at Google

 /*
  * 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.cassandra.spark.reader;

 import java.io.ByteArrayInputStream;
 import java.io.DataInputStream;
 import java.io.EOFException;
 import java.io.IOException;
 import java.math.BigInteger;
 import java.nio.ByteBuffer;
 import java.nio.file.Path;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
 import java.util.Objects;
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;

 import com.google.common.base.Preconditions;
 import org.junit.jupiter.api.Test;

 import org.apache.cassandra.bridge.CassandraBridgeImplementation;
 import org.apache.cassandra.bridge.TokenRange;
 import org.apache.cassandra.dht.IPartitioner;
 import org.apache.cassandra.schema.Schema;
 import org.apache.cassandra.schema.TableMetadata;
 import org.apache.cassandra.spark.data.FileType;
 import org.apache.cassandra.spark.data.SSTable;
 import org.apache.cassandra.spark.data.partitioner.Partitioner;
 import org.apache.cassandra.spark.stats.Stats;
 import org.apache.cassandra.spark.utils.TemporaryDirectory;
 import org.apache.cassandra.spark.utils.test.TestSSTable;
 import org.apache.cassandra.spark.utils.test.TestSchema;
 import org.jetbrains.annotations.NotNull;

 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.quicktheories.QuickTheory.qt;
 import static org.quicktheories.generators.Generate.constant;
 import static org.quicktheories.generators.SourceDSL.arbitrary;
 import static org.quicktheories.generators.SourceDSL.integers;
 import static org.quicktheories.generators.SourceDSL.maps;

 public class IndexDbTests
 {
     private static final CassandraBridgeImplementation BRIDGE = new CassandraBridgeImplementation();

     private static final class IndexRow implements Comparable<IndexRow>
     {
         private final BigInteger token;
         private final int value;
         private int position = 0;

         IndexRow(IPartitioner partitioner, int value)
         {
             this.token = token(partitioner, value);
             this.value = value;
         }

         public int compareTo(@NotNull IndexRow that)
         {
             return this.token.compareTo(that.token);
         }
     }

     @Test
     @SuppressWarnings("static-access")
     public void testSearchIndex()
     {
         qt().forAll(arbitrary().enumValues(Partitioner.class))
             .checkAssert(partitioner -> {
                 try (TemporaryDirectory directory = new TemporaryDirectory())
                 {
                     TestSchema schema = TestSchema.basicBuilder(BRIDGE).withCompression(false).build();
                     IPartitioner iPartitioner = BRIDGE.getPartitioner(partitioner);
                     int numRows = 5000;

                     // Write an SSTable and record token
                     List<BigInteger> tokens = new ArrayList<>(numRows);
                     schema.writeSSTable(directory, BRIDGE, partitioner, writer -> {
                         for (int row = 0; row < numRows; row++)
                         {
                             ByteBuffer key = (ByteBuffer) ByteBuffer.allocate(4).putInt(row).flip();
                             BigInteger token = ReaderUtils.tokenToBigInteger(iPartitioner.decorateKey(key).getToken());
                             tokens.add(token);
                             writer.write(row, 0, row);
                         }
                     });
                     assertEquals(1, TestSSTable.countIn(directory.path()));
                     Collections.sort(tokens);

                     TableMetadata metadata = Schema.instance.getTableMetadata(schema.keyspace, schema.table);
                     assertNotNull(metadata, "Could not find table metadata");

                     Path summaryDb = TestSSTable.firstIn(directory.path(), FileType.SUMMARY);
                     assertNotNull(summaryDb, "Could not find summary");

                     SSTable ssTable = TestSSTable.firstIn(directory.path());
                     assertNotNull(ssTable, "Could not find SSTable");

                     int rowSize = 39;
                     int sample = 4;
                     // Sample the token list and read offset in Index.db for sampled list and verify the
                     // offset matches the expected we sample the list as IndexDbUtils.findStartOffset(...)
                     // returns the previous offset, so we want to test it correctly skips tokens less than
                     // the token we are looking for before returning
                     List<BigInteger> sparseList = IntStream.range(0, tokens.size())
                                                            .filter(index -> index > 0 && index % sample == 0)
                                                            .mapToObj(tokens::get)
                                                            .collect(Collectors.toList());
                     assertEquals((numRows / 4) - 1, sparseList.size());
                     try (DataInputStream in = new DataInputStream(Objects.requireNonNull(ssTable.openPrimaryIndexStream())))
                     {
                         try
                         {
                             for (int index = 0; index < sparseList.size(); index++)
                             {
                                 BigInteger token = sparseList.get(index);
                                 long expectedOffset = (((index + 1L) * sample) - 1) * rowSize;
                                 long offset = IndexDbUtils.findStartOffset(in,
                                                                            iPartitioner,
                                                                            TokenRange.singleton(token),
                                                                            Stats.DoNothingStats.INSTANCE);
                                 assertEquals(expectedOffset, offset);
                                 ReaderUtils.skipRowIndexEntry(in);
                             }
                         }
                         catch (EOFException ignore)
                         {
                         }
                     }
                 }
                 catch (IOException exception)
                 {
                     throw new RuntimeException(exception);
                 }
             });
     }

     @Test
     @SuppressWarnings("static-access")
     public void testFindStartEndOffset()
     {
         int numValues = 5000;
         qt().forAll(arbitrary().enumValues(Partitioner.class),                           // Partitioner
                     maps().of(integers().allPositive(), constant(1)).ofSize(numValues),  // Unique keys (ignore values)
                     integers().between(1, numValues - 1))                                // Start position
             .checkAssert((partitioner, rawValues, startPos) -> {
                     IPartitioner iPartitioner = BRIDGE.getPartitioner(partitioner);
                     int rowSize = 256;

                     // Generate random index row values and sort by token
                     IndexRow[] rows = rawValues.keySet().stream()
                                                         .map(value -> new IndexRow(iPartitioner, value))
                                                         .sorted()
                                                         .toArray(IndexRow[]::new);
                     IntStream.range(0, rows.length).forEach(index -> rows[index].position = index * rowSize);  // Update position offset
                     IndexRow startRow = rows[startPos];
                     int[] valuesAndOffsets = Arrays.stream(rows)
                                                    .map(row -> new int[]{row.value, row.position})
                                                    .flatMapToInt(Arrays::stream)
                                                    .toArray();

                     try (DataInputStream in = mockDataInputStream(valuesAndOffsets))
                     {
                         long startOffset = IndexDbUtils.findStartOffset(in,
                                                                         iPartitioner,
                                                                         TokenRange.singleton(startRow.token),
                                                                         Stats.DoNothingStats.INSTANCE);
                         assertEquals(rows[startPos - 1].position, startOffset);
                         ReaderUtils.skipRowIndexEntry(in);
                     }
                     catch (IOException exception)
                     {
                         throw new RuntimeException(exception);
                     }
                 });
     }

     @Test
     @SuppressWarnings("static-access")
     public void testReadToken()
     {
         qt().withExamples(500)
             .forAll(arbitrary().enumValues(Partitioner.class), integers().all())
             .checkAssert((partitioner, value) -> {
                 IPartitioner iPartitioner = BRIDGE.getPartitioner(partitioner);
                 BigInteger expectedToken = token(iPartitioner, value);
                 try (DataInputStream in = mockDataInputStream(value, 0))
                 {
                     IndexDbUtils.readNextToken(iPartitioner, in, new Stats()
                     {
                         public void readPartitionIndexDb(ByteBuffer key, BigInteger token)
                         {
                             assertEquals(value.intValue(), key.getInt());
                             assertEquals(expectedToken, token);
                         }
                     });
                 }
                 catch (IOException exception)
                 {
                     throw new RuntimeException(exception);
                 }
             }
         );
     }

     @Test
     public void testLessThan()
     {
         assertTrue(IndexDbUtils.isLessThan(BigInteger.valueOf(4L), TokenRange.openClosed(BigInteger.valueOf(5L), BigInteger.valueOf(10L))));
         assertTrue(IndexDbUtils.isLessThan(BigInteger.valueOf(4L), TokenRange.closed(BigInteger.valueOf(5L), BigInteger.valueOf(10L))));

         assertTrue(IndexDbUtils.isLessThan(BigInteger.valueOf(5L), TokenRange.openClosed(BigInteger.valueOf(5L), BigInteger.valueOf(10L))));
         assertFalse(IndexDbUtils.isLessThan(BigInteger.valueOf(5L), TokenRange.closed(BigInteger.valueOf(5L), BigInteger.valueOf(10L))));

         assertFalse(IndexDbUtils.isLessThan(BigInteger.valueOf(6L), TokenRange.openClosed(BigInteger.valueOf(5L), BigInteger.valueOf(10L))));
         assertFalse(IndexDbUtils.isLessThan(BigInteger.valueOf(6L), TokenRange.closed(BigInteger.valueOf(5L), BigInteger.valueOf(10L))));
     }

     private static BigInteger token(IPartitioner iPartitioner, int value)
     {
         // Cast to ByteBuffer required when compiling with Java 8
         return ReaderUtils.tokenToBigInteger(iPartitioner.decorateKey((ByteBuffer) ByteBuffer.allocate(4).putInt(value).flip()).getToken());
     }

     // Creates an in-memory DataInputStream mocking Index.db bytes, with length (short), key (int), position (vint)
     private static DataInputStream mockDataInputStream(int... valuesAndOffsets) throws IOException
     {
         Preconditions.checkArgument(valuesAndOffsets.length % 2 == 0);
         int numValues = valuesAndOffsets.length / 2;

         int size = (numValues * 7);  // 2 bytes short length, 4 bytes partition key value, 1 byte promoted index
         size += IntStream.range(0, valuesAndOffsets.length)  // Variable int for position offset
                          .filter(index -> (index + 1) % 2 == 0)
                          .map(index -> valuesAndOffsets[index])
                          .map(ReaderUtils::vIntSize)
                          .sum();

         ByteBuffer buffer = ByteBuffer.allocate(size);
         for (int index = 0; index < valuesAndOffsets.length; index += 2)
         {
             buffer.putShort((short) 4)
                .putInt(valuesAndOffsets[index]);  // Value
             ReaderUtils.writePosition(valuesAndOffsets[index + 1], buffer);  // Write variable int position offset
             ReaderUtils.writePosition(0L, buffer);  // Promoted index
         }

         buffer.flip();
         byte[] bytes = new byte[buffer.remaining()];
         buffer.get(bytes);

         return new DataInputStream(new ByteArrayInputStream(bytes));
     }
 }
	/*
	* 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.cassandra.spark.reader;

	import java.io.ByteArrayInputStream;
	import java.io.DataInputStream;
	import java.io.EOFException;
	import java.io.IOException;
	import java.math.BigInteger;
	import java.nio.ByteBuffer;
	import java.nio.file.Path;
	import java.util.ArrayList;
	import java.util.Arrays;
	import java.util.Collections;
	import java.util.List;
	import java.util.Objects;
	import java.util.stream.Collectors;
	import java.util.stream.IntStream;

	import com.google.common.base.Preconditions;
	import org.junit.jupiter.api.Test;

	import org.apache.cassandra.bridge.CassandraBridgeImplementation;
	import org.apache.cassandra.bridge.TokenRange;
	import org.apache.cassandra.dht.IPartitioner;
	import org.apache.cassandra.schema.Schema;
	import org.apache.cassandra.schema.TableMetadata;
	import org.apache.cassandra.spark.data.FileType;
	import org.apache.cassandra.spark.data.SSTable;
	import org.apache.cassandra.spark.data.partitioner.Partitioner;
	import org.apache.cassandra.spark.stats.Stats;
	import org.apache.cassandra.spark.utils.TemporaryDirectory;
	import org.apache.cassandra.spark.utils.test.TestSSTable;
	import org.apache.cassandra.spark.utils.test.TestSchema;
	import org.jetbrains.annotations.NotNull;

	import static org.junit.jupiter.api.Assertions.assertEquals;
	import static org.junit.jupiter.api.Assertions.assertFalse;
	import static org.junit.jupiter.api.Assertions.assertNotNull;
	import static org.junit.jupiter.api.Assertions.assertTrue;
	import static org.quicktheories.QuickTheory.qt;
	import static org.quicktheories.generators.Generate.constant;
	import static org.quicktheories.generators.SourceDSL.arbitrary;
	import static org.quicktheories.generators.SourceDSL.integers;
	import static org.quicktheories.generators.SourceDSL.maps;

	public class IndexDbTests
	{
	private static final CassandraBridgeImplementation BRIDGE = new CassandraBridgeImplementation();

	private static final class IndexRow implements Comparable<IndexRow>
	{
	private final BigInteger token;
	private final int value;
	private int position = 0;

	IndexRow(IPartitioner partitioner, int value)
	{
	this.token = token(partitioner, value);
	this.value = value;
	}

	public int compareTo(@NotNull IndexRow that)
	{
	return this.token.compareTo(that.token);
	}
	}

	@Test
	@SuppressWarnings("static-access")
	public void testSearchIndex()
	{
	qt().forAll(arbitrary().enumValues(Partitioner.class))
	.checkAssert(partitioner -> {
	try (TemporaryDirectory directory = new TemporaryDirectory())
	{
	TestSchema schema = TestSchema.basicBuilder(BRIDGE).withCompression(false).build();
	IPartitioner iPartitioner = BRIDGE.getPartitioner(partitioner);
	int numRows = 5000;

	// Write an SSTable and record token
	List<BigInteger> tokens = new ArrayList<>(numRows);
	schema.writeSSTable(directory, BRIDGE, partitioner, writer -> {
	for (int row = 0; row < numRows; row++)
	{
	ByteBuffer key = (ByteBuffer) ByteBuffer.allocate(4).putInt(row).flip();
	BigInteger token = ReaderUtils.tokenToBigInteger(iPartitioner.decorateKey(key).getToken());
	tokens.add(token);
	writer.write(row, 0, row);
	}
	});
	assertEquals(1, TestSSTable.countIn(directory.path()));
	Collections.sort(tokens);

	TableMetadata metadata = Schema.instance.getTableMetadata(schema.keyspace, schema.table);
	assertNotNull(metadata, "Could not find table metadata");

	Path summaryDb = TestSSTable.firstIn(directory.path(), FileType.SUMMARY);
	assertNotNull(summaryDb, "Could not find summary");

	SSTable ssTable = TestSSTable.firstIn(directory.path());
	assertNotNull(ssTable, "Could not find SSTable");

	int rowSize = 39;
	int sample = 4;
	// Sample the token list and read offset in Index.db for sampled list and verify the
	// offset matches the expected we sample the list as IndexDbUtils.findStartOffset(...)
	// returns the previous offset, so we want to test it correctly skips tokens less than
	// the token we are looking for before returning
	List<BigInteger> sparseList = IntStream.range(0, tokens.size())
	.filter(index -> index > 0 && index % sample == 0)
	.mapToObj(tokens::get)
	.collect(Collectors.toList());
	assertEquals((numRows / 4) - 1, sparseList.size());
	try (DataInputStream in = new DataInputStream(Objects.requireNonNull(ssTable.openPrimaryIndexStream())))
	{
	try
	{
	for (int index = 0; index < sparseList.size(); index++)
	{
	BigInteger token = sparseList.get(index);
	long expectedOffset = (((index + 1L) * sample) - 1) * rowSize;
	long offset = IndexDbUtils.findStartOffset(in,
	iPartitioner,
	TokenRange.singleton(token),
	Stats.DoNothingStats.INSTANCE);
	assertEquals(expectedOffset, offset);
	ReaderUtils.skipRowIndexEntry(in);
	}
	}
	catch (EOFException ignore)
	{
	}
	}
	}
	catch (IOException exception)
	{
	throw new RuntimeException(exception);
	}
	});
	}

	@Test
	@SuppressWarnings("static-access")
	public void testFindStartEndOffset()
	{
	int numValues = 5000;
	qt().forAll(arbitrary().enumValues(Partitioner.class), // Partitioner
	maps().of(integers().allPositive(), constant(1)).ofSize(numValues), // Unique keys (ignore values)
	integers().between(1, numValues - 1)) // Start position
	.checkAssert((partitioner, rawValues, startPos) -> {
	IPartitioner iPartitioner = BRIDGE.getPartitioner(partitioner);
	int rowSize = 256;

	// Generate random index row values and sort by token
	IndexRow[] rows = rawValues.keySet().stream()
	.map(value -> new IndexRow(iPartitioner, value))
	.sorted()
	.toArray(IndexRow[]::new);
	IntStream.range(0, rows.length).forEach(index -> rows[index].position = index * rowSize); // Update position offset
	IndexRow startRow = rows[startPos];
	int[] valuesAndOffsets = Arrays.stream(rows)
	.map(row -> new int[]{row.value, row.position})
	.flatMapToInt(Arrays::stream)
	.toArray();

	try (DataInputStream in = mockDataInputStream(valuesAndOffsets))
	{
	long startOffset = IndexDbUtils.findStartOffset(in,
	iPartitioner,
	TokenRange.singleton(startRow.token),
	Stats.DoNothingStats.INSTANCE);
	assertEquals(rows[startPos - 1].position, startOffset);
	ReaderUtils.skipRowIndexEntry(in);
	}
	catch (IOException exception)
	{
	throw new RuntimeException(exception);
	}
	});
	}

	@Test
	@SuppressWarnings("static-access")
	public void testReadToken()
	{
	qt().withExamples(500)
	.forAll(arbitrary().enumValues(Partitioner.class), integers().all())
	.checkAssert((partitioner, value) -> {
	IPartitioner iPartitioner = BRIDGE.getPartitioner(partitioner);
	BigInteger expectedToken = token(iPartitioner, value);
	try (DataInputStream in = mockDataInputStream(value, 0))
	{
	IndexDbUtils.readNextToken(iPartitioner, in, new Stats()
	{
	public void readPartitionIndexDb(ByteBuffer key, BigInteger token)
	{
	assertEquals(value.intValue(), key.getInt());
	assertEquals(expectedToken, token);
	}
	});
	}
	catch (IOException exception)
	{
	throw new RuntimeException(exception);
	}
	}
	);
	}

	@Test
	public void testLessThan()
	{
	assertTrue(IndexDbUtils.isLessThan(BigInteger.valueOf(4L), TokenRange.openClosed(BigInteger.valueOf(5L), BigInteger.valueOf(10L))));
	assertTrue(IndexDbUtils.isLessThan(BigInteger.valueOf(4L), TokenRange.closed(BigInteger.valueOf(5L), BigInteger.valueOf(10L))));

	assertTrue(IndexDbUtils.isLessThan(BigInteger.valueOf(5L), TokenRange.openClosed(BigInteger.valueOf(5L), BigInteger.valueOf(10L))));
	assertFalse(IndexDbUtils.isLessThan(BigInteger.valueOf(5L), TokenRange.closed(BigInteger.valueOf(5L), BigInteger.valueOf(10L))));

	assertFalse(IndexDbUtils.isLessThan(BigInteger.valueOf(6L), TokenRange.openClosed(BigInteger.valueOf(5L), BigInteger.valueOf(10L))));
	assertFalse(IndexDbUtils.isLessThan(BigInteger.valueOf(6L), TokenRange.closed(BigInteger.valueOf(5L), BigInteger.valueOf(10L))));
	}

	private static BigInteger token(IPartitioner iPartitioner, int value)
	{
	// Cast to ByteBuffer required when compiling with Java 8
	return ReaderUtils.tokenToBigInteger(iPartitioner.decorateKey((ByteBuffer) ByteBuffer.allocate(4).putInt(value).flip()).getToken());
	}

	// Creates an in-memory DataInputStream mocking Index.db bytes, with length (short), key (int), position (vint)
	private static DataInputStream mockDataInputStream(int... valuesAndOffsets) throws IOException
	{
	Preconditions.checkArgument(valuesAndOffsets.length % 2 == 0);
	int numValues = valuesAndOffsets.length / 2;

	int size = (numValues * 7); // 2 bytes short length, 4 bytes partition key value, 1 byte promoted index
	size += IntStream.range(0, valuesAndOffsets.length) // Variable int for position offset
	.filter(index -> (index + 1) % 2 == 0)
	.map(index -> valuesAndOffsets[index])
	.map(ReaderUtils::vIntSize)
	.sum();

	ByteBuffer buffer = ByteBuffer.allocate(size);
	for (int index = 0; index < valuesAndOffsets.length; index += 2)
	{
	buffer.putShort((short) 4)
	.putInt(valuesAndOffsets[index]); // Value
	ReaderUtils.writePosition(valuesAndOffsets[index + 1], buffer); // Write variable int position offset
	ReaderUtils.writePosition(0L, buffer); // Promoted index
	}

	buffer.flip();
	byte[] bytes = new byte[buffer.remaining()];
	buffer.get(bytes);

	return new DataInputStream(new ByteArrayInputStream(bytes));
	}
	}