Examples with Rectangle io.trino.geospatial.Rectangle used on opensource projects

Example 1 with Rectangle

use of io.trino.geospatial.Rectangle in project trino by trinodb.

the class TestSpatialPartitioningInternalAggregation method test.

@Test(dataProvider = "partitionCount")
public void test(int partitionCount) {
    TestingAggregationFunction function = getFunction();
    List<OGCGeometry> geometries = makeGeometries();
    Block geometryBlock = makeGeometryBlock(geometries);
    Block partitionCountBlock = BlockAssertions.createRLEBlock(partitionCount, geometries.size());
    Rectangle expectedExtent = new Rectangle(-10, -10, Math.nextUp(10.0), Math.nextUp(10.0));
    String expectedValue = getSpatialPartitioning(expectedExtent, geometries, partitionCount);
    AggregatorFactory aggregatorFactory = function.createAggregatorFactory(SINGLE, Ints.asList(0, 1), OptionalInt.empty());
    Page page = new Page(geometryBlock, partitionCountBlock);
    Aggregator aggregator = aggregatorFactory.createAggregator();
    aggregator.processPage(page);
    String aggregation = (String) BlockAssertions.getOnlyValue(function.getFinalType(), getFinalBlock(function.getFinalType(), aggregator));
    assertEquals(aggregation, expectedValue);
    GroupedAggregator groupedAggregator = aggregatorFactory.createGroupedAggregator();
    groupedAggregator.processPage(createGroupByIdBlock(0, page.getPositionCount()), page);
    String groupValue = (String) getGroupValue(function.getFinalType(), groupedAggregator, 0);
    assertEquals(groupValue, expectedValue);
}

Example 2 with Rectangle

use of io.trino.geospatial.Rectangle in project trino by trinodb.

the class SpatialPartitioningInternalAggregateFunction method input.

@InputFunction
public static void input(SpatialPartitioningState state, @SqlType(GEOMETRY_TYPE_NAME) Slice slice, @SqlType(INTEGER) long partitionCount) {
    Envelope envelope = deserializeEnvelope(slice);
    if (envelope.isEmpty()) {
        return;
    }
    Rectangle extent = new Rectangle(envelope.getXMin(), envelope.getYMin(), envelope.getXMax(), envelope.getYMax());
    if (state.getCount() == 0) {
        state.setPartitionCount(toIntExact(partitionCount));
        state.setExtent(extent);
        state.setSamples(new ArrayList<>());
    } else {
        state.setExtent(state.getExtent().merge(extent));
    }
    // use reservoir sampling
    List<Rectangle> samples = state.getSamples();
    if (samples.size() <= MAX_SAMPLE_COUNT) {
        samples.add(extent);
    } else {
        long sampleIndex = ThreadLocalRandom.current().nextLong(state.getCount());
        if (sampleIndex < MAX_SAMPLE_COUNT) {
            samples.set(toIntExact(sampleIndex), extent);
        }
    }
    state.setCount(state.getCount() + 1);
}

Example 3 with Rectangle

use of io.trino.geospatial.Rectangle in project trino by trinodb.

the class GeoFunctions method spatialPartitions.

private static Block spatialPartitions(KdbTree kdbTree, Rectangle envelope) {
    Map<Integer, Rectangle> partitions = kdbTree.findIntersectingLeaves(envelope);
    if (partitions.isEmpty()) {
        return EMPTY_ARRAY_OF_INTS;
    }
    // For input rectangles that represent a single point, return at most one partition
    // by excluding right and upper sides of partition rectangles. The logic that builds
    // KDB tree needs to make sure to add some padding to the right and upper sides of the
    // overall extent of the tree to avoid missing right-most and top-most points.
    boolean point = (envelope.getWidth() == 0 && envelope.getHeight() == 0);
    if (point) {
        for (Map.Entry<Integer, Rectangle> partition : partitions.entrySet()) {
            if (envelope.getXMin() < partition.getValue().getXMax() && envelope.getYMin() < partition.getValue().getYMax()) {
                BlockBuilder blockBuilder = IntegerType.INTEGER.createFixedSizeBlockBuilder(1);
                blockBuilder.writeInt(partition.getKey());
                return blockBuilder.build();
            }
        }
        throw new VerifyException(format("Cannot find half-open partition extent for a point: (%s, %s)", envelope.getXMin(), envelope.getYMin()));
    }
    BlockBuilder blockBuilder = IntegerType.INTEGER.createFixedSizeBlockBuilder(partitions.size());
    for (int id : partitions.keySet()) {
        blockBuilder.writeInt(id);
    }
    return blockBuilder.build();
}

Example 4 with Rectangle

use of io.trino.geospatial.Rectangle in project trino by trinodb.

the class TestSpatialPartitioningInternalAggregation method getSpatialPartitioning.

private String getSpatialPartitioning(Rectangle extent, List<OGCGeometry> geometries, int partitionCount) {
    ImmutableList.Builder<Rectangle> rectangles = ImmutableList.builder();
    for (OGCGeometry geometry : geometries) {
        Envelope envelope = new Envelope();
        geometry.getEsriGeometry().queryEnvelope(envelope);
        rectangles.add(new Rectangle(envelope.getXMin(), envelope.getYMin(), envelope.getXMax(), envelope.getYMax()));
    }
    return KdbTreeUtils.toJson(buildKdbTree(roundToInt(geometries.size() * 1.0 / partitionCount, CEILING), extent, rectangles.build()));
}

Example 5 with Rectangle

use of io.trino.geospatial.Rectangle in project trino by trinodb.

the class TestKdbTreeType method createTestBlock.

private static Block createTestBlock() {
    BlockBuilder blockBuilder = KDB_TREE.createBlockBuilder(null, 1);
    KdbTree kdbTree = new KdbTree(new Node(new Rectangle(10, 20, 30, 40), OptionalInt.of(42), Optional.empty(), Optional.empty()));
    KDB_TREE.writeObject(blockBuilder, kdbTree);
    return blockBuilder.build();
}