Search in sources :

Example 1 with ArrowColumn

use of edu.iu.dsc.tws.common.table.ArrowColumn in project twister2 by DSC-SPIDAL.

the class ArrowAllToAll method isComplete.

/**
 * Check weather complete
 * @return true if operation is complete
 */
public boolean isComplete() {
    if (completed) {
        return true;
    }
    boolean isAllEmpty = true;
    for (Map.Entry<Integer, PendingSendTable> t : inputs.entrySet()) {
        PendingSendTable pst = t.getValue();
        if (pst.status == ArrowHeader.HEADER_INIT) {
            if (!pst.pending.isEmpty()) {
                pst.currentTable = pst.pending.poll();
                assert !pst.target.isEmpty();
                pst.currentTarget = pst.target.poll();
                pst.status = ArrowHeader.COLUMN_CONTINUE;
            }
        }
        if (pst.status == ArrowHeader.COLUMN_CONTINUE) {
            int noOfColumns = pst.currentTable.getColumns().size();
            boolean canContinue = true;
            while (pst.columnIndex < noOfColumns && canContinue) {
                ArrowColumn col = pst.currentTable.getColumns().get(pst.columnIndex);
                FieldVector vector = col.getVector();
                List<ArrowFieldNode> nodes = new ArrayList<>();
                List<ArrowBuf> bufs = new ArrayList<>();
                appendNodes(vector, nodes, bufs);
                while (pst.bufferIndex < bufs.size()) {
                    ArrowBuf buf = bufs.get(pst.bufferIndex);
                    int[] hdr = new int[HEADER_SIZE];
                    hdr[0] = pst.columnIndex;
                    hdr[1] = pst.bufferIndex;
                    hdr[2] = bufs.size();
                    hdr[3] = vector.getValueCount();
                    int length = (int) buf.capacity();
                    hdr[4] = length;
                    // target
                    hdr[5] = pst.currentTarget;
                    boolean accept = all.insert(buf.nioBuffer(), length, hdr, HEADER_SIZE, t.getKey());
                    if (!accept) {
                        canContinue = false;
                        break;
                    }
                    pst.bufferIndex++;
                }
                if (canContinue) {
                    pst.bufferIndex = 0;
                    pst.columnIndex++;
                }
            }
            if (canContinue) {
                pst.columnIndex = 0;
                pst.bufferIndex = 0;
                pst.status = ArrowHeader.HEADER_INIT;
            }
        }
        if (!pst.pending.isEmpty() || pst.status == ArrowHeader.COLUMN_CONTINUE) {
            isAllEmpty = false;
        }
    }
    if (isAllEmpty && finished && !finishedSent) {
        all.finish();
        finishedSent = true;
    }
    boolean b = isAllEmpty && all.isComplete() && finishedSources.size() == sourceWorkerList.size();
    if (b) {
        completed = true;
    }
    return b;
}
Also used : ArrowBuf(io.netty.buffer.ArrowBuf) ArrayList(java.util.ArrayList) FieldVector(org.apache.arrow.vector.FieldVector) ArrowColumn(edu.iu.dsc.tws.common.table.ArrowColumn) ArrowFieldNode(org.apache.arrow.vector.ipc.message.ArrowFieldNode) HashMap(java.util.HashMap) Map(java.util.Map)

Example 2 with ArrowColumn

use of edu.iu.dsc.tws.common.table.ArrowColumn in project twister2 by DSC-SPIDAL.

the class STPartition method isComplete.

@Override
public boolean isComplete() {
    for (Map.Entry<Integer, Queue<Table>> e : inputs.entrySet()) {
        if (e.getValue().isEmpty()) {
            continue;
        }
        // partition the table, default
        Table t = e.getValue().poll();
        List<ArrowColumn> columns = t.getColumns();
        ArrowColumn col = columns.get(indexes[0]);
        for (int i = 0; i < col.getVector().getValueCount(); i++) {
            Row row = new OneRow(col.get(i));
            int target = selector.next(e.getKey(), row);
            TableBuilder builder = partitionedTables.get(target);
            if (builder == null) {
                builder = new ArrowTableBuilder(schema, allocator);
                this.partitionedTables.put(target, builder);
            }
            for (int j = 0; j < columns.size(); j++) {
                builder.getColumns().get(j).addValue(columns.get(j).get(i));
            }
        }
    }
    if (finished) {
        for (Map.Entry<Integer, TableBuilder> e : partitionedTables.entrySet()) {
            Table t = e.getValue().build();
            allToAll.insert(t, e.getKey());
        }
        // clear the tables, so we won't build the tables again
        partitionedTables.clear();
        for (int s : finishedSources) {
            allToAll.finish(s);
        }
        // clear so, we won't call finish again
        finishedSources.clear();
        return allToAll.isComplete();
    }
    return false;
}
Also used : Table(edu.iu.dsc.tws.common.table.Table) ArrowTableBuilder(edu.iu.dsc.tws.common.table.ArrowTableBuilder) ArrowTableBuilder(edu.iu.dsc.tws.common.table.ArrowTableBuilder) TableBuilder(edu.iu.dsc.tws.common.table.TableBuilder) ArrowColumn(edu.iu.dsc.tws.common.table.ArrowColumn) OneRow(edu.iu.dsc.tws.common.table.OneRow) Row(edu.iu.dsc.tws.common.table.Row) OneRow(edu.iu.dsc.tws.common.table.OneRow) HashMap(java.util.HashMap) Map(java.util.Map) Queue(java.util.Queue)

Example 3 with ArrowColumn

use of edu.iu.dsc.tws.common.table.ArrowColumn in project twister2 by DSC-SPIDAL.

the class ArrowAllToAll method onReceive.

@Override
public void onReceive(int source, ChannelBuffer buffer, int length) {
    PendingReceiveTable table = receives.get(source);
    receivedBuffers++;
    ArrowBuf buf = ((ArrowChannelBuffer) buffer).getArrowBuf();
    table.buffers.add(buf);
    if (table.bufferIndex == 0) {
        table.fieldNodes.add(new ArrowFieldNode(table.noArray, 0));
    }
    VectorSchemaRoot schemaRoot = table.root;
    List<FieldVector> fieldVectors = schemaRoot.getFieldVectors();
    // we received everything for this array
    if (table.noBuffers == table.bufferIndex + 1) {
        FieldVector fieldVector = fieldVectors.get(table.columnIndex);
        loadBuffers(fieldVector, fieldVector.getField(), table.buffers.iterator(), table.fieldNodes.iterator());
        table.arrays.add(fieldVector);
        table.buffers.clear();
        if (table.arrays.size() == schemaRoot.getFieldVectors().size()) {
            List<ArrowColumn> columns = new ArrayList<>();
            // create the table
            for (FieldVector v : fieldVectors) {
                ArrowColumn c;
                if (v instanceof BaseFixedWidthVector) {
                    if (v instanceof IntVector) {
                        c = new Int4Column((IntVector) v);
                    } else if (v instanceof Float4Vector) {
                        c = new Float4Column((Float4Vector) v);
                    } else if (v instanceof Float8Vector) {
                        c = new Float8Column((Float8Vector) v);
                    } else if (v instanceof UInt8Vector) {
                        c = new Int8Column((UInt8Vector) v);
                    } else if (v instanceof UInt2Vector) {
                        c = new UInt2Column((UInt2Vector) v);
                    } else {
                        throw new RuntimeException("Un-supported type : " + v.getClass().getName());
                    }
                } else if (v instanceof BaseVariableWidthVector) {
                    if (v instanceof VarCharVector) {
                        c = new StringColumn((VarCharVector) v);
                    } else if (v instanceof VarBinaryVector) {
                        c = new BinaryColumn((VarBinaryVector) v);
                    } else {
                        throw new RuntimeException("Un-supported type : " + v.getClass().getName());
                    }
                } else {
                    throw new RuntimeException("Un-supported type : " + v.getClass().getName());
                }
                columns.add(c);
            }
            Table t = new ArrowTable(schemaRoot.getSchema(), table.noArray, columns);
            LOG.info("Received table from source " + source + " to " + table.target + " count" + t.rowCount());
            recvCallback.onReceive(source, table.target, t);
            table.clear();
        }
    }
}
Also used : BaseFixedWidthVector(org.apache.arrow.vector.BaseFixedWidthVector) VectorSchemaRoot(org.apache.arrow.vector.VectorSchemaRoot) ArrowBuf(io.netty.buffer.ArrowBuf) Float4Vector(org.apache.arrow.vector.Float4Vector) BinaryColumn(edu.iu.dsc.tws.common.table.arrow.BinaryColumn) ArrayList(java.util.ArrayList) VarBinaryVector(org.apache.arrow.vector.VarBinaryVector) ArrowColumn(edu.iu.dsc.tws.common.table.ArrowColumn) BaseVariableWidthVector(org.apache.arrow.vector.BaseVariableWidthVector) ArrowFieldNode(org.apache.arrow.vector.ipc.message.ArrowFieldNode) Int8Column(edu.iu.dsc.tws.common.table.arrow.Int8Column) StringColumn(edu.iu.dsc.tws.common.table.arrow.StringColumn) Table(edu.iu.dsc.tws.common.table.Table) ArrowTable(edu.iu.dsc.tws.common.table.arrow.ArrowTable) IntVector(org.apache.arrow.vector.IntVector) UInt2Column(edu.iu.dsc.tws.common.table.arrow.UInt2Column) Float8Vector(org.apache.arrow.vector.Float8Vector) VarCharVector(org.apache.arrow.vector.VarCharVector) FieldVector(org.apache.arrow.vector.FieldVector) Float4Column(edu.iu.dsc.tws.common.table.arrow.Float4Column) UInt8Vector(org.apache.arrow.vector.UInt8Vector) Float8Column(edu.iu.dsc.tws.common.table.arrow.Float8Column) Int4Column(edu.iu.dsc.tws.common.table.arrow.Int4Column) ArrowTable(edu.iu.dsc.tws.common.table.arrow.ArrowTable) UInt2Vector(org.apache.arrow.vector.UInt2Vector)

Aggregations

ArrowColumn (edu.iu.dsc.tws.common.table.ArrowColumn)3 Table (edu.iu.dsc.tws.common.table.Table)2 ArrowBuf (io.netty.buffer.ArrowBuf)2 ArrayList (java.util.ArrayList)2 HashMap (java.util.HashMap)2 Map (java.util.Map)2 FieldVector (org.apache.arrow.vector.FieldVector)2 ArrowFieldNode (org.apache.arrow.vector.ipc.message.ArrowFieldNode)2 ArrowTableBuilder (edu.iu.dsc.tws.common.table.ArrowTableBuilder)1 OneRow (edu.iu.dsc.tws.common.table.OneRow)1 Row (edu.iu.dsc.tws.common.table.Row)1 TableBuilder (edu.iu.dsc.tws.common.table.TableBuilder)1 ArrowTable (edu.iu.dsc.tws.common.table.arrow.ArrowTable)1 BinaryColumn (edu.iu.dsc.tws.common.table.arrow.BinaryColumn)1 Float4Column (edu.iu.dsc.tws.common.table.arrow.Float4Column)1 Float8Column (edu.iu.dsc.tws.common.table.arrow.Float8Column)1 Int4Column (edu.iu.dsc.tws.common.table.arrow.Int4Column)1 Int8Column (edu.iu.dsc.tws.common.table.arrow.Int8Column)1 StringColumn (edu.iu.dsc.tws.common.table.arrow.StringColumn)1 UInt2Column (edu.iu.dsc.tws.common.table.arrow.UInt2Column)1