Java Code Examples for org.apache.arrow.vector.ValueVector#setValueCount()

protected void runProjector(int recordCount) {
  if (projector != null) {
    try {
      if (recordCount > 0) {
        splitter.projectRecords(recordCount, javaCodeGenWatch,
            gandivaCodeGenWatch);
      }
      javaCodeGenWatch.start();
      projector.projectRecords(recordCount);
      javaCodeGenWatch.stop();
    } catch (Exception e) {
      throw Throwables.propagate(e);
    }
    for (final ValueVector v : allocationVectors) {
      v.setValueCount(recordCount);
    }
  }
  OperatorStats stats = context.getStats();
  stats.addLongStat(ScanOperator.Metric.JAVA_EXECUTE_TIME, javaCodeGenWatch.elapsed(TimeUnit.NANOSECONDS));
  stats.addLongStat(ScanOperator.Metric.GANDIVA_EXECUTE_TIME, gandivaCodeGenWatch.elapsed(TimeUnit.NANOSECONDS));
  javaCodeGenWatch.reset();
  gandivaCodeGenWatch.reset();
}
 

public void runProjector(int recordCount) {
  if (projector != null) {
    try {
      if (recordCount > 0) {
        splitter.projectRecords(recordCount, javaCodeGenWatch,
          gandivaCodeGenWatch);
      }
      javaCodeGenWatch.start();
      projector.projectRecords(recordCount);
      javaCodeGenWatch.stop();
    } catch (Exception e) {
      throw Throwables.propagate(e);
    }
    for (final ValueVector v : allocationVectors) {
      v.setValueCount(recordCount);
    }
  }
  OperatorStats stats = context.getStats();
  stats.addLongStat(ScanOperator.Metric.JAVA_EXECUTE_TIME, javaCodeGenWatch.elapsed(TimeUnit.NANOSECONDS));
  stats.addLongStat(ScanOperator.Metric.GANDIVA_EXECUTE_TIME, gandivaCodeGenWatch.elapsed(TimeUnit.NANOSECONDS));
  javaCodeGenWatch.reset();
  gandivaCodeGenWatch.reset();
}
 

/**
 * Method is invoked when we have a straight aggregate (no group by expression) and our input is empty.
 * In this case we construct an outgoing batch with record count as 1. For the nullable vectors we don't set anything
 * as we want the output to be NULL. For the required vectors (only for count()) we set the value to be zero since
 * we don't zero out our buffers initially while allocating them.
 */
private void constructSpecialBatch() {
  outgoing.allocateNew();
  List<NamedExpression> exprs = config.getAggrExprs();
  if(outgoing.getNumberOfColumns() != exprs.size()){
    throw new IllegalStateException();
  }
  int exprIndex = 0;
  for (final VectorWrapper<?> vw: outgoing) {
    final ValueVector vv = vw.getValueVector();
    if (!exprs.isEmpty() && isCount(exprs.get(exprIndex))) {
      ((BigIntVector) vv).setSafe(0, 0);
    }
    vv.setValueCount(SPECIAL_BATCH_COUNT);
    exprIndex++;
  }
  outgoing.setRecordCount(SPECIAL_BATCH_COUNT);
}
 

private void setValueCount(int count) {
  for (ValueVector v : allocationVectors) {
    v.setValueCount(count);
  }

  if(flattenVector != null){
    flattenVector.setValueCount(count);
  }

  if (complexWriters != null) {
    for (ComplexWriter writer : complexWriters) {
      writer.setValueCount(count);
    }
  }


}
 

private void setValueCount(int numRecords) {
  for(ValueVector vv : allocationVectors) {
    vv.setValueCount(numRecords);
  }

  for (final ComplexWriter writer : complexWriters) {
    writer.setValueCount(numRecords);
  }
}
 

private void setValueCount(final int count) {
  for (final ValueVector v : allocationVectors) {
    v.setValueCount(count);
  }

  for (final ComplexWriter writer : complexWriters) {
    writer.setValueCount(count);
  }
}
 

private static VectorContainer createBatch(int recordCount, ValueVector... vv) {
  VectorContainer container = new VectorContainer();
  if (recordCount != 0) {
    for (ValueVector v : vv) {
      v.setValueCount(recordCount);
    }
  }
  container.addCollection(asList(vv));
  container.setRecordCount(recordCount);
  container.buildSchema(SelectionVectorMode.NONE);

  return container;
}
 

protected void ensureVectorHasRequiredCapacity(ValueVector vector, int required) {
  while (required > vector.getValueCapacity()) {
    vector.reAlloc();
  }
  // setValueCount helps in setting inputBytes job stat correctly
  vector.setValueCount(required);
}
 

protected void ensureVectorHasRequiredCapacity(ValueVector vector, int required) {
  while (required > vector.getValueCapacity()) {
    vector.reAlloc();
  }
  // setValueCount helps in setting inputBytes job stat correctly
  vector.setValueCount(required);
}
 

@Override
public int next() {
  int count = 0;
  long maxRecordCount = 0;
  try {
    // No columns found in the file were selected, simply return a full batch of null records for each column requested
    if (noColumnsFound) {
      if (mockRecordsRead == footer.getBlocks().get(rowGroupIndex).getRowCount()) {
        return 0;
      }
      long recordsToRead = 0;
      recordsToRead = Math.min(numRowsPerBatch, footer.getBlocks().get(rowGroupIndex).getRowCount() - mockRecordsRead);
      writer.setValueCount((int)recordsToRead);
      mockRecordsRead += recordsToRead;
      totalRead += recordsToRead;
      return (int) recordsToRead;
    }

    maxRecordCount = numRowsPerBatch;
    if (deltas != null) {
      maxRecordCount = deltas.getValueCount();
      vectorizedBasedFilter.reset();
    }
    while (count < maxRecordCount && totalRead < recordCount) {
      recordMaterializer.setPosition(count);
      recordReader.read();
      count++;
      totalRead++;
    }
    writer.setValueCount(count);
    // if we have requested columns that were not found in the file fill their vectors with null
    // (by simply setting the value counts inside of them, as they start null filled)
    if (nullFilledVectors != null) {
      for (final ValueVector vv : nullFilledVectors) {
        vv.setValueCount(count);
      }
    }
    return count;
  } catch (Throwable t) {
    throw UserException.dataWriteError(t)
        .message("Failed to read data from parquet file")
        .addContext("File path", path)
        .addContext("Rowgroup index", rowGroupIndex)
        .addContext("Delta vector present, size (if present)", deltas != null ? "yes, " + deltas.getValueCount() : "no")
        .addContext("Max no. of rows trying to read", maxRecordCount)
        .addContext("No. of rows read so far in current iteration", count)
        .addContext("No. of rows read so far in current rowgroup", totalRead)
        .addContext("Max no. rows in current rowgroup", recordCount)
        .addContext("Footer %s", footer)
        .build(logger);
  }
}
 

private void setValueCount() {
  for (VectorWrapper<?> vw : htContainer) {
    ValueVector vv = vw.getValueVector();
    vv.setValueCount(maxOccupiedIdx + 1);
  }
}
 

@Override
public void getStream(CallContext callContext, Ticket ticket, ServerStreamListener serverStreamListener) {
  /* Note that we do not trim record batches that we receive from the Job Results Store. This may result
   * in sending record that the client does not care about, or in the case of sequential requests, sending
   * duplicate records. We may want to trim the record batches if this presents a problem.
   */
  try {
    final JobsFlightTicket jobsFlightTicket = JobsFlightTicket.from(ticket);
    final JobProtobuf.JobId jobId = JobProtobuf.JobId.newBuilder().setId(jobsFlightTicket.getJobId()).build();
    final int offset = jobsFlightTicket.getOffset();
    final int limit = jobsFlightTicket.getLimit();

    try (final JobDataFragment jobDataFragment = jobsService.get().getJobData(JobsProtoUtil.toStuff(jobId), offset, limit)) {
      final Schema schema = jobDataFragment.getSchema();
      try (final VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) {
        serverStreamListener.start(root);
        for (RecordBatchHolder holder : jobDataFragment.getRecordBatches()) {
          // iterate over the columns
          int numRecords = holder.size();
          for (int i = 0; i < schema.getFields().size(); i++) {
            ValueVector vector = root.getVector(schema.getFields().get(i).getName());
            ValueVector dataVector = holder.getData().getVectors().get(i);
            int k = 0; // index at which value need to written in "vector" from "dataVector"
            // iterate over values in the column to copy data
            for (int j = holder.getStart(); j < holder.getEnd(); j++, k++ ) {
              // Copy value at dataVector[j] into vector[k]
              vector.copyFromSafe(j, k, dataVector);
            }
            vector.setValueCount(numRecords);
            root.setRowCount(numRecords);
          }
          serverStreamListener.putNext();
          root.allocateNew();
        }
      }
      serverStreamListener.completed();
    }
  } catch (UserException ue) {
    serverStreamListener.error(JobsRpcUtils.toStatusRuntimeException(ue));
  } catch (Exception e) {
    serverStreamListener.error(Status.UNKNOWN.withCause(e).withDescription(e.getMessage()).asException());
  }
}
 

public static void evaluate(int recordCount, FunctionContext functionContext, VectorAccessible incoming, ValueVector outVV, LogicalExpression expr) {

    InitVisitor initVisitor = new InitVisitor(functionContext);
    EvalVisitor evalVisitor = new EvalVisitor(incoming, functionContext);

    expr.accept(initVisitor, incoming);

    for (int i = 0; i < recordCount; i++) {
      ValueHolder out = expr.accept(evalVisitor, i);
      TypeHelper.setValueSafe(outVV, i, out);
    }

    outVV.setValueCount(recordCount);

  }