VectorSchemaRoot transform()

in flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/VectorSchemaRootTransformer.java [36:166]

• 76 lines of code
• 6 McCabe index (conditional complexity)

  VectorSchemaRoot transform(VectorSchemaRoot originalRoot, VectorSchemaRoot transformedRoot)
      throws Exception;

  /** Transformer's helper class; builds a new {@link VectorSchemaRoot}. */
  class Builder {

    private final Schema schema;
    private final BufferAllocator bufferAllocator;
    private final List<Field> newFields = new ArrayList<>();
    private final Collection<Task> tasks = new ArrayList<>();

    /**
     * Constructor for the VectorSchemaRootTransformer's Builder.
     *
     * @param schema The Arrow schema.
     * @param bufferAllocator The BufferAllocator to use for allocating memory.
     */
    public Builder(final Schema schema, final BufferAllocator bufferAllocator) {
      this.schema = schema;
      this.bufferAllocator =
          bufferAllocator.newChildAllocator(
              "VectorSchemaRootTransformer", 0, bufferAllocator.getLimit());
    }

    /**
     * Add task to transform a vector to a new vector renaming it. This also adds
     * transformedVectorName to the transformed {@link VectorSchemaRoot} schema.
     *
     * @param originalVectorName Name of the original vector to be transformed.
     * @param transformedVectorName Name of the vector that is the result of the transformation.
     * @return a VectorSchemaRoot instance with a task to rename a field vector.
     */
    public Builder renameFieldVector(
        final String originalVectorName, final String transformedVectorName) {
      tasks.add(
          (originalRoot, transformedRoot) -> {
            final FieldVector originalVector = originalRoot.getVector(originalVectorName);
            final FieldVector transformedVector = transformedRoot.getVector(transformedVectorName);

            final ArrowType originalType = originalVector.getField().getType();
            final ArrowType transformedType = transformedVector.getField().getType();
            if (!originalType.equals(transformedType)) {
              throw new IllegalArgumentException(
                  String.format(
                      "Cannot transfer vector with field type %s to %s",
                      originalType, transformedType));
            }

            if (originalVector instanceof BaseVariableWidthVector) {
              ((BaseVariableWidthVector) originalVector)
                  .transferTo(((BaseVariableWidthVector) transformedVector));
            } else if (originalVector instanceof BaseFixedWidthVector) {
              ((BaseFixedWidthVector) originalVector)
                  .transferTo(((BaseFixedWidthVector) transformedVector));
            } else {
              throw new IllegalStateException(
                  String.format("Cannot transfer vector of type %s", originalVector.getClass()));
            }
          });

      final Field originalField = schema.findField(originalVectorName);
      newFields.add(
          new Field(
              transformedVectorName,
              new FieldType(
                  originalField.isNullable(),
                  originalField.getType(),
                  originalField.getDictionary(),
                  originalField.getMetadata()),
              originalField.getChildren()));

      return this;
    }

    /**
     * Adds an empty field to the transformed {@link VectorSchemaRoot} schema.
     *
     * @param fieldName Name of the field to be added.
     * @param fieldType Type of the field to be added.
     * @return a VectorSchemaRoot instance with the current tasks.
     */
    public Builder addEmptyField(final String fieldName, final Types.MinorType fieldType) {
      newFields.add(Field.nullable(fieldName, fieldType.getType()));

      return this;
    }

    /**
     * Adds an empty field to the transformed {@link VectorSchemaRoot} schema.
     *
     * @param fieldName Name of the field to be added.
     * @param fieldType Type of the field to be added.
     * @return a VectorSchemaRoot instance with the current tasks.
     */
    public Builder addEmptyField(final String fieldName, final ArrowType fieldType) {
      newFields.add(Field.nullable(fieldName, fieldType));

      return this;
    }

    /**
     * Build the {@link VectorSchemaRoot} with applied transformation tasks.
     *
     * @return The built {@link VectorSchemaRoot}.
     */
    public VectorSchemaRootTransformer build() {
      return (originalRoot, transformedRoot) -> {
        if (transformedRoot == null) {
          transformedRoot = VectorSchemaRoot.create(new Schema(newFields), bufferAllocator);
        }

        for (final Task task : tasks) {
          task.run(originalRoot, transformedRoot);
        }

        transformedRoot.setRowCount(originalRoot.getRowCount());

        originalRoot.clear();
        return transformedRoot;
      };
    }

    /**
     * Functional interface used to a task to transform a VectorSchemaRoot into a new
     * VectorSchemaRoot.
     */
    @FunctionalInterface
    interface Task {
      void run(VectorSchemaRoot originalRoot, VectorSchemaRoot transformedRoot);
    }
  }