pass 4

Author: LuciferYang

connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/DataFrameStatFunctions.scala

@@ -652,11 +652,7 @@ final class DataFrameStatFunctions private[sql] (sparkSession: SparkSession, roo
       numBits: Long,
       fpp: Double): BloomFilter = {
 
-    val agg = if (!fpp.isNaN) {
-      Column.fn("bloom_filter_agg", col, lit(expectedNumItems), lit(fpp))
-    } else {
-      Column.fn("bloom_filter_agg", col, lit(expectedNumItems), lit(numBits))
-    }
+    val agg = Column.fn("bloom_filter_agg", col, lit(expectedNumItems), lit(numBits), lit(fpp))
 
     val ds = sparkSession.newDataset(BinaryEncoder) { builder =>
       builder.getProjectBuilder

connector/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/SparkConnectPlanner.scala

@@ -1732,49 +1732,72 @@ class SparkConnectPlanner(val sessionHolder: SessionHolder) extends Logging {
         val ignoreNulls = extractBoolean(children(3), "ignoreNulls")
         Some(Lead(children.head, children(1), children(2), ignoreNulls))
 
-      case "bloom_filter_agg" if fun.getArgumentsCount == 3 =>
-        // [col, expectedNumItems: Long, numBits: Long] or
-        // [col, expectedNumItems: Long, fpp: Double]
+      case "bloom_filter_agg" if fun.getArgumentsCount == 4 =>
+        // [col, expectedNumItems: Long, numBits: Long, fpp: Double]
         val children = fun.getArgumentsList.asScala.map(transformExpression)
 
-        // Check expectedNumItems > 0L
-        val expectedNumItemsExpr = children(1)
-        val expectedNumItems = expectedNumItemsExpr match {
-          case Literal(l: Long, LongType) => l
+        val fpp = children(3) match {
+          case DoubleLiteral(d) => d
           case _ =>
-            throw InvalidPlanInput("Expected insertions must be long literal.")
-        }
-        if (expectedNumItems <= 0L) {
-          throw InvalidPlanInput("Expected insertions must be positive.")
+            throw InvalidPlanInput("False positive must be double literal.")
         }
 
-        val numberBitsOrFpp = children(2)
+        if (fpp.isNaN) {
+          // Use expectedNumItems and numBits when `fpp.isNaN` if true.
+          // Check expectedNumItems > 0L
+          val expectedNumItemsExpr = children(1)
+          expectedNumItemsExpr match {
+            case Literal(l: Long, LongType) =>
+              if (l <= 0L) {
+                throw InvalidPlanInput("Expected insertions must be positive.")
+              }
+            case _ =>
+              throw InvalidPlanInput("Expected insertions must be long literal.")
+          }
+          // Check numBits > 0L
+          val numBitsExpr = children(2)
+          val numBits = numBitsExpr match {
+            case Literal(l: Long, LongType) => l
+            case _ =>
+              throw InvalidPlanInput("Number of bits must be long literal.")
+          }
+          if (numBits <= 0L) {
+            throw InvalidPlanInput("Number of bits must be positive.")
+          }
+          // Create BloomFilterAggregate with expectedNumItemsExpr and numBitsExpr.
+          Some(
+            new BloomFilterAggregate(children.head, expectedNumItemsExpr, numBitsExpr)
+              .toAggregateExpression())
 
-        val numBitsExpr = numberBitsOrFpp match {
-          case Literal(numBits: Long, LongType) =>
-            // Check numBits > 0L
-            if (numBits <= 0L) {
-              throw InvalidPlanInput("Number of bits must be positive.")
-            }
-            numberBitsOrFpp
-          case DoubleLiteral(fpp) =>
-            // Check fpp not NaN and in  (0.0, 1.0).
-            if (fpp.isNaN || fpp <= 0d || fpp >= 1d) {
-              throw InvalidPlanInput("False positive probability must be within range (0.0, 1.0)")
-            }
-            // Calculate numBits through expectedNumItems and fpp,
-            // refer to `BloomFilter.optimalNumOfBits(long, double)`.
-            val numBits = BloomFilterHelper.optimalNumOfBits(expectedNumItems, fpp)
-            if (numBits <= 0L) {
-              throw InvalidPlanInput("Number of bits must be positive")
-            }
-            Literal(numBits, LongType)
-          case _ =>
-            throw InvalidPlanInput("The 3rd parameter must be double or long literal.")
+        } else if (fpp <= 0d || fpp >= 1d) {
+          // fpp must in (0.0, 1.0).
+          throw InvalidPlanInput("False positive probability must be within range (0.0, 1.0)")
+        } else {
+          // Use expectedNumItems and fpp when `fpp.isNaN` if false and `fpp` in (0.0, 1.0).
+          // Check expectedNumItems > 0L and extract expectedNumItems value.
+          val expectedNumItemsExpr = children(1)
+          val expectedNumItems = expectedNumItemsExpr match {
+            case Literal(l: Long, LongType) => l
+            case _ =>
+              throw InvalidPlanInput("Expected insertions must be long literal.")
+          }
+          if (expectedNumItems <= 0L) {
+            throw InvalidPlanInput("Expected insertions must be positive.")
+          }
+
+          // Calculate numBits through expectedNumItems and fpp, numBits must be greater than 0.
+          val numBits = BloomFilterHelper.optimalNumOfBits(expectedNumItems, fpp)
+          if (numBits <= 0L) {
+            throw InvalidPlanInput("Number of bits must be positive")
+          }
+          // Create BloomFilterAggregate with expectedNumItemsExpr and new numBits.
+          Some(
+            new BloomFilterAggregate(
+              children.head,
+              expectedNumItemsExpr,
+              Literal(numBits, LongType))
+              .toAggregateExpression())
         }
-        Some(
-          new BloomFilterAggregate(children.head, expectedNumItemsExpr, numBitsExpr)
-            .toAggregateExpression())
 
       case "window" if Seq(2, 3, 4).contains(fun.getArgumentsCount) =>
         val children = fun.getArgumentsList.asScala.map(transformExpression)