io.datakernel.aggregation.AggregationChunker 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.aggregation;

import io.datakernel.aggregation.ot.AggregationStructure;
import io.datakernel.aggregation.util.PartitionPredicate;
import io.datakernel.async.SettableStage;
import io.datakernel.async.Stage;
import io.datakernel.async.StagesAccumulator;
import io.datakernel.codegen.DefiningClassLoader;
import io.datakernel.stream.*;

import java.util.ArrayList;
import java.util.List;

public final class AggregationChunker extends ForwardingStreamConsumer implements StreamConsumerWithResult> {
	private final StreamConsumerSwitcher switcher;
	private final SettableStage> result = SettableStage.create();

	private final AggregationStructure aggregation;
	private final List fields;
	private final Class recordClass;
	private final PartitionPredicate partitionPredicate;
	private final AggregationChunkStorage storage;
	private final StagesAccumulator> chunksAccumulator;
	private final DefiningClassLoader classLoader;

	private final int chunkSize;

	private AggregationChunker(StreamConsumerSwitcher switcher,
	                           AggregationStructure aggregation, List fields,
	                           Class recordClass, PartitionPredicate partitionPredicate,
	                           AggregationChunkStorage storage,
	                           DefiningClassLoader classLoader,
	                           int chunkSize) {
		super(switcher);
		this.switcher = switcher;
		this.aggregation = aggregation;
		this.fields = fields;
		this.recordClass = recordClass;
		this.partitionPredicate = partitionPredicate;
		this.storage = storage;
		this.classLoader = classLoader;
		this.chunksAccumulator = StagesAccumulator.>create(new ArrayList<>())
				.withStage(switcher.getEndOfStream(), (accumulator, $) -> {});
		this.chunkSize = chunkSize;
		chunksAccumulator.get().whenComplete(result::trySet);
		getEndOfStream().whenException(result::trySetException);
	}

	public static  AggregationChunker create(AggregationStructure aggregation, List fields,
	                                               Class recordClass, PartitionPredicate partitionPredicate,
	                                               AggregationChunkStorage storage,
	                                               DefiningClassLoader classLoader,
	                                               int chunkSize) {
		StreamConsumerSwitcher switcher = StreamConsumerSwitcher.create();
		AggregationChunker chunker = new AggregationChunker<>(switcher, aggregation, fields, recordClass, partitionPredicate, storage, classLoader, chunkSize);
		chunker.startNewChunk();
		return chunker;
	}

	@Override
	public Stage> getResult() {
		return result;
	}

	private class ChunkWriter extends ForwardingStreamConsumer implements StreamConsumerWithResult, StreamDataReceiver {
		private final SettableStage result = SettableStage.create();
		private final long chunkId;
		private final int chunkSize;
		private final PartitionPredicate partitionPredicate;
		private StreamDataReceiver dataReceiver;

		private T first;
		private T last;
		private int count;

		boolean switched;

		public ChunkWriter(StreamConsumerWithResult actualConsumer,
		                   long chunkId, int chunkSize, PartitionPredicate partitionPredicate) {
			super(actualConsumer);
			this.chunkId = chunkId;
			this.chunkSize = chunkSize;
			this.partitionPredicate = partitionPredicate;
			actualConsumer.getResult()
					.thenApply($ -> count == 0 ? null :
							AggregationChunk.create(chunkId,
									fields,
									PrimaryKey.ofObject(first, aggregation.getKeys()),
									PrimaryKey.ofObject(last, aggregation.getKeys()),
									count))
					.whenComplete(result::trySet);
			getEndOfStream().whenException(result::trySetException);
		}

		@Override
		public void setProducer(StreamProducer producer) {
			super.setProducer(new ForwardingStreamProducer(producer) {
				@Override
				public void produce(StreamDataReceiver dataReceiver) {
					ChunkWriter.this.dataReceiver = dataReceiver;
					super.produce(ChunkWriter.this);
				}
			});
		}

		@Override
		public void onData(T item) {
			if (first == null) {
				first = item;
			}
			last = item;
			dataReceiver.onData(item);
			if (++count == chunkSize || (partitionPredicate != null && !partitionPredicate.isSamePartition(last, item))) {
				if (!switched) {
					switched = true;
					startNewChunk();
				}
			}
		}

		@Override
		public Stage getResult() {
			return result;
		}
	}

	private void startNewChunk() {
		StreamConsumerWithResult consumer = StreamConsumerWithResult.ofStage(
				storage.createId()
						.thenCompose(chunkId -> storage.write(aggregation, fields, recordClass, chunkId, classLoader)
								.thenApply(streamConsumer ->
										new ChunkWriter(streamConsumer, chunkId, chunkSize, partitionPredicate)
												.withLateBinding())));

		switcher.switchTo(consumer);

		chunksAccumulator.addStage(consumer.getResult(), (accumulator, newChunk) -> {
			if (newChunk != null && newChunk.getCount() != 0) {
				accumulator.add(newChunk);
			}
		});
	}

}