override date comparision

apache · Nov 3, 2023 · 17c1f34 · 17c1f34
1 parent 81fae1e
commit 17c1f34
Show file tree

Hide file tree

Showing 4 changed files with 376 additions and 15 deletions.
diff --git a/...lickhouse/src/main/scala/io/glutenproject/backendsapi/clickhouse/CHSparkPlanExecApi.scala b/...lickhouse/src/main/scala/io/glutenproject/backendsapi/clickhouse/CHSparkPlanExecApi.scala
@@ -49,6 +49,7 @@ import org.apache.spark.sql.execution.joins.{BuildSideRelation, ClickHouseBuildS
 import org.apache.spark.sql.execution.metric.SQLMetric
 import org.apache.spark.sql.execution.utils.CHExecUtil
 import org.apache.spark.sql.extension.ClickHouseAnalysis
+import org.apache.spark.sql.extension.RewriteDateTimestampComparisonRule
 import org.apache.spark.sql.types.StructType
 import org.apache.spark.sql.vectorized.ColumnarBatch
 
@@ -328,7 +329,12 @@ class CHSparkPlanExecApi extends SparkPlanExecApi {
    * @return
    */
   override def genExtendedAnalyzers(): List[SparkSession => Rule[LogicalPlan]] = {
-    List(spark => new ClickHouseAnalysis(spark, spark.sessionState.conf))
+    val analyzers = List(spark => new ClickHouseAnalysis(spark, spark.sessionState.conf))
+    if (GlutenConfig.getConf.enableDateTimestampComparison) {
+      analyzers :+ (spark => new RewriteDateTimestampComparisonRule(spark, spark.sessionState.conf))
+    } else {
+      analyzers
+    }
   }
 
   /**

diff --git a/...ds-clickhouse/src/test/scala/io/glutenproject/execution/GlutenFunctionValidateSuite.scala b/...ds-clickhouse/src/test/scala/io/glutenproject/execution/GlutenFunctionValidateSuite.scala
@@ -113,26 +113,27 @@ class GlutenFunctionValidateSuite extends GlutenClickHouseWholeStageTransformerS
 
     val dateSchema = StructType(
       Array(
+        StructField("ts", IntegerType, true),
         StructField("day", DateType, true),
         StructField("weekday_abbr", StringType, true)
       )
     )
     val dateRows = sparkContext.parallelize(
       Seq(
-        Row(Date.valueOf("2019-01-01"), "MO"),
-        Row(Date.valueOf("2019-01-01"), "TU"),
-        Row(Date.valueOf("2019-01-01"), "TH"),
-        Row(Date.valueOf("2019-01-01"), "WE"),
-        Row(Date.valueOf("2019-01-01"), "FR"),
-        Row(Date.valueOf("2019-01-01"), "SA"),
-        Row(Date.valueOf("2019-01-01"), "SU"),
-        Row(Date.valueOf("2019-01-01"), "MO"),
-        Row(Date.valueOf("2019-01-02"), "MM"),
-        Row(Date.valueOf("2019-01-03"), "TH"),
-        Row(Date.valueOf("2019-01-04"), "WE"),
-        Row(Date.valueOf("2019-01-05"), "FR"),
-        Row(null, "SA"),
-        Row(Date.valueOf("2019-01-07"), null)
+        Row(1546309380, Date.valueOf("2019-01-01"), "MO"),
+        Row(1546273380, Date.valueOf("2019-01-01"), "TU"),
+        Row(1546358340, Date.valueOf("2019-01-01"), "TH"),
+        Row(1546311540, Date.valueOf("2019-01-01"), "WE"),
+        Row(1546308540, Date.valueOf("2019-01-01"), "FR"),
+        Row(1546319340, Date.valueOf("2019-01-01"), "SA"),
+        Row(1546319940, Date.valueOf("2019-01-01"), "SU"),
+        Row(1546323545, Date.valueOf("2019-01-01"), "MO"),
+        Row(1546409940, Date.valueOf("2019-01-02"), "MM"),
+        Row(1546496340, Date.valueOf("2019-01-03"), "TH"),
+        Row(1546586340, Date.valueOf("2019-01-04"), "WE"),
+        Row(1546676341, Date.valueOf("2019-01-05"), "FR"),
+        Row(null, null, "SA"),
+        Row(1546849141, Date.valueOf("2019-01-07"), null)
       )
     )
     val dateTableFile = Files.createTempFile("", ".parquet").toFile
@@ -466,4 +467,22 @@ class GlutenFunctionValidateSuite extends GlutenClickHouseWholeStageTransformerS
       "select round(0.41875f * id , 4) from range(10);"
     )(checkOperatorMatch[ProjectExecTransformer])
   }
+
+  test("test date comparision expression override") {
+    runQueryAndCompare(
+      "select * from date_table where to_date(from_unixtime(ts)) < '2019-01-02'",
+      noFallBack = true) { _ => }
+    runQueryAndCompare(
+      "select * from date_table where to_date(from_unixtime(ts)) <= '2019-01-02'",
+      noFallBack = true) { _ => }
+    runQueryAndCompare(
+      "select * from date_table where to_date(from_unixtime(ts)) > '2019-01-02'",
+      noFallBack = true) { _ => }
+    runQueryAndCompare(
+      "select * from date_table where to_date(from_unixtime(ts)) >= '2019-01-02'",
+      noFallBack = true) { _ => }
+    runQueryAndCompare(
+      "select * from date_table where to_date(from_unixtime(ts)) = '2019-01-01'",
+      noFallBack = true) { _ => }
+  }
 }