[apacheGH-2565] Fix NULL handling for various aggregation functions in SedonaSpark (apache#2563)

Author: Kontinuation

spark/common/src/main/scala/org/apache/spark/sql/sedona_sql/expressions/AggregateFunctions.scala

@@ -21,13 +21,14 @@ package org.apache.spark.sql.sedona_sql.expressions
 import org.apache.sedona.core.spatialRddTool.AdvancedStatCollector
 import org.apache.spark.sql.catalyst.encoders.{ExpressionEncoder, RowEncoder}
 import org.apache.sedona.common.Functions
+import org.apache.spark.sql.{Encoder, Encoders}
 import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder
 import org.apache.spark.sql.expressions.Aggregator
 import org.apache.spark.sql.sedona_sql.utils.SparkCompatUtil
 import org.apache.spark.sql.types.{DoubleType, LongType, StructField, StructType}
 import org.apache.spark.sql.{Encoder, Row}
 import org.apache.spark.sql.sedona_sql.EncodersShim
-import org.locationtech.jts.geom.{Coordinate, Geometry, GeometryFactory}
+import org.locationtech.jts.geom.{Coordinate, Envelope, Geometry, GeometryFactory}
 import org.locationtech.jts.operation.overlayng.OverlayNGRobust
 
 import scala.collection.JavaConverters._
@@ -38,18 +39,7 @@ import scala.collection.mutable.ListBuffer
  */
 
 trait TraitSTAggregateExec {
-  val initialGeometry: Geometry = {
-    // dummy value for initial value(polygon but )
-    // any other value is ok.
-    val coordinates: Array[Coordinate] = new Array[Coordinate](5)
-    coordinates(0) = new Coordinate(-999999999, -999999999)
-    coordinates(1) = new Coordinate(-999999999, -999999999)
-    coordinates(2) = new Coordinate(-999999999, -999999999)
-    coordinates(3) = new Coordinate(-999999999, -999999999)
-    coordinates(4) = coordinates(0)
-    val geometryFactory = new GeometryFactory()
-    geometryFactory.createPolygon(coordinates)
-  }
+  val initialGeometry: Geometry = null
   val serde = ExpressionEncoder[Geometry]()
 
   def zero: Geometry = initialGeometry
@@ -68,7 +58,9 @@ private[apache] class ST_Union_Aggr(bufferSize: Int = 1000)
   val bufferSerde = ExpressionEncoder[ListBuffer[Geometry]]()
 
   override def reduce(buffer: ListBuffer[Geometry], input: Geometry): ListBuffer[Geometry] = {
-    buffer += input
+    if (input != null) {
+      buffer += input
+    }
     if (buffer.size >= bufferSize) {
       // Perform the union when buffer size is reached
       val unionGeometry = OverlayNGRobust.union(buffer.asJava)
@@ -92,6 +84,9 @@ private[apache] class ST_Union_Aggr(bufferSize: Int = 1000)
   }
 
   override def finish(reduction: ListBuffer[Geometry]): Geometry = {
+    if (reduction.isEmpty) {
+      return null
+    }
     OverlayNGRobust.union(reduction.asJava)
   }
 
@@ -103,81 +98,76 @@ private[apache] class ST_Union_Aggr(bufferSize: Int = 1000)
 }
 
 /**
- * Return the envelope boundary of the entire column
+ * A helper class to store envelope boundary during aggregation. We use this custom case class
+ * instead of JTS Envelope to work with the Spark Encoder.
  */
-private[apache] class ST_Envelope_Aggr
-    extends Aggregator[Geometry, Geometry, Geometry]
-    with TraitSTAggregateExec {
+case class EnvelopeBuffer(minX: Double, maxX: Double, minY: Double, maxY: Double) {
+  def isNull: Boolean = minX > maxX
 
-  def reduce(buffer: Geometry, input: Geometry): Geometry = {
-    val accumulateEnvelope = buffer.getEnvelopeInternal
-    val newEnvelope = input.getEnvelopeInternal
-    val coordinates: Array[Coordinate] = new Array[Coordinate](5)
-    var minX = 0.0
-    var minY = 0.0
-    var maxX = 0.0
-    var maxY = 0.0
-    if (accumulateEnvelope.equals(initialGeometry.getEnvelopeInternal)) {
-      // Found the accumulateEnvelope is the initial value
-      minX = newEnvelope.getMinX
-      minY = newEnvelope.getMinY
-      maxX = newEnvelope.getMaxX
-      maxY = newEnvelope.getMaxY
-    } else if (newEnvelope.equals(initialGeometry.getEnvelopeInternal)) {
-      minX = accumulateEnvelope.getMinX
-      minY = accumulateEnvelope.getMinY
-      maxX = accumulateEnvelope.getMaxX
-      maxY = accumulateEnvelope.getMaxY
+  def toEnvelope: Envelope = {
+    if (isNull) {
+      new Envelope()
     } else {
-      minX = Math.min(accumulateEnvelope.getMinX, newEnvelope.getMinX)
-      minY = Math.min(accumulateEnvelope.getMinY, newEnvelope.getMinY)
-      maxX = Math.max(accumulateEnvelope.getMaxX, newEnvelope.getMaxX)
-      maxY = Math.max(accumulateEnvelope.getMaxY, newEnvelope.getMaxY)
+      new Envelope(minX, maxX, minY, maxY)
     }
-    coordinates(0) = new Coordinate(minX, minY)
-    coordinates(1) = new Coordinate(minX, maxY)
-    coordinates(2) = new Coordinate(maxX, maxY)
-    coordinates(3) = new Coordinate(maxX, minY)
-    coordinates(4) = coordinates(0)
-    val geometryFactory = new GeometryFactory()
-    geometryFactory.createPolygon(coordinates)
-
   }
 
-  def merge(buffer1: Geometry, buffer2: Geometry): Geometry = {
-    val leftEnvelope = buffer1.getEnvelopeInternal
-    val rightEnvelope = buffer2.getEnvelopeInternal
-    val coordinates: Array[Coordinate] = new Array[Coordinate](5)
-    var minX = 0.0
-    var minY = 0.0
-    var maxX = 0.0
-    var maxY = 0.0
-    if (leftEnvelope.equals(initialGeometry.getEnvelopeInternal)) {
-      minX = rightEnvelope.getMinX
-      minY = rightEnvelope.getMinY
-      maxX = rightEnvelope.getMaxX
-      maxY = rightEnvelope.getMaxY
-    } else if (rightEnvelope.equals(initialGeometry.getEnvelopeInternal)) {
-      minX = leftEnvelope.getMinX
-      minY = leftEnvelope.getMinY
-      maxX = leftEnvelope.getMaxX
-      maxY = leftEnvelope.getMaxY
+  def merge(other: EnvelopeBuffer): EnvelopeBuffer = {
+    if (this.isNull) {
+      other
+    } else if (other.isNull) {
+      this
     } else {
-      minX = Math.min(leftEnvelope.getMinX, rightEnvelope.getMinX)
-      minY = Math.min(leftEnvelope.getMinY, rightEnvelope.getMinY)
-      maxX = Math.max(leftEnvelope.getMaxX, rightEnvelope.getMaxX)
-      maxY = Math.max(leftEnvelope.getMaxY, rightEnvelope.getMaxY)
+      EnvelopeBuffer(
+        math.min(this.minX, other.minX),
+        math.max(this.maxX, other.maxX),
+        math.min(this.minY, other.minY),
+        math.max(this.maxY, other.maxY))
     }
+  }
+}
 
-    coordinates(0) = new Coordinate(minX, minY)
-    coordinates(1) = new Coordinate(minX, maxY)
-    coordinates(2) = new Coordinate(maxX, maxY)
-    coordinates(3) = new Coordinate(maxX, minY)
-    coordinates(4) = coordinates(0)
-    val geometryFactory = new GeometryFactory()
-    geometryFactory.createPolygon(coordinates)
+/**
+ * Return the envelope boundary of the entire column
+ */
+private[apache] class ST_Envelope_Aggr
+    extends Aggregator[Geometry, Option[EnvelopeBuffer], Geometry] {
+
+  val serde = ExpressionEncoder[Geometry]()
+
+  def reduce(buffer: Option[EnvelopeBuffer], input: Geometry): Option[EnvelopeBuffer] = {
+    if (input == null) return buffer
+    val env = input.getEnvelopeInternal
+    val envBuffer = EnvelopeBuffer(env.getMinX, env.getMaxX, env.getMinY, env.getMaxY)
+    buffer match {
+      case Some(b) => Some(b.merge(envBuffer))
+      case None => Some(envBuffer)
+    }
+  }
+
+  def merge(
+      buffer1: Option[EnvelopeBuffer],
+      buffer2: Option[EnvelopeBuffer]): Option[EnvelopeBuffer] = {
+    (buffer1, buffer2) match {
+      case (Some(b1), Some(b2)) => Some(b1.merge(b2))
+      case (Some(_), None) => buffer1
+      case (None, Some(_)) => buffer2
+      case (None, None) => None
+    }
+  }
+
+  def finish(reduction: Option[EnvelopeBuffer]): Geometry = {
+    reduction match {
+      case Some(b) => new GeometryFactory().toGeometry(b.toEnvelope)
+      case None => null
+    }
   }
 
+  def bufferEncoder: Encoder[Option[EnvelopeBuffer]] = Encoders.product[Option[EnvelopeBuffer]]
+
+  def outputEncoder: ExpressionEncoder[Geometry] = serde
+
+  def zero: Option[EnvelopeBuffer] = None
 }
 
 /**
@@ -187,16 +177,26 @@ private[apache] class ST_Intersection_Aggr
     extends Aggregator[Geometry, Geometry, Geometry]
     with TraitSTAggregateExec {
   def reduce(buffer: Geometry, input: Geometry): Geometry = {
-    if (buffer.isEmpty) input
-    else if (buffer.equalsExact(initialGeometry)) input
-    else buffer.intersection(input)
+    if (input == null) {
+      return buffer
+    }
+    if (buffer == null) {
+      return input
+    }
+    buffer.intersection(input)
   }
 
   def merge(buffer1: Geometry, buffer2: Geometry): Geometry = {
-    if (buffer1.equalsExact(initialGeometry)) buffer2
-    else if (buffer2.equalsExact(initialGeometry)) buffer1
-    else buffer1.intersection(buffer2)
+    if (buffer1 == null) {
+      return buffer2
+    }
+    if (buffer2 == null) {
+      return buffer1
+    }
+    buffer1.intersection(buffer2)
   }
+
+  override def finish(out: Geometry): Geometry = out
 }
 
 /**
@@ -225,7 +225,7 @@ private[apache] class ST_Collect_Agg
 
   override def finish(reduction: ListBuffer[Geometry]): Geometry = {
     if (reduction.isEmpty) {
-      new GeometryFactory().createGeometryCollection()
+      null
     } else {
       Functions.createMultiGeometry(reduction.toArray)
     }
@@ -303,25 +303,3 @@ class ST_Analyze_Aggr extends Aggregator[Geometry, AdvancedStatCollector, Row] {
         StructField("mean_envelope_height", DoubleType, nullable = false),
         StructField("mean_envelope_area", DoubleType, nullable = false))))
 }
-
-  override def merge(
-      buffer1: ListBuffer[Geometry],
-      buffer2: ListBuffer[Geometry]): ListBuffer[Geometry] = {
-    buffer1 ++= buffer2
-    buffer1
-  }
-
-  override def finish(reduction: ListBuffer[Geometry]): Geometry = {
-    if (reduction.isEmpty) {
-      new GeometryFactory().createGeometryCollection()
-    } else {
-      Functions.createMultiGeometry(reduction.toArray)
-    }
-  }
-
-  def bufferEncoder: ExpressionEncoder[ListBuffer[Geometry]] = bufferSerde
-
-  def outputEncoder: ExpressionEncoder[Geometry] = serde
-
-  override def zero: ListBuffer[Geometry] = ListBuffer.empty
-}