org.apache.flink.runtime.operators.CoGroupRawDriver 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.flink.runtime.operators;

import org.apache.flink.api.common.functions.CoGroupFunction;
import org.apache.flink.runtime.operators.util.TaskConfig;
import org.apache.flink.util.Collector;
import org.apache.flink.util.MutableObjectIterator;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.util.Iterator;

public class CoGroupRawDriver implements Driver, OT> {

    private static final Logger LOG = LoggerFactory.getLogger(CoGroupRawDriver.class);

    private TaskContext, OT> taskContext;

    private SimpleIterable coGroupIterator1;
    private SimpleIterable coGroupIterator2;

    @Override
    public void setup(TaskContext, OT> context) {
        this.taskContext = context;
    }

    @Override
    public int getNumberOfInputs() {
        return 2;
    }

    @Override
    public int getNumberOfDriverComparators() {
        return 0;
    }

    @Override
    public Class> getStubType() {
        @SuppressWarnings("unchecked")
        final Class> clazz =
                (Class>) (Class) CoGroupFunction.class;
        return clazz;
    }

    @Override
    public void prepare() throws Exception {
        final TaskConfig config = this.taskContext.getTaskConfig();
        if (config.getDriverStrategy() != DriverStrategy.CO_GROUP_RAW) {
            throw new Exception(
                    "Unrecognized driver strategy for CoGoup Python driver: "
                            + config.getDriverStrategy().name());
        }

        final MutableObjectIterator in1 = this.taskContext.getInput(0);
        final MutableObjectIterator in2 = this.taskContext.getInput(1);

        IT1 reuse1 = this.taskContext.getInputSerializer(0).getSerializer().createInstance();
        IT2 reuse2 = this.taskContext.getInputSerializer(1).getSerializer().createInstance();

        this.coGroupIterator1 = new SimpleIterable(reuse1, in1);
        this.coGroupIterator2 = new SimpleIterable(reuse2, in2);

        if (LOG.isDebugEnabled()) {
            LOG.debug(this.taskContext.formatLogString("CoGroup task iterator ready."));
        }
    }

    @Override
    public void run() throws Exception {
        final CoGroupFunction coGroupStub = this.taskContext.getStub();
        final Collector collector = this.taskContext.getOutputCollector();
        final SimpleIterable i1 = this.coGroupIterator1;
        final SimpleIterable i2 = this.coGroupIterator2;

        coGroupStub.coGroup(i1, i2, collector);
    }

    @Override
    public void cleanup() throws Exception {}

    @Override
    public void cancel() throws Exception {
        cleanup();
    }

    public static class SimpleIterable implements Iterable {
        private IN reuse;
        private final MutableObjectIterator iterator;

        public SimpleIterable(IN reuse, MutableObjectIterator iterator) throws IOException {
            this.iterator = iterator;
            this.reuse = reuse;
        }

        @Override
        public Iterator iterator() {
            return new SimpleIterator(reuse, iterator);
        }

        protected class SimpleIterator implements Iterator {
            private IN reuse;
            private final MutableObjectIterator iterator;
            private boolean consumed = true;

            public SimpleIterator(IN reuse, MutableObjectIterator iterator) {
                this.iterator = iterator;
                this.reuse = reuse;
            }

            @Override
            public boolean hasNext() {
                try {
                    if (!consumed) {
                        return true;
                    }
                    IN result = iterator.next(reuse);
                    consumed = result == null;
                    return !consumed;
                } catch (IOException ioex) {
                    throw new RuntimeException(
                            "An error occurred while reading the next record: " + ioex.getMessage(),
                            ioex);
                }
            }

            @Override
            public IN next() {
                consumed = true;
                return reuse;
            }

            @Override
            public void remove() { // unused
            }
        }
    }
}