Geospatial datatype support

NEXT_CHANGELOG.md

@@ -3,7 +3,8 @@
 ## [Unreleased]
 
 ### Added
-
+- Enabled direct results by default in SEA mode to improve latency for short and small queries.
+- Added support for geospatial data types.
 ### Updated
 
 ### Fixed

pom.xml

@@ -287,6 +287,11 @@
       <artifactId>resilience4j-core</artifactId>
       <version>${resilience4j.version}</version>
     </dependency>
+    <dependency>
+      <groupId>org.locationtech.jts</groupId>
+      <artifactId>jts-core</artifactId>
+      <version>1.20.0</version>
+    </dependency>
   </dependencies>
 
   <build>

src/main/java/com/databricks/jdbc/api/impl/AbstractDatabricksGeospatial.java

@@ -0,0 +1,109 @@
+package com.databricks.jdbc.api.impl;
+
+import com.databricks.jdbc.api.impl.converters.WKTConverter;
+import com.databricks.jdbc.exception.DatabricksValidationException;
+import com.databricks.jdbc.log.JdbcLogger;
+import com.databricks.jdbc.log.JdbcLoggerFactory;
+import java.util.Objects;
+
+/**
+ * Abstract base class for geospatial data types in Databricks JDBC driver.
+ *
+ * <p>This class provides common functionality for both GEOMETRY and GEOGRAPHY types, including
+ * storage of WKT (Well-Known Text) format data and access to both WKT and WKB representations.
+ */
+public abstract class AbstractDatabricksGeospatial implements DatabricksGeospatial {
+
+  private static final JdbcLogger LOGGER =
+      JdbcLoggerFactory.getLogger(AbstractDatabricksGeospatial.class);
+
+  private final String wkt;
+  private final int srid;
+
+  /**
+   * Constructs an AbstractDatabricksGeospatial with the specified WKT and SRID.
+   *
+   * @param wkt the Well-Known Text representation of the geospatial object
+   * @param srid the Spatial Reference System Identifier
+   * @throws DatabricksValidationException if the WKT is invalid
+   */
+  protected AbstractDatabricksGeospatial(String wkt, int srid)
+      throws DatabricksValidationException {
+    if (wkt == null || wkt.trim().isEmpty()) {
+      throw new DatabricksValidationException("WKT string cannot be null or empty");
+    }
+
+    this.wkt = wkt.trim();
+    this.srid = srid;
+  }
+
+  /**
+   * Returns the Well-Known Binary (WKB) representation of the geospatial object.
+   *
+   * @return the WKB representation as a byte array
+   * @throws DatabricksValidationException if WKT to WKB conversion fails
+   */
+  @Override
+  public byte[] getWkb() throws DatabricksValidationException {
+    return WKTConverter.toWKB(wkt);
+  }
+
+  /**
+   * Returns the Spatial Reference System Identifier (SRID) of the geospatial object.
+   *
+   * @return the SRID value
+   */
+  @Override
+  public int getSrid() {
+    return srid;
+  }
+
+  /**
+   * Returns the Well-Known Text (WKT) representation of the geospatial object.
+   *
+   * @return the WKT string
+   */
+  @Override
+  public String getWkt() {
+    return wkt;
+  }
+
+  /**
+   * Returns a string representation of the geospatial object in EWKT format.
+   *
+   * @return the EWKT string representation
+   */
+  @Override
+  public String toString() {
+    return String.format("SRID=%d;%s", srid, wkt);
+  }
+
+  /**
+   * Checks if this geospatial object is equal to another object.
+   *
+   * @param obj the object to compare
+   * @return true if the objects are equal, false otherwise
+   */
+  @Override
+  public boolean equals(Object obj) {
+    if (this == obj) {
+      return true;
+    }
+    if (obj == null || getClass() != obj.getClass()) {
+      return false;
+    }
+
+    AbstractDatabricksGeospatial that = (AbstractDatabricksGeospatial) obj;
+    return srid == that.srid && wkt.equals(that.wkt);
+  }
+
+  /**
+   * Returns the hash code for this geospatial object.
+   *
+   * @return the hash code
+   */
+  @Override
+  public int hashCode() {
+    return Objects.hash(wkt, srid);
+  }
+}

src/main/java/com/databricks/jdbc/api/impl/DatabricksGeography.java

@@ -0,0 +1,17 @@
+package com.databricks.jdbc.api.impl;
+
+import com.databricks.jdbc.exception.DatabricksValidationException;
+
+public class DatabricksGeography extends AbstractDatabricksGeospatial {
+
+  /**
+   * Constructs a DatabricksGeography with the specified WKT and SRID.
+   *
+   * @param wkt the Well-Known Text representation of the geography
+   * @param srid the Spatial Reference System Identifier
+   * @throws DatabricksValidationException if the WKT is invalid
+   */
+  public DatabricksGeography(String wkt, int srid) throws DatabricksValidationException {
+    super(wkt, srid);
+  }
+}

src/main/java/com/databricks/jdbc/api/impl/DatabricksGeometry.java

@@ -0,0 +1,17 @@
+package com.databricks.jdbc.api.impl;
+
+import com.databricks.jdbc.exception.DatabricksValidationException;
+
+public class DatabricksGeometry extends AbstractDatabricksGeospatial {
+
+  /**
+   * Constructs a DatabricksGeometry with the specified WKT and SRID.
+   *
+   * @param wkt the Well-Known Text representation of the geometry
+   * @param srid the Spatial Reference System Identifier
+   * @throws DatabricksValidationException if the WKT is invalid
+   */
+  public DatabricksGeometry(String wkt, int srid) throws DatabricksValidationException {
+    super(wkt, srid);
+  }
+}

src/main/java/com/databricks/jdbc/api/impl/DatabricksGeospatial.java

@@ -0,0 +1,53 @@
+package com.databricks.jdbc.api.impl;
+
+import com.databricks.jdbc.exception.DatabricksValidationException;
+
+/**
+ * Interface for geospatial data types in Databricks JDBC driver.
+ *
+ * <p>This interface provides common functionality for both GEOMETRY and GEOGRAPHY types, allowing
+ * access to Well-Known Text (WKT), Well-Known Binary (WKB) representation and Spatial Reference
+ * System Identifier (SRID).
+ *
+ * <p>Following the established patterns of DatabricksStruct, DatabricksArray, and DatabricksMap,
+ * this interface enables consistent handling of geospatial data across the JDBC driver.
+ */
+public interface DatabricksGeospatial {
+
+  /**
+   * Returns the Well-Known Binary (WKB) representation of the geospatial object.
+   *
+   * <p>WKB is a binary format for representing geometry data that is compact and suitable for
+   * storage and transmission. This method converts the internal representation to WKB format on
+   * demand.
+   *
+   * @return the WKB representation as a byte array
+   * @throws DatabricksValidationException if WKT to WKB conversion fails
+   */
+  byte[] getWkb() throws DatabricksValidationException;
+
+  /**
+   * Returns the Spatial Reference System Identifier (SRID) of the geospatial object.
+   *
+   * <p>SRID identifies the coordinate system used by the geometry. Common values include:
+   *
+   * <ul>
+   *   <li>4326 - WGS 84 (World Geodetic System 1984)
+   *   <li>3857 - Web Mercator
+   *   <li>0 - No SRID specified
+   * </ul>
+   *
+   * @return the SRID value
+   */
+  int getSrid();
+
+  /**
+   * Returns the Well-Known Text (WKT) representation of the geospatial object.
+   *
+   * <p>WKT is a human-readable text format for representing geometry data. This provides a
+   * complement to the binary WKB format, allowing easy inspection and debugging of geospatial data.
+   *
+   * @return the WKT string representation
+   */
+  String getWkt();
+}

src/main/java/com/databricks/jdbc/api/impl/DatabricksResultSet.java

@@ -1,10 +1,7 @@
 package com.databricks.jdbc.api.impl;
 
 import static com.databricks.jdbc.common.DatabricksJdbcConstants.EMPTY_STRING;
-import static com.databricks.jdbc.common.util.DatabricksTypeUtil.ARRAY;
-import static com.databricks.jdbc.common.util.DatabricksTypeUtil.MAP;
-import static com.databricks.jdbc.common.util.DatabricksTypeUtil.STRUCT;
-import static com.databricks.jdbc.common.util.DatabricksTypeUtil.VARIANT;
+import static com.databricks.jdbc.common.util.DatabricksTypeUtil.*;
 
 import com.databricks.jdbc.api.IDatabricksResultSet;
 import com.databricks.jdbc.api.IExecutionStatus;
@@ -480,13 +477,19 @@ public ResultSetMetaData getMetaData() throws SQLException {
   }
 
   /**
-   * Checks if the given type name represents a complex type (ARRAY, MAP, or STRUCT).
+   * Checks if the given type name represents a complex type (ARRAY, MAP, STRUCT, GEOMETRY, or
+   * GEOGRAPHY).
    *
    * @param typeName The type name to check
-   * @return true if the type name starts with ARRAY, MAP, or STRUCT, false otherwise
+   * @return true if the type name starts with ARRAY, MAP, STRUCT, GEOMETRY, or GEOGRAPHY, false
+   *     otherwise
    */
   private static boolean isComplexType(String typeName) {
-    return typeName.startsWith(ARRAY) || typeName.startsWith(MAP) || typeName.startsWith(STRUCT);
+    return typeName.startsWith(ARRAY)
+        || typeName.startsWith(MAP)
+        || typeName.startsWith(STRUCT)
+        || typeName.startsWith(GEOMETRY)
+        || typeName.startsWith(GEOGRAPHY);
   }
 
   @Override
@@ -531,6 +534,10 @@ private Object handleComplexDataTypesForSEAInline(Object obj, String columnName)
       return parser.parseJsonStringToDbMap(obj.toString(), columnName).toString();
     } else if (columnName.startsWith(STRUCT)) {
       return parser.parseJsonStringToDbStruct(obj.toString(), columnName).toString();
+    } else if (columnName.startsWith(GEOMETRY)) {
+      return obj;
+    } else if (columnName.startsWith(GEOGRAPHY)) {
+      return obj;
     }
     throw new DatabricksParsingException(
         "Unexpected metadata format. Type is not a COMPLEX: " + columnName,
@@ -1969,7 +1976,11 @@ private <T> T getConvertedObject(
       return defaultValue.get();
     }
     int columnType = resultSetMetaData.getColumnType(columnIndex);
-    ObjectConverter converter = ConverterHelper.getConverterForSqlType(columnType);
+    String columnTypeName = resultSetMetaData.getColumnTypeName(columnIndex);
+
+    // Use metadata-aware converter selection for proper handling of databricks-specific types
+    ObjectConverter converter =
+        ConverterHelper.getConverterForColumnType(columnType, columnTypeName);
     return convertMethod.apply(converter, obj);
   }
 

src/main/java/com/databricks/jdbc/api/impl/DatabricksResultSetMetaData.java

@@ -5,10 +5,7 @@
 import static com.databricks.jdbc.common.MetadataResultConstants.LARGE_DISPLAY_COLUMNS;
 import static com.databricks.jdbc.common.MetadataResultConstants.REMARKS_COLUMN;
 import static com.databricks.jdbc.common.util.DatabricksThriftUtil.*;
-import static com.databricks.jdbc.common.util.DatabricksTypeUtil.TIMESTAMP;
-import static com.databricks.jdbc.common.util.DatabricksTypeUtil.TIMESTAMP_NTZ;
-import static com.databricks.jdbc.common.util.DatabricksTypeUtil.VARIANT;
-import static com.databricks.jdbc.common.util.DatabricksTypeUtil.getBasePrecisionAndScale;
+import static com.databricks.jdbc.common.util.DatabricksTypeUtil.*;
 
 import com.databricks.jdbc.api.internal.IDatabricksConnectionContext;
 import com.databricks.jdbc.common.AccessType;
@@ -178,6 +175,7 @@ public DatabricksResultSetMetaData(
             columnIndex < resultManifest.getSchema().getColumnsSize();
             columnIndex++) {
           TColumnDesc columnDesc = resultManifest.getSchema().getColumns().get(columnIndex);
+
           ColumnInfo columnInfo = getColumnInfoFromTColumnDesc(columnDesc);
           int[] precisionAndScale = getPrecisionAndScale(columnInfo);
           int precision = precisionAndScale[0];
@@ -205,6 +203,16 @@ public DatabricksResultSetMetaData(
                 .columnTypeClassName("java.lang.String")
                 .columnType(Types.OTHER)
                 .columnTypeText(VARIANT);
+          } else if (isGeometryColumn(arrowMetadata, columnIndex)) {
+            columnBuilder
+                .columnTypeClassName("com.databricks.jdbc.api.impl.DatabricksGeometry")
+                .columnType(Types.OTHER)
+                .columnTypeText(GEOMETRY);
+          } else if (isGeographyColumn(arrowMetadata, columnIndex)) {
+            columnBuilder
+                .columnTypeClassName("com.databricks.jdbc.api.impl.DatabricksGeography")
+                .columnType(Types.OTHER)
+                .columnTypeText(GEOGRAPHY);
           }
           columnsBuilder.add(columnBuilder.build());
           columnNameToIndexMap.putIfAbsent(columnInfo.getName(), ++currIndex);
@@ -642,6 +650,20 @@ private boolean isVariantColumn(List<String> arrowMetadata, int i) {
         && arrowMetadata.get(i).equalsIgnoreCase(VARIANT);
   }
 
+  private boolean isGeometryColumn(List<String> arrowMetadata, int i) {
+    return arrowMetadata != null
+        && arrowMetadata.size() > i
+        && arrowMetadata.get(i) != null
+        && arrowMetadata.get(i).contains(GEOMETRY);
+  }
+
+  private boolean isGeographyColumn(List<String> arrowMetadata, int i) {
+    return arrowMetadata != null
+        && arrowMetadata.size() > i
+        && arrowMetadata.get(i) != null
+        && arrowMetadata.get(i).contains(GEOGRAPHY);
+  }
+
   private ImmutableDatabricksColumn.Builder getColumnBuilder() {
     return ImmutableDatabricksColumn.builder()
         .isAutoIncrement(false)

src/main/java/com/databricks/jdbc/api/impl/arrow/ArrowStreamResult.java

@@ -160,7 +160,7 @@ public Object getObject(int columnIndex) throws DatabricksSQLException {
   }
 
   /**
-   * Checks if the given type is a complex type (ARRAY, MAP, or STRUCT).
+   * Checks if the given type is a complex type (ARRAY, MAP, STRUCT, GEOMETRY, or GEOGRAPHY).
    *
    * @param type The type to check
    * @return true if the type is a complex type, false otherwise
@@ -169,7 +169,9 @@ public Object getObject(int columnIndex) throws DatabricksSQLException {
   public static boolean isComplexType(ColumnInfoTypeName type) {
     return type == ColumnInfoTypeName.ARRAY
         || type == ColumnInfoTypeName.MAP
-        || type == ColumnInfoTypeName.STRUCT;
+        || type == ColumnInfoTypeName.STRUCT
+        || type == ColumnInfoTypeName.GEOMETRY
+        || type == ColumnInfoTypeName.GEOGRAPHY;
   }
 
   /** {@inheritDoc} */