org.apache.kafka.streams.kstream.internals.KGroupedTableImpl Maven / Gradle / Ivy

Go to download
/**
 * 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.kafka.streams.kstream.internals;

import org.apache.kafka.common.serialization.Deserializer;
import org.apache.kafka.common.serialization.Serde;
import org.apache.kafka.common.serialization.Serdes;
import org.apache.kafka.common.serialization.Serializer;
import org.apache.kafka.streams.kstream.Aggregator;
import org.apache.kafka.streams.kstream.Initializer;
import org.apache.kafka.streams.kstream.KGroupedTable;
import org.apache.kafka.streams.kstream.KStreamBuilder;
import org.apache.kafka.streams.kstream.KTable;
import org.apache.kafka.streams.kstream.Reducer;
import org.apache.kafka.streams.processor.ProcessorSupplier;
import org.apache.kafka.streams.processor.StateStoreSupplier;
import org.apache.kafka.streams.state.Stores;

import java.util.Collections;

/**
 * The implementation class of {@link KGroupedTable}.
 * 
 * @param  the key type
 * @param  the value type
 */
public class KGroupedTableImpl extends AbstractStream implements KGroupedTable {

    private static final String AGGREGATE_NAME = "KTABLE-AGGREGATE-";

    private static final String REDUCE_NAME = "KTABLE-REDUCE-";

    private static final String REPARTITION_TOPIC_SUFFIX = "-repartition";

    protected final Serde keySerde;
    protected final Serde valSerde;

    public KGroupedTableImpl(KStreamBuilder topology,
                             String name,
                             String sourceName,
                             Serde keySerde,
                             Serde valSerde) {
        super(topology, name, Collections.singleton(sourceName));
        this.keySerde = keySerde;
        this.valSerde = valSerde;
    }

    @Override
    public  KTable aggregate(Initializer initializer,
                                      Aggregator adder,
                                      Aggregator subtractor,
                                      Serde aggValueSerde,
                                      String name) {

        ProcessorSupplier> aggregateSupplier = new KTableAggregate<>(name, initializer, adder, subtractor);
        return doAggregate(aggregateSupplier, aggValueSerde, AGGREGATE_NAME, name);
    }

    @Override
    public  KTable aggregate(Initializer initializer,
                            Aggregator adder,
                            Aggregator substractor,
                            String name) {

        return aggregate(initializer, adder, substractor, null, name);
    }

    private  KTable doAggregate(ProcessorSupplier> aggregateSupplier,
                                         Serde aggValueSerde,
                                         String functionName,
                                         String name) {
        String sinkName = topology.newName(KStreamImpl.SINK_NAME);
        String sourceName = topology.newName(KStreamImpl.SOURCE_NAME);
        String funcName = topology.newName(functionName);

        String topic = name + REPARTITION_TOPIC_SUFFIX;

        Serializer keySerializer = keySerde == null ? null : keySerde.serializer();
        Deserializer keyDeserializer = keySerde == null ? null : keySerde.deserializer();
        Serializer valueSerializer = valSerde == null ? null : valSerde.serializer();
        Deserializer valueDeserializer = valSerde == null ? null : valSerde.deserializer();

        ChangedSerializer changedValueSerializer = new ChangedSerializer<>(valueSerializer);
        ChangedDeserializer changedValueDeserializer = new ChangedDeserializer<>(valueDeserializer);

        StateStoreSupplier aggregateStore = Stores.create(name)
            .withKeys(keySerde)
            .withValues(aggValueSerde)
            .persistent()
            .build();

        // send the aggregate key-value pairs to the intermediate topic for partitioning
        topology.addInternalTopic(topic);
        topology.addSink(sinkName, topic, keySerializer, changedValueSerializer, this.name);

        // read the intermediate topic
        topology.addSource(sourceName, keyDeserializer, changedValueDeserializer, topic);

        // aggregate the values with the aggregator and local store
        topology.addProcessor(funcName, aggregateSupplier, sourceName);
        topology.addStateStore(aggregateStore, funcName);

        // return the KTable representation with the intermediate topic as the sources
        return new KTableImpl<>(topology, funcName, aggregateSupplier, Collections.singleton(sourceName));
    }

    @Override
    public KTable reduce(Reducer adder,
                               Reducer subtractor,
                               String name) {
        ProcessorSupplier> aggregateSupplier = new KTableReduce<>(name, adder, subtractor);
        return doAggregate(aggregateSupplier, valSerde, REDUCE_NAME, name);
    }

    @Override
    public KTable count(String name) {
        return this.aggregate(
                new Initializer() {
                    @Override
                    public Long apply() {
                        return 0L;
                    }
                },
                new Aggregator() {
                    @Override
                    public Long apply(K aggKey, V value, Long aggregate) {
                        return aggregate + 1L;
                    }
                }, new Aggregator() {
                    @Override
                    public Long apply(K aggKey, V value, Long aggregate) {
                        return aggregate - 1L;
                    }
                },
                Serdes.Long(), name);
    }

}