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 456e425920 Add `gridDisk` H3 Function (#15349)
456e425920 is described below
commit 456e4259209aeafd63699a83845976a3c4eb24f4
Author: ashishjayamohan <[email protected]>
AuthorDate: Sat Mar 29 17:23:17 2025 -0700
Add `gridDisk` H3 Function (#15349)
---
.../common/function/TransformFunctionType.java | 1 +
.../transform/function/GridDiskFunction.java | 90 +++++++++++++++
.../transform/function/ScalarFunctions.java | 20 ++++
.../function/TransformFunctionFactory.java | 2 +
.../transform/function/GridDiskFunctionTest.java | 122 +++++++++++++++++++++
5 files changed, 235 insertions(+)
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 75693e3719..05ee16a5e2 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
@@ -221,6 +221,7 @@ public enum TransformFunctionType {
OperandTypes.or(OperandTypes.family(SqlTypeFamily.BINARY,
SqlTypeFamily.INTEGER),
OperandTypes.family(SqlTypeFamily.NUMERIC, SqlTypeFamily.NUMERIC,
SqlTypeFamily.INTEGER))),
GRID_DISTANCE("gridDistance", ReturnTypes.BIGINT,
OperandTypes.NUMERIC_NUMERIC),
+ GRID_DISK("gridDisk", ReturnTypes.BIGINT, OperandTypes.NUMERIC_NUMERIC),
// Vector functions
// TODO: Once VECTOR type is defined, we should update here.
diff --git
a/pinot-core/src/main/java/org/apache/pinot/core/geospatial/transform/function/GridDiskFunction.java
b/pinot-core/src/main/java/org/apache/pinot/core/geospatial/transform/function/GridDiskFunction.java
new file mode 100644
index 0000000000..393caa10c5
--- /dev/null
+++
b/pinot-core/src/main/java/org/apache/pinot/core/geospatial/transform/function/GridDiskFunction.java
@@ -0,0 +1,90 @@
+/**
+ * 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 javax.annotation.Nullable;
+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.spi.data.FieldSpec;
+import org.roaringbitmap.RoaringBitmap;
+
+/**
+ * Function that returns all H3 indexes within a specified hexagonal grid
distance from a given origin index.
+ * The function takes two arguments:
+ * - gridDisk(origin, k)
+ */
+public class GridDiskFunction extends BaseTransformFunction {
+ public static final String FUNCTION_NAME = "gridDisk";
+ private TransformFunction _originArgument;
+ private TransformFunction _kArgument;
+
+ @Override
+ public String getName() {
+ return FUNCTION_NAME;
+ }
+
+ @Override
+ public void init(List<TransformFunction> arguments, Map<String,
ColumnContext> columnContextMap) {
+ super.init(arguments, columnContextMap);
+ Preconditions.checkArgument(arguments.size() == 2, "Transform function %s
requires 2 arguments", getName());
+
+ TransformFunction transformFunction = arguments.get(0);
+
Preconditions.checkArgument(transformFunction.getResultMetadata().isSingleValue(),
+ "First argument (origin) must be single-valued for transform function:
%s", getName());
+ _originArgument = transformFunction;
+
+ transformFunction = arguments.get(1);
+
Preconditions.checkArgument(transformFunction.getResultMetadata().isSingleValue(),
+ "Second argument (k) must be single-valued for transform function:
%s", getName());
+ _kArgument = transformFunction;
+ }
+
+ @Override
+ public TransformResultMetadata getResultMetadata() {
+ // Returns a list of H3 indexes, so it's multi-valued
+ return new TransformResultMetadata(FieldSpec.DataType.LONG, false, false);
+ }
+
+ @Nullable
+ public RoaringBitmap getNullBitmap(ValueBlock valueBlock) {
+ return null;
+ }
+
+ @Override
+ public long[][] transformToLongValuesMV(ValueBlock valueBlock) {
+ int numDocs = valueBlock.getNumDocs();
+ initLongValuesMV(numDocs);
+
+ long[] origins = _originArgument.transformToLongValuesSV(valueBlock);
+ int[] ks = _kArgument.transformToIntValuesSV(valueBlock);
+
+ for (int i = 0; i < numDocs; i++) {
+ long[] diskIndexes = ScalarFunctions.gridDisk(origins[i], ks[i]);
+ _longValuesMV[i] = diskIndexes;
+ }
+
+ return _longValuesMV;
+ }
+}
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 b1ffb1396a..b6390afa85 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
@@ -18,6 +18,7 @@
*/
package org.apache.pinot.core.geospatial.transform.function;
+import java.util.List;
import org.apache.pinot.segment.local.utils.GeometrySerializer;
import org.apache.pinot.segment.local.utils.GeometryUtils;
import org.apache.pinot.segment.local.utils.H3Utils;
@@ -267,4 +268,23 @@ public class ScalarFunctions {
}
return H3Utils.H3_CORE.gridDistance(firstH3Index, secondH3Index);
}
+
+ /**
+ * Returns all H3 indexes within a specified hexagonal grid distance from a
given origin index. The function considers
+ * hexagonal traversal rules and accounts for wrapping around pentagons.
+ *
+ * @param origin The starting H3 index from which to compute the grid disk.
+ * @param k The radius of the disk in hexagonal grid steps.
+ * @return An array of H3 indexes within the specified distance from the
origin.
+ */
+ @ScalarFunction
+ public static long[] gridDisk(long origin, int k) {
+ List<Long> diskCells = H3Utils.H3_CORE.gridDisk(origin, k);
+ long[] result = new long[diskCells.size()];
+ int index = 0;
+ for (Long cell : diskCells) {
+ result[index++] = cell;
+ }
+ return result;
+ }
}
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 55b67a4254..b2521e7681 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
@@ -34,6 +34,7 @@ import
org.apache.pinot.common.request.context.FunctionContext;
import org.apache.pinot.common.request.context.LiteralContext;
import org.apache.pinot.common.utils.HashUtil;
import org.apache.pinot.core.geospatial.transform.function.GeoToH3Function;
+import org.apache.pinot.core.geospatial.transform.function.GridDiskFunction;
import
org.apache.pinot.core.geospatial.transform.function.GridDistanceFunction;
import org.apache.pinot.core.geospatial.transform.function.StAreaFunction;
import org.apache.pinot.core.geospatial.transform.function.StAsBinaryFunction;
@@ -212,6 +213,7 @@ public class TransformFunctionFactory {
// geo indexing
typeToImplementation.put(TransformFunctionType.GEO_TO_H3,
GeoToH3Function.class);
typeToImplementation.put(TransformFunctionType.GRID_DISTANCE,
GridDistanceFunction.class);
+ typeToImplementation.put(TransformFunctionType.GRID_DISK,
GridDiskFunction.class);
// tuple selection
typeToImplementation.put(TransformFunctionType.LEAST,
LeastTransformFunction.class);
diff --git
a/pinot-core/src/test/java/org/apache/pinot/core/geospatial/transform/function/GridDiskFunctionTest.java
b/pinot-core/src/test/java/org/apache/pinot/core/geospatial/transform/function/GridDiskFunctionTest.java
new file mode 100644
index 0000000000..4b38a765ad
--- /dev/null
+++
b/pinot-core/src/test/java/org/apache/pinot/core/geospatial/transform/function/GridDiskFunctionTest.java
@@ -0,0 +1,122 @@
+/**
+ * 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.uber.h3core.H3Core;
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import org.apache.pinot.common.request.context.ExpressionContext;
+import org.apache.pinot.common.request.context.FunctionContext;
+import org.apache.pinot.core.operator.ColumnContext;
+import
org.apache.pinot.core.operator.transform.function.BaseTransformFunctionTest;
+import org.apache.pinot.core.operator.transform.function.TransformFunction;
+import
org.apache.pinot.core.operator.transform.function.TransformFunctionFactory;
+import org.apache.pinot.segment.spi.datasource.DataSource;
+import org.apache.pinot.spi.data.FieldSpec;
+import org.testng.annotations.BeforeClass;
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.assertEquals;
+
+public class GridDiskFunctionTest extends BaseTransformFunctionTest {
+ @BeforeClass
+ @Override
+ public void setUp() throws Exception {
+ super.setUp();
+ }
+
+ @Test
+ public void testGridDisk() throws IOException {
+ H3Core h3Core = H3Core.newInstance();
+ // Test point in San Francisco
+ double lat = 37.7749;
+ double lng = -122.4194;
+ long h3Index = h3Core.latLngToCell(lat, lng, 9);
+
+ // Test with k=1 (immediate neighbors)
+ int k = 1;
+
+ // Create expected values
+ List<Long> expectedDisk = h3Core.gridDisk(h3Index, k);
+ long[][] expectedValues = new long[NUM_ROWS][];
+ Arrays.fill(expectedValues,
expectedDisk.stream().mapToLong(Long::longValue).toArray());
+
+ // Create the transform function
+ ExpressionContext h3Context =
ExpressionContext.forLiteral(FieldSpec.DataType.LONG, h3Index);
+ ExpressionContext kContext =
ExpressionContext.forLiteral(FieldSpec.DataType.INT, k);
+ TransformFunction transformFunction = TransformFunctionFactory.get(
+ ExpressionContext.forFunction(new
FunctionContext(FunctionContext.Type.TRANSFORM, "gridDisk",
+ Arrays.asList(h3Context, kContext))), _dataSourceMap);
+
+ Map<String, ColumnContext> columnContextMap = new HashMap<>();
+ for (Map.Entry<String, DataSource> entry : _dataSourceMap.entrySet()) {
+ columnContextMap.put(entry.getKey(),
ColumnContext.fromDataSource(entry.getValue()));
+ }
+ transformFunction.init(Arrays.asList(
+ TransformFunctionFactory.get(h3Context, _dataSourceMap),
+ TransformFunctionFactory.get(kContext, _dataSourceMap)),
columnContextMap);
+
+ long[][] result =
transformFunction.transformToLongValuesMV(_projectionBlock);
+ assertArrayEquals(result, expectedValues);
+
+ // Test with k=0 (only the center)
+ k = 0;
+ expectedDisk = h3Core.gridDisk(h3Index, k);
+ expectedValues = new long[NUM_ROWS][];
+ Arrays.fill(expectedValues,
expectedDisk.stream().mapToLong(Long::longValue).toArray());
+
+ kContext = ExpressionContext.forLiteral(FieldSpec.DataType.INT, k);
+ transformFunction = TransformFunctionFactory.get(
+ ExpressionContext.forFunction(new
FunctionContext(FunctionContext.Type.TRANSFORM, "gridDisk",
+ Arrays.asList(h3Context, kContext))), _dataSourceMap);
+ transformFunction.init(Arrays.asList(
+ TransformFunctionFactory.get(h3Context, _dataSourceMap),
+ TransformFunctionFactory.get(kContext, _dataSourceMap)),
columnContextMap);
+
+ result = transformFunction.transformToLongValuesMV(_projectionBlock);
+ assertArrayEquals(result, expectedValues);
+
+ // Test with k=2 (two rings of neighbors)
+ k = 2;
+ expectedDisk = h3Core.gridDisk(h3Index, k);
+ expectedValues = new long[NUM_ROWS][];
+ Arrays.fill(expectedValues,
expectedDisk.stream().mapToLong(Long::longValue).toArray());
+
+ kContext = ExpressionContext.forLiteral(FieldSpec.DataType.INT, k);
+ transformFunction = TransformFunctionFactory.get(
+ ExpressionContext.forFunction(new
FunctionContext(FunctionContext.Type.TRANSFORM, "gridDisk",
+ Arrays.asList(h3Context, kContext))), _dataSourceMap);
+ transformFunction.init(Arrays.asList(
+ TransformFunctionFactory.get(h3Context, _dataSourceMap),
+ TransformFunctionFactory.get(kContext, _dataSourceMap)),
columnContextMap);
+
+ result = transformFunction.transformToLongValuesMV(_projectionBlock);
+ assertArrayEquals(result, expectedValues);
+ }
+
+ private void assertArrayEquals(long[][] actual, long[][] expected) {
+ assertEquals(actual.length, expected.length);
+ for (int i = 0; i < actual.length; i++) {
+ assertEquals(actual[i], expected[i]);
+ }
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
