com.alibaba.alink.operator.batch.utils.MapBatchOp Maven / Gradle / Ivy

package com.alibaba.alink.operator.batch.utils;

import org.apache.flink.api.java.DataSet;
import org.apache.flink.ml.api.misc.param.Params;
import org.apache.flink.table.api.TableSchema;
import org.apache.flink.types.Row;

import com.alibaba.alink.common.mapper.Mapper;
import com.alibaba.alink.common.mapper.MapperAdapter;
import com.alibaba.alink.operator.batch.BatchOperator;

import java.util.function.BiFunction;

/**
 * class for a flat map {@link BatchOperator}.
 *
 * @param  class type of the {@link MapBatchOp} implementation itself.
 */
public class MapBatchOp> extends BatchOperator {

	private final BiFunction mapperBuilder;

	public MapBatchOp(BiFunction mapperBuilder, Params params) {
		super(params);
		this.mapperBuilder = mapperBuilder;
	}


	@Override
	public T linkFrom(BatchOperator... inputs) {
		BatchOperator in = checkAndGetFirst(inputs);

		try {
			Mapper mapper = this.mapperBuilder.apply(in.getSchema(), this.getParams());
			DataSet resultRows = in.getDataSet().map(new MapperAdapter(mapper));
			TableSchema resultSchema = mapper.getOutputSchema();
			this.setOutput(resultRows, resultSchema);
			return (T) this;
		} catch (Exception ex) {
			throw new RuntimeException(ex);
		}
	}

}