Add ST_GeomFromWKB and ST_AsBinary functions

Author: mpertierra

presto-docs/src/main/sphinx/functions/geospatial.rst

@@ -22,6 +22,10 @@ Presto Geospatial functions support the Well-Known Text (WKT) form of spatial ob
 Constructors
 ------------
 
+.. function:: ST_AsBinary(Geometry) -> varbinary
+
+    Returns the WKB representation of the geometry.
+
 .. function:: ST_AsText(Geometry) -> varchar
 
     Returns the WKT representation of the geometry. For empty geometries,
@@ -32,6 +36,10 @@ Constructors
 
     Returns a geometry type object from WKT representation.
 
+.. function:: ST_GeomFromBinary(varbinary) -> Geometry
+
+    Returns a geometry type object from WKB representation.
+
 .. function:: ST_LineFromText(varchar) -> LineString
 
     Returns a geometry type linestring object from WKT representation.

presto-geospatial/src/main/java/com/facebook/presto/plugin/geospatial/GeoFunctions.java

@@ -80,13 +80,15 @@
 import static com.facebook.presto.spi.type.StandardTypes.INTEGER;
 import static com.google.common.base.Preconditions.checkArgument;
 import static io.airlift.slice.Slices.utf8Slice;
+import static io.airlift.slice.Slices.wrappedBuffer;
 import static java.lang.Math.atan2;
 import static java.lang.Math.cos;
 import static java.lang.Math.sin;
 import static java.lang.Math.sqrt;
 import static java.lang.Math.toIntExact;
 import static java.lang.Math.toRadians;
 import static java.lang.String.format;
+import static java.util.Objects.requireNonNull;
 import static org.locationtech.jts.simplify.TopologyPreservingSimplifier.simplify;
 
 public final class GeoFunctions
@@ -169,6 +171,14 @@ public static Slice stGeometryFromText(@SqlType(StandardTypes.VARCHAR) Slice inp
         return serialize(geometryFromText(input));
     }
 
+    @Description("Returns a Geometry type object from Well-Known Binary representation (WKB)")
+    @ScalarFunction("ST_GeomFromBinary")
+    @SqlType(GEOMETRY_TYPE_NAME)
+    public static Slice stGeomFromBinary(@SqlType(StandardTypes.VARBINARY) Slice input)
+    {
+        return serialize(geomFromBinary(input));
+    }
+
     @Description("Returns the Well-Known Text (WKT) representation of the geometry")
     @ScalarFunction("ST_AsText")
     @SqlType(StandardTypes.VARCHAR)
@@ -177,6 +187,14 @@ public static Slice stAsText(@SqlType(GEOMETRY_TYPE_NAME) Slice input)
         return utf8Slice(deserialize(input).asText());
     }
 
+    @Description("Returns the Well-Known Binary (WKB) representation of the geometry")
+    @ScalarFunction("ST_AsBinary")
+    @SqlType(StandardTypes.VARBINARY)
+    public static Slice stAsBinary(@SqlType(GEOMETRY_TYPE_NAME) Slice input)
+    {
+        return wrappedBuffer(deserialize(input).asBinary());
+    }
+
     @SqlNullable
     @Description("Returns the geometry that represents all points whose distance from the specified geometry is less than or equal to the specified distance")
     @ScalarFunction("ST_Buffer")
@@ -1044,6 +1062,20 @@ private static OGCGeometry geometryFromText(Slice input)
         return geometry;
     }
 
+    private static OGCGeometry geomFromBinary(Slice input)
+    {
+        requireNonNull(input, "input is null");
+        OGCGeometry geometry;
+        try {
+            geometry = OGCGeometry.fromBinary(input.toByteBuffer());
+        }
+        catch (IllegalArgumentException | IndexOutOfBoundsException e) {
+            throw new PrestoException(INVALID_FUNCTION_ARGUMENT, "Invalid WKB", e);
+        }
+        geometry.setSpatialReference(null);
+        return geometry;
+    }
+
     private static void validateType(String function, OGCGeometry geometry, Set<GeometryType> validTypes)
     {
         GeometryType type = GeometryType.getForEsriGeometryType(geometry.geometryType());

presto-geospatial/src/test/java/com/facebook/presto/plugin/geospatial/TestGeoFunctions.java

@@ -1010,4 +1010,41 @@ public void testSTGeometryType()
         assertFunction("ST_GeometryType(ST_GeometryFromText('GEOMETRYCOLLECTION(POINT(4 6),LINESTRING(4 6, 7 10))'))", VARCHAR, "ST_GeomCollection");
         assertFunction("ST_GeometryType(ST_Envelope(ST_GeometryFromText('LINESTRING (1 1, 2 2)')))", VARCHAR, "ST_Polygon");
     }
+
+    @Test
+    public void testSTGeometryFromBinary()
+    {
+        assertFunction("ST_GeomFromBinary(null)", GEOMETRY, null);
+
+        // empty geometries
+        assertGeomFromBinary("POINT EMPTY");
+        assertGeomFromBinary("MULTIPOINT EMPTY");
+        assertGeomFromBinary("LINESTRING EMPTY");
+        assertGeomFromBinary("MULTILINESTRING EMPTY");
+        assertGeomFromBinary("POLYGON EMPTY");
+        assertGeomFromBinary("MULTIPOLYGON EMPTY");
+        assertGeomFromBinary("GEOMETRYCOLLECTION EMPTY");
+
+        // valid nonempty geometries
+        assertGeomFromBinary("POINT (1 2)");
+        assertGeomFromBinary("MULTIPOINT ((1 2), (3 4))");
+        assertGeomFromBinary("LINESTRING (0 0, 1 2, 3 4)");
+        assertGeomFromBinary("MULTILINESTRING ((1 1, 5 1), (2 4, 4 4))");
+        assertGeomFromBinary("POLYGON ((0 0, 1 0, 1 1, 0 1, 0 0))");
+        assertGeomFromBinary("POLYGON ((0 0, 3 0, 3 3, 0 3, 0 0), (1 1, 1 2, 2 2, 2 1, 1 1))");
+        assertGeomFromBinary("MULTIPOLYGON (((1 1, 3 1, 3 3, 1 3, 1 1)), ((2 4, 6 4, 6 6, 2 6, 2 4)))");
+        assertGeomFromBinary("GEOMETRYCOLLECTION (POINT (1 2), LINESTRING (0 0, 1 2, 3 4), POLYGON ((0 0, 1 0, 1 1, 0 1, 0 0)))");
+
+        // invalid geometries
+        assertGeomFromBinary("MULTIPOINT ((0 0), (0 1), (1 1), (0 1))");
+        assertGeomFromBinary("LINESTRING (0 0, 0 1, 0 1, 1 1, 1 0, 0 0)");
+
+        // invalid binary
+        assertInvalidFunction("ST_GeomFromBinary(from_hex('deadbeef'))", "Invalid WKB");
+    }
+
+    private void assertGeomFromBinary(String wkt)
+    {
+        assertFunction(format("ST_AsText(ST_GeomFromBinary(ST_AsBinary(ST_GeometryFromText('%s'))))", wkt), VARCHAR, wkt);
+    }
 }