From c01a877d930ac2774e1bd5fd9f33a3a28896acaf Mon Sep 17 00:00:00 2001 From: Hyukjin Kwon Date: Tue, 17 May 2022 18:32:23 +0900 Subject: [PATCH 1/2] Use function registry in the SparkSession --- .../sql/BloomFilterAggregateQuerySuite.scala | 34 ++++++++++--------- .../org/apache/spark/sql/SQLQuerySuite.scala | 3 ++ 2 files changed, 21 insertions(+), 16 deletions(-) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/BloomFilterAggregateQuerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/BloomFilterAggregateQuerySuite.scala index 7fc89ecc88ba3..05513cddccb86 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/BloomFilterAggregateQuerySuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/BloomFilterAggregateQuerySuite.scala @@ -18,7 +18,6 @@ package org.apache.spark.sql import org.apache.spark.sql.catalyst.FunctionIdentifier -import org.apache.spark.sql.catalyst.analysis.FunctionRegistry import org.apache.spark.sql.catalyst.expressions._ import org.apache.spark.sql.catalyst.expressions.aggregate.BloomFilterAggregate import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanExec @@ -35,23 +34,26 @@ class BloomFilterAggregateQuerySuite extends QueryTest with SharedSparkSession { val funcId_bloom_filter_agg = new FunctionIdentifier("bloom_filter_agg") val funcId_might_contain = new FunctionIdentifier("might_contain") - // Register 'bloom_filter_agg' to builtin. - FunctionRegistry.builtin.registerFunction(funcId_bloom_filter_agg, - new ExpressionInfo(classOf[BloomFilterAggregate].getName, "bloom_filter_agg"), - (children: Seq[Expression]) => children.size match { - case 1 => new BloomFilterAggregate(children.head) - case 2 => new BloomFilterAggregate(children.head, children(1)) - case 3 => new BloomFilterAggregate(children.head, children(1), children(2)) - }) - - // Register 'might_contain' to builtin. - FunctionRegistry.builtin.registerFunction(funcId_might_contain, - new ExpressionInfo(classOf[BloomFilterMightContain].getName, "might_contain"), - (children: Seq[Expression]) => BloomFilterMightContain(children.head, children(1))) + override def beforeAll(): Unit = { + super.beforeAll() + // Register 'bloom_filter_agg' to builtin. + spark.sessionState.functionRegistry.registerFunction(funcId_bloom_filter_agg, + new ExpressionInfo(classOf[BloomFilterAggregate].getName, "bloom_filter_agg"), + (children: Seq[Expression]) => children.size match { + case 1 => new BloomFilterAggregate(children.head) + case 2 => new BloomFilterAggregate(children.head, children(1)) + case 3 => new BloomFilterAggregate(children.head, children(1), children(2)) + }) + + // Register 'might_contain' to builtin. + spark.sessionState.functionRegistry.registerFunction(funcId_might_contain, + new ExpressionInfo(classOf[BloomFilterMightContain].getName, "might_contain"), + (children: Seq[Expression]) => BloomFilterMightContain(children.head, children(1))) + } override def afterAll(): Unit = { - FunctionRegistry.builtin.dropFunction(funcId_bloom_filter_agg) - FunctionRegistry.builtin.dropFunction(funcId_might_contain) + spark.sessionState.functionRegistry.dropFunction(funcId_bloom_filter_agg) + spark.sessionState.functionRegistry.dropFunction(funcId_might_contain) super.afterAll() } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala index 07ec42902501c..806b145297068 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala @@ -147,6 +147,9 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark test("SPARK-14415: All functions should have own descriptions") { for (f <- spark.sessionState.functionRegistry.listFunction()) { if (!Seq("cube", "grouping", "grouping_id", "rollup").contains(f.unquotedString)) { + if (f.unquotedString == "bloom_filter_agg") { + null + } checkKeywordsNotExist(sql(s"describe function $f"), "N/A.") } } From fe1d1eaf6a4aaa9f1647f0f0c77cb2cc054d1061 Mon Sep 17 00:00:00 2001 From: Hyukjin Kwon Date: Tue, 17 May 2022 19:04:31 +0900 Subject: [PATCH 2/2] Remove if-else checkpoint --- .../src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala | 3 --- 1 file changed, 3 deletions(-) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala index 806b145297068..07ec42902501c 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala @@ -147,9 +147,6 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark test("SPARK-14415: All functions should have own descriptions") { for (f <- spark.sessionState.functionRegistry.listFunction()) { if (!Seq("cube", "grouping", "grouping_id", "rollup").contains(f.unquotedString)) { - if (f.unquotedString == "bloom_filter_agg") { - null - } checkKeywordsNotExist(sql(s"describe function $f"), "N/A.") } }