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

// ── Configure these before running ──────────────────────────────────────────
val useHive    = true   // Hive/Parquet tables
val useIceberg = false  // Iceberg tables (requires iceberg-tabbydb-runtime jar)
val createExternalHiveTables = true  // true = external tables pointing to rootDir
                                     // false = internal (managed) tables

val rootDir      = "file:///path/to/tpcds/data"  // where generated data will be written
val databaseName = "tpcds"                        // metastore database name
val scaleFactor  = "50"                           // dataset size in GB
val dsdgenDir    = "/path/to/tpcds-kit/tools"    // tpcds-kit/tools built in Step 1
// ────────────────────────────────────────────────────────────────────────────

val sqlContext = if (useHive) {
  SparkSession.builder()
    .appName("tpcds-generate")
    .enableHiveSupport()
    .getOrCreate()
    .sqlContext
} else if (useIceberg) {
  SparkSession.builder()
    .appName("tpcds-generate")
    .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)

// Generate raw TPC-DS data.
// Non-partitioned tables sorted on the date column per local partition —
// this is 6-7x faster than partitioned and required for TabbyDB's BHJ pushdown gains.
tables.genData(
  location                     = rootDir,
  format                       = "parquet",
  overwrite                    = true,
  partitionTables              = false,
  clusterByPartitionColumns    = false,
  filterOutNullPartitionValues = true,
  tableFilter                  = "",        // "" = all tables
  numPartitions                = 100,
  sortOnCol                    = true)

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")
}