All Downloads are FREE. Search and download functionalities are using the official Maven repository.

org.qcri.rheem.profiler.spark.SparkCollectionSourceProfiler Maven / Gradle / Ivy

The newest version!
package org.qcri.rheem.profiler.spark;

import org.qcri.rheem.core.api.Configuration;
import org.qcri.rheem.core.types.DataSetType;
import org.qcri.rheem.spark.operators.SparkCollectionSource;
import org.qcri.rheem.spark.operators.SparkTextFileSource;

import java.util.ArrayList;
import java.util.function.Supplier;

/**
 * {@link SparkOperatorProfiler} for the {@link SparkTextFileSource}.
 */
public class SparkCollectionSourceProfiler extends SparkSourceProfiler {

    private final ArrayList collection;

    public  SparkCollectionSourceProfiler(Configuration configuration,
                                             Supplier dataQuantumGenerator,
                                             DataSetType outputType) {
        this(new ArrayList<>(), configuration, dataQuantumGenerator, outputType);
    }

    private  SparkCollectionSourceProfiler(ArrayList collection,
                                                             Configuration configuration,
                                                             Supplier dataQuantumGenerator,
                                                             DataSetType outputType) {
        super(() -> new SparkCollectionSource<>(collection, outputType), configuration, dataQuantumGenerator);
        this.collection = (ArrayList) collection;
    }

    @Override
    protected void prepareInput(int inputIndex, long inputCardinality) {
        assert inputIndex == 0;
        assert inputCardinality <= Integer.MAX_VALUE;

        this.collection.clear();
        this.collection.ensureCapacity((int) inputCardinality);
        final Supplier supplier = this.dataQuantumGenerators.get(0);
        for (long i = 0; i < inputCardinality; i++) {
            this.collection.add(supplier.get());
        }
    }

    @Override
    public void cleanUp() {
        super.cleanUp();

        this.collection.clear();
        this.collection.trimToSize();
    }
}