Spark 3.4: Decimal pushdown

ClickHouse · May 22, 2023 · 402f713 · 402f713
1 parent 6ba254a
commit 402f713
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 2 deletions.
diff --git a/...rk-it/src/test/scala/org/apache/spark/sql/clickhouse/single/ClickHouseDataTypeSuite.scala b/...rk-it/src/test/scala/org/apache/spark/sql/clickhouse/single/ClickHouseDataTypeSuite.scala
@@ -159,15 +159,21 @@ class ClickHouseDataTypeSuite extends SparkClickHouseSingleTest {
       testDataType(dataType) { (db, tbl) =>
         runClickHouseSQL(
           s"""INSERT INTO $db.$tbl VALUES
-             |(1, '11.1')
+             |(1, '11.1'),
+             |(2, '22.2')
              |""".stripMargin
         )
       } { df =>
         assert(df.schema.length === 2)
         assert(df.schema.fields(1).dataType === DecimalType(p, s))
         checkAnswer(
           df,
-          Row(1, BigDecimal("11.1", new MathContext(p))) :: Nil
+          Row(1, BigDecimal("11.1", new MathContext(p))) ::
+            Row(2, BigDecimal("22.2", new MathContext(p))) :: Nil
+        )
+        checkAnswer(
+          df.filter("value > 20"),
+          Row(2, BigDecimal("22.2", new MathContext(p))) :: Nil
         )
       }
     }

diff --git a/spark-3.4/clickhouse-spark/src/main/scala/xenon/clickhouse/SQLHelper.scala b/spark-3.4/clickhouse-spark/src/main/scala/xenon/clickhouse/SQLHelper.scala
@@ -20,6 +20,7 @@ import org.apache.commons.lang3.StringUtils
 import org.apache.spark.sql.connector.expressions.aggregate._
 import org.apache.spark.sql.connector.expressions.NamedReference
 import org.apache.spark.sql.sources._
+import org.apache.spark.sql.types.DecimalType
 import org.apache.spark.unsafe.types.UTF8String
 import xenon.clickhouse.Utils._
 
@@ -38,6 +39,9 @@ trait SQLHelper {
     case localDateTime: LocalDateTime => s"'${dateTimeFmt.format(localDateTime)}'"
     case legacyDate: Date => s"'${legacyDateFmt.format(legacyDate)}'"
     case localDate: LocalDate => s"'${dateFmt.format(localDate)}'"
+    case decimal: DecimalType if decimal.precision <= 9 => s"toDecimal32($value, ${decimal.scale})"
+    case decimal: DecimalType if decimal.precision <= 18 => s"toDecimal64($value, ${decimal.scale})"
+    case decimal: DecimalType => s"toDecimal128($value, ${decimal.scale})"
     case array: Array[Any] => array.map(compileValue).mkString(",")
     case _ => value
   }