net.e6tech.elements.cassandra.etl.PartitionOrderByMap Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of elements-cassandra Show documentation
elements-cassandra
The newest version!
/*
 * Copyright 2017 Futeh Kao
 *
 * Licensed 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 net.e6tech.elements.cassandra.etl;

import net.e6tech.elements.cassandra.Sibyl;

import java.util.*;
import java.util.function.BiConsumer;
import java.util.function.Consumer;

/**
 * This is a map that support storing entities with a SINGLE partition key and ONE or more clustering key.  However, only
 * first clustering key is used to sort the list for each partition key.
 */
public class PartitionOrderByMap {

    private Map> partitionMap = new LinkedHashMap<>(); // partition key, list of T object
    private Map> primaryKeys = new LinkedHashMap<>(); // partition key, list of primary keys
    private Sibyl sibyl;
    private Class sourceClass;

    public PartitionOrderByMap(Sibyl sibyl, Class sourceClass) {
        this.sibyl = sibyl;
        this.sourceClass = sourceClass;
    }

    protected Inspector getInspector(Class cls) {
        return sibyl.getInspector(cls);
    }

    @SuppressWarnings("unchecked")
    public PartitionOrderByMap addAll(Collection objects) {
        for (T object : objects) {
            List list = partitionMap.computeIfAbsent((Comparable) getInspector(sourceClass).getPartitionKey(object, 0), key -> new ArrayList<>());
            list.add(object);
        }

        // sort each list using its clustering key
        for (List list : partitionMap.values()) {
            Collections.sort(list, Comparator.comparing(t -> (Comparable) getInspector(sourceClass).getClusteringKey(t, 0)));
        }

        // primary keys
        for (Map.Entry> entry : partitionMap.entrySet()) {
            List list = primaryKeys.computeIfAbsent(entry.getKey(), key -> new ArrayList<>());
            for (T t : entry.getValue()) {
                list.add(getInspector(t.getClass()).getPrimaryKey(t));
            }
        }
        return this;
    }

    public PartitionOrderByMap getValueList(Comparable partitionKey, Consumer> consumer) {
        List list = partitionMap.get(partitionKey);
        if (list != null)
            consumer.accept(list);
        return this;
    }

    public PartitionOrderByMap getPrimaryKeyList(Comparable partitionKey, Consumer> consumer) {
        List list = primaryKeys.get(partitionKey);
        if (list != null)
            consumer.accept(list);
        return this;
    }

    public PartitionOrderByMap forEachValueList(BiConsumer> consumer) {
        for (Map.Entry> entry : partitionMap.entrySet()) {
            consumer.accept(entry.getKey(), entry.getValue());
        }
        return this;
    }

    public PartitionOrderByMap forEachPrimaryKeyList(BiConsumer> consumer) {
        for (Map.Entry> entry : primaryKeys.entrySet()) {
            consumer.accept(entry.getKey(), entry.getValue());
        }
        return this;
    }

    public PartitionOrderByMap forEachValue(Consumer consumer) {
        for (Map.Entry> entry : partitionMap.entrySet()) {
            for (T t : entry.getValue())
             consumer.accept(t);
        }
        return this;
    }

    public PartitionOrderByMap forEachPrimaryKey(Consumer consumer) {
        for (Map.Entry> entry : primaryKeys.entrySet()) {
            for (PrimaryKey t : entry.getValue())
                consumer.accept(t);
        }
        return this;
    }
}