Java Code Examples for org.apache.arrow.vector.VectorSchemaRoot#getFieldVectors()

public static int run( final String[] args ) throws IOException{
  CommandLine cl;
  try{
    CommandLineParser clParser = new GnuParser();
    cl = clParser.parse( createOptions( args ) , args );
  }catch( ParseException e ){
    printHelp( args );
    throw new IOException( e );
  }

  if( cl.hasOption( "help" ) ){
    printHelp( args );
    return 0;
  }

  String input = cl.getOptionValue( "input" , null );
  String output = cl.getOptionValue( "output" , null );

  Configuration config = new Configuration();

  ArrowFileReader arrowReader = new ArrowFileReader( new FileInputStream( input ).getChannel() , new RootAllocator( Integer.MAX_VALUE ) );
  OutputStream out = FileUtil.create( output );
  MDSWriter writer = new MDSWriter( out , config );
  List<ArrowBlock> blockList = arrowReader.getRecordBlocks();
  for( ArrowBlock block : blockList ){
    VectorSchemaRoot root = arrowReader.getVectorSchemaRoot();
    arrowReader.loadRecordBatch(block);
    List<FieldVector> fieldVectorList = root.getFieldVectors();
    Spread spread = ArrowSpreadUtil.toSpread( root.getRowCount() , fieldVectorList );
    writer.append( spread );
  }
  arrowReader.close();
  writer.close();

  return 0;
}
 

public static void transferFromRoot(VectorSchemaRoot root, VectorContainer container, BufferAllocator allocator) {
  container.clear();
  // iterate over and transfer columns
  for (FieldVector fv : root.getFieldVectors()) {
    final TransferPair tp = fv.getTransferPair(allocator);
    tp.transfer();
    container.add(tp.getTo());
  }

  container.setRecordCount(root.getRowCount());
  container.addSchema(root.getSchema());
  container.buildSchema();
}
 

/**
 * Creates an {@link ArrowWriter} for the specified {@link VectorSchemaRoot}.
 */
public static ArrowWriter<Row> createRowArrowWriter(VectorSchemaRoot root, RowType rowType) {
	ArrowFieldWriter<Row>[] fieldWriters = new ArrowFieldWriter[root.getFieldVectors().size()];
	List<FieldVector> vectors = root.getFieldVectors();
	for (int i = 0; i < vectors.size(); i++) {
		FieldVector vector = vectors.get(i);
		vector.allocateNew();
		fieldWriters[i] = createRowArrowFieldWriter(vector, rowType.getTypeAt(i));
	}

	return new ArrowWriter<>(root, fieldWriters);
}
 

/**
 * Creates an {@link ArrowWriter} for blink planner for the specified {@link VectorSchemaRoot}.
 */
public static ArrowWriter<RowData> createRowDataArrowWriter(VectorSchemaRoot root, RowType rowType) {
	ArrowFieldWriter<RowData>[] fieldWriters = new ArrowFieldWriter[root.getFieldVectors().size()];
	List<FieldVector> vectors = root.getFieldVectors();
	for (int i = 0; i < vectors.size(); i++) {
		FieldVector vector = vectors.get(i);
		vector.allocateNew();
		fieldWriters[i] = createArrowFieldWriterForRow(vector, rowType.getTypeAt(i));
	}

	return new ArrowWriter<>(root, fieldWriters);
}
 

/**
 * Creates an {@link ArrowReader} for the specified {@link VectorSchemaRoot}.
 */
public static RowArrowReader createRowArrowReader(VectorSchemaRoot root, RowType rowType) {
	List<ArrowFieldReader> fieldReaders = new ArrayList<>();
	List<FieldVector> fieldVectors = root.getFieldVectors();
	for (int i = 0; i < fieldVectors.size(); i++) {
		fieldReaders.add(createRowArrowFieldReader(fieldVectors.get(i), rowType.getTypeAt(i)));
	}

	return new RowArrowReader(fieldReaders.toArray(new ArrowFieldReader[0]));
}
 

/**
 * Creates an {@link ArrowReader} for blink planner for the specified {@link VectorSchemaRoot}.
 */
public static RowDataArrowReader createRowDataArrowReader(VectorSchemaRoot root, RowType rowType) {
	List<ColumnVector> columnVectors = new ArrayList<>();
	List<FieldVector> fieldVectors = root.getFieldVectors();
	for (int i = 0; i < fieldVectors.size(); i++) {
		columnVectors.add(createColumnVector(fieldVectors.get(i), rowType.getTypeAt(i)));
	}

	return new RowDataArrowReader(columnVectors.toArray(new ColumnVector[0]));
}
 

@Test
public void T_convert_1() throws IOException{
  byte[] mdsFile = createTestData();
  InputStream in = new ByteArrayInputStream( mdsFile );
  MDSReader reader = new MDSReader();
  Configuration config = new Configuration();
  reader.setNewStream( in , mdsFile.length , config );
  MDSArrowReader arrowReader = new MDSArrowReader( reader , config );
  File testFile = new File( "target/TestConvertArrowFormatTool_T_convert_1.mds" );
  if( testFile.exists() ){
    testFile.delete();
  }
  FileOutputStream out = new FileOutputStream( testFile );
  ConvertArrowFormatTool.convert( arrowReader , out , config );

  FileInputStream arrowIn = new FileInputStream( testFile ); 
  ArrowFileReader ar = new ArrowFileReader( arrowIn.getChannel() , new RootAllocator( Integer.MAX_VALUE ) );
  VectorSchemaRoot root  = ar.getVectorSchemaRoot();
  ArrowBlock rbBlock = ar.getRecordBlocks().get(0);
  ar.loadRecordBatch(rbBlock);
  List<FieldVector> fieldVectorList = root.getFieldVectors();
  Map<String,FieldVector> vectorMap = new HashMap<String,FieldVector>();
  for( FieldVector v : fieldVectorList ){
    vectorMap.put( v.getField().getName() , v );
  }

  assertTrue( vectorMap.containsKey( "col1" ) );
  assertTrue( vectorMap.containsKey( "col2" ) );
  assertTrue( vectorMap.containsKey( "col3" ) );

  BigIntVector col1 = (BigIntVector)( vectorMap.get( "col1" ) );
  VarCharVector col2 = (VarCharVector)( vectorMap.get( "col2" ) );
  VarCharVector col3 = (VarCharVector)( vectorMap.get( "col3" ) );


  assertEquals( col1.get(0) , 100L );
  assertEquals( col1.get(1) , 200L );
  assertEquals( col1.get(2) , 300L );

  assertEquals( col2.getObject(0).toString() , "aaa" );
  assertTrue( col2.isNull(1) );
  assertTrue( col2.isNull(2) );

  assertTrue( col3.isNull(0) );
  assertEquals( col3.getObject(1).toString() , "BBB" );
  assertEquals( col3.getObject(2).toString() , "CCC" );

  testFile.delete();
  ar.close();
}