[BLAZE-773] Support long type for floor function (#774)

kwai · Jan 20, 2025 · 829ce1c · 829ce1c
1 parent ef5b706
commit 829ce1c
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 17 deletions.
diff --git a/spark-extension-shims-spark3/src/test/scala/org/apache/spark/sql/blaze/BlazeQuerySuite.scala b/spark-extension-shims-spark3/src/test/scala/org/apache/spark/sql/blaze/BlazeQuerySuite.scala
@@ -86,19 +86,24 @@ class BlazeQuerySuite extends org.apache.spark.sql.QueryTest with BaseBlazeSQLSu
       sql("create table t1(c1 binary, c2 int) using parquet")
       sql(
         "insert into t1 values (to_binary('test1', 'utf-8'), 1), (to_binary('test2', 'utf-8'), 2)")
-      spark.table("t1").printSchema()
       val df = sql("select c2 from t1 order by c1")
       checkAnswer(df, Seq(Row(1), Row(2)))
     }
   }
-  
+
   test("log function with negative input") {
     withTable("t1") {
       sql("create table t1 using parquet as select -1 as c1")
-      spark.table("t1").printSchema()
       val df = sql("select ln(c1) from t1")
-      df.show()
       checkAnswer(df, Seq(Row(null)))
     }
   }
+
+  test("floor function with long input") {
+    withTable("t1") {
+      sql("create table t1 using parquet as select 1L as c1, 2.2 as c2")
+      val df = sql("select floor(c1), floor(c2) from t1")
+      checkAnswer(df, Seq(Row(1, 2)))
+    }
+  }
 }
diff --git a/spark-extension/src/main/scala/org/apache/spark/sql/blaze/NativeConverters.scala b/spark-extension/src/main/scala/org/apache/spark/sql/blaze/NativeConverters.scala
@@ -827,13 +827,17 @@ object NativeConverters extends Logging {
       case e: Log10 =>
         buildScalarFunction(pb.ScalarFunction.Log10, e.children.map(nullIfNegative), e.dataType)
       case e: Floor if !e.dataType.isInstanceOf[DecimalType] =>
-        buildExprNode {
-          _.setTryCast(
-            pb.PhysicalTryCastNode
-              .newBuilder()
-              .setExpr(buildScalarFunction(pb.ScalarFunction.Floor, e.children, e.dataType))
-              .setArrowType(convertDataType(e.dataType))
-              .build())
+        if (e.child.dataType.isInstanceOf[LongType]) {
+          convertExprWithFallback(e.child, isPruningExpr, fallback)
+        } else {
+          buildExprNode {
+            _.setTryCast(
+              pb.PhysicalTryCastNode
+                .newBuilder()
+                .setExpr(buildScalarFunction(pb.ScalarFunction.Floor, e.children, e.dataType))
+                .setArrowType(convertDataType(e.dataType))
+                .build())
+          }
         }
       case e: Ceil if !e.dataType.isInstanceOf[DecimalType] =>
         buildExprNode {

diff --git a/...ension/src/main/scala/org/apache/spark/sql/execution/blaze/arrowio/ArrowFFIExporter.scala b/...ension/src/main/scala/org/apache/spark/sql/execution/blaze/arrowio/ArrowFFIExporter.scala
@@ -67,13 +67,11 @@ class ArrowFFIExporter(rowIter: Iterator[InternalRow], schema: StructType) {
       callRowIter(exportArrowArrayPtr)
     } else {
       // otherwise, process rows as native user
-      nativeCurrentUser.doAs(
-        new PrivilegedExceptionAction[Boolean] {
-          override def run(): Boolean = {
-            callRowIter(exportArrowArrayPtr)
-          }
+      nativeCurrentUser.doAs(new PrivilegedExceptionAction[Boolean] {
+        override def run(): Boolean = {
+          callRowIter(exportArrowArrayPtr)
         }
-      )
+      })
     }
   }