[SEDONA-414] Make ST_MakeLine in sedona-spark work with array inputs. (…

…#1068)
apache · Nov 1, 2023 · 2f4f9dd · 2f4f9dd
1 parent 5bd0172
commit 2f4f9dd
Show file tree

Hide file tree

Showing 4 changed files with 17 additions and 5 deletions.
diff --git a/flink/src/test/java/org/apache/sedona/flink/FunctionTest.java b/flink/src/test/java/org/apache/sedona/flink/FunctionTest.java
@@ -794,6 +794,10 @@ public void testMakeLine() {
         table = table.select(call(Functions.ST_MakeLine.class.getSimpleName(), $("point1"), $("point2")));
         Geometry result = (Geometry) first(table).getField(0);
         assertEquals("LINESTRING (0 0, 1 1)", result.toString());
+
+        table = tableEnv.sqlQuery("SELECT ST_MakeLine(ARRAY[ST_Point(2, 2), ST_Point(3, 3)]) AS line");
+        result = (Geometry) first(table).getField(0);
+        assertEquals("LINESTRING (2 2, 3 3)", result.toString());
     }
 
     @Test

diff --git a/spark/common/src/main/scala/org/apache/spark/sql/sedona_sql/expressions/Functions.scala b/spark/common/src/main/scala/org/apache/spark/sql/sedona_sql/expressions/Functions.scala
@@ -687,7 +687,7 @@ case class ST_SubDivideExplode(children: Seq[Expression])
 }
 
 case class ST_MakeLine(inputExpressions: Seq[Expression])
-  extends InferredExpression(InferrableFunction.allowRightNull(Functions.makeLine _)) {
+  extends InferredExpression(inferrableFunction2(Functions.makeLine), inferrableFunction1(Functions.makeLine)) {
 
   protected def withNewChildrenInternal(newChildren: IndexedSeq[Expression]) = {
     copy(inputExpressions = newChildren)

diff --git a/spark/common/src/test/scala/org/apache/sedona/sql/dataFrameAPITestScala.scala b/spark/common/src/test/scala/org/apache/sedona/sql/dataFrameAPITestScala.scala
@@ -314,6 +314,10 @@ class dataFrameAPITestScala extends TestBaseScala {
       val actualResult = df.take(1)(0).get(0).asInstanceOf[Geometry].toText()
       val expectedResult = "LINESTRING (0 0, 1 1)"
       assert(actualResult == expectedResult)
+
+      val df2 = sparkSession.sql("SELECT ST_MakeLine(ARRAY(ST_Point(0, 0), ST_Point(1, 1), ST_Point(2, 2)))")
+      val actualResult2 = df2.take(1)(0).get(0).asInstanceOf[Geometry].toText()
+      assert(actualResult2 == "LINESTRING (0 0, 1 1, 2 2)")
     }
 
     it("Passed ST_MakeValid On Invalid Polygon") {

diff --git a/spark/common/src/test/scala/org/apache/sedona/sql/functionTestScala.scala b/spark/common/src/test/scala/org/apache/sedona/sql/functionTestScala.scala
@@ -322,11 +322,15 @@ class functionTestScala extends TestBaseScala with Matchers with GeometrySample
     }
 
     it("Passed ST_MakeLine") {
-
-      var testtable = sparkSession.sql(
-        "SELECT ST_MakeLine(ST_GeomFromText('POINT(1 2)'), ST_GeomFromText('POINT(3 4)'))"
+      val testtable = sparkSession.sql(
+        """SELECT
+          |ST_MakeLine(ST_GeomFromText('POINT(1 2)'), ST_GeomFromText('POINT(3 4)')),
+          |ST_MakeLine(ARRAY(ST_Point(5, 6), ST_Point(7, 8), ST_Point(9, 10)))
+          |""".stripMargin
       )
-      assert(testtable.take(1)(0).get(0).asInstanceOf[Geometry].toText.equals("LINESTRING (1 2, 3 4)"))
+      val row = testtable.take(1)(0)
+      assert(row.get(0).asInstanceOf[Geometry].toText.equals("LINESTRING (1 2, 3 4)"))
+      assert(row.get(1).asInstanceOf[Geometry].toText.equals("LINESTRING (5 6, 7 8, 9 10)"))
     }
 
     it("Passed ST_Polygon") {