content: /*
   * Licensed to the Apache Software Foundation (ASF) under one or more
   * contributor license agreements.  See the NOTICE file distributed with
   * this work for additional information regarding copyright ownership.
   * The ASF licenses this file to You under the Apache License, Version 2.0
   * (the "License"); you may not use this file except in compliance with
   * the License.  You may obtain a copy of the License at
   *
   *    http://www.apache.org/licenses/LICENSE-2.0
   *
   * Unless required by applicable law or agreed to in writing, software
   * distributed under the License is distributed on an "AS IS" BASIS,
   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   * See the License for the specific language governing permissions and
   * limitations under the License.
   */
  package org.apache.arrow.gandiva.evaluator;
  
  import java.util.ArrayList;
  import java.util.List;
  import org.apache.arrow.gandiva.exceptions.EvaluatorClosedException;
  import org.apache.arrow.gandiva.exceptions.GandivaException;
  import org.apache.arrow.gandiva.exceptions.UnsupportedTypeException;
  import org.apache.arrow.gandiva.expression.ArrowTypeHelper;
  import org.apache.arrow.gandiva.expression.ExpressionTree;
  import org.apache.arrow.gandiva.ipc.GandivaTypes;
  import org.apache.arrow.gandiva.ipc.GandivaTypes.SelectionVectorType;
  import org.apache.arrow.memory.ArrowBuf;
  import org.apache.arrow.vector.BaseVariableWidthVector;
  import org.apache.arrow.vector.FixedWidthVector;
  import org.apache.arrow.vector.ValueVector;
  import org.apache.arrow.vector.VariableWidthVector;
  import org.apache.arrow.vector.ipc.message.ArrowBuffer;
  import org.apache.arrow.vector.ipc.message.ArrowRecordBatch;
  import org.apache.arrow.vector.types.pojo.Schema;
  
  /**
   * This class provides a mechanism to evaluate a set of expressions against a RecordBatch. Follow
   * these steps to use this class: 1) Use the static method make() to create an instance of this
   * class that evaluates a set of expressions 2) Invoke the method evaluate() to evaluate these
   * expressions against a RecordBatch 3) Invoke close() to release resources
   */
  public class Projector {
    private static final org.slf4j.Logger logger = org.slf4j.LoggerFactory.getLogger(Projector.class);
  
    private JniWrapper wrapper;
    private final long moduleId;
    private final Schema schema;
    private final int numExprs;
    private boolean closed;
  
    private Projector(JniWrapper wrapper, long moduleId, Schema schema, int numExprs) {
      this.wrapper = wrapper;
      this.moduleId = moduleId;
      this.schema = schema;
      this.numExprs = numExprs;
      this.closed = false;
    }
  
    /**
     * Invoke this function to generate LLVM code to evaluate the list of project expressions. Invoke
     * Projector::Evaluate() against a RecordBatch to evaluate the record batch against these
     * projections.
     *
     * @param schema Table schema. The field names in the schema should match the fields used to
     *     create the TreeNodes
     * @param exprs List of expressions to be evaluated against data
     * @return A native evaluator object that can be used to invoke these projections on a RecordBatch
     */
    public static Projector make(Schema schema, List<ExpressionTree> exprs) throws GandivaException {
      return make(schema, exprs, SelectionVectorType.SV_NONE, JniLoader.getDefaultConfiguration());
    }
  
    /**
     * Invoke this function to generate LLVM code to evaluate the list of project expressions. Invoke
     * Projector::Evaluate() against a RecordBatch to evaluate the record batch against these
     * projections.
     *
     * @param schema Table schema. The field names in the schema should match the fields used to
     *     create the TreeNodes
     * @param exprs List of expressions to be evaluated against data
     * @param configOptions ConfigOptions parameter
     * @return A native evaluator object that can be used to invoke these projections on a RecordBatch
     */
    public static Projector make(
        Schema schema, List<ExpressionTree> exprs, ConfigurationBuilder.ConfigOptions configOptions)
        throws GandivaException {
      return make(
          schema, exprs, SelectionVectorType.SV_NONE, JniLoader.getConfiguration(configOptions));
    }
  
    /**
     * Invoke this function to generate LLVM code to evaluate the list of project expressions. Invoke
     * Projector::Evaluate() against a RecordBatch to evaluate the record batch against these
     * projections.
     *
     * @param schema Table schema. The field names in the schema should match the fields used to
     *     create the TreeNodes
     * @param exprs List of expressions to be evaluated against data
     * @param optimize Flag to choose if the generated llvm code is to be optimized
     * @return A native evaluator object that can be used to invoke these projections on a RecordBatch
     */
    @Deprecated
    public static Projector make(Schema schema, List<ExpressionTree> exprs, boolean optimize)
        throws GandivaException {
      return make(
          schema,
          exprs,
          SelectionVectorType.SV_NONE,
          JniLoader.getConfiguration(
              (new ConfigurationBuilder.ConfigOptions()).withOptimize(optimize)));
    }
  
    /**
     * Invoke this function to generate LLVM code to evaluate the list of project expressions. Invoke
     * Projector::Evaluate() against a RecordBatch to evaluate the record batch against these
     * projections.
     *
     * @param schema Table schema. The field names in the schema should match the fields used to
     *     create the TreeNodes
     * @param exprs List of expressions to be evaluated against data
     * @param selectionVectorType type of selection vector
     * @return A native evaluator object that can be used to invoke these projections on a RecordBatch
     */
    public static Projector make(
        Schema schema, List<ExpressionTree> exprs, SelectionVectorType selectionVectorType)
        throws GandivaException {
      return make(schema, exprs, selectionVectorType, JniLoader.getDefaultConfiguration());
    }
  
    /**
     * Invoke this function to generate LLVM code to evaluate the list of project expressions. Invoke
     * Projector::Evaluate() against a RecordBatch to evaluate the record batch against these
     * projections.
     *
     * @param schema Table schema. The field names in the schema should match the fields used to
     *     create the TreeNodes
     * @param exprs List of expressions to be evaluated against data
     * @param selectionVectorType type of selection vector
     * @param configOptions ConfigOptions parameter
     * @return A native evaluator object that can be used to invoke these projections on a RecordBatch
     */
    public static Projector make(
        Schema schema,
        List<ExpressionTree> exprs,
        SelectionVectorType selectionVectorType,
        ConfigurationBuilder.ConfigOptions configOptions)
        throws GandivaException {
      return make(schema, exprs, selectionVectorType, JniLoader.getConfiguration(configOptions));
    }
  
    /**
     * Invoke this function to generate LLVM code to evaluate the list of project expressions. Invoke
     * Projector::Evaluate() against a RecordBatch to evaluate the record batch against these
     * projections.
     *
     * @param schema Table schema. The field names in the schema should match the fields used to
     *     create the TreeNodes
     * @param exprs List of expressions to be evaluated against data
     * @param selectionVectorType type of selection vector
     * @param optimize Flag to choose if the generated llvm code is to be optimized
     * @return A native evaluator object that can be used to invoke these projections on a RecordBatch
     */
    @Deprecated
    public static Projector make(
        Schema schema,
        List<ExpressionTree> exprs,
        SelectionVectorType selectionVectorType,
        boolean optimize)
        throws GandivaException {
      return make(
          schema,
          exprs,
          selectionVectorType,
          JniLoader.getConfiguration(
              (new ConfigurationBuilder.ConfigOptions()).withOptimize(optimize)));
    }
  
    /**
     * Invoke this function to generate LLVM code to evaluate the list of project expressions. Invoke
     * Projector::Evaluate() against a RecordBatch to evaluate the record batch against these
     * projections.
     *
     * @param schema Table schema. The field names in the schema should match the fields used to
     *     create the TreeNodes
     * @param exprs List of expressions to be evaluated against data
     * @param selectionVectorType type of selection vector
     * @param configurationId Custom configuration created through config builder.
     * @return A native evaluator object that can be used to invoke these projections on a RecordBatch
     */
    public static synchronized Projector make(
        Schema schema,
        List<ExpressionTree> exprs,
        SelectionVectorType selectionVectorType,
        long configurationId)
        throws GandivaException {
      // serialize the schema and the list of expressions as a protobuf
      GandivaTypes.ExpressionList.Builder builder = GandivaTypes.ExpressionList.newBuilder();
      for (ExpressionTree expr : exprs) {
        builder.addExprs(expr.toProtobuf());
      }
  
      // Invoke the JNI layer to create the LLVM module representing the expressions
      GandivaTypes.Schema schemaBuf = ArrowTypeHelper.arrowSchemaToProtobuf(schema);
      JniWrapper wrapper = JniLoader.getInstance().getWrapper();
      long moduleId =
          wrapper.buildProjector(
              schemaBuf.toByteArray(),
              builder.build().toByteArray(),
              selectionVectorType.getNumber(),
              configurationId);
      logger.debug("Created module for the projector with id {}", moduleId);
      return new Projector(wrapper, moduleId, schema, exprs.size());
    }
  
    /**
     * Invoke this function to evaluate a set of expressions against a recordBatch.
     *
     * @param recordBatch Record batch including the data
     * @param outColumns Result of applying the project on the data
     */
    public void evaluate(ArrowRecordBatch recordBatch, List<ValueVector> outColumns)
        throws GandivaException {
      evaluate(
          recordBatch.getLength(),
          recordBatch.getBuffers(),
          recordBatch.getBuffersLayout(),
          SelectionVectorType.SV_NONE.getNumber(),
          recordBatch.getLength(),
          0,
          0,
          outColumns);
    }
  
    /**
     * Invoke this function to evaluate a set of expressions against a set of arrow buffers. (this is
     * an optimised version that skips taking references).
     *
     * @param numRows number of rows.
     * @param buffers List of input arrow buffers
     * @param outColumns Result of applying the project on the data
     */
    public void evaluate(int numRows, List<ArrowBuf> buffers, List<ValueVector> outColumns)
        throws GandivaException {
      List<ArrowBuffer> buffersLayout = new ArrayList<>();
      long offset = 0;
      for (ArrowBuf arrowBuf : buffers) {
        long size = arrowBuf.readableBytes();
        buffersLayout.add(new ArrowBuffer(offset, size));
        offset += size;
      }
      evaluate(
          numRows,
          buffers,
          buffersLayout,
          SelectionVectorType.SV_NONE.getNumber(),
          numRows,
          0,
          0,
          outColumns);
    }
  
    /**
     * Invoke this function to evaluate a set of expressions against a {@link ArrowRecordBatch}.
     *
     * @param recordBatch The data to evaluate against.
     * @param selectionVector Selection vector which stores the selected rows.
     * @param outColumns Result of applying the project on the data
     */
    public void evaluate(
        ArrowRecordBatch recordBatch, SelectionVector selectionVector, List<ValueVector> outColumns)
        throws GandivaException {
      evaluate(
          recordBatch.getLength(),
          recordBatch.getBuffers(),
          recordBatch.getBuffersLayout(),
          selectionVector.getType().getNumber(),
          selectionVector.getRecordCount(),
          selectionVector.getBuffer().memoryAddress(),
          selectionVector.getBuffer().capacity(),
          outColumns);
    }
  
    /**
     * Invoke this function to evaluate a set of expressions against a set of arrow buffers on the
     * selected positions. (this is an optimised version that skips taking references).
     *
     * @param numRows number of rows.
     * @param buffers List of input arrow buffers
     * @param selectionVector Selection vector which stores the selected rows.
     * @param outColumns Result of applying the project on the data
     */
    public void evaluate(
        int numRows,
        List<ArrowBuf> buffers,
        SelectionVector selectionVector,
        List<ValueVector> outColumns)
        throws GandivaException {
      List<ArrowBuffer> buffersLayout = new ArrayList<>();
      long offset = 0;
      for (ArrowBuf arrowBuf : buffers) {
        long size = arrowBuf.readableBytes();
        buffersLayout.add(new ArrowBuffer(offset, size));
        offset += size;
      }
      evaluate(
          numRows,
          buffers,
          buffersLayout,
          selectionVector.getType().getNumber(),
          selectionVector.getRecordCount(),
          selectionVector.getBuffer().memoryAddress(),
          selectionVector.getBuffer().capacity(),
          outColumns);
    }
  
    private synchronized void evaluate(
        int numRows,
        List<ArrowBuf> buffers,
        List<ArrowBuffer> buffersLayout,
        int selectionVectorType,
        int selectionVectorRecordCount,
        long selectionVectorAddr,
        long selectionVectorSize,
        List<ValueVector> outColumns)
        throws GandivaException {
      if (this.closed) {
        throw new EvaluatorClosedException();
      }
  
      if (numExprs != outColumns.size()) {
        logger.info("Expected " + numExprs + " columns, got " + outColumns.size());
        throw new GandivaException("Incorrect number of columns for the output vector");
      }
  
      long[] bufAddrs = new long[buffers.size()];
      long[] bufSizes = new long[buffers.size()];
  
      int idx = 0;
      for (ArrowBuf buf : buffers) {
        bufAddrs[idx++] = buf.memoryAddress();
      }
  
      idx = 0;
      for (ArrowBuffer bufLayout : buffersLayout) {
        bufSizes[idx++] = bufLayout.getSize();
      }
  
      boolean hasVariableWidthColumns = false;
      BaseVariableWidthVector[] resizableVectors = new BaseVariableWidthVector[outColumns.size()];
      long[] outAddrs = new long[3 * outColumns.size()];
      long[] outSizes = new long[3 * outColumns.size()];
      idx = 0;
      int outColumnIdx = 0;
      for (ValueVector valueVector : outColumns) {
        boolean isFixedWith = valueVector instanceof FixedWidthVector;
        boolean isVarWidth = valueVector instanceof VariableWidthVector;
        if (!isFixedWith && !isVarWidth) {
          throw new UnsupportedTypeException(
              "Unsupported value vector type " + valueVector.getField().getFieldType());
        }
  
        outAddrs[idx] = valueVector.getValidityBuffer().memoryAddress();
        outSizes[idx++] = valueVector.getValidityBuffer().capacity();
        if (isVarWidth) {
          outAddrs[idx] = valueVector.getOffsetBuffer().memoryAddress();
          outSizes[idx++] = valueVector.getOffsetBuffer().capacity();
          hasVariableWidthColumns = true;
  
          // save vector to allow for resizing.
          resizableVectors[outColumnIdx] = (BaseVariableWidthVector) valueVector;
        }
        outAddrs[idx] = valueVector.getDataBuffer().memoryAddress();
        outSizes[idx++] = valueVector.getDataBuffer().capacity();
  
        valueVector.setValueCount(selectionVectorRecordCount);
        outColumnIdx++;
      }
  
      wrapper.evaluateProjector(
          hasVariableWidthColumns ? new VectorExpander(resizableVectors) : null,
          this.moduleId,
          numRows,
          bufAddrs,
          bufSizes,
          selectionVectorType,
          selectionVectorRecordCount,
          selectionVectorAddr,
          selectionVectorSize,
          outAddrs,
          outSizes);
    }
  
    /** Closes the LLVM module representing this evaluator. */
    public void close() throws GandivaException {
      if (this.closed) {
        return;
      }
  
      wrapper.closeProjector(this.moduleId);
      this.closed = true;
    }
  }
  