|  | 
|  | 1 | +/* | 
|  | 2 | + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one | 
|  | 3 | + * or more contributor license agreements. Licensed under the Elastic License | 
|  | 4 | + * 2.0; you may not use this file except in compliance with the Elastic License | 
|  | 5 | + * 2.0. | 
|  | 6 | + */ | 
|  | 7 | + | 
|  | 8 | +package org.elasticsearch.xpack.esql.spatial; | 
|  | 9 | + | 
|  | 10 | +import org.elasticsearch.ElasticsearchException; | 
|  | 11 | +import org.elasticsearch.action.bulk.BulkRequestBuilder; | 
|  | 12 | +import org.elasticsearch.action.index.IndexRequest; | 
|  | 13 | +import org.elasticsearch.action.support.WriteRequest; | 
|  | 14 | +import org.elasticsearch.xpack.esql.VerificationException; | 
|  | 15 | +import org.elasticsearch.xpack.esql.action.AbstractEsqlIntegTestCase; | 
|  | 16 | +import org.elasticsearch.xpack.esql.action.EsqlCapabilities; | 
|  | 17 | +import org.junit.Before; | 
|  | 18 | + | 
|  | 19 | +import java.util.List; | 
|  | 20 | +import java.util.Locale; | 
|  | 21 | + | 
|  | 22 | +import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertAcked; | 
|  | 23 | +import static org.elasticsearch.xpack.esql.EsqlTestUtils.getValuesList; | 
|  | 24 | +import static org.hamcrest.Matchers.containsString; | 
|  | 25 | +import static org.hamcrest.Matchers.equalTo; | 
|  | 26 | + | 
|  | 27 | +public abstract class SpatialExtentAggregationTestCase extends AbstractEsqlIntegTestCase { | 
|  | 28 | + | 
|  | 29 | +    @Before | 
|  | 30 | +    public void setupIndex() throws Exception { | 
|  | 31 | +        assumeTrue("requires ST_EXTENT_AGG capability", EsqlCapabilities.Cap.ST_EXTENT_AGG.isEnabled()); | 
|  | 32 | +        createAndPopulateIndexes(-10, 10, -10, 10); | 
|  | 33 | +    } | 
|  | 34 | + | 
|  | 35 | +    /** | 
|  | 36 | +     * This test should pass only with an enterprise license | 
|  | 37 | +     */ | 
|  | 38 | +    public abstract void testStExtentAggregationWithShapes(); | 
|  | 39 | + | 
|  | 40 | +    /** | 
|  | 41 | +     * This test should pass with and without enterprise licenses | 
|  | 42 | +     */ | 
|  | 43 | +    public void testStExtentAggregationWithPoints() throws Exception { | 
|  | 44 | +        assertStExtentFromIndex("index_geo_point"); | 
|  | 45 | +    } | 
|  | 46 | + | 
|  | 47 | +    protected void assertStExtentFromIndex(String index) { | 
|  | 48 | +        var query = String.format(Locale.ROOT, """ | 
|  | 49 | +            FROM %s | 
|  | 50 | +            | STATS extent = ST_EXTENT_AGG(location) | 
|  | 51 | +            | EVAL minX = ROUND(ST_XMIN(extent)) | 
|  | 52 | +            | EVAL maxX = ROUND(ST_XMAX(extent)) | 
|  | 53 | +            | EVAL minY = ROUND(ST_YMIN(extent)) | 
|  | 54 | +            | EVAL maxY = ROUND(ST_YMAX(extent)) | 
|  | 55 | +            """, index); | 
|  | 56 | +        try (var resp = run(query)) { | 
|  | 57 | +            assertColumnNames(resp.columns(), List.of("extent", "minX", "maxX", "minY", "maxY")); | 
|  | 58 | +            assertColumnTypes(resp.columns(), List.of("geo_shape", "double", "double", "double", "double")); | 
|  | 59 | +            List<List<Object>> values = getValuesList(resp.values()); | 
|  | 60 | +            assertThat(values.size(), equalTo(1)); | 
|  | 61 | +            List<Object> row = values.get(0); | 
|  | 62 | +            List<Object> expectedValues = List.of(-10.0, 10.0, -10.0, 10.0); | 
|  | 63 | +            assertThat(row.subList(1, row.size()), equalTo(expectedValues)); | 
|  | 64 | +        } | 
|  | 65 | +    } | 
|  | 66 | + | 
|  | 67 | +    protected void assertStExtentFailsWith(String index) { | 
|  | 68 | +        var query = String.format(Locale.ROOT, """ | 
|  | 69 | +            FROM %s | 
|  | 70 | +            | STATS extent = ST_EXTENT_AGG(location) | 
|  | 71 | +            | EVAL minX = ROUND(ST_XMIN(extent)) | 
|  | 72 | +            | EVAL maxX = ROUND(ST_XMAX(extent)) | 
|  | 73 | +            | EVAL minY = ROUND(ST_YMIN(extent)) | 
|  | 74 | +            | EVAL maxY = ROUND(ST_YMAX(extent)) | 
|  | 75 | +            """, index); | 
|  | 76 | +        ElasticsearchException e = expectThrows(VerificationException.class, () -> run(query)); | 
|  | 77 | +        assertThat(e.getMessage(), containsString("current license is non-compliant for [ST_EXTENT_AGG(location)]")); | 
|  | 78 | +    } | 
|  | 79 | + | 
|  | 80 | +    private void createAndPopulateIndexes(double minX, double maxX, double minY, double maxY) throws Exception { | 
|  | 81 | +        int numX = 21; | 
|  | 82 | +        int numY = 21; | 
|  | 83 | +        initIndex("index_", "geo_point"); | 
|  | 84 | +        initIndex("index_", "geo_shape"); | 
|  | 85 | +        BulkRequestBuilder points = client().prepareBulk(); | 
|  | 86 | +        BulkRequestBuilder shapes = client().prepareBulk(); | 
|  | 87 | +        for (int xi = 0; xi < numX; xi++) { | 
|  | 88 | +            for (int yi = 0; yi < numY; yi++) { | 
|  | 89 | +                double x = minX + xi * (maxX - minX) / (numX - 1); | 
|  | 90 | +                double y = minY + yi * (maxY - minY) / (numY - 1); | 
|  | 91 | +                String point = "POINT(" + x + " " + y + ")"; | 
|  | 92 | +                points.add(new IndexRequest("index_geo_point").id(x + ":" + y).source("location", point)); | 
|  | 93 | +                if (xi > 0 && yi > 0) { | 
|  | 94 | +                    double px = minX + (xi - 1) * (maxX - minX) / numX; | 
|  | 95 | +                    double py = minY + (yi - 1) * (maxY - minY) / numY; | 
|  | 96 | +                    String shape = "BBOX(" + px + ", " + x + ", " + y + ", " + py + ")"; | 
|  | 97 | +                    shapes.add(new IndexRequest("index_geo_shape").id(x + ":" + y).source("location", shape)); | 
|  | 98 | +                } | 
|  | 99 | +            } | 
|  | 100 | +        } | 
|  | 101 | +        points.setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE).get(); | 
|  | 102 | +        shapes.setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE).get(); | 
|  | 103 | +        ensureYellow("index_geo_point"); | 
|  | 104 | +        ensureYellow("index_geo_shape"); | 
|  | 105 | +    } | 
|  | 106 | + | 
|  | 107 | +    protected void initIndex(String prefix, String fieldType) { | 
|  | 108 | +        assertAcked(prepareCreate(prefix + fieldType).setMapping(String.format(Locale.ROOT, """ | 
|  | 109 | +            { | 
|  | 110 | +              "properties" : { | 
|  | 111 | +                "location": { "type" : "%s" } | 
|  | 112 | +              } | 
|  | 113 | +            } | 
|  | 114 | +            """, fieldType))); | 
|  | 115 | +    } | 
|  | 116 | +} | 
0 commit comments