test/microbench/org/apache/cassandra/test/microbench/tries/ComparisonReadBench.java

/*
 * 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.test.microbench.tries;

import java.math.BigDecimal;
import java.math.BigInteger;
import java.nio.charset.StandardCharsets;
import java.util.AbstractMap;
import java.util.Comparator;
import java.util.Map;
import java.util.NavigableMap;
import java.util.Random;
import java.util.TreeMap;
import java.util.concurrent.ConcurrentSkipListMap;
import java.util.concurrent.TimeUnit;
import java.util.function.BiConsumer;
import java.util.function.Consumer;

import com.google.common.base.Throwables;
import com.google.common.collect.Iterables;

import org.apache.cassandra.db.marshal.DecimalType;
import org.apache.cassandra.db.marshal.IntegerType;
import org.apache.cassandra.db.tries.InMemoryTrie;
import org.apache.cassandra.io.compress.BufferType;
import org.apache.cassandra.utils.ByteArrayUtil;
import org.apache.cassandra.utils.bytecomparable.ByteComparable;
import org.apache.cassandra.utils.bytecomparable.ByteSource;
import org.apache.cassandra.utils.bytecomparable.ByteSourceInverse;
import org.github.jamm.MemoryMeter;
import org.openjdk.jmh.annotations.Benchmark;
import org.openjdk.jmh.annotations.BenchmarkMode;
import org.openjdk.jmh.annotations.Fork;
import org.openjdk.jmh.annotations.Level;
import org.openjdk.jmh.annotations.Measurement;
import org.openjdk.jmh.annotations.Mode;
import org.openjdk.jmh.annotations.OutputTimeUnit;
import org.openjdk.jmh.annotations.Param;
import org.openjdk.jmh.annotations.Scope;
import org.openjdk.jmh.annotations.Setup;
import org.openjdk.jmh.annotations.State;
import org.openjdk.jmh.annotations.Threads;
import org.openjdk.jmh.annotations.Warmup;

@BenchmarkMode(Mode.AverageTime)
@OutputTimeUnit(TimeUnit.MILLISECONDS)
@Warmup(iterations = 2, time = 1)
@Measurement(iterations = 3, time = 1)
@Fork(value = 1,jvmArgsAppend = { "-Xmx4G", "-Xms4G", "-Djmh.executor=CUSTOM", "-Djmh.executor.class=org.apache.cassandra.test.microbench.FastThreadExecutor"})
@Threads(1) // no concurrent writes
@State(Scope.Benchmark)
public class ComparisonReadBench
{
    // Note: To see a printout of the usage for each object, add .enableDebug() here (most useful with smaller number of
    // partitions).
    static MemoryMeter meter = new MemoryMeter().ignoreKnownSingletons()
                                                .withGuessing(MemoryMeter.Guess.FALLBACK_UNSAFE);

    @Param({"ON_HEAP"})
    BufferType bufferType = BufferType.OFF_HEAP;

    @Param({"1000", "100000", "10000000"})
    int count = 1000;

    @Param({"TREE_MAP", "CSLM", "TRIE"})
    MapOption map = MapOption.TRIE;

    @Param({"LONG"})
    TypeOption type = TypeOption.LONG;

    final static InMemoryTrie.UpsertTransformer<Byte, Byte> resolver = (x, y) -> y;

    Access<?> access;

    @Setup(Level.Trial)
    public void setup() throws Throwable
    {
        switch (map)
        {
            case TREE_MAP:
                access = new NavigableMapAccess(new TreeMap(type.type.comparator()), type.type);
                break;
            case CSLM:
                access = new NavigableMapAccess(new ConcurrentSkipListMap(type.type.comparator()), type.type);
                break;
            case TRIE:
                access = new TrieAccess(type.type);
                break;
        }
        Random rand = new Random(1);

        System.out.format("Putting %,d\n", count);
        long time = System.currentTimeMillis();
        for (long current = 0; current < count; ++current)
        {
            long l = rand.nextLong();
            access.put(l, Byte.valueOf((byte) (l >> 56)));
        }
        time = System.currentTimeMillis() - time;
        System.out.format("Took %.3f seconds\n", time * 0.001);
        access.printSize();
    }

    interface Type<T>
    {
        T fromLong(long l);
        T fromByteComparable(ByteComparable bc);
        ByteComparable longToByteComparable(long l);
        Comparator<T> comparator();
    }

    public enum TypeOption
    {
        LONG(new LongType()),
        BIGINT(new BigIntType()),
        DECIMAL(new BigDecimalType()),
        STRING_BASE16(new StringType(16)),
        STRING_BASE10(new StringType(10)),
        STRING_BASE4(new StringType(4)),
        ARRAY(new ArrayType(1)),
        ARRAY_REP3(new ArrayType(3)),
        ARRAY_REP7(new ArrayType(7));

        final Type<?> type;

        TypeOption(Type<?> type)
        {
            this.type = type;
        }
    }

    static class LongType implements Type<Long>
    {
        public Long fromLong(long l)
        {
            return l;
        }

        public Long fromByteComparable(ByteComparable bc)
        {
            return ByteSourceInverse.getSignedLong(bc.asComparableBytes(ByteComparable.Version.OSS50));
        }

        public ByteComparable longToByteComparable(long l)
        {
            return ByteComparable.of(l);
        }

        public Comparator<Long> comparator()
        {
            return Comparator.naturalOrder();
        }
    }

    static class BigIntType implements Type<BigInteger>
    {
        public BigInteger fromLong(long l)
        {
            return BigInteger.valueOf(l);
        }

        public BigInteger fromByteComparable(ByteComparable bc)
        {
            return IntegerType.instance.compose(IntegerType.instance.fromComparableBytes(ByteSource.peekable(bc.asComparableBytes(ByteComparable.Version.OSS50)),
                                                                                         ByteComparable.Version.OSS50));
        }

        public ByteComparable longToByteComparable(long l)
        {
            return v -> IntegerType.instance.asComparableBytes(IntegerType.instance.decompose(fromLong(l)), v);
        }

        public Comparator<BigInteger> comparator()
        {
            return Comparator.naturalOrder();
        }
    }

    static class BigDecimalType implements Type<BigDecimal>
    {
        public BigDecimal fromLong(long l)
        {
            return BigDecimal.valueOf(l);
        }

        public BigDecimal fromByteComparable(ByteComparable bc)
        {
            return DecimalType.instance.compose(DecimalType.instance.fromComparableBytes(ByteSource.peekable(bc.asComparableBytes(ByteComparable.Version.OSS50)),
                                                                                         ByteComparable.Version.OSS50));
        }

        public ByteComparable longToByteComparable(long l)
        {
            return v -> DecimalType.instance.asComparableBytes(DecimalType.instance.decompose(fromLong(l)), v);
        }

        public Comparator<BigDecimal> comparator()
        {
            return Comparator.naturalOrder();
        }
    }

    static class StringType implements Type<String>
    {
        final int base;

        StringType(int base)
        {
            this.base = base;
        }

        public String fromLong(long l)
        {
            return Long.toString(l, base);
        }

        public String fromByteComparable(ByteComparable bc)
        {
            return new String(ByteSourceInverse.readBytes(bc.asComparableBytes(ByteComparable.Version.OSS50)), StandardCharsets.UTF_8);
        }

        public ByteComparable longToByteComparable(long l)
        {
            return ByteComparable.fixedLength(fromLong(l).getBytes(StandardCharsets.UTF_8));
        }

        public Comparator<String> comparator()
        {
            return Comparator.naturalOrder();
        }
    }

    static class ArrayType implements Type<byte[]>
    {
        final int reps;

        ArrayType(int reps)
        {
            this.reps = reps;
        }

        public byte[] fromLong(long l)
        {
            byte[] value = new byte[8 * reps];
            for (int i = 0; i < 8; ++i)
            {
                for (int j = 0; j < reps; ++j)
                    value[i * reps + j] = (byte)(l >> (56 - i * 8));
            }
            return value;
        }

        public byte[] fromByteComparable(ByteComparable bc)
        {
            return ByteSourceInverse.readBytes(bc.asComparableBytes(ByteComparable.Version.OSS50));
        }

        public ByteComparable longToByteComparable(long l)
        {
            return ByteComparable.fixedLength(fromLong(l));
        }

        public Comparator<byte[]> comparator()
        {
            return ByteArrayUtil::compareUnsigned;
        }
    }

    interface Access<T>
    {
        void put(long v, byte b);
        byte get(long v);
        Iterable<Byte> values();
        Iterable<Byte> valuesSlice(long left, boolean includeLeft, long right, boolean includeRight);
        Iterable<Map.Entry<T, Byte>> entrySet();
        void consumeValues(Consumer<Byte> consumer);
        void consumeEntries(BiConsumer<T, Byte> consumer);
        void printSize();
    }

    public enum MapOption
    {
        TREE_MAP,
        CSLM,
        TRIE
    }

    class TrieAccess<T> implements Access<T>
    {
        final InMemoryTrie<Byte> trie;
        final Type<T> type;

        TrieAccess(Type<T> type)
        {
            this.type = type;
            trie = new InMemoryTrie<>(bufferType);
        }

        public void put(long v, byte b)
        {
            try
            {
                trie.putRecursive(type.longToByteComparable(v), b, resolver);
            }
            catch (InMemoryTrie.SpaceExhaustedException e)
            {
                throw Throwables.propagate(e);
            }
        }

        public byte get(long v)
        {
            return trie.get(type.longToByteComparable(v));
        }

        public Iterable<Byte> values()
        {
            return trie.values();
        }

        public Iterable<Byte> valuesSlice(long left, boolean includeLeft, long right, boolean includeRight)
        {
            return trie.subtrie(type.longToByteComparable(left), includeLeft, type.longToByteComparable(right), includeRight)
                       .values();
        }

        public Iterable<Map.Entry<T, Byte>> entrySet()
        {
            return Iterables.transform(trie.entrySet(),
                    en -> new AbstractMap.SimpleEntry<>(type.fromByteComparable(en.getKey()),
                                                        en.getValue()));
        }

        public void consumeValues(Consumer<Byte> consumer)
        {
            trie.forEachValue(consumer::accept);
        }

        public void consumeEntries(BiConsumer<T, Byte> consumer)
        {
            trie.forEachEntry((key, value) -> consumer.accept(type.fromByteComparable(key), value));
        }

        public void printSize()
        {
            long deepsize = meter.measureDeep(trie);
            System.out.format("Trie size on heap %,d off heap %,d deep size %,d\n",
                              trie.sizeOnHeap(), trie.sizeOffHeap(), deepsize);
            System.out.format("per entry on heap %.2f off heap %.2f deep size %.2f\n",
                              trie.sizeOnHeap() * 1.0 / count, trie.sizeOffHeap() * 1.0 / count, deepsize * 1.0 / count);
        }
    }

    class NavigableMapAccess<T> implements Access<T>
    {
        final NavigableMap<T, Byte> navigableMap;
        final Type<T> type;

        NavigableMapAccess(NavigableMap<T, Byte> navigableMap, Type<T> type)
        {
            this.navigableMap = navigableMap;
            this.type = type;
        }

        public void put(long v, byte b)
        {
            navigableMap.put(type.fromLong(v), b);
        }

        public byte get(long v)
        {
            return navigableMap.get(type.fromLong(v));
        }

        public Iterable<Byte> values()
        {
            return navigableMap.values();
        }

        public Iterable<Byte> valuesSlice(long left, boolean includeLeft, long right, boolean includeRight)
        {
            return navigableMap.subMap(type.fromLong(left), includeLeft, type.fromLong(right), includeRight)
                               .values();
        }

        public Iterable<Map.Entry<T, Byte>> entrySet()
        {
            return navigableMap.entrySet();
        }

        public void consumeValues(Consumer<Byte> consumer)
        {
            navigableMap.values().forEach(consumer);
        }

        public void consumeEntries(BiConsumer<T, Byte> consumer)
        {
            navigableMap.forEach(consumer);
        }

        public void printSize()
        {
            long size = meter.measureDeep(navigableMap);
            System.out.format(map + " size on heap %,d\n", size);
            System.out.format("per entry on heap %.2f\n", size * 1.0 / count);
        }
    }

    @Benchmark
    public void getRandom()
    {
        Random rand = new Random(1);

        for (long current = 0; current < count; ++current)
        {
            long l = rand.nextLong();
            Byte res = access.get(l);
            if (res.byteValue() != l >> 56)
                throw new AssertionError();
        }
    }

    @Benchmark
    public int iterateValues()
    {
        int sum = 0;
        for (byte b : access.values())
            sum += b;
        return sum;
    }

    @Benchmark
    public int consumeValues()
    {
        class Counter implements Consumer<Byte>
        {
            int sum = 0;

            @Override
            public void accept(Byte aByte)
            {
                sum += aByte;
            }
        }
        Counter counter = new Counter();
        access.consumeValues(counter);
        return counter.sum;
    }

    @Benchmark
    public int consumeEntries()
    {
        class Counter<T> implements BiConsumer<T, Byte>
        {
            int sum = 0;

            @Override
            public void accept(T key, Byte aByte)
            {
                sum += aByte;
            }
        }
        Counter counter = new Counter();
        access.consumeEntries(counter);
        return counter.sum;
    }

    @Benchmark
    public int iterateEntries()
    {
        int sum = 0;
        for (Map.Entry<?, Byte> en : access.entrySet())
            sum += en.getValue();
        return sum;
    }


    @Benchmark
    public int getByIterateValueSlice()
    {
        Random rand = new Random(1);
        int sum = 0;
        for (int i = 0; i < count; ++i)
        {
            long v = rand.nextLong();
            Iterable<Byte> values = access.valuesSlice(v, true, v, true);
            for (byte b : values)
                sum += b;
        }
        return sum;
    }

    @Benchmark
    public int iterateValuesLimited()
    {
        int sum = 0;
        Iterable<Byte> values = access.valuesSlice(0L, false, Long.MAX_VALUE / 2, true); // 1/4
        for (byte b : values)
            sum += b;
        return sum;
    }
}