Java Code Examples for org.nd4j.linalg.primitives.Pair#of()

/**
 * Read a datavec schema and record set
 * from the given arrow file.
 * @param input the input to read
 * @return the associated datavec schema and record
 */
public static Pair<Schema,ArrowWritableRecordBatch> readFromFile(FileInputStream input) throws IOException {
    BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE);
    Schema retSchema = null;
    ArrowWritableRecordBatch ret = null;
    SeekableReadChannel channel = new SeekableReadChannel(input.getChannel());
    ArrowFileReader reader = new ArrowFileReader(channel, allocator);
    reader.loadNextBatch();
    retSchema = toDatavecSchema(reader.getVectorSchemaRoot().getSchema());
    //load the batch
    VectorUnloader unloader = new VectorUnloader(reader.getVectorSchemaRoot());
    VectorLoader vectorLoader = new VectorLoader(reader.getVectorSchemaRoot());
    ArrowRecordBatch recordBatch = unloader.getRecordBatch();

    vectorLoader.load(recordBatch);
    ret = asDataVecBatch(recordBatch,retSchema,reader.getVectorSchemaRoot());
    ret.setUnloader(unloader);

    return Pair.of(retSchema,ret);

}
 

/**
 * Read a datavec schema and record set
 * from the given bytes (usually expected to be an arrow format file)
 * @param input the input to read
 * @return the associated datavec schema and record
 */
public static Pair<Schema,ArrowWritableRecordBatch> readFromBytes(byte[] input) throws IOException {
    BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE);
    Schema retSchema = null;
    ArrowWritableRecordBatch ret = null;
    SeekableReadChannel channel = new SeekableReadChannel(new ByteArrayReadableSeekableByteChannel(input));
    ArrowFileReader reader = new ArrowFileReader(channel, allocator);
    reader.loadNextBatch();
    retSchema = toDatavecSchema(reader.getVectorSchemaRoot().getSchema());
    //load the batch
    VectorUnloader unloader = new VectorUnloader(reader.getVectorSchemaRoot());
    VectorLoader vectorLoader = new VectorLoader(reader.getVectorSchemaRoot());
    ArrowRecordBatch recordBatch = unloader.getRecordBatch();

    vectorLoader.load(recordBatch);
    ret = asDataVecBatch(recordBatch,retSchema,reader.getVectorSchemaRoot());
    ret.setUnloader(unloader);

    return Pair.of(retSchema,ret);

}
 

@Override
public Pair<String, List<Writable>> apply(List<Writable> writables) {
    List<String> keyColumns = reducer.getKeyColumns();

    if(keyColumns == null){
        //Global reduction
        return Pair.of(GLOBAL_KEY, writables);
    } else {
        Schema schema = reducer.getInputSchema();
        String key;
        if (keyColumns.size() == 1)
            key = writables.get(schema.getIndexOfColumn(keyColumns.get(0))).toString();
        else {
            StringBuilder sb = new StringBuilder();
            for (int i = 0; i < keyColumns.size(); i++) {
                if (i > 0)
                    sb.append("_");
                sb.append(writables.get(schema.getIndexOfColumn(keyColumns.get(i))).toString());
            }
            key = sb.toString();
        }

        return Pair.of(key, writables);
    }
}
 

private Pair<Schema,List<List<Writable>>> recordToWrite() {
    List<List<Writable>> records = new ArrayList<>();
    records.add(Arrays.<Writable>asList(new DoubleWritable(0.0),new DoubleWritable(0.0)));
    records.add(Arrays.<Writable>asList(new DoubleWritable(0.0),new DoubleWritable(0.0)));
    Schema.Builder schemaBuilder = new Schema.Builder();
    for(int i = 0; i < 2; i++) {
        schemaBuilder.addColumnFloat("col-" + i);
    }

    return Pair.of(schemaBuilder.build(),records);
}
 

@Override
public Pair<List<Writable>, List<Writable>> apply(List<Writable> writables) {

    List<Writable> keyValues;
    if (columnIndexes.length == 1) {
        keyValues = Collections.singletonList(writables.get(columnIndexes[0]));
    } else {
        keyValues = new ArrayList<>(columnIndexes.length);
        for (int i : columnIndexes) {
            keyValues.add(writables.get(i));
        }
    }

    return Pair.of(keyValues, writables);
}
 

@Override
public Pair<List<Writable>, List<Writable>> apply(List<Writable> writables) {
    List<Writable> keyOut = new ArrayList<>(keyColumnIdxs.length);
    for (int keyColumnIdx : keyColumnIdxs) {
        keyOut.add(writables.get(keyColumnIdx));
    }
    return Pair.of(keyOut, writables);
}
 

@Override
public Pair<Text, BytesWritable> apply(Pair<String, InputStream> in) {
    try {
        return Pair.of(new Text(in.getFirst()), new BytesWritable(IOUtils.toByteArray(in.getSecond())));
    } catch (IOException e) {
        throw new IllegalStateException(e);

    }

}
 

@Override
public Pair<Writable, Long> apply(List<Writable> list) {
    return Pair.of(list.get(columnIndex), 1L);
}
 

@Override
public Pair<Writable, List<Writable>> apply(List<Writable> writables) {
    return Pair.of(writables.get(columnIdx), writables);
}
 

@Override
public Pair<Writable, List<Writable>> apply(List<Writable> writables) {
    return Pair.of(writables.get(keyColumnIdx), writables);
}