package water.udf;
import water.MRTask;
import water.fvec.Chunk;
import water.fvec.Frame;
import water.fvec.Vec;
import water.util.fp.Function;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import static water.udf.specialized.Enums.enums;
/**
* Single-column frame that knows its data type and can unfold
*/
public class UnfoldingFrame<X> extends Frame {
protected final ColumnFactory<X> factory;
protected final long len;
protected final Function<Long, List<X>> function;
protected final int width;
/** for deserialization (sigh) */
public UnfoldingFrame() {
factory = null;
len = -1;
function = null;
width = -1;
}
public UnfoldingFrame(ColumnFactory<X> factory, long len, Function<Long, List<X>> function, int width) {
super();
this.factory = factory;
this.len = len;
this.function = function;
this.width = width;
assert len >= 0: "Frame must have a nonnegative length, but found"+len;
assert width >= 0: "Multicolumn frame must have a nonnegative width, but found"+width;
}
public static <X> UnfoldingFrame<X> unfoldingFrame(final ColumnFactory<X> factory, final Column<List<X>> master, int width) {
return new UnfoldingFrame<X>(factory, master.size(), master, width) {
@Override protected Vec buildZeroVec() {
Vec v0 = DataColumns.buildZeroVec(this.len, factory.typeCode());
v0.align(master.vec());
return v0;
}
};
}
static class UnfoldingEnumFrame extends UnfoldingFrame<Integer> {
private final String[] domain;
/** for deserialization */
public UnfoldingEnumFrame() {domain = null; }
public UnfoldingEnumFrame(long length, Function<Long, List<Integer>> function, int width, String[] domain) {
super(enums(domain), length, function, width);
this.domain = domain;
assert domain != null : "An enum must have a domain";
assert domain.length > 0 : "Domain cannot be empty";
}
}
public static <X> UnfoldingEnumFrame UnfoldingEnumFrame(final Column<List<Integer>> master, int width, String[] domain) {
return new UnfoldingEnumFrame(master.size(), master, width, domain) {
@Override protected Vec buildZeroVec() {
Vec v0 = DataColumns.buildZeroVec(this.len, Vec.T_CAT);
v0.align(master.vec());
return v0;
}
};
}
protected Vec buildZeroVec() {
return DataColumns.buildZeroVec(len, factory.typeCode());
}
protected List<Vec> makeVecs() throws IOException {
Vec[] vecs = new Vec[width];
for (int j = 0; j < width; j++) {
vecs[j] = buildZeroVec();
}
MRTask task = new MRTask() {
@Override
public void map(Chunk[] cs) {
int size = cs[0].len(); // TODO(vlad): find a solution for empty
long start = cs[0].start();
for (int r = 0; r < size; r++) {
long i = r + start;
List<X> values = function.apply(i);
for (int j = 0; j < cs.length; j++) {
DataChunk<X> tc = factory.apply(cs[j]);
tc.set(r, j < values.size() ? values.get(j) : null);
}
}
}
};
MRTask mrTask = task.doAll(vecs);
return Arrays.asList(mrTask._fr.vecs());
}
public List<DataColumn<X>> materialize() throws IOException {
List<Vec> vecs = makeVecs();
List<DataColumn<X>> result = new ArrayList<>(width);
for (Vec vec : vecs) result.add(factory.newColumn(vec));
return result;
}
}