Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add WKB Support for PostGIS Geometry Columns #9951

Closed
wants to merge 18 commits into from
Closed
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions plugin/trino-postgresql/pom.xml
Original file line number Diff line number Diff line change
@@ -23,6 +23,11 @@
<artifactId>trino-base-jdbc</artifactId>
</dependency>

<dependency>
<groupId>io.trino</groupId>
<artifactId>trino-geospatial</artifactId>
</dependency>

<dependency>
<groupId>io.trino</groupId>
<artifactId>trino-plugin-toolkit</artifactId>
Original file line number Diff line number Diff line change
@@ -31,6 +31,7 @@
import io.trino.plugin.jdbc.JdbcExpression;
import io.trino.plugin.jdbc.JdbcJoinCondition;
import io.trino.plugin.jdbc.JdbcSortItem;
import io.trino.plugin.jdbc.JdbcSplit;
import io.trino.plugin.jdbc.JdbcTableHandle;
import io.trino.plugin.jdbc.JdbcTypeHandle;
import io.trino.plugin.jdbc.LongReadFunction;
@@ -123,8 +124,11 @@
import static com.google.common.base.Preconditions.checkArgument;
import static com.google.common.base.Verify.verify;
import static io.airlift.slice.Slices.utf8Slice;
import static io.airlift.slice.Slices.wrappedBuffer;
import static io.trino.plugin.base.util.JsonTypeUtil.jsonParse;
import static io.trino.plugin.base.util.JsonTypeUtil.toJsonValue;
import static io.trino.plugin.geospatial.GeoFunctions.stAsBinary;
import static io.trino.plugin.geospatial.GeoFunctions.stGeomFromBinary;
import static io.trino.plugin.jdbc.DecimalConfig.DecimalMapping.ALLOW_OVERFLOW;
import static io.trino.plugin.jdbc.DecimalSessionSessionProperties.getDecimalDefaultScale;
import static io.trino.plugin.jdbc.DecimalSessionSessionProperties.getDecimalRounding;
@@ -230,6 +234,7 @@ public class PostgreSqlClient
private static final DateTimeFormatter DATE_FORMATTER = DateTimeFormatter.ofPattern("y-MM-dd[ G]");
private static final DateTimeFormatter TIME_FORMATTER = DateTimeFormatter.ofPattern("HH:mm:ss.SSSSSS");

private final Type geometryType;
private final Type jsonType;
private final Type uuidType;
private final MapType varcharMapType;
@@ -272,6 +277,7 @@ public PostgreSqlClient(
IdentifierMapping identifierMapping)
{
super(config, "\"", connectionFactory, identifierMapping);
this.geometryType = typeManager.getType(new TypeSignature(StandardTypes.GEOMETRY));
this.jsonType = typeManager.getType(new TypeSignature(JSON));
this.uuidType = typeManager.getType(new TypeSignature(StandardTypes.UUID));
this.varcharMapType = (MapType) typeManager.getType(mapType(VARCHAR.getTypeSignature(), VARCHAR.getTypeSignature()));
@@ -307,6 +313,30 @@ public PostgreSqlClient(
.build());
}

@Override
public PreparedStatement buildSql(ConnectorSession session, Connection connection, JdbcSplit split, JdbcTableHandle table, List<JdbcColumnHandle> columns)
throws SQLException
{
Map<String, String> supposedColumnExpressions = new HashMap<>();
for (JdbcColumnHandle column : columns) {
JdbcTypeHandle jdbcTypeHandle = column.getJdbcTypeHandle();
if (jdbcTypeHandle.getJdbcTypeName().isPresent() && jdbcTypeHandle.getJdbcTypeName().get().equals("geometry")) {
String columnName = column.getColumnName();
log.debug("Find geometry type, changing '%s' to '%s'", columnName, "ST_AsBinary(\"" + columnName + "\")");
supposedColumnExpressions.put(columnName, "ST_AsBinary(\"" + columnName + "\")");
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

cc: @ebyhr since you recently were looking into ways to apply functions to columns on the read path. Do you know of a better alternative?

}
}

ImmutableMap.Builder<String, String> builder = ImmutableMap.builder();
for (Map.Entry<String, String> entry : supposedColumnExpressions.entrySet()) {
builder.put(entry.getKey(), entry.getValue());
}
Map<String, String> columnExpressions = builder.buildOrThrow();

PreparedQuery preparedQuery = prepareQuery(session, connection, table, Optional.empty(), columns, columnExpressions, Optional.of(split));
return new QueryBuilder(this).prepareStatement(session, connection, preparedQuery);
}

@Override
public void createTable(ConnectorSession session, ConnectorTableMetadata tableMetadata)
{
@@ -451,6 +481,8 @@ public Optional<ColumnMapping> toColumnMapping(ConnectorSession session, Connect
return mapping;
}
switch (jdbcTypeName) {
case "geometry":
return Optional.of(geometryColumnMapping());
case "money":
return Optional.of(moneyColumnMapping());
case "uuid":
@@ -1210,6 +1242,36 @@ private static JdbcTypeHandle getArrayElementTypeHandle(Connection connection, J
}
}

public static SliceWriteFunction geometryWriteFunction()
{
String bindExpression = format("ST_GeomFromWKB(?)");
return new SliceWriteFunction()
{
@Override
public String getBindExpression()
{
return bindExpression;
}

@Override
public void set(PreparedStatement statement, int index, Slice slice)
throws SQLException
{
byte[] bytes = stAsBinary(slice).getBytes();
statement.setBytes(index, bytes);
}
};
}

private ColumnMapping geometryColumnMapping()
{
return ColumnMapping.sliceMapping(
geometryType,
(resultSet, columnIndex) -> stGeomFromBinary(wrappedBuffer(resultSet.getBytes(columnIndex))),
geometryWriteFunction(),
DISABLE_PUSHDOWN);
}

private ColumnMapping jsonColumnMapping()
{
return ColumnMapping.sliceMapping(
Original file line number Diff line number Diff line change
@@ -17,6 +17,7 @@
import io.airlift.log.Logger;
import io.airlift.log.Logging;
import io.trino.Session;
import io.trino.plugin.geospatial.GeoPlugin;
import io.trino.plugin.jmx.JmxPlugin;
import io.trino.plugin.tpch.TpchPlugin;
import io.trino.testing.DistributedQueryRunner;
@@ -59,6 +60,7 @@ public static DistributedQueryRunner createPostgreSqlQueryRunner(
connectorProperties.putIfAbsent("postgresql.include-system-tables", "true");
//connectorProperties.putIfAbsent("postgresql.experimental.enable-string-pushdown-with-collate", "true");

queryRunner.installPlugin(new GeoPlugin());
queryRunner.installPlugin(new PostgreSqlPlugin());
queryRunner.createCatalog("postgresql", "postgresql", connectorProperties);

Loading