Skip to content

Commit f81439b

Browse files
daindain
committed
Use JTS Geometry as native stack type
Change the internal representation of geometry values to use JTS Geometry objects directly, avoiding unnecessary serialization cycles between function calls.
1 parent 3042fee commit f81439b

22 files changed

+779
-504
lines changed

‎plugin/trino-geospatial/src/main/java/io/trino/plugin/geospatial/AbstractGeometryType.java‎

Lines changed: 461 additions & 0 deletions
Large diffs are not rendered by default.

‎plugin/trino-geospatial/src/main/java/io/trino/plugin/geospatial/BingTileFunctions.java‎

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,6 @@
1616
import com.google.common.collect.ImmutableList;
1717
import com.google.errorprone.annotations.FormatMethod;
1818
import io.airlift.slice.Slice;
19-
import io.trino.geospatial.serde.JtsGeometrySerde;
2019
import io.trino.spi.TrinoException;
2120
import io.trino.spi.block.Block;
2221
import io.trino.spi.block.BlockBuilder;
@@ -330,23 +329,22 @@ public static Block bingTilesAround(
330329
@Description("Given a Bing tile, returns the polygon representation of the tile")
331330
@ScalarFunction("bing_tile_polygon")
332331
@SqlType(StandardTypes.GEOMETRY)
333-
public static Slice bingTilePolygon(@SqlType(StandardTypes.BING_TILE) long input)
332+
public static Geometry bingTilePolygon(@SqlType(StandardTypes.BING_TILE) long input)
334333
{
335334
BingTile tile = BingTile.decode(input);
336335

337-
return JtsGeometrySerde.serialize(tileToEnvelope(tile));
336+
return GEOMETRY_FACTORY.toGeometry(tileToEnvelope(tile));
338337
}
339338

340339
@Description("Given a geometry and a zoom level, returns the minimum set of Bing tiles that fully covers that geometry")
341340
@ScalarFunction("geometry_to_bing_tiles")
342341
@SqlType("array(" + StandardTypes.BING_TILE + ")")
343-
public static Block geometryToBingTiles(@SqlType(StandardTypes.GEOMETRY) Slice input, @SqlType(StandardTypes.INTEGER) long zoomLevelInput)
342+
public static Block geometryToBingTiles(@SqlType(StandardTypes.GEOMETRY) Geometry geometry, @SqlType(StandardTypes.INTEGER) long zoomLevelInput)
344343
{
345344
checkZoomLevel(zoomLevelInput);
346345

347346
int zoomLevel = toIntExact(zoomLevelInput);
348347

349-
Geometry geometry = JtsGeometrySerde.deserialize(input);
350348
if (geometry.isEmpty()) {
351349
return EMPTY_TILE_ARRAY;
352350
}

‎plugin/trino-geospatial/src/main/java/io/trino/plugin/geospatial/EncodedPolylineFunctions.java‎

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,6 @@
1717
import io.airlift.slice.DynamicSliceOutput;
1818
import io.airlift.slice.Slice;
1919
import io.trino.geospatial.GeometryType;
20-
import io.trino.geospatial.serde.JtsGeometrySerde;
2120
import io.trino.spi.TrinoException;
2221
import io.trino.spi.function.Description;
2322
import io.trino.spi.function.ScalarFunction;
@@ -54,9 +53,9 @@ private EncodedPolylineFunctions() {}
5453
@Description("Decodes a polyline to a linestring")
5554
@ScalarFunction("from_encoded_polyline")
5655
@SqlType(StandardTypes.GEOMETRY)
57-
public static Slice fromEncodedPolyline(@SqlType(StandardTypes.VARCHAR) Slice input)
56+
public static Geometry fromEncodedPolyline(@SqlType(StandardTypes.VARCHAR) Slice input)
5857
{
59-
return JtsGeometrySerde.serialize(decodePolyline(input.toStringUtf8()));
58+
return decodePolyline(input.toStringUtf8());
6059
}
6160

6261
private static LineString decodePolyline(String polyline)
@@ -104,9 +103,8 @@ private static LineString decodePolyline(String polyline)
104103
@Description("Encodes a linestring or multipoint geometry to a polyline")
105104
@ScalarFunction("to_encoded_polyline")
106105
@SqlType(StandardTypes.VARCHAR)
107-
public static Slice toEncodedPolyline(@SqlType(StandardTypes.GEOMETRY) Slice input)
106+
public static Slice toEncodedPolyline(@SqlType(StandardTypes.GEOMETRY) Geometry geometry)
108107
{
109-
Geometry geometry = JtsGeometrySerde.deserialize(input);
110108
validateType("encode_polyline", geometry, Set.of(LINE_STRING, MULTI_POINT));
111109
GeometryType geometryType = GeometryType.getForJtsGeometryType(geometry.getGeometryType());
112110
return switch (geometryType) {

‎plugin/trino-geospatial/src/main/java/io/trino/plugin/geospatial/GeoFunctions.java‎

Lines changed: 208 additions & 256 deletions
Large diffs are not rendered by default.

‎plugin/trino-geospatial/src/main/java/io/trino/plugin/geospatial/GeometryType.java‎

Lines changed: 2 additions & 111 deletions
Original file line numberDiff line numberDiff line change
@@ -13,131 +13,22 @@
1313
*/
1414
package io.trino.plugin.geospatial;
1515

16-
import io.airlift.slice.Slice;
17-
import io.airlift.slice.XxHash64;
18-
import io.trino.spi.block.Block;
19-
import io.trino.spi.block.BlockBuilder;
20-
import io.trino.spi.block.VariableWidthBlock;
21-
import io.trino.spi.block.VariableWidthBlockBuilder;
22-
import io.trino.spi.function.IsNull;
23-
import io.trino.spi.function.ScalarOperator;
24-
import io.trino.spi.type.AbstractVariableWidthType;
25-
import io.trino.spi.type.TypeOperatorDeclaration;
26-
import io.trino.spi.type.TypeOperators;
2716
import io.trino.spi.type.TypeSignature;
2817

29-
import static io.trino.geospatial.serde.JtsGeometrySerde.deserialize;
30-
import static io.trino.spi.function.OperatorType.EQUAL;
31-
import static io.trino.spi.function.OperatorType.HASH_CODE;
32-
import static io.trino.spi.function.OperatorType.IDENTICAL;
33-
import static io.trino.spi.function.OperatorType.XX_HASH_64;
34-
3518
public class GeometryType
36-
extends AbstractVariableWidthType
19+
extends AbstractGeometryType
3720
{
3821
public static final String NAME = "Geometry";
3922
public static final GeometryType GEOMETRY = new GeometryType();
4023

41-
private static final TypeOperatorDeclaration TYPE_OPERATOR_DECLARATION =
42-
TypeOperatorDeclaration.builder(Slice.class)
43-
.addOperators(DEFAULT_READ_OPERATORS)
44-
.addOperators(DEFAULT_COMPARABLE_OPERATORS)
45-
.build();
46-
4724
private GeometryType()
4825
{
49-
super(new TypeSignature(NAME), Slice.class);
50-
}
51-
52-
protected GeometryType(TypeSignature signature)
53-
{
54-
super(signature, Slice.class);
55-
}
56-
57-
@Override
58-
public TypeOperatorDeclaration getTypeOperatorDeclaration(TypeOperators typeOperators)
59-
{
60-
return TYPE_OPERATOR_DECLARATION;
26+
super(new TypeSignature(NAME));
6127
}
6228

6329
@Override
6430
public String getDisplayName()
6531
{
6632
return NAME;
6733
}
68-
69-
@Override
70-
public boolean isComparable()
71-
{
72-
return true;
73-
}
74-
75-
@Override
76-
public Slice getSlice(Block block, int position)
77-
{
78-
VariableWidthBlock valueBlock = (VariableWidthBlock) block.getUnderlyingValueBlock();
79-
int valuePosition = block.getUnderlyingValuePosition(position);
80-
return valueBlock.getSlice(valuePosition);
81-
}
82-
83-
@Override
84-
public void writeSlice(BlockBuilder blockBuilder, Slice value)
85-
{
86-
if (value == null) {
87-
blockBuilder.appendNull();
88-
return;
89-
}
90-
writeSlice(blockBuilder, value, 0, value.length());
91-
}
92-
93-
@Override
94-
public void writeSlice(BlockBuilder blockBuilder, Slice value, int offset, int length)
95-
{
96-
if (value == null) {
97-
blockBuilder.appendNull();
98-
return;
99-
}
100-
((VariableWidthBlockBuilder) blockBuilder).writeEntry(value, offset, length);
101-
}
102-
103-
@Override
104-
public Object getObjectValue(Block block, int position)
105-
{
106-
if (block.isNull(position)) {
107-
return null;
108-
}
109-
try {
110-
return deserialize(getSlice(block, position)).toText();
111-
}
112-
catch (Exception e) {
113-
return "<invalid geometry>";
114-
}
115-
}
116-
117-
@ScalarOperator(HASH_CODE)
118-
private static long hashCodeOperator(Slice value)
119-
{
120-
return value.hashCode();
121-
}
122-
123-
@ScalarOperator(XX_HASH_64)
124-
private static long xxHash64Operator(Slice value)
125-
{
126-
return XxHash64.hash(value);
127-
}
128-
129-
@ScalarOperator(EQUAL)
130-
private static boolean equalOperator(Slice left, Slice right)
131-
{
132-
return left.equals(right);
133-
}
134-
135-
@ScalarOperator(IDENTICAL)
136-
private static boolean identical(Slice left, @IsNull boolean leftNull, Slice right, @IsNull boolean rightNull)
137-
{
138-
if (leftNull || rightNull) {
139-
return leftNull == rightNull;
140-
}
141-
return left.equals(right);
142-
}
14334
}

‎plugin/trino-geospatial/src/main/java/io/trino/plugin/geospatial/SpatialPartitioningAggregateFunction.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -13,13 +13,13 @@
1313
*/
1414
package io.trino.plugin.geospatial;
1515

16-
import io.airlift.slice.Slice;
1716
import io.trino.spi.block.BlockBuilder;
1817
import io.trino.spi.function.AggregationFunction;
1918
import io.trino.spi.function.InputFunction;
2019
import io.trino.spi.function.OutputFunction;
2120
import io.trino.spi.function.SqlType;
2221
import io.trino.spi.type.StandardTypes;
22+
import org.locationtech.jts.geom.Geometry;
2323

2424
import static io.trino.plugin.geospatial.SpatialPartitioningAggregateFunction.NAME;
2525

@@ -31,7 +31,7 @@ public final class SpatialPartitioningAggregateFunction
3131
private SpatialPartitioningAggregateFunction() {}
3232

3333
@InputFunction
34-
public static void input(SpatialPartitioningState state, @SqlType(StandardTypes.GEOMETRY) Slice slice)
34+
public static void input(SpatialPartitioningState state, @SqlType(StandardTypes.GEOMETRY) Geometry geometry)
3535
{
3636
throw new UnsupportedOperationException("spatial_partitioning(geometry) aggregate function should be re-written into spatial_partitioning(geometry, partitionCount)");
3737
}

‎plugin/trino-geospatial/src/main/java/io/trino/plugin/geospatial/SpatialPartitioningInternalAggregateFunction.java‎

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,6 @@
1313
*/
1414
package io.trino.plugin.geospatial;
1515

16-
import io.airlift.slice.Slice;
1716
import io.trino.geospatial.KdbTreeUtils;
1817
import io.trino.geospatial.Rectangle;
1918
import io.trino.spi.block.BlockBuilder;
@@ -23,13 +22,13 @@
2322
import io.trino.spi.function.SqlType;
2423
import io.trino.spi.type.StandardTypes;
2524
import org.locationtech.jts.geom.Envelope;
25+
import org.locationtech.jts.geom.Geometry;
2626

2727
import java.util.ArrayList;
2828
import java.util.List;
2929
import java.util.concurrent.ThreadLocalRandom;
3030

3131
import static io.trino.geospatial.KdbTree.buildKdbTree;
32-
import static io.trino.geospatial.serde.JtsGeometrySerde.deserializeEnvelope;
3332
import static io.trino.plugin.geospatial.SpatialPartitioningAggregateFunction.NAME;
3433
import static io.trino.spi.type.StandardTypes.INTEGER;
3534
import static io.trino.spi.type.VarcharType.VARCHAR;
@@ -43,9 +42,9 @@ public final class SpatialPartitioningInternalAggregateFunction
4342
private SpatialPartitioningInternalAggregateFunction() {}
4443

4544
@InputFunction
46-
public static void input(SpatialPartitioningState state, @SqlType(StandardTypes.GEOMETRY) Slice slice, @SqlType(INTEGER) long partitionCount)
45+
public static void input(SpatialPartitioningState state, @SqlType(StandardTypes.GEOMETRY) Geometry geometry, @SqlType(INTEGER) long partitionCount)
4746
{
48-
Envelope envelope = deserializeEnvelope(slice);
47+
Envelope envelope = geometry.getEnvelopeInternal();
4948
if (envelope.isNull()) {
5049
return;
5150
}

‎plugin/trino-geospatial/src/main/java/io/trino/plugin/geospatial/SphericalGeographyType.java‎

Lines changed: 2 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -13,64 +13,22 @@
1313
*/
1414
package io.trino.plugin.geospatial;
1515

16-
import io.airlift.slice.Slice;
17-
import io.trino.spi.block.Block;
18-
import io.trino.spi.block.BlockBuilder;
19-
import io.trino.spi.block.VariableWidthBlock;
20-
import io.trino.spi.block.VariableWidthBlockBuilder;
21-
import io.trino.spi.type.AbstractVariableWidthType;
2216
import io.trino.spi.type.TypeSignature;
2317

24-
import static io.trino.geospatial.serde.JtsGeometrySerde.deserialize;
25-
2618
public class SphericalGeographyType
27-
extends AbstractVariableWidthType
19+
extends AbstractGeometryType
2820
{
2921
public static final String NAME = "SphericalGeography";
3022
public static final SphericalGeographyType SPHERICAL_GEOGRAPHY = new SphericalGeographyType();
3123

3224
private SphericalGeographyType()
3325
{
34-
super(new TypeSignature(NAME), Slice.class);
35-
}
36-
37-
@Override
38-
public Slice getSlice(Block block, int position)
39-
{
40-
VariableWidthBlock valueBlock = (VariableWidthBlock) block.getUnderlyingValueBlock();
41-
int valuePosition = block.getUnderlyingValuePosition(position);
42-
return valueBlock.getSlice(valuePosition);
43-
}
44-
45-
@Override
46-
public void writeSlice(BlockBuilder blockBuilder, Slice value)
47-
{
48-
writeSlice(blockBuilder, value, 0, value.length());
49-
}
50-
51-
@Override
52-
public void writeSlice(BlockBuilder blockBuilder, Slice value, int offset, int length)
53-
{
54-
((VariableWidthBlockBuilder) blockBuilder).writeEntry(value, offset, length);
26+
super(new TypeSignature(NAME));
5527
}
5628

5729
@Override
5830
public String getDisplayName()
5931
{
6032
return NAME;
6133
}
62-
63-
@Override
64-
public Object getObjectValue(Block block, int position)
65-
{
66-
if (block.isNull(position)) {
67-
return null;
68-
}
69-
try {
70-
return deserialize(getSlice(block, position)).toText();
71-
}
72-
catch (Exception e) {
73-
return "<invalid geometry>";
74-
}
75-
}
7634
}

‎plugin/trino-geospatial/src/main/java/io/trino/plugin/geospatial/aggregation/ConvexHullAggregation.java‎

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -13,8 +13,6 @@
1313
*/
1414
package io.trino.plugin.geospatial.aggregation;
1515

16-
import io.airlift.slice.Slice;
17-
import io.trino.geospatial.serde.JtsGeometrySerde;
1816
import io.trino.spi.block.BlockBuilder;
1917
import io.trino.spi.function.AggregationFunction;
2018
import io.trino.spi.function.AggregationState;
@@ -42,9 +40,8 @@ private ConvexHullAggregation() {}
4240

4341
@InputFunction
4442
public static void input(@AggregationState GeometryState state,
45-
@SqlType(StandardTypes.GEOMETRY) Slice input)
43+
@SqlType(StandardTypes.GEOMETRY) Geometry geometry)
4644
{
47-
Geometry geometry = JtsGeometrySerde.deserialize(input);
4845
if (state.getGeometry() == null) {
4946
Geometry result = geometry.convexHull();
5047
result.setSRID(geometry.getSRID());
@@ -80,7 +77,7 @@ public static void output(@AggregationState GeometryState state, BlockBuilder ou
8077
out.appendNull();
8178
}
8279
else {
83-
GEOMETRY.writeSlice(out, JtsGeometrySerde.serialize(state.getGeometry()));
80+
GEOMETRY.writeObject(out, state.getGeometry());
8481
}
8582
}
8683
}

‎plugin/trino-geospatial/src/main/java/io/trino/plugin/geospatial/aggregation/GeometryUnionAgg.java‎

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -13,8 +13,6 @@
1313
*/
1414
package io.trino.plugin.geospatial.aggregation;
1515

16-
import io.airlift.slice.Slice;
17-
import io.trino.geospatial.serde.JtsGeometrySerde;
1816
import io.trino.spi.block.BlockBuilder;
1917
import io.trino.spi.function.AggregationFunction;
2018
import io.trino.spi.function.AggregationState;
@@ -41,9 +39,8 @@ public final class GeometryUnionAgg
4139
private GeometryUnionAgg() {}
4240

4341
@InputFunction
44-
public static void input(@AggregationState GeometryState state, @SqlType(StandardTypes.GEOMETRY) Slice input)
42+
public static void input(@AggregationState GeometryState state, @SqlType(StandardTypes.GEOMETRY) Geometry geometry)
4543
{
46-
Geometry geometry = JtsGeometrySerde.deserialize(input);
4744
if (state.getGeometry() == null) {
4845
state.setGeometry(geometry);
4946
}
@@ -76,7 +73,7 @@ public static void output(@AggregationState GeometryState state, BlockBuilder ou
7673
out.appendNull();
7774
}
7875
else {
79-
GEOMETRY.writeSlice(out, JtsGeometrySerde.serialize(state.getGeometry()));
76+
GEOMETRY.writeObject(out, state.getGeometry());
8077
}
8178
}
8279
}

0 commit comments

Comments
 (0)