io.datakernel.datastream.processor.AbstractStreamReducer Maven / Gradle / Ivy

Go to download
/*
 * Copyright (C) 2015 SoftIndex LLC.
 *
 * 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 io.datakernel.datastream.processor;

import io.datakernel.datastream.*;
import io.datakernel.datastream.processor.StreamReducers.Reducer;
import io.datakernel.promise.Promise;
import org.jetbrains.annotations.Nullable;

import java.util.*;
import java.util.function.Function;

import static io.datakernel.common.Preconditions.checkArgument;

/**
 * Perform aggregative functions on the elements from input streams. Searches key of item
 * with key function, selects elements with some key, reductions it and streams result sorted by key.
 * Elements from stream to input must be sorted by keys. It is Stream Transformer
 * because it represents few consumers and one supplier.
 *
 * @param  type of key of element
 * @param  type of output data
 * @param  type of accumulator
 */
@SuppressWarnings({"rawtypes", "unchecked"})
public abstract class AbstractStreamReducer implements StreamInputs, StreamOutput {
	public static final int DEFAULT_BUFFER_SIZE = 2000;

	private final List inputs = new ArrayList<>();
	private final Output output;

	private int bufferSize = DEFAULT_BUFFER_SIZE;

	@Nullable
	private Input lastInput;
	@Nullable
	private K key = null;
	@Nullable
	private A accumulator;

	private final PriorityQueue priorityQueue;
	private int streamsAwaiting;
	private int streamsOpen;

	/**
	 * Creates a new instance of AbstractStreamReducer
	 *
	 * @param keyComparator comparator for compare keys
	 */
	public AbstractStreamReducer(Comparator keyComparator) {
		this.output = new Output();
		this.priorityQueue = new PriorityQueue<>(1, (o1, o2) -> {
			int compare = ((Comparator) keyComparator).compare(o1.headKey, o2.headKey);
			if (compare != 0)
				return compare;
			return o1.index - o2.index;
		});
	}

	protected AbstractStreamReducer withBufferSize(int bufferSize) {
		checkArgument(bufferSize >= 0, "bufferSize must be positive value, got %s", bufferSize);
		this.bufferSize = bufferSize;
		return this;
	}

	protected  StreamConsumer newInput(Function keyFunction, Reducer reducer) {
		Input input = new Input(inputs.size(), priorityQueue, keyFunction, reducer, bufferSize);
		inputs.add(input);
		streamsAwaiting++;
		streamsOpen++;
		return input;
	}

	@Override
	public List> getInputs() {
		return (List) inputs;
	}

	@Override
	public StreamSupplier getOutput() {
		return output;
	}

	private final class Input extends AbstractStreamConsumer implements StreamDataAcceptor {
		private I headItem;
		private K headKey;
		private final int index;
		private final PriorityQueue priorityQueue;
		private final ArrayDeque deque = new ArrayDeque<>();
		private final int bufferSize;

		private final Function keyFunction;
		private final Reducer reducer;

		private Input(int index,
				PriorityQueue priorityQueue, Function keyFunction, Reducer reducer, int bufferSize) {
			this.index = index;
			this.priorityQueue = priorityQueue;
			this.keyFunction = keyFunction;
			this.reducer = reducer;
			this.bufferSize = bufferSize;
		}

		@Override
		protected void onWired() {
			super.onWired();
		}

		@Override
		protected void onStarted() {
			getSupplier().resume(this);
		}

		/**
		 * Processes received item. Adds item to deque, if deque size is buffer size or it is last
		 * input begins to reduce streams
		 *
		 * @param item item to receive
		 */
		@Override
		public void accept(I item) {
			if (headItem == null) {
				headItem = item;
				headKey = keyFunction.apply(headItem);
				priorityQueue.offer(this);
				streamsAwaiting--;
			} else {
				deque.offer(item);
				if (deque.size() == bufferSize) {
					getSupplier().suspend();
					produce();
				}
			}
		}

		@Override
		protected Promise onEndOfStream() {
			streamsOpen--;
			if (headItem == null) {
				streamsAwaiting--;
			}
			produce();
			assert output.getConsumer() != null;
			return output.getConsumer().getAcknowledgement();
		}

		@Override
		protected void onError(Throwable e) {
			output.close(e);
		}
	}

	private final class Output extends AbstractStreamSupplier {
		@Override
		protected void onError(Throwable e) {
			inputs.forEach(input -> input.close(e));
		}

		@Override
		protected void produce(AsyncProduceController async) {
			AbstractStreamReducer.this.produce();
		}
	}

	private void produce() {
		StreamDataAcceptor dataAcceptor = output.getCurrentDataAcceptor();
		if (dataAcceptor == null)
			return;
		while (streamsAwaiting == 0) {
			Input