SNOW-1183003 Support Geometry Type

src/main/java/com/snowflake/snowpark_java/Row.java

@@ -2,6 +2,7 @@
 
 import com.snowflake.snowpark.internal.JavaUtils;
 import com.snowflake.snowpark_java.types.Geography;
+import com.snowflake.snowpark_java.types.Geometry;
 import com.snowflake.snowpark_java.types.InternalUtils;
 import com.snowflake.snowpark_java.types.Variant;
 import java.io.Serializable;
@@ -53,6 +54,8 @@ private static Object[] javaObjectToScalaObject(Object[] input) {
       } else if (result[i] instanceof Geography) {
         result[i] =
             com.snowflake.snowpark.types.Geography.fromGeoJSON(((Geography) result[i]).asGeoJSON());
+      } else if (result[i] instanceof Geometry) {
+        result[i] = com.snowflake.snowpark.types.Geometry.fromGeoJSON(result[i].toString());
       }
     }
     return result;
@@ -134,6 +137,8 @@ public Object get(int index) {
       return InternalUtils.createVariant((com.snowflake.snowpark.types.Variant) result);
     } else if (result instanceof com.snowflake.snowpark.types.Geography) {
       return Geography.fromGeoJSON(((com.snowflake.snowpark.types.Geography) result).asGeoJSON());
+    } else if (result instanceof com.snowflake.snowpark.types.Geometry) {
+      return Geometry.fromGeoJSON(result.toString());
     } else if (result instanceof com.snowflake.snowpark.types.Variant[]) {
       com.snowflake.snowpark.types.Variant[] scalaVariantArray =
           (com.snowflake.snowpark.types.Variant[]) result;
@@ -325,6 +330,17 @@ public Geography getGeography(int index) {
     return Geography.fromGeoJSON(scalaRow.getGeography(index).asGeoJSON());
   }
 
+  /**
+   * Retrieves the value of the column at the given index as a Geometry value.
+   *
+   * @param index The index of target column
+   * @return The Geometry value of the column at the given index
+   * @since 1.12.0
+   */
+  public Geometry getGeometry(int index) {
+    return Geometry.fromGeoJSON(scalaRow.getGeometry(index).toString());
+  }
+
   /**
    * Retrieves the value of the column at the given index as a list of Variant.
    *

src/main/java/com/snowflake/snowpark_java/SessionBuilder.java

@@ -82,8 +82,8 @@ public Session getOrCreate() {
   }
 
   /**
-   * Adds the app name to set in the query_tag after session creation.
-   * The query tag will be set with this format 'APPNAME=${appName}'.
+   * Adds the app name to set in the query_tag after session creation. The query tag will be set
+   * with this format 'APPNAME=${appName}'.
    *
    * @param appName Name of the app.
    * @return A {@code SessionBuilder} object

src/main/java/com/snowflake/snowpark_java/types/DataTypes.java

@@ -80,6 +80,13 @@ private DataTypes() {}
    */
   public static final GeographyType GeographyType = new GeographyType();
 
+  /**
+   * Retrieves the GeometryType object.
+   *
+   * @since 1.12.0
+   */
+  public static final GeometryType GeometryType = new GeometryType();
+
   /**
    * Retrieves the StringType object.
    *

src/main/java/com/snowflake/snowpark_java/types/Geometry.java

@@ -0,0 +1,70 @@
+package com.snowflake.snowpark_java.types;
+
+import java.io.IOException;
+import java.io.Serializable;
+import java.io.UncheckedIOException;
+
+/**
+ * Java representation of Snowflake Geometry data.
+ *
+ * @since 1.12.0
+ */
+public class Geometry implements Serializable {
+  private final String data;
+
+  private Geometry(String data) {
+    if (data == null) {
+      throw new UncheckedIOException(
+          new IOException("Cannot create geometry object from null input"));
+    }
+    this.data = data;
+  }
+
+  /**
+   * Checks whether two Geometry object are equal.
+   *
+   * @param other A Geometry object
+   * @return true if these two object are equal
+   * @since 1.12.0
+   */
+  @Override
+  public boolean equals(Object other) {
+    if (other instanceof Geometry) {
+      return data.equals(((Geometry) other).data);
+    }
+    return false;
+  }
+
+  /**
+   * Calculates the hash code of this Geometry Object.
+   *
+   * @return An int number representing the hash code value
+   * @since 1.12.0
+   */
+  @Override
+  public int hashCode() {
+    return this.data.hashCode();
+  }
+
+  /**
+   * Converts this Geometry object to a String value.
+   *
+   * @return A String value.
+   * @since 1.12.0
+   */
+  @Override
+  public String toString() {
+    return this.data;
+  }
+
+  /**
+   * Creates a Geometry object from a GeoJSON string.
+   *
+   * @param g GeoJSON String
+   * @return a new Geometry object
+   * @since 1.12.0
+   */
+  public static Geometry fromGeoJSON(String g) {
+    return new Geometry(g);
+  }
+}

src/main/java/com/snowflake/snowpark_java/types/GeometryType.java

@@ -0,0 +1,10 @@
+package com.snowflake.snowpark_java.types;
+
+/**
+ * Geography data type. This maps to GEOMETRY data type in Snowflake.
+ *
+ * @since 1.12.0
+ */
+public class GeometryType extends AtomicType {
+  GeometryType() {}
+}

src/main/scala/com/snowflake/snowpark/Row.scala

@@ -1,9 +1,8 @@
 package com.snowflake.snowpark
 
 import java.sql.{Date, Time, Timestamp}
-
 import com.snowflake.snowpark.internal.ErrorMessage
-import com.snowflake.snowpark.types.{Geography, Variant}
+import com.snowflake.snowpark.types.{Geography, Geometry, Variant}
 
 import scala.util.hashing.MurmurHash3
 
@@ -245,6 +244,7 @@ class Row private (values: Array[Any]) extends Serializable {
     get(index) match {
       case variant: Variant => variant.toString
       case geo: Geography => geo.toString
+      case geo: Geometry => geo.toString
       case array: Array[_] => new Variant(array).toString
       case seq: Seq[_] => new Variant(seq).toString
       case map: Map[_, _] => new Variant(map).toString
@@ -301,6 +301,14 @@ class Row private (values: Array[Any]) extends Serializable {
    */
   def getGeography(index: Int): Geography = getAs[Geography](index)
 
+  /**
+   * Returns the value of the column at the given index as Geometry class
+   *
+   * @since 1.12.0
+   * @group getter
+   */
+  def getGeometry(index: Int): Geometry = getAs[Geometry](index)
+
   /**
    * Returns the value of the column at the given index as a Seq of Variant
    * @since 0.2.0

src/main/scala/com/snowflake/snowpark/Session.scala

@@ -733,8 +733,8 @@ class Session private (private[snowpark] val conn: ServerConnection) extends Log
     val spAttrs = schema.map { field =>
       {
         val sfType = field.dataType match {
-          case _ @(VariantType | _: ArrayType | _: MapType | GeographyType | TimeType | DateType |
-              TimestampType) =>
+          case _ @(VariantType | _: ArrayType | _: MapType | GeographyType | GeometryType |
+              TimeType | DateType | TimestampType) =>
             StringType
           case other => other
         }
@@ -764,6 +764,7 @@ class Session private (private[snowpark] val conn: ServerConnection) extends Log
         case (value, _: AtomicType) => value
         case (value: Variant, VariantType) => value.asJsonString()
         case (value: Geography, GeographyType) => value.asGeoJSON()
+        case (value: Geometry, GeometryType) => value.toString
         case (value: Array[_], _: ArrayType) =>
           new Variant(value.toSeq).asJsonString()
         case (value: Map[_, _], _: MapType) => new Variant(value).asJsonString()
@@ -784,6 +785,7 @@ class Session private (private[snowpark] val conn: ServerConnection) extends Log
         case TimestampType => callUDF("to_timestamp", column(field.name)).as(field.name)
         case VariantType => to_variant(parse_json(column(field.name))).as(field.name)
         case GeographyType => callUDF("to_geography", column(field.name)).as(field.name)
+        case GeometryType => callUDF("to_geometry", column(field.name)).as(field.name)
         case _: ArrayType => to_array(parse_json(column(field.name))).as(field.name)
         case _: MapType => to_object(parse_json(column(field.name))).as(field.name)
         case _ => column(field.name)
@@ -1408,7 +1410,6 @@ object Session extends Logging {
       this
     }
 
-
     /**
      * Adds the app name to set in the query_tag after session creation.
      * The query tag will be set with this format 'APPNAME=${appName}'.

src/main/scala/com/snowflake/snowpark/internal/ErrorMessage.scala

@@ -158,7 +158,8 @@ private[snowpark] object ErrorMessage {
     "0424" ->
       """Invalid input argument type, the input argument type of Explode function should be either Map or Array types.
         |The input argument type: %s
-        |""".stripMargin)
+        |""".stripMargin,
+    "0425" -> "Unsupported Geometry output format: %s. Please set session parameter GEOMETRY_OUTPUT_FORMAT to GeoJSON.")
   // scalastyle:on
 
   /*
@@ -405,6 +406,9 @@ private[snowpark] object ErrorMessage {
   def MISC_INVALID_EXPLODE_ARGUMENT_TYPE(argumentType: String): SnowparkClientException =
     createException("0424", argumentType)
 
+  def MISC_UNSUPPORTED_GEOMETRY_FORMAT(typeName: String): SnowparkClientException =
+    createException("0425", typeName)
+
   /**
    * Create Snowpark client Exception.
    *

src/main/scala/com/snowflake/snowpark/internal/JavaDataTypeUtils.scala

@@ -13,6 +13,7 @@ import com.snowflake.snowpark_java.types.{
   DoubleType => JDoubleType,
   FloatType => JFloatType,
   GeographyType => JGeographyType,
+  GeometryType => JGeometryType,
   IntegerType => JIntegerType,
   LongType => JLongType,
   MapType => JMapType,
@@ -36,6 +37,7 @@ object JavaDataTypeUtils {
       case DoubleType => JDataTypes.DoubleType
       case FloatType => JDataTypes.FloatType
       case GeographyType => JDataTypes.GeographyType
+      case GeometryType => JDataTypes.GeometryType
       case IntegerType => JDataTypes.IntegerType
       case LongType => JDataTypes.LongType
       case MapType(keyType, valueType) =>
@@ -58,6 +60,7 @@ object JavaDataTypeUtils {
       case _: JDoubleType => DoubleType
       case _: JFloatType => FloatType
       case _: JGeographyType => GeographyType
+      case _: JGeometryType => GeometryType
       case _: JIntegerType => IntegerType
       case _: JLongType => LongType
       case mp: JMapType =>

src/main/scala/com/snowflake/snowpark/internal/JavaUtils.scala

@@ -25,7 +25,7 @@ import com.snowflake.snowpark.{
 }
 
 import java.io._
-import com.snowflake.snowpark.types.{Geography, Variant}
+import com.snowflake.snowpark.types.{Geography, Geometry, Variant}
 import com.snowflake.snowpark_java.types.InternalUtils
 import com.snowflake.snowpark_java.udtf._
 
@@ -180,14 +180,24 @@ object JavaUtils {
 
   def geographyToString(g: Geography): String = if (g == null) null else g.asGeoJSON()
 
+  def geometryToString(g: Geometry): String = if (g == null) null else g.toString()
+
   def geographyToString(g: com.snowflake.snowpark_java.types.Geography): String =
     if (g == null) null else g.asGeoJSON()
 
+  def geometryToString(g: com.snowflake.snowpark_java.types.Geometry): String =
+    if (g == null) null else g.toString
+
   def stringToGeography(g: String): Geography = if (g == null) null else Geography.fromGeoJSON(g)
 
+  def stringToGeometry(g: String): Geometry = if (g == null) null else Geometry.fromGeoJSON(g)
+
   def stringToJavaGeography(g: String): com.snowflake.snowpark_java.types.Geography =
     if (g == null) null else com.snowflake.snowpark_java.types.Geography.fromGeoJSON(g)
 
+  def stringToJavaGeometry(g: String): com.snowflake.snowpark_java.types.Geometry =
+    if (g == null) null else com.snowflake.snowpark_java.types.Geometry.fromGeoJSON(g)
+
   def variantToString(v: Variant): String = if (v == null) null else v.asJsonString()
 
   def variantToString(v: com.snowflake.snowpark_java.types.Variant): String =

src/main/scala/com/snowflake/snowpark/internal/ParameterUtils.scala

@@ -19,6 +19,7 @@ private[snowpark] object ParameterUtils extends Logging {
   // client parameters
   private[snowpark] val SnowparkLazyAnalysis: String = "snowpark_lazy_analysis"
   private[snowpark] val GeographyOutputFormat: String = "geography_output_format"
+  private[snowpark] val GeometryOutputFormat: String = "geometry_output_format"
   private[snowpark] val SnowparkUseScopedTempObjects: String = "snowpark_use_scoped_temp_objects"
   private[snowpark] val SnowparkEnableClosureCleaner: String = "snowpark_enable_closure_cleaner"
   private[snowpark] val SnowparkRequestTimeoutInSeconds: String =

src/main/scala/com/snowflake/snowpark/internal/ScalaFunctions.scala

@@ -55,6 +55,7 @@ object ScalaFunctions {
     case t if t =:= typeOf[scala.collection.mutable.Map[String, String]] => true
     case t if t =:= typeOf[scala.collection.mutable.Map[String, Variant]] => true
     case t if t =:= typeOf[Geography] => true
+    case t if t =:= typeOf[Geometry] => true
     case t if t =:= typeOf[Variant] => true
     case t if t <:< typeOf[scala.collection.Iterable[_]] =>
       throw new UnsupportedOperationException(
@@ -97,6 +98,7 @@ object ScalaFunctions {
     case t if t =:= typeOf[scala.collection.mutable.Map[String, Variant]] =>
       UdfColumnSchema(MapType(StringType, VariantType))
     case t if t =:= typeOf[Geography] => UdfColumnSchema(GeographyType)
+    case t if t =:= typeOf[Geometry] => UdfColumnSchema(GeometryType)
     case t if t =:= typeOf[Variant] => UdfColumnSchema(VariantType)
     case t => throw new UnsupportedOperationException(s"Unsupported type $t")
   }

src/main/scala/com/snowflake/snowpark/internal/ServerConnection.scala

@@ -80,6 +80,7 @@ private[snowpark] object ServerConnection {
       case "VARIANT" => VariantType
       case "OBJECT" => MapType(StringType, StringType)
       case "GEOGRAPHY" => GeographyType
+      case "GEOMETRY" => GeometryType
       case _ => getTypeFromJDBCType(sqlType, precision, scale, signed)
     }
   }
@@ -258,6 +259,7 @@ private[snowpark] class ServerConnection(
       val schema = ServerConnection.convertResultMetaToAttribute(data.getMetaData)
 
       lazy val geographyOutputFormat = getParameterValue(ParameterUtils.GeographyOutputFormat)
+      lazy val geometryOutputFormat = getParameterValue(ParameterUtils.GeometryOutputFormat)
 
       val iterator = new CloseableIterator[Row] {
         private var _currentRow: Row = _
@@ -299,6 +301,13 @@ private[snowpark] class ServerConnection(
                           throw ErrorMessage.MISC_UNSUPPORTED_GEOGRAPHY_FORMAT(
                             geographyOutputFormat)
                       }
+                    case GeometryType =>
+                      geometryOutputFormat match {
+                        case "GeoJSON" => Geometry.fromGeoJSON(data.getString(resultIndex))
+                        case _ =>
+                          throw ErrorMessage.MISC_UNSUPPORTED_GEOMETRY_FORMAT(
+                            geometryOutputFormat)
+                      }
                     case _ =>
                       // ArrayType, StructType, MapType
                       throw new UnsupportedOperationException(

src/main/scala/com/snowflake/snowpark/internal/TypeToSchemaConverter.scala

@@ -84,6 +84,7 @@ object TypeToSchemaConverter {
 
       case t if t =:= typeOf[Variant] => (VariantType, true)
       case t if t =:= typeOf[Geography] => (GeographyType, true)
+      case t if t =:= typeOf[Geometry] => (GeometryType, true)
       case t if t =:= typeOf[Date] => (DateType, true)
       case t if t =:= typeOf[Timestamp] => (TimestampType, true)
       case t if t =:= typeOf[Time] => (TimeType, true)