From 2be4cd983fb13a2903ad61ea2c6212aa31e39364 Mon Sep 17 00:00:00 2001 From: Gal Lalouche Date: Fri, 13 Dec 2024 12:41:24 +0200 Subject: [PATCH] ESQL: Support ST_EXTENT_AGG (#117451) This PR adds support for ST_EXTENT_AGG aggregation, i.e., computing a bounding box over a set of points/shapes (Cartesian or geo). Note the difference between this aggregation and the already implemented scalar function ST_EXTENT. This isn't a very efficient implementation, and future PRs will attempt to read these extents directly from the doc values. We currently always use longitude wrapping, i.e., we may wrap around the dateline for a smaller bounding box. Future PRs will let the user control this behavior. Fixes #104659. --- docs/changelog/117451.yaml | 6 + .../functions/aggregation-functions.asciidoc | 2 + .../description/st_extent_agg.asciidoc | 5 + .../functions/examples/st_extent_agg.asciidoc | 13 + .../kibana/definition/st_extent_agg.json | 61 +++++ .../functions/kibana/docs/st_extent_agg.md | 12 + .../functions/layout/st_extent_agg.asciidoc | 15 ++ .../parameters/st_extent_agg.asciidoc | 6 + .../functions/signature/st_extent_agg.svg | 1 + .../functions/types/st_extent_agg.asciidoc | 12 + .../utils/SpatialEnvelopeVisitor.java | 72 ++--- .../utils/SpatialEnvelopeVisitorTests.java | 12 +- ...esianPointDocValuesAggregatorFunction.java | 187 +++++++++++++ ...ntDocValuesAggregatorFunctionSupplier.java | 41 +++ ...ntDocValuesGroupingAggregatorFunction.java | 230 ++++++++++++++++ ...anPointSourceValuesAggregatorFunction.java | 192 ++++++++++++++ ...ourceValuesAggregatorFunctionSupplier.java | 41 +++ ...ourceValuesGroupingAggregatorFunction.java | 235 +++++++++++++++++ ...xtentCartesianShapeAggregatorFunction.java | 192 ++++++++++++++ ...tesianShapeAggregatorFunctionSupplier.java | 40 +++ ...tesianShapeGroupingAggregatorFunction.java | 235 +++++++++++++++++ ...ntGeoPointDocValuesAggregatorFunction.java | 201 ++++++++++++++ ...ntDocValuesAggregatorFunctionSupplier.java | 40 +++ ...ntDocValuesGroupingAggregatorFunction.java | 242 +++++++++++++++++ ...eoPointSourceValuesAggregatorFunction.java | 206 +++++++++++++++ ...ourceValuesAggregatorFunctionSupplier.java | 41 +++ ...ourceValuesGroupingAggregatorFunction.java | 247 ++++++++++++++++++ ...atialExtentGeoShapeAggregatorFunction.java | 206 +++++++++++++++ ...entGeoShapeAggregatorFunctionSupplier.java | 40 +++ ...entGeoShapeGroupingAggregatorFunction.java | 247 ++++++++++++++++++ .../aggregation/AbstractArrayState.java | 4 +- .../spatial/CentroidPointAggregator.java | 7 + .../spatial/GeoPointEnvelopeVisitor.java | 63 +++++ .../aggregation/spatial/PointType.java | 107 ++++++++ .../spatial/SpatialAggregationUtils.java | 88 +++++++ ...roidCartesianPointDocValuesAggregator.java | 22 +- ...dCartesianPointSourceValuesAggregator.java | 20 +- ...alCentroidGeoPointDocValuesAggregator.java | 26 +- ...entroidGeoPointSourceValuesAggregator.java | 20 +- .../spatial/SpatialExtentAggregator.java | 36 +++ ...tentCartesianPointDocValuesAggregator.java | 42 +++ ...tCartesianPointSourceValuesAggregator.java | 45 ++++ ...SpatialExtentCartesianShapeAggregator.java | 43 +++ ...tialExtentGeoPointDocValuesAggregator.java | 45 ++++ ...lExtentGeoPointSourceValuesAggregator.java | 48 ++++ .../SpatialExtentGeoShapeAggregator.java | 46 ++++ .../spatial/SpatialExtentGroupingState.java | 154 +++++++++++ ...entGroupingStateWrappedLongitudeState.java | 182 +++++++++++++ ...tialExtentLongitudeWrappingAggregator.java | 62 +++++ .../spatial/SpatialExtentState.java | 82 ++++++ ...atialExtentStateWrappedLongitudeState.java | 91 +++++++ .../mapping-airports_no_doc_values.json | 4 +- .../src/main/resources/spatial.csv-spec | 99 +++++++ .../xpack/esql/action/EsqlCapabilities.java | 3 + .../function/EsqlFunctionRegistry.java | 2 + .../aggregate/AggregateWritables.java | 1 + .../aggregate/SpatialAggregateFunction.java | 31 ++- .../function/aggregate/SpatialCentroid.java | 43 ++- .../function/aggregate/SpatialExtent.java | 119 +++++++++ .../function/scalar/spatial/StEnvelope.java | 3 +- .../function/scalar/spatial/StXMax.java | 3 +- .../function/scalar/spatial/StXMin.java | 3 +- .../function/scalar/spatial/StYMax.java | 3 +- .../function/scalar/spatial/StYMin.java | 3 +- .../xpack/esql/planner/AggregateMapper.java | 52 +++- .../esql/expression/RectangleMatcher.java | 61 +++++ .../WellKnownBinaryBytesRefMatcher.java | 45 ++++ .../function/AbstractAggregationTestCase.java | 8 +- .../function/MultiRowTestCaseSupplier.java | 85 +++--- .../function/aggregate/CountTests.java | 5 +- .../aggregate/SpatialCentroidTests.java | 5 +- .../aggregate/SpatialExtentTests.java | 102 ++++++++ .../scalar/spatial/StEnvelopeTests.java | 5 +- .../function/scalar/spatial/StXMaxTests.java | 5 +- .../function/scalar/spatial/StXMinTests.java | 5 +- .../function/scalar/spatial/StYMaxTests.java | 5 +- .../function/scalar/spatial/StYMinTests.java | 5 +- .../optimizer/PhysicalPlanOptimizerTests.java | 135 +++++++++- .../rest-api-spec/test/esql/60_usage.yml | 4 +- 79 files changed, 4923 insertions(+), 234 deletions(-) create mode 100644 docs/changelog/117451.yaml create mode 100644 docs/reference/esql/functions/description/st_extent_agg.asciidoc create mode 100644 docs/reference/esql/functions/examples/st_extent_agg.asciidoc create mode 100644 docs/reference/esql/functions/kibana/definition/st_extent_agg.json create mode 100644 docs/reference/esql/functions/kibana/docs/st_extent_agg.md create mode 100644 docs/reference/esql/functions/layout/st_extent_agg.asciidoc create mode 100644 docs/reference/esql/functions/parameters/st_extent_agg.asciidoc create mode 100644 docs/reference/esql/functions/signature/st_extent_agg.svg create mode 100644 docs/reference/esql/functions/types/st_extent_agg.asciidoc create mode 100644 x-pack/plugin/esql/compute/src/main/generated/org/elasticsearch/compute/aggregation/spatial/SpatialExtentCartesianPointDocValuesAggregatorFunction.java create mode 100644 x-pack/plugin/esql/compute/src/main/generated/org/elasticsearch/compute/aggregation/spatial/SpatialExtentCartesianPointDocValuesAggregatorFunctionSupplier.java create mode 100644 x-pack/plugin/esql/compute/src/main/generated/org/elasticsearch/compute/aggregation/spatial/SpatialExtentCartesianPointDocValuesGroupingAggregatorFunction.java create mode 100644 x-pack/plugin/esql/compute/src/main/generated/org/elasticsearch/compute/aggregation/spatial/SpatialExtentCartesianPointSourceValuesAggregatorFunction.java create mode 100644 x-pack/plugin/esql/compute/src/main/generated/org/elasticsearch/compute/aggregation/spatial/SpatialExtentCartesianPointSourceValuesAggregatorFunctionSupplier.java create mode 100644 x-pack/plugin/esql/compute/src/main/generated/org/elasticsearch/compute/aggregation/spatial/SpatialExtentCartesianPointSourceValuesGroupingAggregatorFunction.java create mode 100644 x-pack/plugin/esql/compute/src/main/generated/org/elasticsearch/compute/aggregation/spatial/SpatialExtentCartesianShapeAggregatorFunction.java create mode 100644 x-pack/plugin/esql/compute/src/main/generated/org/elasticsearch/compute/aggregation/spatial/SpatialExtentCartesianShapeAggregatorFunctionSupplier.java create mode 100644 x-pack/plugin/esql/compute/src/main/generated/org/elasticsearch/compute/aggregation/spatial/SpatialExtentCartesianShapeGroupingAggregatorFunction.java create mode 100644 x-pack/plugin/esql/compute/src/main/generated/org/elasticsearch/compute/aggregation/spatial/SpatialExtentGeoPointDocValuesAggregatorFunction.java create mode 100644 x-pack/plugin/esql/compute/src/main/generated/org/elasticsearch/compute/aggregation/spatial/SpatialExtentGeoPointDocValuesAggregatorFunctionSupplier.java create mode 100644 x-pack/plugin/esql/compute/src/main/generated/org/elasticsearch/compute/aggregation/spatial/SpatialExtentGeoPointDocValuesGroupingAggregatorFunction.java create mode 100644 x-pack/plugin/esql/compute/src/main/generated/org/elasticsearch/compute/aggregation/spatial/SpatialExtentGeoPointSourceValuesAggregatorFunction.java create mode 100644 x-pack/plugin/esql/compute/src/main/generated/org/elasticsearch/compute/aggregation/spatial/SpatialExtentGeoPointSourceValuesAggregatorFunctionSupplier.java create mode 100644 x-pack/plugin/esql/compute/src/main/generated/org/elasticsearch/compute/aggregation/spatial/SpatialExtentGeoPointSourceValuesGroupingAggregatorFunction.java create mode 100644 x-pack/plugin/esql/compute/src/main/generated/org/elasticsearch/compute/aggregation/spatial/SpatialExtentGeoShapeAggregatorFunction.java create mode 100644 x-pack/plugin/esql/compute/src/main/generated/org/elasticsearch/compute/aggregation/spatial/SpatialExtentGeoShapeAggregatorFunctionSupplier.java create mode 100644 x-pack/plugin/esql/compute/src/main/generated/org/elasticsearch/compute/aggregation/spatial/SpatialExtentGeoShapeGroupingAggregatorFunction.java create mode 100644 x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/spatial/GeoPointEnvelopeVisitor.java create mode 100644 x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/spatial/PointType.java create mode 100644 x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/spatial/SpatialAggregationUtils.java create mode 100644 x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/spatial/SpatialExtentAggregator.java create mode 100644 x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/spatial/SpatialExtentCartesianPointDocValuesAggregator.java create mode 100644 x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/spatial/SpatialExtentCartesianPointSourceValuesAggregator.java create mode 100644 x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/spatial/SpatialExtentCartesianShapeAggregator.java create mode 100644 x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/spatial/SpatialExtentGeoPointDocValuesAggregator.java create mode 100644 x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/spatial/SpatialExtentGeoPointSourceValuesAggregator.java create mode 100644 x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/spatial/SpatialExtentGeoShapeAggregator.java create mode 100644 x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/spatial/SpatialExtentGroupingState.java create mode 100644 x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/spatial/SpatialExtentGroupingStateWrappedLongitudeState.java create mode 100644 x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/spatial/SpatialExtentLongitudeWrappingAggregator.java create mode 100644 x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/spatial/SpatialExtentState.java create mode 100644 x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/spatial/SpatialExtentStateWrappedLongitudeState.java create mode 100644 x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/aggregate/SpatialExtent.java create mode 100644 x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/RectangleMatcher.java create mode 100644 x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/WellKnownBinaryBytesRefMatcher.java create mode 100644 x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/aggregate/SpatialExtentTests.java diff --git a/docs/changelog/117451.yaml b/docs/changelog/117451.yaml new file mode 100644 index 0000000000000..bda0ca59e0953 --- /dev/null +++ b/docs/changelog/117451.yaml @@ -0,0 +1,6 @@ +pr: 117451 +summary: ST_EXTENT aggregation +area: ES|QL +type: feature +issues: + - 104659 diff --git a/docs/reference/esql/functions/aggregation-functions.asciidoc b/docs/reference/esql/functions/aggregation-functions.asciidoc index c2c2508ad5de2..24b42a6efd831 100644 --- a/docs/reference/esql/functions/aggregation-functions.asciidoc +++ b/docs/reference/esql/functions/aggregation-functions.asciidoc @@ -17,6 +17,7 @@ The <> command supports these aggregate functions: * <> * <> * experimental:[] <> +* experimental:[] <> * <> * <> * <> @@ -33,6 +34,7 @@ include::layout/median_absolute_deviation.asciidoc[] include::layout/min.asciidoc[] include::layout/percentile.asciidoc[] include::layout/st_centroid_agg.asciidoc[] +include::layout/st_extent_agg.asciidoc[] include::layout/std_dev.asciidoc[] include::layout/sum.asciidoc[] include::layout/top.asciidoc[] diff --git a/docs/reference/esql/functions/description/st_extent_agg.asciidoc b/docs/reference/esql/functions/description/st_extent_agg.asciidoc new file mode 100644 index 0000000000000..a9e1acfb0e6fb --- /dev/null +++ b/docs/reference/esql/functions/description/st_extent_agg.asciidoc @@ -0,0 +1,5 @@ +// This is generated by ESQL's AbstractFunctionTestCase. Do no edit it. See ../README.md for how to regenerate it. + +*Description* + +Calculate the spatial extent over a field with geometry type. Returns a bounding box for all values of the field. diff --git a/docs/reference/esql/functions/examples/st_extent_agg.asciidoc b/docs/reference/esql/functions/examples/st_extent_agg.asciidoc new file mode 100644 index 0000000000000..179be82103641 --- /dev/null +++ b/docs/reference/esql/functions/examples/st_extent_agg.asciidoc @@ -0,0 +1,13 @@ +// This is generated by ESQL's AbstractFunctionTestCase. Do no edit it. See ../README.md for how to regenerate it. + +*Example* + +[source.merge.styled,esql] +---- +include::{esql-specs}/spatial.csv-spec[tag=st_extent_agg-airports] +---- +[%header.monospaced.styled,format=dsv,separator=|] +|=== +include::{esql-specs}/spatial.csv-spec[tag=st_extent_agg-airports-result] +|=== + diff --git a/docs/reference/esql/functions/kibana/definition/st_extent_agg.json b/docs/reference/esql/functions/kibana/definition/st_extent_agg.json new file mode 100644 index 0000000000000..19afcc59e38a4 --- /dev/null +++ b/docs/reference/esql/functions/kibana/definition/st_extent_agg.json @@ -0,0 +1,61 @@ +{ + "comment" : "This is generated by ESQL's AbstractFunctionTestCase. Do no edit it. See ../README.md for how to regenerate it.", + "type" : "agg", + "name" : "st_extent_agg", + "description" : "Calculate the spatial extent over a field with geometry type. Returns a bounding box for all values of the field.", + "signatures" : [ + { + "params" : [ + { + "name" : "field", + "type" : "cartesian_point", + "optional" : false, + "description" : "" + } + ], + "variadic" : false, + "returnType" : "cartesian_shape" + }, + { + "params" : [ + { + "name" : "field", + "type" : "cartesian_shape", + "optional" : false, + "description" : "" + } + ], + "variadic" : false, + "returnType" : "cartesian_shape" + }, + { + "params" : [ + { + "name" : "field", + "type" : "geo_point", + "optional" : false, + "description" : "" + } + ], + "variadic" : false, + "returnType" : "geo_shape" + }, + { + "params" : [ + { + "name" : "field", + "type" : "geo_shape", + "optional" : false, + "description" : "" + } + ], + "variadic" : false, + "returnType" : "geo_shape" + } + ], + "examples" : [ + "FROM airports\n| WHERE country == \"India\"\n| STATS extent = ST_EXTENT_AGG(location)" + ], + "preview" : false, + "snapshot_only" : false +} diff --git a/docs/reference/esql/functions/kibana/docs/st_extent_agg.md b/docs/reference/esql/functions/kibana/docs/st_extent_agg.md new file mode 100644 index 0000000000000..a2e307c5b2c55 --- /dev/null +++ b/docs/reference/esql/functions/kibana/docs/st_extent_agg.md @@ -0,0 +1,12 @@ + + +### ST_EXTENT_AGG +Calculate the spatial extent over a field with geometry type. Returns a bounding box for all values of the field. + +``` +FROM airports +| WHERE country == "India" +| STATS extent = ST_EXTENT_AGG(location) +``` diff --git a/docs/reference/esql/functions/layout/st_extent_agg.asciidoc b/docs/reference/esql/functions/layout/st_extent_agg.asciidoc new file mode 100644 index 0000000000000..946bef661e70c --- /dev/null +++ b/docs/reference/esql/functions/layout/st_extent_agg.asciidoc @@ -0,0 +1,15 @@ +// This is generated by ESQL's AbstractFunctionTestCase. Do no edit it. See ../README.md for how to regenerate it. + +[discrete] +[[esql-st_extent_agg]] +=== `ST_EXTENT_AGG` + +*Syntax* + +[.text-center] +image::esql/functions/signature/st_extent_agg.svg[Embedded,opts=inline] + +include::../parameters/st_extent_agg.asciidoc[] +include::../description/st_extent_agg.asciidoc[] +include::../types/st_extent_agg.asciidoc[] +include::../examples/st_extent_agg.asciidoc[] diff --git a/docs/reference/esql/functions/parameters/st_extent_agg.asciidoc b/docs/reference/esql/functions/parameters/st_extent_agg.asciidoc new file mode 100644 index 0000000000000..8903aa1a472a3 --- /dev/null +++ b/docs/reference/esql/functions/parameters/st_extent_agg.asciidoc @@ -0,0 +1,6 @@ +// This is generated by ESQL's AbstractFunctionTestCase. Do no edit it. See ../README.md for how to regenerate it. + +*Parameters* + +`field`:: + diff --git a/docs/reference/esql/functions/signature/st_extent_agg.svg b/docs/reference/esql/functions/signature/st_extent_agg.svg new file mode 100644 index 0000000000000..bb19b68bfb08b --- /dev/null +++ b/docs/reference/esql/functions/signature/st_extent_agg.svg @@ -0,0 +1 @@ +ST_EXTENT_AGG(field) \ No newline at end of file diff --git a/docs/reference/esql/functions/types/st_extent_agg.asciidoc b/docs/reference/esql/functions/types/st_extent_agg.asciidoc new file mode 100644 index 0000000000000..c836aa1896f07 --- /dev/null +++ b/docs/reference/esql/functions/types/st_extent_agg.asciidoc @@ -0,0 +1,12 @@ +// This is generated by ESQL's AbstractFunctionTestCase. Do no edit it. See ../README.md for how to regenerate it. + +*Supported types* + +[%header.monospaced.styled,format=dsv,separator=|] +|=== +field | result +cartesian_point | cartesian_shape +cartesian_shape | cartesian_shape +geo_point | geo_shape +geo_shape | geo_shape +|=== diff --git a/libs/geo/src/main/java/org/elasticsearch/geometry/utils/SpatialEnvelopeVisitor.java b/libs/geo/src/main/java/org/elasticsearch/geometry/utils/SpatialEnvelopeVisitor.java index eee4a62c7d588..696be2808ed1f 100644 --- a/libs/geo/src/main/java/org/elasticsearch/geometry/utils/SpatialEnvelopeVisitor.java +++ b/libs/geo/src/main/java/org/elasticsearch/geometry/utils/SpatialEnvelopeVisitor.java @@ -83,10 +83,15 @@ public static Optional visitCartesian(Geometry geometry) { return Optional.empty(); } + public enum WrapLongitude { + NO_WRAP, + WRAP + } + /** * Determine the BBOX assuming the CRS is geographic (eg WGS84) and optionally wrapping the longitude around the dateline. */ - public static Optional visitGeo(Geometry geometry, boolean wrapLongitude) { + public static Optional visitGeo(Geometry geometry, WrapLongitude wrapLongitude) { var visitor = new SpatialEnvelopeVisitor(new GeoPointVisitor(wrapLongitude)); if (geometry.visit(visitor)) { return Optional.of(visitor.getResult()); @@ -181,40 +186,16 @@ public Rectangle getResult() { * */ public static class GeoPointVisitor implements PointVisitor { - private double minY = Double.POSITIVE_INFINITY; - private double maxY = Double.NEGATIVE_INFINITY; - private double minNegX = Double.POSITIVE_INFINITY; - private double maxNegX = Double.NEGATIVE_INFINITY; - private double minPosX = Double.POSITIVE_INFINITY; - private double maxPosX = Double.NEGATIVE_INFINITY; - - public double getMinY() { - return minY; - } - - public double getMaxY() { - return maxY; - } - - public double getMinNegX() { - return minNegX; - } + protected double minY = Double.POSITIVE_INFINITY; + protected double maxY = Double.NEGATIVE_INFINITY; + protected double minNegX = Double.POSITIVE_INFINITY; + protected double maxNegX = Double.NEGATIVE_INFINITY; + protected double minPosX = Double.POSITIVE_INFINITY; + protected double maxPosX = Double.NEGATIVE_INFINITY; - public double getMaxNegX() { - return maxNegX; - } - - public double getMinPosX() { - return minPosX; - } + private final WrapLongitude wrapLongitude; - public double getMaxPosX() { - return maxPosX; - } - - private final boolean wrapLongitude; - - public GeoPointVisitor(boolean wrapLongitude) { + public GeoPointVisitor(WrapLongitude wrapLongitude) { this.wrapLongitude = wrapLongitude; } @@ -253,32 +234,35 @@ public Rectangle getResult() { return getResult(minNegX, minPosX, maxNegX, maxPosX, maxY, minY, wrapLongitude); } - private static Rectangle getResult( + protected static Rectangle getResult( double minNegX, double minPosX, double maxNegX, double maxPosX, double maxY, double minY, - boolean wrapLongitude + WrapLongitude wrapLongitude ) { assert Double.isFinite(maxY); if (Double.isInfinite(minPosX)) { return new Rectangle(minNegX, maxNegX, maxY, minY); } else if (Double.isInfinite(minNegX)) { return new Rectangle(minPosX, maxPosX, maxY, minY); - } else if (wrapLongitude) { - double unwrappedWidth = maxPosX - minNegX; - double wrappedWidth = (180 - minPosX) - (-180 - maxNegX); - if (unwrappedWidth <= wrappedWidth) { - return new Rectangle(minNegX, maxPosX, maxY, minY); - } else { - return new Rectangle(minPosX, maxNegX, maxY, minY); - } } else { - return new Rectangle(minNegX, maxPosX, maxY, minY); + return switch (wrapLongitude) { + case NO_WRAP -> new Rectangle(minNegX, maxPosX, maxY, minY); + case WRAP -> maybeWrap(minNegX, minPosX, maxNegX, maxPosX, maxY, minY); + }; } } + + private static Rectangle maybeWrap(double minNegX, double minPosX, double maxNegX, double maxPosX, double maxY, double minY) { + double unwrappedWidth = maxPosX - minNegX; + double wrappedWidth = 360 + maxNegX - minPosX; + return unwrappedWidth <= wrappedWidth + ? new Rectangle(minNegX, maxPosX, maxY, minY) + : new Rectangle(minPosX, maxNegX, maxY, minY); + } } private boolean isValid() { diff --git a/libs/geo/src/test/java/org/elasticsearch/geometry/utils/SpatialEnvelopeVisitorTests.java b/libs/geo/src/test/java/org/elasticsearch/geometry/utils/SpatialEnvelopeVisitorTests.java index fc35df295e566..893a1700616a6 100644 --- a/libs/geo/src/test/java/org/elasticsearch/geometry/utils/SpatialEnvelopeVisitorTests.java +++ b/libs/geo/src/test/java/org/elasticsearch/geometry/utils/SpatialEnvelopeVisitorTests.java @@ -13,6 +13,7 @@ import org.elasticsearch.geo.ShapeTestUtils; import org.elasticsearch.geometry.Point; import org.elasticsearch.geometry.Rectangle; +import org.elasticsearch.geometry.utils.SpatialEnvelopeVisitor.WrapLongitude; import org.elasticsearch.test.ESTestCase; import static org.hamcrest.Matchers.equalTo; @@ -36,7 +37,7 @@ public void testVisitCartesianShape() { public void testVisitGeoShapeNoWrap() { for (int i = 0; i < 1000; i++) { var geometry = GeometryTestUtils.randomGeometryWithoutCircle(0, false); - var bbox = SpatialEnvelopeVisitor.visitGeo(geometry, false); + var bbox = SpatialEnvelopeVisitor.visitGeo(geometry, WrapLongitude.NO_WRAP); assertNotNull(bbox); assertTrue(i + ": " + geometry, bbox.isPresent()); var result = bbox.get(); @@ -48,7 +49,8 @@ public void testVisitGeoShapeNoWrap() { public void testVisitGeoShapeWrap() { for (int i = 0; i < 1000; i++) { var geometry = GeometryTestUtils.randomGeometryWithoutCircle(0, true); - var bbox = SpatialEnvelopeVisitor.visitGeo(geometry, false); + // TODO this should be WRAP instead + var bbox = SpatialEnvelopeVisitor.visitGeo(geometry, WrapLongitude.NO_WRAP); assertNotNull(bbox); assertTrue(i + ": " + geometry, bbox.isPresent()); var result = bbox.get(); @@ -81,7 +83,7 @@ public void testVisitCartesianPoints() { } public void testVisitGeoPointsNoWrapping() { - var visitor = new SpatialEnvelopeVisitor(new SpatialEnvelopeVisitor.GeoPointVisitor(false)); + var visitor = new SpatialEnvelopeVisitor(new SpatialEnvelopeVisitor.GeoPointVisitor(WrapLongitude.NO_WRAP)); double minY = Double.MAX_VALUE; double maxY = -Double.MAX_VALUE; double minX = Double.MAX_VALUE; @@ -103,7 +105,7 @@ public void testVisitGeoPointsNoWrapping() { } public void testVisitGeoPointsWrapping() { - var visitor = new SpatialEnvelopeVisitor(new SpatialEnvelopeVisitor.GeoPointVisitor(true)); + var visitor = new SpatialEnvelopeVisitor(new SpatialEnvelopeVisitor.GeoPointVisitor(WrapLongitude.WRAP)); double minY = Double.POSITIVE_INFINITY; double maxY = Double.NEGATIVE_INFINITY; double minNegX = Double.POSITIVE_INFINITY; @@ -145,7 +147,7 @@ public void testVisitGeoPointsWrapping() { } public void testWillCrossDateline() { - var visitor = new SpatialEnvelopeVisitor(new SpatialEnvelopeVisitor.GeoPointVisitor(true)); + var visitor = new SpatialEnvelopeVisitor(new SpatialEnvelopeVisitor.GeoPointVisitor(WrapLongitude.WRAP)); visitor.visit(new Point(-90.0, 0.0)); visitor.visit(new Point(90.0, 0.0)); assertCrossesDateline(visitor, false); diff --git a/x-pack/plugin/esql/compute/src/main/generated/org/elasticsearch/compute/aggregation/spatial/SpatialExtentCartesianPointDocValuesAggregatorFunction.java b/x-pack/plugin/esql/compute/src/main/generated/org/elasticsearch/compute/aggregation/spatial/SpatialExtentCartesianPointDocValuesAggregatorFunction.java new file mode 100644 index 0000000000000..21306036fbf50 --- /dev/null +++ b/x-pack/plugin/esql/compute/src/main/generated/org/elasticsearch/compute/aggregation/spatial/SpatialExtentCartesianPointDocValuesAggregatorFunction.java @@ -0,0 +1,187 @@ +// Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one +// or more contributor license agreements. Licensed under the Elastic License +// 2.0; you may not use this file except in compliance with the Elastic License +// 2.0. +package org.elasticsearch.compute.aggregation.spatial; + +import java.lang.Integer; +import java.lang.Override; +import java.lang.String; +import java.lang.StringBuilder; +import java.util.List; +import org.elasticsearch.compute.aggregation.AggregatorFunction; +import org.elasticsearch.compute.aggregation.IntermediateStateDesc; +import org.elasticsearch.compute.data.Block; +import org.elasticsearch.compute.data.BooleanVector; +import org.elasticsearch.compute.data.ElementType; +import org.elasticsearch.compute.data.IntBlock; +import org.elasticsearch.compute.data.IntVector; +import org.elasticsearch.compute.data.LongBlock; +import org.elasticsearch.compute.data.LongVector; +import org.elasticsearch.compute.data.Page; +import org.elasticsearch.compute.operator.DriverContext; + +/** + * {@link AggregatorFunction} implementation for {@link SpatialExtentCartesianPointDocValuesAggregator}. + * This class is generated. Do not edit it. + */ +public final class SpatialExtentCartesianPointDocValuesAggregatorFunction implements AggregatorFunction { + private static final List INTERMEDIATE_STATE_DESC = List.of( + new IntermediateStateDesc("minX", ElementType.INT), + new IntermediateStateDesc("maxX", ElementType.INT), + new IntermediateStateDesc("maxY", ElementType.INT), + new IntermediateStateDesc("minY", ElementType.INT) ); + + private final DriverContext driverContext; + + private final SpatialExtentState state; + + private final List channels; + + public SpatialExtentCartesianPointDocValuesAggregatorFunction(DriverContext driverContext, + List channels, SpatialExtentState state) { + this.driverContext = driverContext; + this.channels = channels; + this.state = state; + } + + public static SpatialExtentCartesianPointDocValuesAggregatorFunction create( + DriverContext driverContext, List channels) { + return new SpatialExtentCartesianPointDocValuesAggregatorFunction(driverContext, channels, SpatialExtentCartesianPointDocValuesAggregator.initSingle()); + } + + public static List intermediateStateDesc() { + return INTERMEDIATE_STATE_DESC; + } + + @Override + public int intermediateBlockCount() { + return INTERMEDIATE_STATE_DESC.size(); + } + + @Override + public void addRawInput(Page page, BooleanVector mask) { + if (mask.allFalse()) { + // Entire page masked away + return; + } + if (mask.allTrue()) { + // No masking + LongBlock block = page.getBlock(channels.get(0)); + LongVector vector = block.asVector(); + if (vector != null) { + addRawVector(vector); + } else { + addRawBlock(block); + } + return; + } + // Some positions masked away, others kept + LongBlock block = page.getBlock(channels.get(0)); + LongVector vector = block.asVector(); + if (vector != null) { + addRawVector(vector, mask); + } else { + addRawBlock(block, mask); + } + } + + private void addRawVector(LongVector vector) { + for (int i = 0; i < vector.getPositionCount(); i++) { + SpatialExtentCartesianPointDocValuesAggregator.combine(state, vector.getLong(i)); + } + } + + private void addRawVector(LongVector vector, BooleanVector mask) { + for (int i = 0; i < vector.getPositionCount(); i++) { + if (mask.getBoolean(i) == false) { + continue; + } + SpatialExtentCartesianPointDocValuesAggregator.combine(state, vector.getLong(i)); + } + } + + private void addRawBlock(LongBlock block) { + for (int p = 0; p < block.getPositionCount(); p++) { + if (block.isNull(p)) { + continue; + } + int start = block.getFirstValueIndex(p); + int end = start + block.getValueCount(p); + for (int i = start; i < end; i++) { + SpatialExtentCartesianPointDocValuesAggregator.combine(state, block.getLong(i)); + } + } + } + + private void addRawBlock(LongBlock block, BooleanVector mask) { + for (int p = 0; p < block.getPositionCount(); p++) { + if (mask.getBoolean(p) == false) { + continue; + } + if (block.isNull(p)) { + continue; + } + int start = block.getFirstValueIndex(p); + int end = start + block.getValueCount(p); + for (int i = start; i < end; i++) { + SpatialExtentCartesianPointDocValuesAggregator.combine(state, block.getLong(i)); + } + } + } + + @Override + public void addIntermediateInput(Page page) { + assert channels.size() == intermediateBlockCount(); + assert page.getBlockCount() >= channels.get(0) + intermediateStateDesc().size(); + Block minXUncast = page.getBlock(channels.get(0)); + if (minXUncast.areAllValuesNull()) { + return; + } + IntVector minX = ((IntBlock) minXUncast).asVector(); + assert minX.getPositionCount() == 1; + Block maxXUncast = page.getBlock(channels.get(1)); + if (maxXUncast.areAllValuesNull()) { + return; + } + IntVector maxX = ((IntBlock) maxXUncast).asVector(); + assert maxX.getPositionCount() == 1; + Block maxYUncast = page.getBlock(channels.get(2)); + if (maxYUncast.areAllValuesNull()) { + return; + } + IntVector maxY = ((IntBlock) maxYUncast).asVector(); + assert maxY.getPositionCount() == 1; + Block minYUncast = page.getBlock(channels.get(3)); + if (minYUncast.areAllValuesNull()) { + return; + } + IntVector minY = ((IntBlock) minYUncast).asVector(); + assert minY.getPositionCount() == 1; + SpatialExtentCartesianPointDocValuesAggregator.combineIntermediate(state, minX.getInt(0), maxX.getInt(0), maxY.getInt(0), minY.getInt(0)); + } + + @Override + public void evaluateIntermediate(Block[] blocks, int offset, DriverContext driverContext) { + state.toIntermediate(blocks, offset, driverContext); + } + + @Override + public void evaluateFinal(Block[] blocks, int offset, DriverContext driverContext) { + blocks[offset] = SpatialExtentCartesianPointDocValuesAggregator.evaluateFinal(state, driverContext); + } + + @Override + public String toString() { + StringBuilder sb = new StringBuilder(); + sb.append(getClass().getSimpleName()).append("["); + sb.append("channels=").append(channels); + sb.append("]"); + return sb.toString(); + } + + @Override + public void close() { + state.close(); + } +} diff --git a/x-pack/plugin/esql/compute/src/main/generated/org/elasticsearch/compute/aggregation/spatial/SpatialExtentCartesianPointDocValuesAggregatorFunctionSupplier.java b/x-pack/plugin/esql/compute/src/main/generated/org/elasticsearch/compute/aggregation/spatial/SpatialExtentCartesianPointDocValuesAggregatorFunctionSupplier.java new file mode 100644 index 0000000000000..751ea3b4c4a9d --- /dev/null +++ b/x-pack/plugin/esql/compute/src/main/generated/org/elasticsearch/compute/aggregation/spatial/SpatialExtentCartesianPointDocValuesAggregatorFunctionSupplier.java @@ -0,0 +1,41 @@ +// Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one +// or more contributor license agreements. Licensed under the Elastic License +// 2.0; you may not use this file except in compliance with the Elastic License +// 2.0. +package org.elasticsearch.compute.aggregation.spatial; + +import java.lang.Integer; +import java.lang.Override; +import java.lang.String; +import java.util.List; +import org.elasticsearch.compute.aggregation.AggregatorFunctionSupplier; +import org.elasticsearch.compute.operator.DriverContext; + +/** + * {@link AggregatorFunctionSupplier} implementation for {@link SpatialExtentCartesianPointDocValuesAggregator}. + * This class is generated. Do not edit it. + */ +public final class SpatialExtentCartesianPointDocValuesAggregatorFunctionSupplier implements AggregatorFunctionSupplier { + private final List channels; + + public SpatialExtentCartesianPointDocValuesAggregatorFunctionSupplier(List channels) { + this.channels = channels; + } + + @Override + public SpatialExtentCartesianPointDocValuesAggregatorFunction aggregator( + DriverContext driverContext) { + return SpatialExtentCartesianPointDocValuesAggregatorFunction.create(driverContext, channels); + } + + @Override + public SpatialExtentCartesianPointDocValuesGroupingAggregatorFunction groupingAggregator( + DriverContext driverContext) { + return SpatialExtentCartesianPointDocValuesGroupingAggregatorFunction.create(channels, driverContext); + } + + @Override + public String describe() { + return "spatial_extent_cartesian_point_doc of valuess"; + } +} diff --git a/x-pack/plugin/esql/compute/src/main/generated/org/elasticsearch/compute/aggregation/spatial/SpatialExtentCartesianPointDocValuesGroupingAggregatorFunction.java b/x-pack/plugin/esql/compute/src/main/generated/org/elasticsearch/compute/aggregation/spatial/SpatialExtentCartesianPointDocValuesGroupingAggregatorFunction.java new file mode 100644 index 0000000000000..a5191e57959b8 --- /dev/null +++ b/x-pack/plugin/esql/compute/src/main/generated/org/elasticsearch/compute/aggregation/spatial/SpatialExtentCartesianPointDocValuesGroupingAggregatorFunction.java @@ -0,0 +1,230 @@ +// Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one +// or more contributor license agreements. Licensed under the Elastic License +// 2.0; you may not use this file except in compliance with the Elastic License +// 2.0. +package org.elasticsearch.compute.aggregation.spatial; + +import java.lang.Integer; +import java.lang.Override; +import java.lang.String; +import java.lang.StringBuilder; +import java.util.List; +import org.elasticsearch.compute.aggregation.GroupingAggregatorFunction; +import org.elasticsearch.compute.aggregation.IntermediateStateDesc; +import org.elasticsearch.compute.aggregation.SeenGroupIds; +import org.elasticsearch.compute.data.Block; +import org.elasticsearch.compute.data.ElementType; +import org.elasticsearch.compute.data.IntBlock; +import org.elasticsearch.compute.data.IntVector; +import org.elasticsearch.compute.data.LongBlock; +import org.elasticsearch.compute.data.LongVector; +import org.elasticsearch.compute.data.Page; +import org.elasticsearch.compute.operator.DriverContext; + +/** + * {@link GroupingAggregatorFunction} implementation for {@link SpatialExtentCartesianPointDocValuesAggregator}. + * This class is generated. Do not edit it. + */ +public final class SpatialExtentCartesianPointDocValuesGroupingAggregatorFunction implements GroupingAggregatorFunction { + private static final List INTERMEDIATE_STATE_DESC = List.of( + new IntermediateStateDesc("minX", ElementType.INT), + new IntermediateStateDesc("maxX", ElementType.INT), + new IntermediateStateDesc("maxY", ElementType.INT), + new IntermediateStateDesc("minY", ElementType.INT) ); + + private final SpatialExtentGroupingState state; + + private final List channels; + + private final DriverContext driverContext; + + public SpatialExtentCartesianPointDocValuesGroupingAggregatorFunction(List channels, + SpatialExtentGroupingState state, DriverContext driverContext) { + this.channels = channels; + this.state = state; + this.driverContext = driverContext; + } + + public static SpatialExtentCartesianPointDocValuesGroupingAggregatorFunction create( + List channels, DriverContext driverContext) { + return new SpatialExtentCartesianPointDocValuesGroupingAggregatorFunction(channels, SpatialExtentCartesianPointDocValuesAggregator.initGrouping(), driverContext); + } + + public static List intermediateStateDesc() { + return INTERMEDIATE_STATE_DESC; + } + + @Override + public int intermediateBlockCount() { + return INTERMEDIATE_STATE_DESC.size(); + } + + @Override + public GroupingAggregatorFunction.AddInput prepareProcessPage(SeenGroupIds seenGroupIds, + Page page) { + LongBlock valuesBlock = page.getBlock(channels.get(0)); + LongVector valuesVector = valuesBlock.asVector(); + if (valuesVector == null) { + if (valuesBlock.mayHaveNulls()) { + state.enableGroupIdTracking(seenGroupIds); + } + return new GroupingAggregatorFunction.AddInput() { + @Override + public void add(int positionOffset, IntBlock groupIds) { + addRawInput(positionOffset, groupIds, valuesBlock); + } + + @Override + public void add(int positionOffset, IntVector groupIds) { + addRawInput(positionOffset, groupIds, valuesBlock); + } + + @Override + public void close() { + } + }; + } + return new GroupingAggregatorFunction.AddInput() { + @Override + public void add(int positionOffset, IntBlock groupIds) { + addRawInput(positionOffset, groupIds, valuesVector); + } + + @Override + public void add(int positionOffset, IntVector groupIds) { + addRawInput(positionOffset, groupIds, valuesVector); + } + + @Override + public void close() { + } + }; + } + + private void addRawInput(int positionOffset, IntVector groups, LongBlock values) { + for (int groupPosition = 0; groupPosition < groups.getPositionCount(); groupPosition++) { + int groupId = groups.getInt(groupPosition); + if (values.isNull(groupPosition + positionOffset)) { + continue; + } + int valuesStart = values.getFirstValueIndex(groupPosition + positionOffset); + int valuesEnd = valuesStart + values.getValueCount(groupPosition + positionOffset); + for (int v = valuesStart; v < valuesEnd; v++) { + SpatialExtentCartesianPointDocValuesAggregator.combine(state, groupId, values.getLong(v)); + } + } + } + + private void addRawInput(int positionOffset, IntVector groups, LongVector values) { + for (int groupPosition = 0; groupPosition < groups.getPositionCount(); groupPosition++) { + int groupId = groups.getInt(groupPosition); + SpatialExtentCartesianPointDocValuesAggregator.combine(state, groupId, values.getLong(groupPosition + positionOffset)); + } + } + + private void addRawInput(int positionOffset, IntBlock groups, LongBlock values) { + for (int groupPosition = 0; groupPosition < groups.getPositionCount(); groupPosition++) { + if (groups.isNull(groupPosition)) { + continue; + } + int groupStart = groups.getFirstValueIndex(groupPosition); + int groupEnd = groupStart + groups.getValueCount(groupPosition); + for (int g = groupStart; g < groupEnd; g++) { + int groupId = groups.getInt(g); + if (values.isNull(groupPosition + positionOffset)) { + continue; + } + int valuesStart = values.getFirstValueIndex(groupPosition + positionOffset); + int valuesEnd = valuesStart + values.getValueCount(groupPosition + positionOffset); + for (int v = valuesStart; v < valuesEnd; v++) { + SpatialExtentCartesianPointDocValuesAggregator.combine(state, groupId, values.getLong(v)); + } + } + } + } + + private void addRawInput(int positionOffset, IntBlock groups, LongVector values) { + for (int groupPosition = 0; groupPosition < groups.getPositionCount(); groupPosition++) { + if (groups.isNull(groupPosition)) { + continue; + } + int groupStart = groups.getFirstValueIndex(groupPosition); + int groupEnd = groupStart + groups.getValueCount(groupPosition); + for (int g = groupStart; g < groupEnd; g++) { + int groupId = groups.getInt(g); + SpatialExtentCartesianPointDocValuesAggregator.combine(state, groupId, values.getLong(groupPosition + positionOffset)); + } + } + } + + @Override + public void selectedMayContainUnseenGroups(SeenGroupIds seenGroupIds) { + state.enableGroupIdTracking(seenGroupIds); + } + + @Override + public void addIntermediateInput(int positionOffset, IntVector groups, Page page) { + state.enableGroupIdTracking(new SeenGroupIds.Empty()); + assert channels.size() == intermediateBlockCount(); + Block minXUncast = page.getBlock(channels.get(0)); + if (minXUncast.areAllValuesNull()) { + return; + } + IntVector minX = ((IntBlock) minXUncast).asVector(); + Block maxXUncast = page.getBlock(channels.get(1)); + if (maxXUncast.areAllValuesNull()) { + return; + } + IntVector maxX = ((IntBlock) maxXUncast).asVector(); + Block maxYUncast = page.getBlock(channels.get(2)); + if (maxYUncast.areAllValuesNull()) { + return; + } + IntVector maxY = ((IntBlock) maxYUncast).asVector(); + Block minYUncast = page.getBlock(channels.get(3)); + if (minYUncast.areAllValuesNull()) { + return; + } + IntVector minY = ((IntBlock) minYUncast).asVector(); + assert minX.getPositionCount() == maxX.getPositionCount() && minX.getPositionCount() == maxY.getPositionCount() && minX.getPositionCount() == minY.getPositionCount(); + for (int groupPosition = 0; groupPosition < groups.getPositionCount(); groupPosition++) { + int groupId = groups.getInt(groupPosition); + SpatialExtentCartesianPointDocValuesAggregator.combineIntermediate(state, groupId, minX.getInt(groupPosition + positionOffset), maxX.getInt(groupPosition + positionOffset), maxY.getInt(groupPosition + positionOffset), minY.getInt(groupPosition + positionOffset)); + } + } + + @Override + public void addIntermediateRowInput(int groupId, GroupingAggregatorFunction input, int position) { + if (input.getClass() != getClass()) { + throw new IllegalArgumentException("expected " + getClass() + "; got " + input.getClass()); + } + SpatialExtentGroupingState inState = ((SpatialExtentCartesianPointDocValuesGroupingAggregatorFunction) input).state; + state.enableGroupIdTracking(new SeenGroupIds.Empty()); + SpatialExtentCartesianPointDocValuesAggregator.combineStates(state, groupId, inState, position); + } + + @Override + public void evaluateIntermediate(Block[] blocks, int offset, IntVector selected) { + state.toIntermediate(blocks, offset, selected, driverContext); + } + + @Override + public void evaluateFinal(Block[] blocks, int offset, IntVector selected, + DriverContext driverContext) { + blocks[offset] = SpatialExtentCartesianPointDocValuesAggregator.evaluateFinal(state, selected, driverContext); + } + + @Override + public String toString() { + StringBuilder sb = new StringBuilder(); + sb.append(getClass().getSimpleName()).append("["); + sb.append("channels=").append(channels); + sb.append("]"); + return sb.toString(); + } + + @Override + public void close() { + state.close(); + } +} diff --git a/x-pack/plugin/esql/compute/src/main/generated/org/elasticsearch/compute/aggregation/spatial/SpatialExtentCartesianPointSourceValuesAggregatorFunction.java b/x-pack/plugin/esql/compute/src/main/generated/org/elasticsearch/compute/aggregation/spatial/SpatialExtentCartesianPointSourceValuesAggregatorFunction.java new file mode 100644 index 0000000000000..6610168e1df21 --- /dev/null +++ b/x-pack/plugin/esql/compute/src/main/generated/org/elasticsearch/compute/aggregation/spatial/SpatialExtentCartesianPointSourceValuesAggregatorFunction.java @@ -0,0 +1,192 @@ +// Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one +// or more contributor license agreements. Licensed under the Elastic License +// 2.0; you may not use this file except in compliance with the Elastic License +// 2.0. +package org.elasticsearch.compute.aggregation.spatial; + +import java.lang.Integer; +import java.lang.Override; +import java.lang.String; +import java.lang.StringBuilder; +import java.util.List; +import org.apache.lucene.util.BytesRef; +import org.elasticsearch.compute.aggregation.AggregatorFunction; +import org.elasticsearch.compute.aggregation.IntermediateStateDesc; +import org.elasticsearch.compute.data.Block; +import org.elasticsearch.compute.data.BooleanVector; +import org.elasticsearch.compute.data.BytesRefBlock; +import org.elasticsearch.compute.data.BytesRefVector; +import org.elasticsearch.compute.data.ElementType; +import org.elasticsearch.compute.data.IntBlock; +import org.elasticsearch.compute.data.IntVector; +import org.elasticsearch.compute.data.Page; +import org.elasticsearch.compute.operator.DriverContext; + +/** + * {@link AggregatorFunction} implementation for {@link SpatialExtentCartesianPointSourceValuesAggregator}. + * This class is generated. Do not edit it. + */ +public final class SpatialExtentCartesianPointSourceValuesAggregatorFunction implements AggregatorFunction { + private static final List INTERMEDIATE_STATE_DESC = List.of( + new IntermediateStateDesc("minX", ElementType.INT), + new IntermediateStateDesc("maxX", ElementType.INT), + new IntermediateStateDesc("maxY", ElementType.INT), + new IntermediateStateDesc("minY", ElementType.INT) ); + + private final DriverContext driverContext; + + private final SpatialExtentState state; + + private final List channels; + + public SpatialExtentCartesianPointSourceValuesAggregatorFunction(DriverContext driverContext, + List channels, SpatialExtentState state) { + this.driverContext = driverContext; + this.channels = channels; + this.state = state; + } + + public static SpatialExtentCartesianPointSourceValuesAggregatorFunction create( + DriverContext driverContext, List channels) { + return new SpatialExtentCartesianPointSourceValuesAggregatorFunction(driverContext, channels, SpatialExtentCartesianPointSourceValuesAggregator.initSingle()); + } + + public static List intermediateStateDesc() { + return INTERMEDIATE_STATE_DESC; + } + + @Override + public int intermediateBlockCount() { + return INTERMEDIATE_STATE_DESC.size(); + } + + @Override + public void addRawInput(Page page, BooleanVector mask) { + if (mask.allFalse()) { + // Entire page masked away + return; + } + if (mask.allTrue()) { + // No masking + BytesRefBlock block = page.getBlock(channels.get(0)); + BytesRefVector vector = block.asVector(); + if (vector != null) { + addRawVector(vector); + } else { + addRawBlock(block); + } + return; + } + // Some positions masked away, others kept + BytesRefBlock block = page.getBlock(channels.get(0)); + BytesRefVector vector = block.asVector(); + if (vector != null) { + addRawVector(vector, mask); + } else { + addRawBlock(block, mask); + } + } + + private void addRawVector(BytesRefVector vector) { + BytesRef scratch = new BytesRef(); + for (int i = 0; i < vector.getPositionCount(); i++) { + SpatialExtentCartesianPointSourceValuesAggregator.combine(state, vector.getBytesRef(i, scratch)); + } + } + + private void addRawVector(BytesRefVector vector, BooleanVector mask) { + BytesRef scratch = new BytesRef(); + for (int i = 0; i < vector.getPositionCount(); i++) { + if (mask.getBoolean(i) == false) { + continue; + } + SpatialExtentCartesianPointSourceValuesAggregator.combine(state, vector.getBytesRef(i, scratch)); + } + } + + private void addRawBlock(BytesRefBlock block) { + BytesRef scratch = new BytesRef(); + for (int p = 0; p < block.getPositionCount(); p++) { + if (block.isNull(p)) { + continue; + } + int start = block.getFirstValueIndex(p); + int end = start + block.getValueCount(p); + for (int i = start; i < end; i++) { + SpatialExtentCartesianPointSourceValuesAggregator.combine(state, block.getBytesRef(i, scratch)); + } + } + } + + private void addRawBlock(BytesRefBlock block, BooleanVector mask) { + BytesRef scratch = new BytesRef(); + for (int p = 0; p < block.getPositionCount(); p++) { + if (mask.getBoolean(p) == false) { + continue; + } + if (block.isNull(p)) { + continue; + } + int start = block.getFirstValueIndex(p); + int end = start + block.getValueCount(p); + for (int i = start; i < end; i++) { + SpatialExtentCartesianPointSourceValuesAggregator.combine(state, block.getBytesRef(i, scratch)); + } + } + } + + @Override + public void addIntermediateInput(Page page) { + assert channels.size() == intermediateBlockCount(); + assert page.getBlockCount() >= channels.get(0) + intermediateStateDesc().size(); + Block minXUncast = page.getBlock(channels.get(0)); + if (minXUncast.areAllValuesNull()) { + return; + } + IntVector minX = ((IntBlock) minXUncast).asVector(); + assert minX.getPositionCount() == 1; + Block maxXUncast = page.getBlock(channels.get(1)); + if (maxXUncast.areAllValuesNull()) { + return; + } + IntVector maxX = ((IntBlock) maxXUncast).asVector(); + assert maxX.getPositionCount() == 1; + Block maxYUncast = page.getBlock(channels.get(2)); + if (maxYUncast.areAllValuesNull()) { + return; + } + IntVector maxY = ((IntBlock) maxYUncast).asVector(); + assert maxY.getPositionCount() == 1; + Block minYUncast = page.getBlock(channels.get(3)); + if (minYUncast.areAllValuesNull()) { + return; + } + IntVector minY = ((IntBlock) minYUncast).asVector(); + assert minY.getPositionCount() == 1; + SpatialExtentCartesianPointSourceValuesAggregator.combineIntermediate(state, minX.getInt(0), maxX.getInt(0), maxY.getInt(0), minY.getInt(0)); + } + + @Override + public void evaluateIntermediate(Block[] blocks, int offset, DriverContext driverContext) { + state.toIntermediate(blocks, offset, driverContext); + } + + @Override + public void evaluateFinal(Block[] blocks, int offset, DriverContext driverContext) { + blocks[offset] = SpatialExtentCartesianPointSourceValuesAggregator.evaluateFinal(state, driverContext); + } + + @Override + public String toString() { + StringBuilder sb = new StringBuilder(); + sb.append(getClass().getSimpleName()).append("["); + sb.append("channels=").append(channels); + sb.append("]"); + return sb.toString(); + } + + @Override + public void close() { + state.close(); + } +} diff --git a/x-pack/plugin/esql/compute/src/main/generated/org/elasticsearch/compute/aggregation/spatial/SpatialExtentCartesianPointSourceValuesAggregatorFunctionSupplier.java b/x-pack/plugin/esql/compute/src/main/generated/org/elasticsearch/compute/aggregation/spatial/SpatialExtentCartesianPointSourceValuesAggregatorFunctionSupplier.java new file mode 100644 index 0000000000000..7f4d1d69ae928 --- /dev/null +++ b/x-pack/plugin/esql/compute/src/main/generated/org/elasticsearch/compute/aggregation/spatial/SpatialExtentCartesianPointSourceValuesAggregatorFunctionSupplier.java @@ -0,0 +1,41 @@ +// Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one +// or more contributor license agreements. Licensed under the Elastic License +// 2.0; you may not use this file except in compliance with the Elastic License +// 2.0. +package org.elasticsearch.compute.aggregation.spatial; + +import java.lang.Integer; +import java.lang.Override; +import java.lang.String; +import java.util.List; +import org.elasticsearch.compute.aggregation.AggregatorFunctionSupplier; +import org.elasticsearch.compute.operator.DriverContext; + +/** + * {@link AggregatorFunctionSupplier} implementation for {@link SpatialExtentCartesianPointSourceValuesAggregator}. + * This class is generated. Do not edit it. + */ +public final class SpatialExtentCartesianPointSourceValuesAggregatorFunctionSupplier implements AggregatorFunctionSupplier { + private final List channels; + + public SpatialExtentCartesianPointSourceValuesAggregatorFunctionSupplier(List channels) { + this.channels = channels; + } + + @Override + public SpatialExtentCartesianPointSourceValuesAggregatorFunction aggregator( + DriverContext driverContext) { + return SpatialExtentCartesianPointSourceValuesAggregatorFunction.create(driverContext, channels); + } + + @Override + public SpatialExtentCartesianPointSourceValuesGroupingAggregatorFunction groupingAggregator( + DriverContext driverContext) { + return SpatialExtentCartesianPointSourceValuesGroupingAggregatorFunction.create(channels, driverContext); + } + + @Override + public String describe() { + return "spatial_extent_cartesian_point_source of valuess"; + } +} diff --git a/x-pack/plugin/esql/compute/src/main/generated/org/elasticsearch/compute/aggregation/spatial/SpatialExtentCartesianPointSourceValuesGroupingAggregatorFunction.java b/x-pack/plugin/esql/compute/src/main/generated/org/elasticsearch/compute/aggregation/spatial/SpatialExtentCartesianPointSourceValuesGroupingAggregatorFunction.java new file mode 100644 index 0000000000000..4e06158952fc3 --- /dev/null +++ b/x-pack/plugin/esql/compute/src/main/generated/org/elasticsearch/compute/aggregation/spatial/SpatialExtentCartesianPointSourceValuesGroupingAggregatorFunction.java @@ -0,0 +1,235 @@ +// Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one +// or more contributor license agreements. Licensed under the Elastic License +// 2.0; you may not use this file except in compliance with the Elastic License +// 2.0. +package org.elasticsearch.compute.aggregation.spatial; + +import java.lang.Integer; +import java.lang.Override; +import java.lang.String; +import java.lang.StringBuilder; +import java.util.List; +import org.apache.lucene.util.BytesRef; +import org.elasticsearch.compute.aggregation.GroupingAggregatorFunction; +import org.elasticsearch.compute.aggregation.IntermediateStateDesc; +import org.elasticsearch.compute.aggregation.SeenGroupIds; +import org.elasticsearch.compute.data.Block; +import org.elasticsearch.compute.data.BytesRefBlock; +import org.elasticsearch.compute.data.BytesRefVector; +import org.elasticsearch.compute.data.ElementType; +import org.elasticsearch.compute.data.IntBlock; +import org.elasticsearch.compute.data.IntVector; +import org.elasticsearch.compute.data.Page; +import org.elasticsearch.compute.operator.DriverContext; + +/** + * {@link GroupingAggregatorFunction} implementation for {@link SpatialExtentCartesianPointSourceValuesAggregator}. + * This class is generated. Do not edit it. + */ +public final class SpatialExtentCartesianPointSourceValuesGroupingAggregatorFunction implements GroupingAggregatorFunction { + private static final List INTERMEDIATE_STATE_DESC = List.of( + new IntermediateStateDesc("minX", ElementType.INT), + new IntermediateStateDesc("maxX", ElementType.INT), + new IntermediateStateDesc("maxY", ElementType.INT), + new IntermediateStateDesc("minY", ElementType.INT) ); + + private final SpatialExtentGroupingState state; + + private final List channels; + + private final DriverContext driverContext; + + public SpatialExtentCartesianPointSourceValuesGroupingAggregatorFunction(List channels, + SpatialExtentGroupingState state, DriverContext driverContext) { + this.channels = channels; + this.state = state; + this.driverContext = driverContext; + } + + public static SpatialExtentCartesianPointSourceValuesGroupingAggregatorFunction create( + List channels, DriverContext driverContext) { + return new SpatialExtentCartesianPointSourceValuesGroupingAggregatorFunction(channels, SpatialExtentCartesianPointSourceValuesAggregator.initGrouping(), driverContext); + } + + public static List intermediateStateDesc() { + return INTERMEDIATE_STATE_DESC; + } + + @Override + public int intermediateBlockCount() { + return INTERMEDIATE_STATE_DESC.size(); + } + + @Override + public GroupingAggregatorFunction.AddInput prepareProcessPage(SeenGroupIds seenGroupIds, + Page page) { + BytesRefBlock valuesBlock = page.getBlock(channels.get(0)); + BytesRefVector valuesVector = valuesBlock.asVector(); + if (valuesVector == null) { + if (valuesBlock.mayHaveNulls()) { + state.enableGroupIdTracking(seenGroupIds); + } + return new GroupingAggregatorFunction.AddInput() { + @Override + public void add(int positionOffset, IntBlock groupIds) { + addRawInput(positionOffset, groupIds, valuesBlock); + } + + @Override + public void add(int positionOffset, IntVector groupIds) { + addRawInput(positionOffset, groupIds, valuesBlock); + } + + @Override + public void close() { + } + }; + } + return new GroupingAggregatorFunction.AddInput() { + @Override + public void add(int positionOffset, IntBlock groupIds) { + addRawInput(positionOffset, groupIds, valuesVector); + } + + @Override + public void add(int positionOffset, IntVector groupIds) { + addRawInput(positionOffset, groupIds, valuesVector); + } + + @Override + public void close() { + } + }; + } + + private void addRawInput(int positionOffset, IntVector groups, BytesRefBlock values) { + BytesRef scratch = new BytesRef(); + for (int groupPosition = 0; groupPosition < groups.getPositionCount(); groupPosition++) { + int groupId = groups.getInt(groupPosition); + if (values.isNull(groupPosition + positionOffset)) { + continue; + } + int valuesStart = values.getFirstValueIndex(groupPosition + positionOffset); + int valuesEnd = valuesStart + values.getValueCount(groupPosition + positionOffset); + for (int v = valuesStart; v < valuesEnd; v++) { + SpatialExtentCartesianPointSourceValuesAggregator.combine(state, groupId, values.getBytesRef(v, scratch)); + } + } + } + + private void addRawInput(int positionOffset, IntVector groups, BytesRefVector values) { + BytesRef scratch = new BytesRef(); + for (int groupPosition = 0; groupPosition < groups.getPositionCount(); groupPosition++) { + int groupId = groups.getInt(groupPosition); + SpatialExtentCartesianPointSourceValuesAggregator.combine(state, groupId, values.getBytesRef(groupPosition + positionOffset, scratch)); + } + } + + private void addRawInput(int positionOffset, IntBlock groups, BytesRefBlock values) { + BytesRef scratch = new BytesRef(); + for (int groupPosition = 0; groupPosition < groups.getPositionCount(); groupPosition++) { + if (groups.isNull(groupPosition)) { + continue; + } + int groupStart = groups.getFirstValueIndex(groupPosition); + int groupEnd = groupStart + groups.getValueCount(groupPosition); + for (int g = groupStart; g < groupEnd; g++) { + int groupId = groups.getInt(g); + if (values.isNull(groupPosition + positionOffset)) { + continue; + } + int valuesStart = values.getFirstValueIndex(groupPosition + positionOffset); + int valuesEnd = valuesStart + values.getValueCount(groupPosition + positionOffset); + for (int v = valuesStart; v < valuesEnd; v++) { + SpatialExtentCartesianPointSourceValuesAggregator.combine(state, groupId, values.getBytesRef(v, scratch)); + } + } + } + } + + private void addRawInput(int positionOffset, IntBlock groups, BytesRefVector values) { + BytesRef scratch = new BytesRef(); + for (int groupPosition = 0; groupPosition < groups.getPositionCount(); groupPosition++) { + if (groups.isNull(groupPosition)) { + continue; + } + int groupStart = groups.getFirstValueIndex(groupPosition); + int groupEnd = groupStart + groups.getValueCount(groupPosition); + for (int g = groupStart; g < groupEnd; g++) { + int groupId = groups.getInt(g); + SpatialExtentCartesianPointSourceValuesAggregator.combine(state, groupId, values.getBytesRef(groupPosition + positionOffset, scratch)); + } + } + } + + @Override + public void selectedMayContainUnseenGroups(SeenGroupIds seenGroupIds) { + state.enableGroupIdTracking(seenGroupIds); + } + + @Override + public void addIntermediateInput(int positionOffset, IntVector groups, Page page) { + state.enableGroupIdTracking(new SeenGroupIds.Empty()); + assert channels.size() == intermediateBlockCount(); + Block minXUncast = page.getBlock(channels.get(0)); + if (minXUncast.areAllValuesNull()) { + return; + } + IntVector minX = ((IntBlock) minXUncast).asVector(); + Block maxXUncast = page.getBlock(channels.get(1)); + if (maxXUncast.areAllValuesNull()) { + return; + } + IntVector maxX = ((IntBlock) maxXUncast).asVector(); + Block maxYUncast = page.getBlock(channels.get(2)); + if (maxYUncast.areAllValuesNull()) { + return; + } + IntVector maxY = ((IntBlock) maxYUncast).asVector(); + Block minYUncast = page.getBlock(channels.get(3)); + if (minYUncast.areAllValuesNull()) { + return; + } + IntVector minY = ((IntBlock) minYUncast).asVector(); + assert minX.getPositionCount() == maxX.getPositionCount() && minX.getPositionCount() == maxY.getPositionCount() && minX.getPositionCount() == minY.getPositionCount(); + for (int groupPosition = 0; groupPosition < groups.getPositionCount(); groupPosition++) { + int groupId = groups.getInt(groupPosition); + SpatialExtentCartesianPointSourceValuesAggregator.combineIntermediate(state, groupId, minX.getInt(groupPosition + positionOffset), maxX.getInt(groupPosition + positionOffset), maxY.getInt(groupPosition + positionOffset), minY.getInt(groupPosition + positionOffset)); + } + } + + @Override + public void addIntermediateRowInput(int groupId, GroupingAggregatorFunction input, int position) { + if (input.getClass() != getClass()) { + throw new IllegalArgumentException("expected " + getClass() + "; got " + input.getClass()); + } + SpatialExtentGroupingState inState = ((SpatialExtentCartesianPointSourceValuesGroupingAggregatorFunction) input).state; + state.enableGroupIdTracking(new SeenGroupIds.Empty()); + SpatialExtentCartesianPointSourceValuesAggregator.combineStates(state, groupId, inState, position); + } + + @Override + public void evaluateIntermediate(Block[] blocks, int offset, IntVector selected) { + state.toIntermediate(blocks, offset, selected, driverContext); + } + + @Override + public void evaluateFinal(Block[] blocks, int offset, IntVector selected, + DriverContext driverContext) { + blocks[offset] = SpatialExtentCartesianPointSourceValuesAggregator.evaluateFinal(state, selected, driverContext); + } + + @Override + public String toString() { + StringBuilder sb = new StringBuilder(); + sb.append(getClass().getSimpleName()).append("["); + sb.append("channels=").append(channels); + sb.append("]"); + return sb.toString(); + } + + @Override + public void close() { + state.close(); + } +} diff --git a/x-pack/plugin/esql/compute/src/main/generated/org/elasticsearch/compute/aggregation/spatial/SpatialExtentCartesianShapeAggregatorFunction.java b/x-pack/plugin/esql/compute/src/main/generated/org/elasticsearch/compute/aggregation/spatial/SpatialExtentCartesianShapeAggregatorFunction.java new file mode 100644 index 0000000000000..19aa4f7ca78a2 --- /dev/null +++ b/x-pack/plugin/esql/compute/src/main/generated/org/elasticsearch/compute/aggregation/spatial/SpatialExtentCartesianShapeAggregatorFunction.java @@ -0,0 +1,192 @@ +// Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one +// or more contributor license agreements. Licensed under the Elastic License +// 2.0; you may not use this file except in compliance with the Elastic License +// 2.0. +package org.elasticsearch.compute.aggregation.spatial; + +import java.lang.Integer; +import java.lang.Override; +import java.lang.String; +import java.lang.StringBuilder; +import java.util.List; +import org.apache.lucene.util.BytesRef; +import org.elasticsearch.compute.aggregation.AggregatorFunction; +import org.elasticsearch.compute.aggregation.IntermediateStateDesc; +import org.elasticsearch.compute.data.Block; +import org.elasticsearch.compute.data.BooleanVector; +import org.elasticsearch.compute.data.BytesRefBlock; +import org.elasticsearch.compute.data.BytesRefVector; +import org.elasticsearch.compute.data.ElementType; +import org.elasticsearch.compute.data.IntBlock; +import org.elasticsearch.compute.data.IntVector; +import org.elasticsearch.compute.data.Page; +import org.elasticsearch.compute.operator.DriverContext; + +/** + * {@link AggregatorFunction} implementation for {@link SpatialExtentCartesianShapeAggregator}. + * This class is generated. Do not edit it. + */ +public final class SpatialExtentCartesianShapeAggregatorFunction implements AggregatorFunction { + private static final List INTERMEDIATE_STATE_DESC = List.of( + new IntermediateStateDesc("minX", ElementType.INT), + new IntermediateStateDesc("maxX", ElementType.INT), + new IntermediateStateDesc("maxY", ElementType.INT), + new IntermediateStateDesc("minY", ElementType.INT) ); + + private final DriverContext driverContext; + + private final SpatialExtentState state; + + private final List channels; + + public SpatialExtentCartesianShapeAggregatorFunction(DriverContext driverContext, + List channels, SpatialExtentState state) { + this.driverContext = driverContext; + this.channels = channels; + this.state = state; + } + + public static SpatialExtentCartesianShapeAggregatorFunction create(DriverContext driverContext, + List channels) { + return new SpatialExtentCartesianShapeAggregatorFunction(driverContext, channels, SpatialExtentCartesianShapeAggregator.initSingle()); + } + + public static List intermediateStateDesc() { + return INTERMEDIATE_STATE_DESC; + } + + @Override + public int intermediateBlockCount() { + return INTERMEDIATE_STATE_DESC.size(); + } + + @Override + public void addRawInput(Page page, BooleanVector mask) { + if (mask.allFalse()) { + // Entire page masked away + return; + } + if (mask.allTrue()) { + // No masking + BytesRefBlock block = page.getBlock(channels.get(0)); + BytesRefVector vector = block.asVector(); + if (vector != null) { + addRawVector(vector); + } else { + addRawBlock(block); + } + return; + } + // Some positions masked away, others kept + BytesRefBlock block = page.getBlock(channels.get(0)); + BytesRefVector vector = block.asVector(); + if (vector != null) { + addRawVector(vector, mask); + } else { + addRawBlock(block, mask); + } + } + + private void addRawVector(BytesRefVector vector) { + BytesRef scratch = new BytesRef(); + for (int i = 0; i < vector.getPositionCount(); i++) { + SpatialExtentCartesianShapeAggregator.combine(state, vector.getBytesRef(i, scratch)); + } + } + + private void addRawVector(BytesRefVector vector, BooleanVector mask) { + BytesRef scratch = new BytesRef(); + for (int i = 0; i < vector.getPositionCount(); i++) { + if (mask.getBoolean(i) == false) { + continue; + } + SpatialExtentCartesianShapeAggregator.combine(state, vector.getBytesRef(i, scratch)); + } + } + + private void addRawBlock(BytesRefBlock block) { + BytesRef scratch = new BytesRef(); + for (int p = 0; p < block.getPositionCount(); p++) { + if (block.isNull(p)) { + continue; + } + int start = block.getFirstValueIndex(p); + int end = start + block.getValueCount(p); + for (int i = start; i < end; i++) { + SpatialExtentCartesianShapeAggregator.combine(state, block.getBytesRef(i, scratch)); + } + } + } + + private void addRawBlock(BytesRefBlock block, BooleanVector mask) { + BytesRef scratch = new BytesRef(); + for (int p = 0; p < block.getPositionCount(); p++) { + if (mask.getBoolean(p) == false) { + continue; + } + if (block.isNull(p)) { + continue; + } + int start = block.getFirstValueIndex(p); + int end = start + block.getValueCount(p); + for (int i = start; i < end; i++) { + SpatialExtentCartesianShapeAggregator.combine(state, block.getBytesRef(i, scratch)); + } + } + } + + @Override + public void addIntermediateInput(Page page) { + assert channels.size() == intermediateBlockCount(); + assert page.getBlockCount() >= channels.get(0) + intermediateStateDesc().size(); + Block minXUncast = page.getBlock(channels.get(0)); + if (minXUncast.areAllValuesNull()) { + return; + } + IntVector minX = ((IntBlock) minXUncast).asVector(); + assert minX.getPositionCount() == 1; + Block maxXUncast = page.getBlock(channels.get(1)); + if (maxXUncast.areAllValuesNull()) { + return; + } + IntVector maxX = ((IntBlock) maxXUncast).asVector(); + assert maxX.getPositionCount() == 1; + Block maxYUncast = page.getBlock(channels.get(2)); + if (maxYUncast.areAllValuesNull()) { + return; + } + IntVector maxY = ((IntBlock) maxYUncast).asVector(); + assert maxY.getPositionCount() == 1; + Block minYUncast = page.getBlock(channels.get(3)); + if (minYUncast.areAllValuesNull()) { + return; + } + IntVector minY = ((IntBlock) minYUncast).asVector(); + assert minY.getPositionCount() == 1; + SpatialExtentCartesianShapeAggregator.combineIntermediate(state, minX.getInt(0), maxX.getInt(0), maxY.getInt(0), minY.getInt(0)); + } + + @Override + public void evaluateIntermediate(Block[] blocks, int offset, DriverContext driverContext) { + state.toIntermediate(blocks, offset, driverContext); + } + + @Override + public void evaluateFinal(Block[] blocks, int offset, DriverContext driverContext) { + blocks[offset] = SpatialExtentCartesianShapeAggregator.evaluateFinal(state, driverContext); + } + + @Override + public String toString() { + StringBuilder sb = new StringBuilder(); + sb.append(getClass().getSimpleName()).append("["); + sb.append("channels=").append(channels); + sb.append("]"); + return sb.toString(); + } + + @Override + public void close() { + state.close(); + } +} diff --git a/x-pack/plugin/esql/compute/src/main/generated/org/elasticsearch/compute/aggregation/spatial/SpatialExtentCartesianShapeAggregatorFunctionSupplier.java b/x-pack/plugin/esql/compute/src/main/generated/org/elasticsearch/compute/aggregation/spatial/SpatialExtentCartesianShapeAggregatorFunctionSupplier.java new file mode 100644 index 0000000000000..9e4b292a0ea29 --- /dev/null +++ b/x-pack/plugin/esql/compute/src/main/generated/org/elasticsearch/compute/aggregation/spatial/SpatialExtentCartesianShapeAggregatorFunctionSupplier.java @@ -0,0 +1,40 @@ +// Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one +// or more contributor license agreements. Licensed under the Elastic License +// 2.0; you may not use this file except in compliance with the Elastic License +// 2.0. +package org.elasticsearch.compute.aggregation.spatial; + +import java.lang.Integer; +import java.lang.Override; +import java.lang.String; +import java.util.List; +import org.elasticsearch.compute.aggregation.AggregatorFunctionSupplier; +import org.elasticsearch.compute.operator.DriverContext; + +/** + * {@link AggregatorFunctionSupplier} implementation for {@link SpatialExtentCartesianShapeAggregator}. + * This class is generated. Do not edit it. + */ +public final class SpatialExtentCartesianShapeAggregatorFunctionSupplier implements AggregatorFunctionSupplier { + private final List channels; + + public SpatialExtentCartesianShapeAggregatorFunctionSupplier(List channels) { + this.channels = channels; + } + + @Override + public SpatialExtentCartesianShapeAggregatorFunction aggregator(DriverContext driverContext) { + return SpatialExtentCartesianShapeAggregatorFunction.create(driverContext, channels); + } + + @Override + public SpatialExtentCartesianShapeGroupingAggregatorFunction groupingAggregator( + DriverContext driverContext) { + return SpatialExtentCartesianShapeGroupingAggregatorFunction.create(channels, driverContext); + } + + @Override + public String describe() { + return "spatial_extent_cartesian of shapes"; + } +} diff --git a/x-pack/plugin/esql/compute/src/main/generated/org/elasticsearch/compute/aggregation/spatial/SpatialExtentCartesianShapeGroupingAggregatorFunction.java b/x-pack/plugin/esql/compute/src/main/generated/org/elasticsearch/compute/aggregation/spatial/SpatialExtentCartesianShapeGroupingAggregatorFunction.java new file mode 100644 index 0000000000000..c55c3d9c66946 --- /dev/null +++ b/x-pack/plugin/esql/compute/src/main/generated/org/elasticsearch/compute/aggregation/spatial/SpatialExtentCartesianShapeGroupingAggregatorFunction.java @@ -0,0 +1,235 @@ +// Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one +// or more contributor license agreements. Licensed under the Elastic License +// 2.0; you may not use this file except in compliance with the Elastic License +// 2.0. +package org.elasticsearch.compute.aggregation.spatial; + +import java.lang.Integer; +import java.lang.Override; +import java.lang.String; +import java.lang.StringBuilder; +import java.util.List; +import org.apache.lucene.util.BytesRef; +import org.elasticsearch.compute.aggregation.GroupingAggregatorFunction; +import org.elasticsearch.compute.aggregation.IntermediateStateDesc; +import org.elasticsearch.compute.aggregation.SeenGroupIds; +import org.elasticsearch.compute.data.Block; +import org.elasticsearch.compute.data.BytesRefBlock; +import org.elasticsearch.compute.data.BytesRefVector; +import org.elasticsearch.compute.data.ElementType; +import org.elasticsearch.compute.data.IntBlock; +import org.elasticsearch.compute.data.IntVector; +import org.elasticsearch.compute.data.Page; +import org.elasticsearch.compute.operator.DriverContext; + +/** + * {@link GroupingAggregatorFunction} implementation for {@link SpatialExtentCartesianShapeAggregator}. + * This class is generated. Do not edit it. + */ +public final class SpatialExtentCartesianShapeGroupingAggregatorFunction implements GroupingAggregatorFunction { + private static final List INTERMEDIATE_STATE_DESC = List.of( + new IntermediateStateDesc("minX", ElementType.INT), + new IntermediateStateDesc("maxX", ElementType.INT), + new IntermediateStateDesc("maxY", ElementType.INT), + new IntermediateStateDesc("minY", ElementType.INT) ); + + private final SpatialExtentGroupingState state; + + private final List channels; + + private final DriverContext driverContext; + + public SpatialExtentCartesianShapeGroupingAggregatorFunction(List channels, + SpatialExtentGroupingState state, DriverContext driverContext) { + this.channels = channels; + this.state = state; + this.driverContext = driverContext; + } + + public static SpatialExtentCartesianShapeGroupingAggregatorFunction create(List channels, + DriverContext driverContext) { + return new SpatialExtentCartesianShapeGroupingAggregatorFunction(channels, SpatialExtentCartesianShapeAggregator.initGrouping(), driverContext); + } + + public static List intermediateStateDesc() { + return INTERMEDIATE_STATE_DESC; + } + + @Override + public int intermediateBlockCount() { + return INTERMEDIATE_STATE_DESC.size(); + } + + @Override + public GroupingAggregatorFunction.AddInput prepareProcessPage(SeenGroupIds seenGroupIds, + Page page) { + BytesRefBlock valuesBlock = page.getBlock(channels.get(0)); + BytesRefVector valuesVector = valuesBlock.asVector(); + if (valuesVector == null) { + if (valuesBlock.mayHaveNulls()) { + state.enableGroupIdTracking(seenGroupIds); + } + return new GroupingAggregatorFunction.AddInput() { + @Override + public void add(int positionOffset, IntBlock groupIds) { + addRawInput(positionOffset, groupIds, valuesBlock); + } + + @Override + public void add(int positionOffset, IntVector groupIds) { + addRawInput(positionOffset, groupIds, valuesBlock); + } + + @Override + public void close() { + } + }; + } + return new GroupingAggregatorFunction.AddInput() { + @Override + public void add(int positionOffset, IntBlock groupIds) { + addRawInput(positionOffset, groupIds, valuesVector); + } + + @Override + public void add(int positionOffset, IntVector groupIds) { + addRawInput(positionOffset, groupIds, valuesVector); + } + + @Override + public void close() { + } + }; + } + + private void addRawInput(int positionOffset, IntVector groups, BytesRefBlock values) { + BytesRef scratch = new BytesRef(); + for (int groupPosition = 0; groupPosition < groups.getPositionCount(); groupPosition++) { + int groupId = groups.getInt(groupPosition); + if (values.isNull(groupPosition + positionOffset)) { + continue; + } + int valuesStart = values.getFirstValueIndex(groupPosition + positionOffset); + int valuesEnd = valuesStart + values.getValueCount(groupPosition + positionOffset); + for (int v = valuesStart; v < valuesEnd; v++) { + SpatialExtentCartesianShapeAggregator.combine(state, groupId, values.getBytesRef(v, scratch)); + } + } + } + + private void addRawInput(int positionOffset, IntVector groups, BytesRefVector values) { + BytesRef scratch = new BytesRef(); + for (int groupPosition = 0; groupPosition < groups.getPositionCount(); groupPosition++) { + int groupId = groups.getInt(groupPosition); + SpatialExtentCartesianShapeAggregator.combine(state, groupId, values.getBytesRef(groupPosition + positionOffset, scratch)); + } + } + + private void addRawInput(int positionOffset, IntBlock groups, BytesRefBlock values) { + BytesRef scratch = new BytesRef(); + for (int groupPosition = 0; groupPosition < groups.getPositionCount(); groupPosition++) { + if (groups.isNull(groupPosition)) { + continue; + } + int groupStart = groups.getFirstValueIndex(groupPosition); + int groupEnd = groupStart + groups.getValueCount(groupPosition); + for (int g = groupStart; g < groupEnd; g++) { + int groupId = groups.getInt(g); + if (values.isNull(groupPosition + positionOffset)) { + continue; + } + int valuesStart = values.getFirstValueIndex(groupPosition + positionOffset); + int valuesEnd = valuesStart + values.getValueCount(groupPosition + positionOffset); + for (int v = valuesStart; v < valuesEnd; v++) { + SpatialExtentCartesianShapeAggregator.combine(state, groupId, values.getBytesRef(v, scratch)); + } + } + } + } + + private void addRawInput(int positionOffset, IntBlock groups, BytesRefVector values) { + BytesRef scratch = new BytesRef(); + for (int groupPosition = 0; groupPosition < groups.getPositionCount(); groupPosition++) { + if (groups.isNull(groupPosition)) { + continue; + } + int groupStart = groups.getFirstValueIndex(groupPosition); + int groupEnd = groupStart + groups.getValueCount(groupPosition); + for (int g = groupStart; g < groupEnd; g++) { + int groupId = groups.getInt(g); + SpatialExtentCartesianShapeAggregator.combine(state, groupId, values.getBytesRef(groupPosition + positionOffset, scratch)); + } + } + } + + @Override + public void selectedMayContainUnseenGroups(SeenGroupIds seenGroupIds) { + state.enableGroupIdTracking(seenGroupIds); + } + + @Override + public void addIntermediateInput(int positionOffset, IntVector groups, Page page) { + state.enableGroupIdTracking(new SeenGroupIds.Empty()); + assert channels.size() == intermediateBlockCount(); + Block minXUncast = page.getBlock(channels.get(0)); + if (minXUncast.areAllValuesNull()) { + return; + } + IntVector minX = ((IntBlock) minXUncast).asVector(); + Block maxXUncast = page.getBlock(channels.get(1)); + if (maxXUncast.areAllValuesNull()) { + return; + } + IntVector maxX = ((IntBlock) maxXUncast).asVector(); + Block maxYUncast = page.getBlock(channels.get(2)); + if (maxYUncast.areAllValuesNull()) { + return; + } + IntVector maxY = ((IntBlock) maxYUncast).asVector(); + Block minYUncast = page.getBlock(channels.get(3)); + if (minYUncast.areAllValuesNull()) { + return; + } + IntVector minY = ((IntBlock) minYUncast).asVector(); + assert minX.getPositionCount() == maxX.getPositionCount() && minX.getPositionCount() == maxY.getPositionCount() && minX.getPositionCount() == minY.getPositionCount(); + for (int groupPosition = 0; groupPosition < groups.getPositionCount(); groupPosition++) { + int groupId = groups.getInt(groupPosition); + SpatialExtentCartesianShapeAggregator.combineIntermediate(state, groupId, minX.getInt(groupPosition + positionOffset), maxX.getInt(groupPosition + positionOffset), maxY.getInt(groupPosition + positionOffset), minY.getInt(groupPosition + positionOffset)); + } + } + + @Override + public void addIntermediateRowInput(int groupId, GroupingAggregatorFunction input, int position) { + if (input.getClass() != getClass()) { + throw new IllegalArgumentException("expected " + getClass() + "; got " + input.getClass()); + } + SpatialExtentGroupingState inState = ((SpatialExtentCartesianShapeGroupingAggregatorFunction) input).state; + state.enableGroupIdTracking(new SeenGroupIds.Empty()); + SpatialExtentCartesianShapeAggregator.combineStates(state, groupId, inState, position); + } + + @Override + public void evaluateIntermediate(Block[] blocks, int offset, IntVector selected) { + state.toIntermediate(blocks, offset, selected, driverContext); + } + + @Override + public void evaluateFinal(Block[] blocks, int offset, IntVector selected, + DriverContext driverContext) { + blocks[offset] = SpatialExtentCartesianShapeAggregator.evaluateFinal(state, selected, driverContext); + } + + @Override + public String toString() { + StringBuilder sb = new StringBuilder(); + sb.append(getClass().getSimpleName()).append("["); + sb.append("channels=").append(channels); + sb.append("]"); + return sb.toString(); + } + + @Override + public void close() { + state.close(); + } +} diff --git a/x-pack/plugin/esql/compute/src/main/generated/org/elasticsearch/compute/aggregation/spatial/SpatialExtentGeoPointDocValuesAggregatorFunction.java b/x-pack/plugin/esql/compute/src/main/generated/org/elasticsearch/compute/aggregation/spatial/SpatialExtentGeoPointDocValuesAggregatorFunction.java new file mode 100644 index 0000000000000..c883e82d45989 --- /dev/null +++ b/x-pack/plugin/esql/compute/src/main/generated/org/elasticsearch/compute/aggregation/spatial/SpatialExtentGeoPointDocValuesAggregatorFunction.java @@ -0,0 +1,201 @@ +// Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one +// or more contributor license agreements. Licensed under the Elastic License +// 2.0; you may not use this file except in compliance with the Elastic License +// 2.0. +package org.elasticsearch.compute.aggregation.spatial; + +import java.lang.Integer; +import java.lang.Override; +import java.lang.String; +import java.lang.StringBuilder; +import java.util.List; +import org.elasticsearch.compute.aggregation.AggregatorFunction; +import org.elasticsearch.compute.aggregation.IntermediateStateDesc; +import org.elasticsearch.compute.data.Block; +import org.elasticsearch.compute.data.BooleanVector; +import org.elasticsearch.compute.data.ElementType; +import org.elasticsearch.compute.data.IntBlock; +import org.elasticsearch.compute.data.IntVector; +import org.elasticsearch.compute.data.LongBlock; +import org.elasticsearch.compute.data.LongVector; +import org.elasticsearch.compute.data.Page; +import org.elasticsearch.compute.operator.DriverContext; + +/** + * {@link AggregatorFunction} implementation for {@link SpatialExtentGeoPointDocValuesAggregator}. + * This class is generated. Do not edit it. + */ +public final class SpatialExtentGeoPointDocValuesAggregatorFunction implements AggregatorFunction { + private static final List INTERMEDIATE_STATE_DESC = List.of( + new IntermediateStateDesc("minNegX", ElementType.INT), + new IntermediateStateDesc("minPosX", ElementType.INT), + new IntermediateStateDesc("maxNegX", ElementType.INT), + new IntermediateStateDesc("maxPosX", ElementType.INT), + new IntermediateStateDesc("maxY", ElementType.INT), + new IntermediateStateDesc("minY", ElementType.INT) ); + + private final DriverContext driverContext; + + private final SpatialExtentStateWrappedLongitudeState state; + + private final List channels; + + public SpatialExtentGeoPointDocValuesAggregatorFunction(DriverContext driverContext, + List channels, SpatialExtentStateWrappedLongitudeState state) { + this.driverContext = driverContext; + this.channels = channels; + this.state = state; + } + + public static SpatialExtentGeoPointDocValuesAggregatorFunction create(DriverContext driverContext, + List channels) { + return new SpatialExtentGeoPointDocValuesAggregatorFunction(driverContext, channels, SpatialExtentGeoPointDocValuesAggregator.initSingle()); + } + + public static List intermediateStateDesc() { + return INTERMEDIATE_STATE_DESC; + } + + @Override + public int intermediateBlockCount() { + return INTERMEDIATE_STATE_DESC.size(); + } + + @Override + public void addRawInput(Page page, BooleanVector mask) { + if (mask.allFalse()) { + // Entire page masked away + return; + } + if (mask.allTrue()) { + // No masking + LongBlock block = page.getBlock(channels.get(0)); + LongVector vector = block.asVector(); + if (vector != null) { + addRawVector(vector); + } else { + addRawBlock(block); + } + return; + } + // Some positions masked away, others kept + LongBlock block = page.getBlock(channels.get(0)); + LongVector vector = block.asVector(); + if (vector != null) { + addRawVector(vector, mask); + } else { + addRawBlock(block, mask); + } + } + + private void addRawVector(LongVector vector) { + for (int i = 0; i < vector.getPositionCount(); i++) { + SpatialExtentGeoPointDocValuesAggregator.combine(state, vector.getLong(i)); + } + } + + private void addRawVector(LongVector vector, BooleanVector mask) { + for (int i = 0; i < vector.getPositionCount(); i++) { + if (mask.getBoolean(i) == false) { + continue; + } + SpatialExtentGeoPointDocValuesAggregator.combine(state, vector.getLong(i)); + } + } + + private void addRawBlock(LongBlock block) { + for (int p = 0; p < block.getPositionCount(); p++) { + if (block.isNull(p)) { + continue; + } + int start = block.getFirstValueIndex(p); + int end = start + block.getValueCount(p); + for (int i = start; i < end; i++) { + SpatialExtentGeoPointDocValuesAggregator.combine(state, block.getLong(i)); + } + } + } + + private void addRawBlock(LongBlock block, BooleanVector mask) { + for (int p = 0; p < block.getPositionCount(); p++) { + if (mask.getBoolean(p) == false) { + continue; + } + if (block.isNull(p)) { + continue; + } + int start = block.getFirstValueIndex(p); + int end = start + block.getValueCount(p); + for (int i = start; i < end; i++) { + SpatialExtentGeoPointDocValuesAggregator.combine(state, block.getLong(i)); + } + } + } + + @Override + public void addIntermediateInput(Page page) { + assert channels.size() == intermediateBlockCount(); + assert page.getBlockCount() >= channels.get(0) + intermediateStateDesc().size(); + Block minNegXUncast = page.getBlock(channels.get(0)); + if (minNegXUncast.areAllValuesNull()) { + return; + } + IntVector minNegX = ((IntBlock) minNegXUncast).asVector(); + assert minNegX.getPositionCount() == 1; + Block minPosXUncast = page.getBlock(channels.get(1)); + if (minPosXUncast.areAllValuesNull()) { + return; + } + IntVector minPosX = ((IntBlock) minPosXUncast).asVector(); + assert minPosX.getPositionCount() == 1; + Block maxNegXUncast = page.getBlock(channels.get(2)); + if (maxNegXUncast.areAllValuesNull()) { + return; + } + IntVector maxNegX = ((IntBlock) maxNegXUncast).asVector(); + assert maxNegX.getPositionCount() == 1; + Block maxPosXUncast = page.getBlock(channels.get(3)); + if (maxPosXUncast.areAllValuesNull()) { + return; + } + IntVector maxPosX = ((IntBlock) maxPosXUncast).asVector(); + assert maxPosX.getPositionCount() == 1; + Block maxYUncast = page.getBlock(channels.get(4)); + if (maxYUncast.areAllValuesNull()) { + return; + } + IntVector maxY = ((IntBlock) maxYUncast).asVector(); + assert maxY.getPositionCount() == 1; + Block minYUncast = page.getBlock(channels.get(5)); + if (minYUncast.areAllValuesNull()) { + return; + } + IntVector minY = ((IntBlock) minYUncast).asVector(); + assert minY.getPositionCount() == 1; + SpatialExtentGeoPointDocValuesAggregator.combineIntermediate(state, minNegX.getInt(0), minPosX.getInt(0), maxNegX.getInt(0), maxPosX.getInt(0), maxY.getInt(0), minY.getInt(0)); + } + + @Override + public void evaluateIntermediate(Block[] blocks, int offset, DriverContext driverContext) { + state.toIntermediate(blocks, offset, driverContext); + } + + @Override + public void evaluateFinal(Block[] blocks, int offset, DriverContext driverContext) { + blocks[offset] = SpatialExtentGeoPointDocValuesAggregator.evaluateFinal(state, driverContext); + } + + @Override + public String toString() { + StringBuilder sb = new StringBuilder(); + sb.append(getClass().getSimpleName()).append("["); + sb.append("channels=").append(channels); + sb.append("]"); + return sb.toString(); + } + + @Override + public void close() { + state.close(); + } +} diff --git a/x-pack/plugin/esql/compute/src/main/generated/org/elasticsearch/compute/aggregation/spatial/SpatialExtentGeoPointDocValuesAggregatorFunctionSupplier.java b/x-pack/plugin/esql/compute/src/main/generated/org/elasticsearch/compute/aggregation/spatial/SpatialExtentGeoPointDocValuesAggregatorFunctionSupplier.java new file mode 100644 index 0000000000000..f72a4cc648ec8 --- /dev/null +++ b/x-pack/plugin/esql/compute/src/main/generated/org/elasticsearch/compute/aggregation/spatial/SpatialExtentGeoPointDocValuesAggregatorFunctionSupplier.java @@ -0,0 +1,40 @@ +// Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one +// or more contributor license agreements. Licensed under the Elastic License +// 2.0; you may not use this file except in compliance with the Elastic License +// 2.0. +package org.elasticsearch.compute.aggregation.spatial; + +import java.lang.Integer; +import java.lang.Override; +import java.lang.String; +import java.util.List; +import org.elasticsearch.compute.aggregation.AggregatorFunctionSupplier; +import org.elasticsearch.compute.operator.DriverContext; + +/** + * {@link AggregatorFunctionSupplier} implementation for {@link SpatialExtentGeoPointDocValuesAggregator}. + * This class is generated. Do not edit it. + */ +public final class SpatialExtentGeoPointDocValuesAggregatorFunctionSupplier implements AggregatorFunctionSupplier { + private final List channels; + + public SpatialExtentGeoPointDocValuesAggregatorFunctionSupplier(List channels) { + this.channels = channels; + } + + @Override + public SpatialExtentGeoPointDocValuesAggregatorFunction aggregator(DriverContext driverContext) { + return SpatialExtentGeoPointDocValuesAggregatorFunction.create(driverContext, channels); + } + + @Override + public SpatialExtentGeoPointDocValuesGroupingAggregatorFunction groupingAggregator( + DriverContext driverContext) { + return SpatialExtentGeoPointDocValuesGroupingAggregatorFunction.create(channels, driverContext); + } + + @Override + public String describe() { + return "spatial_extent_geo_point_doc of valuess"; + } +} diff --git a/x-pack/plugin/esql/compute/src/main/generated/org/elasticsearch/compute/aggregation/spatial/SpatialExtentGeoPointDocValuesGroupingAggregatorFunction.java b/x-pack/plugin/esql/compute/src/main/generated/org/elasticsearch/compute/aggregation/spatial/SpatialExtentGeoPointDocValuesGroupingAggregatorFunction.java new file mode 100644 index 0000000000000..eee5bc5df41a4 --- /dev/null +++ b/x-pack/plugin/esql/compute/src/main/generated/org/elasticsearch/compute/aggregation/spatial/SpatialExtentGeoPointDocValuesGroupingAggregatorFunction.java @@ -0,0 +1,242 @@ +// Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one +// or more contributor license agreements. Licensed under the Elastic License +// 2.0; you may not use this file except in compliance with the Elastic License +// 2.0. +package org.elasticsearch.compute.aggregation.spatial; + +import java.lang.Integer; +import java.lang.Override; +import java.lang.String; +import java.lang.StringBuilder; +import java.util.List; +import org.elasticsearch.compute.aggregation.GroupingAggregatorFunction; +import org.elasticsearch.compute.aggregation.IntermediateStateDesc; +import org.elasticsearch.compute.aggregation.SeenGroupIds; +import org.elasticsearch.compute.data.Block; +import org.elasticsearch.compute.data.ElementType; +import org.elasticsearch.compute.data.IntBlock; +import org.elasticsearch.compute.data.IntVector; +import org.elasticsearch.compute.data.LongBlock; +import org.elasticsearch.compute.data.LongVector; +import org.elasticsearch.compute.data.Page; +import org.elasticsearch.compute.operator.DriverContext; + +/** + * {@link GroupingAggregatorFunction} implementation for {@link SpatialExtentGeoPointDocValuesAggregator}. + * This class is generated. Do not edit it. + */ +public final class SpatialExtentGeoPointDocValuesGroupingAggregatorFunction implements GroupingAggregatorFunction { + private static final List INTERMEDIATE_STATE_DESC = List.of( + new IntermediateStateDesc("minNegX", ElementType.INT), + new IntermediateStateDesc("minPosX", ElementType.INT), + new IntermediateStateDesc("maxNegX", ElementType.INT), + new IntermediateStateDesc("maxPosX", ElementType.INT), + new IntermediateStateDesc("maxY", ElementType.INT), + new IntermediateStateDesc("minY", ElementType.INT) ); + + private final SpatialExtentGroupingStateWrappedLongitudeState state; + + private final List channels; + + private final DriverContext driverContext; + + public SpatialExtentGeoPointDocValuesGroupingAggregatorFunction(List channels, + SpatialExtentGroupingStateWrappedLongitudeState state, DriverContext driverContext) { + this.channels = channels; + this.state = state; + this.driverContext = driverContext; + } + + public static SpatialExtentGeoPointDocValuesGroupingAggregatorFunction create( + List channels, DriverContext driverContext) { + return new SpatialExtentGeoPointDocValuesGroupingAggregatorFunction(channels, SpatialExtentGeoPointDocValuesAggregator.initGrouping(), driverContext); + } + + public static List intermediateStateDesc() { + return INTERMEDIATE_STATE_DESC; + } + + @Override + public int intermediateBlockCount() { + return INTERMEDIATE_STATE_DESC.size(); + } + + @Override + public GroupingAggregatorFunction.AddInput prepareProcessPage(SeenGroupIds seenGroupIds, + Page page) { + LongBlock valuesBlock = page.getBlock(channels.get(0)); + LongVector valuesVector = valuesBlock.asVector(); + if (valuesVector == null) { + if (valuesBlock.mayHaveNulls()) { + state.enableGroupIdTracking(seenGroupIds); + } + return new GroupingAggregatorFunction.AddInput() { + @Override + public void add(int positionOffset, IntBlock groupIds) { + addRawInput(positionOffset, groupIds, valuesBlock); + } + + @Override + public void add(int positionOffset, IntVector groupIds) { + addRawInput(positionOffset, groupIds, valuesBlock); + } + + @Override + public void close() { + } + }; + } + return new GroupingAggregatorFunction.AddInput() { + @Override + public void add(int positionOffset, IntBlock groupIds) { + addRawInput(positionOffset, groupIds, valuesVector); + } + + @Override + public void add(int positionOffset, IntVector groupIds) { + addRawInput(positionOffset, groupIds, valuesVector); + } + + @Override + public void close() { + } + }; + } + + private void addRawInput(int positionOffset, IntVector groups, LongBlock values) { + for (int groupPosition = 0; groupPosition < groups.getPositionCount(); groupPosition++) { + int groupId = groups.getInt(groupPosition); + if (values.isNull(groupPosition + positionOffset)) { + continue; + } + int valuesStart = values.getFirstValueIndex(groupPosition + positionOffset); + int valuesEnd = valuesStart + values.getValueCount(groupPosition + positionOffset); + for (int v = valuesStart; v < valuesEnd; v++) { + SpatialExtentGeoPointDocValuesAggregator.combine(state, groupId, values.getLong(v)); + } + } + } + + private void addRawInput(int positionOffset, IntVector groups, LongVector values) { + for (int groupPosition = 0; groupPosition < groups.getPositionCount(); groupPosition++) { + int groupId = groups.getInt(groupPosition); + SpatialExtentGeoPointDocValuesAggregator.combine(state, groupId, values.getLong(groupPosition + positionOffset)); + } + } + + private void addRawInput(int positionOffset, IntBlock groups, LongBlock values) { + for (int groupPosition = 0; groupPosition < groups.getPositionCount(); groupPosition++) { + if (groups.isNull(groupPosition)) { + continue; + } + int groupStart = groups.getFirstValueIndex(groupPosition); + int groupEnd = groupStart + groups.getValueCount(groupPosition); + for (int g = groupStart; g < groupEnd; g++) { + int groupId = groups.getInt(g); + if (values.isNull(groupPosition + positionOffset)) { + continue; + } + int valuesStart = values.getFirstValueIndex(groupPosition + positionOffset); + int valuesEnd = valuesStart + values.getValueCount(groupPosition + positionOffset); + for (int v = valuesStart; v < valuesEnd; v++) { + SpatialExtentGeoPointDocValuesAggregator.combine(state, groupId, values.getLong(v)); + } + } + } + } + + private void addRawInput(int positionOffset, IntBlock groups, LongVector values) { + for (int groupPosition = 0; groupPosition < groups.getPositionCount(); groupPosition++) { + if (groups.isNull(groupPosition)) { + continue; + } + int groupStart = groups.getFirstValueIndex(groupPosition); + int groupEnd = groupStart + groups.getValueCount(groupPosition); + for (int g = groupStart; g < groupEnd; g++) { + int groupId = groups.getInt(g); + SpatialExtentGeoPointDocValuesAggregator.combine(state, groupId, values.getLong(groupPosition + positionOffset)); + } + } + } + + @Override + public void selectedMayContainUnseenGroups(SeenGroupIds seenGroupIds) { + state.enableGroupIdTracking(seenGroupIds); + } + + @Override + public void addIntermediateInput(int positionOffset, IntVector groups, Page page) { + state.enableGroupIdTracking(new SeenGroupIds.Empty()); + assert channels.size() == intermediateBlockCount(); + Block minNegXUncast = page.getBlock(channels.get(0)); + if (minNegXUncast.areAllValuesNull()) { + return; + } + IntVector minNegX = ((IntBlock) minNegXUncast).asVector(); + Block minPosXUncast = page.getBlock(channels.get(1)); + if (minPosXUncast.areAllValuesNull()) { + return; + } + IntVector minPosX = ((IntBlock) minPosXUncast).asVector(); + Block maxNegXUncast = page.getBlock(channels.get(2)); + if (maxNegXUncast.areAllValuesNull()) { + return; + } + IntVector maxNegX = ((IntBlock) maxNegXUncast).asVector(); + Block maxPosXUncast = page.getBlock(channels.get(3)); + if (maxPosXUncast.areAllValuesNull()) { + return; + } + IntVector maxPosX = ((IntBlock) maxPosXUncast).asVector(); + Block maxYUncast = page.getBlock(channels.get(4)); + if (maxYUncast.areAllValuesNull()) { + return; + } + IntVector maxY = ((IntBlock) maxYUncast).asVector(); + Block minYUncast = page.getBlock(channels.get(5)); + if (minYUncast.areAllValuesNull()) { + return; + } + IntVector minY = ((IntBlock) minYUncast).asVector(); + assert minNegX.getPositionCount() == minPosX.getPositionCount() && minNegX.getPositionCount() == maxNegX.getPositionCount() && minNegX.getPositionCount() == maxPosX.getPositionCount() && minNegX.getPositionCount() == maxY.getPositionCount() && minNegX.getPositionCount() == minY.getPositionCount(); + for (int groupPosition = 0; groupPosition < groups.getPositionCount(); groupPosition++) { + int groupId = groups.getInt(groupPosition); + SpatialExtentGeoPointDocValuesAggregator.combineIntermediate(state, groupId, minNegX.getInt(groupPosition + positionOffset), minPosX.getInt(groupPosition + positionOffset), maxNegX.getInt(groupPosition + positionOffset), maxPosX.getInt(groupPosition + positionOffset), maxY.getInt(groupPosition + positionOffset), minY.getInt(groupPosition + positionOffset)); + } + } + + @Override + public void addIntermediateRowInput(int groupId, GroupingAggregatorFunction input, int position) { + if (input.getClass() != getClass()) { + throw new IllegalArgumentException("expected " + getClass() + "; got " + input.getClass()); + } + SpatialExtentGroupingStateWrappedLongitudeState inState = ((SpatialExtentGeoPointDocValuesGroupingAggregatorFunction) input).state; + state.enableGroupIdTracking(new SeenGroupIds.Empty()); + SpatialExtentGeoPointDocValuesAggregator.combineStates(state, groupId, inState, position); + } + + @Override + public void evaluateIntermediate(Block[] blocks, int offset, IntVector selected) { + state.toIntermediate(blocks, offset, selected, driverContext); + } + + @Override + public void evaluateFinal(Block[] blocks, int offset, IntVector selected, + DriverContext driverContext) { + blocks[offset] = SpatialExtentGeoPointDocValuesAggregator.evaluateFinal(state, selected, driverContext); + } + + @Override + public String toString() { + StringBuilder sb = new StringBuilder(); + sb.append(getClass().getSimpleName()).append("["); + sb.append("channels=").append(channels); + sb.append("]"); + return sb.toString(); + } + + @Override + public void close() { + state.close(); + } +} diff --git a/x-pack/plugin/esql/compute/src/main/generated/org/elasticsearch/compute/aggregation/spatial/SpatialExtentGeoPointSourceValuesAggregatorFunction.java b/x-pack/plugin/esql/compute/src/main/generated/org/elasticsearch/compute/aggregation/spatial/SpatialExtentGeoPointSourceValuesAggregatorFunction.java new file mode 100644 index 0000000000000..cf65fbdde594c --- /dev/null +++ b/x-pack/plugin/esql/compute/src/main/generated/org/elasticsearch/compute/aggregation/spatial/SpatialExtentGeoPointSourceValuesAggregatorFunction.java @@ -0,0 +1,206 @@ +// Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one +// or more contributor license agreements. Licensed under the Elastic License +// 2.0; you may not use this file except in compliance with the Elastic License +// 2.0. +package org.elasticsearch.compute.aggregation.spatial; + +import java.lang.Integer; +import java.lang.Override; +import java.lang.String; +import java.lang.StringBuilder; +import java.util.List; +import org.apache.lucene.util.BytesRef; +import org.elasticsearch.compute.aggregation.AggregatorFunction; +import org.elasticsearch.compute.aggregation.IntermediateStateDesc; +import org.elasticsearch.compute.data.Block; +import org.elasticsearch.compute.data.BooleanVector; +import org.elasticsearch.compute.data.BytesRefBlock; +import org.elasticsearch.compute.data.BytesRefVector; +import org.elasticsearch.compute.data.ElementType; +import org.elasticsearch.compute.data.IntBlock; +import org.elasticsearch.compute.data.IntVector; +import org.elasticsearch.compute.data.Page; +import org.elasticsearch.compute.operator.DriverContext; + +/** + * {@link AggregatorFunction} implementation for {@link SpatialExtentGeoPointSourceValuesAggregator}. + * This class is generated. Do not edit it. + */ +public final class SpatialExtentGeoPointSourceValuesAggregatorFunction implements AggregatorFunction { + private static final List INTERMEDIATE_STATE_DESC = List.of( + new IntermediateStateDesc("minNegX", ElementType.INT), + new IntermediateStateDesc("minPosX", ElementType.INT), + new IntermediateStateDesc("maxNegX", ElementType.INT), + new IntermediateStateDesc("maxPosX", ElementType.INT), + new IntermediateStateDesc("maxY", ElementType.INT), + new IntermediateStateDesc("minY", ElementType.INT) ); + + private final DriverContext driverContext; + + private final SpatialExtentStateWrappedLongitudeState state; + + private final List channels; + + public SpatialExtentGeoPointSourceValuesAggregatorFunction(DriverContext driverContext, + List channels, SpatialExtentStateWrappedLongitudeState state) { + this.driverContext = driverContext; + this.channels = channels; + this.state = state; + } + + public static SpatialExtentGeoPointSourceValuesAggregatorFunction create( + DriverContext driverContext, List channels) { + return new SpatialExtentGeoPointSourceValuesAggregatorFunction(driverContext, channels, SpatialExtentGeoPointSourceValuesAggregator.initSingle()); + } + + public static List intermediateStateDesc() { + return INTERMEDIATE_STATE_DESC; + } + + @Override + public int intermediateBlockCount() { + return INTERMEDIATE_STATE_DESC.size(); + } + + @Override + public void addRawInput(Page page, BooleanVector mask) { + if (mask.allFalse()) { + // Entire page masked away + return; + } + if (mask.allTrue()) { + // No masking + BytesRefBlock block = page.getBlock(channels.get(0)); + BytesRefVector vector = block.asVector(); + if (vector != null) { + addRawVector(vector); + } else { + addRawBlock(block); + } + return; + } + // Some positions masked away, others kept + BytesRefBlock block = page.getBlock(channels.get(0)); + BytesRefVector vector = block.asVector(); + if (vector != null) { + addRawVector(vector, mask); + } else { + addRawBlock(block, mask); + } + } + + private void addRawVector(BytesRefVector vector) { + BytesRef scratch = new BytesRef(); + for (int i = 0; i < vector.getPositionCount(); i++) { + SpatialExtentGeoPointSourceValuesAggregator.combine(state, vector.getBytesRef(i, scratch)); + } + } + + private void addRawVector(BytesRefVector vector, BooleanVector mask) { + BytesRef scratch = new BytesRef(); + for (int i = 0; i < vector.getPositionCount(); i++) { + if (mask.getBoolean(i) == false) { + continue; + } + SpatialExtentGeoPointSourceValuesAggregator.combine(state, vector.getBytesRef(i, scratch)); + } + } + + private void addRawBlock(BytesRefBlock block) { + BytesRef scratch = new BytesRef(); + for (int p = 0; p < block.getPositionCount(); p++) { + if (block.isNull(p)) { + continue; + } + int start = block.getFirstValueIndex(p); + int end = start + block.getValueCount(p); + for (int i = start; i < end; i++) { + SpatialExtentGeoPointSourceValuesAggregator.combine(state, block.getBytesRef(i, scratch)); + } + } + } + + private void addRawBlock(BytesRefBlock block, BooleanVector mask) { + BytesRef scratch = new BytesRef(); + for (int p = 0; p < block.getPositionCount(); p++) { + if (mask.getBoolean(p) == false) { + continue; + } + if (block.isNull(p)) { + continue; + } + int start = block.getFirstValueIndex(p); + int end = start + block.getValueCount(p); + for (int i = start; i < end; i++) { + SpatialExtentGeoPointSourceValuesAggregator.combine(state, block.getBytesRef(i, scratch)); + } + } + } + + @Override + public void addIntermediateInput(Page page) { + assert channels.size() == intermediateBlockCount(); + assert page.getBlockCount() >= channels.get(0) + intermediateStateDesc().size(); + Block minNegXUncast = page.getBlock(channels.get(0)); + if (minNegXUncast.areAllValuesNull()) { + return; + } + IntVector minNegX = ((IntBlock) minNegXUncast).asVector(); + assert minNegX.getPositionCount() == 1; + Block minPosXUncast = page.getBlock(channels.get(1)); + if (minPosXUncast.areAllValuesNull()) { + return; + } + IntVector minPosX = ((IntBlock) minPosXUncast).asVector(); + assert minPosX.getPositionCount() == 1; + Block maxNegXUncast = page.getBlock(channels.get(2)); + if (maxNegXUncast.areAllValuesNull()) { + return; + } + IntVector maxNegX = ((IntBlock) maxNegXUncast).asVector(); + assert maxNegX.getPositionCount() == 1; + Block maxPosXUncast = page.getBlock(channels.get(3)); + if (maxPosXUncast.areAllValuesNull()) { + return; + } + IntVector maxPosX = ((IntBlock) maxPosXUncast).asVector(); + assert maxPosX.getPositionCount() == 1; + Block maxYUncast = page.getBlock(channels.get(4)); + if (maxYUncast.areAllValuesNull()) { + return; + } + IntVector maxY = ((IntBlock) maxYUncast).asVector(); + assert maxY.getPositionCount() == 1; + Block minYUncast = page.getBlock(channels.get(5)); + if (minYUncast.areAllValuesNull()) { + return; + } + IntVector minY = ((IntBlock) minYUncast).asVector(); + assert minY.getPositionCount() == 1; + SpatialExtentGeoPointSourceValuesAggregator.combineIntermediate(state, minNegX.getInt(0), minPosX.getInt(0), maxNegX.getInt(0), maxPosX.getInt(0), maxY.getInt(0), minY.getInt(0)); + } + + @Override + public void evaluateIntermediate(Block[] blocks, int offset, DriverContext driverContext) { + state.toIntermediate(blocks, offset, driverContext); + } + + @Override + public void evaluateFinal(Block[] blocks, int offset, DriverContext driverContext) { + blocks[offset] = SpatialExtentGeoPointSourceValuesAggregator.evaluateFinal(state, driverContext); + } + + @Override + public String toString() { + StringBuilder sb = new StringBuilder(); + sb.append(getClass().getSimpleName()).append("["); + sb.append("channels=").append(channels); + sb.append("]"); + return sb.toString(); + } + + @Override + public void close() { + state.close(); + } +} diff --git a/x-pack/plugin/esql/compute/src/main/generated/org/elasticsearch/compute/aggregation/spatial/SpatialExtentGeoPointSourceValuesAggregatorFunctionSupplier.java b/x-pack/plugin/esql/compute/src/main/generated/org/elasticsearch/compute/aggregation/spatial/SpatialExtentGeoPointSourceValuesAggregatorFunctionSupplier.java new file mode 100644 index 0000000000000..1af20d72d08b0 --- /dev/null +++ b/x-pack/plugin/esql/compute/src/main/generated/org/elasticsearch/compute/aggregation/spatial/SpatialExtentGeoPointSourceValuesAggregatorFunctionSupplier.java @@ -0,0 +1,41 @@ +// Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one +// or more contributor license agreements. Licensed under the Elastic License +// 2.0; you may not use this file except in compliance with the Elastic License +// 2.0. +package org.elasticsearch.compute.aggregation.spatial; + +import java.lang.Integer; +import java.lang.Override; +import java.lang.String; +import java.util.List; +import org.elasticsearch.compute.aggregation.AggregatorFunctionSupplier; +import org.elasticsearch.compute.operator.DriverContext; + +/** + * {@link AggregatorFunctionSupplier} implementation for {@link SpatialExtentGeoPointSourceValuesAggregator}. + * This class is generated. Do not edit it. + */ +public final class SpatialExtentGeoPointSourceValuesAggregatorFunctionSupplier implements AggregatorFunctionSupplier { + private final List channels; + + public SpatialExtentGeoPointSourceValuesAggregatorFunctionSupplier(List channels) { + this.channels = channels; + } + + @Override + public SpatialExtentGeoPointSourceValuesAggregatorFunction aggregator( + DriverContext driverContext) { + return SpatialExtentGeoPointSourceValuesAggregatorFunction.create(driverContext, channels); + } + + @Override + public SpatialExtentGeoPointSourceValuesGroupingAggregatorFunction groupingAggregator( + DriverContext driverContext) { + return SpatialExtentGeoPointSourceValuesGroupingAggregatorFunction.create(channels, driverContext); + } + + @Override + public String describe() { + return "spatial_extent_geo_point_source of valuess"; + } +} diff --git a/x-pack/plugin/esql/compute/src/main/generated/org/elasticsearch/compute/aggregation/spatial/SpatialExtentGeoPointSourceValuesGroupingAggregatorFunction.java b/x-pack/plugin/esql/compute/src/main/generated/org/elasticsearch/compute/aggregation/spatial/SpatialExtentGeoPointSourceValuesGroupingAggregatorFunction.java new file mode 100644 index 0000000000000..bf8ab2554c7b7 --- /dev/null +++ b/x-pack/plugin/esql/compute/src/main/generated/org/elasticsearch/compute/aggregation/spatial/SpatialExtentGeoPointSourceValuesGroupingAggregatorFunction.java @@ -0,0 +1,247 @@ +// Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one +// or more contributor license agreements. Licensed under the Elastic License +// 2.0; you may not use this file except in compliance with the Elastic License +// 2.0. +package org.elasticsearch.compute.aggregation.spatial; + +import java.lang.Integer; +import java.lang.Override; +import java.lang.String; +import java.lang.StringBuilder; +import java.util.List; +import org.apache.lucene.util.BytesRef; +import org.elasticsearch.compute.aggregation.GroupingAggregatorFunction; +import org.elasticsearch.compute.aggregation.IntermediateStateDesc; +import org.elasticsearch.compute.aggregation.SeenGroupIds; +import org.elasticsearch.compute.data.Block; +import org.elasticsearch.compute.data.BytesRefBlock; +import org.elasticsearch.compute.data.BytesRefVector; +import org.elasticsearch.compute.data.ElementType; +import org.elasticsearch.compute.data.IntBlock; +import org.elasticsearch.compute.data.IntVector; +import org.elasticsearch.compute.data.Page; +import org.elasticsearch.compute.operator.DriverContext; + +/** + * {@link GroupingAggregatorFunction} implementation for {@link SpatialExtentGeoPointSourceValuesAggregator}. + * This class is generated. Do not edit it. + */ +public final class SpatialExtentGeoPointSourceValuesGroupingAggregatorFunction implements GroupingAggregatorFunction { + private static final List INTERMEDIATE_STATE_DESC = List.of( + new IntermediateStateDesc("minNegX", ElementType.INT), + new IntermediateStateDesc("minPosX", ElementType.INT), + new IntermediateStateDesc("maxNegX", ElementType.INT), + new IntermediateStateDesc("maxPosX", ElementType.INT), + new IntermediateStateDesc("maxY", ElementType.INT), + new IntermediateStateDesc("minY", ElementType.INT) ); + + private final SpatialExtentGroupingStateWrappedLongitudeState state; + + private final List channels; + + private final DriverContext driverContext; + + public SpatialExtentGeoPointSourceValuesGroupingAggregatorFunction(List channels, + SpatialExtentGroupingStateWrappedLongitudeState state, DriverContext driverContext) { + this.channels = channels; + this.state = state; + this.driverContext = driverContext; + } + + public static SpatialExtentGeoPointSourceValuesGroupingAggregatorFunction create( + List channels, DriverContext driverContext) { + return new SpatialExtentGeoPointSourceValuesGroupingAggregatorFunction(channels, SpatialExtentGeoPointSourceValuesAggregator.initGrouping(), driverContext); + } + + public static List intermediateStateDesc() { + return INTERMEDIATE_STATE_DESC; + } + + @Override + public int intermediateBlockCount() { + return INTERMEDIATE_STATE_DESC.size(); + } + + @Override + public GroupingAggregatorFunction.AddInput prepareProcessPage(SeenGroupIds seenGroupIds, + Page page) { + BytesRefBlock valuesBlock = page.getBlock(channels.get(0)); + BytesRefVector valuesVector = valuesBlock.asVector(); + if (valuesVector == null) { + if (valuesBlock.mayHaveNulls()) { + state.enableGroupIdTracking(seenGroupIds); + } + return new GroupingAggregatorFunction.AddInput() { + @Override + public void add(int positionOffset, IntBlock groupIds) { + addRawInput(positionOffset, groupIds, valuesBlock); + } + + @Override + public void add(int positionOffset, IntVector groupIds) { + addRawInput(positionOffset, groupIds, valuesBlock); + } + + @Override + public void close() { + } + }; + } + return new GroupingAggregatorFunction.AddInput() { + @Override + public void add(int positionOffset, IntBlock groupIds) { + addRawInput(positionOffset, groupIds, valuesVector); + } + + @Override + public void add(int positionOffset, IntVector groupIds) { + addRawInput(positionOffset, groupIds, valuesVector); + } + + @Override + public void close() { + } + }; + } + + private void addRawInput(int positionOffset, IntVector groups, BytesRefBlock values) { + BytesRef scratch = new BytesRef(); + for (int groupPosition = 0; groupPosition < groups.getPositionCount(); groupPosition++) { + int groupId = groups.getInt(groupPosition); + if (values.isNull(groupPosition + positionOffset)) { + continue; + } + int valuesStart = values.getFirstValueIndex(groupPosition + positionOffset); + int valuesEnd = valuesStart + values.getValueCount(groupPosition + positionOffset); + for (int v = valuesStart; v < valuesEnd; v++) { + SpatialExtentGeoPointSourceValuesAggregator.combine(state, groupId, values.getBytesRef(v, scratch)); + } + } + } + + private void addRawInput(int positionOffset, IntVector groups, BytesRefVector values) { + BytesRef scratch = new BytesRef(); + for (int groupPosition = 0; groupPosition < groups.getPositionCount(); groupPosition++) { + int groupId = groups.getInt(groupPosition); + SpatialExtentGeoPointSourceValuesAggregator.combine(state, groupId, values.getBytesRef(groupPosition + positionOffset, scratch)); + } + } + + private void addRawInput(int positionOffset, IntBlock groups, BytesRefBlock values) { + BytesRef scratch = new BytesRef(); + for (int groupPosition = 0; groupPosition < groups.getPositionCount(); groupPosition++) { + if (groups.isNull(groupPosition)) { + continue; + } + int groupStart = groups.getFirstValueIndex(groupPosition); + int groupEnd = groupStart + groups.getValueCount(groupPosition); + for (int g = groupStart; g < groupEnd; g++) { + int groupId = groups.getInt(g); + if (values.isNull(groupPosition + positionOffset)) { + continue; + } + int valuesStart = values.getFirstValueIndex(groupPosition + positionOffset); + int valuesEnd = valuesStart + values.getValueCount(groupPosition + positionOffset); + for (int v = valuesStart; v < valuesEnd; v++) { + SpatialExtentGeoPointSourceValuesAggregator.combine(state, groupId, values.getBytesRef(v, scratch)); + } + } + } + } + + private void addRawInput(int positionOffset, IntBlock groups, BytesRefVector values) { + BytesRef scratch = new BytesRef(); + for (int groupPosition = 0; groupPosition < groups.getPositionCount(); groupPosition++) { + if (groups.isNull(groupPosition)) { + continue; + } + int groupStart = groups.getFirstValueIndex(groupPosition); + int groupEnd = groupStart + groups.getValueCount(groupPosition); + for (int g = groupStart; g < groupEnd; g++) { + int groupId = groups.getInt(g); + SpatialExtentGeoPointSourceValuesAggregator.combine(state, groupId, values.getBytesRef(groupPosition + positionOffset, scratch)); + } + } + } + + @Override + public void selectedMayContainUnseenGroups(SeenGroupIds seenGroupIds) { + state.enableGroupIdTracking(seenGroupIds); + } + + @Override + public void addIntermediateInput(int positionOffset, IntVector groups, Page page) { + state.enableGroupIdTracking(new SeenGroupIds.Empty()); + assert channels.size() == intermediateBlockCount(); + Block minNegXUncast = page.getBlock(channels.get(0)); + if (minNegXUncast.areAllValuesNull()) { + return; + } + IntVector minNegX = ((IntBlock) minNegXUncast).asVector(); + Block minPosXUncast = page.getBlock(channels.get(1)); + if (minPosXUncast.areAllValuesNull()) { + return; + } + IntVector minPosX = ((IntBlock) minPosXUncast).asVector(); + Block maxNegXUncast = page.getBlock(channels.get(2)); + if (maxNegXUncast.areAllValuesNull()) { + return; + } + IntVector maxNegX = ((IntBlock) maxNegXUncast).asVector(); + Block maxPosXUncast = page.getBlock(channels.get(3)); + if (maxPosXUncast.areAllValuesNull()) { + return; + } + IntVector maxPosX = ((IntBlock) maxPosXUncast).asVector(); + Block maxYUncast = page.getBlock(channels.get(4)); + if (maxYUncast.areAllValuesNull()) { + return; + } + IntVector maxY = ((IntBlock) maxYUncast).asVector(); + Block minYUncast = page.getBlock(channels.get(5)); + if (minYUncast.areAllValuesNull()) { + return; + } + IntVector minY = ((IntBlock) minYUncast).asVector(); + assert minNegX.getPositionCount() == minPosX.getPositionCount() && minNegX.getPositionCount() == maxNegX.getPositionCount() && minNegX.getPositionCount() == maxPosX.getPositionCount() && minNegX.getPositionCount() == maxY.getPositionCount() && minNegX.getPositionCount() == minY.getPositionCount(); + for (int groupPosition = 0; groupPosition < groups.getPositionCount(); groupPosition++) { + int groupId = groups.getInt(groupPosition); + SpatialExtentGeoPointSourceValuesAggregator.combineIntermediate(state, groupId, minNegX.getInt(groupPosition + positionOffset), minPosX.getInt(groupPosition + positionOffset), maxNegX.getInt(groupPosition + positionOffset), maxPosX.getInt(groupPosition + positionOffset), maxY.getInt(groupPosition + positionOffset), minY.getInt(groupPosition + positionOffset)); + } + } + + @Override + public void addIntermediateRowInput(int groupId, GroupingAggregatorFunction input, int position) { + if (input.getClass() != getClass()) { + throw new IllegalArgumentException("expected " + getClass() + "; got " + input.getClass()); + } + SpatialExtentGroupingStateWrappedLongitudeState inState = ((SpatialExtentGeoPointSourceValuesGroupingAggregatorFunction) input).state; + state.enableGroupIdTracking(new SeenGroupIds.Empty()); + SpatialExtentGeoPointSourceValuesAggregator.combineStates(state, groupId, inState, position); + } + + @Override + public void evaluateIntermediate(Block[] blocks, int offset, IntVector selected) { + state.toIntermediate(blocks, offset, selected, driverContext); + } + + @Override + public void evaluateFinal(Block[] blocks, int offset, IntVector selected, + DriverContext driverContext) { + blocks[offset] = SpatialExtentGeoPointSourceValuesAggregator.evaluateFinal(state, selected, driverContext); + } + + @Override + public String toString() { + StringBuilder sb = new StringBuilder(); + sb.append(getClass().getSimpleName()).append("["); + sb.append("channels=").append(channels); + sb.append("]"); + return sb.toString(); + } + + @Override + public void close() { + state.close(); + } +} diff --git a/x-pack/plugin/esql/compute/src/main/generated/org/elasticsearch/compute/aggregation/spatial/SpatialExtentGeoShapeAggregatorFunction.java b/x-pack/plugin/esql/compute/src/main/generated/org/elasticsearch/compute/aggregation/spatial/SpatialExtentGeoShapeAggregatorFunction.java new file mode 100644 index 0000000000000..abee9a1cee284 --- /dev/null +++ b/x-pack/plugin/esql/compute/src/main/generated/org/elasticsearch/compute/aggregation/spatial/SpatialExtentGeoShapeAggregatorFunction.java @@ -0,0 +1,206 @@ +// Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one +// or more contributor license agreements. Licensed under the Elastic License +// 2.0; you may not use this file except in compliance with the Elastic License +// 2.0. +package org.elasticsearch.compute.aggregation.spatial; + +import java.lang.Integer; +import java.lang.Override; +import java.lang.String; +import java.lang.StringBuilder; +import java.util.List; +import org.apache.lucene.util.BytesRef; +import org.elasticsearch.compute.aggregation.AggregatorFunction; +import org.elasticsearch.compute.aggregation.IntermediateStateDesc; +import org.elasticsearch.compute.data.Block; +import org.elasticsearch.compute.data.BooleanVector; +import org.elasticsearch.compute.data.BytesRefBlock; +import org.elasticsearch.compute.data.BytesRefVector; +import org.elasticsearch.compute.data.ElementType; +import org.elasticsearch.compute.data.IntBlock; +import org.elasticsearch.compute.data.IntVector; +import org.elasticsearch.compute.data.Page; +import org.elasticsearch.compute.operator.DriverContext; + +/** + * {@link AggregatorFunction} implementation for {@link SpatialExtentGeoShapeAggregator}. + * This class is generated. Do not edit it. + */ +public final class SpatialExtentGeoShapeAggregatorFunction implements AggregatorFunction { + private static final List INTERMEDIATE_STATE_DESC = List.of( + new IntermediateStateDesc("minNegX", ElementType.INT), + new IntermediateStateDesc("minPosX", ElementType.INT), + new IntermediateStateDesc("maxNegX", ElementType.INT), + new IntermediateStateDesc("maxPosX", ElementType.INT), + new IntermediateStateDesc("maxY", ElementType.INT), + new IntermediateStateDesc("minY", ElementType.INT) ); + + private final DriverContext driverContext; + + private final SpatialExtentStateWrappedLongitudeState state; + + private final List channels; + + public SpatialExtentGeoShapeAggregatorFunction(DriverContext driverContext, + List channels, SpatialExtentStateWrappedLongitudeState state) { + this.driverContext = driverContext; + this.channels = channels; + this.state = state; + } + + public static SpatialExtentGeoShapeAggregatorFunction create(DriverContext driverContext, + List channels) { + return new SpatialExtentGeoShapeAggregatorFunction(driverContext, channels, SpatialExtentGeoShapeAggregator.initSingle()); + } + + public static List intermediateStateDesc() { + return INTERMEDIATE_STATE_DESC; + } + + @Override + public int intermediateBlockCount() { + return INTERMEDIATE_STATE_DESC.size(); + } + + @Override + public void addRawInput(Page page, BooleanVector mask) { + if (mask.allFalse()) { + // Entire page masked away + return; + } + if (mask.allTrue()) { + // No masking + BytesRefBlock block = page.getBlock(channels.get(0)); + BytesRefVector vector = block.asVector(); + if (vector != null) { + addRawVector(vector); + } else { + addRawBlock(block); + } + return; + } + // Some positions masked away, others kept + BytesRefBlock block = page.getBlock(channels.get(0)); + BytesRefVector vector = block.asVector(); + if (vector != null) { + addRawVector(vector, mask); + } else { + addRawBlock(block, mask); + } + } + + private void addRawVector(BytesRefVector vector) { + BytesRef scratch = new BytesRef(); + for (int i = 0; i < vector.getPositionCount(); i++) { + SpatialExtentGeoShapeAggregator.combine(state, vector.getBytesRef(i, scratch)); + } + } + + private void addRawVector(BytesRefVector vector, BooleanVector mask) { + BytesRef scratch = new BytesRef(); + for (int i = 0; i < vector.getPositionCount(); i++) { + if (mask.getBoolean(i) == false) { + continue; + } + SpatialExtentGeoShapeAggregator.combine(state, vector.getBytesRef(i, scratch)); + } + } + + private void addRawBlock(BytesRefBlock block) { + BytesRef scratch = new BytesRef(); + for (int p = 0; p < block.getPositionCount(); p++) { + if (block.isNull(p)) { + continue; + } + int start = block.getFirstValueIndex(p); + int end = start + block.getValueCount(p); + for (int i = start; i < end; i++) { + SpatialExtentGeoShapeAggregator.combine(state, block.getBytesRef(i, scratch)); + } + } + } + + private void addRawBlock(BytesRefBlock block, BooleanVector mask) { + BytesRef scratch = new BytesRef(); + for (int p = 0; p < block.getPositionCount(); p++) { + if (mask.getBoolean(p) == false) { + continue; + } + if (block.isNull(p)) { + continue; + } + int start = block.getFirstValueIndex(p); + int end = start + block.getValueCount(p); + for (int i = start; i < end; i++) { + SpatialExtentGeoShapeAggregator.combine(state, block.getBytesRef(i, scratch)); + } + } + } + + @Override + public void addIntermediateInput(Page page) { + assert channels.size() == intermediateBlockCount(); + assert page.getBlockCount() >= channels.get(0) + intermediateStateDesc().size(); + Block minNegXUncast = page.getBlock(channels.get(0)); + if (minNegXUncast.areAllValuesNull()) { + return; + } + IntVector minNegX = ((IntBlock) minNegXUncast).asVector(); + assert minNegX.getPositionCount() == 1; + Block minPosXUncast = page.getBlock(channels.get(1)); + if (minPosXUncast.areAllValuesNull()) { + return; + } + IntVector minPosX = ((IntBlock) minPosXUncast).asVector(); + assert minPosX.getPositionCount() == 1; + Block maxNegXUncast = page.getBlock(channels.get(2)); + if (maxNegXUncast.areAllValuesNull()) { + return; + } + IntVector maxNegX = ((IntBlock) maxNegXUncast).asVector(); + assert maxNegX.getPositionCount() == 1; + Block maxPosXUncast = page.getBlock(channels.get(3)); + if (maxPosXUncast.areAllValuesNull()) { + return; + } + IntVector maxPosX = ((IntBlock) maxPosXUncast).asVector(); + assert maxPosX.getPositionCount() == 1; + Block maxYUncast = page.getBlock(channels.get(4)); + if (maxYUncast.areAllValuesNull()) { + return; + } + IntVector maxY = ((IntBlock) maxYUncast).asVector(); + assert maxY.getPositionCount() == 1; + Block minYUncast = page.getBlock(channels.get(5)); + if (minYUncast.areAllValuesNull()) { + return; + } + IntVector minY = ((IntBlock) minYUncast).asVector(); + assert minY.getPositionCount() == 1; + SpatialExtentGeoShapeAggregator.combineIntermediate(state, minNegX.getInt(0), minPosX.getInt(0), maxNegX.getInt(0), maxPosX.getInt(0), maxY.getInt(0), minY.getInt(0)); + } + + @Override + public void evaluateIntermediate(Block[] blocks, int offset, DriverContext driverContext) { + state.toIntermediate(blocks, offset, driverContext); + } + + @Override + public void evaluateFinal(Block[] blocks, int offset, DriverContext driverContext) { + blocks[offset] = SpatialExtentGeoShapeAggregator.evaluateFinal(state, driverContext); + } + + @Override + public String toString() { + StringBuilder sb = new StringBuilder(); + sb.append(getClass().getSimpleName()).append("["); + sb.append("channels=").append(channels); + sb.append("]"); + return sb.toString(); + } + + @Override + public void close() { + state.close(); + } +} diff --git a/x-pack/plugin/esql/compute/src/main/generated/org/elasticsearch/compute/aggregation/spatial/SpatialExtentGeoShapeAggregatorFunctionSupplier.java b/x-pack/plugin/esql/compute/src/main/generated/org/elasticsearch/compute/aggregation/spatial/SpatialExtentGeoShapeAggregatorFunctionSupplier.java new file mode 100644 index 0000000000000..09f210c7085f8 --- /dev/null +++ b/x-pack/plugin/esql/compute/src/main/generated/org/elasticsearch/compute/aggregation/spatial/SpatialExtentGeoShapeAggregatorFunctionSupplier.java @@ -0,0 +1,40 @@ +// Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one +// or more contributor license agreements. Licensed under the Elastic License +// 2.0; you may not use this file except in compliance with the Elastic License +// 2.0. +package org.elasticsearch.compute.aggregation.spatial; + +import java.lang.Integer; +import java.lang.Override; +import java.lang.String; +import java.util.List; +import org.elasticsearch.compute.aggregation.AggregatorFunctionSupplier; +import org.elasticsearch.compute.operator.DriverContext; + +/** + * {@link AggregatorFunctionSupplier} implementation for {@link SpatialExtentGeoShapeAggregator}. + * This class is generated. Do not edit it. + */ +public final class SpatialExtentGeoShapeAggregatorFunctionSupplier implements AggregatorFunctionSupplier { + private final List channels; + + public SpatialExtentGeoShapeAggregatorFunctionSupplier(List channels) { + this.channels = channels; + } + + @Override + public SpatialExtentGeoShapeAggregatorFunction aggregator(DriverContext driverContext) { + return SpatialExtentGeoShapeAggregatorFunction.create(driverContext, channels); + } + + @Override + public SpatialExtentGeoShapeGroupingAggregatorFunction groupingAggregator( + DriverContext driverContext) { + return SpatialExtentGeoShapeGroupingAggregatorFunction.create(channels, driverContext); + } + + @Override + public String describe() { + return "spatial_extent_geo of shapes"; + } +} diff --git a/x-pack/plugin/esql/compute/src/main/generated/org/elasticsearch/compute/aggregation/spatial/SpatialExtentGeoShapeGroupingAggregatorFunction.java b/x-pack/plugin/esql/compute/src/main/generated/org/elasticsearch/compute/aggregation/spatial/SpatialExtentGeoShapeGroupingAggregatorFunction.java new file mode 100644 index 0000000000000..1200259ea6c41 --- /dev/null +++ b/x-pack/plugin/esql/compute/src/main/generated/org/elasticsearch/compute/aggregation/spatial/SpatialExtentGeoShapeGroupingAggregatorFunction.java @@ -0,0 +1,247 @@ +// Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one +// or more contributor license agreements. Licensed under the Elastic License +// 2.0; you may not use this file except in compliance with the Elastic License +// 2.0. +package org.elasticsearch.compute.aggregation.spatial; + +import java.lang.Integer; +import java.lang.Override; +import java.lang.String; +import java.lang.StringBuilder; +import java.util.List; +import org.apache.lucene.util.BytesRef; +import org.elasticsearch.compute.aggregation.GroupingAggregatorFunction; +import org.elasticsearch.compute.aggregation.IntermediateStateDesc; +import org.elasticsearch.compute.aggregation.SeenGroupIds; +import org.elasticsearch.compute.data.Block; +import org.elasticsearch.compute.data.BytesRefBlock; +import org.elasticsearch.compute.data.BytesRefVector; +import org.elasticsearch.compute.data.ElementType; +import org.elasticsearch.compute.data.IntBlock; +import org.elasticsearch.compute.data.IntVector; +import org.elasticsearch.compute.data.Page; +import org.elasticsearch.compute.operator.DriverContext; + +/** + * {@link GroupingAggregatorFunction} implementation for {@link SpatialExtentGeoShapeAggregator}. + * This class is generated. Do not edit it. + */ +public final class SpatialExtentGeoShapeGroupingAggregatorFunction implements GroupingAggregatorFunction { + private static final List INTERMEDIATE_STATE_DESC = List.of( + new IntermediateStateDesc("minNegX", ElementType.INT), + new IntermediateStateDesc("minPosX", ElementType.INT), + new IntermediateStateDesc("maxNegX", ElementType.INT), + new IntermediateStateDesc("maxPosX", ElementType.INT), + new IntermediateStateDesc("maxY", ElementType.INT), + new IntermediateStateDesc("minY", ElementType.INT) ); + + private final SpatialExtentGroupingStateWrappedLongitudeState state; + + private final List channels; + + private final DriverContext driverContext; + + public SpatialExtentGeoShapeGroupingAggregatorFunction(List channels, + SpatialExtentGroupingStateWrappedLongitudeState state, DriverContext driverContext) { + this.channels = channels; + this.state = state; + this.driverContext = driverContext; + } + + public static SpatialExtentGeoShapeGroupingAggregatorFunction create(List channels, + DriverContext driverContext) { + return new SpatialExtentGeoShapeGroupingAggregatorFunction(channels, SpatialExtentGeoShapeAggregator.initGrouping(), driverContext); + } + + public static List intermediateStateDesc() { + return INTERMEDIATE_STATE_DESC; + } + + @Override + public int intermediateBlockCount() { + return INTERMEDIATE_STATE_DESC.size(); + } + + @Override + public GroupingAggregatorFunction.AddInput prepareProcessPage(SeenGroupIds seenGroupIds, + Page page) { + BytesRefBlock valuesBlock = page.getBlock(channels.get(0)); + BytesRefVector valuesVector = valuesBlock.asVector(); + if (valuesVector == null) { + if (valuesBlock.mayHaveNulls()) { + state.enableGroupIdTracking(seenGroupIds); + } + return new GroupingAggregatorFunction.AddInput() { + @Override + public void add(int positionOffset, IntBlock groupIds) { + addRawInput(positionOffset, groupIds, valuesBlock); + } + + @Override + public void add(int positionOffset, IntVector groupIds) { + addRawInput(positionOffset, groupIds, valuesBlock); + } + + @Override + public void close() { + } + }; + } + return new GroupingAggregatorFunction.AddInput() { + @Override + public void add(int positionOffset, IntBlock groupIds) { + addRawInput(positionOffset, groupIds, valuesVector); + } + + @Override + public void add(int positionOffset, IntVector groupIds) { + addRawInput(positionOffset, groupIds, valuesVector); + } + + @Override + public void close() { + } + }; + } + + private void addRawInput(int positionOffset, IntVector groups, BytesRefBlock values) { + BytesRef scratch = new BytesRef(); + for (int groupPosition = 0; groupPosition < groups.getPositionCount(); groupPosition++) { + int groupId = groups.getInt(groupPosition); + if (values.isNull(groupPosition + positionOffset)) { + continue; + } + int valuesStart = values.getFirstValueIndex(groupPosition + positionOffset); + int valuesEnd = valuesStart + values.getValueCount(groupPosition + positionOffset); + for (int v = valuesStart; v < valuesEnd; v++) { + SpatialExtentGeoShapeAggregator.combine(state, groupId, values.getBytesRef(v, scratch)); + } + } + } + + private void addRawInput(int positionOffset, IntVector groups, BytesRefVector values) { + BytesRef scratch = new BytesRef(); + for (int groupPosition = 0; groupPosition < groups.getPositionCount(); groupPosition++) { + int groupId = groups.getInt(groupPosition); + SpatialExtentGeoShapeAggregator.combine(state, groupId, values.getBytesRef(groupPosition + positionOffset, scratch)); + } + } + + private void addRawInput(int positionOffset, IntBlock groups, BytesRefBlock values) { + BytesRef scratch = new BytesRef(); + for (int groupPosition = 0; groupPosition < groups.getPositionCount(); groupPosition++) { + if (groups.isNull(groupPosition)) { + continue; + } + int groupStart = groups.getFirstValueIndex(groupPosition); + int groupEnd = groupStart + groups.getValueCount(groupPosition); + for (int g = groupStart; g < groupEnd; g++) { + int groupId = groups.getInt(g); + if (values.isNull(groupPosition + positionOffset)) { + continue; + } + int valuesStart = values.getFirstValueIndex(groupPosition + positionOffset); + int valuesEnd = valuesStart + values.getValueCount(groupPosition + positionOffset); + for (int v = valuesStart; v < valuesEnd; v++) { + SpatialExtentGeoShapeAggregator.combine(state, groupId, values.getBytesRef(v, scratch)); + } + } + } + } + + private void addRawInput(int positionOffset, IntBlock groups, BytesRefVector values) { + BytesRef scratch = new BytesRef(); + for (int groupPosition = 0; groupPosition < groups.getPositionCount(); groupPosition++) { + if (groups.isNull(groupPosition)) { + continue; + } + int groupStart = groups.getFirstValueIndex(groupPosition); + int groupEnd = groupStart + groups.getValueCount(groupPosition); + for (int g = groupStart; g < groupEnd; g++) { + int groupId = groups.getInt(g); + SpatialExtentGeoShapeAggregator.combine(state, groupId, values.getBytesRef(groupPosition + positionOffset, scratch)); + } + } + } + + @Override + public void selectedMayContainUnseenGroups(SeenGroupIds seenGroupIds) { + state.enableGroupIdTracking(seenGroupIds); + } + + @Override + public void addIntermediateInput(int positionOffset, IntVector groups, Page page) { + state.enableGroupIdTracking(new SeenGroupIds.Empty()); + assert channels.size() == intermediateBlockCount(); + Block minNegXUncast = page.getBlock(channels.get(0)); + if (minNegXUncast.areAllValuesNull()) { + return; + } + IntVector minNegX = ((IntBlock) minNegXUncast).asVector(); + Block minPosXUncast = page.getBlock(channels.get(1)); + if (minPosXUncast.areAllValuesNull()) { + return; + } + IntVector minPosX = ((IntBlock) minPosXUncast).asVector(); + Block maxNegXUncast = page.getBlock(channels.get(2)); + if (maxNegXUncast.areAllValuesNull()) { + return; + } + IntVector maxNegX = ((IntBlock) maxNegXUncast).asVector(); + Block maxPosXUncast = page.getBlock(channels.get(3)); + if (maxPosXUncast.areAllValuesNull()) { + return; + } + IntVector maxPosX = ((IntBlock) maxPosXUncast).asVector(); + Block maxYUncast = page.getBlock(channels.get(4)); + if (maxYUncast.areAllValuesNull()) { + return; + } + IntVector maxY = ((IntBlock) maxYUncast).asVector(); + Block minYUncast = page.getBlock(channels.get(5)); + if (minYUncast.areAllValuesNull()) { + return; + } + IntVector minY = ((IntBlock) minYUncast).asVector(); + assert minNegX.getPositionCount() == minPosX.getPositionCount() && minNegX.getPositionCount() == maxNegX.getPositionCount() && minNegX.getPositionCount() == maxPosX.getPositionCount() && minNegX.getPositionCount() == maxY.getPositionCount() && minNegX.getPositionCount() == minY.getPositionCount(); + for (int groupPosition = 0; groupPosition < groups.getPositionCount(); groupPosition++) { + int groupId = groups.getInt(groupPosition); + SpatialExtentGeoShapeAggregator.combineIntermediate(state, groupId, minNegX.getInt(groupPosition + positionOffset), minPosX.getInt(groupPosition + positionOffset), maxNegX.getInt(groupPosition + positionOffset), maxPosX.getInt(groupPosition + positionOffset), maxY.getInt(groupPosition + positionOffset), minY.getInt(groupPosition + positionOffset)); + } + } + + @Override + public void addIntermediateRowInput(int groupId, GroupingAggregatorFunction input, int position) { + if (input.getClass() != getClass()) { + throw new IllegalArgumentException("expected " + getClass() + "; got " + input.getClass()); + } + SpatialExtentGroupingStateWrappedLongitudeState inState = ((SpatialExtentGeoShapeGroupingAggregatorFunction) input).state; + state.enableGroupIdTracking(new SeenGroupIds.Empty()); + SpatialExtentGeoShapeAggregator.combineStates(state, groupId, inState, position); + } + + @Override + public void evaluateIntermediate(Block[] blocks, int offset, IntVector selected) { + state.toIntermediate(blocks, offset, selected, driverContext); + } + + @Override + public void evaluateFinal(Block[] blocks, int offset, IntVector selected, + DriverContext driverContext) { + blocks[offset] = SpatialExtentGeoShapeAggregator.evaluateFinal(state, selected, driverContext); + } + + @Override + public String toString() { + StringBuilder sb = new StringBuilder(); + sb.append(getClass().getSimpleName()).append("["); + sb.append("channels=").append(channels); + sb.append("]"); + return sb.toString(); + } + + @Override + public void close() { + state.close(); + } +} diff --git a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/AbstractArrayState.java b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/AbstractArrayState.java index 45a45f4337beb..5fa1394e8cf96 100644 --- a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/AbstractArrayState.java +++ b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/AbstractArrayState.java @@ -28,7 +28,7 @@ public AbstractArrayState(BigArrays bigArrays) { this.bigArrays = bigArrays; } - boolean hasValue(int groupId) { + public boolean hasValue(int groupId) { return seen == null || seen.get(groupId); } @@ -37,7 +37,7 @@ boolean hasValue(int groupId) { * idempotent and fast if already tracking so it's safe to, say, call it once * for every block of values that arrives containing {@code null}. */ - final void enableGroupIdTracking(SeenGroupIds seenGroupIds) { + public final void enableGroupIdTracking(SeenGroupIds seenGroupIds) { if (seen == null) { seen = seenGroupIds.seenGroupIds(bigArrays); } diff --git a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/spatial/CentroidPointAggregator.java b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/spatial/CentroidPointAggregator.java index c66c960dd8a99..47d927fda91b5 100644 --- a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/spatial/CentroidPointAggregator.java +++ b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/spatial/CentroidPointAggregator.java @@ -32,6 +32,13 @@ * This requires that the planner has planned that points are loaded from the index as doc-values. */ abstract class CentroidPointAggregator { + public static CentroidState initSingle() { + return new CentroidState(); + } + + public static GroupingCentroidState initGrouping(BigArrays bigArrays) { + return new GroupingCentroidState(bigArrays); + } public static void combine(CentroidState current, double xVal, double xDel, double yVal, double yDel, long count) { current.add(xVal, xDel, yVal, yDel, count); diff --git a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/spatial/GeoPointEnvelopeVisitor.java b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/spatial/GeoPointEnvelopeVisitor.java new file mode 100644 index 0000000000000..6bdd028f3d6ee --- /dev/null +++ b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/spatial/GeoPointEnvelopeVisitor.java @@ -0,0 +1,63 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +package org.elasticsearch.compute.aggregation.spatial; + +import org.elasticsearch.geometry.Rectangle; +import org.elasticsearch.geometry.utils.SpatialEnvelopeVisitor; +import org.elasticsearch.geometry.utils.SpatialEnvelopeVisitor.WrapLongitude; + +class GeoPointEnvelopeVisitor extends SpatialEnvelopeVisitor.GeoPointVisitor { + GeoPointEnvelopeVisitor() { + super(WrapLongitude.WRAP); + } + + void reset() { + minY = Double.POSITIVE_INFINITY; + maxY = Double.NEGATIVE_INFINITY; + minNegX = Double.POSITIVE_INFINITY; + maxNegX = Double.NEGATIVE_INFINITY; + minPosX = Double.POSITIVE_INFINITY; + maxPosX = Double.NEGATIVE_INFINITY; + } + + double getMinNegX() { + return minNegX; + } + + double getMinPosX() { + return minPosX; + } + + double getMaxNegX() { + return maxNegX; + } + + double getMaxPosX() { + return maxPosX; + } + + double getMaxY() { + return maxY; + } + + double getMinY() { + return minY; + } + + static Rectangle asRectangle( + double minNegX, + double minPosX, + double maxNegX, + double maxPosX, + double maxY, + double minY, + WrapLongitude wrapLongitude + ) { + return SpatialEnvelopeVisitor.GeoPointVisitor.getResult(minNegX, minPosX, maxNegX, maxPosX, maxY, minY, wrapLongitude); + } +} diff --git a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/spatial/PointType.java b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/spatial/PointType.java new file mode 100644 index 0000000000000..5395ca0b85163 --- /dev/null +++ b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/spatial/PointType.java @@ -0,0 +1,107 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +package org.elasticsearch.compute.aggregation.spatial; + +import org.apache.lucene.geo.GeoEncodingUtils; +import org.apache.lucene.geo.XYEncodingUtils; +import org.elasticsearch.geometry.Geometry; +import org.elasticsearch.geometry.Rectangle; +import org.elasticsearch.geometry.utils.SpatialEnvelopeVisitor; +import org.elasticsearch.geometry.utils.SpatialEnvelopeVisitor.WrapLongitude; + +import java.util.Optional; + +public enum PointType { + GEO { + @Override + public Optional computeEnvelope(Geometry geo) { + return SpatialEnvelopeVisitor.visitGeo(geo, WrapLongitude.WRAP); + } + + @Override + public double decodeX(int encoded) { + return GeoEncodingUtils.decodeLongitude(encoded); + } + + @Override + public double decodeY(int encoded) { + return GeoEncodingUtils.decodeLatitude(encoded); + } + + @Override + public int encodeX(double decoded) { + return GeoEncodingUtils.encodeLongitude(decoded); + } + + @Override + public int encodeY(double decoded) { + return GeoEncodingUtils.encodeLatitude(decoded); + } + + // Geo encodes the longitude in the lower 32 bits and the latitude in the upper 32 bits. + @Override + public int extractX(long encoded) { + return SpatialAggregationUtils.extractSecond(encoded); + } + + @Override + public int extractY(long encoded) { + return SpatialAggregationUtils.extractFirst(encoded); + } + }, + CARTESIAN { + @Override + public Optional computeEnvelope(Geometry geo) { + return SpatialEnvelopeVisitor.visitCartesian(geo); + } + + @Override + public double decodeX(int encoded) { + return XYEncodingUtils.decode(encoded); + } + + @Override + public double decodeY(int encoded) { + return XYEncodingUtils.decode(encoded); + } + + @Override + public int encodeX(double decoded) { + return XYEncodingUtils.encode((float) decoded); + } + + @Override + public int encodeY(double decoded) { + return XYEncodingUtils.encode((float) decoded); + } + + @Override + public int extractX(long encoded) { + return SpatialAggregationUtils.extractFirst(encoded); + } + + @Override + public int extractY(long encoded) { + return SpatialAggregationUtils.extractSecond(encoded); + } + }; + + public abstract Optional computeEnvelope(Geometry geo); + + public abstract double decodeX(int encoded); + + public abstract double decodeY(int encoded); + + public abstract int encodeX(double decoded); + + public abstract int encodeY(double decoded); + + public abstract int extractX(long encoded); + + public abstract int extractY(long encoded); +} diff --git a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/spatial/SpatialAggregationUtils.java b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/spatial/SpatialAggregationUtils.java new file mode 100644 index 0000000000000..6b29b20601dae --- /dev/null +++ b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/spatial/SpatialAggregationUtils.java @@ -0,0 +1,88 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +package org.elasticsearch.compute.aggregation.spatial; + +import org.apache.lucene.geo.GeoEncodingUtils; +import org.apache.lucene.geo.XYEncodingUtils; +import org.apache.lucene.util.BytesRef; +import org.elasticsearch.geometry.Geometry; +import org.elasticsearch.geometry.Point; +import org.elasticsearch.geometry.Rectangle; +import org.elasticsearch.geometry.utils.GeometryValidator; +import org.elasticsearch.geometry.utils.SpatialEnvelopeVisitor.WrapLongitude; +import org.elasticsearch.geometry.utils.WellKnownBinary; + +class SpatialAggregationUtils { + private SpatialAggregationUtils() { /* Utility class */ } + + public static Geometry decode(BytesRef wkb) { + return WellKnownBinary.fromWKB(GeometryValidator.NOOP, false /* coerce */, wkb.bytes, wkb.offset, wkb.length); + } + + public static Point decodePoint(BytesRef wkb) { + return (Point) decode(wkb); + } + + public static double decodeX(long encoded) { + return XYEncodingUtils.decode(extractFirst(encoded)); + } + + public static int extractFirst(long encoded) { + return (int) (encoded >>> 32); + } + + public static double decodeY(long encoded) { + return XYEncodingUtils.decode(extractSecond(encoded)); + } + + public static int extractSecond(long encoded) { + return (int) (encoded & 0xFFFFFFFFL); + } + + public static double decodeLongitude(long encoded) { + return GeoEncodingUtils.decodeLongitude((int) (encoded & 0xFFFFFFFFL)); + } + + public static double decodeLatitude(long encoded) { + return GeoEncodingUtils.decodeLatitude((int) (encoded >>> 32)); + } + + public static int encodeNegativeLongitude(double d) { + return Double.isFinite(d) ? GeoEncodingUtils.encodeLongitude(d) : DEFAULT_NEG; + } + + public static int encodePositiveLongitude(double d) { + return Double.isFinite(d) ? GeoEncodingUtils.encodeLongitude(d) : DEFAULT_POS; + } + + public static Rectangle asRectangle(int minNegX, int minPosX, int maxNegX, int maxPosX, int maxY, int minY) { + assert minNegX <= 0 == maxNegX <= 0; + assert minPosX >= 0 == maxPosX >= 0; + return GeoPointEnvelopeVisitor.asRectangle( + minNegX <= 0 ? decodeLongitude(minNegX) : Double.POSITIVE_INFINITY, + minPosX >= 0 ? decodeLongitude(minPosX) : Double.POSITIVE_INFINITY, + maxNegX <= 0 ? decodeLongitude(maxNegX) : Double.NEGATIVE_INFINITY, + maxPosX >= 0 ? decodeLongitude(maxPosX) : Double.NEGATIVE_INFINITY, + GeoEncodingUtils.decodeLatitude(maxY), + GeoEncodingUtils.decodeLatitude(minY), + WrapLongitude.WRAP + ); + } + + public static int maxNeg(int a, int b) { + return a <= 0 && b <= 0 ? Math.max(a, b) : Math.min(a, b); + } + + public static int minPos(int a, int b) { + return a >= 0 && b >= 0 ? Math.min(a, b) : Math.max(a, b); + } + + // The default values are intentionally non-negative/non-positive, so we can mark unassigned values. + public static final int DEFAULT_POS = -1; + public static final int DEFAULT_NEG = 1; +} diff --git a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/spatial/SpatialCentroidCartesianPointDocValuesAggregator.java b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/spatial/SpatialCentroidCartesianPointDocValuesAggregator.java index 0bafb6f8112de..891c22b71c7e9 100644 --- a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/spatial/SpatialCentroidCartesianPointDocValuesAggregator.java +++ b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/spatial/SpatialCentroidCartesianPointDocValuesAggregator.java @@ -7,12 +7,13 @@ package org.elasticsearch.compute.aggregation.spatial; -import org.apache.lucene.geo.XYEncodingUtils; -import org.elasticsearch.common.util.BigArrays; import org.elasticsearch.compute.ann.Aggregator; import org.elasticsearch.compute.ann.GroupingAggregator; import org.elasticsearch.compute.ann.IntermediateState; +import static org.elasticsearch.compute.aggregation.spatial.SpatialAggregationUtils.decodeX; +import static org.elasticsearch.compute.aggregation.spatial.SpatialAggregationUtils.decodeY; + /** * This aggregator calculates the centroid of a set of cartesian points. * It is assumes that the cartesian points are encoded as longs. @@ -28,15 +29,6 @@ ) @GroupingAggregator class SpatialCentroidCartesianPointDocValuesAggregator extends CentroidPointAggregator { - - public static CentroidState initSingle() { - return new CentroidState(); - } - - public static GroupingCentroidState initGrouping(BigArrays bigArrays) { - return new GroupingCentroidState(bigArrays); - } - public static void combine(CentroidState current, long v) { current.add(decodeX(v), decodeY(v)); } @@ -44,12 +36,4 @@ public static void combine(CentroidState current, long v) { public static void combine(GroupingCentroidState current, int groupId, long encoded) { current.add(decodeX(encoded), 0d, decodeY(encoded), 0d, 1, groupId); } - - private static double decodeX(long encoded) { - return XYEncodingUtils.decode((int) (encoded >>> 32)); - } - - private static double decodeY(long encoded) { - return XYEncodingUtils.decode((int) (encoded & 0xFFFFFFFFL)); - } } diff --git a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/spatial/SpatialCentroidCartesianPointSourceValuesAggregator.java b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/spatial/SpatialCentroidCartesianPointSourceValuesAggregator.java index 5673892be4bf0..700721e3ea9d4 100644 --- a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/spatial/SpatialCentroidCartesianPointSourceValuesAggregator.java +++ b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/spatial/SpatialCentroidCartesianPointSourceValuesAggregator.java @@ -8,13 +8,10 @@ package org.elasticsearch.compute.aggregation.spatial; import org.apache.lucene.util.BytesRef; -import org.elasticsearch.common.util.BigArrays; import org.elasticsearch.compute.ann.Aggregator; import org.elasticsearch.compute.ann.GroupingAggregator; import org.elasticsearch.compute.ann.IntermediateState; import org.elasticsearch.geometry.Point; -import org.elasticsearch.geometry.utils.GeometryValidator; -import org.elasticsearch.geometry.utils.WellKnownBinary; /** * This aggregator calculates the centroid of a set of cartesian points. @@ -33,26 +30,13 @@ ) @GroupingAggregator class SpatialCentroidCartesianPointSourceValuesAggregator extends CentroidPointAggregator { - - public static CentroidState initSingle() { - return new CentroidState(); - } - - public static GroupingCentroidState initGrouping(BigArrays bigArrays) { - return new GroupingCentroidState(bigArrays); - } - public static void combine(CentroidState current, BytesRef wkb) { - Point point = decode(wkb); + Point point = SpatialAggregationUtils.decodePoint(wkb); current.add(point.getX(), point.getY()); } public static void combine(GroupingCentroidState current, int groupId, BytesRef wkb) { - Point point = decode(wkb); + Point point = SpatialAggregationUtils.decodePoint(wkb); current.add(point.getX(), 0d, point.getY(), 0d, 1, groupId); } - - private static Point decode(BytesRef wkb) { - return (Point) WellKnownBinary.fromWKB(GeometryValidator.NOOP, false, wkb.bytes, wkb.offset, wkb.length); - } } diff --git a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/spatial/SpatialCentroidGeoPointDocValuesAggregator.java b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/spatial/SpatialCentroidGeoPointDocValuesAggregator.java index ee5ab0e292547..431e25a03779e 100644 --- a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/spatial/SpatialCentroidGeoPointDocValuesAggregator.java +++ b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/spatial/SpatialCentroidGeoPointDocValuesAggregator.java @@ -7,12 +7,13 @@ package org.elasticsearch.compute.aggregation.spatial; -import org.apache.lucene.geo.GeoEncodingUtils; -import org.elasticsearch.common.util.BigArrays; import org.elasticsearch.compute.ann.Aggregator; import org.elasticsearch.compute.ann.GroupingAggregator; import org.elasticsearch.compute.ann.IntermediateState; +import static org.elasticsearch.compute.aggregation.spatial.SpatialAggregationUtils.decodeLatitude; +import static org.elasticsearch.compute.aggregation.spatial.SpatialAggregationUtils.decodeLongitude; + /** * This aggregator calculates the centroid of a set of geo points. It is assumes that the geo points are encoded as longs. * This requires that the planner has planned that points are loaded from the index as doc-values. @@ -27,28 +28,11 @@ ) @GroupingAggregator class SpatialCentroidGeoPointDocValuesAggregator extends CentroidPointAggregator { - - public static CentroidState initSingle() { - return new CentroidState(); - } - - public static GroupingCentroidState initGrouping(BigArrays bigArrays) { - return new GroupingCentroidState(bigArrays); - } - public static void combine(CentroidState current, long v) { - current.add(decodeX(v), decodeY(v)); + current.add(decodeLongitude(v), decodeLatitude(v)); } public static void combine(GroupingCentroidState current, int groupId, long encoded) { - current.add(decodeX(encoded), 0d, decodeY(encoded), 0d, 1, groupId); - } - - private static double decodeX(long encoded) { - return GeoEncodingUtils.decodeLongitude((int) (encoded & 0xFFFFFFFFL)); - } - - private static double decodeY(long encoded) { - return GeoEncodingUtils.decodeLatitude((int) (encoded >>> 32)); + current.add(decodeLongitude(encoded), 0d, decodeLatitude(encoded), 0d, 1, groupId); } } diff --git a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/spatial/SpatialCentroidGeoPointSourceValuesAggregator.java b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/spatial/SpatialCentroidGeoPointSourceValuesAggregator.java index caf55dcc2f4e1..90563b33b8abb 100644 --- a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/spatial/SpatialCentroidGeoPointSourceValuesAggregator.java +++ b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/spatial/SpatialCentroidGeoPointSourceValuesAggregator.java @@ -8,13 +8,10 @@ package org.elasticsearch.compute.aggregation.spatial; import org.apache.lucene.util.BytesRef; -import org.elasticsearch.common.util.BigArrays; import org.elasticsearch.compute.ann.Aggregator; import org.elasticsearch.compute.ann.GroupingAggregator; import org.elasticsearch.compute.ann.IntermediateState; import org.elasticsearch.geometry.Point; -import org.elasticsearch.geometry.utils.GeometryValidator; -import org.elasticsearch.geometry.utils.WellKnownBinary; /** * This aggregator calculates the centroid of a set of geo points. @@ -33,26 +30,13 @@ ) @GroupingAggregator class SpatialCentroidGeoPointSourceValuesAggregator extends CentroidPointAggregator { - - public static CentroidState initSingle() { - return new CentroidState(); - } - - public static GroupingCentroidState initGrouping(BigArrays bigArrays) { - return new GroupingCentroidState(bigArrays); - } - public static void combine(CentroidState current, BytesRef wkb) { - Point point = decode(wkb); + Point point = SpatialAggregationUtils.decodePoint(wkb); current.add(point.getX(), point.getY()); } public static void combine(GroupingCentroidState current, int groupId, BytesRef wkb) { - Point point = decode(wkb); + Point point = SpatialAggregationUtils.decodePoint(wkb); current.add(point.getX(), 0d, point.getY(), 0d, 1, groupId); } - - private static Point decode(BytesRef wkb) { - return (Point) WellKnownBinary.fromWKB(GeometryValidator.NOOP, false, wkb.bytes, wkb.offset, wkb.length); - } } diff --git a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/spatial/SpatialExtentAggregator.java b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/spatial/SpatialExtentAggregator.java new file mode 100644 index 0000000000000..91e0f098d795e --- /dev/null +++ b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/spatial/SpatialExtentAggregator.java @@ -0,0 +1,36 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +package org.elasticsearch.compute.aggregation.spatial; + +import org.elasticsearch.compute.data.Block; +import org.elasticsearch.compute.data.IntVector; +import org.elasticsearch.compute.operator.DriverContext; + +// A bit of abuse of notation here, since we're extending this class to "inherit" its static methods. +// Unfortunately, this is the way it has to be done, since the generated code invokes these methods statically. +abstract class SpatialExtentAggregator { + public static void combineIntermediate(SpatialExtentState current, int minX, int maxX, int maxY, int minY) { + current.add(minX, maxX, maxY, minY); + } + + public static void combineIntermediate(SpatialExtentGroupingState current, int groupId, int minX, int maxX, int maxY, int minY) { + current.add(groupId, minX, maxX, maxY, minY); + } + + public static Block evaluateFinal(SpatialExtentState state, DriverContext driverContext) { + return state.toBlock(driverContext); + } + + public static Block evaluateFinal(SpatialExtentGroupingState state, IntVector selected, DriverContext driverContext) { + return state.toBlock(selected, driverContext); + } + + public static void combineStates(SpatialExtentGroupingState current, int groupId, SpatialExtentGroupingState inState, int inPosition) { + current.add(groupId, inState, inPosition); + } +} diff --git a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/spatial/SpatialExtentCartesianPointDocValuesAggregator.java b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/spatial/SpatialExtentCartesianPointDocValuesAggregator.java new file mode 100644 index 0000000000000..f64949b77707c --- /dev/null +++ b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/spatial/SpatialExtentCartesianPointDocValuesAggregator.java @@ -0,0 +1,42 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +package org.elasticsearch.compute.aggregation.spatial; + +import org.elasticsearch.compute.ann.Aggregator; +import org.elasticsearch.compute.ann.GroupingAggregator; +import org.elasticsearch.compute.ann.IntermediateState; + +/** + * Computes the extent of a set of cartesian points. It is assumed the points are encoded as longs. + * This requires that the planner has planned that points are loaded from the index as doc-values. + */ +@Aggregator( + { + @IntermediateState(name = "minX", type = "INT"), + @IntermediateState(name = "maxX", type = "INT"), + @IntermediateState(name = "maxY", type = "INT"), + @IntermediateState(name = "minY", type = "INT") } +) +@GroupingAggregator +class SpatialExtentCartesianPointDocValuesAggregator extends SpatialExtentAggregator { + public static SpatialExtentState initSingle() { + return new SpatialExtentState(PointType.CARTESIAN); + } + + public static SpatialExtentGroupingState initGrouping() { + return new SpatialExtentGroupingState(PointType.CARTESIAN); + } + + public static void combine(SpatialExtentState current, long v) { + current.add(v); + } + + public static void combine(SpatialExtentGroupingState current, int groupId, long v) { + current.add(groupId, v); + } +} diff --git a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/spatial/SpatialExtentCartesianPointSourceValuesAggregator.java b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/spatial/SpatialExtentCartesianPointSourceValuesAggregator.java new file mode 100644 index 0000000000000..3488af4525dcb --- /dev/null +++ b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/spatial/SpatialExtentCartesianPointSourceValuesAggregator.java @@ -0,0 +1,45 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +package org.elasticsearch.compute.aggregation.spatial; + +import org.apache.lucene.util.BytesRef; +import org.elasticsearch.compute.ann.Aggregator; +import org.elasticsearch.compute.ann.GroupingAggregator; +import org.elasticsearch.compute.ann.IntermediateState; + +/** + * Computes the extent of a set of cartesian points. It is assumed that the cartesian points are encoded as WKB BytesRef. + * This requires that the planner has NOT planned that points are loaded from the index as doc-values, but from source instead. + * This is also used for final aggregations and aggregations in the coordinator node, + * even if the local node partial aggregation is done with {@link SpatialExtentCartesianPointDocValuesAggregator}. + */ +@Aggregator( + { + @IntermediateState(name = "minX", type = "INT"), + @IntermediateState(name = "maxX", type = "INT"), + @IntermediateState(name = "maxY", type = "INT"), + @IntermediateState(name = "minY", type = "INT") } +) +@GroupingAggregator +class SpatialExtentCartesianPointSourceValuesAggregator extends SpatialExtentAggregator { + public static SpatialExtentState initSingle() { + return new SpatialExtentState(PointType.CARTESIAN); + } + + public static SpatialExtentGroupingState initGrouping() { + return new SpatialExtentGroupingState(PointType.CARTESIAN); + } + + public static void combine(SpatialExtentState current, BytesRef bytes) { + current.add(SpatialAggregationUtils.decode(bytes)); + } + + public static void combine(SpatialExtentGroupingState current, int groupId, BytesRef bytes) { + current.add(groupId, SpatialAggregationUtils.decode(bytes)); + } +} diff --git a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/spatial/SpatialExtentCartesianShapeAggregator.java b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/spatial/SpatialExtentCartesianShapeAggregator.java new file mode 100644 index 0000000000000..6d50d27aa5a2d --- /dev/null +++ b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/spatial/SpatialExtentCartesianShapeAggregator.java @@ -0,0 +1,43 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +package org.elasticsearch.compute.aggregation.spatial; + +import org.apache.lucene.util.BytesRef; +import org.elasticsearch.compute.ann.Aggregator; +import org.elasticsearch.compute.ann.GroupingAggregator; +import org.elasticsearch.compute.ann.IntermediateState; + +/** + * Computes the extent of a set of cartesian shapes. It is assumed that the cartesian shapes are encoded as WKB BytesRef. + * We do not currently support reading shape values or extents from doc values. + */ +@Aggregator( + { + @IntermediateState(name = "minX", type = "INT"), + @IntermediateState(name = "maxX", type = "INT"), + @IntermediateState(name = "maxY", type = "INT"), + @IntermediateState(name = "minY", type = "INT") } +) +@GroupingAggregator +class SpatialExtentCartesianShapeAggregator extends SpatialExtentAggregator { + public static SpatialExtentState initSingle() { + return new SpatialExtentState(PointType.CARTESIAN); + } + + public static SpatialExtentGroupingState initGrouping() { + return new SpatialExtentGroupingState(PointType.CARTESIAN); + } + + public static void combine(SpatialExtentState current, BytesRef bytes) { + current.add(SpatialAggregationUtils.decode(bytes)); + } + + public static void combine(SpatialExtentGroupingState current, int groupId, BytesRef bytes) { + current.add(groupId, SpatialAggregationUtils.decode(bytes)); + } +} diff --git a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/spatial/SpatialExtentGeoPointDocValuesAggregator.java b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/spatial/SpatialExtentGeoPointDocValuesAggregator.java new file mode 100644 index 0000000000000..b9b8bf65e116b --- /dev/null +++ b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/spatial/SpatialExtentGeoPointDocValuesAggregator.java @@ -0,0 +1,45 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +package org.elasticsearch.compute.aggregation.spatial; + +import org.elasticsearch.compute.ann.Aggregator; +import org.elasticsearch.compute.ann.GroupingAggregator; +import org.elasticsearch.compute.ann.IntermediateState; + +/** + * Computes the extent of a set of geo points. It is assumed the points are encoded as longs. + * This requires that the planner has planned that points are loaded from the index as doc-values. + */ +@Aggregator( + { + @IntermediateState(name = "minNegX", type = "INT"), + @IntermediateState(name = "minPosX", type = "INT"), + @IntermediateState(name = "maxNegX", type = "INT"), + @IntermediateState(name = "maxPosX", type = "INT"), + @IntermediateState(name = "maxY", type = "INT"), + @IntermediateState(name = "minY", type = "INT") } +) +@GroupingAggregator +class SpatialExtentGeoPointDocValuesAggregator extends SpatialExtentLongitudeWrappingAggregator { + // TODO support non-longitude wrapped geo shapes. + public static SpatialExtentStateWrappedLongitudeState initSingle() { + return new SpatialExtentStateWrappedLongitudeState(); + } + + public static SpatialExtentGroupingStateWrappedLongitudeState initGrouping() { + return new SpatialExtentGroupingStateWrappedLongitudeState(); + } + + public static void combine(SpatialExtentStateWrappedLongitudeState current, long encoded) { + current.add(encoded); + } + + public static void combine(SpatialExtentGroupingStateWrappedLongitudeState current, int groupId, long encoded) { + current.add(groupId, encoded); + } +} diff --git a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/spatial/SpatialExtentGeoPointSourceValuesAggregator.java b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/spatial/SpatialExtentGeoPointSourceValuesAggregator.java new file mode 100644 index 0000000000000..36a4e359f23fc --- /dev/null +++ b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/spatial/SpatialExtentGeoPointSourceValuesAggregator.java @@ -0,0 +1,48 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +package org.elasticsearch.compute.aggregation.spatial; + +import org.apache.lucene.util.BytesRef; +import org.elasticsearch.compute.ann.Aggregator; +import org.elasticsearch.compute.ann.GroupingAggregator; +import org.elasticsearch.compute.ann.IntermediateState; + +/** + * Computes the extent of a set of geo points. It is assumed that the geo points are encoded as WKB BytesRef. + * This requires that the planner has NOT planned that points are loaded from the index as doc-values, but from source instead. + * This is also used for final aggregations and aggregations in the coordinator node, + * even if the local node partial aggregation is done with {@link SpatialExtentGeoPointDocValuesAggregator}. + */ +@Aggregator( + { + @IntermediateState(name = "minNegX", type = "INT"), + @IntermediateState(name = "minPosX", type = "INT"), + @IntermediateState(name = "maxNegX", type = "INT"), + @IntermediateState(name = "maxPosX", type = "INT"), + @IntermediateState(name = "maxY", type = "INT"), + @IntermediateState(name = "minY", type = "INT") } +) +@GroupingAggregator +class SpatialExtentGeoPointSourceValuesAggregator extends SpatialExtentLongitudeWrappingAggregator { + // TODO support non-longitude wrapped geo shapes. + public static SpatialExtentStateWrappedLongitudeState initSingle() { + return new SpatialExtentStateWrappedLongitudeState(); + } + + public static SpatialExtentGroupingStateWrappedLongitudeState initGrouping() { + return new SpatialExtentGroupingStateWrappedLongitudeState(); + } + + public static void combine(SpatialExtentStateWrappedLongitudeState current, BytesRef bytes) { + current.add(SpatialAggregationUtils.decode(bytes)); + } + + public static void combine(SpatialExtentGroupingStateWrappedLongitudeState current, int groupId, BytesRef bytes) { + current.add(groupId, SpatialAggregationUtils.decode(bytes)); + } +} diff --git a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/spatial/SpatialExtentGeoShapeAggregator.java b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/spatial/SpatialExtentGeoShapeAggregator.java new file mode 100644 index 0000000000000..3d1b9b6300c9d --- /dev/null +++ b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/spatial/SpatialExtentGeoShapeAggregator.java @@ -0,0 +1,46 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +package org.elasticsearch.compute.aggregation.spatial; + +import org.apache.lucene.util.BytesRef; +import org.elasticsearch.compute.ann.Aggregator; +import org.elasticsearch.compute.ann.GroupingAggregator; +import org.elasticsearch.compute.ann.IntermediateState; + +/** + * Computes the extent of a set of geo shapes. It is assumed that the geo shapes are encoded as WKB BytesRef. + * We do not currently support reading shape values or extents from doc values. + */ +@Aggregator( + { + @IntermediateState(name = "minNegX", type = "INT"), + @IntermediateState(name = "minPosX", type = "INT"), + @IntermediateState(name = "maxNegX", type = "INT"), + @IntermediateState(name = "maxPosX", type = "INT"), + @IntermediateState(name = "maxY", type = "INT"), + @IntermediateState(name = "minY", type = "INT") } +) +@GroupingAggregator +class SpatialExtentGeoShapeAggregator extends SpatialExtentLongitudeWrappingAggregator { + // TODO support non-longitude wrapped geo shapes. + public static SpatialExtentStateWrappedLongitudeState initSingle() { + return new SpatialExtentStateWrappedLongitudeState(); + } + + public static SpatialExtentGroupingStateWrappedLongitudeState initGrouping() { + return new SpatialExtentGroupingStateWrappedLongitudeState(); + } + + public static void combine(SpatialExtentStateWrappedLongitudeState current, BytesRef bytes) { + current.add(SpatialAggregationUtils.decode(bytes)); + } + + public static void combine(SpatialExtentGroupingStateWrappedLongitudeState current, int groupId, BytesRef bytes) { + current.add(groupId, SpatialAggregationUtils.decode(bytes)); + } +} diff --git a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/spatial/SpatialExtentGroupingState.java b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/spatial/SpatialExtentGroupingState.java new file mode 100644 index 0000000000000..9ce0ccdda0ff5 --- /dev/null +++ b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/spatial/SpatialExtentGroupingState.java @@ -0,0 +1,154 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +package org.elasticsearch.compute.aggregation.spatial; + +import org.apache.lucene.util.BytesRef; +import org.elasticsearch.common.util.BigArrays; +import org.elasticsearch.common.util.IntArray; +import org.elasticsearch.compute.aggregation.AbstractArrayState; +import org.elasticsearch.compute.aggregation.SeenGroupIds; +import org.elasticsearch.compute.data.Block; +import org.elasticsearch.compute.data.IntVector; +import org.elasticsearch.compute.operator.DriverContext; +import org.elasticsearch.geometry.Geometry; +import org.elasticsearch.geometry.Rectangle; +import org.elasticsearch.geometry.utils.WellKnownBinary; + +import java.nio.ByteOrder; + +final class SpatialExtentGroupingState extends AbstractArrayState { + private final PointType pointType; + private IntArray minXs; + private IntArray maxXs; + private IntArray maxYs; + private IntArray minYs; + + SpatialExtentGroupingState(PointType pointType) { + this(pointType, BigArrays.NON_RECYCLING_INSTANCE); + } + + SpatialExtentGroupingState(PointType pointType, BigArrays bigArrays) { + super(bigArrays); + this.pointType = pointType; + this.minXs = bigArrays.newIntArray(0, false); + this.maxXs = bigArrays.newIntArray(0, false); + this.maxYs = bigArrays.newIntArray(0, false); + this.minYs = bigArrays.newIntArray(0, false); + enableGroupIdTracking(new SeenGroupIds.Empty()); + } + + @Override + public void toIntermediate(Block[] blocks, int offset, IntVector selected, DriverContext driverContext) { + assert blocks.length >= offset; + try ( + var minXsBuilder = driverContext.blockFactory().newIntBlockBuilder(selected.getPositionCount()); + var maxXsBuilder = driverContext.blockFactory().newIntBlockBuilder(selected.getPositionCount()); + var maxYsBuilder = driverContext.blockFactory().newIntBlockBuilder(selected.getPositionCount()); + var minYsBuilder = driverContext.blockFactory().newIntBlockBuilder(selected.getPositionCount()); + ) { + for (int i = 0; i < selected.getPositionCount(); i++) { + int group = selected.getInt(i); + assert hasValue(group); + minXsBuilder.appendInt(minXs.get(group)); + maxXsBuilder.appendInt(maxXs.get(group)); + maxYsBuilder.appendInt(maxYs.get(group)); + minYsBuilder.appendInt(minYs.get(group)); + } + blocks[offset + 0] = minXsBuilder.build(); + blocks[offset + 1] = maxXsBuilder.build(); + blocks[offset + 2] = maxYsBuilder.build(); + blocks[offset + 3] = minYsBuilder.build(); + } + } + + public void add(int groupId, Geometry geometry) { + ensureCapacity(groupId); + pointType.computeEnvelope(geometry) + .ifPresent( + r -> add( + groupId, + pointType.encodeX(r.getMinX()), + pointType.encodeX(r.getMaxX()), + pointType.encodeY(r.getMaxY()), + pointType.encodeY(r.getMinY()) + ) + ); + } + + public void add(int groupId, long encoded) { + int x = pointType.extractX(encoded); + int y = pointType.extractY(encoded); + add(groupId, x, x, y, y); + } + + public void add(int groupId, int minX, int maxX, int maxY, int minY) { + ensureCapacity(groupId); + if (hasValue(groupId)) { + minXs.set(groupId, Math.min(minXs.get(groupId), minX)); + maxXs.set(groupId, Math.max(maxXs.get(groupId), maxX)); + maxYs.set(groupId, Math.max(maxYs.get(groupId), maxY)); + minYs.set(groupId, Math.min(minYs.get(groupId), minY)); + } else { + minXs.set(groupId, minX); + maxXs.set(groupId, maxX); + maxYs.set(groupId, maxY); + minYs.set(groupId, minY); + } + trackGroupId(groupId); + } + + private void ensureCapacity(int groupId) { + long requiredSize = groupId + 1; + if (minXs.size() < requiredSize) { + assert minXs.size() == maxXs.size() && minXs.size() == maxYs.size() && minXs.size() == minYs.size(); + minXs = bigArrays.grow(minXs, requiredSize); + maxXs = bigArrays.grow(maxXs, requiredSize); + maxYs = bigArrays.grow(maxYs, requiredSize); + minYs = bigArrays.grow(minYs, requiredSize); + } + } + + public Block toBlock(IntVector selected, DriverContext driverContext) { + try (var builder = driverContext.blockFactory().newBytesRefBlockBuilder(selected.getPositionCount())) { + for (int i = 0; i < selected.getPositionCount(); i++) { + int si = selected.getInt(i); + if (hasValue(si)) { + builder.appendBytesRef( + new BytesRef( + WellKnownBinary.toWKB( + new Rectangle( + pointType.decodeX(minXs.get(si)), + pointType.decodeX(maxXs.get(si)), + pointType.decodeY(maxYs.get(si)), + pointType.decodeY(minYs.get(si)) + ), + ByteOrder.LITTLE_ENDIAN + ) + ) + ); + } else { + builder.appendNull(); + } + } + return builder.build(); + } + } + + public void add(int groupId, SpatialExtentGroupingState inState, int inPosition) { + ensureCapacity(groupId); + if (inState.hasValue(inPosition)) { + add( + groupId, + inState.minXs.get(inPosition), + inState.maxXs.get(inPosition), + inState.maxYs.get(inPosition), + inState.minYs.get(inPosition) + ); + } + } +} diff --git a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/spatial/SpatialExtentGroupingStateWrappedLongitudeState.java b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/spatial/SpatialExtentGroupingStateWrappedLongitudeState.java new file mode 100644 index 0000000000000..3dd7a6d4acde2 --- /dev/null +++ b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/spatial/SpatialExtentGroupingStateWrappedLongitudeState.java @@ -0,0 +1,182 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +package org.elasticsearch.compute.aggregation.spatial; + +import org.apache.lucene.util.BytesRef; +import org.elasticsearch.common.util.BigArrays; +import org.elasticsearch.common.util.IntArray; +import org.elasticsearch.compute.aggregation.AbstractArrayState; +import org.elasticsearch.compute.aggregation.GroupingAggregatorState; +import org.elasticsearch.compute.aggregation.SeenGroupIds; +import org.elasticsearch.compute.data.Block; +import org.elasticsearch.compute.data.IntVector; +import org.elasticsearch.compute.operator.DriverContext; +import org.elasticsearch.geometry.Geometry; +import org.elasticsearch.geometry.utils.SpatialEnvelopeVisitor; +import org.elasticsearch.geometry.utils.WellKnownBinary; + +import java.nio.ByteOrder; + +final class SpatialExtentGroupingStateWrappedLongitudeState extends AbstractArrayState implements GroupingAggregatorState { + // Only geo points support longitude wrapping. + private static final PointType POINT_TYPE = PointType.GEO; + private IntArray minNegXs; + private IntArray minPosXs; + private IntArray maxNegXs; + private IntArray maxPosXs; + private IntArray maxYs; + private IntArray minYs; + + private GeoPointEnvelopeVisitor geoPointVisitor = new GeoPointEnvelopeVisitor(); + + SpatialExtentGroupingStateWrappedLongitudeState() { + this(BigArrays.NON_RECYCLING_INSTANCE); + } + + SpatialExtentGroupingStateWrappedLongitudeState(BigArrays bigArrays) { + super(bigArrays); + this.minNegXs = bigArrays.newIntArray(0, false); + this.minPosXs = bigArrays.newIntArray(0, false); + this.maxNegXs = bigArrays.newIntArray(0, false); + this.maxPosXs = bigArrays.newIntArray(0, false); + this.maxYs = bigArrays.newIntArray(0, false); + this.minYs = bigArrays.newIntArray(0, false); + enableGroupIdTracking(new SeenGroupIds.Empty()); + } + + @Override + public void toIntermediate(Block[] blocks, int offset, IntVector selected, DriverContext driverContext) { + assert blocks.length >= offset; + try ( + var minNegXsBuilder = driverContext.blockFactory().newIntBlockBuilder(selected.getPositionCount()); + var minPosXsBuilder = driverContext.blockFactory().newIntBlockBuilder(selected.getPositionCount()); + var maxNegXsBuilder = driverContext.blockFactory().newIntBlockBuilder(selected.getPositionCount()); + var maxPosXsBuilder = driverContext.blockFactory().newIntBlockBuilder(selected.getPositionCount()); + var maxYsBuilder = driverContext.blockFactory().newIntBlockBuilder(selected.getPositionCount()); + var minYsBuilder = driverContext.blockFactory().newIntBlockBuilder(selected.getPositionCount()); + ) { + for (int i = 0; i < selected.getPositionCount(); i++) { + int group = selected.getInt(i); + assert hasValue(group); + assert minNegXs.get(group) <= 0 == maxNegXs.get(group) <= 0; + assert minPosXs.get(group) >= 0 == maxPosXs.get(group) >= 0; + minNegXsBuilder.appendInt(minNegXs.get(group)); + minPosXsBuilder.appendInt(minPosXs.get(group)); + maxNegXsBuilder.appendInt(maxNegXs.get(group)); + maxPosXsBuilder.appendInt(maxPosXs.get(group)); + maxYsBuilder.appendInt(maxYs.get(group)); + minYsBuilder.appendInt(minYs.get(group)); + } + blocks[offset + 0] = minNegXsBuilder.build(); + blocks[offset + 1] = minPosXsBuilder.build(); + blocks[offset + 2] = maxNegXsBuilder.build(); + blocks[offset + 3] = maxPosXsBuilder.build(); + blocks[offset + 4] = maxYsBuilder.build(); + blocks[offset + 5] = minYsBuilder.build(); + } + } + + public void add(int groupId, Geometry geo) { + ensureCapacity(groupId); + geoPointVisitor.reset(); + if (geo.visit(new SpatialEnvelopeVisitor(geoPointVisitor))) { + add( + groupId, + SpatialAggregationUtils.encodeNegativeLongitude(geoPointVisitor.getMinNegX()), + SpatialAggregationUtils.encodePositiveLongitude(geoPointVisitor.getMinPosX()), + SpatialAggregationUtils.encodeNegativeLongitude(geoPointVisitor.getMaxNegX()), + SpatialAggregationUtils.encodePositiveLongitude(geoPointVisitor.getMaxPosX()), + POINT_TYPE.encodeY(geoPointVisitor.getMaxY()), + POINT_TYPE.encodeY(geoPointVisitor.getMinY()) + ); + } + } + + public void add(int groupId, SpatialExtentGroupingStateWrappedLongitudeState inState, int inPosition) { + ensureCapacity(groupId); + if (inState.hasValue(inPosition)) { + add( + groupId, + inState.minNegXs.get(inPosition), + inState.minPosXs.get(inPosition), + inState.maxNegXs.get(inPosition), + inState.maxPosXs.get(inPosition), + inState.maxYs.get(inPosition), + inState.minYs.get(inPosition) + ); + } + } + + public void add(int groupId, long encoded) { + int x = POINT_TYPE.extractX(encoded); + int y = POINT_TYPE.extractY(encoded); + add(groupId, x, x, x, x, y, y); + } + + public void add(int groupId, int minNegX, int minPosX, int maxNegX, int maxPosX, int maxY, int minY) { + ensureCapacity(groupId); + if (hasValue(groupId)) { + minNegXs.set(groupId, Math.min(minNegXs.get(groupId), minNegX)); + minPosXs.set(groupId, SpatialAggregationUtils.minPos(minPosXs.get(groupId), minPosX)); + maxNegXs.set(groupId, SpatialAggregationUtils.maxNeg(maxNegXs.get(groupId), maxNegX)); + maxPosXs.set(groupId, Math.max(maxPosXs.get(groupId), maxPosX)); + maxYs.set(groupId, Math.max(maxYs.get(groupId), maxY)); + minYs.set(groupId, Math.min(minYs.get(groupId), minY)); + } else { + minNegXs.set(groupId, minNegX); + minPosXs.set(groupId, minPosX); + maxNegXs.set(groupId, maxNegX); + maxPosXs.set(groupId, maxPosX); + maxYs.set(groupId, maxY); + minYs.set(groupId, minY); + } + assert minNegX <= 0 == maxNegX <= 0 : "minNegX=" + minNegX + " maxNegX=" + maxNegX; + assert minPosX >= 0 == maxPosX >= 0 : "minPosX=" + minPosX + " maxPosX=" + maxPosX; + trackGroupId(groupId); + } + + private void ensureCapacity(int groupId) { + long requiredSize = groupId + 1; + if (minNegXs.size() < requiredSize) { + minNegXs = bigArrays.grow(minNegXs, requiredSize); + minPosXs = bigArrays.grow(minPosXs, requiredSize); + maxNegXs = bigArrays.grow(maxNegXs, requiredSize); + maxPosXs = bigArrays.grow(maxPosXs, requiredSize); + minYs = bigArrays.grow(minYs, requiredSize); + maxYs = bigArrays.grow(maxYs, requiredSize); + } + } + + public Block toBlock(IntVector selected, DriverContext driverContext) { + try (var builder = driverContext.blockFactory().newBytesRefBlockBuilder(selected.getPositionCount())) { + for (int i = 0; i < selected.getPositionCount(); i++) { + int si = selected.getInt(i); + if (hasValue(si)) { + builder.appendBytesRef( + new BytesRef( + WellKnownBinary.toWKB( + SpatialAggregationUtils.asRectangle( + minNegXs.get(si), + minPosXs.get(si), + maxNegXs.get(si), + maxPosXs.get(si), + maxYs.get(si), + minYs.get(si) + ), + ByteOrder.LITTLE_ENDIAN + ) + ) + ); + } else { + builder.appendNull(); + } + } + return builder.build(); + } + } +} diff --git a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/spatial/SpatialExtentLongitudeWrappingAggregator.java b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/spatial/SpatialExtentLongitudeWrappingAggregator.java new file mode 100644 index 0000000000000..80ba2d5e45658 --- /dev/null +++ b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/spatial/SpatialExtentLongitudeWrappingAggregator.java @@ -0,0 +1,62 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +package org.elasticsearch.compute.aggregation.spatial; + +import org.elasticsearch.compute.data.Block; +import org.elasticsearch.compute.data.IntVector; +import org.elasticsearch.compute.operator.DriverContext; + +// A bit of abuse of notation here, since we're extending this class to "inherit" its static methods. +// Unfortunately, this is the way it has to be done, since the generated code invokes these methods statically. +abstract class SpatialExtentLongitudeWrappingAggregator { + public static void combineIntermediate( + SpatialExtentStateWrappedLongitudeState current, + int minNegX, + int minPosX, + int maxNegX, + int maxPosX, + int maxY, + int minY + ) { + current.add(minNegX, minPosX, maxNegX, maxPosX, maxY, minY); + } + + public static void combineIntermediate( + SpatialExtentGroupingStateWrappedLongitudeState current, + int groupId, + int minNegX, + int minPosX, + int maxNegX, + int maxPosX, + int maxY, + int minY + ) { + current.add(groupId, minNegX, minPosX, maxNegX, maxPosX, maxY, minY); + } + + public static Block evaluateFinal(SpatialExtentStateWrappedLongitudeState state, DriverContext driverContext) { + return state.toBlock(driverContext); + } + + public static Block evaluateFinal( + SpatialExtentGroupingStateWrappedLongitudeState state, + IntVector selected, + DriverContext driverContext + ) { + return state.toBlock(selected, driverContext); + } + + public static void combineStates( + SpatialExtentGroupingStateWrappedLongitudeState current, + int groupId, + SpatialExtentGroupingStateWrappedLongitudeState inState, + int inPosition + ) { + current.add(groupId, inState, inPosition); + } +} diff --git a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/spatial/SpatialExtentState.java b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/spatial/SpatialExtentState.java new file mode 100644 index 0000000000000..0eea9b79f73ea --- /dev/null +++ b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/spatial/SpatialExtentState.java @@ -0,0 +1,82 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +package org.elasticsearch.compute.aggregation.spatial; + +import org.apache.lucene.util.BytesRef; +import org.elasticsearch.compute.aggregation.AggregatorState; +import org.elasticsearch.compute.data.Block; +import org.elasticsearch.compute.operator.DriverContext; +import org.elasticsearch.geometry.Geometry; +import org.elasticsearch.geometry.Rectangle; +import org.elasticsearch.geometry.utils.WellKnownBinary; + +import java.nio.ByteOrder; + +final class SpatialExtentState implements AggregatorState { + private final PointType pointType; + private boolean seen = false; + private int minX = Integer.MAX_VALUE; + private int maxX = Integer.MIN_VALUE; + private int maxY = Integer.MIN_VALUE; + private int minY = Integer.MAX_VALUE; + + SpatialExtentState(PointType pointType) { + this.pointType = pointType; + } + + @Override + public void close() {} + + @Override + public void toIntermediate(Block[] blocks, int offset, DriverContext driverContext) { + assert blocks.length >= offset + 4; + var blockFactory = driverContext.blockFactory(); + blocks[offset + 0] = blockFactory.newConstantIntBlockWith(minX, 1); + blocks[offset + 1] = blockFactory.newConstantIntBlockWith(maxX, 1); + blocks[offset + 2] = blockFactory.newConstantIntBlockWith(maxY, 1); + blocks[offset + 3] = blockFactory.newConstantIntBlockWith(minY, 1); + } + + public void add(Geometry geo) { + pointType.computeEnvelope(geo) + .ifPresent( + r -> add( + pointType.encodeX(r.getMinX()), + pointType.encodeX(r.getMaxX()), + pointType.encodeY(r.getMaxY()), + pointType.encodeY(r.getMinY()) + ) + ); + } + + public void add(int minX, int maxX, int maxY, int minY) { + seen = true; + this.minX = Math.min(this.minX, minX); + this.maxX = Math.max(this.maxX, maxX); + this.maxY = Math.max(this.maxY, maxY); + this.minY = Math.min(this.minY, minY); + } + + public void add(long encoded) { + int x = pointType.extractX(encoded); + int y = pointType.extractY(encoded); + add(x, x, y, y); + } + + public Block toBlock(DriverContext driverContext) { + var factory = driverContext.blockFactory(); + return seen ? factory.newConstantBytesRefBlockWith(new BytesRef(toWKB()), 1) : factory.newConstantNullBlock(1); + } + + private byte[] toWKB() { + return WellKnownBinary.toWKB( + new Rectangle(pointType.decodeX(minX), pointType.decodeX(maxX), pointType.decodeY(maxY), pointType.decodeY(minY)), + ByteOrder.LITTLE_ENDIAN + ); + } +} diff --git a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/spatial/SpatialExtentStateWrappedLongitudeState.java b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/spatial/SpatialExtentStateWrappedLongitudeState.java new file mode 100644 index 0000000000000..99200d2ed99f5 --- /dev/null +++ b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/spatial/SpatialExtentStateWrappedLongitudeState.java @@ -0,0 +1,91 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +package org.elasticsearch.compute.aggregation.spatial; + +import org.apache.lucene.util.BytesRef; +import org.elasticsearch.compute.aggregation.AggregatorState; +import org.elasticsearch.compute.data.Block; +import org.elasticsearch.compute.operator.DriverContext; +import org.elasticsearch.geometry.Geometry; +import org.elasticsearch.geometry.utils.SpatialEnvelopeVisitor; +import org.elasticsearch.geometry.utils.WellKnownBinary; + +import java.nio.ByteOrder; + +final class SpatialExtentStateWrappedLongitudeState implements AggregatorState { + // Only geo points support longitude wrapping. + private static final PointType POINT_TYPE = PointType.GEO; + private boolean seen = false; + private int minNegX = SpatialAggregationUtils.DEFAULT_NEG; + private int minPosX = SpatialAggregationUtils.DEFAULT_POS; + private int maxNegX = SpatialAggregationUtils.DEFAULT_NEG; + private int maxPosX = SpatialAggregationUtils.DEFAULT_POS; + private int maxY = Integer.MIN_VALUE; + private int minY = Integer.MAX_VALUE; + + private GeoPointEnvelopeVisitor geoPointVisitor = new GeoPointEnvelopeVisitor(); + + @Override + public void close() {} + + @Override + public void toIntermediate(Block[] blocks, int offset, DriverContext driverContext) { + assert blocks.length >= offset + 6; + var blockFactory = driverContext.blockFactory(); + blocks[offset + 0] = blockFactory.newConstantIntBlockWith(minNegX, 1); + blocks[offset + 1] = blockFactory.newConstantIntBlockWith(minPosX, 1); + blocks[offset + 2] = blockFactory.newConstantIntBlockWith(maxNegX, 1); + blocks[offset + 3] = blockFactory.newConstantIntBlockWith(maxPosX, 1); + blocks[offset + 4] = blockFactory.newConstantIntBlockWith(maxY, 1); + blocks[offset + 5] = blockFactory.newConstantIntBlockWith(minY, 1); + } + + public void add(Geometry geo) { + geoPointVisitor.reset(); + if (geo.visit(new SpatialEnvelopeVisitor(geoPointVisitor))) { + add( + SpatialAggregationUtils.encodeNegativeLongitude(geoPointVisitor.getMinNegX()), + SpatialAggregationUtils.encodePositiveLongitude(geoPointVisitor.getMinPosX()), + SpatialAggregationUtils.encodeNegativeLongitude(geoPointVisitor.getMaxNegX()), + SpatialAggregationUtils.encodePositiveLongitude(geoPointVisitor.getMaxPosX()), + POINT_TYPE.encodeY(geoPointVisitor.getMaxY()), + POINT_TYPE.encodeY(geoPointVisitor.getMinY()) + ); + } + } + + public void add(int minNegX, int minPosX, int maxNegX, int maxPosX, int maxY, int minY) { + seen = true; + this.minNegX = Math.min(this.minNegX, minNegX); + this.minPosX = SpatialAggregationUtils.minPos(this.minPosX, minPosX); + this.maxNegX = SpatialAggregationUtils.maxNeg(this.maxNegX, maxNegX); + this.maxPosX = Math.max(this.maxPosX, maxPosX); + this.maxY = Math.max(this.maxY, maxY); + this.minY = Math.min(this.minY, minY); + assert this.minNegX <= 0 == this.maxNegX <= 0 : "minNegX=" + this.minNegX + " maxNegX=" + this.maxNegX; + assert this.minPosX >= 0 == this.maxPosX >= 0 : "minPosX=" + this.minPosX + " maxPosX=" + this.maxPosX; + } + + public void add(long encoded) { + int x = POINT_TYPE.extractX(encoded); + int y = POINT_TYPE.extractY(encoded); + add(x, x, x, x, y, y); + } + + public Block toBlock(DriverContext driverContext) { + var factory = driverContext.blockFactory(); + return seen ? factory.newConstantBytesRefBlockWith(new BytesRef(toWKB()), 1) : factory.newConstantNullBlock(1); + } + + private byte[] toWKB() { + return WellKnownBinary.toWKB( + SpatialAggregationUtils.asRectangle(minNegX, minPosX, maxNegX, maxPosX, maxY, minY), + ByteOrder.LITTLE_ENDIAN + ); + } +} diff --git a/x-pack/plugin/esql/qa/testFixtures/src/main/resources/mapping-airports_no_doc_values.json b/x-pack/plugin/esql/qa/testFixtures/src/main/resources/mapping-airports_no_doc_values.json index d7097f89a17df..782fd40712f43 100644 --- a/x-pack/plugin/esql/qa/testFixtures/src/main/resources/mapping-airports_no_doc_values.json +++ b/x-pack/plugin/esql/qa/testFixtures/src/main/resources/mapping-airports_no_doc_values.json @@ -24,7 +24,9 @@ "type": "keyword" }, "city_location": { - "type": "geo_point" + "type": "geo_point", + "index": true, + "doc_values": false } } } diff --git a/x-pack/plugin/esql/qa/testFixtures/src/main/resources/spatial.csv-spec b/x-pack/plugin/esql/qa/testFixtures/src/main/resources/spatial.csv-spec index ac9948c90f5e9..8694c973448e9 100644 --- a/x-pack/plugin/esql/qa/testFixtures/src/main/resources/spatial.csv-spec +++ b/x-pack/plugin/esql/qa/testFixtures/src/main/resources/spatial.csv-spec @@ -519,6 +519,63 @@ centroid:geo_point | count:long POINT (42.97109629958868 14.7552534006536) | 1 ; +############################################### +# Tests for ST_EXTENT_AGG on GEO_POINT type + +stExtentSingleGeoPoint +required_capability: st_extent_agg +ROW point = TO_GEOPOINT("POINT(42.97109629958868 14.7552534006536)") +| STATS extent = ST_EXTENT_AGG(point) +; + +extent:geo_shape +BBOX(42.97109629958868, 42.97109629958868, 14.7552534006536, 14.7552534006536) +; + +stExtentMultipleGeoPoints +required_capability: st_extent_agg +// tag::st_extent_agg-airports[] +FROM airports +| WHERE country == "India" +| STATS extent = ST_EXTENT_AGG(location) +// end::st_extent_agg-airports[] +; + +// tag::st_extent_agg-airports-result[] +extent:geo_shape +BBOX (70.77995480038226, 91.5882289968431, 33.9830909203738, 8.47650992218405) +// end::st_extent_agg-airports-result[] +; + +stExtentMultipleGeoPointsNoDocValues +required_capability: st_extent_agg +FROM airports_no_doc_values | WHERE country == "India" | STATS extent = ST_EXTENT_AGG(location) +; + +extent:geo_shape +BBOX (70.77995480038226, 91.5882289968431, 33.9830909203738, 8.47650992218405) +; + +stExtentMultipleGeoPointGrouping +required_capability: st_extent_agg +FROM airports | STATS extent = ST_EXTENT_AGG(location) BY country | SORT country | LIMIT 3 +; + +extent:geo_shape | country:keyword +BBOX (69.2100736219436, 69.2100736219436, 34.56339786294848, 34.56339786294848) | Afghanistan +BBOX (19.715032372623682, 19.715032372623682, 41.4208514476195, 41.4208514476195) | Albania +BBOX (-0.6067969836294651, 6.621946580708027, 36.69972063973546, 35.62027471605688) | Algeria +; + +stExtentGeoShapes +required_capability: st_extent_agg +FROM airport_city_boundaries | WHERE region == "City of New York" | STATS extent = ST_EXTENT_AGG(city_boundary) +; + +extent:geo_shape +BBOX (-74.25880000926554, -73.70020005851984, 40.91759996954352, 40.47659996431321) +; + ############################################### # Tests for ST_INTERSECTS on GEO_POINT type @@ -1698,6 +1755,48 @@ centroid:cartesian_point | count:long POINT (726480.0130685265 3359566.331716279) | 849 ; +############################################### +# Tests for ST_EXTENT_AGG on CARTESIAN_POINT type + +stExtentSingleCartesianPoint +required_capability: st_extent_agg +ROW point = TO_CARTESIANPOINT("POINT(429.7109629958868 147.552534006536)") +| STATS extent = ST_EXTENT_AGG(point) +; + +extent:cartesian_shape +BBOX (429.7109680175781, 429.7109680175781, 147.5525360107422, 147.5525360107422) +; + +stExtentMultipleCartesianPoints +required_capability: st_extent_agg +FROM airports_web | WHERE scalerank == 9 | STATS extent = ST_EXTENT_AGG(location) +; + +extent:cartesian_shape +BBOX (4783520.5, 1.6168486E7, 8704352.0, -584415.9375) +; + +stExtentMultipleCartesianPointGrouping +required_capability: st_extent_agg +FROM airports_web | STATS extent = ST_EXTENT_AGG(location) BY scalerank | SORT scalerank DESC | LIMIT 3 +; + +extent:cartesian_shape | scalerank:integer +BBOX (4783520.5, 1.6168486E7, 8704352.0, -584415.9375) | 9 +BBOX (-1.936604E7, 1.8695374E7, 1.4502138E7, -3943067.25) | 8 +BBOX (-1.891609E7, 1.9947946E7, 8455470.0, -7128878.5) | 7 +; + +stExtentCartesianShapes +required_capability: st_extent_agg +FROM cartesian_multipolygons | STATS extent = ST_EXTENT_AGG(shape) +; + +extent:cartesian_shape +BBOX (0.0, 3.0, 3.0, 0.0) +; + ############################################### # Tests for ST_INTERSECTS on CARTESIAN_POINT type diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/EsqlCapabilities.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/EsqlCapabilities.java index ddabb3e937dc2..4cf3162fcca3b 100644 --- a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/EsqlCapabilities.java +++ b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/EsqlCapabilities.java @@ -189,6 +189,9 @@ public enum Cap { */ ST_DISTANCE, + /** Support for function {@code ST_EXTENT}. */ + ST_EXTENT_AGG, + /** * Fix determination of CRS types in spatial functions when folding. */ diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/EsqlFunctionRegistry.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/EsqlFunctionRegistry.java index 3749b46879354..50d0d2438d8a1 100644 --- a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/EsqlFunctionRegistry.java +++ b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/EsqlFunctionRegistry.java @@ -28,6 +28,7 @@ import org.elasticsearch.xpack.esql.expression.function.aggregate.Percentile; import org.elasticsearch.xpack.esql.expression.function.aggregate.Rate; import org.elasticsearch.xpack.esql.expression.function.aggregate.SpatialCentroid; +import org.elasticsearch.xpack.esql.expression.function.aggregate.SpatialExtent; import org.elasticsearch.xpack.esql.expression.function.aggregate.StdDev; import org.elasticsearch.xpack.esql.expression.function.aggregate.Sum; import org.elasticsearch.xpack.esql.expression.function.aggregate.Top; @@ -353,6 +354,7 @@ private static FunctionDefinition[][] functions() { new FunctionDefinition[] { def(SpatialCentroid.class, SpatialCentroid::new, "st_centroid_agg"), def(SpatialContains.class, SpatialContains::new, "st_contains"), + def(SpatialExtent.class, SpatialExtent::new, "st_extent_agg"), def(SpatialDisjoint.class, SpatialDisjoint::new, "st_disjoint"), def(SpatialIntersects.class, SpatialIntersects::new, "st_intersects"), def(SpatialWithin.class, SpatialWithin::new, "st_within"), diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/aggregate/AggregateWritables.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/aggregate/AggregateWritables.java index d74b5c8b386b8..db1d2a9e6f254 100644 --- a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/aggregate/AggregateWritables.java +++ b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/aggregate/AggregateWritables.java @@ -25,6 +25,7 @@ public static List getNamedWriteables() { Percentile.ENTRY, Rate.ENTRY, SpatialCentroid.ENTRY, + SpatialExtent.ENTRY, StdDev.ENTRY, Sum.ENTRY, Top.ENTRY, diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/aggregate/SpatialAggregateFunction.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/aggregate/SpatialAggregateFunction.java index 87eec540932b1..35f99e4b648df 100644 --- a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/aggregate/SpatialAggregateFunction.java +++ b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/aggregate/SpatialAggregateFunction.java @@ -8,6 +8,9 @@ package org.elasticsearch.xpack.esql.expression.function.aggregate; import org.elasticsearch.common.io.stream.StreamInput; +import org.elasticsearch.index.mapper.MappedFieldType.FieldExtractPreference; +import org.elasticsearch.license.License; +import org.elasticsearch.license.XPackLicenseState; import org.elasticsearch.xpack.esql.core.expression.Expression; import org.elasticsearch.xpack.esql.core.tree.Source; @@ -22,26 +25,34 @@ * select the best one. */ public abstract class SpatialAggregateFunction extends AggregateFunction { - protected final boolean useDocValues; + protected final FieldExtractPreference fieldExtractPreference; - protected SpatialAggregateFunction(Source source, Expression field, Expression filter, boolean useDocValues) { + protected SpatialAggregateFunction(Source source, Expression field, Expression filter, FieldExtractPreference fieldExtractPreference) { super(source, field, filter, emptyList()); - this.useDocValues = useDocValues; + this.fieldExtractPreference = fieldExtractPreference; } - protected SpatialAggregateFunction(StreamInput in, boolean useDocValues) throws IOException { + protected SpatialAggregateFunction(StreamInput in, FieldExtractPreference fieldExtractPreference) throws IOException { super(in); - // The useDocValues field is only used on data nodes local planning, and therefor never serialized - this.useDocValues = useDocValues; + // The fieldExtractPreference field is only used on data nodes local planning, and therefore never serialized + this.fieldExtractPreference = fieldExtractPreference; } public abstract SpatialAggregateFunction withDocValues(); + @Override + public boolean checkLicense(XPackLicenseState state) { + return switch (field().dataType()) { + case GEO_SHAPE, CARTESIAN_SHAPE -> state.isAllowedByLicense(License.OperationMode.PLATINUM); + default -> true; + }; + } + @Override public int hashCode() { // NB: the hashcode is currently used for key generation so // to avoid clashes between aggs with the same arguments, add the class name as variation - return Objects.hash(getClass(), children(), useDocValues); + return Objects.hash(getClass(), children(), fieldExtractPreference); } @Override @@ -50,12 +61,12 @@ public boolean equals(Object obj) { SpatialAggregateFunction other = (SpatialAggregateFunction) obj; return Objects.equals(other.field(), field()) && Objects.equals(other.parameters(), parameters()) - && Objects.equals(other.useDocValues, useDocValues); + && Objects.equals(other.fieldExtractPreference, fieldExtractPreference); } return false; } - public boolean useDocValues() { - return useDocValues; + public FieldExtractPreference fieldExtractPreference() { + return fieldExtractPreference; } } diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/aggregate/SpatialCentroid.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/aggregate/SpatialCentroid.java index aad95c07e3492..84915d024ea82 100644 --- a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/aggregate/SpatialCentroid.java +++ b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/aggregate/SpatialCentroid.java @@ -13,6 +13,7 @@ import org.elasticsearch.compute.aggregation.spatial.SpatialCentroidCartesianPointSourceValuesAggregatorFunctionSupplier; import org.elasticsearch.compute.aggregation.spatial.SpatialCentroidGeoPointDocValuesAggregatorFunctionSupplier; import org.elasticsearch.compute.aggregation.spatial.SpatialCentroidGeoPointSourceValuesAggregatorFunctionSupplier; +import org.elasticsearch.index.mapper.MappedFieldType.FieldExtractPreference; import org.elasticsearch.xpack.esql.EsqlIllegalArgumentException; import org.elasticsearch.xpack.esql.core.expression.Expression; import org.elasticsearch.xpack.esql.core.expression.Literal; @@ -27,6 +28,7 @@ import java.io.IOException; import java.util.List; +import static org.elasticsearch.index.mapper.MappedFieldType.FieldExtractPreference.NONE; import static org.elasticsearch.xpack.esql.core.expression.TypeResolutions.ParamOrdinal.DEFAULT; import static org.elasticsearch.xpack.esql.expression.EsqlTypeResolutions.isSpatialPoint; @@ -47,15 +49,15 @@ public class SpatialCentroid extends SpatialAggregateFunction implements ToAggre examples = @Example(file = "spatial", tag = "st_centroid_agg-airports") ) public SpatialCentroid(Source source, @Param(name = "field", type = { "geo_point", "cartesian_point" }) Expression field) { - this(source, field, Literal.TRUE, false); + this(source, field, Literal.TRUE, NONE); } - private SpatialCentroid(Source source, Expression field, Expression filter, boolean useDocValues) { - super(source, field, filter, useDocValues); + private SpatialCentroid(Source source, Expression field, Expression filter, FieldExtractPreference preference) { + super(source, field, filter, preference); } private SpatialCentroid(StreamInput in) throws IOException { - super(in, false); + super(in, NONE); } @Override @@ -65,12 +67,12 @@ public String getWriteableName() { @Override public SpatialCentroid withFilter(Expression filter) { - return new SpatialCentroid(source(), field(), filter, useDocValues); + return new SpatialCentroid(source(), field(), filter, fieldExtractPreference); } @Override public SpatialCentroid withDocValues() { - return new SpatialCentroid(source(), field(), filter(), true); + return new SpatialCentroid(source(), field(), filter(), FieldExtractPreference.DOC_VALUES); } @Override @@ -98,23 +100,16 @@ public SpatialCentroid replaceChildren(List newChildren) { @Override public AggregatorFunctionSupplier supplier(List inputChannels) { DataType type = field().dataType(); - if (useDocValues) { - // When the points are read as doc-values (eg. from the index), feed them into the doc-values aggregator - if (type == DataType.GEO_POINT) { - return new SpatialCentroidGeoPointDocValuesAggregatorFunctionSupplier(inputChannels); - } - if (type == DataType.CARTESIAN_POINT) { - return new SpatialCentroidCartesianPointDocValuesAggregatorFunctionSupplier(inputChannels); - } - } else { - // When the points are read as WKB from source or as point literals, feed them into the source-values aggregator - if (type == DataType.GEO_POINT) { - return new SpatialCentroidGeoPointSourceValuesAggregatorFunctionSupplier(inputChannels); - } - if (type == DataType.CARTESIAN_POINT) { - return new SpatialCentroidCartesianPointSourceValuesAggregatorFunctionSupplier(inputChannels); - } - } - throw EsqlIllegalArgumentException.illegalDataType(type); + return switch (type) { + case DataType.GEO_POINT -> switch (fieldExtractPreference) { + case DOC_VALUES -> new SpatialCentroidGeoPointDocValuesAggregatorFunctionSupplier(inputChannels); + case NONE -> new SpatialCentroidGeoPointSourceValuesAggregatorFunctionSupplier(inputChannels); + }; + case DataType.CARTESIAN_POINT -> switch (fieldExtractPreference) { + case DOC_VALUES -> new SpatialCentroidCartesianPointDocValuesAggregatorFunctionSupplier(inputChannels); + case NONE -> new SpatialCentroidCartesianPointSourceValuesAggregatorFunctionSupplier(inputChannels); + }; + default -> throw EsqlIllegalArgumentException.illegalDataType(type); + }; } } diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/aggregate/SpatialExtent.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/aggregate/SpatialExtent.java new file mode 100644 index 0000000000000..5cc1701faf13a --- /dev/null +++ b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/aggregate/SpatialExtent.java @@ -0,0 +1,119 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ +package org.elasticsearch.xpack.esql.expression.function.aggregate; + +import org.elasticsearch.common.io.stream.NamedWriteableRegistry; +import org.elasticsearch.common.io.stream.StreamInput; +import org.elasticsearch.compute.aggregation.AggregatorFunctionSupplier; +import org.elasticsearch.compute.aggregation.spatial.SpatialExtentCartesianPointDocValuesAggregatorFunctionSupplier; +import org.elasticsearch.compute.aggregation.spatial.SpatialExtentCartesianPointSourceValuesAggregatorFunctionSupplier; +import org.elasticsearch.compute.aggregation.spatial.SpatialExtentCartesianShapeAggregatorFunctionSupplier; +import org.elasticsearch.compute.aggregation.spatial.SpatialExtentGeoPointDocValuesAggregatorFunctionSupplier; +import org.elasticsearch.compute.aggregation.spatial.SpatialExtentGeoPointSourceValuesAggregatorFunctionSupplier; +import org.elasticsearch.compute.aggregation.spatial.SpatialExtentGeoShapeAggregatorFunctionSupplier; +import org.elasticsearch.index.mapper.MappedFieldType.FieldExtractPreference; +import org.elasticsearch.xpack.esql.EsqlIllegalArgumentException; +import org.elasticsearch.xpack.esql.core.expression.Expression; +import org.elasticsearch.xpack.esql.core.expression.Literal; +import org.elasticsearch.xpack.esql.core.tree.NodeInfo; +import org.elasticsearch.xpack.esql.core.tree.Source; +import org.elasticsearch.xpack.esql.core.type.DataType; +import org.elasticsearch.xpack.esql.expression.function.Example; +import org.elasticsearch.xpack.esql.expression.function.FunctionInfo; +import org.elasticsearch.xpack.esql.expression.function.Param; +import org.elasticsearch.xpack.esql.planner.ToAggregator; + +import java.io.IOException; +import java.util.List; + +import static org.elasticsearch.xpack.esql.core.expression.TypeResolutions.ParamOrdinal.DEFAULT; +import static org.elasticsearch.xpack.esql.expression.EsqlTypeResolutions.isSpatial; + +/** + * Calculate spatial extent of all values of a field in matching documents. + */ +public final class SpatialExtent extends SpatialAggregateFunction implements ToAggregator { + public static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry( + Expression.class, + "SpatialExtent", + SpatialExtent::new + ); + + @FunctionInfo( + returnType = { "geo_shape", "cartesian_shape" }, + description = "Calculate the spatial extent over a field with geometry type. Returns a bounding box for all values of the field.", + isAggregation = true, + examples = @Example(file = "spatial", tag = "st_extent_agg-airports") + ) + public SpatialExtent( + Source source, + @Param(name = "field", type = { "geo_point", "cartesian_point", "geo_shape", "cartesian_shape" }) Expression field + ) { + this(source, field, Literal.TRUE, FieldExtractPreference.NONE); + } + + private SpatialExtent(Source source, Expression field, Expression filter, FieldExtractPreference preference) { + super(source, field, filter, preference); + } + + private SpatialExtent(StreamInput in) throws IOException { + super(in, FieldExtractPreference.NONE); + } + + @Override + public String getWriteableName() { + return ENTRY.name; + } + + @Override + public SpatialExtent withFilter(Expression filter) { + return new SpatialExtent(source(), field(), filter, fieldExtractPreference); + } + + @Override + public org.elasticsearch.xpack.esql.expression.function.aggregate.SpatialExtent withDocValues() { + return new SpatialExtent(source(), field(), filter(), FieldExtractPreference.DOC_VALUES); + } + + @Override + protected TypeResolution resolveType() { + return isSpatial(field(), sourceText(), DEFAULT); + } + + @Override + public DataType dataType() { + return DataType.isSpatialGeo(field().dataType()) ? DataType.GEO_SHAPE : DataType.CARTESIAN_SHAPE; + } + + @Override + protected NodeInfo info() { + return NodeInfo.create(this, SpatialExtent::new, field()); + } + + @Override + public SpatialExtent replaceChildren(List newChildren) { + return new SpatialExtent(source(), newChildren.get(0)); + } + + @Override + public AggregatorFunctionSupplier supplier(List inputChannels) { + return switch (field().dataType()) { + case DataType.GEO_POINT -> switch (fieldExtractPreference) { + case DOC_VALUES -> new SpatialExtentGeoPointDocValuesAggregatorFunctionSupplier(inputChannels); + case NONE -> new SpatialExtentGeoPointSourceValuesAggregatorFunctionSupplier(inputChannels); + }; + case DataType.CARTESIAN_POINT -> switch (fieldExtractPreference) { + case DOC_VALUES -> new SpatialExtentCartesianPointDocValuesAggregatorFunctionSupplier(inputChannels); + case NONE -> new SpatialExtentCartesianPointSourceValuesAggregatorFunctionSupplier(inputChannels); + }; + // Shapes don't differentiate between source and doc values. + case DataType.GEO_SHAPE -> new SpatialExtentGeoShapeAggregatorFunctionSupplier(inputChannels); + case DataType.CARTESIAN_SHAPE -> new SpatialExtentCartesianShapeAggregatorFunctionSupplier(inputChannels); + default -> throw EsqlIllegalArgumentException.illegalDataType(field().dataType()); + }; + } +} diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/spatial/StEnvelope.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/spatial/StEnvelope.java index 934991f3a8088..ca243efcc2851 100644 --- a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/spatial/StEnvelope.java +++ b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/spatial/StEnvelope.java @@ -14,6 +14,7 @@ import org.elasticsearch.compute.operator.EvalOperator; import org.elasticsearch.geometry.Point; import org.elasticsearch.geometry.utils.SpatialEnvelopeVisitor; +import org.elasticsearch.geometry.utils.SpatialEnvelopeVisitor.WrapLongitude; import org.elasticsearch.xpack.esql.core.expression.Expression; import org.elasticsearch.xpack.esql.core.expression.TypeResolutions; import org.elasticsearch.xpack.esql.core.tree.NodeInfo; @@ -129,7 +130,7 @@ static BytesRef fromWellKnownBinaryGeo(BytesRef wkb) { if (geometry instanceof Point) { return wkb; } - var envelope = SpatialEnvelopeVisitor.visitGeo(geometry, true); + var envelope = SpatialEnvelopeVisitor.visitGeo(geometry, WrapLongitude.WRAP); if (envelope.isPresent()) { return UNSPECIFIED.asWkb(envelope.get()); } diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/spatial/StXMax.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/spatial/StXMax.java index d6d710b175113..69eede1c5fac5 100644 --- a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/spatial/StXMax.java +++ b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/spatial/StXMax.java @@ -14,6 +14,7 @@ import org.elasticsearch.compute.operator.EvalOperator; import org.elasticsearch.geometry.Point; import org.elasticsearch.geometry.utils.SpatialEnvelopeVisitor; +import org.elasticsearch.geometry.utils.SpatialEnvelopeVisitor.WrapLongitude; import org.elasticsearch.xpack.esql.core.expression.Expression; import org.elasticsearch.xpack.esql.core.expression.TypeResolutions; import org.elasticsearch.xpack.esql.core.tree.NodeInfo; @@ -114,7 +115,7 @@ static double fromWellKnownBinaryGeo(BytesRef wkb) { if (geometry instanceof Point point) { return point.getX(); } - var envelope = SpatialEnvelopeVisitor.visitGeo(geometry, true); + var envelope = SpatialEnvelopeVisitor.visitGeo(geometry, WrapLongitude.WRAP); if (envelope.isPresent()) { return envelope.get().getMaxX(); } diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/spatial/StXMin.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/spatial/StXMin.java index a5fa11bc11b0f..b29a547ab0af6 100644 --- a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/spatial/StXMin.java +++ b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/spatial/StXMin.java @@ -14,6 +14,7 @@ import org.elasticsearch.compute.operator.EvalOperator; import org.elasticsearch.geometry.Point; import org.elasticsearch.geometry.utils.SpatialEnvelopeVisitor; +import org.elasticsearch.geometry.utils.SpatialEnvelopeVisitor.WrapLongitude; import org.elasticsearch.xpack.esql.core.expression.Expression; import org.elasticsearch.xpack.esql.core.expression.TypeResolutions; import org.elasticsearch.xpack.esql.core.tree.NodeInfo; @@ -114,7 +115,7 @@ static double fromWellKnownBinaryGeo(BytesRef wkb) { if (geometry instanceof Point point) { return point.getX(); } - var envelope = SpatialEnvelopeVisitor.visitGeo(geometry, true); + var envelope = SpatialEnvelopeVisitor.visitGeo(geometry, WrapLongitude.WRAP); if (envelope.isPresent()) { return envelope.get().getMinX(); } diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/spatial/StYMax.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/spatial/StYMax.java index fbbea8e024a6b..981b500bcaef7 100644 --- a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/spatial/StYMax.java +++ b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/spatial/StYMax.java @@ -14,6 +14,7 @@ import org.elasticsearch.compute.operator.EvalOperator; import org.elasticsearch.geometry.Point; import org.elasticsearch.geometry.utils.SpatialEnvelopeVisitor; +import org.elasticsearch.geometry.utils.SpatialEnvelopeVisitor.WrapLongitude; import org.elasticsearch.xpack.esql.core.expression.Expression; import org.elasticsearch.xpack.esql.core.expression.TypeResolutions; import org.elasticsearch.xpack.esql.core.tree.NodeInfo; @@ -114,7 +115,7 @@ static double fromWellKnownBinaryGeo(BytesRef wkb) { if (geometry instanceof Point point) { return point.getY(); } - var envelope = SpatialEnvelopeVisitor.visitGeo(geometry, true); + var envelope = SpatialEnvelopeVisitor.visitGeo(geometry, WrapLongitude.WRAP); if (envelope.isPresent()) { return envelope.get().getMaxY(); } diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/spatial/StYMin.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/spatial/StYMin.java index 1707d3b4f2fb9..882aeb30afaee 100644 --- a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/spatial/StYMin.java +++ b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/spatial/StYMin.java @@ -14,6 +14,7 @@ import org.elasticsearch.compute.operator.EvalOperator; import org.elasticsearch.geometry.Point; import org.elasticsearch.geometry.utils.SpatialEnvelopeVisitor; +import org.elasticsearch.geometry.utils.SpatialEnvelopeVisitor.WrapLongitude; import org.elasticsearch.xpack.esql.core.expression.Expression; import org.elasticsearch.xpack.esql.core.expression.TypeResolutions; import org.elasticsearch.xpack.esql.core.tree.NodeInfo; @@ -114,7 +115,7 @@ static double fromWellKnownBinaryGeo(BytesRef wkb) { if (geometry instanceof Point point) { return point.getY(); } - var envelope = SpatialEnvelopeVisitor.visitGeo(geometry, true); + var envelope = SpatialEnvelopeVisitor.visitGeo(geometry, WrapLongitude.WRAP); if (envelope.isPresent()) { return envelope.get().getMinY(); } diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/AggregateMapper.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/AggregateMapper.java index 1f55e293b8e75..1918e3036e2b0 100644 --- a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/AggregateMapper.java +++ b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/AggregateMapper.java @@ -35,6 +35,7 @@ import org.elasticsearch.xpack.esql.expression.function.aggregate.Rate; import org.elasticsearch.xpack.esql.expression.function.aggregate.SpatialAggregateFunction; import org.elasticsearch.xpack.esql.expression.function.aggregate.SpatialCentroid; +import org.elasticsearch.xpack.esql.expression.function.aggregate.SpatialExtent; import org.elasticsearch.xpack.esql.expression.function.aggregate.StdDev; import org.elasticsearch.xpack.esql.expression.function.aggregate.Sum; import org.elasticsearch.xpack.esql.expression.function.aggregate.ToPartial; @@ -66,7 +67,7 @@ final class AggregateMapper { private static final List NUMERIC = List.of("Int", "Long", "Double"); - private static final List SPATIAL = List.of("GeoPoint", "CartesianPoint"); + private static final List SPATIAL_EXTRA_CONFIGS = List.of("SourceValues", "DocValues"); /** List of all mappable ESQL agg functions (excludes surrogates like AVG = SUM/COUNT). */ private static final List> AGG_FUNCTIONS = List.of( @@ -77,6 +78,7 @@ final class AggregateMapper { Min.class, Percentile.class, SpatialCentroid.class, + SpatialExtent.class, StdDev.class, Sum.class, Values.class, @@ -89,7 +91,11 @@ final class AggregateMapper { ); /** Record of agg Class, type, and grouping (or non-grouping). */ - private record AggDef(Class aggClazz, String type, String extra, boolean grouping) {} + private record AggDef(Class aggClazz, String type, String extra, boolean grouping) { + public AggDef withoutExtra() { + return new AggDef(aggClazz, type, "", grouping); + } + } /** Map of AggDef types to intermediate named expressions. */ private static final Map> MAPPER = AGG_FUNCTIONS.stream() @@ -145,7 +151,7 @@ private static List entryForAgg(String aggAlias, AggregateFunct var aggDef = new AggDef( aggregateFunction.getClass(), dataTypeToString(aggregateFunction.field().dataType(), aggregateFunction.getClass()), - aggregateFunction instanceof SpatialCentroid ? "SourceValues" : "", + aggregateFunction instanceof SpatialAggregateFunction ? "SourceValues" : "", grouping ); var is = getNonNull(aggDef); @@ -154,7 +160,7 @@ private static List entryForAgg(String aggAlias, AggregateFunct /** Gets the agg from the mapper - wrapper around map::get for more informative failure.*/ private static List getNonNull(AggDef aggDef) { - var l = MAPPER.get(aggDef); + var l = MAPPER.getOrDefault(aggDef, MAPPER.get(aggDef.withoutExtra())); if (l == null) { throw new EsqlIllegalArgumentException("Cannot find intermediate state for: " + aggDef); } @@ -170,9 +176,14 @@ private static Stream, Tuple>> typeAndNames(Class types = List.of("Boolean", "Int", "Long", "Double", "Ip", "BytesRef"); } else if (clazz == Count.class) { types = List.of(""); // no extra type distinction - } else if (SpatialAggregateFunction.class.isAssignableFrom(clazz)) { - types = SPATIAL; - extraConfigs = List.of("SourceValues", "DocValues"); + } else if (clazz == SpatialCentroid.class) { + types = List.of("GeoPoint", "CartesianPoint"); + extraConfigs = SPATIAL_EXTRA_CONFIGS; + } else if (clazz == SpatialExtent.class) { + return Stream.concat( + combine(clazz, List.of("GeoPoint", "CartesianPoint"), SPATIAL_EXTRA_CONFIGS), + combine(clazz, List.of("GeoShape", "CartesianShape"), List.of("")) + ); } else if (Values.class.isAssignableFrom(clazz)) { // TODO can't we figure this out from the function itself? types = List.of("Int", "Long", "Double", "Boolean", "BytesRef"); @@ -188,6 +199,10 @@ private static Stream, Tuple>> typeAndNames(Class assert false : "unknown aggregate type " + clazz; throw new IllegalArgumentException("unknown aggregate type " + clazz); } + return combine(clazz, types, extraConfigs); + } + + private static Stream, Tuple>> combine(Class clazz, List types, List extraConfigs) { return combinations(types, extraConfigs).map(combo -> new Tuple<>(clazz, combo)); } @@ -219,6 +234,15 @@ private static List lookupIntermediateState(AggDef aggDef /** Looks up the intermediate state method for a given class, type, and grouping. */ private static MethodHandle lookup(Class clazz, String type, String extra, boolean grouping) { + try { + return lookupRetry(clazz, type, extra, grouping); + } catch (IllegalAccessException | NoSuchMethodException | ClassNotFoundException e) { + throw new EsqlIllegalArgumentException(e); + } + } + + private static MethodHandle lookupRetry(Class clazz, String type, String extra, boolean grouping) throws IllegalAccessException, + NoSuchMethodException, ClassNotFoundException { try { return MethodHandles.lookup() .findStatic( @@ -226,8 +250,14 @@ private static MethodHandle lookup(Class clazz, String type, String extra, bo "intermediateStateDesc", MethodType.methodType(List.class) ); - } catch (IllegalAccessException | NoSuchMethodException | ClassNotFoundException e) { - throw new EsqlIllegalArgumentException(e); + } catch (NoSuchMethodException ignore) { + // Retry without the extra information. + return MethodHandles.lookup() + .findStatic( + Class.forName(determineAggName(clazz, type, "", grouping)), + "intermediateStateDesc", + MethodType.methodType(List.class) + ); } } @@ -301,8 +331,10 @@ private static String dataTypeToString(DataType type, Class aggClass) { case DataType.KEYWORD, DataType.IP, DataType.VERSION, DataType.TEXT, DataType.SEMANTIC_TEXT -> "BytesRef"; case GEO_POINT -> "GeoPoint"; case CARTESIAN_POINT -> "CartesianPoint"; + case GEO_SHAPE -> "GeoShape"; + case CARTESIAN_SHAPE -> "CartesianShape"; case UNSUPPORTED, NULL, UNSIGNED_LONG, SHORT, BYTE, FLOAT, HALF_FLOAT, SCALED_FLOAT, OBJECT, SOURCE, DATE_PERIOD, TIME_DURATION, - CARTESIAN_SHAPE, GEO_SHAPE, DOC_DATA_TYPE, TSID_DATA_TYPE, PARTIAL_AGG -> throw new EsqlIllegalArgumentException( + DOC_DATA_TYPE, TSID_DATA_TYPE, PARTIAL_AGG -> throw new EsqlIllegalArgumentException( "illegal agg type: " + type.typeName() ); }; diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/RectangleMatcher.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/RectangleMatcher.java new file mode 100644 index 0000000000000..48fbc9c8e0378 --- /dev/null +++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/RectangleMatcher.java @@ -0,0 +1,61 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +package org.elasticsearch.xpack.esql.expression; + +import org.elasticsearch.compute.aggregation.spatial.PointType; +import org.elasticsearch.geometry.Rectangle; +import org.hamcrest.Description; +import org.hamcrest.Matchers; +import org.hamcrest.TypeSafeMatcher; + +/** + * Example usage: assertThat(actualRectangle, RectangleMatcher.closeTo(expectedRectangle, 0.0001, PointType.CARTESIAN));, or it + * can be used as a parameter to {@link WellKnownBinaryBytesRefMatcher}. + */ +public class RectangleMatcher extends TypeSafeMatcher { + private final Rectangle r; + private final PointType pointType; + private final double error; + + public static TypeSafeMatcher closeTo(Rectangle r, double error, PointType pointType) { + return new RectangleMatcher(r, error, pointType); + } + + private RectangleMatcher(Rectangle r, double error, PointType pointType) { + this.r = r; + this.pointType = pointType; + this.error = error; + } + + @Override + protected boolean matchesSafely(Rectangle other) { + // For geo bounds, longitude of (-180, 180) and (epsilon, -epsilon) are actually very close, since both encompass the entire globe. + boolean wrapAroundWorkAround = pointType == PointType.GEO && r.getMinX() >= r.getMaxX(); + boolean matchMinX = Matchers.closeTo(r.getMinX(), error).matches(other.getMinX()) + || (wrapAroundWorkAround && Matchers.closeTo(r.getMinX() - 180, error).matches(other.getMinX())) + || (wrapAroundWorkAround && Matchers.closeTo(r.getMinX(), error).matches(other.getMinX() - 180)); + boolean matchMaxX = Matchers.closeTo(r.getMaxX(), error).matches(other.getMaxX()) + || (wrapAroundWorkAround && Matchers.closeTo(r.getMaxX() + 180, error).matches(other.getMaxX())) + || (wrapAroundWorkAround && Matchers.closeTo(r.getMaxX(), error).matches(other.getMaxX() + 180)); + + return matchMinX + && matchMaxX + && Matchers.closeTo(r.getMaxY(), error).matches(other.getMaxY()) + && Matchers.closeTo(r.getMinY(), error).matches(other.getMinY()); + } + + @Override + public void describeMismatchSafely(Rectangle rectangle, Description description) { + description.appendText("was ").appendValue(rectangle); + } + + @Override + public void describeTo(Description description) { + description.appendValue(" " + r); + } +} diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/WellKnownBinaryBytesRefMatcher.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/WellKnownBinaryBytesRefMatcher.java new file mode 100644 index 0000000000000..535bb820458cd --- /dev/null +++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/WellKnownBinaryBytesRefMatcher.java @@ -0,0 +1,45 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +package org.elasticsearch.xpack.esql.expression; + +import org.apache.lucene.util.BytesRef; +import org.elasticsearch.geometry.Geometry; +import org.elasticsearch.geometry.utils.GeometryValidator; +import org.elasticsearch.geometry.utils.WellKnownBinary; +import org.hamcrest.Description; +import org.hamcrest.Matcher; +import org.hamcrest.TypeSafeMatcher; + +/** A wrapper for matching geometries encoded as WKB in a BytesRef. */ +public class WellKnownBinaryBytesRefMatcher extends TypeSafeMatcher { + private final Matcher matcher; + + public WellKnownBinaryBytesRefMatcher(Matcher matcher) { + this.matcher = matcher; + } + + @Override + public boolean matchesSafely(BytesRef bytesRef) { + return matcher.matches(fromBytesRef(bytesRef)); + } + + @Override + public void describeMismatchSafely(BytesRef bytesRef, Description description) { + matcher.describeMismatch(fromBytesRef(bytesRef), description); + } + + @SuppressWarnings("unchecked") + private G fromBytesRef(BytesRef bytesRef) { + return (G) WellKnownBinary.fromWKB(GeometryValidator.NOOP, false /* coerce */, bytesRef.bytes, bytesRef.offset, bytesRef.length); + } + + @Override + public void describeTo(Description description) { + matcher.describeTo(description); + } +} diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/AbstractAggregationTestCase.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/AbstractAggregationTestCase.java index df1675ba22568..c086245d6fd61 100644 --- a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/AbstractAggregationTestCase.java +++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/AbstractAggregationTestCase.java @@ -277,9 +277,11 @@ private void evaluate(Expression evaluableExpression) { } private void resolveExpression(Expression expression, Consumer onAggregator, Consumer onEvaluableExpression) { - logger.info( - "Test Values: " + testCase.getData().stream().map(TestCaseSupplier.TypedData::toString).collect(Collectors.joining(",")) - ); + String valuesString = testCase.getData().stream().map(TestCaseSupplier.TypedData::toString).collect(Collectors.joining(",")); + if (valuesString.length() > 200) { + valuesString = valuesString.substring(0, 200) + "..."; + } + logger.info("Test Values: " + valuesString); if (testCase.getExpectedTypeError() != null) { assertTypeResolutionFailure(expression); return; diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/MultiRowTestCaseSupplier.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/MultiRowTestCaseSupplier.java index 775ca45bfa124..bb0d2e57c3440 100644 --- a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/MultiRowTestCaseSupplier.java +++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/MultiRowTestCaseSupplier.java @@ -9,9 +9,11 @@ import org.apache.lucene.document.InetAddressPoint; import org.apache.lucene.util.BytesRef; +import org.elasticsearch.common.Strings; import org.elasticsearch.common.network.InetAddresses; import org.elasticsearch.geo.GeometryTestUtils; import org.elasticsearch.geo.ShapeTestUtils; +import org.elasticsearch.geometry.Geometry; import org.elasticsearch.test.ESTestCase; import org.elasticsearch.xpack.esql.core.type.DataType; import org.elasticsearch.xpack.versionfield.Version; @@ -19,11 +21,11 @@ import java.math.BigInteger; import java.util.ArrayList; import java.util.List; +import java.util.function.Function; import java.util.function.Supplier; import static org.elasticsearch.test.ESTestCase.randomBoolean; import static org.elasticsearch.test.ESTestCase.randomList; -import static org.elasticsearch.xpack.esql.core.util.SpatialCoordinateTypes.CARTESIAN; import static org.elasticsearch.xpack.esql.core.util.SpatialCoordinateTypes.GEO; import static org.elasticsearch.xpack.esql.expression.function.TestCaseSupplier.TypedDataSupplier; @@ -263,9 +265,7 @@ public static List dateCases(int minRows, int maxRows) { } /** - * * Generate cases for {@link DataType#DATE_NANOS}. - * */ public static List dateNanosCases(int minRows, int maxRows) { List cases = new ArrayList<>(); @@ -370,53 +370,58 @@ public static List versionCases(int minRows, int maxRows) { return cases; } - public static List geoPointCases(int minRows, int maxRows, boolean withAltitude) { - List cases = new ArrayList<>(); + public enum IncludingAltitude { + YES, + NO + } - addSuppliers( - cases, + public static List geoPointCases(int minRows, int maxRows, IncludingAltitude withAltitude) { + return spatialCases(minRows, maxRows, withAltitude, "geo_point", DataType.GEO_POINT, GeometryTestUtils::randomPoint); + } + + public static List geoShapeCasesWithoutCircle(int minRows, int maxRows, IncludingAltitude includingAltitude) { + return spatialCases( minRows, maxRows, - "", - DataType.GEO_POINT, - () -> GEO.asWkb(GeometryTestUtils.randomPoint(false)) + includingAltitude, + "geo_shape", + DataType.GEO_SHAPE, + b -> GeometryTestUtils.randomGeometryWithoutCircle(0, b) ); - - if (withAltitude) { - addSuppliers( - cases, - minRows, - maxRows, - "", - DataType.GEO_POINT, - () -> GEO.asWkb(GeometryTestUtils.randomPoint(true)) - ); - } - - return cases; } - public static List cartesianPointCases(int minRows, int maxRows, boolean withAltitude) { - List cases = new ArrayList<>(); - - addSuppliers( - cases, + public static List cartesianShapeCasesWithoutCircle(int minRows, int maxRows, IncludingAltitude includingAltitude) { + return spatialCases( minRows, maxRows, - "", - DataType.CARTESIAN_POINT, - () -> CARTESIAN.asWkb(ShapeTestUtils.randomPoint(false)) + includingAltitude, + "geo_shape", + DataType.CARTESIAN_SHAPE, + b -> ShapeTestUtils.randomGeometryWithoutCircle(0, b) ); + } - if (withAltitude) { - addSuppliers( - cases, - minRows, - maxRows, - "", - DataType.CARTESIAN_POINT, - () -> CARTESIAN.asWkb(ShapeTestUtils.randomPoint(true)) - ); + public static List cartesianPointCases(int minRows, int maxRows, IncludingAltitude includingAltitude) { + return spatialCases(minRows, maxRows, includingAltitude, "cartesian_point", DataType.CARTESIAN_POINT, ShapeTestUtils::randomPoint); + } + + @SuppressWarnings("fallthrough") + private static List spatialCases( + int minRows, + int maxRows, + IncludingAltitude includingAltitude, + String name, + DataType type, + Function gen + ) { + List cases = new ArrayList<>(); + + switch (includingAltitude) { + case YES: + addSuppliers(cases, minRows, maxRows, Strings.format("", name), type, () -> GEO.asWkb(gen.apply(true))); + // Explicit fallthrough: always generate a case without altitude. + case NO: + addSuppliers(cases, minRows, maxRows, Strings.format("", name), type, () -> GEO.asWkb(gen.apply(false))); } return cases; diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/aggregate/CountTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/aggregate/CountTests.java index 131072acff870..0485714959f63 100644 --- a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/aggregate/CountTests.java +++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/aggregate/CountTests.java @@ -15,6 +15,7 @@ import org.elasticsearch.xpack.esql.core.type.DataType; import org.elasticsearch.xpack.esql.expression.function.AbstractAggregationTestCase; import org.elasticsearch.xpack.esql.expression.function.MultiRowTestCaseSupplier; +import org.elasticsearch.xpack.esql.expression.function.MultiRowTestCaseSupplier.IncludingAltitude; import org.elasticsearch.xpack.esql.expression.function.TestCaseSupplier; import java.math.BigInteger; @@ -44,8 +45,8 @@ public static Iterable parameters() { MultiRowTestCaseSupplier.booleanCases(1, 1000), MultiRowTestCaseSupplier.ipCases(1, 1000), MultiRowTestCaseSupplier.versionCases(1, 1000), - MultiRowTestCaseSupplier.geoPointCases(1, 1000, true), - MultiRowTestCaseSupplier.cartesianPointCases(1, 1000, true), + MultiRowTestCaseSupplier.geoPointCases(1, 1000, IncludingAltitude.YES), + MultiRowTestCaseSupplier.cartesianPointCases(1, 1000, IncludingAltitude.YES), MultiRowTestCaseSupplier.stringCases(1, 1000, DataType.KEYWORD), MultiRowTestCaseSupplier.stringCases(1, 1000, DataType.TEXT), MultiRowTestCaseSupplier.stringCases(1, 1000, DataType.SEMANTIC_TEXT) diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/aggregate/SpatialCentroidTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/aggregate/SpatialCentroidTests.java index 15ea029a05554..b92b32aa7ad09 100644 --- a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/aggregate/SpatialCentroidTests.java +++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/aggregate/SpatialCentroidTests.java @@ -21,6 +21,7 @@ import org.elasticsearch.xpack.esql.expression.function.AbstractAggregationTestCase; import org.elasticsearch.xpack.esql.expression.function.FunctionName; import org.elasticsearch.xpack.esql.expression.function.MultiRowTestCaseSupplier; +import org.elasticsearch.xpack.esql.expression.function.MultiRowTestCaseSupplier.IncludingAltitude; import org.elasticsearch.xpack.esql.expression.function.TestCaseSupplier; import org.hamcrest.BaseMatcher; import org.hamcrest.Description; @@ -41,8 +42,8 @@ public SpatialCentroidTests(@Name("TestCase") Supplier parameters() { var suppliers = Stream.of( - MultiRowTestCaseSupplier.geoPointCases(1, 1000, true), - MultiRowTestCaseSupplier.cartesianPointCases(1, 1000, true) + MultiRowTestCaseSupplier.geoPointCases(1, 1000, IncludingAltitude.NO), + MultiRowTestCaseSupplier.cartesianPointCases(1, 1000, IncludingAltitude.NO) ).flatMap(List::stream).map(SpatialCentroidTests::makeSupplier).toList(); // The withNoRowsExpectingNull() cases don't work here, as this aggregator doesn't return nulls. diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/aggregate/SpatialExtentTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/aggregate/SpatialExtentTests.java new file mode 100644 index 0000000000000..a1faa537ba052 --- /dev/null +++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/aggregate/SpatialExtentTests.java @@ -0,0 +1,102 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +package org.elasticsearch.xpack.esql.expression.function.aggregate; + +import com.carrotsearch.randomizedtesting.annotations.Name; +import com.carrotsearch.randomizedtesting.annotations.ParametersFactory; + +import org.apache.lucene.util.BytesRef; +import org.elasticsearch.compute.aggregation.spatial.PointType; +import org.elasticsearch.geometry.Rectangle; +import org.elasticsearch.geometry.utils.GeometryValidator; +import org.elasticsearch.geometry.utils.SpatialEnvelopeVisitor; +import org.elasticsearch.geometry.utils.SpatialEnvelopeVisitor.WrapLongitude; +import org.elasticsearch.geometry.utils.WellKnownBinary; +import org.elasticsearch.xpack.esql.core.expression.Expression; +import org.elasticsearch.xpack.esql.core.tree.Source; +import org.elasticsearch.xpack.esql.core.type.DataType; +import org.elasticsearch.xpack.esql.expression.RectangleMatcher; +import org.elasticsearch.xpack.esql.expression.WellKnownBinaryBytesRefMatcher; +import org.elasticsearch.xpack.esql.expression.function.AbstractAggregationTestCase; +import org.elasticsearch.xpack.esql.expression.function.FunctionName; +import org.elasticsearch.xpack.esql.expression.function.MultiRowTestCaseSupplier; +import org.elasticsearch.xpack.esql.expression.function.MultiRowTestCaseSupplier.IncludingAltitude; +import org.elasticsearch.xpack.esql.expression.function.TestCaseSupplier; + +import java.util.List; +import java.util.function.Supplier; +import java.util.stream.Stream; + +@FunctionName("st_extent_agg") +public class SpatialExtentTests extends AbstractAggregationTestCase { + public SpatialExtentTests(@Name("TestCase") Supplier testCaseSupplier) { + this.testCase = testCaseSupplier.get(); + } + + @ParametersFactory + public static Iterable parameters() { + var suppliers = Stream.of( + MultiRowTestCaseSupplier.geoPointCases(1, 1000, IncludingAltitude.NO), + MultiRowTestCaseSupplier.cartesianPointCases(1, 1000, IncludingAltitude.NO), + MultiRowTestCaseSupplier.geoShapeCasesWithoutCircle(1, 1000, IncludingAltitude.NO), + MultiRowTestCaseSupplier.cartesianShapeCasesWithoutCircle(1, 1000, IncludingAltitude.NO) + ).flatMap(List::stream).map(SpatialExtentTests::makeSupplier).toList(); + + // The withNoRowsExpectingNull() cases don't work here, as this aggregator doesn't return nulls. + // return parameterSuppliersFromTypedDataWithDefaultChecks(suppliers); + return parameterSuppliersFromTypedData(randomizeBytesRefsOffset(suppliers)); + } + + @Override + protected Expression build(Source source, List args) { + return new SpatialExtent(source, args.get(0)); + } + + private static TestCaseSupplier makeSupplier(TestCaseSupplier.TypedDataSupplier fieldSupplier) { + return new TestCaseSupplier(List.of(fieldSupplier.type()), () -> { + PointType pointType = switch (fieldSupplier.type()) { + case DataType.CARTESIAN_POINT, DataType.CARTESIAN_SHAPE -> PointType.CARTESIAN; + case DataType.GEO_POINT, DataType.GEO_SHAPE -> PointType.GEO; + default -> throw new IllegalArgumentException("Unsupported type: " + fieldSupplier.type()); + }; + var pointVisitor = switch (pointType) { + case CARTESIAN -> new SpatialEnvelopeVisitor.CartesianPointVisitor(); + case GEO -> new SpatialEnvelopeVisitor.GeoPointVisitor(WrapLongitude.WRAP); + }; + + var fieldTypedData = fieldSupplier.get(); + DataType expectedType = DataType.isSpatialGeo(fieldTypedData.type()) ? DataType.GEO_SHAPE : DataType.CARTESIAN_SHAPE; + fieldTypedData.multiRowData() + .stream() + .map(value -> (BytesRef) value) + .map(value -> WellKnownBinary.fromWKB(GeometryValidator.NOOP, false, value.bytes, value.offset, value.length)) + .forEach(g -> g.visit(new SpatialEnvelopeVisitor(pointVisitor))); + assert pointVisitor.isValid(); + Rectangle result = pointVisitor.getResult(); + return new TestCaseSupplier.TestCase( + List.of(fieldTypedData), + "SpatialExtent[field=Attribute[channel=0]]", + expectedType, + new WellKnownBinaryBytesRefMatcher<>( + RectangleMatcher.closeTo( + new Rectangle( + // Since we use integers locally which are later decoded to doubles, all computation is effectively done using + // floats, not doubles. + (float) result.getMinX(), + (float) result.getMaxX(), + (float) result.getMaxY(), + (float) result.getMinY() + ), + 1e-3, + pointType + ) + ) + ); + }); + } +} diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/spatial/StEnvelopeTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/spatial/StEnvelopeTests.java index ac87d45491447..9f629d9127673 100644 --- a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/spatial/StEnvelopeTests.java +++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/spatial/StEnvelopeTests.java @@ -13,6 +13,7 @@ import org.apache.lucene.util.BytesRef; import org.elasticsearch.geometry.Point; import org.elasticsearch.geometry.utils.SpatialEnvelopeVisitor; +import org.elasticsearch.geometry.utils.SpatialEnvelopeVisitor.WrapLongitude; import org.elasticsearch.xpack.esql.core.expression.Expression; import org.elasticsearch.xpack.esql.core.tree.Source; import org.elasticsearch.xpack.esql.expression.function.AbstractScalarFunctionTestCase; @@ -74,7 +75,9 @@ private static BytesRef valueOf(BytesRef wkb, boolean geo) { if (geometry instanceof Point) { return wkb; } - var envelope = geo ? SpatialEnvelopeVisitor.visitGeo(geometry, true) : SpatialEnvelopeVisitor.visitCartesian(geometry); + var envelope = geo + ? SpatialEnvelopeVisitor.visitGeo(geometry, WrapLongitude.WRAP) + : SpatialEnvelopeVisitor.visitCartesian(geometry); if (envelope.isPresent()) { return UNSPECIFIED.asWkb(envelope.get()); } diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/spatial/StXMaxTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/spatial/StXMaxTests.java index dc6e61e44f599..9205879fa1cb9 100644 --- a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/spatial/StXMaxTests.java +++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/spatial/StXMaxTests.java @@ -13,6 +13,7 @@ import org.apache.lucene.util.BytesRef; import org.elasticsearch.geometry.Point; import org.elasticsearch.geometry.utils.SpatialEnvelopeVisitor; +import org.elasticsearch.geometry.utils.SpatialEnvelopeVisitor.WrapLongitude; import org.elasticsearch.xpack.esql.core.expression.Expression; import org.elasticsearch.xpack.esql.core.tree.Source; import org.elasticsearch.xpack.esql.expression.function.AbstractScalarFunctionTestCase; @@ -61,7 +62,9 @@ private static double valueOf(BytesRef wkb, boolean geo) { if (geometry instanceof Point point) { return point.getX(); } - var envelope = geo ? SpatialEnvelopeVisitor.visitGeo(geometry, true) : SpatialEnvelopeVisitor.visitCartesian(geometry); + var envelope = geo + ? SpatialEnvelopeVisitor.visitGeo(geometry, WrapLongitude.WRAP) + : SpatialEnvelopeVisitor.visitCartesian(geometry); if (envelope.isPresent()) { return envelope.get().getMaxX(); } diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/spatial/StXMinTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/spatial/StXMinTests.java index 8c06d18b1e281..3603bff9656fe 100644 --- a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/spatial/StXMinTests.java +++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/spatial/StXMinTests.java @@ -13,6 +13,7 @@ import org.apache.lucene.util.BytesRef; import org.elasticsearch.geometry.Point; import org.elasticsearch.geometry.utils.SpatialEnvelopeVisitor; +import org.elasticsearch.geometry.utils.SpatialEnvelopeVisitor.WrapLongitude; import org.elasticsearch.xpack.esql.core.expression.Expression; import org.elasticsearch.xpack.esql.core.tree.Source; import org.elasticsearch.xpack.esql.expression.function.AbstractScalarFunctionTestCase; @@ -61,7 +62,9 @@ private static double valueOf(BytesRef wkb, boolean geo) { if (geometry instanceof Point point) { return point.getX(); } - var envelope = geo ? SpatialEnvelopeVisitor.visitGeo(geometry, true) : SpatialEnvelopeVisitor.visitCartesian(geometry); + var envelope = geo + ? SpatialEnvelopeVisitor.visitGeo(geometry, WrapLongitude.WRAP) + : SpatialEnvelopeVisitor.visitCartesian(geometry); if (envelope.isPresent()) { return envelope.get().getMinX(); } diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/spatial/StYMaxTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/spatial/StYMaxTests.java index 7222d7517f7ff..cb2a03c3a9473 100644 --- a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/spatial/StYMaxTests.java +++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/spatial/StYMaxTests.java @@ -13,6 +13,7 @@ import org.apache.lucene.util.BytesRef; import org.elasticsearch.geometry.Point; import org.elasticsearch.geometry.utils.SpatialEnvelopeVisitor; +import org.elasticsearch.geometry.utils.SpatialEnvelopeVisitor.WrapLongitude; import org.elasticsearch.xpack.esql.core.expression.Expression; import org.elasticsearch.xpack.esql.core.tree.Source; import org.elasticsearch.xpack.esql.expression.function.AbstractScalarFunctionTestCase; @@ -61,7 +62,9 @@ private static double valueOf(BytesRef wkb, boolean geo) { if (geometry instanceof Point point) { return point.getY(); } - var envelope = geo ? SpatialEnvelopeVisitor.visitGeo(geometry, true) : SpatialEnvelopeVisitor.visitCartesian(geometry); + var envelope = geo + ? SpatialEnvelopeVisitor.visitGeo(geometry, WrapLongitude.WRAP) + : SpatialEnvelopeVisitor.visitCartesian(geometry); if (envelope.isPresent()) { return envelope.get().getMaxY(); } diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/spatial/StYMinTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/spatial/StYMinTests.java index 843c7bb649114..0c191f6dc4c5b 100644 --- a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/spatial/StYMinTests.java +++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/spatial/StYMinTests.java @@ -13,6 +13,7 @@ import org.apache.lucene.util.BytesRef; import org.elasticsearch.geometry.Point; import org.elasticsearch.geometry.utils.SpatialEnvelopeVisitor; +import org.elasticsearch.geometry.utils.SpatialEnvelopeVisitor.WrapLongitude; import org.elasticsearch.xpack.esql.core.expression.Expression; import org.elasticsearch.xpack.esql.core.tree.Source; import org.elasticsearch.xpack.esql.expression.function.AbstractScalarFunctionTestCase; @@ -61,7 +62,9 @@ private static double valueOf(BytesRef wkb, boolean geo) { if (geometry instanceof Point point) { return point.getY(); } - var envelope = geo ? SpatialEnvelopeVisitor.visitGeo(geometry, true) : SpatialEnvelopeVisitor.visitCartesian(geometry); + var envelope = geo + ? SpatialEnvelopeVisitor.visitGeo(geometry, WrapLongitude.WRAP) + : SpatialEnvelopeVisitor.visitCartesian(geometry); if (envelope.isPresent()) { return envelope.get().getMinY(); } diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/PhysicalPlanOptimizerTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/PhysicalPlanOptimizerTests.java index ec1d55a0fc58f..dc3ae0a3388cb 100644 --- a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/PhysicalPlanOptimizerTests.java +++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/PhysicalPlanOptimizerTests.java @@ -21,6 +21,7 @@ import org.elasticsearch.geometry.Polygon; import org.elasticsearch.geometry.ShapeType; import org.elasticsearch.index.IndexMode; +import org.elasticsearch.index.mapper.MappedFieldType.FieldExtractPreference; import org.elasticsearch.index.query.BoolQueryBuilder; import org.elasticsearch.index.query.ExistsQueryBuilder; import org.elasticsearch.index.query.QueryBuilder; @@ -64,6 +65,7 @@ import org.elasticsearch.xpack.esql.expression.function.aggregate.Count; import org.elasticsearch.xpack.esql.expression.function.aggregate.SpatialAggregateFunction; import org.elasticsearch.xpack.esql.expression.function.aggregate.SpatialCentroid; +import org.elasticsearch.xpack.esql.expression.function.aggregate.SpatialExtent; import org.elasticsearch.xpack.esql.expression.function.aggregate.Sum; import org.elasticsearch.xpack.esql.expression.function.fulltext.Match; import org.elasticsearch.xpack.esql.expression.function.scalar.math.Round; @@ -253,7 +255,7 @@ public void init() { "mapping-airports_no_doc_values.json", functionRegistry, enrichResolution, - new TestConfigurableSearchStats().exclude(Config.DOC_VALUES, "location") + new TestConfigurableSearchStats().exclude(Config.DOC_VALUES, "location").exclude(Config.DOC_VALUES, "city_location") ); this.airportsNotIndexed = makeTestDataSource( "airports-not-indexed", @@ -2804,7 +2806,7 @@ public void testPartialAggFoldingOutputForSyntheticAgg() { * Also note that the type converting function is removed when it does not actually convert the type, * ensuring that ReferenceAttributes are not created for the same field, and the optimization can still work. */ - public void testSpatialTypesAndStatsUseDocValues() { + public void testSpatialTypesAndStatsCentroidUseDocValues() { for (String query : new String[] { "from airports | stats centroid = st_centroid_agg(location)", "from airports | stats centroid = st_centroid_agg(to_geopoint(location))", @@ -2838,6 +2840,129 @@ public void testSpatialTypesAndStatsUseDocValues() { } } + /** + * Before local optimizations: + * + * LimitExec[1000[INTEGER]] + * \_AggregateExec[[],[SPATIALSTEXTENT(location{f}#48,true[BOOLEAN]) AS extent],FINAL,[minNegX{r}#52, minPosX{r}#53, maxNegX{r}#54, + * maxPosX{r}#55, maxY{r}#56, minY{r}#57],null] + * \_ExchangeExec[[minNegX{r}#52, minPosX{r}#53, maxNegX{r}#54, maxPosX{r}#55, maxY{r}#56, minY{r}#57],true] + * \_FragmentExec[filter=null, estimatedRowSize=0, reducer=[], fragment=[ + * Aggregate[STANDARD,[],[SPATIALSTEXTENT(location{f}#48,true[BOOLEAN]) AS extent]] + * \_EsRelation[airports][abbrev{f}#44, city{f}#50, city_location{f}#51, coun..]]] + * + * After local optimizations: + * + * LimitExec[1000[INTEGER]] + * \_AggregateExec[[],[SPATIALSTEXTENT(location{f}#48,true[BOOLEAN]) AS extent],FINAL,[minNegX{r}#52, minPosX{r}#53, maxNegX{r}#54, + * maxPosX{r}#55, maxY{r}#56, minY{r}#57],21] + * \_ExchangeExec[[minNegX{r}#52, minPosX{r}#53, maxNegX{r}#54, maxPosX{r}#55, maxY{r}#56, minY{r}#57],true] + * \_AggregateExec[[],[SPATIALSTEXTENT(location{f}#48,true[BOOLEAN]) AS extent],INITIAL,[ + * minNegX{r}#73, minPosX{r}#74, maxNegX{rb#75, maxPosX{r}#76, maxY{r}#77, minY{r}#78],21] + * \_FieldExtractExec[location{f}#48][location{f}#48] + * \_EsQueryExec[airports], indexMode[standard], query[{"exists":{"field":"location","boost":1.0}}][ + * _doc{f}#79], limit[], sort[] estimatedRowSize[25] + * + * Note the FieldExtractExec has 'location' set for stats: FieldExtractExec[location{f}#9][location{f}#9] + *

+ * Also note that the type converting function is removed when it does not actually convert the type, + * ensuring that ReferenceAttributes are not created for the same field, and the optimization can still work. + */ + public void testSpatialTypesAndStatsExtentUseDocValues() { + for (String query : new String[] { + "from airports | stats extent = st_extent_agg(location)", + "from airports | stats extent = st_extent_agg(to_geopoint(location))", + "from airports | eval location = to_geopoint(location) | stats extent = st_extent_agg(location)" }) { + for (boolean withDocValues : new boolean[] { false, true }) { + var testData = withDocValues ? airports : airportsNoDocValues; + var plan = physicalPlan(query, testData); + + var limit = as(plan, LimitExec.class); + var agg = as(limit.child(), AggregateExec.class); + // Before optimization the aggregation does not use doc-values + assertAggregation(agg, "extent", SpatialExtent.class, GEO_POINT, false); + + var exchange = as(agg.child(), ExchangeExec.class); + var fragment = as(exchange.child(), FragmentExec.class); + var fAgg = as(fragment.fragment(), Aggregate.class); + as(fAgg.child(), EsRelation.class); + + // Now optimize the plan and assert the aggregation uses doc-values + var optimized = optimizedPlan(plan, testData.stats); + limit = as(optimized, LimitExec.class); + agg = as(limit.child(), AggregateExec.class); + // Above the exchange (in coordinator) the aggregation is not using doc-values + assertAggregation(agg, "extent", SpatialExtent.class, GEO_POINT, false); + exchange = as(agg.child(), ExchangeExec.class); + agg = as(exchange.child(), AggregateExec.class); + // below the exchange (in data node) the aggregation is using doc-values + assertAggregation(agg, "extent", SpatialExtent.class, GEO_POINT, withDocValues); + assertChildIsGeoPointExtract(withDocValues ? agg : as(agg.child(), FilterExec.class), withDocValues); + } + } + } + + /** + * Before local optimizations: + * + * LimitExec[1000[INTEGER]] + * \_AggregateExec[[],[SPATIALSTEXTENT(location{f}#48,true[BOOLEAN]) AS extent],FINAL,[minNegX{r}#52, minPosX{r}#53, maxNegX{r}#54, + * maxPosX{r}#55, maxY{r}#56, minY{r}#57],null] + * \_ExchangeExec[[minNegX{r}#52, minPosX{r}#53, maxNegX{r}#54, maxPosX{r}#55, maxY{r}#56, minY{r}#57],true] + * \_FragmentExec[filter=null, estimatedRowSize=0, reducer=[], fragment=[ + * Aggregate[STANDARD,[],[SPATIALSTEXTENT(location{f}#48,true[BOOLEAN]) AS extent]] + * \_EsRelation[airports][abbrev{f}#44, city{f}#50, city_location{f}#51, coun..]]] + * + * After local optimizations: + * + * LimitExec[1000[INTEGER]] + * \_AggregateExec[[],[SPATIALSTEXTENT(location{f}#48,true[BOOLEAN]) AS extent],FINAL,[minNegX{r}#52, minPosX{r}#53, maxNegX{r}#54, + * maxPosX{r}#55, maxY{r}#56, minY{r}#57],21] + * \_ExchangeExec[[minNegX{r}#52, minPosX{r}#53, maxNegX{r}#54, maxPosX{r}#55, maxY{r}#56, minY{r}#57],true] + * \_AggregateExec[[],[SPATIALSTEXTENT(location{f}#48,true[BOOLEAN]) AS extent],INITIAL,[ + * minNegX{r}#73, minPosX{r}#74, maxNegX{rb#75, maxPosX{r}#76, maxY{r}#77, minY{r}#78],21] + * \_FieldExtractExec[location{f}#48][location{f}#48] + * \_EsQueryExec[airports], indexMode[standard], query[{"exists":{"field":"location","boost":1.0}}][ + * _doc{f}#79], limit[], sort[] estimatedRowSize[25] + * + * Note the FieldExtractExec has 'location' set for stats: FieldExtractExec[location{f}#9][location{f}#9] + *

+ * Also note that the type converting function is removed when it does not actually convert the type, + * ensuring that ReferenceAttributes are not created for the same field, and the optimization can still work. + */ + public void testSpatialTypesAndStatsExtentAndCentroidUseDocValues() { + for (String query : new String[] { + "from airports | stats extent = st_extent_agg(location), centroid = st_centroid_agg(location)", + "from airports | stats extent = st_extent_agg(location), centroid = st_centroid_agg(city_location)", }) { + for (boolean withDocValues : new boolean[] { false, true }) { + var testData = withDocValues ? airports : airportsNoDocValues; + var plan = physicalPlan(query, testData); + + var limit = as(plan, LimitExec.class); + var agg = as(limit.child(), AggregateExec.class); + // Before optimization the aggregation does not use doc-values + assertAggregation(agg, "extent", SpatialExtent.class, GEO_POINT, false); + + var exchange = as(agg.child(), ExchangeExec.class); + var fragment = as(exchange.child(), FragmentExec.class); + var fAgg = as(fragment.fragment(), Aggregate.class); + as(fAgg.child(), EsRelation.class); + + // Now optimize the plan and assert the aggregation uses doc-values + var optimized = optimizedPlan(plan, testData.stats); + limit = as(optimized, LimitExec.class); + agg = as(limit.child(), AggregateExec.class); + // Above the exchange (in coordinator) the aggregation is not using doc-values + assertAggregation(agg, "extent", SpatialExtent.class, GEO_POINT, false); + exchange = as(agg.child(), ExchangeExec.class); + agg = as(exchange.child(), AggregateExec.class); + // below the exchange (in data node) the aggregation is using doc-values + assertAggregation(agg, "extent", SpatialExtent.class, GEO_POINT, withDocValues); + assertChildIsGeoPointExtract(withDocValues ? agg : as(agg.child(), FilterExec.class), withDocValues); + } + } + } + /** * This test does not have real index fields, and therefor asserts that doc-values field extraction does NOT occur. * Before local optimizations: @@ -6805,7 +6930,11 @@ private static void assertAggregation( var aggFunc = assertAggregation(plan, aliasName, aggClass); var aggField = as(aggFunc.field(), Attribute.class); var spatialAgg = as(aggFunc, SpatialAggregateFunction.class); - assertThat("Expected spatial aggregation to use doc-values", spatialAgg.useDocValues(), equalTo(useDocValues)); + assertThat( + "Expected spatial aggregation to use doc-values", + spatialAgg.fieldExtractPreference(), + equalTo(useDocValues ? FieldExtractPreference.DOC_VALUES : FieldExtractPreference.NONE) + ); assertThat("", aggField.dataType(), equalTo(fieldType)); } diff --git a/x-pack/plugin/src/yamlRestTest/resources/rest-api-spec/test/esql/60_usage.yml b/x-pack/plugin/src/yamlRestTest/resources/rest-api-spec/test/esql/60_usage.yml index 81f65668722fc..2a4cde9a680e9 100644 --- a/x-pack/plugin/src/yamlRestTest/resources/rest-api-spec/test/esql/60_usage.yml +++ b/x-pack/plugin/src/yamlRestTest/resources/rest-api-spec/test/esql/60_usage.yml @@ -92,7 +92,7 @@ setup: - gt: {esql.functions.to_long: $functions_to_long} - match: {esql.functions.coalesce: $functions_coalesce} # Testing for the entire function set isn't feasbile, so we just check that we return the correct count as an approximation. - - length: {esql.functions: 128} # check the "sister" test below for a likely update to the same esql.functions length check + - length: {esql.functions: 129} # check the "sister" test below for a likely update to the same esql.functions length check --- "Basic ESQL usage output (telemetry) non-snapshot version": @@ -163,4 +163,4 @@ setup: - match: {esql.functions.cos: $functions_cos} - gt: {esql.functions.to_long: $functions_to_long} - match: {esql.functions.coalesce: $functions_coalesce} - - length: {esql.functions: 124} # check the "sister" test above for a likely update to the same esql.functions length check + - length: {esql.functions: 125} # check the "sister" test above for a likely update to the same esql.functions length check