Spark 3.3: Implement SupportsRuntimeFiltering (#276)

Author: zhouyifan279

spark-3.3/clickhouse-spark-it/src/test/scala/org/apache/spark/sql/clickhouse/cluster/ClickHouseClusterReadSuite.scala

@@ -16,6 +16,8 @@ package org.apache.spark.sql.clickhouse.cluster
 
 import org.apache.spark.sql.clickhouse.ClickHouseSQLConf.READ_DISTRIBUTED_CONVERT_LOCAL
 import org.apache.spark.sql.{AnalysisException, Row}
+import org.apache.spark.sql.catalyst.TableIdentifier
+import org.apache.spark.sql.execution.datasources.v2.BatchScanExec
 
 class ClickHouseClusterReadSuite extends SparkClickHouseClusterTest {
 
@@ -84,4 +86,31 @@ class ClickHouseClusterReadSuite extends SparkClickHouseClusterTest {
       )
     }
   }
+
+  test("runtime filter - distributed table") {
+    withSimpleDistTable("single_replica", "runtime_db", "runtime_tbl", true) { (_, db, tbl_dist, _) =>
+      spark.sql("set spark.clickhouse.read.runtimeFilter.enabled=false")
+      checkAnswer(
+        spark.sql(s"SELECT id FROM $db.$tbl_dist " +
+          s"WHERE id IN (" +
+          s"  SELECT id FROM $db.$tbl_dist " +
+          s"  WHERE DATE_FORMAT(create_time, 'yyyy-MM-dd') between '2021-01-01' and '2022-01-01'" +
+          s")"),
+        Row(1)
+      )
+
+      spark.sql("set spark.clickhouse.read.runtimeFilter.enabled=true")
+      val df = spark.sql(s"SELECT id FROM $db.$tbl_dist " +
+        s"WHERE id IN (" +
+        s"  SELECT id FROM $db.$tbl_dist " +
+        s"  WHERE DATE_FORMAT(create_time, 'yyyy-MM-dd') between '2021-01-01' and '2022-01-01'" +
+        s")")
+      checkAnswer(df, Row(1))
+      val runtimeFilterExists = df.queryExecution.sparkPlan.exists {
+        case BatchScanExec(_, _, runtimeFilters, _) if runtimeFilters.nonEmpty => true
+        case _ => false
+      }
+      assert(runtimeFilterExists)
+    }
+  }
 }

spark-3.3/clickhouse-spark-it/src/test/scala/org/apache/spark/sql/clickhouse/single/ClickHouseSingleSuite.scala

@@ -15,6 +15,8 @@
 package org.apache.spark.sql.clickhouse.single
 
 import org.apache.spark.sql.Row
+import org.apache.spark.sql.catalyst.TableIdentifier
+import org.apache.spark.sql.execution.datasources.v2.BatchScanExec
 import org.apache.spark.sql.types._
 
 class ClickHouseSingleSuite extends SparkClickHouseSingleTest {
@@ -451,4 +453,34 @@ class ClickHouseSingleSuite extends SparkClickHouseSingleTest {
         spark.sql(s"UNCACHE TABLE $db.$tbl")
     }
   }
+
+  test("runtime filter") {
+    val db = "runtime_db"
+    val tbl = "runtime_tbl"
+
+    withSimpleTable(db, tbl, true) {
+      spark.sql("set spark.clickhouse.read.runtimeFilter.enabled=false")
+      checkAnswer(
+        spark.sql(s"SELECT id FROM $db.$tbl " +
+          s"WHERE id IN (" +
+          s"  SELECT id FROM $db.$tbl " +
+          s"  WHERE DATE_FORMAT(create_time, 'yyyy-MM-dd') between '2021-01-01' and '2022-01-01'" +
+          s")"),
+        Row(1)
+      )
+
+      spark.sql("set spark.clickhouse.read.runtimeFilter.enabled=true")
+      val df = spark.sql(s"SELECT id FROM $db.$tbl " +
+        s"WHERE id IN (" +
+        s"  SELECT id FROM $db.$tbl " +
+        s"  WHERE DATE_FORMAT(create_time, 'yyyy-MM-dd') between '2021-01-01' and '2022-01-01'" +
+        s")")
+      checkAnswer(df, Row(1))
+      val runtimeFilterExists = df.queryExecution.sparkPlan.exists {
+        case BatchScanExec(_, _, runtimeFilters, _) if runtimeFilters.nonEmpty => true
+        case _ => false
+      }
+      assert(runtimeFilterExists)
+    }
+  }
 }

spark-3.3/clickhouse-spark/src/main/scala/org/apache/spark/sql/clickhouse/ClickHouseSQLConf.scala

@@ -173,6 +173,13 @@ object ClickHouseSQLConf {
       .transform(_.toLowerCase)
       .createWithDefault("json")
 
+  val RUNTIME_FILTER_ENABLED: ConfigEntry[Boolean] =
+    buildConf("spark.clickhouse.read.runtimeFilter.enabled")
+      .doc("Enable runtime filter for reading.")
+      .version("0.8.0")
+      .booleanConf
+      .createWithDefault(false)
+
   val WRITE_FORMAT: ConfigEntry[String] =
     buildConf("spark.clickhouse.write.format")
       .doc("Serialize format for writing. Supported formats: json, arrow")

spark-3.3/clickhouse-spark/src/main/scala/org/apache/spark/sql/clickhouse/SparkOptions.scala

@@ -48,6 +48,9 @@ class ReadOptions(_options: JMap[String, String]) extends SparkOptions {
 
   def format: String =
     eval(READ_FORMAT.key, READ_FORMAT)
+
+  def runtimeFilterEnabled: Boolean =
+    eval(RUNTIME_FILTER_ENABLED.key, RUNTIME_FILTER_ENABLED)
 }
 
 class WriteOptions(_options: JMap[String, String]) extends SparkOptions {

spark-3.3/clickhouse-spark/src/main/scala/xenon/clickhouse/read/ClickHouseRead.scala

@@ -16,7 +16,7 @@ package xenon.clickhouse.read
 
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.clickhouse.ClickHouseSQLConf._
-import org.apache.spark.sql.connector.expressions.Transform
+import org.apache.spark.sql.connector.expressions.{Expressions, NamedReference, Transform}
 import org.apache.spark.sql.connector.expressions.aggregate.Aggregation
 import org.apache.spark.sql.connector.metric.CustomMetric
 import org.apache.spark.sql.connector.read._
@@ -127,8 +127,14 @@ class ClickHouseScanBuilder(
 
 class ClickHouseBatchScan(scanJob: ScanJobDescription) extends Scan with Batch
     with SupportsReportPartitioning
+    with SupportsRuntimeFiltering
     with PartitionReaderFactory
-    with ClickHouseHelper {
+    with ClickHouseHelper
+    with SQLHelper {
+
+  implicit private val tz: ZoneId = scanJob.tz
+
+  private var runtimeFilters: Array[Filter] = Array.empty
 
   val database: String = scanJob.database
   val table: String = scanJob.table
@@ -187,9 +193,13 @@ class ClickHouseBatchScan(scanJob: ScanJobDescription) extends Scan with Batch
   override def createReader(_partition: InputPartition): PartitionReader[InternalRow] = {
     val format = scanJob.readOptions.format
     val partition = _partition.asInstanceOf[ClickHouseInputPartition]
+    val finalScanJob = scanJob.copy(filtersExpr =
+      scanJob.filtersExpr + " AND "
+        + compileFilters(AlwaysTrue :: runtimeFilters.toList)
+    )
     format match {
-      case "json" => new ClickHouseJsonReader(scanJob, partition)
-      case "binary" => new ClickHouseBinaryReader(scanJob, partition)
+      case "json" => new ClickHouseJsonReader(finalScanJob, partition)
+      case "binary" => new ClickHouseBinaryReader(finalScanJob, partition)
       case unsupported => throw CHClientException(s"Unsupported read format: $unsupported")
     }
   }
@@ -198,4 +208,14 @@ class ClickHouseBatchScan(scanJob: ScanJobDescription) extends Scan with Batch
     BlocksReadMetric(),
     BytesReadMetric()
   )
+
+  override def filterAttributes(): Array[NamedReference] =
+    if (scanJob.readOptions.runtimeFilterEnabled) {
+      scanJob.readSchema.fields.map(field => Expressions.column(field.name))
+    } else {
+      Array.empty
+    }
+
+  override def filter(filters: Array[Filter]): Unit =
+    runtimeFilters = filters
 }