Java Code Examples for org.apache.arrow.vector.FieldVector#getMinorType()

The following examples show how to use org.apache.arrow.vector.FieldVector#getMinorType() . You can vote up the ones you like or vote down the ones you don't like, and go to the original project or source file by following the links above each example. You may check out the related API usage on the sidebar.
Example 1
Source File: SFArrowResultSetIT.java    From snowflake-jdbc with Apache License 2.0 5 votes vote down vote up
private File createArrowFile(String fileName, Schema schema, Object[][] data,
                             int rowsPerRecordBatch)
throws IOException
{
  File file = resultFolder.newFile(fileName);
  VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator);

  try (ArrowWriter writer = new ArrowStreamWriter(
      root, new DictionaryProvider.MapDictionaryProvider(),
      new FileOutputStream(file)))
  {
    writer.start();

    for (int i = 0; i < data[0].length; )
    {
      int rowsToAppend = Math.min(rowsPerRecordBatch, data[0].length - i);
      root.setRowCount(rowsToAppend);

      for (int j = 0; j < data.length; j++)
      {
        FieldVector vector = root.getFieldVectors().get(j);

        switch (vector.getMinorType())
        {
          case INT:
            writeIntToField(vector, data[j], i, rowsToAppend);
            break;
        }
      }

      writer.writeBatch();
      i += rowsToAppend;
    }
  }

  return file;
}
 
Example 2
Source File: BlockUtils.java    From aws-athena-query-federation with Apache License 2.0 4 votes vote down vote up
/**
 * Used to mark a particular cell as null.
 *
 * @param vector The FieldVector to write the null value to.
 * @param pos The position (row) in the FieldVector to mark as null.
 */
private static void setNullValue(FieldVector vector, int pos)
{
    switch (vector.getMinorType()) {
        case TIMESTAMPMILLITZ:
            ((TimeStampMilliTZVector) vector).setNull(pos);
            break;
        case DATEMILLI:
            ((DateMilliVector) vector).setNull(pos);
            break;
        case DATEDAY:
            ((DateDayVector) vector).setNull(pos);
            break;
        case FLOAT8:
            ((Float8Vector) vector).setNull(pos);
            break;
        case FLOAT4:
            ((Float4Vector) vector).setNull(pos);
            break;
        case INT:
            ((IntVector) vector).setNull(pos);
            break;
        case TINYINT:
            ((TinyIntVector) vector).setNull(pos);
            break;
        case SMALLINT:
            ((SmallIntVector) vector).setNull(pos);
            break;
        case UINT1:
            ((UInt1Vector) vector).setNull(pos);
            break;
        case UINT2:
            ((UInt2Vector) vector).setNull(pos);
            break;
        case UINT4:
            ((UInt4Vector) vector).setNull(pos);
            break;
        case UINT8:
            ((UInt8Vector) vector).setNull(pos);
            break;
        case BIGINT:
            ((BigIntVector) vector).setNull(pos);
            break;
        case VARBINARY:
            ((VarBinaryVector) vector).setNull(pos);
            break;
        case DECIMAL:
            ((DecimalVector) vector).setNull(pos);
            break;
        case VARCHAR:
            ((VarCharVector) vector).setNull(pos);
            break;
        case BIT:
            ((BitVector) vector).setNull(pos);
            break;
        default:
            throw new IllegalArgumentException("Unknown type " + vector.getMinorType());
    }
}
 
Example 3
Source File: BlockUtils.java    From aws-athena-query-federation with Apache License 2.0 4 votes vote down vote up
/**
 * In some filtering situations it can be useful to 'unset' a row as an indication to a later processing stage
 * that the row is irrelevant. The mechanism by which we 'unset' a row is actually field type specific and as such
 * this method is not supported for all field types.
 *
 * @param row The row number to unset in the provided Block.
 * @param block The Block where we'd like to unset the specified row.
 */
public static void unsetRow(int row, Block block)
{
    for (FieldVector vector : block.getFieldVectors()) {
        switch (vector.getMinorType()) {
            case TIMESTAMPMILLITZ:
                ((TimeStampMilliTZVector) vector).setNull(row);
                break;
            case DATEDAY:
                ((DateDayVector) vector).setNull(row);
                break;
            case DATEMILLI:
                ((DateMilliVector) vector).setNull(row);
                break;
            case TINYINT:
                ((TinyIntVector) vector).setNull(row);
                break;
            case UINT1:
                ((UInt1Vector) vector).setNull(row);
                break;
            case SMALLINT:
                ((SmallIntVector) vector).setNull(row);
                break;
            case UINT2:
                ((UInt2Vector) vector).setNull(row);
                break;
            case UINT4:
                ((UInt4Vector) vector).setNull(row);
                break;
            case INT:
                ((IntVector) vector).setNull(row);
                break;
            case UINT8:
                ((UInt8Vector) vector).setNull(row);
                break;
            case BIGINT:
                ((BigIntVector) vector).setNull(row);
                break;
            case FLOAT4:
                ((Float4Vector) vector).setNull(row);
                break;
            case FLOAT8:
                ((Float8Vector) vector).setNull(row);
                break;
            case DECIMAL:
                ((DecimalVector) vector).setNull(row);
                break;
            case VARBINARY:
                ((VarBinaryVector) vector).setNull(row);
                break;
            case VARCHAR:
                ((VarCharVector) vector).setNull(row);
                break;
            case BIT:
                ((BitVector) vector).setNull(row);
                break;
            case STRUCT:
                ((StructVector) vector).setNull(row);
                break;
            case LIST:
                UnionListWriter writer = ((ListVector) vector).getWriter();
                writer.setPosition(row);
                writer.startList();
                writer.endList();
                writer.setValueCount(0);
                break;
            default:
                throw new IllegalArgumentException("Unknown type " + vector.getMinorType());
        }
    }
}
 
Example 4
Source File: VectorizedHashAggPartition.java    From dremio-oss with Apache License 2.0 4 votes vote down vote up
private void updateSumAccumulator(final FieldVector deserializedAccumulator,
                                  final Accumulator[] partitionAccumulators,
                                  final int index,
                                  final BufferAllocator computationVectorAllocator) {
  final Accumulator partitionAccumulator = partitionAccumulators[index];
  Types.MinorType type = deserializedAccumulator.getMinorType();
  switch (type) {
    case BIGINT: {
      if (partitionAccumulator.getInput() instanceof BigIntVector) {
        /* We started with BigIntSumAccumulator so type doesn't change.
         * Just set the input vector to the accumulator vector read from
         * spilled batch
         */
        Preconditions.checkArgument(partitionAccumulator instanceof SumAccumulators.BigIntSumAccumulator,
                                    "Error: expecting bigint sum accumulator");
        partitionAccumulator.setInput(deserializedAccumulator);
      } else {
        /* We started with IntSumAccumulator that has input vector of type INT
         * and accumulator vector of type BIGINT and the output vector into which
         * we will eventually transfer contents also BIGINT. We spilled the
         * accumulator vector. So IntSumAccumulator now becomes BigIntSumAccumulator
         * for post-spill processing with the accumulator vector read from spilled batch
         * acting as new input vector for BigIntSumAccumulator
         */
        Preconditions.checkArgument(partitionAccumulator instanceof SumAccumulators.IntSumAccumulator,
                                    "Error: expecting int sum accumulator");
        partitionAccumulators[index] =
          new SumAccumulators.BigIntSumAccumulator((SumAccumulators.IntSumAccumulator)partitionAccumulator,
                                                   deserializedAccumulator,
                                                   hashTable.getActualValuesPerBatch(),
                                                   computationVectorAllocator);
      }

      break;
    }

    case FLOAT8: {
      if (partitionAccumulator.getInput() instanceof Float8Vector) {
        /* We started with DoubleSumAccumulator so type doesn't change.
         * Just set the input vector to the accumulator vector read from
         * spilled batch
         */
        Preconditions.checkArgument(partitionAccumulator instanceof SumAccumulators.DoubleSumAccumulator,
                                    "Error: expecting double sum accumulator");
        partitionAccumulator.setInput(deserializedAccumulator);
      } else if (partitionAccumulator.getInput() instanceof Float4Vector) {
        /* We started with FloatSumAccumulator that has input vector of type FLOAT
         * and accumulator vector of type DOUBLE and the output vector into which
         * we will eventually transfer contents also DOUBLE. We spilled the
         * accumulator vector. So FloatSumAccumulator now becomes DoubleSumAccumulator
         * for post-spill processing with the accumulator vector read from spilled batch
         * acting as new input vector for DoubleSumAccumulator
         */
        Preconditions.checkArgument(partitionAccumulator instanceof SumAccumulators.FloatSumAccumulator,
                                    "Error: expecting float sum accumulator");
        partitionAccumulators[index] =
          new SumAccumulators.DoubleSumAccumulator((SumAccumulators.FloatSumAccumulator)partitionAccumulator,
                                                   deserializedAccumulator,
                                                   hashTable.getActualValuesPerBatch(),
                                                   computationVectorAllocator);
      } else if (partitionAccumulator.getInput() instanceof DecimalVector) {
         /* We started with DecimalSumAccumulator that has input vector of type DECIMAL
         * and accumulator vector of type DOUBLE and the output vector into which
         * we will eventually transfer contents also DOUBLE. We spilled the
         * accumulator vector. So DecimalSumAccumulator now becomes DoubleSumAccumulator
         * for post-spill processing with the accumulator vector read from spilled batch
         * acting as new input vector for DoubleSumAccumulator
         */
         // ensure that if this happens decimal complete is turned off.
         Preconditions.checkArgument(decimalV2Enabled == false);
         Preconditions.checkArgument(partitionAccumulator instanceof SumAccumulators.DecimalSumAccumulator,
                                     "Error: expecting decimal sum accumulator");
        partitionAccumulators[index] =
          new SumAccumulators.DoubleSumAccumulator((SumAccumulators.DecimalSumAccumulator)partitionAccumulator,
                                                   deserializedAccumulator,
                                                   hashTable.getActualValuesPerBatch(),
                                                   computationVectorAllocator);
      }

      break;
    }

    case DECIMAL: {
      Preconditions.checkArgument(partitionAccumulator instanceof SumAccumulators.DecimalSumAccumulatorV2,
        "Error: expecting decimal sum accumulator");
      partitionAccumulator.setInput(deserializedAccumulator);
      break;
    }

    default: {
      /* we should not be here */
      throw new IllegalStateException("Error: incorrect type of deserialized sum accumulator");
    }
  }
}
 
Example 5
Source File: VectorizedHashAggPartition.java    From dremio-oss with Apache License 2.0 4 votes vote down vote up
private void updateSumZeroAccumulator(final FieldVector deserializedAccumulator,
                                      final Accumulator[] partitionAccumulators,
                                      final int index,
                                      final BufferAllocator computationVectorAllocator) {
  final Accumulator partitionAccumulator = partitionAccumulators[index];
  Types.MinorType type = deserializedAccumulator.getMinorType();
  switch (type) {
    case BIGINT: {
      if (partitionAccumulator.getInput() instanceof BigIntVector) {
        /* We started with BigIntSumZeroAccumulator so type doesn't change.
         * Just set the input vector to the accumulator vector read from
         * spilled batch
         */
        Preconditions.checkArgument(partitionAccumulator instanceof SumZeroAccumulators.BigIntSumZeroAccumulator,
                                    "Error: expecting bigint sumzero accumulator");
        partitionAccumulator.setInput(deserializedAccumulator);
      } else {
        /* We started with IntSumZeroAccumulator that has input vector of type INT
         * and accumulator vector of type BIGINT and the output vector into which
         * we will eventually transfer contents also BIGINT. We spilled the
         * accumulator vector. So IntSumZeroAccumulator now becomes BigIntSumZeroAccumulator
         * for post-spill processing with the accumulator vector read from spilled batch
         * acting as new input vector for BigIntSumZeroAccumulator
         */
        Preconditions.checkArgument(partitionAccumulator instanceof SumZeroAccumulators.IntSumZeroAccumulator,
                                    "Error: expecting int sumzero accumulator");
        partitionAccumulators[index] =
          new SumZeroAccumulators.BigIntSumZeroAccumulator((SumZeroAccumulators.IntSumZeroAccumulator)partitionAccumulator,
                                                           deserializedAccumulator,
                                                           hashTable.getActualValuesPerBatch(),
                                                           computationVectorAllocator);
      }

      break;
    }

    case FLOAT8: {
      if (partitionAccumulator.getInput() instanceof Float8Vector) {
        /* We started with DoubleSumZeroAccumulator so type doesn't change.
         * Just set the input vector to the accumulator vector read from
         * spilled batch
         */
        Preconditions.checkArgument(partitionAccumulator instanceof SumZeroAccumulators.DoubleSumZeroAccumulator,
                                    "Error: expecting double sum zero accumulator");
        partitionAccumulator.setInput(deserializedAccumulator);
      } else if (partitionAccumulator.getInput() instanceof Float4Vector) {
        /* We started with FloatSumZeroAccumulator that has input vector of type FLOAT
         * and accumulator vector of type DOUBLE and the output vector into which
         * we will eventually transfer contents also DOUBLE. We spilled the
         * accumulator vector. So FloatZeroAccumulator now becomes DoubleSumZeroAccumulator
         * for post-spill processing with the accumulator vector read from spilled batch
         * acting as new input vector for DoubleSumZeroAccumulator
         */
        Preconditions.checkArgument(partitionAccumulator instanceof SumZeroAccumulators.FloatSumZeroAccumulator,
                                    "Error: expecting float sum zero accumulator");
        partitionAccumulators[index] =
        new SumZeroAccumulators.DoubleSumZeroAccumulator((SumZeroAccumulators.FloatSumZeroAccumulator)partitionAccumulator,
                                                         deserializedAccumulator,
                                                         hashTable.getActualValuesPerBatch(),
                                                         computationVectorAllocator);
      } else if (partitionAccumulator.getInput() instanceof DecimalVector) {
         /* We started with DecimalSumZeroAccumulator that has input vector of type DECIMAL
         * and accumulator vector of type DOUBLE and the output vector into which
         * we will eventually transfer contents also DOUBLE. We spilled the
         * accumulator vector. So DecimalSumZeroAccumulator now becomes DoubleSumAccumulator
         * for post-spill processing with the accumulator vector read from spilled batch
         * acting as new input vector for DoubleSumAccumulator
         */
        // ensure that if this happens decimal complete is turned off.
        Preconditions.checkArgument(decimalV2Enabled == false);
        Preconditions.checkArgument(partitionAccumulator instanceof SumZeroAccumulators.DecimalSumZeroAccumulator,
                                    "Error: expecting decimal sum zero accumulator");
        partitionAccumulators[index] =
          new SumZeroAccumulators.DoubleSumZeroAccumulator((SumZeroAccumulators.DecimalSumZeroAccumulator) partitionAccumulator,
                                                           deserializedAccumulator,
                                                           hashTable.getActualValuesPerBatch(),
                                                           computationVectorAllocator);
      }
      break;
    }

    case DECIMAL: {
      Preconditions.checkArgument(partitionAccumulator instanceof SumZeroAccumulators.DecimalSumZeroAccumulatorV2,"Error: expecting decimal sum zero accumulator");
      partitionAccumulator.setInput(deserializedAccumulator);
      break;
    }

    default: {
      /* we should not be here */
      throw new IllegalStateException("Error: incorrect type of deserialized sumzero accumulator");
    }
  }
}