org.apache.arrow.memory.BufferAllocator Java Examples

public static List<FieldVector> toArrowColumns(BufferAllocator bufferAllocator, Schema schema, List<List<Writable>> dataVecRecord) {
    int numRows = dataVecRecord.size();
    List<FieldVector> ret = createFieldVectors(bufferAllocator, schema, numRows);

    for (int j = 0; j < schema.numColumns(); ++j) {
        FieldVector fieldVector = ret.get(j);
        int row = 0;

        for (Iterator var8 = dataVecRecord.iterator(); var8.hasNext(); ++row) {
            List<Writable> record = (List) var8.next();
            Writable writable = record.get(j);
            setValue(schema.getType(j), fieldVector, writable, row);
        }
    }

    return ret;
}
 

public RecordBatchData(VectorAccessible batch, BufferAllocator allocator, boolean applySVMode) {
  this.recordCount = batch.getRecordCount();

  final List<ValueVector> vectors = Lists.newArrayList();

  if (applySVMode && batch.getSchema().getSelectionVectorMode() == SelectionVectorMode.TWO_BYTE) {
    this.sv2 = batch.getSelectionVector2().clone();
  } else {
    this.sv2 = null;
  }

  for (VectorWrapper<?> v : batch) {
    if (v.isHyper()) {
      throw new UnsupportedOperationException("Record batch data can't be created based on a hyper batch.");
    }
    TransferPair tp = v.getValueVector().getTransferPair(allocator);
    tp.transfer();
    vectors.add(tp.getTo());
  }

  container.addCollection(vectors);
  container.setRecordCount(recordCount);
  container.buildSchema(applySVMode ? batch.getSchema().getSelectionVectorMode() : SelectionVectorMode.NONE);
}
 

/**
 * @return next updateId
 */
private static UpdateId getUpdateId(final JobId jobId, final JobsService jobsService, BufferAllocator allocator) {
  final int fetchLimit = 1000;
  UpdateIdWrapper updateIdWrapper = new UpdateIdWrapper();

  int offset = 0;
  while (true) {
    try (final JobDataFragment data = JobDataClientUtils.getJobData(jobsService, allocator, jobId, offset, fetchLimit)) {
      if (data.getReturnedRowCount() <= 0) {
        break;
      }
      for (int i = 0; i < data.getReturnedRowCount(); i++) {
        byte[] b = (byte[]) data.extractValue(RecordWriter.METADATA_COLUMN, i);
        if(b == null) {
          throw new IllegalStateException("Didn't find metadata output for job " + jobId.getId());
        }
        updateIdWrapper.update(UpdateIdWrapper.deserialize(b));
      }
      offset += data.getReturnedRowCount();
    }
  }
  return updateIdWrapper.getUpdateId();
}
 

public PDFSService(
    Provider<FabricService> fabricService,
    Provider<NodeEndpoint> identityProvider,
    Provider<Iterable<NodeEndpoint>> nodeProvider,
    Tracer tracer,
    SabotConfig config,
    BufferAllocator allocator,
    PDFSMode mode) {
  this.fabricService = fabricService;
  this.identityProvider = identityProvider;
  this.nodeProvider = nodeProvider;
  this.tracer = tracer;
  this.config = config;
  this.allocator = allocator.newChildAllocator("pdfs-allocator", 0, Long.MAX_VALUE);
  this.allowLocalAccess = mode == PDFSMode.DATA;
}
 

/**
 * Creates a phase ticket (along with a phase-level allocator) for a given phase (major fragment) of this query, if
 * one has not already been created. The created phase ticket is tracked by this query ticket.
 *
 * Multi-thread safe`
 */
public PhaseTicket getOrCreatePhaseTicket(int majorFragmentId, long maxAllocation) {
  PhaseTicket phaseTicket = phaseTickets.get(majorFragmentId);
  if (phaseTicket == null) {
    final BufferAllocator phaseAllocator = getAllocator().newChildAllocator("phase-" + majorFragmentId, 0, maxAllocation);
    phaseTicket = new PhaseTicket(this, majorFragmentId, phaseAllocator);
    PhaseTicket insertedTicket = phaseTickets.putIfAbsent(majorFragmentId, phaseTicket);
    if (insertedTicket == null) {
      this.reserve();
    } else {
      // Race condition: another user managed to insert a phase ticket. Let's close ours and use theirs
      Preconditions.checkState(insertedTicket != phaseTicket);
      try {
        AutoCloseables.close(phaseTicket);  // NB: closing the ticket will close the phaseAllocator
      } catch (Exception e) {
        // Ignored
      }
      phaseTicket = insertedTicket;
    }
  }
  return phaseTicket;
}
 

/**
 * Create a new exec service. Note that at start time, the provider to one of
 * the handlers may be a noop implementation. This is allowed if this is a
 * single role node.
 *  @param config
 * @param allocator
 * @param fabricService
 * @param executorService
 * @param execResults
 * @param execStatus
 */
public CoordExecService(
  SabotConfig config,
  BufferAllocator allocator,
  Provider<FabricService> fabricService,
  Provider<ExecutorService> executorService,
  Provider<ExecToCoordResultsHandler> execResults,
  Provider<ExecToCoordStatusHandler> execStatus,
  Provider<CoordinationProtos.NodeEndpoint> selfEndpoint,
  Provider<JobTelemetryClient> jobTelemetryClient) {

  super();

  this.fabricService = fabricService;
  this.allocator =  allocator.newChildAllocator("coord-exec-rpc",
          config.getLong("dremio.exec.rpc.bit.server.memory.data.reservation"),
          config.getLong("dremio.exec.rpc.bit.server.memory.data.maximum"));
  this.executorService = executorService;
  this.execResults = execResults;
  this.execStatus = execStatus;
  this.selfEndpoint = selfEndpoint;
  this.jobTelemetryClient = jobTelemetryClient;
  this.config = getMapping(config);
}
 

/**
 * Get a set of transferred clones of this container. Note that this guarantees that the vectors in the cloned
 * container have the same TypedFieldIds as the existing container, allowing interchangeability in generated code. In
 * the case of hyper vectors, this container actually doesn't do a full transfer, rather creating a clone vector
 * wrapper only.
 *
 * @param incoming
 *          The RecordBatch iterator the contains the batch we should take over.
 * @return A cloned vector container.
 */
@SuppressWarnings("rawtypes")
public static VectorContainer getTransferClone(VectorAccessible incoming, BufferAllocator allocator) {
  final VectorContainer vc = new VectorContainer(allocator);
  for (VectorWrapper<?> w : incoming) {
    vc.cloneAndTransfer(w);
  }
  vc.setRecordCount(incoming.getRecordCount());
  vc.buildSchema(SelectionVectorMode.NONE);
  return vc;
}
 

@Test
public void T_convert_1() throws IOException{
  BufferAllocator allocator = new RootAllocator( 1024 * 1024 * 10 );
  TinyIntVector vector = new TinyIntVector( "test" , allocator );
  vector.allocateNew();
  vector.setSafe( 0 , (byte)0 );  
  vector.setSafe( 1 , (byte)1 );  
  vector.setSafe( 2 , (byte)0 );  
  vector.setNull( 3 );  
  vector.setSafe( 4 , (byte)1 );  
  vector.setSafe( 5 , (byte)1 );  
  vector.setSafe( 6 , (byte)1 );  
  vector.setNull( 7 );  
  vector.setValueCount( 8 );

  IColumn column = ArrowColumnFactory.convert( "test" , vector );
  assertEquals( column.getColumnName() , "test" );
  assertEquals( column.size() , 8 );
  assertTrue( ( column.getColumnType() == ColumnType.BYTE ) );
  assertEquals( ( (PrimitiveObject)( column.get(0).getRow() ) ).getByte() , (byte)0  );
  assertEquals( ( (PrimitiveObject)( column.get(1).getRow() ) ).getByte() , (byte)1  );
  assertEquals( ( (PrimitiveObject)( column.get(2).getRow() ) ).getByte() , (byte)0  );
  assertEquals( column.get(3).getRow() , null  );
  assertEquals( ( (PrimitiveObject)( column.get(4).getRow() ) ).getByte() , (byte)1 );
  assertEquals( ( (PrimitiveObject)( column.get(5).getRow() ) ).getByte() , (byte)1 );
  assertEquals( ( (PrimitiveObject)( column.get(6).getRow() ) ).getByte() , (byte)1 );
  assertEquals( column.get(7).getRow() , null  );
}
 

/**
 * Create a new untitled dataset, and load preview data.
 *
 * @param from Source from where the dataset is created (can be a query or other dataset)
 * @param version Initial version of the new dataset
 * @param context Dataset context or current schema
 * @return
 * @throws DatasetNotFoundException
 * @throws DatasetVersionNotFoundException
 * @throws NamespaceException
 */
public InitialPreviewResponse newUntitled(
  BufferAllocator allocator,
  FromBase from,
  DatasetVersion version,
  List<String> context,
  DatasetSummary parentSummary,
  boolean prepare,
  Integer limit)
    throws DatasetNotFoundException, DatasetVersionNotFoundException, NamespaceException, NewDatasetQueryException {
  return newUntitled(allocator, from, version, context, parentSummary, prepare, limit, false);
}
 

@Test
public void writeAndReadEmptyListVectors() throws Exception {
  try (final BufferAllocator allocator = allocatorRule.newAllocator("test-arrow-file-reader", 0, Long.MAX_VALUE);
       final VectorContainer batchData = createBatch(1, testEmptyListVector(allocator))) {

    final Path basePath = new Path(dateGenFolder.getRoot().getPath());
    final ArrowFileMetadata metadata = writeArrowFile(batchData);
    try (final ArrowFileReader reader =
             new ArrowFileReader(HadoopFileSystem.getLocal(FS_CONF), com.dremio.io.file.Path.of(basePath.toUri()), metadata, allocator)) {

      final List<RecordBatchHolder> batchHolders = reader.read(0, 1);
      assertEquals(1, batchHolders.size());
      assertNotNull(batchHolders.get(0).getData());
      assertEquals(0, batchHolders.get(0).getStart());
      assertEquals(1, batchHolders.get(0).getEnd());

      final BatchSchema schema = batchHolders.get(0).getData().getContainer().getSchema();
      assertEquals(1, schema.getFieldCount());
      assertEquals("emptyListVector", schema.getColumn(0).getName());
      assertEquals(MinorType.LIST, Types.getMinorTypeForArrowType(schema.getColumn(0).getType()));

      final VectorContainer batchContainer = batchHolders.get(0).getData().getContainer();
      assertTrue(Iterators.size(batchContainer.iterator()) == 1);
      for (final VectorWrapper<?> wrapper : batchContainer) {
        assertTrue(wrapper.getValueVector() instanceof ListVector);
        assertTrue(((ListVector) (wrapper.getValueVector())).getDataVector() instanceof NullVector);
      }

      releaseBatches(batchHolders);
    }
  }
}
 

/**
 * Set the vector of Array and initialize it.
 */
public ArrowFixedSchemaArrayMemoryAllocator(
    final ArrayContainerField schema ,
    final BufferAllocator allocator ,
    final ListVector vector , final int rowCount ) {
  this.allocator = allocator;
  this.vector = vector;
  vector.allocateNew();
  childSchema = schema.getField();
}
 

@Test
public void T_setString_2() throws IOException{
  IColumn column = new PrimitiveColumn( ColumnType.STRING , "boolean" );
  column.add( ColumnType.STRING , new StringObj( "a" ) , 0 );
  column.add( ColumnType.STRING , new StringObj( "b" ) , 1 );
  column.add( ColumnType.STRING , new StringObj( "c" ) , 5 );

  ColumnBinaryMakerConfig defaultConfig = new ColumnBinaryMakerConfig();
  ColumnBinaryMakerCustomConfigNode configNode = new ColumnBinaryMakerCustomConfigNode( "root" , defaultConfig );

  IColumnBinaryMaker maker = new UnsafeOptimizeDumpStringColumnBinaryMaker();
  ColumnBinary columnBinary = maker.toBinary( defaultConfig , null , new CompressResultNode() , column );

  BufferAllocator allocator = new RootAllocator( 1024 * 1024 * 10 );
  SchemaChangeCallBack callBack = new SchemaChangeCallBack();
  StructVector parent = new StructVector("root", allocator, new FieldType(false, Struct.INSTANCE, null, null), callBack);
  parent.allocateNew();
  IMemoryAllocator memoryAllocator = ArrowMemoryAllocatorFactory.getFromStructVector( ColumnType.STRING , "target" , allocator , parent , 3 );

  maker.loadInMemoryStorage( columnBinary , memoryAllocator );

  StructReader rootReader = parent.getReader();
  FieldReader reader = rootReader.reader( "target" );
  reader.setPosition( 0 );
  assertEquals( reader.readText().toString() , "a" );
  reader.setPosition( 1 );
  assertEquals( reader.readText().toString() , "b" );
  reader.setPosition( 5 );
  assertEquals( reader.readText().toString() , "c" );
  reader.setPosition( 2 );
  assertEquals( reader.readText() , null );
  reader.setPosition( 3 );
  assertEquals( reader.readText() , null );
  reader.setPosition( 4 );
  assertEquals( reader.readText() , null );
}
 

private BigIntSumAccumulator(final FieldVector input, final FieldVector output,
                             final FieldVector transferVector, final int maxValuesPerBatch,
                             final BufferAllocator computationVectorAllocator,
                             final long[] bitAddresses, final long[] valueAddresses,
                             final FieldVector[] accumulators) {
  super(input, output, transferVector,
        AccumulatorBuilder.AccumulatorType.SUM,
        maxValuesPerBatch, computationVectorAllocator,
        bitAddresses,
        valueAddresses,
        accumulators);
}
 

private static CompleteType getLiteralSchema(QueryContext context, byte[] bytes) {
  try(
      BufferAllocator allocator = context.getAllocator().newChildAllocator("convert-from-json-sampling", 0, 1024*1024);
      BufferManager bufferManager = new BufferManagerImpl(allocator);
      ArrowBuf data = allocator.buffer(bytes.length);
      VectorContainer container = new VectorContainer(allocator);
      VectorAccessibleComplexWriter vc = new VectorAccessibleComplexWriter(container)

      ){
    data.writeBytes(bytes);
    final int sizeLimit = Math.toIntExact(context.getOptions().getOption(ExecConstants.LIMIT_FIELD_SIZE_BYTES));
    final int maxLeafLimit = Math.toIntExact(context.getOptions().getOption(CatalogOptions.METADATA_LEAF_COLUMN_MAX));
    JsonReader jsonReader = new JsonReader(bufferManager.getManagedBuffer(), sizeLimit, maxLeafLimit,
      context.getOptions().getOption(ExecConstants.JSON_READER_ALL_TEXT_MODE_VALIDATOR), false, false);
    jsonReader.setSource(bytes);

      ComplexWriter writer = new ComplexWriterImpl("dummy", vc);
      writer.setPosition(0);
      ReadState state = jsonReader.write(writer);
      if(state == ReadState.END_OF_STREAM){
        throw new EOFException("Unexpected arrival at end of JSON literal stream");
      }

    container.buildSchema();
    return CompleteType.fromField(container.getSchema().getFields().get(0));
  }catch(Exception ex){
    throw UserException.validationError(ex).message("Failure while trying to parse JSON literal.").build(logger);
  }
}
 

public void allocate(BufferAllocator allocator) throws Exception {
  try (RollbackCloseable rbc = new RollbackCloseable()){
    probeOffsets = rbc.add(allocator.buffer(maxOutputRange * PROBE_OUTPUT_SIZE));
    buildOffsets = rbc.add(allocator.buffer(maxOutputRange * BUILD_OUTPUT_SIZE));
    rbc.commit();
  }
}
 

@Test
public void T_setShort_1() throws IOException{
  BufferAllocator allocator = new RootAllocator( 1024 * 1024 * 10 );
  SchemaChangeCallBack callBack = new SchemaChangeCallBack();
  StructVector parent = new StructVector("root", allocator, new FieldType(false, Struct.INSTANCE, null, null), callBack);
  parent.allocateNew();
  IMemoryAllocator memoryAllocator = ArrowMemoryAllocatorFactory.getFromStructVector( ColumnType.SHORT , "target" , allocator , parent , 1001 );

  memoryAllocator.setShort( 0 , (short)100 );
  memoryAllocator.setShort( 1 , (short)200 );
  memoryAllocator.setShort( 5 , (short)255 );
  memoryAllocator.setShort( 1000 , (short)10 );

  StructReader rootReader = parent.getReader();
  FieldReader reader = rootReader.reader( "target" );
  reader.setPosition( 0 );
  assertEquals( (short)( reader.readShort() ) , (short)100 );
  reader.setPosition( 1 );
  assertEquals( (short)( reader.readShort() ) , (short)200 );
  reader.setPosition( 5 );
  assertEquals( (short)( reader.readShort() ) , (short)255 );
  for( int i = 6 ; i < 1000 ; i++ ){
    reader.setPosition( i );
    assertEquals( reader.readShort() , null );
  }
  reader.setPosition( 1000 );
  assertEquals( (short)( reader.readShort() ) , (short)10 );
}
 

public PartitionToLoadSpilledData(final BufferAllocator allocator,
                                  final int fixedDataLength,
                                  final int variableDataLength,
                                  final List<Field> postSpillAccumulatorVectorTypes,
                                  final int batchSize) throws Exception {
  Preconditions.checkArgument(allocator != null, "Error: need a valid allocator to pre-allocate memory");
  this.allocator = allocator;
  /* we use Numbers.nextPowerOfTwo because that is how memory allocation happens
   * inside FixedBlockVector and VariableBlockVector when inserting into hashtable.
   * if we don't use nextPowerOfTwo for actual allocation size, we might run into
   * situation where a spilled batch has more data than the capacity we have pre-allocated
   * here to load the spilled batch into memory.
   */
  try(AutoCloseables.RollbackCloseable rollbackable = new AutoCloseables.RollbackCloseable()) {
    fixedKeyColPivotedData = allocator.buffer(Numbers.nextPowerOfTwo(fixedDataLength));
    rollbackable.add(fixedKeyColPivotedData);
    variableKeyColPivotedData = allocator.buffer(Numbers.nextPowerOfTwo(variableDataLength));
    rollbackable.add(variableKeyColPivotedData);
    this.postSpillAccumulatorVectors = new FieldVector[postSpillAccumulatorVectorTypes.size()];
    this.fixedDataLength = fixedDataLength;
    this.variableDataLength = variableDataLength;
    this.batchSize = batchSize;
    this.recordsInBatch = 0;
    this.accumulatorTypes = new byte[postSpillAccumulatorVectorTypes.size()];
    initBuffers();
    initPostSpillAccumulatorVectors(postSpillAccumulatorVectorTypes, batchSize, rollbackable);
    rollbackable.commit();
    logger.debug("Extra Partition Pre-allocation, fixed-data length: {}, variable-data length: {}, actual fixed-data capacity: {}, actual variable-data capacty: {}, batchSize: {}",
                 fixedDataLength, variableDataLength, fixedKeyColPivotedData.capacity(), variableKeyColPivotedData.capacity(), batchSize);
  }
}
 

ParquetResultListener(BufferAllocator allocator, ParquetTestProperties props,
    int numberOfTimesRead, boolean testValues) {
  this.allocator = allocator;
  this.props = props;
  this.totalRecords = props.recordsPerRowGroup * props.numberRowGroups * numberOfTimesRead;
  this.testValues = testValues;
}
 

public List<Card<ExtractMapRule>> recommendExtractMap(MapSelection selection, SqlQuery datasetSql, BufferAllocator allocator)
    throws DatasetVersionNotFoundException {

  List<TransformRuleWrapper<ExtractMapRule>> ruleWrappers = getRuleWrappers(extractMap, selection, DataType.MAP);
  return cardGenerator.generateCards(datasetSql, selection.getColName(), ruleWrappers,
      Recommenders.<ExtractMapRule>genericComparator(), allocator);
}
 

@Override
public IMemoryAllocator create(
    final BufferAllocator allocator ,
    final StructVector rootVector ,
    final int rowCount ) throws IOException {
  return new ArrowMapMemoryAllocator( allocator , rootVector , rowCount );
}
 

/**
 * Convert a set of input strings to arrow columns
 * for a time series.
 * @param bufferAllocator the buffer allocator to use
 * @param schema the schema to use
 * @param dataVecRecord the collection of input strings to process
 * @return the created vectors
 */
public static <T>  List<FieldVector> toArrowColumnsTimeSeriesHelper(final BufferAllocator bufferAllocator,
                                                                    final Schema schema,
                                                                    List<List<List<T>>> dataVecRecord) {
    //time series length * number of columns
    int numRows = 0;
    for(List<List<T>> timeStep : dataVecRecord) {
        numRows += timeStep.get(0).size() * timeStep.size();
    }

    numRows /= schema.numColumns();


    List<FieldVector> ret = createFieldVectors(bufferAllocator,schema,numRows);
    Map<Integer,Integer> currIndex = new HashMap<>(ret.size());
    for(int i = 0; i < ret.size(); i++) {
        currIndex.put(i,0);
    }
    for(int i = 0; i < dataVecRecord.size(); i++) {
        List<List<T>> record = dataVecRecord.get(i);
        for(int j = 0; j < record.size(); j++) {
            List<T> curr = record.get(j);
            for(int k = 0; k < curr.size(); k++) {
                Integer idx = currIndex.get(k);
                FieldVector fieldVector = ret.get(k);
                T writable = curr.get(k);
                setValue(schema.getType(k), fieldVector, writable, idx);
                currIndex.put(k,idx + 1);
            }
        }
    }

    return ret;
}
 

@Test
public void T_convert_1() throws IOException{
  BufferAllocator allocator = new RootAllocator( 1024 * 1024 * 10 );
  IntVector vector = new IntVector( "test" , allocator );
  vector.allocateNew();
  vector.setSafe( 0 , 0 );
  vector.setSafe( 1 , 1 );
  vector.setSafe( 2 , 0 );
  vector.setNull( 3 );
  vector.setSafe( 4 , 1 );
  vector.setSafe( 5 , 1 );
  vector.setSafe( 6 , 1 );
  vector.setNull( 7 );
  vector.setValueCount( 8 );

  IColumn column = ArrowColumnFactory.convert( "test" , vector );
  assertEquals( column.getColumnName() , "test" );
  assertEquals( column.size() , 8 );
  assertTrue( ( column.getColumnType() == ColumnType.INTEGER ) );
  assertEquals( ( (PrimitiveObject)( column.get(0).getRow() ) ).getInt() , 0  );
  assertEquals( ( (PrimitiveObject)( column.get(1).getRow() ) ).getInt() , 1  );
  assertEquals( ( (PrimitiveObject)( column.get(2).getRow() ) ).getInt() , 0  );
  assertEquals( column.get(3).getRow() , null  );
  assertEquals( ( (PrimitiveObject)( column.get(4).getRow() ) ).getInt() , 1 );
  assertEquals( ( (PrimitiveObject)( column.get(5).getRow() ) ).getInt() , 1 );
  assertEquals( ( (PrimitiveObject)( column.get(6).getRow() ) ).getInt() , 1 );
  assertEquals( column.get(7).getRow() , null  );
}
 

public FragmentStatusReporter(
    final FragmentHandle handle,
    final FragmentStats stats,
    final MaestroProxy maestroProxy,
    final BufferAllocator fragmentAllocator) {
  this.handle = handle;
  this.stats = stats;
  this.maestroProxy = maestroProxy;
  this.fragmentAllocator = fragmentAllocator;
}
 

private void createBigSchemaAndInputContainer(final BufferAllocator allocator) {
  final BatchSchema schema = BatchSchema.newBuilder()
    .addField(DECIMAL_KEY)
    .addField(DECIMAL_MEASURE)
    .build();

  container = VectorContainer.create(allocator, schema);

  decimalKey = container.addOrGet(DECIMAL_KEY);

  decimalMeasure = container.addOrGet(DECIMAL_MEASURE);
}
 

@Test
public void T_convert_1() throws IOException{
  BufferAllocator allocator = new RootAllocator( 1024 * 1024 * 10 );
  Float4Vector vector = new Float4Vector( "test" , allocator );
  vector.allocateNew();
  vector.setSafe( 0 , (float)0 ); 
  vector.setSafe( 1 , (float)1 ); 
  vector.setSafe( 2 , (float)0 ); 
  vector.setNull( 3 );
  vector.setSafe( 4 , (float)1 );
  vector.setSafe( 5 , (float)1 );
  vector.setSafe( 6 , (float)1 );
  vector.setNull( 7 );
  vector.setValueCount( 8 );

  IColumn column = ArrowColumnFactory.convert( "test" , vector );
  assertEquals( column.getColumnName() , "test" );
  assertEquals( column.size() , 8 );
  assertTrue( ( column.getColumnType() == ColumnType.FLOAT ) );
  assertEquals( ( (PrimitiveObject)( column.get(0).getRow() ) ).getFloat() , (float)0  );
  assertEquals( ( (PrimitiveObject)( column.get(1).getRow() ) ).getFloat() , (float)1  );
  assertEquals( ( (PrimitiveObject)( column.get(2).getRow() ) ).getFloat() , (float)0  );
  assertEquals( column.get(3).getRow() , null  );
  assertEquals( ( (PrimitiveObject)( column.get(4).getRow() ) ).getFloat() , (float)1 );
  assertEquals( ( (PrimitiveObject)( column.get(5).getRow() ) ).getFloat() , (float)1 );
  assertEquals( ( (PrimitiveObject)( column.get(6).getRow() ) ).getFloat() , (float)1 );
  assertEquals( column.get(7).getRow() , null  );
}
 

public void addBatch(RecordBatchData data, BufferAllocator copyTargetAllocator) throws SchemaChangeException {
  // We now generate an sv2 for the local sort. We do this even if the incoming
  // batch has an sv2. This is because we need to treat that one as immutable.
  //
  // Note that we shouldn't have an issue with allocation here since we'll use
  // the copyTargetAllocator provided by the caller (it is used for the final
  // sort but it isn't yet used and is guaranteed to be larger than the size of
  // this // ephemeral allocation).
  try (SelectionVector2 localSortVector = new SelectionVector2(copyTargetAllocator)) {

    final int recordCount = data.getRecordCount();
    localSortVector.allocateNew(recordCount);
    final SelectionVector2 incomingSv2 = data.getSv2();
    if (incomingSv2 != null) {
      // just copy the sv2.
      NettyArrowBuf buffer = localSortVector.getBuffer(false).asNettyBuffer();
      buffer.arrowBuf().setBytes(buffer.writerIndex(), incomingSv2.getBuffer(false), 0,
        recordCount * 2);
    } else {
      for (int i = 0; i < recordCount; i++) {
        localSortVector.setIndex(i * 2, i);
      }
    }

    // Quicksort for cache-local and SplayTree performance benefits (includes resetting vector references)
    localSorter.setup(classProducer.getFunctionContext(), localSortVector, data.getContainer());
    localSorter.sort(localSortVector);

    // now we need to insert the values into the splay tree.
    treeManager.add(localSortVector, data);
  }
}
 

/**
 *
 * @param allocator
 * @param name
 * @param data
 * @return
 */
public static IntVector vectorFor(BufferAllocator allocator,String name,int[] data) {
    IntVector float8Vector = new IntVector(name,FieldType.nullable(new ArrowType.Int(32,true)),allocator);
    float8Vector.allocateNew(data.length);
    for(int i = 0; i < data.length; i++) {
        float8Vector.setSafe(i,data[i]);
    }

    float8Vector.setValueCount(data.length);

    return float8Vector;
}
 

/**
 * Creates a wrapper around batch and sv2 for writing to a stream. sv2 will never be released by this class, and ownership
 * is maintained by caller. Also indicates whether compression is to be used for spilling the batch.
 * @param batch
 * @param sv2
 * @param allocator
 */
public VectorAccessibleSerializable(WritableBatch batch, SelectionVector2 sv2, BufferAllocator allocator, boolean useCodec) {
  this.allocator = allocator;
  this.batch = batch;
  if (sv2 != null) {
    this.sv2 = sv2;
    svMode = BatchSchema.SelectionVectorMode.TWO_BYTE;
  }
  this.useCodec = useCodec;
}
 

/**
 * Helper method to create {@link InitialPreviewResponse} for existing dataset.
 * @param newDataset
 * @param tipVersion - a top level history version for a dataset.
 * @return
 * @throws DatasetVersionNotFoundException
 */
InitialPreviewResponse createPreviewResponseForExistingDataset (
    BufferAllocator allocator,
    VirtualDatasetUI newDataset,
    DatasetVersionResourcePath tipVersion,
    Integer limit
    ) throws DatasetVersionNotFoundException, NamespaceException, JobNotFoundException {

  SqlQuery query = new SqlQuery(newDataset.getSql(), newDataset.getState().getContextList(), username());
  JobData jobData = executor.runQueryWithListener(query, QueryType.UI_PREVIEW, tipVersion.getDataset(), newDataset.getVersion(), JobStatusListener.NO_OP);

  return createPreviewResponse(newDataset, jobData, tipVersion, allocator, limit, true);
}
 

@Test
public void T_setDouble_2() throws IOException{
  IColumn column = new PrimitiveColumn( ColumnType.DOUBLE , "boolean" );
  column.add( ColumnType.DOUBLE , new DoubleObj( (double)100 ) , 0 );
  column.add( ColumnType.DOUBLE , new DoubleObj( (double)200 ) , 1 );
  column.add( ColumnType.DOUBLE , new DoubleObj( (double)255 ) , 5 );

  ColumnBinaryMakerConfig defaultConfig = new ColumnBinaryMakerConfig();
  ColumnBinaryMakerCustomConfigNode configNode = new ColumnBinaryMakerCustomConfigNode( "root" , defaultConfig );

  IColumnBinaryMaker maker = new OptimizeDoubleColumnBinaryMaker();
  ColumnBinary columnBinary = maker.toBinary( defaultConfig , null , column );

  BufferAllocator allocator = new RootAllocator( 1024 * 1024 * 10 );
  SchemaChangeCallBack callBack = new SchemaChangeCallBack();
  StructVector parent = new StructVector("root", allocator, new FieldType(false, Struct.INSTANCE, null, null), callBack);
  parent.allocateNew();
  IMemoryAllocator memoryAllocator = ArrowMemoryAllocatorFactory.getFromStructVector( ColumnType.DOUBLE , "target" , allocator , parent , 3 );

  maker.loadInMemoryStorage( columnBinary , memoryAllocator );

  StructReader rootReader = parent.getReader();
  FieldReader reader = rootReader.reader( "target" );
  reader.setPosition( 0 );
  assertEquals( reader.readDouble().doubleValue() , (double)100 );
  reader.setPosition( 1 );
  assertEquals( reader.readDouble().doubleValue() , (double)200 );
  reader.setPosition( 5 );
  assertEquals( reader.readDouble().doubleValue() , (double)255 );
  reader.setPosition( 2 );
  assertEquals( reader.readDouble() , null );
  reader.setPosition( 3 );
  assertEquals( reader.readDouble() , null );
  reader.setPosition( 4 );
  assertEquals( reader.readDouble() , null );
}