[Spark] Lightweight compaction for clustered tables

docs/src/content/docs/delta-clustering.mdx

@@ -146,6 +146,14 @@ In Delta Lake 3.3 and above, you can force reclustering of all records in a tabl
 
 Run `OPTIMIZE FULL` when you change clustering columns. If you have previously run `OPTIMIZE FULL` and there has been no change to clustering columns, `OPTIMIZE FULL` runs the same as `OPTIMIZE`. Always use `OPTIMIZE FULL` to ensure that data layout reflects the current clustering columns.
 
+### Lightweight compaction
+
+Each `OPTIMIZE` clusters the new data together with previously clustered data that has not yet reached the target size, and sorts all of it. On tables that receive small amounts of data between frequent `OPTIMIZE` runs, the same data can be rewritten many times. To reduce this cost, set `spark.databricks.delta.optimize.clustering.lightweight.enabled` to `true`. While the table has less unclustered data than `spark.databricks.delta.optimize.clustering.lightweight.maxUnclusteredBytes` (10 GB by default), `OPTIMIZE` and auto compaction then only compact small unclustered files, and leave clustered files untouched:
+
+- Files are grouped by the minimum and maximum values of one clustering column, so each compacted file covers a narrow range of that column. Rows are not sorted.
+- Set `spark.databricks.delta.optimize.clustering.lightweight.column` to choose the column. Otherwise, the clustering column whose files cover the narrowest part of its value range is used, such as a date column on a table written in daily batches. If the files of every clustering column cover more than `spark.databricks.delta.optimize.clustering.lightweight.maxRelativeFileRange` (0.5 by default) of its value range on average, files are grouped by size.
+- The compacted files remain unclustered. Once the table has accumulated `maxUnclusteredBytes` of unclustered data, `OPTIMIZE` clusters it. `OPTIMIZE FULL` always clusters.
+
 ## Read data from a clustered table
 
 You can read data in a clustered table using any Delta Lake client. For best query results, include clustering columns in your query filters, as in the following example:

spark/src/main/resources/error/delta-error-classes.json

@@ -1995,6 +1995,12 @@
     ],
     "sqlState" : "25000"
   },
+  "DELTA_INVALID_LIGHTWEIGHT_CLUSTERING_COLUMN" : {
+    "message" : [
+      "Column '<column>' in '<config>' is not a clustering column of the table. Clustering columns: <clusteringColumns>."
+    ],
+    "sqlState" : "42703"
+  },
   "DELTA_INVALID_LOGSTORE_CONF" : {
     "message" : [
       "(`<classConfig>`) and (`<schemeConfig>`) cannot be set at the same time. Please set only one group of them."

spark/src/main/scala/org/apache/spark/sql/delta/DeltaErrors.scala

@@ -4156,6 +4156,17 @@ trait DeltaErrorsBase
     )
   }
 
+  def invalidLightweightClusteringColumn(
+      column: String, clusteringColumns: Seq[String]): Throwable = {
+    new DeltaAnalysisException(
+      errorClass = "DELTA_INVALID_LIGHTWEIGHT_CLUSTERING_COLUMN",
+      messageParameters = Array(
+        column,
+        DeltaSQLConf.DELTA_OPTIMIZE_CLUSTERING_LIGHTWEIGHT_COLUMN.key,
+        clusteringColumns.mkString(", "))
+    )
+  }
+
   def clusteringColumnUnsupportedDataTypes(clusteringColumnsWithDataTypes: String): Throwable = {
     new DeltaAnalysisException(
       errorClass = "DELTA_CLUSTERING_COLUMNS_DATATYPE_NOT_SUPPORTED",

spark/src/main/scala/org/apache/spark/sql/delta/commands/OptimizeTableCommand.scala

@@ -284,6 +284,9 @@ class OptimizeExecutor(
 
   private val isClusteredTable = ClusteredTableUtils.isSupported(snapshot.protocol)
 
+  private val isLightweightClustering =
+    optimizeStrategy.isInstanceOf[LightweightClusteringStrategy]
+
   private val isMultiDimClustering =
     optimizeStrategy.isInstanceOf[ClusteringStrategy] ||
     optimizeStrategy.isInstanceOf[ZOrderStrategy]
@@ -314,12 +317,12 @@ class OptimizeExecutor(
       // transactions later
       val candidateFiles = snapshot.filesForScan(partitionPredicate, keepNumRecords = true).files
 
-      val filesToProcess = optimizeContext.reorg match {
+      val filesToProcess = optimizeStrategy.prepareFiles(optimizeContext.reorg match {
         case Some(reorgOperation) =>
           reorgOperation.filterFilesToReorg(sparkSession, snapshot, candidateFiles)
         case None =>
           filterCandidateFileList(minFileSize, maxDeletedRowsRatio, candidateFiles)
-      }
+      })
       // Group files by their normalized (typed) partition values so that logically equivalent
       // but differently formatted values (e.g. timestamp variants) end up in the same group.
       val partitionsToCompact = filesToProcess
@@ -608,7 +611,11 @@ class OptimizeExecutor(
         predicate = partitionPredicate,
         zOrderBy = zOrderByColumns,
         auto = isAutoCompact,
-        clusterBy = if (isClusteredTable) Option(clusteringColumns).filter(_.nonEmpty) else None,
+        clusterBy = if (isClusteredTable && !isLightweightClustering) {
+          Option(clusteringColumns).filter(_.nonEmpty)
+        } else {
+          None
+        },
         isFull = optimizeContext.isFull)
     }
   }

spark/src/main/scala/org/apache/spark/sql/delta/commands/OptimizeTableStrategy.scala

@@ -21,13 +21,15 @@ import org.apache.spark.sql.delta.skipping.clustering.ZCube
 import org.apache.spark.sql.delta.{DeltaConfigs, DeltaErrors, OptimisticTransaction, Snapshot}
 import org.apache.spark.sql.delta.actions.{AddFile, DeletionVectorDescriptor, FileAction, RemoveFile}
 import org.apache.spark.sql.delta.commands.OptimizeTableStrategy.DummyBinInfo
-import org.apache.spark.sql.delta.commands.optimize.{AddFileWithNumRecords, DeletionVectorStats, OptimizeStats, ZOrderFileStats, ZOrderStats}
+import org.apache.spark.sql.delta.commands.optimize.{AddFileWithNumRecords, ClusteringColumnFileOrdering, DeletionVectorStats, OptimizeStats, ZOrderFileStats, ZOrderStats}
+import org.apache.spark.sql.delta.metering.DeltaLogging
 import org.apache.spark.sql.delta.sources.DeltaSQLConf
 import org.apache.spark.sql.delta.sources.DeltaSQLConf.{DELTA_OPTIMIZE_CLUSTERING_MIN_CUBE_SIZE, DELTA_OPTIMIZE_CLUSTERING_TARGET_CUBE_SIZE}
 import org.apache.spark.sql.delta.zorder.ZCubeInfo
 
 import org.apache.spark.sql.SparkSession
 import org.apache.spark.sql.catalyst.expressions.Literal
+import org.apache.spark.sql.functions.{coalesce, col, lit, sum}
 
 object OptimizeTableMode extends Enumeration {
   type OptimizeTableMode = Value
@@ -46,6 +48,14 @@ trait OptimizeTableStrategy {
    */
   def maxBinSize: Long
 
+  /**
+   * Utility method to prepare all files selected for optimization, before they are grouped by
+   * partition and passed to [[prepareFilesPerPartition]].
+   *
+   * @return Prepared files for the subsequent optimization.
+   */
+  def prepareFiles(inputFiles: Seq[AddFile]): Seq[AddFile] = inputFiles
+
   /**
    * Utility method to prepare files in a partition for optimization.
    *
@@ -93,7 +103,7 @@ trait OptimizeTableStrategy {
       bins: Seq[Bin]): Unit
 }
 
-object OptimizeTableStrategy {
+object OptimizeTableStrategy extends DeltaLogging {
   // A trait representing the context for a Bin.
   sealed trait BinInfo
 
@@ -109,16 +119,47 @@ object OptimizeTableStrategy {
       optimizeContext: DeltaOptimizeContext,
       zOrderBy: Seq[String]): OptimizeTableStrategy = getMode(snapshot, zOrderBy) match {
     case OptimizeTableMode.CLUSTERING =>
-      ClusteringStrategy(
-        sparkSession,
-        ClusteringColumnInfo.extractLogicalNames(snapshot),
-        optimizeContext)
+      val clusteringColumns = ClusteringColumnInfo.extractLogicalNames(snapshot)
+      if (useLightweightClustering(sparkSession, snapshot, optimizeContext)) {
+        LightweightClusteringStrategy(sparkSession, snapshot, clusteringColumns, optimizeContext)
+      } else {
+        ClusteringStrategy(sparkSession, clusteringColumns, optimizeContext)
+      }
     case OptimizeTableMode.ZORDER => ZOrderStrategy(sparkSession, zOrderBy)
     case OptimizeTableMode.COMPACTION =>
       CompactionStrategy(sparkSession, optimizeContext)
     case other => throw new UnsupportedOperationException(s"Unsupported mode $other")
   }
 
+  /**
+   * Whether OPTIMIZE on a clustered table should run lightweight compaction instead of
+   * clustering: the table has some unclustered data, but less than
+   * [[DeltaSQLConf.DELTA_OPTIMIZE_CLUSTERING_LIGHTWEIGHT_MAX_UNCLUSTERED_BYTES]]. Without any
+   * unclustered data, clustering still runs so that partial Z-cubes can be merged.
+   */
+  private def useLightweightClustering(
+      sparkSession: SparkSession,
+      snapshot: Snapshot,
+      optimizeContext: DeltaOptimizeContext): Boolean = {
+    val conf = sparkSession.sessionState.conf
+    if (!conf.getConf(DeltaSQLConf.DELTA_OPTIMIZE_CLUSTERING_LIGHTWEIGHT_ENABLED) ||
+        optimizeContext.isFull || optimizeContext.reorg.nonEmpty) {
+      return false
+    }
+    val unclusteredBytes = snapshot.allFiles
+      .where(col("clusteringProvider").isNull)
+      .select(coalesce(sum(col("size")), lit(0L)))
+      .head()
+      .getLong(0)
+    val threshold =
+      conf.getConf(DeltaSQLConf.DELTA_OPTIMIZE_CLUSTERING_LIGHTWEIGHT_MAX_UNCLUSTERED_BYTES)
+    val useLightweight = unclusteredBytes > 0 && unclusteredBytes < threshold
+    logInfo(s"Clustered table has $unclusteredBytes bytes of unclustered data, threshold for " +
+      s"clustering is $threshold bytes: " +
+      (if (useLightweight) "running lightweight compaction" else "clustering"))
+    useLightweight
+  }
+
   private def getMode(snapshot: Snapshot, zOrderBy: Seq[String]): OptimizeTableMode.Value = {
     val isClusteredTable = ClusteredTableUtils.isSupported(snapshot.protocol)
     val hasClusteringColumns = ClusteringColumnInfo.extractLogicalNames(snapshot).nonEmpty
@@ -157,6 +198,59 @@ case class CompactionStrategy(
       bins: Seq[Bin]): Unit = {}
 }
 
+/**
+ * Implements lightweight compaction for clustered tables, used instead of [[ClusteringStrategy]]
+ * while the table has little unclustered data (see
+ * [[DeltaSQLConf.DELTA_OPTIMIZE_CLUSTERING_LIGHTWEIGHT_ENABLED]]).
+ *
+ * Clustering rewrites the new data together with all partial Z-cubes and sorts every row, so
+ * running it often on small amounts of new data rewrites the same data many times. Instead, this
+ * strategy only compacts small unclustered files, like [[CompactionStrategy]], and leaves
+ * clustered files untouched. The files are grouped by the min/max statistics of one clustering
+ * column (see [[ClusteringColumnFileOrdering]]) so that each output file covers a narrow range of
+ * that column, and rows are not sorted. The output files stay unclustered, so they are clustered
+ * once the table has accumulated enough unclustered data.
+ */
+case class LightweightClusteringStrategy(
+    override val sparkSession: SparkSession,
+    snapshot: Snapshot,
+    clusteringColumns: Seq[String],
+    optimizeContext: DeltaOptimizeContext) extends OptimizeTableStrategy with DeltaLogging {
+
+  override val optimizeTableMode: OptimizeTableMode.Value = OptimizeTableMode.COMPACTION
+
+  override val maxBinSize: Long = {
+    optimizeContext.maxFileSize.getOrElse(
+      sparkSession.sessionState.conf.getConf(DeltaSQLConf.DELTA_OPTIMIZE_MAX_FILE_SIZE))
+  }
+
+  override def curve: String = {
+    throw new UnsupportedOperationException("Lightweight compaction doesn't sort rows.")
+  }
+
+  override def prepareFiles(inputFiles: Seq[AddFile]): Seq[AddFile] = {
+    val conf = sparkSession.sessionState.conf
+    val (orderedFiles, column) = ClusteringColumnFileOrdering.order(
+      sparkSession,
+      snapshot,
+      inputFiles.filter(_.clusteringProvider.isEmpty),
+      clusteringColumns,
+      conf.getConf(DeltaSQLConf.DELTA_OPTIMIZE_CLUSTERING_LIGHTWEIGHT_COLUMN),
+      conf.getConf(DeltaSQLConf.DELTA_OPTIMIZE_CLUSTERING_LIGHTWEIGHT_MAX_RELATIVE_FILE_RANGE))
+    logInfo(s"Lightweight compaction of ${orderedFiles.size} unclustered files, grouped by " +
+      column.map(c => s"the statistics of column $c").getOrElse("size"))
+    orderedFiles
+  }
+
+  // The files were already ordered by prepareFiles. Clustered tables are unpartitioned.
+  override def prepareFilesPerPartition(inputFiles: Seq[AddFile]): Seq[AddFile] = inputFiles
+
+  override def updateOptimizeStats(
+      optimizeStats: OptimizeStats,
+      removedFiles: Seq[RemoveFile],
+      bins: Seq[Bin]): Unit = {}
+}
+
 /** Implements ZOrder strategy */
 case class ZOrderStrategy(
     override val sparkSession: SparkSession,

spark/src/main/scala/org/apache/spark/sql/delta/commands/optimize/ClusteringColumnFileOrdering.scala

@@ -0,0 +1,161 @@
+/*
+ * Copyright (2021) The Delta Lake Project Authors.
+ *
+ * Licensed 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.sql.delta.commands.optimize
+
+import java.util.{Arrays, Comparator}
+
+import org.apache.spark.sql.delta.{DeltaErrors, Snapshot}
+import org.apache.spark.sql.delta.actions.AddFile
+import org.apache.spark.sql.delta.stats.DeltaStatistics.{MAX, MIN}
+
+import org.apache.spark.sql.{Column, Encoders, SparkSession}
+import org.apache.spark.sql.catalyst.analysis.UnresolvedAttribute
+import org.apache.spark.sql.functions.col
+
+/**
+ * Orders the files of a clustered table for lightweight compaction by the per-file min/max
+ * statistics of a single clustering column.
+ *
+ * Compaction packs files into bins sequentially in the given order and writes each bin out as one
+ * file without sorting rows. Ordering the files by the (min, max) values of a column makes each
+ * bin cover a contiguous range of that column, so data skipping on the column keeps working after
+ * compaction. Only one column is used: grouping whole files cannot keep several columns narrow at
+ * once, and clustering takes care of all clustering columns later.
+ */
+object ClusteringColumnFileOrdering {
+
+  /** The min and max values of a column in a file. Both are null without statistics. */
+  case class ValueRange(min: Any, max: Any) {
+    def isKnown: Boolean = min != null && max != null
+  }
+
+  /**
+   * Returns the files ordered for lightweight compaction, and the clustering column used to
+   * order them. If `column` is not set, the clustering column whose files cover the narrowest
+   * part of its value range is used, provided that [[relativeFileRange]] is at most
+   * `maxRelativeFileRange`. If no column qualifies, the files are ordered by size.
+   */
+  def order(
+      spark: SparkSession,
+      snapshot: Snapshot,
+      files: Seq[AddFile],
+      clusteringColumns: Seq[String],
+      column: Option[String],
+      maxRelativeFileRange: Double): (Seq[AddFile], Option[String]) = {
+    val columns = column match {
+      case Some(c) => Seq(resolve(c, clusteringColumns))
+      case None => clusteringColumns
+    }
+    if (files.size < 2 || columns.isEmpty) return (files.sortBy(_.size), None)
+
+    val ranges = loadRanges(spark, snapshot, files, columns.map(nameParts))
+    def rangesOf(i: Int): Seq[ValueRange] =
+      files.map(f => ranges.get(f.path).map(_(i)).getOrElse(ValueRange(null, null)))
+
+    val chosen = if (column.isDefined) {
+      Some(0)
+    } else {
+      columns.indices
+        .flatMap(i => relativeFileRange(rangesOf(i)).map(i -> _))
+        .filter(_._2 <= maxRelativeFileRange)
+        .sortBy(_._2)
+        .headOption
+        .map(_._1)
+    }
+    chosen match {
+      case Some(i) => (sort(files.zip(rangesOf(i))), Some(columns(i)))
+      case None => (files.sortBy(_.size), None)
+    }
+  }
+
+  /**
+   * Returns how much of the value range of a column a file covers on average, from 0 when no two
+   * files overlap to 1 when every file covers all values. Values are replaced by their rank among
+   * the distinct min/max values of all files, so the measure works for any orderable type and
+   * reflects how much files overlap rather than how values are distributed. Returns None if fewer
+   * than two files have statistics or all values are equal.
+   */
+  private[delta] def relativeFileRange(ranges: Seq[ValueRange]): Option[Double] = {
+    val known = ranges.filter(_.isKnown)
+    if (known.size < 2) return None
+    val values = known.flatMap(r => Seq(r.min, r.max)).map(_.asInstanceOf[AnyRef]).toArray
+    Arrays.sort(values, valueComparator)
+    val distinct = values.foldLeft(Vector.empty[AnyRef]) { (acc, v) =>
+      if (acc.nonEmpty && valueComparator.compare(acc.last, v) == 0) acc else acc :+ v
+    }.toArray
+    if (distinct.length < 2) return None
+    def rank(v: Any): Int = Arrays.binarySearch(distinct, v.asInstanceOf[AnyRef], valueComparator)
+    val totalWidth = known.map(r => (rank(r.max) - rank(r.min)).toLong).sum
+    Some(totalWidth.toDouble / known.size / (distinct.length - 1))
+  }
+
+  /** Orders files by (min, max), with files without statistics last, ordered by size. */
+  private def sort(files: Seq[(AddFile, ValueRange)]): Seq[AddFile] = {
+    files.sortWith { case ((f1, r1), (f2, r2)) =>
+      val c = compareRanges(r1, r2)
+      if (c != 0) c < 0 else if (f1.size != f2.size) f1.size < f2.size else f1.path < f2.path
+    }.map(_._1)
+  }
+
+  private[delta] def compareRanges(r1: ValueRange, r2: ValueRange): Int = {
+    if (!r1.isKnown || !r2.isKnown) return java.lang.Boolean.compare(!r1.isKnown, !r2.isKnown)
+    val c = valueComparator.compare(r1.min.asInstanceOf[AnyRef], r2.min.asInstanceOf[AnyRef])
+    if (c != 0) c else {
+      valueComparator.compare(r1.max.asInstanceOf[AnyRef], r2.max.asInstanceOf[AnyRef])
+    }
+  }
+
+  private val valueComparator: Comparator[AnyRef] = new Comparator[AnyRef] {
+    override def compare(a: AnyRef, b: AnyRef): Int =
+      a.asInstanceOf[Comparable[AnyRef]].compareTo(b)
+  }
+
+  /** Loads the min/max values of the given columns, by file path. */
+  private def loadRanges(
+      spark: SparkSession,
+      snapshot: Snapshot,
+      files: Seq[AddFile],
+      columns: Seq[Seq[String]]): Map[String, IndexedSeq[ValueRange]] = {
+    val statsColumns: Seq[Column] = columns.flatMap { parts =>
+      Seq(MIN, MAX).map(snapshot.getStatsColumnOrNullLiteral(_, parts.reverse))
+    }
+    val paths = spark.createDataset(files.map(_.path).distinct)(Encoders.STRING).toDF("path")
+    snapshot.withStats
+      .join(paths, "path")
+      .select(col("path") +: statsColumns: _*)
+      .collect()
+      .map { row =>
+        row.getString(0) ->
+          columns.indices.map(i => ValueRange(row.get(1 + 2 * i), row.get(2 + 2 * i)))
+      }
+      .toMap
+  }
+
+  private def nameParts(column: String): Seq[String] =
+    UnresolvedAttribute.quotedString(column).nameParts
+
+  private def resolve(column: String, clusteringColumns: Seq[String]): String = {
+    val parts = nameParts(column)
+    def matches(c: String): Boolean = {
+      val other = nameParts(c)
+      other.size == parts.size && other.zip(parts).forall { case (a, b) => a.equalsIgnoreCase(b) }
+    }
+    clusteringColumns.find(matches).getOrElse {
+      throw DeltaErrors.invalidLightweightClusteringColumn(column, clusteringColumns)
+    }
+  }
+}