From 4d15916d07430efce42b97cd26d5d6dd7a3fb4d0 Mon Sep 17 00:00:00 2001 From: Mariam Almesfer Date: Tue, 3 Jun 2025 15:25:09 +0300 Subject: [PATCH] Add read support for geometry and geography data types in postgres connector Co-authored-by: pratyakshsharma 3@gmail.com> Co-authored-by: agrawalreetika --- presto-base-jdbc/pom.xml | 11 +++ .../presto/plugin/jdbc/GeometryUtils.java | 50 ++++++++++ .../jdbc/mapping/StandardColumnMappings.java | 8 ++ .../src/main/sphinx/connector/postgresql.rst | 5 +- presto-mysql/pom.xml | 5 - .../presto/plugin/mysql/MySqlClient.java | 38 +------- .../presto/plugin/mysql/TestMySqlClient.java | 4 +- presto-postgresql/pom.xml | 17 ++++ .../plugin/postgresql/PostgreSqlClient.java | 39 ++++++++ .../postgresql/TestPostgreSqlClient.java | 96 +++++++++++++++++++ 10 files changed, 228 insertions(+), 45 deletions(-) create mode 100644 presto-base-jdbc/src/main/java/com/facebook/presto/plugin/jdbc/GeometryUtils.java create mode 100644 presto-postgresql/src/test/java/com/facebook/presto/plugin/postgresql/TestPostgreSqlClient.java diff --git a/presto-base-jdbc/pom.xml b/presto-base-jdbc/pom.xml index e1ca899ecd3a3..d1bf195b9ea88 100644 --- a/presto-base-jdbc/pom.xml +++ b/presto-base-jdbc/pom.xml @@ -110,6 +110,17 @@ jmxutils + + com.esri.geometry + esri-geometry-api + + + + com.facebook.presto + presto-geospatial-toolkit + provided + + diff --git a/presto-base-jdbc/src/main/java/com/facebook/presto/plugin/jdbc/GeometryUtils.java b/presto-base-jdbc/src/main/java/com/facebook/presto/plugin/jdbc/GeometryUtils.java new file mode 100644 index 0000000000000..8e7f2fc5d8134 --- /dev/null +++ b/presto-base-jdbc/src/main/java/com/facebook/presto/plugin/jdbc/GeometryUtils.java @@ -0,0 +1,50 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.plugin.jdbc; + +import com.esri.core.geometry.ogc.OGCGeometry; +import com.facebook.presto.spi.PrestoException; +import io.airlift.slice.Slice; + +import static com.esri.core.geometry.ogc.OGCGeometry.fromBinary; +import static com.facebook.presto.geospatial.GeometryUtils.wktFromJtsGeometry; +import static com.facebook.presto.geospatial.serde.EsriGeometrySerde.serialize; +import static com.facebook.presto.geospatial.serde.JtsGeometrySerde.deserialize; +import static com.facebook.presto.spi.StandardErrorCode.INVALID_FUNCTION_ARGUMENT; +import static io.airlift.slice.Slices.utf8Slice; +import static java.util.Objects.requireNonNull; + +public class GeometryUtils +{ + private GeometryUtils() {} + + public static Slice getAsText(Slice input) + { + return utf8Slice(wktFromJtsGeometry(deserialize(input))); + } + + public static Slice stGeomFromBinary(Slice input) + { + requireNonNull(input, "input is null"); + OGCGeometry geometry; + try { + geometry = fromBinary(input.toByteBuffer().slice()); + } + catch (IllegalArgumentException | IndexOutOfBoundsException e) { + throw new PrestoException(INVALID_FUNCTION_ARGUMENT, "Invalid Well-Known Binary (WKB)", e); + } + geometry.setSpatialReference(null); + return serialize(geometry); + } +} diff --git a/presto-base-jdbc/src/main/java/com/facebook/presto/plugin/jdbc/mapping/StandardColumnMappings.java b/presto-base-jdbc/src/main/java/com/facebook/presto/plugin/jdbc/mapping/StandardColumnMappings.java index 8eb53ec5ccb17..3d705e4a2e94e 100644 --- a/presto-base-jdbc/src/main/java/com/facebook/presto/plugin/jdbc/mapping/StandardColumnMappings.java +++ b/presto-base-jdbc/src/main/java/com/facebook/presto/plugin/jdbc/mapping/StandardColumnMappings.java @@ -61,8 +61,11 @@ import static com.facebook.presto.common.type.UuidType.UUID; import static com.facebook.presto.common.type.UuidType.prestoUuidToJavaUuid; import static com.facebook.presto.common.type.VarbinaryType.VARBINARY; +import static com.facebook.presto.common.type.VarcharType.VARCHAR; import static com.facebook.presto.common.type.VarcharType.createUnboundedVarcharType; import static com.facebook.presto.common.type.VarcharType.createVarcharType; +import static com.facebook.presto.plugin.jdbc.GeometryUtils.getAsText; +import static com.facebook.presto.plugin.jdbc.GeometryUtils.stGeomFromBinary; import static com.facebook.presto.plugin.jdbc.mapping.ReadMapping.createBooleanReadMapping; import static com.facebook.presto.plugin.jdbc.mapping.ReadMapping.createDoubleReadMapping; import static com.facebook.presto.plugin.jdbc.mapping.ReadMapping.createLongReadMapping; @@ -445,4 +448,9 @@ else if (type instanceof UuidType) { } return Optional.empty(); } + public static ReadMapping geometryReadMapping() + { + return createSliceReadMapping(VARCHAR, + (resultSet, columnIndex) -> getAsText(stGeomFromBinary(wrappedBuffer(resultSet.getBytes(columnIndex))))); + } } diff --git a/presto-docs/src/main/sphinx/connector/postgresql.rst b/presto-docs/src/main/sphinx/connector/postgresql.rst index 7d0a2a4a49feb..256f7f56e72bb 100644 --- a/presto-docs/src/main/sphinx/connector/postgresql.rst +++ b/presto-docs/src/main/sphinx/connector/postgresql.rst @@ -145,7 +145,10 @@ The connector maps PostgreSQL types to the corresponding PrestoDB types: - ``JSON`` * - ``JSONB`` - ``JSON`` - + * - ``GEOMETRY`` + - ``VARCHAR`` + * - ``GEOGRAPHY`` + - ``VARCHAR`` No other types are supported. PrestoDB to PostgreSQL type mapping diff --git a/presto-mysql/pom.xml b/presto-mysql/pom.xml index 9c238229c0bb9..c1a18a351f7dc 100644 --- a/presto-mysql/pom.xml +++ b/presto-mysql/pom.xml @@ -77,11 +77,6 @@ provided - - com.facebook.presto - presto-geospatial-toolkit - - com.facebook.drift drift-api diff --git a/presto-mysql/src/main/java/com/facebook/presto/plugin/mysql/MySqlClient.java b/presto-mysql/src/main/java/com/facebook/presto/plugin/mysql/MySqlClient.java index ade9d25b27f81..3277bc25221a2 100644 --- a/presto-mysql/src/main/java/com/facebook/presto/plugin/mysql/MySqlClient.java +++ b/presto-mysql/src/main/java/com/facebook/presto/plugin/mysql/MySqlClient.java @@ -13,7 +13,6 @@ */ package com.facebook.presto.plugin.mysql; -import com.esri.core.geometry.ogc.OGCGeometry; import com.facebook.presto.common.type.TimestampType; import com.facebook.presto.common.type.Type; import com.facebook.presto.common.type.VarcharType; @@ -36,7 +35,6 @@ import com.google.common.collect.ImmutableSet; import com.mysql.cj.jdbc.JdbcStatement; import com.mysql.jdbc.Driver; -import io.airlift.slice.Slice; import javax.inject.Inject; @@ -52,31 +50,22 @@ import java.util.Optional; import java.util.Properties; -import static com.esri.core.geometry.ogc.OGCGeometry.fromBinary; import static com.facebook.presto.common.type.RealType.REAL; import static com.facebook.presto.common.type.StandardTypes.GEOMETRY; import static com.facebook.presto.common.type.TimeWithTimeZoneType.TIME_WITH_TIME_ZONE; import static com.facebook.presto.common.type.TimestampWithTimeZoneType.TIMESTAMP_WITH_TIME_ZONE; import static com.facebook.presto.common.type.VarbinaryType.VARBINARY; -import static com.facebook.presto.common.type.VarcharType.VARCHAR; import static com.facebook.presto.common.type.Varchars.isVarcharType; -import static com.facebook.presto.geospatial.GeometryUtils.wktFromJtsGeometry; -import static com.facebook.presto.geospatial.serde.EsriGeometrySerde.serialize; -import static com.facebook.presto.geospatial.serde.JtsGeometrySerde.deserialize; import static com.facebook.presto.plugin.jdbc.DriverConnectionFactory.basicConnectionProperties; import static com.facebook.presto.plugin.jdbc.JdbcErrorCode.JDBC_ERROR; import static com.facebook.presto.plugin.jdbc.QueryBuilder.quote; -import static com.facebook.presto.plugin.jdbc.mapping.ReadMapping.createSliceReadMapping; +import static com.facebook.presto.plugin.jdbc.mapping.StandardColumnMappings.geometryReadMapping; import static com.facebook.presto.spi.StandardErrorCode.ALREADY_EXISTS; -import static com.facebook.presto.spi.StandardErrorCode.INVALID_FUNCTION_ARGUMENT; import static com.facebook.presto.spi.StandardErrorCode.NOT_SUPPORTED; import static com.google.common.collect.ImmutableMap.toImmutableMap; import static com.google.common.util.concurrent.MoreExecutors.directExecutor; -import static io.airlift.slice.Slices.utf8Slice; -import static io.airlift.slice.Slices.wrappedBuffer; import static java.lang.String.format; import static java.util.Locale.ENGLISH; -import static java.util.Objects.requireNonNull; import static java.util.function.Function.identity; public class MySqlClient @@ -253,31 +242,6 @@ public Optional toPrestoType(ConnectorSession session, JdbcTypeHand return super.toPrestoType(session, typeHandle); } - protected static ReadMapping geometryReadMapping() - { - return createSliceReadMapping(VARCHAR, - (resultSet, columnIndex) -> getAsText(stGeomFromBinary(wrappedBuffer(resultSet.getBytes(columnIndex))))); - } - - protected static Slice getAsText(Slice input) - { - return utf8Slice(wktFromJtsGeometry(deserialize(input))); - } - - private static Slice stGeomFromBinary(Slice input) - { - requireNonNull(input, "input is null"); - OGCGeometry geometry; - try { - geometry = fromBinary(input.toByteBuffer().slice()); - } - catch (IllegalArgumentException | IndexOutOfBoundsException e) { - throw new PrestoException(INVALID_FUNCTION_ARGUMENT, "Invalid Well-Known Binary (WKB)", e); - } - geometry.setSpatialReference(null); - return serialize(geometry); - } - @Override public void createTable(ConnectorSession session, ConnectorTableMetadata tableMetadata) { diff --git a/presto-mysql/src/test/java/com/facebook/presto/plugin/mysql/TestMySqlClient.java b/presto-mysql/src/test/java/com/facebook/presto/plugin/mysql/TestMySqlClient.java index e1e67d3cf3729..10096bde0fba5 100644 --- a/presto-mysql/src/test/java/com/facebook/presto/plugin/mysql/TestMySqlClient.java +++ b/presto-mysql/src/test/java/com/facebook/presto/plugin/mysql/TestMySqlClient.java @@ -27,8 +27,8 @@ import java.sql.SQLException; import static com.facebook.presto.geospatial.GeoFunctions.stGeomFromBinary; -import static com.facebook.presto.plugin.mysql.MySqlClient.geometryReadMapping; -import static com.facebook.presto.plugin.mysql.MySqlClient.getAsText; +import static com.facebook.presto.plugin.jdbc.GeometryUtils.getAsText; +import static com.facebook.presto.plugin.jdbc.mapping.StandardColumnMappings.geometryReadMapping; import static org.testng.Assert.assertEquals; import static org.testng.Assert.fail; diff --git a/presto-postgresql/pom.xml b/presto-postgresql/pom.xml index 57e0435561299..4cd35d47850a7 100644 --- a/presto-postgresql/pom.xml +++ b/presto-postgresql/pom.xml @@ -100,7 +100,24 @@ provided + + com.esri.geometry + esri-geometry-api + + + + org.openjdk.jol + jol-core + provided + + + + com.h2database + h2 + test + + com.facebook.presto presto-testng-services diff --git a/presto-postgresql/src/main/java/com/facebook/presto/plugin/postgresql/PostgreSqlClient.java b/presto-postgresql/src/main/java/com/facebook/presto/plugin/postgresql/PostgreSqlClient.java index a9cc35257f62c..8de1c4ade3e84 100644 --- a/presto-postgresql/src/main/java/com/facebook/presto/plugin/postgresql/PostgreSqlClient.java +++ b/presto-postgresql/src/main/java/com/facebook/presto/plugin/postgresql/PostgreSqlClient.java @@ -21,9 +21,12 @@ import com.facebook.presto.plugin.jdbc.BaseJdbcClient; import com.facebook.presto.plugin.jdbc.BaseJdbcConfig; import com.facebook.presto.plugin.jdbc.DriverConnectionFactory; +import com.facebook.presto.plugin.jdbc.JdbcColumnHandle; import com.facebook.presto.plugin.jdbc.JdbcConnectorId; import com.facebook.presto.plugin.jdbc.JdbcIdentity; +import com.facebook.presto.plugin.jdbc.JdbcSplit; import com.facebook.presto.plugin.jdbc.JdbcTypeHandle; +import com.facebook.presto.plugin.jdbc.QueryBuilder; import com.facebook.presto.plugin.jdbc.mapping.ReadMapping; import com.facebook.presto.spi.ConnectorSession; import com.facebook.presto.spi.ConnectorTableMetadata; @@ -48,22 +51,30 @@ import java.sql.PreparedStatement; import java.sql.ResultSet; import java.sql.SQLException; +import java.util.List; +import java.util.Map; import java.util.Optional; import java.util.UUID; +import static com.facebook.presto.common.type.StandardTypes.GEOMETRY; +import static com.facebook.presto.common.type.StandardTypes.SPHERICAL_GEOGRAPHY; import static com.facebook.presto.common.type.VarbinaryType.VARBINARY; import static com.facebook.presto.plugin.jdbc.JdbcErrorCode.JDBC_ERROR; +import static com.facebook.presto.plugin.jdbc.QueryBuilder.quote; import static com.facebook.presto.plugin.jdbc.mapping.ReadMapping.createSliceReadMapping; +import static com.facebook.presto.plugin.jdbc.mapping.StandardColumnMappings.geometryReadMapping; import static com.facebook.presto.spi.StandardErrorCode.ALREADY_EXISTS; import static com.facebook.presto.spi.StandardErrorCode.INVALID_FUNCTION_ARGUMENT; import static com.fasterxml.jackson.core.JsonFactory.Feature.CANONICALIZE_FIELD_NAMES; import static com.fasterxml.jackson.databind.SerializationFeature.ORDER_MAP_ENTRIES_BY_KEYS; +import static com.google.common.collect.ImmutableMap.toImmutableMap; import static io.airlift.slice.Slices.utf8Slice; import static io.airlift.slice.Slices.wrappedLongArray; import static java.lang.Long.reverseBytes; import static java.lang.String.format; import static java.nio.charset.StandardCharsets.UTF_8; import static java.util.Locale.ENGLISH; +import static java.util.function.Function.identity; public class PostgreSqlClient extends BaseJdbcClient @@ -116,9 +127,34 @@ protected String toSqlType(Type type) return super.toSqlType(type); } + @Override + public PreparedStatement buildSql(ConnectorSession session, Connection connection, JdbcSplit split, List columnHandles) throws SQLException + { + Map columnExpressions = columnHandles.stream() + .filter(handle -> handle.getJdbcTypeHandle().getJdbcTypeName().equalsIgnoreCase(GEOMETRY)) + .map(JdbcColumnHandle::getColumnName) + .collect(toImmutableMap( + identity(), + columnName -> "ST_AsBinary(" + quote(identifierQuote, columnName) + ")")); + + return new QueryBuilder(identifierQuote).buildSql( + this, + session, + connection, + split.getCatalogName(), + split.getSchemaName(), + split.getTableName(), + columnHandles, + columnExpressions, + split.getTupleDomain(), + split.getAdditionalPredicate()); + } + @Override public Optional toPrestoType(ConnectorSession session, JdbcTypeHandle typeHandle) { + String typeName = typeHandle.getJdbcTypeName(); + if (typeHandle.getJdbcTypeName().equals("jsonb") || typeHandle.getJdbcTypeName().equals("json")) { return Optional.of(jsonReadMapping()); } @@ -126,6 +162,9 @@ public Optional toPrestoType(ConnectorSession session, JdbcTypeHand else if (typeHandle.getJdbcTypeName().equals("uuid")) { return Optional.of(uuidReadMapping()); } + else if (typeName.equalsIgnoreCase(GEOMETRY) || typeName.equalsIgnoreCase(SPHERICAL_GEOGRAPHY)) { + return Optional.of(geometryReadMapping()); + } return super.toPrestoType(session, typeHandle); } diff --git a/presto-postgresql/src/test/java/com/facebook/presto/plugin/postgresql/TestPostgreSqlClient.java b/presto-postgresql/src/test/java/com/facebook/presto/plugin/postgresql/TestPostgreSqlClient.java new file mode 100644 index 0000000000000..11f61804a1190 --- /dev/null +++ b/presto-postgresql/src/test/java/com/facebook/presto/plugin/postgresql/TestPostgreSqlClient.java @@ -0,0 +1,96 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.plugin.postgresql; + +import com.esri.core.geometry.Point; +import com.esri.core.geometry.ogc.OGCGeometry; +import com.esri.core.geometry.ogc.OGCPoint; +import com.facebook.presto.plugin.jdbc.mapping.functions.SliceReadFunction; +import com.facebook.presto.spi.PrestoException; +import io.airlift.slice.Slice; +import io.airlift.slice.Slices; +import org.h2.tools.SimpleResultSet; +import org.testng.annotations.Test; + +import java.nio.ByteBuffer; +import java.sql.SQLException; + +import static com.facebook.presto.geospatial.GeoFunctions.stGeomFromBinary; +import static com.facebook.presto.plugin.jdbc.GeometryUtils.getAsText; +import static com.facebook.presto.plugin.jdbc.mapping.StandardColumnMappings.geometryReadMapping; +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.fail; + +public class TestPostgreSqlClient +{ + @Test + public void testValidGeometryReadMapping() + { + OGCGeometry geometry = new OGCPoint(new Point(1.0, 2.0), null); + ByteBuffer buffer = geometry.asBinary(); + + Slice value = getAsText(stGeomFromBinary(Slices.wrappedBuffer(buffer))); + + assertEquals(value.toStringUtf8(), "POINT (1 2)"); + } + + @Test + public void testInvalidGeometryReadMapping() + { + byte[] invalidWkb = new byte[]{0x00, 0x01, 0x02}; + try { + getAsText(stGeomFromBinary(Slices.wrappedBuffer(invalidWkb))); + fail("stGeomFromBinary should throw"); + } + catch (PrestoException e) { + assertEquals(e.getMessage(), "Invalid WKB"); + } + } + + @Test + public void testReadMapping() throws SQLException + { + SliceReadFunction fn = (SliceReadFunction) geometryReadMapping().getReadFunction(); + OGCGeometry geometry = new OGCPoint(new Point(1.0, 2.0), null); + ByteBuffer buffer = geometry.asBinary(); + Slice value = fn.readSlice(new MockResultSet(buffer.array()), 1); + assertEquals(value.toStringUtf8(), "POINT (1 2)"); + + byte[] invalid = new byte[] {0x01, 0x02, 0x03}; + try { + fn.readSlice(new MockResultSet(invalid), 1); + fail("stGeomFromBinary should throw"); + } + catch (PrestoException e) { + assertEquals(e.getMessage(), "Invalid Well-Known Binary (WKB)"); + } + } + + private static class MockResultSet + extends SimpleResultSet + { + private final byte[] bytes; + + public MockResultSet(byte[] bytes) + { + this.bytes = bytes; + } + + @Override + public byte[] getBytes(int columnIndex) + { + return bytes; + } + } +}