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.asterix.runtime.aggregates.std;
  
  import java.io.DataOutput;
  import java.io.IOException;
  import java.util.ArrayList;
  import java.util.Arrays;
  import java.util.Comparator;
  import java.util.List;
  
  import org.apache.asterix.formats.nontagged.BinaryComparatorFactoryProvider;
  import org.apache.asterix.formats.nontagged.SerializerDeserializerProvider;
  import org.apache.asterix.om.base.ABinary;
  import org.apache.asterix.om.base.AMutableBinary;
  import org.apache.asterix.om.functions.BuiltinFunctions;
  import org.apache.asterix.om.functions.IFunctionDescriptor;
  import org.apache.asterix.om.functions.IFunctionDescriptorFactory;
  import org.apache.asterix.om.functions.IFunctionTypeInferer;
  import org.apache.asterix.om.types.ATypeTag;
  import org.apache.asterix.om.types.BuiltinType;
  import org.apache.asterix.om.types.IAType;
  import org.apache.asterix.runtime.aggregates.base.AbstractAggregateFunctionDynamicDescriptor;
  import org.apache.asterix.runtime.functions.FunctionTypeInferers;
  import org.apache.hyracks.algebricks.core.algebra.functions.FunctionIdentifier;
  import org.apache.hyracks.algebricks.runtime.base.IAggregateEvaluator;
  import org.apache.hyracks.algebricks.runtime.base.IAggregateEvaluatorFactory;
  import org.apache.hyracks.algebricks.runtime.base.IScalarEvaluator;
  import org.apache.hyracks.algebricks.runtime.base.IScalarEvaluatorFactory;
  import org.apache.hyracks.api.context.IEvaluatorContext;
  import org.apache.hyracks.api.dataflow.value.IBinaryComparator;
  import org.apache.hyracks.api.dataflow.value.ISerializerDeserializer;
  import org.apache.hyracks.api.exceptions.HyracksDataException;
  import org.apache.hyracks.api.exceptions.SourceLocation;
  import org.apache.hyracks.data.std.api.IPointable;
  import org.apache.hyracks.data.std.primitive.ByteArrayPointable;
  import org.apache.hyracks.data.std.primitive.IntegerPointable;
  import org.apache.hyracks.data.std.primitive.VoidPointable;
  import org.apache.hyracks.data.std.util.ArrayBackedValueStorage;
  import org.apache.hyracks.dataflow.common.data.accessors.IFrameTupleReference;
  import org.apache.hyracks.dataflow.common.data.marshalling.ByteArraySerializerDeserializer;
  import org.apache.hyracks.dataflow.common.data.marshalling.DoubleArraySerializerDeserializer;
  import org.apache.hyracks.dataflow.common.data.marshalling.IntArraySerializerDeserializer;
  import org.apache.hyracks.dataflow.common.data.marshalling.IntegerSerializerDeserializer;
  
  /**
   * <pre>
   * {@see {@link org.apache.hyracks.dataflow.common.data.partition.range.RangeMap}} for some description of the range map
   * structure that is produced by this function. Given a list of samples and a number of partitions "k", the algorithm
   * of this function operates as follows (s = sample):
   * It picks (k - 1) split points out of the samples by dividing num_samples/num_partitions. For 4 partitions, it's 3:
   * s0,s1,s2,s3,s4,s5,s6,s7,s8,s9,s10,s11,s12,s13,s14,s15; 16/4 = 4; range map = [s3, s7, s11]
   *           |           |             |
   *
   * s0,s1,s2,s3,s4,s5,s6; 7/4 = 2; range map = [s1, s3, s5]
   *     |     |     |
   *
   * s0,s1,s2,s3,s4; 5/4 = 2; range map = [s1, s3, s4]; if we go out of bound for the last split, we pick the last item.
   *     |     |  |
   *
   * s0,s1,s2,s3; if #_samples <= #_partitions, we sweep from the beginning (should be rare). range map = [s0, s1, s2]
   *  |  |  |
   *
   * s0,s1; if there are way less samples, we sweep and repeat the last item; range map = [s0, s1, s1];
   * Note: a sample (and therefore also a split point) could be single-column or multi-column.
   * </pre>
   */
  public class RangeMapAggregateDescriptor extends AbstractAggregateFunctionDynamicDescriptor {
      private static final long serialVersionUID = 1L;
      private boolean[] ascFlags;
      private int numPartitions;
      private int numOrderFields;
      private IAType[] argsTypes;
  
      public static final IFunctionDescriptorFactory FACTORY = new IFunctionDescriptorFactory() {
          @Override
          public IFunctionDescriptor createFunctionDescriptor() {
              return new RangeMapAggregateDescriptor();
          }
  
          @Override
          public IFunctionTypeInferer createFunctionTypeInferer() {
              return FunctionTypeInferers.SET_SORTING_PARAMETERS;
          }
      };
  
      @Override
      public FunctionIdentifier getIdentifier() {
          return BuiltinFunctions.RANGE_MAP;
      }
  
      /**
       * The sampling function, which generates the splitting vector, needs to know the number of partitions in order to
       * determine how many split points to pick out of the samples. It also needs to know the ascending/descending of
       * each sort field so that it can sort the samples accordingly first and then pick the (number of partitions - 1)
       * split points out of the sorted samples.
       * @param states states[0]: number of partitions, states[1]: ascending flags, states[2]: inputs types
       */
      @Override
      public void setImmutableStates(Object... states) {
          numPartitions = (int) states[0];
          ascFlags = (boolean[]) states[1];
          numOrderFields = ascFlags.length;
          argsTypes = (IAType[]) states[2];
      }
  
      @Override
      public IAggregateEvaluatorFactory createAggregateEvaluatorFactory(final IScalarEvaluatorFactory[] args) {
          return new IAggregateEvaluatorFactory() {
              private static final long serialVersionUID = 1L;
  
              @Override
              public IAggregateEvaluator createAggregateEvaluator(final IEvaluatorContext ctx)
                      throws HyracksDataException {
                  return new RangeMapFunction(args, ctx, ascFlags, numPartitions, numOrderFields, sourceLoc, argsTypes);
              }
          };
      }
  
      private static class RangeMapFunction extends AbstractAggregateFunction {
          @SuppressWarnings("unchecked")
          private ISerializerDeserializer<ABinary> binarySerde =
                  SerializerDeserializerProvider.INSTANCE.getSerializerDeserializer(BuiltinType.ABINARY);
          private final AMutableBinary binary = new AMutableBinary(null, 0, 0);
          private final List<List<byte[]>> finalSamples = new ArrayList<>();
          private final ArrayBackedValueStorage storage = new ArrayBackedValueStorage();
          private final IPointable input = new VoidPointable();
          private final ByteArrayPointable rangeMapPointable = new ByteArrayPointable();
          private final IScalarEvaluator localSamplesEval;
          private final Comparator<List<byte[]>> comparator;
          private final int numOfPartitions;
          private final int numOrderByFields;
          private final int[] splitPoints;
          private final double[] percentages;
  
          @SuppressWarnings("unchecked")
          private RangeMapFunction(IScalarEvaluatorFactory[] args, IEvaluatorContext context, boolean[] ascending,
                  int numOfPartitions, int numOrderByFields, SourceLocation sourceLocation, IAType[] argsTypes)
                  throws HyracksDataException {
              super(sourceLocation);
              this.localSamplesEval = args[0].createScalarEvaluator(context);
              this.comparator = createComparator(ascending, argsTypes);
              this.numOfPartitions = numOfPartitions;
              this.numOrderByFields = numOrderByFields;
              this.splitPoints = new int[numOfPartitions - 1];
              this.percentages = new double[numOfPartitions - 1];
          }
  
          @Override
          public void init() throws HyracksDataException {
              finalSamples.clear();
          }
  
          /**
           * Receives the local samples and appends them to the final list of samples.
           * @param tuple the partition's samples
           * @throws HyracksDataException IO Exception
           */
          @Override
          public void step(IFrameTupleReference tuple) throws HyracksDataException {
              // check if empty stream (system_null), i.e. partition is empty, so no samples
              localSamplesEval.evaluate(tuple, input);
              if (input.getByteArray()[input.getStartOffset()] == ATypeTag.SERIALIZED_SYSTEM_NULL_TYPE_TAG) {
                  return;
              }
              rangeMapPointable.set(input.getByteArray(), input.getStartOffset() + 1, input.getLength() - 1);
              byte[] rangeMapBytes = rangeMapPointable.getByteArray();
              int pointer = rangeMapPointable.getContentStartOffset();
              int numSamples = IntegerPointable.getInteger(rangeMapBytes, pointer);
              pointer += Integer.BYTES; // eat the 4 bytes of the integer (number of samples)
              for (int i = 0; i < numSamples; i++) {
                  List<byte[]> oneSample = new ArrayList<>(numOrderByFields);
                  for (int j = 0; j < numOrderByFields; j++) {
                      int valueLength = IntegerPointable.getInteger(rangeMapBytes, pointer);
                      pointer += Integer.BYTES; // eat the 4 bytes of the integer and move to the value
                      oneSample.add(Arrays.copyOfRange(rangeMapBytes, pointer, pointer + valueLength));
                      pointer += valueLength; // eat the length of the value and move to the next pair length:value
                  }
                  finalSamples.add(oneSample);
              }
          }
  
          /**
           * Produces the range map out of the collected samples from each partition. The final list of samples is sorted
           * first. Then, we select the split points by dividing the samples evenly.
           * @param result contains the serialized range map.
           * @throws HyracksDataException IO Exception
           */
          @Override
          public void finish(IPointable result) throws HyracksDataException {
              // storage == all serialized split values of all split points
              storage.reset();
              DataOutput allSplitValuesOut = storage.getDataOutput();
              int[] endOffsets;
              try {
                  // check if empty dataset, i.e. no samples have been received from any partition
                  if (finalSamples.isEmpty()) {
                      // a range map with null values
                      endOffsets = new int[numOrderByFields];
                      for (int sortField = 0; sortField < numOrderByFields; sortField++) {
                          allSplitValuesOut.write(ATypeTag.SERIALIZED_NULL_TYPE_TAG);
                          endOffsets[sortField] = storage.getLength();
                      }
                  } else {
                      finalSamples.sort(comparator);
                      calculateSplitIndexes();
                      calculatePercentSplit();
  
                      List<byte[]> sample;
                      int endOffsetsCounter = 0;
                      endOffsets = new int[splitPoints.length * numOrderByFields];
                      for (int i = 0; i < splitPoints.length; i++) {
                          sample = finalSamples.get(splitPoints[i]);
                          for (byte[] column : sample) {
                              allSplitValuesOut.write(column);
                              endOffsets[endOffsetsCounter++] = storage.getLength();
                          }
                      }
                  }
              } catch (IOException e) {
                  throw HyracksDataException.create(e);
              }
              serializeRangeMap(numOrderByFields, storage.getByteArray(), endOffsets, result);
          }
  
          private void calculateSplitIndexes() {
              int nextSplitOffset = (int) Math.ceil(finalSamples.size() / (double) numOfPartitions);
              int nextSplitIndex = nextSplitOffset - 1;
  
              for (int split = 0; split < splitPoints.length; split++) {
                  splitPoints[split] = nextSplitIndex;
                  nextSplitIndex += nextSplitOffset;
                  // in case we go beyond the boundary of samples, we pick the last sample repeatedly
                  if (nextSplitIndex >= finalSamples.size()) {
                      nextSplitIndex = finalSamples.size() - 1;
                  }
              }
          }
  
          private void calculatePercentSplit() {
              for (int i = 0; i < splitPoints.length; i++) {
                  List<byte[]> sampleAtSplit = finalSamples.get(splitPoints[i]);
                  int smallestIndexEqualToSample = splitPoints[i];
                  while (smallestIndexEqualToSample >= 0
                          && comparator.compare(sampleAtSplit, finalSamples.get(smallestIndexEqualToSample)) == 0) {
                      smallestIndexEqualToSample--;
                  }
                  smallestIndexEqualToSample++;
  
                  int largestSplitIncludingSample = i;
                  while (largestSplitIncludingSample < splitPoints.length && comparator.compare(sampleAtSplit,
                          finalSamples.get(splitPoints[largestSplitIncludingSample])) == 0) {
                      largestSplitIncludingSample++;
                  }
                  largestSplitIncludingSample--;
  
                  int largestIndexEqualToSample = splitPoints[largestSplitIncludingSample];
                  while (largestIndexEqualToSample < finalSamples.size()
                          && comparator.compare(sampleAtSplit, finalSamples.get(largestIndexEqualToSample)) == 0) {
                      largestIndexEqualToSample++;
                  }
                  largestIndexEqualToSample--;
  
                  double count = largestIndexEqualToSample - smallestIndexEqualToSample + 1;
                  double waterMark = smallestIndexEqualToSample - 1;
                  for (int j = i; j <= largestSplitIncludingSample; j++) {
                      percentages[j] = ((splitPoints[j] - waterMark) * 100) / (count);
                      waterMark = splitPoints[j];
                  }
                  i = largestSplitIncludingSample;
              }
          }
  
          @Override
          public void finishPartial(IPointable result) throws HyracksDataException {
              finish(result);
          }
  
          /**
           * Creates the comparator that sorts all the collected samples before picking split points.
           * @param asc ascending or descending flag for each sort column.
           * @param types types of inputs to range map function produced by the local step and holding sort fields types
           * @return the described comparator
           */
          private static Comparator<List<byte[]>> createComparator(boolean[] asc, IAType[] types) {
              // create the generic comparator for each sort field, sort fields types start at index 1
              IBinaryComparator[] fieldsComparators = new IBinaryComparator[asc.length];
              for (int i = 0, fieldIdx = 1; fieldIdx < types.length; i++, fieldIdx++) {
                  fieldsComparators[i] = BinaryComparatorFactoryProvider.INSTANCE
                          .getBinaryComparatorFactory(types[fieldIdx], types[fieldIdx], asc[i]).createBinaryComparator();
              }
              return (splitPoint1, splitPoint2) -> {
                  try {
                      // two split points must have the same num of fields
                      int numFields = splitPoint1.size();
                      int result = 0;
                      byte[] field1;
                      byte[] field2;
                      for (int i = 0; i < numFields; i++) {
                          field1 = splitPoint1.get(i);
                          field2 = splitPoint2.get(i);
                          result = fieldsComparators[i].compare(field1, 0, field1.length, field2, 0, field2.length);
                          if (result != 0) {
                              return result;
                          }
                      }
                      return result;
                  } catch (HyracksDataException e) {
                      throw new IllegalStateException(e);
                  }
              };
          }
  
          /**
           * Serializes the range map object defined by the below attributes into the "result". The range map object is
           * serialized as binary data.
           * @param numberFields the number of order-by fields (the sort fields)
           * @param splitValues the serialized split values stored one after the other
           * @param endOffsets the end offsets of each split value
           * @param result where the range map object is serialized
           * @throws HyracksDataException IO Exception
           */
          private void serializeRangeMap(int numberFields, byte[] splitValues, int[] endOffsets, IPointable result)
                  throws HyracksDataException {
              ArrayBackedValueStorage serRangeMap = new ArrayBackedValueStorage();
              IntegerSerializerDeserializer.write(numberFields, serRangeMap.getDataOutput());
              ByteArraySerializerDeserializer.write(splitValues, serRangeMap.getDataOutput());
              IntArraySerializerDeserializer.write(endOffsets, serRangeMap.getDataOutput());
              DoubleArraySerializerDeserializer.write(percentages, serRangeMap.getDataOutput());
              binary.setValue(serRangeMap.getByteArray(), serRangeMap.getStartOffset(), serRangeMap.getLength());
              storage.reset();
              binarySerde.serialize(binary, storage.getDataOutput());
              result.set(storage);
          }
      }
  }
  