text: /*
   * Copyright (c) 2018 Uber Technologies, Inc.
   * Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated
   * documentation files (the "Software"), to deal in the Software without restriction, including without limitation the
   * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to
   * permit persons to whom the Software is furnished to do so, subject to the following conditions:
   *
   * The above copyright notice and this permission notice shall be included in all copies or substantial portions
   * of the Software.
   *
   * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO
   * THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
   * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
   * TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
   * IN THE SOFTWARE.
   */
  package com.uber.marmaray.common.converters;
  
  import com.google.common.base.Optional;
  import com.uber.marmaray.common.AvroPayload;
  import com.uber.marmaray.common.configuration.Configuration;
  import com.uber.marmaray.common.converters.data.SparkSourceDataConverter;
  import com.uber.marmaray.common.converters.schema.DataFrameSchemaConverter;
  import com.uber.marmaray.common.util.SparkTestUtil;
  import com.uber.marmaray.utilities.ErrorExtractor;
  import org.apache.avro.LogicalType;
  import org.apache.avro.LogicalTypes;
  import org.apache.avro.Schema;
  import org.apache.avro.SchemaBuilder;
  import org.apache.avro.generic.GenericRecord;
  import org.apache.avro.generic.GenericRecordBuilder;
  import org.apache.spark.SparkConf;
  import org.apache.spark.api.java.JavaRDD;
  import org.apache.spark.api.java.JavaSparkContext;
  import org.apache.spark.sql.Dataset;
  import org.apache.spark.sql.Row;
  import org.apache.spark.sql.RowFactory;
  import org.apache.spark.sql.SparkSession;
  import org.apache.spark.sql.types.DataTypes;
  import org.apache.spark.sql.types.StructField;
  import org.apache.spark.sql.types.StructType;
  import org.junit.After;
  import org.junit.Assert;
  import org.junit.Before;
  import org.junit.Test;
  
  import java.sql.Timestamp;
  import java.util.ArrayList;
  import java.util.Collections;
  import java.util.List;
  
  /**
   * Tests for converting data from DataFrame (Spark) -> Avro schema with data
   */
  public class TestDataFrameDataConverter {
  
      private static final String STRING_FIELD = "string_field";
      private static final String INT_FIELD = "int_field";
      private static final String BOOL_FIELD = "bool_field";
      private static final String DOUBLE_FIELD = "double_field";
  
      private Optional<SparkSession> spark = Optional.absent();
      private Optional<JavaSparkContext> jsc = Optional.absent();
  
      @Before
      public void setupTest() {
          final SparkConf sparkConf = SparkTestUtil.getSparkConf(TestCassandraDataFrameConverter.class.getName());
          this.spark = Optional.of(SparkTestUtil.getSparkSession(sparkConf));
          this.jsc = Optional.of(new JavaSparkContext(spark.get().sparkContext()));
      }
  
      @After
      public void teardownTest() {
          if (this.spark.isPresent()) {
              this.spark.get().stop();
              this.spark = Optional.absent();
          }
  
          if (this.jsc.isPresent()) {
              this.jsc.get().stop();
              this.jsc = Optional.absent();
          }
      }
  
      @Test
      public void convertDataFrameToCommonDataSchemaTest() {
          final List<Object[]> dataList = new ArrayList<>();
          dataList.add(new Object[] {"data0", Integer.MIN_VALUE, true, Double.MIN_VALUE});
          dataList.add(new Object[] {"data1", Integer.MAX_VALUE, false, Double.MAX_VALUE});
  
          final JavaRDD<Row> rowRDDNoSchema = this.jsc.get()
                  .parallelize(dataList).map((Object[] row) -> RowFactory.create(row));
  
          final StructType schema = DataTypes
                  .createStructType(new StructField[] {
                          DataTypes.createStructField(STRING_FIELD, DataTypes.StringType, false),
                          DataTypes.createStructField(INT_FIELD, DataTypes.IntegerType, false),
                          DataTypes.createStructField(BOOL_FIELD, DataTypes.BooleanType, false),
                          DataTypes.createStructField(DOUBLE_FIELD, DataTypes.DoubleType, false)});
  
          final Dataset<Row> df = this.spark.get().createDataFrame(rowRDDNoSchema, schema);
          final DataFrameSchemaConverter dfsc = new DataFrameSchemaConverter();
          final Schema commonSchema = dfsc.convertToCommonSchema(df.schema());
  
          final SparkSourceDataConverter converter =
              new SparkSourceDataConverter(schema, commonSchema, new Configuration(), Collections.singleton(STRING_FIELD),
                                              new ErrorExtractor());
          final JavaRDD<AvroPayload> payloadRDD = converter.map(df.javaRDD()).getData();
  
          final List<AvroPayload> payloads = payloadRDD.collect();
  
          Assert.assertEquals(2, payloads.size());
  
          final AvroPayload ap = payloads.get(0);
          final GenericRecord gr = ap.getData();
  
          Assert.assertEquals("data0", gr.get(STRING_FIELD).toString());
          Assert.assertEquals(Integer.MIN_VALUE, gr.get(INT_FIELD));
          Assert.assertEquals(true, gr.get(BOOL_FIELD));
          Assert.assertEquals(Double.MIN_VALUE, gr.get(DOUBLE_FIELD));
  
          final AvroPayload ap2 = payloads.get(1);
          final GenericRecord gr2 = ap2.getData();
  
          Assert.assertEquals("data1", gr2.get(STRING_FIELD).toString());
          Assert.assertEquals(Integer.MAX_VALUE, gr2.get(INT_FIELD));
          Assert.assertEquals(false, gr2.get(BOOL_FIELD));
          Assert.assertEquals(Double.MAX_VALUE, gr2.get(DOUBLE_FIELD));
      }
  
      @Test
      public void convertDataFrameToCommonDataNotAllColumnsPopulatedTest() throws Exception {
          final List<Object[]> dataList = new ArrayList<>();
          dataList.add(new Object[] {"data0", Integer.MIN_VALUE, false});
  
          final JavaRDD<Row> rowRDDNoSchema = this.jsc.get()
                  .parallelize(dataList).map((Object[] row) -> RowFactory.create(row));
  
          final StructType schemaWithMissingField = DataTypes
                  .createStructType(new StructField[] {
                          DataTypes.createStructField(STRING_FIELD, DataTypes.StringType, false),
                          DataTypes.createStructField(INT_FIELD, DataTypes.IntegerType, false),
                          DataTypes.createStructField(BOOL_FIELD, DataTypes.BooleanType, false)});
  
          final StructType expectedSchema = DataTypes
                  .createStructType(new StructField[] {
                          DataTypes.createStructField(STRING_FIELD, DataTypes.StringType, false),
                          DataTypes.createStructField(INT_FIELD, DataTypes.IntegerType, false),
                          DataTypes.createStructField(BOOL_FIELD, DataTypes.BooleanType, false),
                          DataTypes.createStructField(DOUBLE_FIELD, DataTypes.DoubleType, true)});
  
          final Dataset<Row> df = this.spark.get().createDataFrame(rowRDDNoSchema, schemaWithMissingField);
          final DataFrameSchemaConverter dfsc = new DataFrameSchemaConverter();
          final Schema commonSchema = dfsc.convertToCommonSchema(expectedSchema);
  
          final List<Row> rows = df.javaRDD().collect();
  
          final SparkSourceDataConverter converter =
                  new SparkSourceDataConverter(expectedSchema, commonSchema,
                          new Configuration(), Collections.singleton(STRING_FIELD), new ErrorExtractor());
  
          // The missing field should be handled gracefully without failing since rows can have sparse data
          final AvroPayload payload = converter.convert(rows.get(0)).get(0).getSuccessData().get().getData();
          final GenericRecord gr = payload.getData();
  
          Assert.assertEquals("data0", gr.get(STRING_FIELD).toString());
          Assert.assertEquals(Integer.MIN_VALUE, gr.get(INT_FIELD));
          Assert.assertEquals(false, gr.get(BOOL_FIELD));
          Assert.assertNull(gr.get(DOUBLE_FIELD));
      }
  }
  