This is an automated email from the ASF dual-hosted git repository. jackie pushed a commit to branch master in repository https://gitbox.apache.org/repos/asf/pinot.git
The following commit(s) were added to refs/heads/master by this push: new b34e805ced Add GeoJSON support. (#14405) b34e805ced is described below commit b34e805cedb1b03eb0677c4010f801c5d518df5a Author: Bolek Ziobrowski <26925920+bziobrow...@users.noreply.github.com> AuthorDate: Tue Nov 26 20:46:09 2024 +0100 Add GeoJSON support. (#14405) --- .../common/function/TransformFunctionType.java | 6 + .../function/BaseBinaryGeoTransformFunction.java | 57 ++++----- ...tion.java => ConstructFromGeoJsonFunction.java} | 46 ++++---- .../function/ConstructFromTextFunction.java | 18 +-- .../function/ConstructFromWKBFunction.java | 13 +-- .../transform/function/GeoToH3Function.java | 18 ++- .../transform/function/ScalarFunctions.java | 29 +++++ .../transform/function/StAreaFunction.java | 14 +-- .../transform/function/StAsBinaryFunction.java | 13 +-- ...sTextFunction.java => StAsGeoJsonFunction.java} | 44 ++++--- .../transform/function/StAsTextFunction.java | 27 +++-- .../function/StGeogFromGeoJsonFunction.java | 41 +++++++ .../function/StGeomFromGeoJsonFunction.java | 41 +++++++ .../transform/function/StGeometryTypeFunction.java | 14 +-- .../transform/function/StPointFunction.java | 13 +-- .../transform/function/StPolygonFunction.java | 16 ++- .../transform/function/StringBuilderWriter.java | 103 +++++++++++++++++ .../transform/function/StringReader.java | 127 +++++++++++++++++++++ .../function/TransformFunctionFactory.java | 8 ++ .../geospatial/transform/GeoInputOutputTest.java | 83 ++++++++++++++ .../transform/StGeometryTypeFunctionTest.java | 104 +++++++++++++++++ .../pinot/segment/local/utils/GeometryUtils.java | 11 ++ 22 files changed, 696 insertions(+), 150 deletions(-) diff --git a/pinot-common/src/main/java/org/apache/pinot/common/function/TransformFunctionType.java b/pinot-common/src/main/java/org/apache/pinot/common/function/TransformFunctionType.java index c9be27d66b..cf8b018e8a 100644 --- a/pinot-common/src/main/java/org/apache/pinot/common/function/TransformFunctionType.java +++ b/pinot-common/src/main/java/org/apache/pinot/common/function/TransformFunctionType.java @@ -176,8 +176,13 @@ public enum TransformFunctionType { // Geo constructors ST_GEOG_FROM_TEXT("ST_GeogFromText", ReturnTypes.VARBINARY, OperandTypes.CHARACTER), ST_GEOM_FROM_TEXT("ST_GeomFromText", ReturnTypes.VARBINARY, OperandTypes.CHARACTER), + + ST_GEOG_FROM_GEO_JSON("ST_GeogFromGeoJSON", ReturnTypes.VARBINARY, OperandTypes.CHARACTER), + ST_GEOM_FROM_GEO_JSON("ST_GeomFromGeoJSON", ReturnTypes.VARBINARY, OperandTypes.CHARACTER), + ST_GEOG_FROM_WKB("ST_GeogFromWKB", ReturnTypes.VARBINARY, OperandTypes.BINARY), ST_GEOM_FROM_WKB("ST_GeomFromWKB", ReturnTypes.VARBINARY, OperandTypes.BINARY), + ST_POINT("ST_Point", ReturnTypes.VARBINARY, OperandTypes.family(List.of(SqlTypeFamily.NUMERIC, SqlTypeFamily.NUMERIC, SqlTypeFamily.ANY), i -> i == 2)), ST_POLYGON("ST_Polygon", ReturnTypes.VARBINARY, OperandTypes.CHARACTER), @@ -190,6 +195,7 @@ public enum TransformFunctionType { // Geo outputs ST_AS_BINARY("ST_AsBinary", ReturnTypes.VARBINARY, OperandTypes.BINARY), ST_AS_TEXT("ST_AsText", ReturnTypes.VARCHAR, OperandTypes.BINARY), + ST_AS_GEO_JSON("ST_AsGeoJSON", ReturnTypes.VARCHAR, OperandTypes.BINARY), // Geo relationship // TODO: Revisit whether we should return BOOLEAN instead diff --git a/pinot-core/src/main/java/org/apache/pinot/core/geospatial/transform/function/BaseBinaryGeoTransformFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/geospatial/transform/function/BaseBinaryGeoTransformFunction.java index 9990d25286..231fc7d8df 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/geospatial/transform/function/BaseBinaryGeoTransformFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/geospatial/transform/function/BaseBinaryGeoTransformFunction.java @@ -19,7 +19,6 @@ package org.apache.pinot.core.geospatial.transform.function; import com.google.common.base.Preconditions; -import java.util.Arrays; import java.util.List; import java.util.Map; import org.apache.pinot.core.operator.ColumnContext; @@ -27,7 +26,6 @@ import org.apache.pinot.core.operator.blocks.ValueBlock; import org.apache.pinot.core.operator.transform.function.BaseTransformFunction; import org.apache.pinot.core.operator.transform.function.LiteralTransformFunction; import org.apache.pinot.core.operator.transform.function.TransformFunction; -import org.apache.pinot.core.plan.DocIdSetPlanNode; import org.apache.pinot.segment.local.utils.GeometrySerializer; import org.apache.pinot.spi.data.FieldSpec; import org.locationtech.jts.geom.Geometry; @@ -41,8 +39,6 @@ public abstract class BaseBinaryGeoTransformFunction extends BaseTransformFuncti private TransformFunction _secondArgument; private Geometry _firstLiteral; private Geometry _secondLiteral; - private int[] _intResults; - private double[] _doubleResults; @Override public void init(List<TransformFunction> arguments, Map<String, ColumnContext> columnContextMap) { @@ -74,63 +70,68 @@ public abstract class BaseBinaryGeoTransformFunction extends BaseTransformFuncti } protected int[] transformGeometryToIntValuesSV(ValueBlock valueBlock) { - if (_intResults == null) { - _intResults = new int[DocIdSetPlanNode.MAX_DOC_PER_CALL]; - } + int numDocs = valueBlock.getNumDocs(); + initIntValuesSV(numDocs); byte[][] firstValues; byte[][] secondValues; + if (_firstArgument == null && _secondArgument == null) { - _intResults = new int[Math.min(valueBlock.getNumDocs(), DocIdSetPlanNode.MAX_DOC_PER_CALL)]; - Arrays.fill(_intResults, transformGeometryToInt(_firstLiteral, _secondLiteral)); + int value = transformGeometryToInt(_firstLiteral, _secondLiteral); + for (int i = 0; i < numDocs; i++) { + _intValuesSV[i] = value; + } } else if (_firstArgument == null) { secondValues = _secondArgument.transformToBytesValuesSV(valueBlock); - for (int i = 0; i < valueBlock.getNumDocs(); i++) { - _intResults[i] = transformGeometryToInt(_firstLiteral, GeometrySerializer.deserialize(secondValues[i])); + for (int i = 0; i < numDocs; i++) { + _intValuesSV[i] = transformGeometryToInt(_firstLiteral, GeometrySerializer.deserialize(secondValues[i])); } } else if (_secondArgument == null) { firstValues = _firstArgument.transformToBytesValuesSV(valueBlock); - for (int i = 0; i < valueBlock.getNumDocs(); i++) { - _intResults[i] = transformGeometryToInt(GeometrySerializer.deserialize(firstValues[i]), _secondLiteral); + for (int i = 0; i < numDocs; i++) { + _intValuesSV[i] = transformGeometryToInt(GeometrySerializer.deserialize(firstValues[i]), _secondLiteral); } } else { firstValues = _firstArgument.transformToBytesValuesSV(valueBlock); secondValues = _secondArgument.transformToBytesValuesSV(valueBlock); - for (int i = 0; i < valueBlock.getNumDocs(); i++) { - _intResults[i] = transformGeometryToInt(GeometrySerializer.deserialize(firstValues[i]), + for (int i = 0; i < numDocs; i++) { + _intValuesSV[i] = transformGeometryToInt(GeometrySerializer.deserialize(firstValues[i]), GeometrySerializer.deserialize(secondValues[i])); } } - return _intResults; + return _intValuesSV; } protected double[] transformGeometryToDoubleValuesSV(ValueBlock valueBlock) { - if (_doubleResults == null) { - _doubleResults = new double[DocIdSetPlanNode.MAX_DOC_PER_CALL]; - } + int numDocs = valueBlock.getNumDocs(); + initDoubleValuesSV(numDocs); + byte[][] firstValues; byte[][] secondValues; + if (_firstArgument == null && _secondArgument == null) { - _doubleResults = new double[Math.min(valueBlock.getNumDocs(), DocIdSetPlanNode.MAX_DOC_PER_CALL)]; - Arrays.fill(_doubleResults, transformGeometryToDouble(_firstLiteral, _secondLiteral)); + double value = transformGeometryToDouble(_firstLiteral, _secondLiteral); + for (int i = 0; i < numDocs; i++) { + _doubleValuesSV[i] = value; + } } else if (_firstArgument == null) { secondValues = _secondArgument.transformToBytesValuesSV(valueBlock); - for (int i = 0; i < valueBlock.getNumDocs(); i++) { - _doubleResults[i] = transformGeometryToDouble(_firstLiteral, GeometrySerializer.deserialize(secondValues[i])); + for (int i = 0; i < numDocs; i++) { + _doubleValuesSV[i] = transformGeometryToDouble(_firstLiteral, GeometrySerializer.deserialize(secondValues[i])); } } else if (_secondArgument == null) { firstValues = _firstArgument.transformToBytesValuesSV(valueBlock); - for (int i = 0; i < valueBlock.getNumDocs(); i++) { - _doubleResults[i] = transformGeometryToDouble(GeometrySerializer.deserialize(firstValues[i]), _secondLiteral); + for (int i = 0; i < numDocs; i++) { + _doubleValuesSV[i] = transformGeometryToDouble(GeometrySerializer.deserialize(firstValues[i]), _secondLiteral); } } else { firstValues = _firstArgument.transformToBytesValuesSV(valueBlock); secondValues = _secondArgument.transformToBytesValuesSV(valueBlock); - for (int i = 0; i < valueBlock.getNumDocs(); i++) { - _doubleResults[i] = transformGeometryToDouble(GeometrySerializer.deserialize(firstValues[i]), + for (int i = 0; i < numDocs; i++) { + _doubleValuesSV[i] = transformGeometryToDouble(GeometrySerializer.deserialize(firstValues[i]), GeometrySerializer.deserialize(secondValues[i])); } } - return _doubleResults; + return _doubleValuesSV; } public int transformGeometryToInt(Geometry firstGeometry, Geometry secondGeometry) { diff --git a/pinot-core/src/main/java/org/apache/pinot/core/geospatial/transform/function/ConstructFromTextFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/geospatial/transform/function/ConstructFromGeoJsonFunction.java similarity index 72% copy from pinot-core/src/main/java/org/apache/pinot/core/geospatial/transform/function/ConstructFromTextFunction.java copy to pinot-core/src/main/java/org/apache/pinot/core/geospatial/transform/function/ConstructFromGeoJsonFunction.java index ae69b5ca84..4d634299c7 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/geospatial/transform/function/ConstructFromTextFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/geospatial/transform/function/ConstructFromGeoJsonFunction.java @@ -21,43 +21,45 @@ package org.apache.pinot.core.geospatial.transform.function; import com.google.common.base.Preconditions; import java.util.List; import java.util.Map; -import org.apache.pinot.common.Utils; import org.apache.pinot.core.operator.ColumnContext; import org.apache.pinot.core.operator.blocks.ValueBlock; import org.apache.pinot.core.operator.transform.TransformResultMetadata; import org.apache.pinot.core.operator.transform.function.BaseTransformFunction; import org.apache.pinot.core.operator.transform.function.TransformFunction; -import org.apache.pinot.core.plan.DocIdSetPlanNode; import org.apache.pinot.segment.local.utils.GeometrySerializer; import org.apache.pinot.spi.data.FieldSpec; import org.locationtech.jts.geom.Geometry; import org.locationtech.jts.io.ParseException; -import org.locationtech.jts.io.WKTReader; +import org.locationtech.jts.io.geojson.GeoJsonReader; /** - * An abstract class for implementing the geo constructor functions from text. + * An abstract class for implementing the geo constructor functions from GEO JSON. */ -abstract class ConstructFromTextFunction extends BaseTransformFunction { +abstract class ConstructFromGeoJsonFunction extends BaseTransformFunction { + protected TransformFunction _transformFunction; - protected byte[][] _results; - protected WKTReader _reader; + protected GeoJsonReader _reader; @Override public void init(List<TransformFunction> arguments, Map<String, ColumnContext> columnContextMap) { super.init(arguments, columnContextMap); - Preconditions.checkArgument(arguments.size() == 1, "Exactly 1 argument is required for transform function: %s", - getName()); + + Preconditions.checkArgument(arguments.size() == 1, + "Exactly 1 argument is required for transform function: " + getName()); + TransformFunction transformFunction = arguments.get(0); + Preconditions.checkArgument(transformFunction.getResultMetadata().isSingleValue(), - "The argument must be single-valued for transform function: %s", getName()); + "The argument must be single-valued for transform function: " + getName()); Preconditions.checkArgument(transformFunction.getResultMetadata().getDataType() == FieldSpec.DataType.STRING, "The argument must be of string type"); + _transformFunction = transformFunction; - _reader = getWKTReader(); + _reader = getGeoJsonReader(); } - abstract protected WKTReader getWKTReader(); + abstract protected GeoJsonReader getGeoJsonReader(); @Override public TransformResultMetadata getResultMetadata() { @@ -66,20 +68,20 @@ abstract class ConstructFromTextFunction extends BaseTransformFunction { @Override public byte[][] transformToBytesValuesSV(ValueBlock valueBlock) { - if (_results == null) { - _results = new byte[DocIdSetPlanNode.MAX_DOC_PER_CALL][]; - } + int numDocs = valueBlock.getNumDocs(); + initBytesValuesSV(numDocs); String[] argumentValues = _transformFunction.transformToStringValuesSV(valueBlock); - int length = valueBlock.getNumDocs(); - for (int i = 0; i < length; i++) { + // use single reader instead of allocating separate instance per row + StringReader reader = new StringReader(); + for (int i = 0; i < numDocs; i++) { try { - Geometry geometry = _reader.read(argumentValues[i]); - _results[i] = GeometrySerializer.serialize(geometry); + reader.setString(argumentValues[i]); + Geometry geometry = _reader.read(reader); + _bytesValuesSV[i] = GeometrySerializer.serialize(geometry); } catch (ParseException e) { - Utils.rethrowException( - new RuntimeException(String.format("Failed to parse geometry from string: %s", argumentValues[i]))); + throw new RuntimeException(String.format("Failed to parse geometry from string: %s", argumentValues[i])); } } - return _results; + return _bytesValuesSV; } } diff --git a/pinot-core/src/main/java/org/apache/pinot/core/geospatial/transform/function/ConstructFromTextFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/geospatial/transform/function/ConstructFromTextFunction.java index ae69b5ca84..9a35a3ecc5 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/geospatial/transform/function/ConstructFromTextFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/geospatial/transform/function/ConstructFromTextFunction.java @@ -21,13 +21,11 @@ package org.apache.pinot.core.geospatial.transform.function; import com.google.common.base.Preconditions; import java.util.List; import java.util.Map; -import org.apache.pinot.common.Utils; import org.apache.pinot.core.operator.ColumnContext; import org.apache.pinot.core.operator.blocks.ValueBlock; import org.apache.pinot.core.operator.transform.TransformResultMetadata; import org.apache.pinot.core.operator.transform.function.BaseTransformFunction; import org.apache.pinot.core.operator.transform.function.TransformFunction; -import org.apache.pinot.core.plan.DocIdSetPlanNode; import org.apache.pinot.segment.local.utils.GeometrySerializer; import org.apache.pinot.spi.data.FieldSpec; import org.locationtech.jts.geom.Geometry; @@ -40,7 +38,6 @@ import org.locationtech.jts.io.WKTReader; */ abstract class ConstructFromTextFunction extends BaseTransformFunction { protected TransformFunction _transformFunction; - protected byte[][] _results; protected WKTReader _reader; @Override @@ -66,20 +63,17 @@ abstract class ConstructFromTextFunction extends BaseTransformFunction { @Override public byte[][] transformToBytesValuesSV(ValueBlock valueBlock) { - if (_results == null) { - _results = new byte[DocIdSetPlanNode.MAX_DOC_PER_CALL][]; - } + int numDocs = valueBlock.getNumDocs(); + initBytesValuesSV(numDocs); String[] argumentValues = _transformFunction.transformToStringValuesSV(valueBlock); - int length = valueBlock.getNumDocs(); - for (int i = 0; i < length; i++) { + for (int i = 0; i < numDocs; i++) { try { Geometry geometry = _reader.read(argumentValues[i]); - _results[i] = GeometrySerializer.serialize(geometry); + _bytesValuesSV[i] = GeometrySerializer.serialize(geometry); } catch (ParseException e) { - Utils.rethrowException( - new RuntimeException(String.format("Failed to parse geometry from string: %s", argumentValues[i]))); + throw new RuntimeException(String.format("Failed to parse geometry from string: %s", argumentValues[i])); } } - return _results; + return _bytesValuesSV; } } diff --git a/pinot-core/src/main/java/org/apache/pinot/core/geospatial/transform/function/ConstructFromWKBFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/geospatial/transform/function/ConstructFromWKBFunction.java index 007d912e04..259f4be6c1 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/geospatial/transform/function/ConstructFromWKBFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/geospatial/transform/function/ConstructFromWKBFunction.java @@ -26,7 +26,6 @@ import org.apache.pinot.core.operator.blocks.ValueBlock; import org.apache.pinot.core.operator.transform.TransformResultMetadata; import org.apache.pinot.core.operator.transform.function.BaseTransformFunction; import org.apache.pinot.core.operator.transform.function.TransformFunction; -import org.apache.pinot.core.plan.DocIdSetPlanNode; import org.apache.pinot.segment.local.utils.GeometrySerializer; import org.apache.pinot.spi.data.FieldSpec; import org.apache.pinot.spi.utils.BytesUtils; @@ -40,7 +39,6 @@ import org.locationtech.jts.io.WKBReader; */ abstract class ConstructFromWKBFunction extends BaseTransformFunction { private TransformFunction _transformFunction; - private byte[][] _results; private WKBReader _reader; @Override @@ -66,19 +64,18 @@ abstract class ConstructFromWKBFunction extends BaseTransformFunction { @Override public byte[][] transformToBytesValuesSV(ValueBlock valueBlock) { - if (_results == null) { - _results = new byte[DocIdSetPlanNode.MAX_DOC_PER_CALL][]; - } + int numDocs = valueBlock.getNumDocs(); + initBytesValuesSV(numDocs); byte[][] argumentValues = _transformFunction.transformToBytesValuesSV(valueBlock); - for (int i = 0; i < valueBlock.getNumDocs(); i++) { + for (int i = 0; i < numDocs; i++) { try { Geometry geometry = _reader.read(argumentValues[i]); - _results[i] = GeometrySerializer.serialize(geometry); + _bytesValuesSV[i] = GeometrySerializer.serialize(geometry); } catch (ParseException e) { throw new RuntimeException( String.format("Failed to parse geometry from bytes %s", BytesUtils.toHexString(argumentValues[i]))); } } - return _results; + return _bytesValuesSV; } } diff --git a/pinot-core/src/main/java/org/apache/pinot/core/geospatial/transform/function/GeoToH3Function.java b/pinot-core/src/main/java/org/apache/pinot/core/geospatial/transform/function/GeoToH3Function.java index 237a5619e5..d336530a77 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/geospatial/transform/function/GeoToH3Function.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/geospatial/transform/function/GeoToH3Function.java @@ -27,7 +27,6 @@ import org.apache.pinot.core.operator.transform.TransformResultMetadata; import org.apache.pinot.core.operator.transform.function.BaseTransformFunction; import org.apache.pinot.core.operator.transform.function.LiteralTransformFunction; import org.apache.pinot.core.operator.transform.function.TransformFunction; -import org.apache.pinot.core.plan.DocIdSetPlanNode; import org.apache.pinot.segment.local.utils.GeometrySerializer; import org.apache.pinot.spi.data.FieldSpec; import org.locationtech.jts.geom.Geometry; @@ -43,7 +42,6 @@ public class GeoToH3Function extends BaseTransformFunction { private TransformFunction _firstArgument; private TransformFunction _secondArgument; private TransformFunction _thirdArgument; - private long[] _results; @Override public String getName() { @@ -98,26 +96,26 @@ public class GeoToH3Function extends BaseTransformFunction { @Override public long[] transformToLongValuesSV(ValueBlock valueBlock) { - if (_results == null) { - _results = new long[DocIdSetPlanNode.MAX_DOC_PER_CALL]; - } + int numDocs = valueBlock.getNumDocs(); + initLongValuesSV(numDocs); if (_thirdArgument == null) { byte[][] geoValues = _firstArgument.transformToBytesValuesSV(valueBlock); int[] resValues = _secondArgument.transformToIntValuesSV(valueBlock); - for (int i = 0; i < valueBlock.getNumDocs(); i++) { + for (int i = 0; i < numDocs; i++) { Geometry geometry = GeometrySerializer.deserialize(geoValues[i]); - _results[i] = ScalarFunctions.geoToH3(geometry.getCoordinate().x, geometry.getCoordinate().y, resValues[i]); + _longValuesSV[i] = + ScalarFunctions.geoToH3(geometry.getCoordinate().x, geometry.getCoordinate().y, resValues[i]); } } else { double[] lonValues = _firstArgument.transformToDoubleValuesSV(valueBlock); double[] latValues = _secondArgument.transformToDoubleValuesSV(valueBlock); int[] resValues = _thirdArgument.transformToIntValuesSV(valueBlock); - for (int i = 0; i < valueBlock.getNumDocs(); i++) { - _results[i] = ScalarFunctions.geoToH3(lonValues[i], latValues[i], resValues[i]); + for (int i = 0; i < numDocs; i++) { + _longValuesSV[i] = ScalarFunctions.geoToH3(lonValues[i], latValues[i], resValues[i]); } } - return _results; + return _longValuesSV; } } diff --git a/pinot-core/src/main/java/org/apache/pinot/core/geospatial/transform/function/ScalarFunctions.java b/pinot-core/src/main/java/org/apache/pinot/core/geospatial/transform/function/ScalarFunctions.java index fc6d256649..da52f0bf18 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/geospatial/transform/function/ScalarFunctions.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/geospatial/transform/function/ScalarFunctions.java @@ -83,6 +83,24 @@ public class ScalarFunctions { return GeometrySerializer.serialize(GeometryUtils.GEOGRAPHY_WKT_READER.read(wkt)); } + /** + * Reads a geometry object from the GeoJSON format. + */ + @ScalarFunction + public static byte[] stGeomFromGeoJson(String geoJson) + throws ParseException { + return GeometrySerializer.serialize(GeometryUtils.GEOMETRY_GEO_JSON_READER.read(geoJson)); + } + + /** + * Reads a geography object from the GeoJSon format. + */ + @ScalarFunction + public static byte[] stGeogFromGeoJson(String geoJson) + throws ParseException { + return GeometrySerializer.serialize(GeometryUtils.GEOGRAPHY_GEO_JSON_READER.read(geoJson)); + } + /** * Reads a geometry object from the WKB format. */ @@ -112,6 +130,17 @@ public class ScalarFunctions { return GeometryUtils.WKT_WRITER.write(GeometrySerializer.deserialize(bytes)); } + /** + * Saves the geometry object in GeoJSON format. + * + * @param bytes the serialized geometry object + * @return the geometry in GeoJSON + */ + @ScalarFunction + public static String stAsGeoJson(byte[] bytes) { + return GeometryUtils.GEO_JSON_WRITER.write(GeometrySerializer.deserialize(bytes)); + } + /** * Saves the geometry object as WKB format. * diff --git a/pinot-core/src/main/java/org/apache/pinot/core/geospatial/transform/function/StAreaFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/geospatial/transform/function/StAreaFunction.java index ef51a30dd5..653b3ee9f4 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/geospatial/transform/function/StAreaFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/geospatial/transform/function/StAreaFunction.java @@ -27,7 +27,6 @@ import org.apache.pinot.core.operator.transform.TransformResultMetadata; import org.apache.pinot.core.operator.transform.function.BaseTransformFunction; import org.apache.pinot.core.operator.transform.function.LiteralTransformFunction; import org.apache.pinot.core.operator.transform.function.TransformFunction; -import org.apache.pinot.core.plan.DocIdSetPlanNode; import org.apache.pinot.segment.local.utils.GeometrySerializer; import org.apache.pinot.segment.local.utils.GeometryUtils; import org.apache.pinot.spi.data.FieldSpec; @@ -49,7 +48,6 @@ import static java.lang.Math.toRadians; public class StAreaFunction extends BaseTransformFunction { private TransformFunction _transformFunction; public static final String FUNCTION_NAME = "ST_Area"; - private double[] _results; @Override public String getName() { @@ -78,17 +76,15 @@ public class StAreaFunction extends BaseTransformFunction { @Override public double[] transformToDoubleValuesSV(ValueBlock valueBlock) { - if (_results == null) { - _results = new double[DocIdSetPlanNode.MAX_DOC_PER_CALL]; - } - - byte[][] values = _transformFunction.transformToBytesValuesSV(valueBlock); int numDocs = valueBlock.getNumDocs(); + initDoubleValuesSV(numDocs); + byte[][] values = _transformFunction.transformToBytesValuesSV(valueBlock); + for (int i = 0; i < numDocs; i++) { Geometry geometry = GeometrySerializer.deserialize(values[i]); - _results[i] = GeometryUtils.isGeography(geometry) ? calculateGeographyArea(geometry) : geometry.getArea(); + _doubleValuesSV[i] = GeometryUtils.isGeography(geometry) ? calculateGeographyArea(geometry) : geometry.getArea(); } - return _results; + return _doubleValuesSV; } private double calculateGeographyArea(Geometry geometry) { diff --git a/pinot-core/src/main/java/org/apache/pinot/core/geospatial/transform/function/StAsBinaryFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/geospatial/transform/function/StAsBinaryFunction.java index 5b7d8da836..a64294ab83 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/geospatial/transform/function/StAsBinaryFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/geospatial/transform/function/StAsBinaryFunction.java @@ -26,7 +26,6 @@ import org.apache.pinot.core.operator.blocks.ValueBlock; import org.apache.pinot.core.operator.transform.TransformResultMetadata; import org.apache.pinot.core.operator.transform.function.BaseTransformFunction; import org.apache.pinot.core.operator.transform.function.TransformFunction; -import org.apache.pinot.core.plan.DocIdSetPlanNode; import org.apache.pinot.segment.local.utils.GeometrySerializer; import org.apache.pinot.segment.local.utils.GeometryUtils; import org.apache.pinot.spi.data.FieldSpec; @@ -40,7 +39,6 @@ public class StAsBinaryFunction extends BaseTransformFunction { public static final String FUNCTION_NAME = "ST_AsBinary"; private TransformFunction _transformFunction; - private byte[][] _results; @Override public String getName() { @@ -67,15 +65,14 @@ public class StAsBinaryFunction extends BaseTransformFunction { @Override public byte[][] transformToBytesValuesSV(ValueBlock valueBlock) { - if (_results == null) { - _results = new byte[DocIdSetPlanNode.MAX_DOC_PER_CALL][]; - } + int numDocs = valueBlock.getNumDocs(); + initBytesValuesSV(numDocs); byte[][] values = _transformFunction.transformToBytesValuesSV(valueBlock); Geometry geometry; - for (int i = 0; i < valueBlock.getNumDocs(); i++) { + for (int i = 0; i < numDocs; i++) { geometry = GeometrySerializer.deserialize(values[i]); - _results[i] = GeometryUtils.WKB_WRITER.write(geometry); + _bytesValuesSV[i] = GeometryUtils.WKB_WRITER.write(geometry); } - return _results; + return _bytesValuesSV; } } diff --git a/pinot-core/src/main/java/org/apache/pinot/core/geospatial/transform/function/StAsTextFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/geospatial/transform/function/StAsGeoJsonFunction.java similarity index 71% copy from pinot-core/src/main/java/org/apache/pinot/core/geospatial/transform/function/StAsTextFunction.java copy to pinot-core/src/main/java/org/apache/pinot/core/geospatial/transform/function/StAsGeoJsonFunction.java index 8d620deb4f..49ccde9d4d 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/geospatial/transform/function/StAsTextFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/geospatial/transform/function/StAsGeoJsonFunction.java @@ -19,6 +19,7 @@ package org.apache.pinot.core.geospatial.transform.function; import com.google.common.base.Preconditions; +import java.io.IOException; import java.util.List; import java.util.Map; import org.apache.pinot.core.operator.ColumnContext; @@ -26,7 +27,6 @@ import org.apache.pinot.core.operator.blocks.ValueBlock; import org.apache.pinot.core.operator.transform.TransformResultMetadata; import org.apache.pinot.core.operator.transform.function.BaseTransformFunction; import org.apache.pinot.core.operator.transform.function.TransformFunction; -import org.apache.pinot.core.plan.DocIdSetPlanNode; import org.apache.pinot.segment.local.utils.GeometrySerializer; import org.apache.pinot.segment.local.utils.GeometryUtils; import org.apache.pinot.spi.data.FieldSpec; @@ -34,15 +34,14 @@ import org.locationtech.jts.geom.Geometry; /** - * Returns the text representation of the geometry object. + * Returns the GEOJson representation of the geometry object . */ -public class StAsTextFunction extends BaseTransformFunction { - public static final String FUNCTION_NAME = "ST_AsText"; +public class StAsGeoJsonFunction extends BaseTransformFunction { + + public static final String FUNCTION_NAME = "ST_AsGeoJSON"; private TransformFunction _transformFunction; - private String[] _results; - @Override public String getName() { return FUNCTION_NAME; } @@ -50,13 +49,15 @@ public class StAsTextFunction extends BaseTransformFunction { @Override public void init(List<TransformFunction> arguments, Map<String, ColumnContext> columnContextMap) { super.init(arguments, columnContextMap); - Preconditions.checkArgument(arguments.size() == 1, "Exactly 1 argument is required for transform function: %s", - getName()); + Preconditions.checkArgument(arguments.size() == 1, + "Exactly 1 argument is required for transform function: " + FUNCTION_NAME); + TransformFunction transformFunction = arguments.get(0); Preconditions.checkArgument(transformFunction.getResultMetadata().isSingleValue(), - "Argument must be single-valued for transform function: %s", getName()); + "Argument must be single-valued for transform function: " + FUNCTION_NAME); Preconditions.checkArgument(transformFunction.getResultMetadata().getDataType() == FieldSpec.DataType.BYTES, "The argument must be of bytes type"); + _transformFunction = transformFunction; } @@ -65,17 +66,24 @@ public class StAsTextFunction extends BaseTransformFunction { return STRING_SV_NO_DICTIONARY_METADATA; } - @Override public String[] transformToStringValuesSV(ValueBlock valueBlock) { - if (_results == null) { - _results = new String[DocIdSetPlanNode.MAX_DOC_PER_CALL]; - } + int numDocs = valueBlock.getNumDocs(); + initStringValuesSV(numDocs); byte[][] values = _transformFunction.transformToBytesValuesSV(valueBlock); - Geometry geometry; - for (int i = 0; i < valueBlock.getNumDocs(); i++) { - geometry = GeometrySerializer.deserialize(values[i]); - _results[i] = GeometryUtils.WKT_WRITER.write(geometry); + // use single buffer instead of allocating separate StringBuffer per row + StringBuilderWriter buffer = new StringBuilderWriter(); + + try { + for (int i = 0; i < numDocs; i++) { + Geometry geometry = GeometrySerializer.deserialize(values[i]); + GeometryUtils.GEO_JSON_WRITER.write(geometry, buffer); + _stringValuesSV[i] = buffer.getString(); + buffer.clear(); + } + } catch (IOException ioe) { + // should never happen + throw new RuntimeException(ioe); } - return _results; + return _stringValuesSV; } } diff --git a/pinot-core/src/main/java/org/apache/pinot/core/geospatial/transform/function/StAsTextFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/geospatial/transform/function/StAsTextFunction.java index 8d620deb4f..c6d943c1c4 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/geospatial/transform/function/StAsTextFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/geospatial/transform/function/StAsTextFunction.java @@ -19,6 +19,7 @@ package org.apache.pinot.core.geospatial.transform.function; import com.google.common.base.Preconditions; +import java.io.IOException; import java.util.List; import java.util.Map; import org.apache.pinot.core.operator.ColumnContext; @@ -26,7 +27,6 @@ import org.apache.pinot.core.operator.blocks.ValueBlock; import org.apache.pinot.core.operator.transform.TransformResultMetadata; import org.apache.pinot.core.operator.transform.function.BaseTransformFunction; import org.apache.pinot.core.operator.transform.function.TransformFunction; -import org.apache.pinot.core.plan.DocIdSetPlanNode; import org.apache.pinot.segment.local.utils.GeometrySerializer; import org.apache.pinot.segment.local.utils.GeometryUtils; import org.apache.pinot.spi.data.FieldSpec; @@ -40,7 +40,6 @@ public class StAsTextFunction extends BaseTransformFunction { public static final String FUNCTION_NAME = "ST_AsText"; private TransformFunction _transformFunction; - private String[] _results; @Override public String getName() { @@ -67,15 +66,23 @@ public class StAsTextFunction extends BaseTransformFunction { @Override public String[] transformToStringValuesSV(ValueBlock valueBlock) { - if (_results == null) { - _results = new String[DocIdSetPlanNode.MAX_DOC_PER_CALL]; - } + int numDocs = valueBlock.getNumDocs(); + initStringValuesSV(numDocs); byte[][] values = _transformFunction.transformToBytesValuesSV(valueBlock); - Geometry geometry; - for (int i = 0; i < valueBlock.getNumDocs(); i++) { - geometry = GeometrySerializer.deserialize(values[i]); - _results[i] = GeometryUtils.WKT_WRITER.write(geometry); + // use single buffer instead of allocating separate StringBuffer per row + StringBuilderWriter buffer = new StringBuilderWriter(); + + try { + for (int i = 0; i < numDocs; i++) { + Geometry geometry = GeometrySerializer.deserialize(values[i]); + GeometryUtils.WKT_WRITER.write(geometry, buffer); + _stringValuesSV[i] = buffer.getString(); + buffer.clear(); + } + } catch (IOException ioe) { + // should never happen + throw new RuntimeException(ioe); } - return _results; + return _stringValuesSV; } } diff --git a/pinot-core/src/main/java/org/apache/pinot/core/geospatial/transform/function/StGeogFromGeoJsonFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/geospatial/transform/function/StGeogFromGeoJsonFunction.java new file mode 100644 index 0000000000..43c334f320 --- /dev/null +++ b/pinot-core/src/main/java/org/apache/pinot/core/geospatial/transform/function/StGeogFromGeoJsonFunction.java @@ -0,0 +1,41 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 org.apache.pinot.core.geospatial.transform.function; + +import org.apache.pinot.segment.local.utils.GeometryUtils; +import org.locationtech.jts.io.geojson.GeoJsonReader; + + +/** + * Constructor function for geography object from GEO JSON. + */ +public class StGeogFromGeoJsonFunction extends ConstructFromGeoJsonFunction { + + public static final String FUNCTION_NAME = "ST_GeogFromGeoJSON"; + + @Override + protected GeoJsonReader getGeoJsonReader() { + return GeometryUtils.GEOGRAPHY_GEO_JSON_READER; + } + + @Override + public String getName() { + return FUNCTION_NAME; + } +} diff --git a/pinot-core/src/main/java/org/apache/pinot/core/geospatial/transform/function/StGeomFromGeoJsonFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/geospatial/transform/function/StGeomFromGeoJsonFunction.java new file mode 100644 index 0000000000..5fadd7dd4d --- /dev/null +++ b/pinot-core/src/main/java/org/apache/pinot/core/geospatial/transform/function/StGeomFromGeoJsonFunction.java @@ -0,0 +1,41 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 org.apache.pinot.core.geospatial.transform.function; + +import org.apache.pinot.segment.local.utils.GeometryUtils; +import org.locationtech.jts.io.geojson.GeoJsonReader; + + +/** + * Constructor function for geometry object from GEO JSON. + */ +public class StGeomFromGeoJsonFunction extends ConstructFromGeoJsonFunction { + + public static final String FUNCTION_NAME = "ST_GeomFromGeoJSON"; + + @Override + protected GeoJsonReader getGeoJsonReader() { + return GeometryUtils.GEOMETRY_GEO_JSON_READER; + } + + @Override + public String getName() { + return FUNCTION_NAME; + } +} diff --git a/pinot-core/src/main/java/org/apache/pinot/core/geospatial/transform/function/StGeometryTypeFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/geospatial/transform/function/StGeometryTypeFunction.java index c5273dcd3a..6d1ca5183e 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/geospatial/transform/function/StGeometryTypeFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/geospatial/transform/function/StGeometryTypeFunction.java @@ -26,7 +26,6 @@ import org.apache.pinot.core.operator.blocks.ValueBlock; import org.apache.pinot.core.operator.transform.TransformResultMetadata; import org.apache.pinot.core.operator.transform.function.BaseTransformFunction; import org.apache.pinot.core.operator.transform.function.TransformFunction; -import org.apache.pinot.core.plan.DocIdSetPlanNode; import org.apache.pinot.segment.local.utils.GeometrySerializer; import org.apache.pinot.spi.data.FieldSpec; import org.locationtech.jts.geom.Geometry; @@ -37,7 +36,6 @@ import org.locationtech.jts.geom.Geometry; */ public class StGeometryTypeFunction extends BaseTransformFunction { private TransformFunction _transformFunction; - private String[] _results; public static final String FUNCTION_NAME = "ST_GEOMETRY_TYPE"; @Override @@ -65,15 +63,15 @@ public class StGeometryTypeFunction extends BaseTransformFunction { @Override public String[] transformToStringValuesSV(ValueBlock valueBlock) { - if (_results == null) { - _results = new String[DocIdSetPlanNode.MAX_DOC_PER_CALL]; - } + int numDocs = valueBlock.getNumDocs(); + initStringValuesSV(numDocs); byte[][] values = _transformFunction.transformToBytesValuesSV(valueBlock); Geometry geometry; - for (int i = 0; i < valueBlock.getNumDocs(); i++) { + + for (int i = 0; i < numDocs; i++) { geometry = GeometrySerializer.deserialize(values[i]); - _results[i] = geometry.getGeometryType(); + _stringValuesSV[i] = geometry.getGeometryType(); } - return _results; + return _stringValuesSV; } } diff --git a/pinot-core/src/main/java/org/apache/pinot/core/geospatial/transform/function/StPointFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/geospatial/transform/function/StPointFunction.java index 0568c1b8da..0c755a90cd 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/geospatial/transform/function/StPointFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/geospatial/transform/function/StPointFunction.java @@ -27,7 +27,6 @@ import org.apache.pinot.core.operator.transform.TransformResultMetadata; import org.apache.pinot.core.operator.transform.function.BaseTransformFunction; import org.apache.pinot.core.operator.transform.function.LiteralTransformFunction; import org.apache.pinot.core.operator.transform.function.TransformFunction; -import org.apache.pinot.core.plan.DocIdSetPlanNode; import org.apache.pinot.segment.local.utils.GeometrySerializer; import org.apache.pinot.segment.local.utils.GeometryUtils; import org.locationtech.jts.geom.Coordinate; @@ -41,7 +40,6 @@ public class StPointFunction extends BaseTransformFunction { public static final String FUNCTION_NAME = "ST_Point"; private TransformFunction _firstArgument; private TransformFunction _secondArgument; - private byte[][] _results; private boolean _isGeography; @Override @@ -77,18 +75,17 @@ public class StPointFunction extends BaseTransformFunction { @Override public byte[][] transformToBytesValuesSV(ValueBlock valueBlock) { - if (_results == null) { - _results = new byte[DocIdSetPlanNode.MAX_DOC_PER_CALL][]; - } + int numDocs = valueBlock.getNumDocs(); + initBytesValuesSV(numDocs); double[] firstValues = _firstArgument.transformToDoubleValuesSV(valueBlock); double[] secondValues = _secondArgument.transformToDoubleValuesSV(valueBlock); - for (int i = 0; i < valueBlock.getNumDocs(); i++) { + for (int i = 0; i < numDocs; i++) { Point point = GeometryUtils.GEOMETRY_FACTORY.createPoint(new Coordinate(firstValues[i], secondValues[i])); if (_isGeography) { GeometryUtils.setGeography(point); } - _results[i] = GeometrySerializer.serialize(point); + _bytesValuesSV[i] = GeometrySerializer.serialize(point); } - return _results; + return _bytesValuesSV; } } diff --git a/pinot-core/src/main/java/org/apache/pinot/core/geospatial/transform/function/StPolygonFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/geospatial/transform/function/StPolygonFunction.java index 7680eafa53..484fb10936 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/geospatial/transform/function/StPolygonFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/geospatial/transform/function/StPolygonFunction.java @@ -20,7 +20,6 @@ package org.apache.pinot.core.geospatial.transform.function; import com.google.common.base.Preconditions; import org.apache.pinot.core.operator.blocks.ValueBlock; -import org.apache.pinot.core.plan.DocIdSetPlanNode; import org.apache.pinot.segment.local.utils.GeometrySerializer; import org.apache.pinot.segment.local.utils.GeometryUtils; import org.locationtech.jts.geom.Geometry; @@ -47,20 +46,19 @@ public class StPolygonFunction extends ConstructFromTextFunction { @Override public byte[][] transformToBytesValuesSV(ValueBlock valueBlock) { - if (_results == null) { - _results = new byte[DocIdSetPlanNode.MAX_DOC_PER_CALL][]; - } + int numDocs = valueBlock.getNumDocs(); + initBytesValuesSV(numDocs); String[] argumentValues = _transformFunction.transformToStringValuesSV(valueBlock); - int length = valueBlock.getNumDocs(); - for (int i = 0; i < length; i++) { + + for (int i = 0; i < numDocs; i++) { try { Geometry geometry = _reader.read(argumentValues[i]); Preconditions.checkArgument(geometry instanceof Polygon, "The geometry object must be polygon"); - _results[i] = GeometrySerializer.serialize(geometry); + _bytesValuesSV[i] = GeometrySerializer.serialize(geometry); } catch (ParseException e) { - new RuntimeException(String.format("Failed to parse geometry from string: %s", argumentValues[i])); + throw new RuntimeException(String.format("Failed to parse geometry from string: %s", argumentValues[i])); } } - return _results; + return _bytesValuesSV; } } diff --git a/pinot-core/src/main/java/org/apache/pinot/core/geospatial/transform/function/StringBuilderWriter.java b/pinot-core/src/main/java/org/apache/pinot/core/geospatial/transform/function/StringBuilderWriter.java new file mode 100644 index 0000000000..77a377c2fd --- /dev/null +++ b/pinot-core/src/main/java/org/apache/pinot/core/geospatial/transform/function/StringBuilderWriter.java @@ -0,0 +1,103 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 org.apache.pinot.core.geospatial.transform.function; + +import java.io.IOException; +import java.io.Writer; + + +/** + * Simple non-synchronized writer implementation using a StringBuilder + * and skipping char buffer to reduce overhead + */ +public class StringBuilderWriter extends Writer { + + private final StringBuilder _buffer = new StringBuilder(); + + @Override + public void write(char[] cbuf, int off, int len) + throws IOException { + _buffer.append(cbuf, off, len); + } + + @Override + public void write(int c) + throws IOException { + _buffer.append((char) c); + } + + @Override + public void write(String str) + throws IOException { + _buffer.append(str); + } + + @Override + public void write(char[] cbuf) + throws IOException { + _buffer.append(cbuf); + } + + @Override + public void write(String str, int off, int len) + throws IOException { + _buffer.append(str, off, off + len); + } + + @Override + public Writer append(char c) + throws IOException { + _buffer.append(c); + return this; + } + + @Override + public Writer append(CharSequence csq) + throws IOException { + _buffer.append(csq); + return this; + } + + @Override + public Writer append(CharSequence csq, int start, int end) + throws IOException { + _buffer.append(csq, start, end); + return this; + } + + @Override + public void flush() + throws IOException { + // nothing to do + } + + @Override + public void close() + throws IOException { + // nothing to do + } + + public void clear() { + _buffer.setLength(0); + } + + public String getString() { + return _buffer.toString(); + } +} diff --git a/pinot-core/src/main/java/org/apache/pinot/core/geospatial/transform/function/StringReader.java b/pinot-core/src/main/java/org/apache/pinot/core/geospatial/transform/function/StringReader.java new file mode 100644 index 0000000000..6da2ed2810 --- /dev/null +++ b/pinot-core/src/main/java/org/apache/pinot/core/geospatial/transform/function/StringReader.java @@ -0,0 +1,127 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 org.apache.pinot.core.geospatial.transform.function; + +import java.io.IOException; +import java.io.Reader; +import java.io.Writer; +import java.nio.CharBuffer; + + +/** + * Re-usable string reader. + * Note: at the moment it implements only methods required by GeoJsonReader, that is: + * - read(char[], int, int) + * - read() + * - close() + */ +public class StringReader extends Reader { + + private int _length; + private String _str; + private int _next = 0; + + public StringReader() { + } + + public void setString(String str) { + _str = str; + _length = str.length(); + _next = 0; + } + + @Override + public int read(char[] cbuf, int off, int len) + throws IOException { + if (len == 0) { + return 0; + } + if (_next >= _length) { + return -1; + } + int n = Math.min(_length - _next, len); + _str.getChars(_next, _next + n, cbuf, off); + _next += n; + return n; + } + + @Override + public int read() + throws IOException { + if (_next >= _length) { + return -1; + } + return _str.charAt(_next++); + } + + @Override + public int read(char[] cbuf) + throws IOException { + return super.read(cbuf); + } + + @Override + public int read(CharBuffer target) + throws IOException { + return super.read(target); + } + + @Override + public void mark(int readAheadLimit) + throws IOException { + super.mark(readAheadLimit); + } + + @Override + public boolean markSupported() { + return super.markSupported(); + } + + @Override + public void close() + throws IOException { + _length = 0; + _str = null; + _next = 0; + } + + @Override + public void reset() + throws IOException { + super.reset(); + } + + @Override + public boolean ready() + throws IOException { + return super.ready(); + } + + @Override + public long skip(long n) + throws IOException { + return super.skip(n); + } + + @Override + public long transferTo(Writer out) + throws IOException { + return super.transferTo(out); + } +} diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/TransformFunctionFactory.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/TransformFunctionFactory.java index 97f0ca775d..3709d08f2b 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/TransformFunctionFactory.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/TransformFunctionFactory.java @@ -36,12 +36,15 @@ import org.apache.pinot.common.utils.HashUtil; import org.apache.pinot.core.geospatial.transform.function.GeoToH3Function; import org.apache.pinot.core.geospatial.transform.function.StAreaFunction; import org.apache.pinot.core.geospatial.transform.function.StAsBinaryFunction; +import org.apache.pinot.core.geospatial.transform.function.StAsGeoJsonFunction; import org.apache.pinot.core.geospatial.transform.function.StAsTextFunction; import org.apache.pinot.core.geospatial.transform.function.StContainsFunction; import org.apache.pinot.core.geospatial.transform.function.StDistanceFunction; import org.apache.pinot.core.geospatial.transform.function.StEqualsFunction; +import org.apache.pinot.core.geospatial.transform.function.StGeogFromGeoJsonFunction; import org.apache.pinot.core.geospatial.transform.function.StGeogFromTextFunction; import org.apache.pinot.core.geospatial.transform.function.StGeogFromWKBFunction; +import org.apache.pinot.core.geospatial.transform.function.StGeomFromGeoJsonFunction; import org.apache.pinot.core.geospatial.transform.function.StGeomFromTextFunction; import org.apache.pinot.core.geospatial.transform.function.StGeomFromWKBFunction; import org.apache.pinot.core.geospatial.transform.function.StGeometryTypeFunction; @@ -180,8 +183,12 @@ public class TransformFunctionFactory { // geo constructors typeToImplementation.put(TransformFunctionType.ST_GEOG_FROM_TEXT, StGeogFromTextFunction.class); typeToImplementation.put(TransformFunctionType.ST_GEOG_FROM_WKB, StGeogFromWKBFunction.class); + typeToImplementation.put(TransformFunctionType.ST_GEOG_FROM_GEO_JSON, StGeogFromGeoJsonFunction.class); + typeToImplementation.put(TransformFunctionType.ST_GEOM_FROM_TEXT, StGeomFromTextFunction.class); typeToImplementation.put(TransformFunctionType.ST_GEOM_FROM_WKB, StGeomFromWKBFunction.class); + typeToImplementation.put(TransformFunctionType.ST_GEOM_FROM_GEO_JSON, StGeomFromGeoJsonFunction.class); + typeToImplementation.put(TransformFunctionType.ST_POINT, StPointFunction.class); typeToImplementation.put(TransformFunctionType.ST_POLYGON, StPolygonFunction.class); @@ -193,6 +200,7 @@ public class TransformFunctionFactory { // geo outputs typeToImplementation.put(TransformFunctionType.ST_AS_BINARY, StAsBinaryFunction.class); typeToImplementation.put(TransformFunctionType.ST_AS_TEXT, StAsTextFunction.class); + typeToImplementation.put(TransformFunctionType.ST_AS_GEO_JSON, StAsGeoJsonFunction.class); // geo relationship typeToImplementation.put(TransformFunctionType.ST_CONTAINS, StContainsFunction.class); diff --git a/pinot-core/src/test/java/org/apache/pinot/core/geospatial/transform/GeoInputOutputTest.java b/pinot-core/src/test/java/org/apache/pinot/core/geospatial/transform/GeoInputOutputTest.java index d7774b6498..82871e52ce 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/geospatial/transform/GeoInputOutputTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/geospatial/transform/GeoInputOutputTest.java @@ -59,4 +59,87 @@ public class GeoInputOutputTest extends GeoFunctionTest { assertStringFunction(String.format("ST_AsText(ST_GeogFromWKB(ST_AsBinary(ST_GeogFromText(%s))))", STRING_SV_COLUMN), new String[]{wkt}, Arrays.asList(new Column(STRING_SV_COLUMN, FieldSpec.DataType.STRING, new String[]{wkt}))); } + + @Test + public void testGeoJsonInputOutput() + throws Exception { + // empty geometries + assertAsGeoJsonAndBinary( + "{\"type\":\"Point\",\"coordinates\":[]," + + "\"crs\":{\"type\":\"name\",\"properties\":{\"name\":\"EPSG:#EPSG#\"}}}"); + + assertAsGeoJsonAndBinary( + "{\"type\":\"LineString\",\"coordinates\":[]," + + "\"crs\":{\"type\":\"name\",\"properties\":{\"name\":\"EPSG:#EPSG#\"}}}"); + + assertAsGeoJsonAndBinary( + "{\"type\":\"MultiLineString\",\"coordinates\":[]," + + "\"crs\":{\"type\":\"name\",\"properties\":{\"name\":\"EPSG:#EPSG#\"}}}"); + + assertAsGeoJsonAndBinary("{\"type\":\"Polygon\",\"coordinates\":[]," + + "\"crs\":{\"type\":\"name\",\"properties\":{\"name\":\"EPSG:#EPSG#\"}}}"); + + assertAsGeoJsonAndBinary("{\"type\":\"MultiPolygon\",\"coordinates\":[]," + + "\"crs\":{\"type\":\"name\",\"properties\":{\"name\":\"EPSG:#EPSG#\"}}}"); + + assertAsGeoJsonAndBinary("{\"type\":\"GeometryCollection\",\"geometries\":[]," + + "\"crs\":{\"type\":\"name\",\"properties\":{\"name\":\"EPSG:#EPSG#\"}}}"); + + // valid nonempty geometries + assertAsGeoJsonAndBinary( + "{\"type\":\"Point\"," + + "\"coordinates\":[1,2]," + + "\"crs\":{\"type\":\"name\",\"properties\":{\"name\":\"EPSG:#EPSG#\"}}}"); + + assertAsGeoJsonAndBinary("{\"type\":\"MultiPoint\"," + + "\"coordinates\":[[1,2],[3,4]]," + + "\"crs\":{\"type\":\"name\",\"properties\":{\"name\":\"EPSG:#EPSG#\"}}}"); + + assertAsGeoJsonAndBinary("{\"type\":\"LineString\"," + + "\"coordinates\":[[0.0,0.0],[1,2],[3,4]]," + + "\"crs\":{\"type\":\"name\",\"properties\":{\"name\":\"EPSG:#EPSG#\"}}}"); + + assertAsGeoJsonAndBinary("{\"type\":\"MultiLineString\"," + + "\"coordinates\":[" + + "[[100,0.0],[101,1]]," + + "[[102,2],[103,3]]" + + "]," + + "\"crs\":{\"type\":\"name\",\"properties\":{\"name\":\"EPSG:#EPSG#\"}}}"); + + assertAsGeoJsonAndBinary("{\"type\":\"Polygon\"," + + "\"coordinates\":[" + + "[[100,0.0],[100,1],[101,1],[101,0.0],[100,0.0]]," + + "[[100.8,0.8],[100.2,0.8],[100.2,0.2],[100.8,0.2],[100.8,0.8]]" + + "]," + + "\"crs\":{\"type\":\"name\",\"properties\":{\"name\":\"EPSG:#EPSG#\"}}}"); + + assertAsGeoJsonAndBinary("{\"type\":\"MultiPolygon\"," + + "\"coordinates\":[" + + "[[[0.0,0.0],[0.0,100],[100,100],[100,0.0],[0.0,0.0]],[[1,1],[10,1],[10,10],[1,10],[1,1]]]," + + "[[[200,200],[200,250],[250,250],[250,200],[200,200]]]" + + "]," + + "\"crs\":{\"type\":\"name\",\"properties\":{\"name\":\"EPSG:#EPSG#\"}}}"); + + assertAsGeoJsonAndBinary("{\"type\":\"GeometryCollection\"," + + "\"geometries\":[" + + "{\"type\":\"Point\",\"coordinates\":[100,0.0]}," + + "{\"type\":\"LineString\",\"coordinates\":[[101,0.0],[102,1]]" + + "}]," + + "\"crs\":{\"type\":\"name\",\"properties\":{\"name\":\"EPSG:#EPSG#\"}}}"); + } + + private void assertAsGeoJsonAndBinary(String geoJson) + throws Exception { + // assert geometry + assertStringFunction( + String.format("ST_AsGeoJSON(ST_GeomFromWKB(ST_AsBinary(ST_GeomFromGeoJSON(%s))))", STRING_SV_COLUMN), + new String[]{geoJson.replace("#EPSG#", "0")}, + Arrays.asList(new Column(STRING_SV_COLUMN, FieldSpec.DataType.STRING, new String[]{geoJson}))); + + // assert geography + assertStringFunction( + String.format("ST_AsGeoJSON(ST_GeogFromWKB(ST_AsBinary(ST_GeogFromGeoJSON(%s))))", STRING_SV_COLUMN), + new String[]{geoJson.replace("#EPSG#", "4326")}, + Arrays.asList(new Column(STRING_SV_COLUMN, FieldSpec.DataType.STRING, new String[]{geoJson}))); + } } diff --git a/pinot-core/src/test/java/org/apache/pinot/core/geospatial/transform/StGeometryTypeFunctionTest.java b/pinot-core/src/test/java/org/apache/pinot/core/geospatial/transform/StGeometryTypeFunctionTest.java index 977bb4cbfe..b17e412093 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/geospatial/transform/StGeometryTypeFunctionTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/geospatial/transform/StGeometryTypeFunctionTest.java @@ -42,8 +42,112 @@ public class StGeometryTypeFunctionTest extends GeoFunctionTest { // assert geometry assertStringFunction(String.format("ST_GeometryType(ST_GeomFromText(%s))", STRING_SV_COLUMN), new String[]{type}, Collections.singletonList(new Column(STRING_SV_COLUMN, DataType.STRING, new String[]{wkt}))); + // assert geography assertStringFunction(String.format("ST_GeometryType(ST_GeogFromText(%s))", STRING_SV_COLUMN), new String[]{type}, Collections.singletonList(new Column(STRING_SV_COLUMN, DataType.STRING, new String[]{wkt}))); } + + @Test + public void testGeoJsonConversion() + throws Exception { + assertGeoType("{\n \"type\": \"Point\",\n \"coordinates\": [100.0, 0.0]\n }", "Point"); + + assertGeoType("{\n" + + " \"type\": \"LineString\",\n" + + " \"coordinates\": [\n" + + " [100.0, 0.0],\n" + + " [101.0, 1.0]\n" + + " ]\n" + + " }", "LineString"); + + //no holes + assertGeoType("{\n" + + " \"type\": \"Polygon\",\n" + + " \"coordinates\": [\n" + + " [\n" + + " [100.0, 0.0],\n" + + " [101.0, 0.0],\n" + + " [101.0, 1.0],\n" + + " [100.0, 1.0],\n" + + " [100.0, 0.0]\n" + + " ]\n" + + " ]\n" + + " }", "Polygon"); + + //with holes + assertGeoType("{\n" + + " \"type\": \"Polygon\",\n" + + " \"coordinates\": [\n" + + " [\n" + + " [100.0, 0.0],\n" + + " [101.0, 0.0],\n" + + " [101.0, 1.0],\n" + + " [100.0, 1.0],\n" + + " [100.0, 0.0]\n" + + " ],\n" + + " [\n" + + " [100.8, 0.8],\n" + + " [100.8, 0.2],\n" + + " [100.2, 0.2],\n" + + " [100.2, 0.8],\n" + + " [100.8, 0.8]\n" + + " ]\n" + + " ]\n" + + "}", "Polygon"); + + assertGeoType("{\n" + + " \"type\": \"MultiPoint\",\n" + + " \"coordinates\": [\n" + + " [100.0, 0.0],\n" + + " [101.0, 1.0]\n" + + " ]\n" + + "}", "MultiPoint"); + + assertGeoType("{\n" + + " \"type\": \"MultiLineString\",\n" + + " \"coordinates\": [\n" + + " [\n" + + " [100.0, 0.0],\n" + + " [101.0, 1.0]\n" + + " ],\n" + + " [\n" + + " [102.0, 2.0],\n" + + " [103.0, 3.0]\n" + + " ]\n" + + " ]\n" + + "}", "MultiLineString"); + + assertGeoType( + "{\"type\":\"MultiPolygon\"," + + "\"coordinates\":[" + + "[[[0.0,0.0],[100,0.0],[100,100],[0.0,100],[0.0,0.0]],[[1,1],[1,10],[10,10],[10,1],[1,1]]]," + + "[[[200,200],[200,250],[250,250],[250,200],[200,200]]]" + + "]}", "MultiPolygon"); + + assertGeoType("{\n" + + " \"type\": \"GeometryCollection\",\n" + + " \"geometries\": [{\n" + + " \"type\": \"Point\",\n" + + " \"coordinates\": [100.0, 0.0]\n" + + " }, {\n" + + " \"type\": \"LineString\",\n" + + " \"coordinates\": [\n" + + " [101.0, 0.0],\n" + + " [102.0, 1.0]\n" + + " ]\n" + + " }]\n" + + "}", "GeometryCollection"); + } + + private void assertGeoType(String geoJson, String type) + throws Exception { + // assert geometry + assertStringFunction(String.format("ST_GeometryType(ST_GeomFromGeoJSON(%s))", STRING_SV_COLUMN), new String[]{type}, + Collections.singletonList(new Column(STRING_SV_COLUMN, DataType.STRING, new String[]{geoJson}))); + + // assert geography + assertStringFunction(String.format("ST_GeometryType(ST_GeogFromGeoJSON(%s))", STRING_SV_COLUMN), new String[]{type}, + Collections.singletonList(new Column(STRING_SV_COLUMN, DataType.STRING, new String[]{geoJson}))); + } } diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/GeometryUtils.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/GeometryUtils.java index aa46c3e4b1..42fdb63ad0 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/GeometryUtils.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/GeometryUtils.java @@ -26,6 +26,8 @@ import org.locationtech.jts.io.WKBReader; import org.locationtech.jts.io.WKBWriter; import org.locationtech.jts.io.WKTReader; import org.locationtech.jts.io.WKTWriter; +import org.locationtech.jts.io.geojson.GeoJsonReader; +import org.locationtech.jts.io.geojson.GeoJsonWriter; /** @@ -41,14 +43,23 @@ public class GeometryUtils { public static final int GEOGRAPHY_SRID = 4326; public static final byte GEOGRAPHY_SET_MASK = (byte) 0x80; public static final byte GEOGRAPHY_GET_MASK = (byte) 0x7f; + public static final GeometryFactory GEOMETRY_FACTORY = new GeometryFactory(); public static final GeometryFactory GEOGRAPHY_FACTORY = new GeometryFactory(new PrecisionModel(), GEOGRAPHY_SRID); + public static final WKTReader GEOMETRY_WKT_READER = new WKTReader(GEOMETRY_FACTORY); public static final WKTReader GEOGRAPHY_WKT_READER = new WKTReader(GEOGRAPHY_FACTORY); + + public static final GeoJsonReader GEOMETRY_GEO_JSON_READER = new GeoJsonReader(GEOMETRY_FACTORY); + public static final GeoJsonReader GEOGRAPHY_GEO_JSON_READER = new GeoJsonReader(GEOGRAPHY_FACTORY); + public static final WKBReader GEOMETRY_WKB_READER = new WKBReader(GEOMETRY_FACTORY); public static final WKBReader GEOGRAPHY_WKB_READER = new WKBReader(GEOGRAPHY_FACTORY); + public static final WKTWriter WKT_WRITER = new WKTWriter(); public static final WKBWriter WKB_WRITER = new WKBWriter(); + public static final GeoJsonWriter GEO_JSON_WRITER = new GeoJsonWriter(); + public static final double EARTH_RADIUS_KM = 6371.01; public static final double EARTH_RADIUS_M = EARTH_RADIUS_KM * 1000.0; public static final Joiner OR_JOINER = Joiner.on(" or "); --------------------------------------------------------------------- To unsubscribe, e-mail: commits-unsubscr...@pinot.apache.org For additional commands, e-mail: commits-h...@pinot.apache.org