hyracks-dataflow-std/src/main/java/edu/uci/ics/hyracks/dataflow/std/aggregators/MultiAggregatorFactory.java

/*
 * Copyright 2009-2010 by The Regents of the University of California
 * 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 from
 * 
 *     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 edu.uci.ics.hyracks.dataflow.std.aggregators;

import java.io.DataOutput;

import edu.uci.ics.hyracks.api.comm.IFrameTupleAccessor;
import edu.uci.ics.hyracks.api.context.IHyracksStageletContext;
import edu.uci.ics.hyracks.api.dataflow.value.RecordDescriptor;
import edu.uci.ics.hyracks.api.exceptions.HyracksDataException;
import edu.uci.ics.hyracks.dataflow.common.comm.io.ArrayTupleBuilder;
import edu.uci.ics.hyracks.dataflow.common.comm.io.FrameTupleAppender;
import edu.uci.ics.hyracks.dataflow.std.group.IAccumulatingAggregator;
import edu.uci.ics.hyracks.dataflow.std.group.IAccumulatingAggregatorFactory;
import edu.uci.ics.hyracks.dataflow.std.group.ISpillableAccumulatingAggregator;

public class MultiAggregatorFactory implements IAccumulatingAggregatorFactory {
    private static final long serialVersionUID = 1L;

    private IFieldValueResultingAggregatorFactory[] aFactories;

    public MultiAggregatorFactory(IFieldValueResultingAggregatorFactory[] aFactories) {
        this.aFactories = aFactories;
    }

    @Override
    public ISpillableAccumulatingAggregator createSpillableAggregator(IHyracksStageletContext ctx,
            RecordDescriptor inRecordDesc, final RecordDescriptor outRecordDescriptor) {
        final ISpillableFieldValueResultingAggregator aggregators[] = new ISpillableFieldValueResultingAggregator[aFactories.length];
        for (int i = 0; i < aFactories.length; ++i) {
            aggregators[i] = aFactories[i].createSpillableFieldValueResultingAggregator();
        }
        final ArrayTupleBuilder tb = new ArrayTupleBuilder(outRecordDescriptor.getFields().length);
        return new ISpillableAccumulatingAggregator() {
            private boolean pending;

            @Override
            public boolean output(FrameTupleAppender appender, IFrameTupleAccessor accessor, int tIndex,
                    int[] keyFieldIndexes) throws HyracksDataException {
                if (!pending) {
                    tb.reset();
                    for (int i = 0; i < keyFieldIndexes.length; ++i) {
                        tb.addField(accessor, tIndex, keyFieldIndexes[i]);
                    }
                    DataOutput dos = tb.getDataOutput();
                    for (int i = 0; i < aggregators.length; ++i) {
                        aggregators[i].output(dos);
                        tb.addFieldEndOffset();
                    }
                }
                if (!appender.append(tb.getFieldEndOffsets(), tb.getByteArray(), 0, tb.getSize())) {
                    pending = true;
                    return false;
                }
                return true;
            }

            @Override
            public void init(IFrameTupleAccessor accessor, int tIndex) throws HyracksDataException {
                tb.reset();
                for (int i = 0; i < aggregators.length; ++i) {
                    aggregators[i].init(accessor, tIndex);
                }
                pending = false;
            }

            @Override
            public void accumulate(IFrameTupleAccessor accessor, int tIndex) throws HyracksDataException {
                for (int i = 0; i < aggregators.length; ++i) {
                    aggregators[i].accumulate(accessor, tIndex);
                }
            }

            @Override
            public void initFromPartial(IFrameTupleAccessor accessor, int tIndex, int[] keyFieldIndexes)
                    throws HyracksDataException {
                tb.reset();
                for (int i = 0; i < aggregators.length; i++) {
                    aggregators[i].initFromPartial(accessor, tIndex, keyFieldIndexes.length + i);
                }
                pending = false;
            }

            @Override
            public void accumulatePartialResult(IFrameTupleAccessor accessor, int tIndex, int[] keyFieldIndexes)
                    throws HyracksDataException {
                for (int i = 0; i < aggregators.length; i++) {
                    aggregators[i].accumulatePartialResult(accessor, tIndex, keyFieldIndexes.length + i);
                }

            }

            @Override
            public boolean output(FrameTupleAppender appender, ArrayTupleBuilder tbder) throws HyracksDataException {
                if (!pending) {
                    // TODO Here to be fixed:
                    DataOutput dos = tbder.getDataOutput();
                    for (int i = 0; i < aggregators.length; ++i) {
                        aggregators[i].output(dos);
                        tbder.addFieldEndOffset();
                    }
                }
                if (!appender.append(tbder.getFieldEndOffsets(), tbder.getByteArray(), 0, tbder.getSize())) {
                    pending = true;
                    return false;
                }
                return true;
            }
        };
    }

    @Override
    public IAccumulatingAggregator createAggregator(IHyracksStageletContext ctx, RecordDescriptor inRecordDesc,
            RecordDescriptor outRecordDescriptor) throws HyracksDataException {
        final IFieldValueResultingAggregator aggregators[] = new IFieldValueResultingAggregator[aFactories.length];
        for (int i = 0; i < aFactories.length; ++i) {
            aggregators[i] = aFactories[i].createFieldValueResultingAggregator();
        }
        final ArrayTupleBuilder tb = new ArrayTupleBuilder(outRecordDescriptor.getFields().length);
        return new IAccumulatingAggregator() {

            private boolean pending;

            @Override
            public void init(IFrameTupleAccessor accessor, int tIndex) throws HyracksDataException {
                tb.reset();
                for (int i = 0; i < aggregators.length; ++i) {
                    aggregators[i].init(accessor, tIndex);
                }
                pending = false;
            }

            @Override
            public void accumulate(IFrameTupleAccessor accessor, int tIndex) throws HyracksDataException {
                for (int i = 0; i < aggregators.length; ++i) {
                    aggregators[i].accumulate(accessor, tIndex);
                }
            }

            @Override
            public boolean output(FrameTupleAppender appender, IFrameTupleAccessor accessor, int tIndex,
                    int[] keyFieldIndexes) throws HyracksDataException {
                if (!pending) {
                    tb.reset();
                    for (int i = 0; i < keyFieldIndexes.length; ++i) {
                        tb.addField(accessor, tIndex, keyFieldIndexes[i]);
                    }
                    DataOutput dos = tb.getDataOutput();
                    for (int i = 0; i < aggregators.length; ++i) {
                        aggregators[i].output(dos);
                        tb.addFieldEndOffset();
                    }
                }
                if (!appender.append(tb.getFieldEndOffsets(), tb.getByteArray(), 0, tb.getSize())) {
                    pending = true;
                    return false;
                }
                return true;
            }

        };
    }
}