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.apex.malhar.kudu.partitioner;
  
  import java.util.ArrayList;
  import java.util.Collection;
  import java.util.List;
  import java.util.Map;
  
  import org.codehaus.jackson.annotate.JsonIgnore;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import org.apache.apex.malhar.kudu.AbstractKuduInputOperator;
  import org.apache.apex.malhar.kudu.ApexKuduConnection;
  import org.apache.apex.malhar.kudu.scanner.AbstractKuduPartitionScanner;
  import org.apache.apex.malhar.kudu.scanner.KuduPartitionScanAssignmentMeta;
  import org.apache.apex.malhar.kudu.sqltranslator.SQLToKuduPredicatesTranslator;
  import org.apache.apex.malhar.lib.util.KryoCloneUtils;
  import org.apache.kudu.ColumnSchema;
  import org.apache.kudu.client.KuduClient;
  import org.apache.kudu.client.KuduScanToken;
  import org.apache.kudu.client.KuduTable;
  
  import com.datatorrent.api.DefaultPartition;
  import com.datatorrent.api.Partitioner;
  
  /**
   * An abstract class that contains logic that is common across all partitioners available for the Kudu input operator.
   *
   * @since 3.8.0
   */
  public abstract class AbstractKuduInputPartitioner implements Partitioner<AbstractKuduInputOperator>
  {
    private static final Logger LOG = LoggerFactory.getLogger(AbstractKuduInputPartitioner.class);
  
    @JsonIgnore
    protected AbstractKuduInputOperator prototypeKuduInputOperator;
  
    public AbstractKuduInputPartitioner(AbstractKuduInputOperator prototypeOperator)
    {
      prototypeKuduInputOperator = prototypeOperator;
    }
  
    /**
     * Builds a set of scan tokens. The list of scan tokens are generated as if the entire table is being scanned
     * i.e. a SELECT * FROM TABLE equivalent expression. This list is used to assign the partition pie assignments
     * for all of the planned partition of operators. Each operator gets a part of the PIE as if all columns were
     * selected. Subsequently when a query is to be processed, the query is used to generate the scan tokens applicable
     * for that query. Given that partition pie represents the entire data set, the scan assignments for the current
     * query will be a subset.
     * @return The list of scan tokens as if the entire table is getting scanned.
     * @throws Exception in cases when the connection to kudu cluster cannot be closed.
     */
    public List<KuduScanToken> getKuduScanTokensForSelectAllColumns() throws Exception
    {
      // We are not using the current query for deciding the partition strategy but a SELECT * as
      // we do not want to want to optimize on just the current query. This prevents rapid throttling of operator
      // instances when the scan patterns are erratic. On the other hand, this might result on under utilized
      // operator resources in the DAG but will be consistent at a minimum.
      ApexKuduConnection apexKuduConnection = prototypeKuduInputOperator.getApexKuduConnectionInfo().build();
      KuduClient clientHandle = apexKuduConnection.getKuduClient();
      KuduTable table = apexKuduConnection.getKuduTable();
      KuduScanToken.KuduScanTokenBuilder builder = clientHandle.newScanTokenBuilder(table);
      List<String> allColumns = new ArrayList<>();
      List<ColumnSchema> columnList = apexKuduConnection.getKuduTable().getSchema().getColumns();
      for ( ColumnSchema column : columnList) {
        allColumns.add(column.getName());
      }
      builder.setProjectedColumnNames(allColumns);
      LOG.debug("Building the partition pie assignments for the input operator");
      List<KuduScanToken> allPossibleTokens = builder.build();
      apexKuduConnection.close();
      return allPossibleTokens;
    }
  
    /***
     * Returns the number of partitions. The configuration is used as the source of truth for
     * number of partitions. The config is then overridden if the code in the client driver code explicitly sets
     * the number of partitions.
     * @param context The context as provided by the launcher
     * @return The number of partitions that are planned for this input operator. At a minimum it is 1.
     */
    public int getNumberOfPartitions(PartitioningContext context)
    {
      int proposedPartitionCount = context.getParallelPartitionCount();
      if ( prototypeKuduInputOperator.getNumberOfPartitions() != -1 ) { // -1 is the default
        // There is a manual override of partitions from code. Use this
        proposedPartitionCount = prototypeKuduInputOperator.getNumberOfPartitions();
        LOG.info(" Set the partition count based on the code as opposed to configuration ");
      }
      if ( proposedPartitionCount <= 0)  {
        // Parallel partitions not enabled. But the design is to use one to many mapping. Hence defaulting to one
        LOG.info(" Defaulting to one partition as parallel partitioning is not enabled");
        proposedPartitionCount = 1;
      }
      LOG.info(" Planning to use " + proposedPartitionCount + " partitions");
      return proposedPartitionCount;
    }
  
    /***
     * Builds a list of scan assignment metadata instances from raw kudu scan tokens as returned by the Kudu Query planner
     *  assuming all of the columns and rows are to be scanned
     * @param partitions The current set of partitions
     * @param context The current partitioning context
     * @return The new set of partitions
     * @throws Exception if the Kudu connection opened for generating the scan plan cannot be closed
     */
    public List<KuduPartitionScanAssignmentMeta> getListOfPartitionAssignments(
        Collection<Partition<AbstractKuduInputOperator>> partitions, PartitioningContext context) throws Exception
    {
      List<KuduPartitionScanAssignmentMeta> returnList = new ArrayList<>();
      List<KuduScanToken> allColumnsScanTokens = new ArrayList<>();
      // we are looking at a first time invocation scenario
      try {
        allColumnsScanTokens.addAll(getKuduScanTokensForSelectAllColumns());
      } catch (Exception e) {
        LOG.error(" Error while calculating the number of scan tokens for all column projections " + e.getMessage(),e);
      }
      if ( allColumnsScanTokens.size() == 0 ) {
        LOG.error("No column information could be extracted from the Kudu table");
        throw new Exception("No column information could be extracted from the Kudu table");
      }
      int totalPartitionCount = allColumnsScanTokens.size();
      LOG.info("Determined maximum as " + totalPartitionCount + " tablets for this table");
      for (int i = 0; i < totalPartitionCount; i++) {
        KuduPartitionScanAssignmentMeta aMeta = new KuduPartitionScanAssignmentMeta();
        aMeta.setOrdinal(i);
        aMeta.setTotalSize(totalPartitionCount);
        returnList.add(aMeta);
        LOG.info("A planned scan meta of the total partitions " + aMeta);
      }
      LOG.info("Total kudu partition size is " + returnList.size());
      return returnList;
    }
  
    /***
     * Clones the original operator and sets the partition pie assignments for this operator. Kryo is used for cloning
     * @param scanAssignmentsForThisPartition The partition pie that is assigned to the operator according to the
     *                                        configured partitioner
     * @return The partition that is used by the runtime to launch the new physical operator instance
     */
    public Partitioner.Partition<AbstractKuduInputOperator> clonePartitionAndAssignScanMeta(
        List<KuduPartitionScanAssignmentMeta> scanAssignmentsForThisPartition)
    {
      Partitioner.Partition<AbstractKuduInputOperator> clonedKuduInputOperator =
          new DefaultPartition<>(KryoCloneUtils.cloneObject(prototypeKuduInputOperator));
      clonedKuduInputOperator.getPartitionedInstance().setPartitionPieAssignment(scanAssignmentsForThisPartition);
      return clonedKuduInputOperator;
    }
  
    /***
     * Implements the main logic for defining partitions. The logic to create a partition is pretty straightforward. A
     * SELECT * expression is used to generate a plan ( a list of kudu scan tokens). The list is then evenly
     * distributed to each of the physical operator instances as defined by the configuration. At runtime when a new query
     * needs to be processed, a plan specific to the query is generated and the planned list of kudu scan tokens
     * for that query are distributed among the available operator instances. Note that we define the partitions only once
     * during the first invocation. This is because t is ideal we do not optimize on the current query that is running
     * before we can dynamically partition and cause too much throttling based on the query patterns. Please see
     * {@link AbstractKuduPartitionScanner#preparePlanForScanners(SQLToKuduPredicatesTranslator)} for details as to how
     * a query based planning is done.
     * @param partitions The current set of partitions
     * @param context The partitioning context as provided by the runtime platform.
     * @return The collection of planned partitions as implemented by the partitioner that is configured by the user
     */
    @Override
    public Collection<Partition<AbstractKuduInputOperator>> definePartitions(
        Collection<Partition<AbstractKuduInputOperator>> partitions, PartitioningContext context)
    {
      if ( partitions != null) {
        LOG.info("The current partitioner plan has " + partitions.size() + " operators before redefining");
      }
      List<Partition<AbstractKuduInputOperator>> partitionsForInputOperator = new ArrayList<>();
      List<KuduPartitionScanAssignmentMeta> assignmentMetaList = new ArrayList<>();
      try {
        assignmentMetaList.addAll(getListOfPartitionAssignments(partitions,context));
      } catch (Exception e) {
        throw new RuntimeException("Aborting partition planning as Kudu meta data could not be obtained", e);
      }
      LOG.info("Maximum possible Kudu input operator partition count is " + assignmentMetaList.size());
      Map<Integer,List<KuduPartitionScanAssignmentMeta>> assignments = assign(assignmentMetaList,context);
      boolean requiresRepartitioning = false;
      if ( partitions == null) {
        requiresRepartitioning = true;
      } else {
        for ( Partition<AbstractKuduInputOperator> aPartition : partitions) {
          if ( !aPartition.getPartitionedInstance().isPartitioned()) {
            requiresRepartitioning = true;
            break;
          }
        }
      }
      if ( requiresRepartitioning ) {
        partitions.clear();
        LOG.info("Clearing all of the current partitions and setting up new ones");
        partitions.clear();
        for ( int i = 0; i < assignments.size(); i++) {
          List<KuduPartitionScanAssignmentMeta> assignmentForThisOperator = assignments.get(i);
          partitionsForInputOperator.add(clonePartitionAndAssignScanMeta(assignmentForThisOperator));
          LOG.info("Assigned apex operator " + i + " with " + assignmentForThisOperator.size() + " kudu mappings");
        }
        LOG.info("Returning " + partitionsForInputOperator.size() + " partitions for the input operator");
        return partitionsForInputOperator;
      } else {
        LOG.info("Not making any changes to the partitions");
        return partitions;
      }
    }
  
    @Override
    public void partitioned(Map<Integer, Partition<AbstractKuduInputOperator>> partitions)
    {
  
    }
  
    /***
     * Abstract method that will be implemented by the individual partition planners basing on the functionality they
     * provide. Please see {@link KuduOneToManyPartitioner} and {@link KuduOneToOnePartitioner} for specific
     * implementations
     * @param totalList The ltotal list of possible tablet scans for all queries
     * @param context The context that is provided by the framework when repartitioning is to be executed
     * @return A map that gives the operator number as key and the list of {@link KuduPartitionScanAssignmentMeta} that
     * are assigned for the operator with that number.
     */
    public abstract  Map<Integer, List<KuduPartitionScanAssignmentMeta>> assign(
        List<KuduPartitionScanAssignmentMeta> totalList, PartitioningContext context);
  
    public AbstractKuduInputOperator getPrototypeKuduInputOperator()
    {
      return prototypeKuduInputOperator;
    }
  
    public void setPrototypeKuduInputOperator(AbstractKuduInputOperator prototypeKuduInputOperator)
    {
      this.prototypeKuduInputOperator = prototypeKuduInputOperator;
    }
  }
  