[VL] support uuid function

apache · Mar 19, 2024 · 4968b83 · 4968b83
1 parent a1e0367
commit 4968b83
Show file tree

Hide file tree

Showing 7 changed files with 38 additions and 2 deletions.
diff --git a/backends-velox/src/main/scala/io/glutenproject/backendsapi/velox/SparkPlanExecApiImpl.scala b/backends-velox/src/main/scala/io/glutenproject/backendsapi/velox/SparkPlanExecApiImpl.scala
@@ -36,7 +36,7 @@ import org.apache.spark.sql.catalyst.{AggregateFunctionRewriteRule, FlushableHas
 import org.apache.spark.sql.catalyst.analysis.FunctionRegistry.FunctionBuilder
 import org.apache.spark.sql.catalyst.catalog.BucketSpec
 import org.apache.spark.sql.catalyst.catalog.CatalogTypes.TablePartitionSpec
-import org.apache.spark.sql.catalyst.expressions.{Alias, Ascending, Attribute, Cast, CreateNamedStruct, ElementAt, Expression, ExpressionInfo, Generator, GetArrayItem, GetMapValue, GetStructField, If, IsNaN, Literal, Murmur3Hash, NamedExpression, NaNvl, PosExplode, Round, SortOrder, StringSplit, StringTrim}
+import org.apache.spark.sql.catalyst.expressions.{Alias, Ascending, Attribute, Cast, CreateNamedStruct, ElementAt, Expression, ExpressionInfo, Generator, GetArrayItem, GetMapValue, GetStructField, If, IsNaN, Literal, Murmur3Hash, NamedExpression, NaNvl, PosExplode, Round, SortOrder, StringSplit, StringTrim, Uuid}
 import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, HLLAdapter}
 import org.apache.spark.sql.catalyst.optimizer.BuildSide
 import org.apache.spark.sql.catalyst.plans.JoinType
@@ -128,6 +128,13 @@ class SparkPlanExecApiImpl extends SparkPlanExecApi {
     )
   }
 
+  /** Transform Uuid to Substrait. */
+  override def genUuidTransformer(
+      substraitExprName: String,
+      original: Uuid): ExpressionTransformer = {
+    UuidTransformer(substraitExprName, original)
+  }
+
   /** Transform map_entries to Substrait. */
   override def genMapEntriesTransformer(
       substraitExprName: String,

diff --git a/backends-velox/src/main/scala/io/glutenproject/backendsapi/velox/VeloxBackend.scala b/backends-velox/src/main/scala/io/glutenproject/backendsapi/velox/VeloxBackend.scala
@@ -26,7 +26,7 @@ import io.glutenproject.substrait.rel.LocalFilesNode.ReadFileFormat
 import io.glutenproject.substrait.rel.LocalFilesNode.ReadFileFormat.{DwrfReadFormat, OrcReadFormat, ParquetReadFormat}
 
 import org.apache.spark.sql.catalyst.catalog.BucketSpec
-import org.apache.spark.sql.catalyst.expressions.{Alias, CumeDist, DenseRank, Descending, Expression, Lag, Lead, Literal, NamedExpression, NthValue, NTile, PercentRank, Rand, RangeFrame, Rank, RowNumber, SortOrder, SpecialFrameBoundary, SpecifiedWindowFrame}
+import org.apache.spark.sql.catalyst.expressions.{Alias, CumeDist, DenseRank, Descending, Expression, Lag, Lead, Literal, NamedExpression, NthValue, NTile, PercentRank, Rand, RangeFrame, Rank, RowNumber, SortOrder, SpecialFrameBoundary, SpecifiedWindowFrame, Uuid}
 import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, Count, Sum}
 import org.apache.spark.sql.catalyst.plans.JoinType
 import org.apache.spark.sql.catalyst.util.CharVarcharUtils
@@ -389,6 +389,7 @@ object BackendSettings extends BackendSettingsApi {
         // Block directly falling back the below functions by FallbackEmptySchemaRelation.
         case alias: Alias => checkExpr(alias.child)
         case _: Rand => true
+        case _: Uuid => true
         case _ => false
       }
     }

diff --git a/backends-velox/src/test/scala/io/glutenproject/execution/VeloxFunctionsValidateSuite.scala b/backends-velox/src/test/scala/io/glutenproject/execution/VeloxFunctionsValidateSuite.scala
@@ -469,4 +469,10 @@ class VeloxFunctionsValidateSuite extends VeloxWholeStageTransformerSuite {
       checkOperatorMatch[ProjectExecTransformer]
     }
   }
+
+  test("Test uuid function") {
+    runQueryAndCompare("""SELECT uuid() from lineitem limit 100""".stripMargin, false) {
+      checkOperatorMatch[ProjectExecTransformer]
+    }
+  }
 }
diff --git a/docs/velox-backend-support-progress.md b/docs/velox-backend-support-progress.md
@@ -428,3 +428,4 @@ Gluten supports 199 functions. (Draw to right to see all data types)
 | spark_partition_id            |                        |                       | S      |                        |         |      |       |     |      |       |        |      |           |        |         |      |        |          |       |      |        |      |
 | stack                         |                        |                       |        |                        |         |      |       |     |      |       |        |      |           |        |         |      |        |          |       |      |        |      |
 | xxhash64                      | xxhash64               | xxhash64              |        |                        |         |      |       |     |      |       |        |      |           |        |         |      |        |          |       |      |        |      |
+| uuid                      | uuid               | uuid              | S       |                        |         |      |       |     |      |       |        |      |           |        |         |      |        |          |       |      |        |      |
diff --git a/gluten-core/src/main/scala/io/glutenproject/backendsapi/SparkPlanExecApi.scala b/gluten-core/src/main/scala/io/glutenproject/backendsapi/SparkPlanExecApi.scala
@@ -199,6 +199,10 @@ trait SparkPlanExecApi {
     throw new UnsupportedOperationException("NaNvl is not supported")
   }
 
+  def genUuidTransformer(substraitExprName: String, original: Uuid): ExpressionTransformer = {
+    GenericExpressionTransformer(substraitExprName, Seq(), original)
+  }
+
   /** Transform map_entries to Substrait. */
   def genMapEntriesTransformer(
       substraitExprName: String,

diff --git a/gluten-core/src/main/scala/io/glutenproject/expression/ExpressionConverter.scala b/gluten-core/src/main/scala/io/glutenproject/expression/ExpressionConverter.scala
@@ -526,6 +526,8 @@ object ExpressionConverter extends SQLConfHelper with Logging {
         val childrenTransformers =
           e.children.map(replaceWithExpressionTransformerInternal(_, attributeSeq, expressionsMap))
         e.getTransformer(childrenTransformers)
+      case u: Uuid =>
+        BackendsApiManager.getSparkPlanExecApiInstance.genUuidTransformer(substraitExprName, u)
       case expr =>
         GenericExpressionTransformer(
           substraitExprName,

diff --git a/gluten-core/src/main/scala/io/glutenproject/expression/UnaryExpressionTransformer.scala b/gluten-core/src/main/scala/io/glutenproject/expression/UnaryExpressionTransformer.scala
@@ -239,3 +239,18 @@ case class GetArrayStructFieldsTransformer(
     ExpressionBuilder.makeScalarFunction(functionId, inputNodes, typeNode)
   }
 }
+
+case class UuidTransformer(substraitExprName: String, original: Uuid)
+  extends ExpressionTransformer {
+
+  override def doTransform(args: java.lang.Object): ExpressionNode = {
+    val functionMap = args.asInstanceOf[java.util.HashMap[String, java.lang.Long]]
+    val functionId = ExpressionBuilder.newScalarFunction(
+      functionMap,
+      ConverterUtils.makeFuncName(substraitExprName, Seq(LongType)))
+    val seed = Literal(original.randomSeed.get)
+    val inputNodes = Lists.newArrayList[ExpressionNode](LiteralTransformer(seed).doTransform(args))
+    val typeNode = ConverterUtils.getTypeNode(original.dataType, original.nullable)
+    ExpressionBuilder.makeScalarFunction(functionId, inputNodes, typeNode)
+  }
+}