/*
 * 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 static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertTrue;

import java.io.File;
import java.util.UUID;

import org.apache.spark.HashPartitioner;
import org.apache.spark.TaskContext;
import org.apache.spark.serializer.KryoSerializer;
import org.apache.spark.serializer.Serializer;
import org.apache.spark.serializer.SerializerInstance;
import org.apache.spark.shuffle.ShuffleWriteMetricsReporter;
import org.apache.spark.shuffle.ShuffleWriter;
import org.apache.spark.sql.execution.UnsafeRowSerializer;
import org.apache.spark.sql.execution.columnar.CelebornBatchBuilder;
import org.apache.spark.sql.execution.columnar.CelebornColumnarBatchSerializer;
import org.apache.spark.sql.types.IntegerType$;
import org.apache.spark.sql.types.StringType$;
import org.apache.spark.sql.types.StructType;
import org.junit.Test;
import org.mockito.MockedStatic;
import org.mockito.Mockito;

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

public class ColumnarHashBasedShuffleWriterSuiteJ extends CelebornShuffleWriterSuiteBase {

  private final StructType schema =
      new StructType().add("key", IntegerType$.MODULE$).add("value", StringType$.MODULE$);

  @Test
  public void createColumnarShuffleWriter() throws Exception {
    Mockito.doReturn(new HashPartitioner(numPartitions)).when(dependency).partitioner();
    final CelebornConf conf = new CelebornConf();
    final File tempFile = new File(tempDir, UUID.randomUUID().toString());
    final DummyShuffleClient client = new DummyShuffleClient(conf, tempFile);
    client.initReducePartitionMap(shuffleId, numPartitions, 1);

    // Create ColumnarHashBasedShuffleWriter with handle of which dependency has null schema.
    Mockito.doReturn(new KryoSerializer(sparkConf)).when(dependency).serializer();
    ShuffleWriter<Integer, String> writer =
        createShuffleWriterWithoutSchema(
            new CelebornShuffleHandle<>(
                "appId", "host", 0, this.userIdentifier, 0, false, 10, this.dependency),
            taskContext,
            conf,
            client,
            metrics.shuffleWriteMetrics());
    assertTrue(writer instanceof ColumnarHashBasedShuffleWriter);
    assertFalse(((ColumnarHashBasedShuffleWriter<?, ?, ?>) writer).isColumnarShuffle());

    // Create ColumnarHashBasedShuffleWriter with handle of which dependency has non-null schema.
    Mockito.doReturn(new UnsafeRowSerializer(2, null)).when(dependency).serializer();
    writer =
        createShuffleWriter(
            new CelebornShuffleHandle<>(
                "appId", "host", 0, this.userIdentifier, 0, false, 10, this.dependency),
            taskContext,
            conf,
            client,
            metrics.shuffleWriteMetrics());
    assertTrue(((ColumnarHashBasedShuffleWriter<?, ?, ?>) writer).isColumnarShuffle());
  }

  @Override
  protected SerializerInstance newSerializerInstance(Serializer serializer) {
    if (serializer instanceof UnsafeRowSerializer
        && CelebornBatchBuilder.supportsColumnarType(schema)) {
      CelebornConf conf = new CelebornConf();
      return new CelebornColumnarBatchSerializer(schema, conf.columnarShuffleOffHeapEnabled(), null)
          .newInstance();
    } else {
      return serializer.newInstance();
    }
  }

  @Override
  protected ShuffleWriter<Integer, String> createShuffleWriter(
      CelebornShuffleHandle handle,
      TaskContext context,
      CelebornConf conf,
      ShuffleClient client,
      ShuffleWriteMetricsReporter metrics) {
    try (MockedStatic<CustomShuffleDependencyUtils> utils =
        Mockito.mockStatic(CustomShuffleDependencyUtils.class)) {
      utils
          .when(() -> CustomShuffleDependencyUtils.getSchema(handle.dependency()))
          .thenReturn(schema);
      return SparkUtils.createColumnarHashBasedShuffleWriter(
          SparkUtils.celebornShuffleId(client, handle, context, true),
          handle,
          context,
          conf,
          client,
          metrics,
          SendBufferPool.get(1, 30, 60));
    }
  }

  private ShuffleWriter<Integer, String> createShuffleWriterWithoutSchema(
      CelebornShuffleHandle handle,
      TaskContext context,
      CelebornConf conf,
      ShuffleClient client,
      ShuffleWriteMetricsReporter metrics) {
    return SparkUtils.createColumnarHashBasedShuffleWriter(
        SparkUtils.celebornShuffleId(client, handle, context, true),
        handle,
        context,
        conf,
        client,
        metrics,
        SendBufferPool.get(1, 30, 60));
  }
}
