ST_Within function

pinot-common/src/main/java/org/apache/pinot/common/function/TransformFunctionType.java

@@ -93,6 +93,7 @@ public enum TransformFunctionType {
   // Geo relationship
   ST_CONTAINS("ST_Contains"),
   ST_EQUALS("ST_Equals"),
+  ST_WITHIN("ST_Within"),
 
   // Geo indexing
   GEOTOH3("geoToH3");

pinot-core/src/main/java/org/apache/pinot/core/geospatial/transform/function/StWithinFunction.java

@@ -0,0 +1,91 @@
+/**
+ * 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 com.google.common.base.Preconditions;
+import java.util.List;
+import java.util.Map;
+import org.apache.pinot.core.operator.blocks.ProjectionBlock;
+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.segment.local.utils.GeometrySerializer;
+import org.apache.pinot.segment.local.utils.GeometryUtils;
+import org.apache.pinot.segment.spi.datasource.DataSource;
+import org.apache.pinot.spi.data.FieldSpec;
+import org.locationtech.jts.geom.Geometry;
+
+
+/**
+ * Function that checks the containment of the two geo-spatial objects. It returns true if and only if
+ * first geometry is completely inside second geometry.
+ */
+public class StWithinFunction extends BaseTransformFunction {
+  public static final String FUNCTION_NAME = "ST_Within";
+  private TransformFunction _firstArgument;
+  private TransformFunction _secondArgument;
+  private int[] _results;
+
+  @Override
+  public String getName() {
+    return FUNCTION_NAME;
+  }
+
+  @Override
+  public void init(List<TransformFunction> arguments, Map<String, DataSource> dataSourceMap) {
+    Preconditions
+        .checkArgument(arguments.size() == 2, "2 arguments are required for transform function: %s", getName());
+    TransformFunction transformFunction = arguments.get(0);
+    Preconditions.checkArgument(transformFunction.getResultMetadata().isSingleValue(),
+        "First argument must be single-valued for transform function: %s", getName());
+    Preconditions.checkArgument(transformFunction.getResultMetadata().getDataType() == FieldSpec.DataType.BYTES
+        || transformFunction instanceof LiteralTransformFunction, "The first argument must be of bytes type");
+    _firstArgument = transformFunction;
+    transformFunction = arguments.get(1);
+    Preconditions.checkArgument(transformFunction.getResultMetadata().isSingleValue(),
+        "Second argument must be single-valued for transform function: %s", getName());
+    Preconditions.checkArgument(transformFunction.getResultMetadata().getDataType() == FieldSpec.DataType.BYTES
+        || transformFunction instanceof LiteralTransformFunction, "The second argument must be of bytes type");
+    _secondArgument = transformFunction;
+  }
+
+  @Override
+  public TransformResultMetadata getResultMetadata() {
+    return INT_SV_NO_DICTIONARY_METADATA;
+  }
+
+  @Override
+  public int[] transformToIntValuesSV(ProjectionBlock projectionBlock) {
+    if (_results == null) {
+      _results = new int[projectionBlock.getNumDocs()];
+    }
+    byte[][] firstValues = _firstArgument.transformToBytesValuesSV(projectionBlock);
+    byte[][] secondValues = _secondArgument.transformToBytesValuesSV(projectionBlock);
+    for (int i = 0; i < projectionBlock.getNumDocs(); i++) {
+      Geometry firstGeometry = GeometrySerializer.deserialize(firstValues[i]);
+      Geometry secondGeometry = GeometrySerializer.deserialize(secondValues[i]);
+      if (GeometryUtils.isGeography(firstGeometry) || GeometryUtils.isGeography(secondGeometry)) {
+        throw new RuntimeException(String.format("%s is available for Geometry objects only", FUNCTION_NAME));
+      }
+      _results[i] = firstGeometry.within(secondGeometry) ? 1 : 0;
+    }
+    return _results;
+  }
+}

pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/TransformFunctionFactory.java

@@ -44,6 +44,7 @@
 import org.apache.pinot.core.geospatial.transform.function.StGeometryTypeFunction;
 import org.apache.pinot.core.geospatial.transform.function.StPointFunction;
 import org.apache.pinot.core.geospatial.transform.function.StPolygonFunction;
+import org.apache.pinot.core.geospatial.transform.function.StWithinFunction;
 import org.apache.pinot.core.operator.transform.function.SingleParamMathTransformFunction.AbsTransformFunction;
 import org.apache.pinot.core.operator.transform.function.SingleParamMathTransformFunction.CeilTransformFunction;
 import org.apache.pinot.core.operator.transform.function.SingleParamMathTransformFunction.ExpTransformFunction;
@@ -157,6 +158,7 @@ private TransformFunctionFactory() {
           // geo relationship
           put(canonicalize(TransformFunctionType.ST_CONTAINS.getName().toLowerCase()), StContainsFunction.class);
           put(canonicalize(TransformFunctionType.ST_EQUALS.getName().toLowerCase()), StEqualsFunction.class);
+          put(canonicalize(TransformFunctionType.ST_WITHIN.getName().toLowerCase()), StWithinFunction.class);
 
           // geo indexing
           put(canonicalize(TransformFunctionType.GEOTOH3.getName().toLowerCase()), GeoToH3Function.class);

pinot-core/src/test/java/org/apache/pinot/core/geospatial/transform/StWithinFunctionTest.java

@@ -0,0 +1,44 @@
+/**
+ * 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;
+
+import org.testng.annotations.Test;
+
+
+public class StWithinFunctionTest extends GeoFunctionTest {
+  @Test
+  public void testWithin()
+      throws Exception {
+    assertRelation("ST_Within", "POINT (25 25)", "POINT (20 20)", false);
+    assertRelation("ST_Within", "POINT (25 25)", "MULTIPOINT (20 20, 25 25)", true);
+    assertRelation("ST_Within", "POINT (25 25)", "LINESTRING (20 20, 30 30)", true);
+    assertRelation("ST_Within", "MULTIPOINT (25 25, 31 31)", "LINESTRING (20 20, 30 30)", false);
+    assertRelation("ST_Within", "LINESTRING (25 25, 27 27)", "LINESTRING (20 20, 30 30)", true);
+    assertRelation("ST_Within", "MULTILINESTRING ((3 4, 4 4), (2 1, 6 1))",
+        "MULTILINESTRING ((1 1, 5 1), (2 4, 4 4))", false);
+    assertRelation("ST_Within", "POLYGON ((1 1, 1 2, 2 2, 2 1, 1 1))", "POLYGON ((0 0, 0 4, 4 4, 4 0, 0 0))", true);
+    assertRelation("ST_Within", "POLYGON ((-1 -1, -1 2, 2 2, 2 -1, -1 -1))", "POLYGON ((0 0, 0 4, 4 4, 4 0, 0 0))",
+        false);
+    assertRelation("ST_Within", "POLYGON ((2 2, 2 3, 3 3, 3 2, 2 2))",
+        "MULTIPOLYGON (((0 0, 0 2, 2 2, 2 0, 0 0)), ((2 2, 2 4, 4 4, 4 2, 2 2)))", true);
+    assertRelation("ST_Within", "POLYGON ((0 0, 0 4, 4 4, 4 0, 0 0))", "LINESTRING (20 20, 30 30)", false);
+    assertRelation("ST_Within", "POLYGON ((0 0, 0 4, 4 4, 4 0, 0 0))", "LINESTRING EMPTY", false);
+    assertRelation("ST_Within", "POLYGON EMPTY", "LINESTRING (20 20, 30 30)", false);
+  }
+}