import com.databricks.spark.sql.perf.tpcds.TPCDSTables
import org.apache.spark.sql.SparkSession

// Use this script when data has already been generated (i.e. generate-data.scala has
// been run before) and you only need to re-register the metastore tables — for example
// after switching from stock Spark to TabbyDB, or after a fresh cluster restart.

// ── Configure these before running ──────────────────────────────────────────
val useHive    = true
val useIceberg = false
val createExternalHiveTables = true

val rootDir      = "file:///path/to/tpcds/data"  // must point to already-generated data
val databaseName = "tpcds"
val scaleFactor  = "50"                           // must match what was used to generate data
val dsdgenDir    = "/path/to/tpcds-kit/tools"    // still required by TPCDSTables constructor
// ────────────────────────────────────────────────────────────────────────────

val sqlContext = if (useHive) {
  SparkSession.builder()
    .appName("tpcds-create-tables")
    .enableHiveSupport()
    .getOrCreate()
    .sqlContext
} else if (useIceberg) {
  SparkSession.builder()
    .appName("tpcds-create-tables")
    .config("spark.sql.extensions", "org.apache.iceberg.spark.extensions.IcebergSparkSessionExtensions")
    .config("spark.sql.catalog.spark_catalog.type", "hadoop")
    .config("spark.sql.catalog.spark_catalog.warehouse", "file:///tmp/iceberg_warehouse")
    .getOrCreate()
    .sqlContext
} else {
  throw new RuntimeException("Set either useHive or useIceberg to true")
}

val tables = new TPCDSTables(sqlContext,
  dsdgenDir           = dsdgenDir,
  scaleFactor         = scaleFactor,
  useDoubleForDecimal = true,
  useStringForDate    = false)

val numSplits: Option[Int] = None

if (useHive) {
  if (createExternalHiveTables)
    tables.createExternalTables(rootDir, "parquet", databaseName,
      overwrite = true, discoverPartitions = false)
  else
    tables.createInternalTables(rootDir, "parquet", databaseName,
      overwrite = true, discoverPartitions = true, numSplits = numSplits, sortOnCol = false)
} else if (useIceberg) {
  tables.createIcebergTables(rootDir, "parquet", databaseName,
    overwrite = true, discoverPartitions = false, tableFilter = "", numSplits, sortOnCol = true)
} else {
  throw new RuntimeException("Set either useHive or useIceberg to true")
}