Skip to content

Commit

Permalink
[ESQL] Remove spatial MultiValuesCombiner abstraction (#112914)
Browse files Browse the repository at this point in the history 
Replaced with static functions
  • Loading branch information
@iverase
iverase authored Sep 17, 2024
1 parent 8e289f2 commit 64127f2
Show file tree
Hide file tree
Showing 3 changed files with 124 additions and 146 deletions.
49 changes: 34 additions & 15 deletions ...java/org/elasticsearch/xpack/esql/expression/function/scalar/spatial/SpatialContains.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -91,9 +91,23 @@ protected boolean geometryRelatesGeometry(BytesRef left, BytesRef right) throws
}

@Override
protected boolean geometryRelatesGeometries(MultiValuesCombiner left, MultiValuesCombiner right) throws IOException {
Component2D[] rightComponent2Ds = asLuceneComponent2Ds(crsType, right.combined());
return geometryRelatesGeometries(left, rightComponent2Ds);
protected void processSourceAndSource(BooleanBlock.Builder builder, int position, BytesRefBlock left, BytesRefBlock right)
throws IOException {
if (right.getValueCount(position) < 1) {
builder.appendNull();
} else {
processSourceAndConstant(builder, position, left, asLuceneComponent2Ds(crsType, right, position));
}
}

@Override
protected void processPointDocValuesAndSource(
BooleanBlock.Builder builder,
int position,
LongBlock leftValue,
BytesRefBlock rightValue
) throws IOException {
processPointDocValuesAndConstant(builder, position, leftValue, asLuceneComponent2Ds(crsType, rightValue, position));
}

private boolean geometryRelatesGeometries(BytesRef left, Component2D[] rightComponent2Ds) throws IOException {
Expand All @@ -102,11 +116,6 @@ private boolean geometryRelatesGeometries(BytesRef left, Component2D[] rightComp
return geometryRelatesGeometries(leftDocValueReader, rightComponent2Ds);
}

private boolean geometryRelatesGeometries(MultiValuesCombiner left, Component2D[] rightComponent2Ds) throws IOException {
GeometryDocValueReader leftDocValueReader = asGeometryDocValueReader(coordinateEncoder, shapeIndexer, left.combined());
return geometryRelatesGeometries(leftDocValueReader, rightComponent2Ds);
}

private boolean geometryRelatesGeometries(GeometryDocValueReader leftDocValueReader, Component2D[] rightComponent2Ds)
throws IOException {
for (Component2D rightComponent2D : rightComponent2Ds) {
Expand All @@ -123,18 +132,28 @@ private void processSourceAndConstant(BooleanBlock.Builder builder, int position
if (left.getValueCount(position) < 1) {
builder.appendNull();
} else {
MultiValuesBytesRef leftValues = new MultiValuesBytesRef(left, position);
builder.appendBoolean(geometryRelatesGeometries(leftValues, right));
final GeometryDocValueReader reader = asGeometryDocValueReader(coordinateEncoder, shapeIndexer, left, position);
builder.appendBoolean(geometryRelatesGeometries(reader, right));
}
}

private void processPointDocValuesAndConstant(BooleanBlock.Builder builder, int p, LongBlock left, @Fixed Component2D[] right)
throws IOException {
if (left.getValueCount(p) < 1) {
private void processPointDocValuesAndConstant(
BooleanBlock.Builder builder,
int position,
LongBlock left,
@Fixed Component2D[] right
) throws IOException {
if (left.getValueCount(position) < 1) {
builder.appendNull();
} else {
MultiValuesLong leftValues = new MultiValuesLong(left, p, spatialCoordinateType::longAsPoint);
builder.appendBoolean(geometryRelatesGeometries(leftValues, right));
final GeometryDocValueReader reader = asGeometryDocValueReader(
coordinateEncoder,
shapeIndexer,
left,
position,
spatialCoordinateType::longAsPoint
);
builder.appendBoolean(geometryRelatesGeometries(reader, right));
}
}
}
Expand Down
138 changes: 22 additions & 116 deletions ...g/elasticsearch/xpack/esql/expression/function/scalar/spatial/SpatialRelatesFunction.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,10 +17,6 @@
import org.elasticsearch.compute.data.BytesRefBlock;
import org.elasticsearch.compute.data.LongBlock;
import org.elasticsearch.compute.operator.EvalOperator;
import org.elasticsearch.geometry.Geometry;
import org.elasticsearch.geometry.GeometryCollection;
import org.elasticsearch.geometry.MultiPoint;
import org.elasticsearch.geometry.Point;
import org.elasticsearch.index.mapper.ShapeIndexer;
import org.elasticsearch.lucene.spatial.Component2DVisitor;
import org.elasticsearch.lucene.spatial.CoordinateEncoder;
Expand All @@ -33,8 +29,6 @@
import org.elasticsearch.xpack.esql.evaluator.mapper.EvaluatorMapper;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Set;
Expand Down Expand Up @@ -169,23 +163,13 @@ protected boolean geometryRelatesGeometry(GeometryDocValueReader reader, Compone
return visitor.matches();
}

protected boolean geometryRelatesGeometries(MultiValuesCombiner left, MultiValuesCombiner right) throws IOException {
Component2D rightComponent2D = asLuceneComponent2D(crsType, right.combined());
return geometryRelatesGeometry(left, rightComponent2D);
}

private boolean geometryRelatesGeometry(MultiValuesCombiner left, Component2D rightComponent2D) throws IOException {
GeometryDocValueReader leftDocValueReader = asGeometryDocValueReader(coordinateEncoder, shapeIndexer, left.combined());
return geometryRelatesGeometry(leftDocValueReader, rightComponent2D);
}

protected void processSourceAndConstant(BooleanBlock.Builder builder, int position, BytesRefBlock left, @Fixed Component2D right)
throws IOException {
if (left.getValueCount(position) < 1) {
builder.appendNull();
} else {
MultiValuesBytesRef leftValues = new MultiValuesBytesRef(left, position);
builder.appendBoolean(geometryRelatesGeometry(leftValues, right));
final GeometryDocValueReader reader = asGeometryDocValueReader(coordinateEncoder, shapeIndexer, left, position);
builder.appendBoolean(geometryRelatesGeometry(reader, right));
}
}

Expand All @@ -194,9 +178,9 @@ protected void processSourceAndSource(BooleanBlock.Builder builder, int position
if (left.getValueCount(position) < 1 || right.getValueCount(position) < 1) {
builder.appendNull();
} else {
MultiValuesBytesRef leftValues = new MultiValuesBytesRef(left, position);
MultiValuesBytesRef rightValues = new MultiValuesBytesRef(right, position);
builder.appendBoolean(geometryRelatesGeometries(leftValues, rightValues));
final GeometryDocValueReader reader = asGeometryDocValueReader(coordinateEncoder, shapeIndexer, left, position);
final Component2D component2D = asLuceneComponent2D(crsType, right, position);
builder.appendBoolean(geometryRelatesGeometry(reader, component2D));
}
}

Expand All @@ -209,8 +193,14 @@ protected void processPointDocValuesAndConstant(
if (leftValue.getValueCount(position) < 1) {
builder.appendNull();
} else {
MultiValuesLong leftValues = new MultiValuesLong(leftValue, position, spatialCoordinateType::longAsPoint);
builder.appendBoolean(geometryRelatesGeometry(leftValues, rightValue));
final GeometryDocValueReader reader = asGeometryDocValueReader(
coordinateEncoder,
shapeIndexer,
leftValue,
position,
spatialCoordinateType::longAsPoint
);
builder.appendBoolean(geometryRelatesGeometry(reader, rightValue));
}
}

Expand All @@ -223,100 +213,16 @@ protected void processPointDocValuesAndSource(
if (leftValue.getValueCount(position) < 1 || rightValue.getValueCount(position) < 1) {
builder.appendNull();
} else {
MultiValuesLong leftValues = new MultiValuesLong(leftValue, position, spatialCoordinateType::longAsPoint);
MultiValuesBytesRef rightValues = new MultiValuesBytesRef(rightValue, position);
builder.appendBoolean(geometryRelatesGeometries(leftValues, rightValues));
}
}
}

/**
* When dealing with ST_CONTAINS and ST_WITHIN we need to pre-combine the field geometries for multi-values in order
* to perform the relationship check. This means instead of relying on the generated evaluators to iterate over all
* values in a multi-value field, the entire block is passed into the spatial function, and we combine the values into
* a geometry collection or multipoint.
*/
protected interface MultiValuesCombiner {
Geometry combined();
}

/**
* Values read from source will be encoded as WKB in BytesRefBlock. The block contains multiple rows, and within
* each row multiple values, so we need to efficiently iterate over only the values required for the requested row.
* This class works for point and shape fields, because both are extracted into the same block encoding.
* However, we do detect if all values in the field are actually points and create a MultiPoint instead of a GeometryCollection.
*/
protected static class MultiValuesBytesRef implements MultiValuesCombiner {
private final BytesRefBlock valueBlock;
private final int valueCount;
private final BytesRef scratch = new BytesRef();
private final int firstValue;

MultiValuesBytesRef(BytesRefBlock valueBlock, int position) {
this.valueBlock = valueBlock;
this.firstValue = valueBlock.getFirstValueIndex(position);
this.valueCount = valueBlock.getValueCount(position);
}

@Override
public Geometry combined() {
int valueIndex = firstValue;
boolean allPoints = true;
if (valueCount == 1) {
return fromBytesRef(valueBlock.getBytesRef(valueIndex, scratch));
}
List<Geometry> geometries = new ArrayList<>();
while (valueIndex < firstValue + valueCount) {
geometries.add(fromBytesRef(valueBlock.getBytesRef(valueIndex++, scratch)));
if (geometries.getLast() instanceof Point == false) {
allPoints = false;
}
}
return allPoints ? new MultiPoint(asPointList(geometries)) : new GeometryCollection<>(geometries);
}

private List<Point> asPointList(List<Geometry> geometries) {
List<Point> points = new ArrayList<>(geometries.size());
for (Geometry geometry : geometries) {
points.add((Point) geometry);
}
return points;
}

protected Geometry fromBytesRef(BytesRef bytesRef) {
return SpatialCoordinateTypes.UNSPECIFIED.wkbToGeometry(bytesRef);
}
}

/**
* Point values read from doc-values will be encoded as in LogBlock. The block contains multiple rows, and within
* each row multiple values, so we need to efficiently iterate over only the values required for the requested row.
* Since the encoding differs for GEO and CARTESIAN, we need the decoder function to be passed in the constructor.
*/
protected static class MultiValuesLong implements MultiValuesCombiner {
private final LongBlock valueBlock;
private final Function<Long, Point> decoder;
private final int valueCount;
private final int firstValue;

MultiValuesLong(LongBlock valueBlock, int position, Function<Long, Point> decoder) {
this.valueBlock = valueBlock;
this.decoder = decoder;
this.firstValue = valueBlock.getFirstValueIndex(position);
this.valueCount = valueBlock.getValueCount(position);
}

@Override
public Geometry combined() {
int valueIndex = firstValue;
if (valueCount == 1) {
return decoder.apply(valueBlock.getLong(valueIndex));
}
List<Point> points = new ArrayList<>();
while (valueIndex < firstValue + valueCount) {
points.add(decoder.apply(valueBlock.getLong(valueIndex++)));
final GeometryDocValueReader reader = asGeometryDocValueReader(
coordinateEncoder,
shapeIndexer,
leftValue,
position,
spatialCoordinateType::longAsPoint
);
final Component2D component2D = asLuceneComponent2D(crsType, rightValue, position);
builder.appendBoolean(geometryRelatesGeometry(reader, component2D));
}
return new MultiPoint(points);
}
}
}
Loading

0 comments on commit 64127f2

Please sign in to comment.