measure/src/test/scala/org/apache/griffin/measure/datasource/connector/batch/FileBasedDataConnectorTest.scala

/*
 * 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.griffin.measure.datasource.connector.batch

import org.apache.spark.sql.SaveMode
import org.apache.spark.sql.types.{IntegerType, LongType, StringType, StructType}
import org.scalatest._

import org.apache.griffin.measure.SparkSuiteBase
import org.apache.griffin.measure.configuration.dqdefinition.DataConnectorParam
import org.apache.griffin.measure.datasource.TimestampStorage
import org.apache.griffin.measure.step.builder.ConstantColumns

class FileBasedDataConnectorTest extends SparkSuiteBase with Matchers {

  override def beforeAll(): Unit = {
    super.beforeAll()

    createDataSets(s"file://${getClass.getResource("/").getPath}")
  }

  private def createDataSets(basePath: String): Unit = {
    val formats = Seq("parquet", "orc", "csv", "tsv")
    val schema = new StructType().add("name", StringType).add("age", IntegerType, nullable = true)

    val df = spark.read.schema(schema).csv(s"${basePath}hive/person_table.csv")

    df.cache()
    formats.foreach(f => {
      val delimiter = if (f.matches("csv")) "," else if (f.matches("tsv")) "\t" else ""
      df.write
        .mode(SaveMode.Overwrite)
        .option("delimiter", delimiter)
        .option("header", "true")
        .format(if (f.matches("tsv")) "csv" else f)
        .save(s"${basePath}files/person_table.$f")
    })

    df.unpersist()
  }

  private final val dcParam = DataConnectorParam("file", "1", "test_df", Map.empty[String, String], Nil)
  private final val timestampStorage = TimestampStorage()

  // Regarding Local FileSystem

  "file based data connector" should "be able to read from local filesystem" in {
    val configs = Map(
      "format" -> "csv",
      "paths" -> Seq(
        s"file://${getClass.getResource("/hive/person_table.csv").getPath}"
      ),
      "options" -> Map(
        "header" -> "false"
      )
    )

    val dc = FileBasedDataConnector(spark, dcParam.copy(config = configs), timestampStorage)
    val result = dc.data(1000L)

    assert(result._1.isDefined)
    assert(result._1.get.collect().length == 2)
  }

  // Regarding User Defined Schema

  it should "respect the provided schema, if any" in {
    val configs = Map(
      "format" -> "csv",
      "paths" -> Seq(
        s"file://${getClass.getResource("/hive/person_table.csv").getPath}"
      )
    )

    // no schema
    assertThrows[IllegalArgumentException](
      FileBasedDataConnector(spark, dcParam.copy(config = configs), timestampStorage)
    )

    // invalid schema
    assertThrows[IllegalStateException](
      FileBasedDataConnector(spark, dcParam.copy(config = configs + (("schema", ""))), timestampStorage)
    )

    // valid schema
    val result1 = FileBasedDataConnector(spark,
      dcParam.copy(config = configs + (("schema",
        Seq(Map("name" -> "name", "type" -> "string"), Map("name" -> "age", "type" -> "int", "nullable" -> "true"))
      )))
      , timestampStorage)
      .data(1L)

    val expSchema = new StructType()
      .add("name", StringType)
      .add("age", IntegerType, nullable = true)
      .add(ConstantColumns.tmst, LongType, nullable = false)

    assert(result1._1.isDefined)
    assert(result1._1.get.collect().length == 2)
    assert(result1._1.get.schema == expSchema)

    // valid headers
    val result2 = FileBasedDataConnector(spark,
      dcParam.copy(config = configs + (("options", Map(
        "header" -> "true"
      )
      )))
      , timestampStorage)
      .data(1L)

    assert(result2._1.isDefined)
    assert(result2._1.get.collect().length == 1)
    result2._1.get.columns should contain theSameElementsAs Seq("Joey", "14", ConstantColumns.tmst)
  }

  // skip on erroneous paths

  it should "respect options if an erroneous path is encountered" in {
    val configs = Map(
      "format" -> "csv",
      "paths" -> Seq(
        s"file://${getClass.getResource("/hive/person_table.csv").getPath}",
        s"${java.util.UUID.randomUUID().toString}/"
      ),
      "skipErrorPaths" -> true,
      "options" -> Map(
        "header" -> "true"
      )
    )

    // valid paths
    val result1 = FileBasedDataConnector(spark, dcParam.copy(config = configs), timestampStorage).data(1L)

    assert(result1._1.isDefined)
    assert(result1._1.get.collect().length == 1)

    // non existent path
    assertThrows[IllegalArgumentException](
      FileBasedDataConnector(spark, dcParam.copy(config = configs - "skipErrorPaths"), timestampStorage).data(1L)
    )

    // no path
    assertThrows[AssertionError](
      FileBasedDataConnector(spark, dcParam.copy(config = configs - "paths"), timestampStorage).data(1L)
    )
  }

  // Regarding various formats
  it should "be able to read all supported file types" in {

    val formats = Seq("parquet", "orc", "csv", "tsv")
    formats.map(f => {
      val configs = Map(
        "format" -> f,
        "paths" -> Seq(
          s"file://${getClass.getResource(s"/files/person_table.$f").getPath}"
        ),
        "options" -> Map(
          "header" -> "true",
          "inferSchema" -> "true"
        )
      )

      val result = FileBasedDataConnector(spark, dcParam.copy(config = configs), timestampStorage).data(1L)

      assert(result._1.isDefined)

      val df = result._1.get
      val expSchema = new StructType()
        .add("name", StringType)
        .add("age", IntegerType, nullable = true)
        .add(ConstantColumns.tmst, LongType, nullable = false)

      assert(df.collect().length == 2)
      assert(df.schema == expSchema)
    })
  }

  it should "apply schema to all formats if provided" in {
    val formats = Seq("parquet", "orc", "csv", "tsv")
    formats.map(f => {
      val configs = Map(
        "format" -> f,
        "paths" -> Seq(
          s"file://${getClass.getResource(s"/files/person_table.$f").getPath}"
        ),
        "options" -> Map(
          "header" -> "true"
        ),
        "schema" -> Seq(Map("name" -> "name", "type" -> "string"))
      )

      val result = FileBasedDataConnector(spark, dcParam.copy(config = configs), timestampStorage).data(1L)

      assert(result._1.isDefined)

      val df = result._1.get
      val expSchema = new StructType()
        .add("name", StringType)
        .add(ConstantColumns.tmst, LongType, nullable = false)

      assert(df.collect().length == 2)
      assert(df.schema == expSchema)
    })
  }

}