[SPARK-53957][Geo][SQL] Support GEOGRAPHY and GEOMETRY in the SpatialReferenceSystemMapper

sql/api/src/main/scala/org/apache/spark/sql/internal/types/SpatialReferenceSystemMapper.java

@@ -18,6 +18,7 @@
 package org.apache.spark.sql.internal.types;
 
 import java.util.HashMap;
+import java.util.HashSet;
 
 /*
  * Class for maintaining mappings between supported SRID values and the string ID of the
@@ -35,20 +36,38 @@ public static SpatialReferenceSystemMapper get() {
     return Instance;
   }
 
-  // Hash maps defining the mappings to/from SRID and string ID for a CRS.
+  // The `SpatialReferenceSystemMapper` class can be used for both GEOGRAPHY and GEOMETRY types.
+  // This enumeration defines the two spatial type options, in the context of SRS mapping.
+  public enum Type {
+    GEOGRAPHY,
+    GEOMETRY
+  }
+
+  // Hash maps for defining the mappings to/from SRID and string ID for a CRS.
   private static final HashMap<Integer, String> sridToStringId = buildSridToStringIdMap();
   private static final HashMap<String, Integer> stringIdToSrid = buildStringIdToSridMap();
 
+  // Hash set for keeping track of the supported SRID values for geographic CRS.
+  public static final HashSet<Integer> geographicSrids = buildGeographicSridSet();
+
   // Returns the string ID corresponding to the input SRID. If the input SRID is not supported,
   // `null` is returned.
-  public String getStringId(int srid) {
-    return sridToStringId.get(srid);
+  public String getStringId(int srid, Type type) {
+    String stringId = sridToStringId.get(srid);
+    return switch (type) {
+      case GEOGRAPHY -> geographicSrids.contains(srid) ? stringId : null;
+      case GEOMETRY -> stringId;
+    };
   }
 
   // Returns the SRID corresponding to the input string ID. If the input string ID is not
   // supported, `null` is returned.
-  public Integer getSrid(String stringId) {
-    return stringIdToSrid.get(stringId);
+  public Integer getSrid(String stringId, Type type) {
+    Integer srid = stringIdToSrid.get(stringId);
+      return switch (type) {
+        case GEOGRAPHY -> geographicSrids.contains(srid) ? srid : null;
+        case GEOMETRY -> srid;
+      };
   }
 
   // Currently, we only support a limited set of SRID / CRS mappings. However, we will soon extend
@@ -72,4 +91,11 @@ private static HashMap<String, Integer> buildStringIdToSridMap() {
     map.put("OGC:CRS84", 4326); // WGS84
     return map;
   }
+
+  // Helper method for building the set of supported geographic SRID values.
+  private static HashSet<Integer> buildGeographicSridSet() {
+    HashSet<Integer> set = new HashSet<>();
+    set.add(4326); // WGS84
+    return set;
+  }
 }

sql/api/src/main/scala/org/apache/spark/sql/types/GeographyType.scala

@@ -21,6 +21,7 @@ import org.json4s.JsonAST.{JString, JValue}
 
 import org.apache.spark.SparkIllegalArgumentException
 import org.apache.spark.annotation.Experimental
+import org.apache.spark.sql.internal.types.SpatialReferenceSystemMapper
 
 /**
  * The data type representing GEOGRAPHY values which are spatial objects, as defined in the Open
@@ -156,12 +157,17 @@ object GeographyType extends SpatialType {
   private final val GEOGRAPHY_MIXED_TYPE: GeographyType =
     GeographyType(MIXED_CRS, GEOGRAPHY_DEFAULT_ALGORITHM)
 
+  /**
+   * Type enum alias for geographic reference system mapping.
+   */
+  private final val GEOG = SpatialReferenceSystemMapper.Type.GEOGRAPHY
+
   /**
    * Constructors for GeographyType.
    */
   def apply(srid: Int): GeographyType = {
-    if (!isValidSrid(srid)) {
-      // Limited geographic SRID values are allowed.
+    val crs = SpatialReferenceSystemMapper.get().getStringId(srid, GEOG)
+    if (crs == null) {
       throw new SparkIllegalArgumentException(
         errorClass = "ST_INVALID_SRID_VALUE",
         messageParameters = Map("srid" -> srid.toString))
@@ -193,35 +199,9 @@ object GeographyType extends SpatialType {
   }
 
   def apply(crs: String, algorithm: EdgeInterpolationAlgorithm): GeographyType = {
-    if (!isValidCrs(crs)) {
-      // Limited geographic CRS values are allowed.
-      throw new SparkIllegalArgumentException(
-        errorClass = "ST_INVALID_CRS_VALUE",
-        messageParameters = Map("crs" -> crs))
-    }
     new GeographyType(crs, algorithm)
   }
 
-  /**
-   * Helper method to validate the CRS value. Limited geographic CRS values are allowed.
-   */
-  private def isValidCrs(crs: String): Boolean = {
-    // Currently, we only support "OGC:CRS84" / "EPSG:4326" / "SRID:ANY".
-    // In the future, we may support others.
-    crs.equalsIgnoreCase(GEOGRAPHY_DEFAULT_CRS) ||
-    crs.equalsIgnoreCase(GEOGRAPHY_DEFAULT_EPSG_CRS) ||
-    crs.equalsIgnoreCase(MIXED_CRS)
-  }
-
-  /**
-   * Helper method to validate the SRID value. Only geographic SRID values are allowed.
-   */
-
-  private def isValidSrid(srid: Int): Boolean = {
-    // Currently, we only support 4326. In the future, we may support others.
-    srid == GEOGRAPHY_DEFAULT_SRID
-  }
-
   override private[sql] def defaultConcreteType: DataType = GEOGRAPHY_MIXED_TYPE
 
   override private[sql] def acceptsType(other: DataType): Boolean =
@@ -235,17 +215,17 @@ object GeographyType extends SpatialType {
   private[types] def toSrid(crs: String): Int = {
     // The special value "SRID:ANY" is used to represent mixed SRID values.
     if (crs.equalsIgnoreCase(GeographyType.MIXED_CRS)) {
-      GeographyType.MIXED_SRID
+      return GeographyType.MIXED_SRID
     }
-    // As for other valid CRS values, we currently offer limited support.
-    else if (crs.equalsIgnoreCase(GeographyType.GEOGRAPHY_DEFAULT_CRS) ||
-      crs.equalsIgnoreCase(GeographyType.GEOGRAPHY_DEFAULT_EPSG_CRS)) {
-      GeographyType.GEOGRAPHY_DEFAULT_SRID
-    } else {
+    // For all other CRS values, we need to look up the corresponding SRID.
+    val srid = SpatialReferenceSystemMapper.get().getSrid(crs, GEOG)
+    if (srid == null) {
+      // If the CRS value is not recognized, we throw an exception.
       throw new SparkIllegalArgumentException(
         errorClass = "ST_INVALID_CRS_VALUE",
         messageParameters = Map("crs" -> crs))
     }
+    srid
   }
 }
 

sql/api/src/main/scala/org/apache/spark/sql/types/GeometryType.scala

@@ -149,11 +149,16 @@ object GeometryType extends SpatialType {
   private final val GEOMETRY_MIXED_TYPE: GeometryType =
     GeometryType(MIXED_CRS)
 
+  /**
+   * Type enum alias for geometric reference system mapping.
+   */
+  private final val GEOM = SpatialReferenceSystemMapper.Type.GEOMETRY
+
   /**
    * Constructors for GeometryType.
    */
   def apply(srid: Int): GeometryType = {
-    val crs = SpatialReferenceSystemMapper.get().getStringId(srid)
+    val crs = SpatialReferenceSystemMapper.get().getStringId(srid, GEOM)
     if (crs == null) {
       throw new SparkIllegalArgumentException(
         errorClass = "ST_INVALID_SRID_VALUE",
@@ -191,7 +196,7 @@ object GeometryType extends SpatialType {
       return GeometryType.MIXED_SRID
     }
     // For all other CRS values, we need to look up the corresponding SRID.
-    val srid = SpatialReferenceSystemMapper.get().getSrid(crs)
+    val srid = SpatialReferenceSystemMapper.get().getSrid(crs, GEOM)
     if (srid == null) {
       // If the CRS value is not recognized, we throw an exception.
       throw new SparkIllegalArgumentException(

sql/core/src/test/scala/org/apache/spark/sql/internal/types/SpatialReferenceSystemMapperSuite.java

@@ -22,33 +22,62 @@
 
 public class SpatialReferenceSystemMapperSuite {
 
+  // Singleton instance of the SpatialReferenceSystemMapper to use across tests.
+  private final SpatialReferenceSystemMapper srsMapper = SpatialReferenceSystemMapper.get();
+
+  // Spatial types for GEOMETRY and GEOGRAPHY reference system mapping.
+  private final SpatialReferenceSystemMapper.Type GEOM =
+    SpatialReferenceSystemMapper.Type.GEOMETRY;
+  private final SpatialReferenceSystemMapper.Type GEOG =
+    SpatialReferenceSystemMapper.Type.GEOGRAPHY;
+
   @Test
   public void getStringIdReturnsCorrectStringIdForValidSrid() {
-    SpatialReferenceSystemMapper srMapper = SpatialReferenceSystemMapper.get();
-    Assertions.assertEquals("SRID:0", srMapper.getStringId(0));
-    Assertions.assertEquals("EPSG:3857", srMapper.getStringId(3857));
-    Assertions.assertEquals("OGC:CRS84", srMapper.getStringId(4326));
+    // GEOMETRY.
+    Assertions.assertEquals("SRID:0", srsMapper.getStringId(0, GEOM));
+    Assertions.assertEquals("EPSG:3857", srsMapper.getStringId(3857, GEOM));
+    Assertions.assertEquals("OGC:CRS84", srsMapper.getStringId(4326, GEOM));
+    // GEOGRAPHY.
+    Assertions.assertEquals("OGC:CRS84", srsMapper.getStringId(4326, GEOG));
   }
 
   @Test
   public void getStringIdReturnsNullForInvalidSrid() {
-    SpatialReferenceSystemMapper srMapper = SpatialReferenceSystemMapper.get();
-    Assertions.assertNull(srMapper.getStringId(-1));
-    Assertions.assertNull(srMapper.getStringId(9999));
+    // GEOMETRY.
+    Assertions.assertNull(srsMapper.getStringId(-2, GEOM));
+    Assertions.assertNull(srsMapper.getStringId(-1, GEOM));
+    Assertions.assertNull(srsMapper.getStringId(1, GEOM));
+    Assertions.assertNull(srsMapper.getStringId(2, GEOM));
+    Assertions.assertNull(srsMapper.getStringId(9999, GEOM));
+    // GEOGRAPHY.
+    Assertions.assertNull(srsMapper.getStringId(-2, GEOG));
+    Assertions.assertNull(srsMapper.getStringId(-1, GEOG));
+    Assertions.assertNull(srsMapper.getStringId(0, GEOG));
+    Assertions.assertNull(srsMapper.getStringId(1, GEOG));
+    Assertions.assertNull(srsMapper.getStringId(2, GEOG));
+    Assertions.assertNull(srsMapper.getStringId(3857, GEOG));
+    Assertions.assertNull(srsMapper.getStringId(9999, GEOG));
   }
 
   @Test
   public void getSridReturnsCorrectSridForValidStringId() {
-    SpatialReferenceSystemMapper srMapper = SpatialReferenceSystemMapper.get();
-    Assertions.assertEquals(0, srMapper.getSrid("SRID:0"));
-    Assertions.assertEquals(3857, srMapper.getSrid("EPSG:3857"));
-    Assertions.assertEquals(4326, srMapper.getSrid("OGC:CRS84"));
+    // GEOMETRY.
+    Assertions.assertEquals(0, srsMapper.getSrid("SRID:0", GEOM));
+    Assertions.assertEquals(3857, srsMapper.getSrid("EPSG:3857", GEOM));
+    Assertions.assertEquals(4326, srsMapper.getSrid("OGC:CRS84", GEOM));
+    // GEOGRAPHY.
+    Assertions.assertEquals(4326, srsMapper.getSrid("OGC:CRS84", GEOG));
   }
 
   @Test
   public void getSridReturnsNullForInvalidStringId() {
-    SpatialReferenceSystemMapper srMapper = SpatialReferenceSystemMapper.get();
-    Assertions.assertNull(srMapper.getSrid("INVALID:ID"));
-    Assertions.assertNull(srMapper.getSrid("EPSG:9999"));
+    // GEOMETRY.
+    Assertions.assertNull(srsMapper.getSrid("INVALID:ID", GEOM));
+    Assertions.assertNull(srsMapper.getSrid("EPSG:9999", GEOM));
+    // GEOGRAPHY.
+    Assertions.assertNull(srsMapper.getSrid("INVALID:ID", GEOG));
+    Assertions.assertNull(srsMapper.getSrid("EPSG:9999", GEOG));
+    Assertions.assertNull(srsMapper.getSrid("SRID:0", GEOG));
+    Assertions.assertNull(srsMapper.getSrid("EPSG:3857", GEOG));
   }
 }

sql/core/src/test/scala/org/apache/spark/sql/types/GeographyTypeSuite.scala

@@ -141,13 +141,13 @@ class GeographyTypeSuite extends SparkFunSuite {
     val validGeographies = Seq(
       "\"geography\"",
       "\"geography(OGC:CRS84)\"",
-      "\"geography(ogc:CRS84)\"",
-      "\"geography( ogc:CRS84  )\"",
-      "\"geography(EPSG:4326)\"",
+      "\"geography(   OGC:CRS84 )\"",
       "\"geography(spherical)\"",
+      "\"geography(SPHERICAL)\"",
       "\"geography(  spherical)\"",
       "\"geography(OGC:CRS84,    spherical    )\"",
-      "\"geography( OGC:CRS84   , spherical )\""
+      "\"geography( OGC:CRS84   , spherical )\"",
+      "\"geography( OGC:CRS84   , SPHERICAL )\""
     )
     validGeographies.foreach { geog =>
       DataType.fromJson(geog).isInstanceOf[GeographyType]