/*
 * 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.spark.shuffle.celeborn;

import java.io.IOException;

import scala.Product2;

import com.google.common.annotations.VisibleForTesting;
import org.apache.spark.ShuffleDependency;
import org.apache.spark.TaskContext;
import org.apache.spark.annotation.Private;
import org.apache.spark.serializer.Serializer;
import org.apache.spark.shuffle.ShuffleWriteMetricsReporter;
import org.apache.spark.sql.catalyst.expressions.UnsafeRow;
import org.apache.spark.sql.execution.UnsafeRowSerializer;
import org.apache.spark.sql.execution.columnar.CelebornBatchBuilder;
import org.apache.spark.sql.execution.columnar.CelebornColumnarBatchBuilder;
import org.apache.spark.sql.execution.columnar.CelebornColumnarBatchCodeGenBuild;
import org.apache.spark.sql.execution.metric.SQLMetric;
import org.apache.spark.sql.types.StructType;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import org.apache.celeborn.client.ShuffleClient;
import org.apache.celeborn.common.CelebornConf;

@Private
public class ColumnarHashBasedShuffleWriter<K, V, C> extends HashBasedShuffleWriter<K, V, C> {

  private static final Logger logger =
      LoggerFactory.getLogger(ColumnarHashBasedShuffleWriter.class);

  private final int stageId;
  private final int shuffleId;
  private final CelebornBatchBuilder[] celebornBatchBuilders;
  private final StructType schema;
  private final Serializer depSerializer;
  private final boolean isColumnarShuffle;
  private final int columnarShuffleBatchSize;
  private final boolean columnarShuffleCodeGenEnabled;
  private final boolean columnarShuffleDictionaryEnabled;
  private final double columnarShuffleDictionaryMaxFactor;

  public ColumnarHashBasedShuffleWriter(
      int shuffleId,
      CelebornShuffleHandle<K, V, C> handle,
      TaskContext taskContext,
      CelebornConf conf,
      ShuffleClient client,
      ShuffleWriteMetricsReporter metrics,
      SendBufferPool sendBufferPool)
      throws IOException {
    super(shuffleId, handle, taskContext, conf, client, metrics, sendBufferPool);
    columnarShuffleBatchSize = conf.columnarShuffleBatchSize();
    columnarShuffleCodeGenEnabled = conf.columnarShuffleCodeGenEnabled();
    columnarShuffleDictionaryEnabled = conf.columnarShuffleDictionaryEnabled();
    columnarShuffleDictionaryMaxFactor = conf.columnarShuffleDictionaryMaxFactor();
    ShuffleDependency<?, ?, ?> shuffleDependency = handle.dependency();
    this.stageId = taskContext.stageId();
    this.shuffleId = shuffleDependency.shuffleId();
    this.schema = CustomShuffleDependencyUtils.getSchema(shuffleDependency);
    this.depSerializer = handle.dependency().serializer();
    this.celebornBatchBuilders =
        new CelebornBatchBuilder[handle.dependency().partitioner().numPartitions()];
    this.isColumnarShuffle = schema != null && CelebornBatchBuilder.supportsColumnarType(schema);
  }

  @Override
  protected void fastWrite0(scala.collection.Iterator iterator)
      throws IOException, InterruptedException {
    if (isColumnarShuffle) {
      logger.info("Fast columnar write of columnar shuffle {} for stage {}.", shuffleId, stageId);
      fastColumnarWrite0(iterator);
    } else {
      super.fastWrite0(iterator);
    }
  }

  private void fastColumnarWrite0(scala.collection.Iterator iterator) throws IOException {
    final scala.collection.Iterator<Product2<Integer, UnsafeRow>> records = iterator;

    SQLMetric dataSize = SparkUtils.getDataSize((UnsafeRowSerializer) depSerializer);
    while (records.hasNext()) {
      final Product2<Integer, UnsafeRow> record = records.next();
      final int partitionId = record._1();
      final UnsafeRow row = record._2();

      if (celebornBatchBuilders[partitionId] == null) {
        CelebornBatchBuilder columnBuilders;
        if (columnarShuffleCodeGenEnabled && !columnarShuffleDictionaryEnabled) {
          columnBuilders =
              new CelebornColumnarBatchCodeGenBuild().create(schema, columnarShuffleBatchSize);
        } else {
          columnBuilders =
              new CelebornColumnarBatchBuilder(
                  schema,
                  columnarShuffleBatchSize,
                  columnarShuffleDictionaryMaxFactor,
                  columnarShuffleDictionaryEnabled);
        }
        columnBuilders.newBuilders();
        celebornBatchBuilders[partitionId] = columnBuilders;
      }

      celebornBatchBuilders[partitionId].writeRow(row);
      if (celebornBatchBuilders[partitionId].getRowCnt() >= columnarShuffleBatchSize) {
        byte[] arr = celebornBatchBuilders[partitionId].buildColumnBytes();
        pushGiantRecord(partitionId, arr, arr.length);
        if (dataSize != null) {
          dataSize.add(arr.length);
        }
        celebornBatchBuilders[partitionId].newBuilders();
      }
      tmpRecordsWritten++;
    }
  }

  @Override
  protected void closeWrite() throws IOException {
    if (canUseFastWrite() && isColumnarShuffle) {
      closeColumnarWrite();
    } else {
      super.closeWrite();
    }
  }

  private void closeColumnarWrite() throws IOException {
    SQLMetric dataSize = SparkUtils.getDataSize((UnsafeRowSerializer) depSerializer);
    for (int i = 0; i < celebornBatchBuilders.length; i++) {
      final CelebornBatchBuilder builders = celebornBatchBuilders[i];
      if (builders != null && builders.getRowCnt() > 0) {
        byte[] buffers = builders.buildColumnBytes();
        if (dataSize != null) {
          dataSize.add(buffers.length);
        }
        mergeData(i, buffers, 0, buffers.length);
        // free buffer
        celebornBatchBuilders[i] = null;
      }
    }
  }

  @VisibleForTesting
  public boolean isColumnarShuffle() {
    return isColumnarShuffle;
  }
}
