Skip to content

Commit

Permalink
Use push-down bbox filter for shapefiles (#757)
Browse files Browse the repository at this point in the history
  • Loading branch information
msbarry authored Dec 18, 2023
1 parent 289bbc6 commit a78e628
Show file tree
Hide file tree
Showing 3 changed files with 81 additions and 11 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -17,14 +17,15 @@
public class Bounds {

private static final Logger LOGGER = LoggerFactory.getLogger(Bounds.class);
public static final Bounds WORLD = new Bounds(null);

private Envelope latLon;
private Envelope world;
private TileExtents tileExtents;

private Geometry shape;

Bounds(Envelope latLon) {
public Bounds(Envelope latLon) {
set(latLon);
}

Expand All @@ -36,6 +37,10 @@ public Envelope world() {
return world == null ? GeoUtils.WORLD_BOUNDS : world;
}

public boolean isWorld() {
return latLon == null || latLon.equals(GeoUtils.WORLD_LAT_LON_BOUNDS);
}

public TileExtents tileExtents() {
if (tileExtents == null) {
tileExtents = TileExtents.computeFromWorldBounds(PlanetilerConfig.MAX_MAXZOOM, world(), shape);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@

import com.onthegomap.planetiler.Profile;
import com.onthegomap.planetiler.collection.FeatureGroup;
import com.onthegomap.planetiler.config.Bounds;
import com.onthegomap.planetiler.config.PlanetilerConfig;
import com.onthegomap.planetiler.stats.Stats;
import java.io.IOException;
Expand All @@ -19,9 +20,13 @@
import org.geotools.api.referencing.operation.OperationNotFoundException;
import org.geotools.api.referencing.operation.TransformException;
import org.geotools.data.shapefile.ShapefileDataStore;
import org.geotools.factory.CommonFactoryFinder;
import org.geotools.feature.FeatureCollection;
import org.geotools.geometry.jts.JTS;
import org.geotools.geometry.jts.ReferencedEnvelope;
import org.geotools.referencing.CRS;
import org.geotools.util.factory.GeoTools;
import org.locationtech.jts.geom.Envelope;
import org.locationtech.jts.geom.Geometry;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
Expand All @@ -45,22 +50,44 @@ public class ShapefileReader extends SimpleReader<SimpleFeature> {
private MathTransform transformToLatLon;

public ShapefileReader(String sourceProjection, String sourceName, Path input) {
this(sourceProjection, sourceName, input, Bounds.WORLD);
}

public ShapefileReader(String sourceProjection, String sourceName, Path input, Bounds bounds) {
super(sourceName);
this.layer = input.getFileName().toString().replaceAll("\\.shp$", "");
dataStore = open(input);
try {
String typeName = dataStore.getTypeNames()[0];
FeatureSource<SimpleFeatureType, org.geotools.api.feature.simple.SimpleFeature> source = dataStore
.getFeatureSource(typeName);

inputSource = source.getFeatures(Filter.INCLUDE);
CoordinateReferenceSystem src =
sourceProjection == null ? source.getSchema().getCoordinateReferenceSystem() : CRS.decode(sourceProjection);
CoordinateReferenceSystem dest = CRS.decode("EPSG:4326", true);
transformToLatLon = findMathTransform(input, src, dest);
if (transformToLatLon.isIdentity()) {
transformToLatLon = null;
}

Filter filter = Filter.INCLUDE;

Envelope env = bounds.latLon();
if (!bounds.isWorld()) {
var ff = CommonFactoryFinder.getFilterFactory(GeoTools.getDefaultHints());
var schema = source.getSchema();

String geometryPropertyName = schema.getGeometryDescriptor().getLocalName();

var bbox = new ReferencedEnvelope(env.getMinX(), env.getMaxX(), env.getMinY(), env.getMaxY(), dest);
try {
var bbox2 = bbox.transform(schema.getGeometryDescriptor().getCoordinateReferenceSystem(), true);
filter = ff.bbox(ff.property(geometryPropertyName), bbox2);
} catch (TransformException e) {
// just use include filter
}
}

inputSource = source.getFeatures(filter);
attributeNames = new String[inputSource.getSchema().getAttributeCount()];
for (int i = 0; i < attributeNames.length; i++) {
attributeNames[i] = inputSource.getSchema().getDescriptor(i).getLocalName();
Expand Down Expand Up @@ -105,7 +132,7 @@ public static void processWithProjection(String sourceProjection, String sourceN
SourceFeatureProcessor.processFiles(
sourceName,
sourcePaths,
path -> new ShapefileReader(sourceProjection, sourceName, path),
path -> new ShapefileReader(sourceProjection, sourceName, path, config.bounds()),
writer, config, profile, stats
);
}
Expand Down
Original file line number Diff line number Diff line change
@@ -1,19 +1,21 @@
package com.onthegomap.planetiler.reader;

import static com.onthegomap.planetiler.TestUtils.newPoint;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertTrue;

import com.onthegomap.planetiler.TestUtils;
import com.onthegomap.planetiler.config.Bounds;
import com.onthegomap.planetiler.geo.GeoUtils;
import com.onthegomap.planetiler.stats.Stats;
import com.onthegomap.planetiler.util.FileUtils;
import com.onthegomap.planetiler.worker.WorkerPipeline;
import java.io.IOException;
import java.nio.file.FileSystems;
import java.nio.file.Path;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CopyOnWriteArrayList;
import org.geotools.api.data.SimpleFeatureStore;
import org.geotools.api.referencing.FactoryException;
import org.geotools.api.referencing.operation.TransformException;
Expand All @@ -29,12 +31,18 @@
import org.junit.jupiter.api.condition.DisabledOnOs;
import org.junit.jupiter.api.condition.OS;
import org.junit.jupiter.api.io.TempDir;
import org.locationtech.jts.geom.Envelope;
import org.locationtech.jts.geom.Geometry;
import org.locationtech.jts.geom.Point;

class ShapefileReaderTest {
@TempDir
private Path tempDir;
private static final Envelope env = newPoint(-77.12911152370515, 38.79930767201779).getEnvelopeInternal();
private static final int numInEnv = 18;
static {
env.expandBy(0.1);
}

@Test
@Timeout(30)
Expand All @@ -55,6 +63,35 @@ void testReadShapefileUnzipped() throws IOException {
testReadShapefile(dest.resolve("shapefile").resolve("stations.shp"));
}

@Test
@Timeout(30)
void testReadShapefileWithBoundingBox() {
var dest = tempDir.resolve("shapefile.zip");
FileUtils.unzipResource("/shapefile.zip", dest);
try (
var reader = new ShapefileReader(null, "test", dest.resolve("shapefile").resolve("stations.shp"), new Bounds(env))
) {
for (int i = 1; i <= 2; i++) {
assertEquals(numInEnv, reader.getFeatureCount());
List<Geometry> points = new CopyOnWriteArrayList<>();
WorkerPipeline.start("test", Stats.inMemory())
.fromGenerator("source", reader::readFeatures)
.addBuffer("reader_queue", 100, 1)
.sinkToConsumer("counter", 1, elem -> {
assertTrue(elem.getTag("name") instanceof String);
assertEquals("test", elem.getSource());
assertEquals("stations", elem.getSourceLayer());
points.add(elem.latLonGeometry());
}).await();
assertEquals(numInEnv, points.size());
var gc = GeoUtils.JTS_FACTORY.createGeometryCollection(points.toArray(new Geometry[0]));
var centroid = gc.getCentroid();
assertEquals(-77.0934256, centroid.getX(), 1e-5, "iter " + i);
assertEquals(38.8509022, centroid.getY(), 1e-5, "iter " + i);
}
}
}

@Test
void testReadShapefileLeniently(@TempDir Path dir) throws IOException, TransformException, FactoryException {
var shpPath = dir.resolve("test.shp");
Expand Down Expand Up @@ -82,7 +119,7 @@ void testReadShapefileLeniently(@TempDir Path dir) throws IOException, Transform
featureStore.setTransaction(transaction);
var collection = new DefaultFeatureCollection();
var featureBuilder = new SimpleFeatureBuilder(type);
featureBuilder.add(TestUtils.newPoint(1, 2));
featureBuilder.add(newPoint(1, 2));
featureBuilder.add(3);
var feature = featureBuilder.buildFeature(null);
collection.add(feature);
Expand All @@ -92,7 +129,7 @@ void testReadShapefileLeniently(@TempDir Path dir) throws IOException, Transform

try (var reader = new ShapefileReader(null, "test", shpPath)) {
assertEquals(1, reader.getFeatureCount());
List<SimpleFeature> features = new ArrayList<>();
List<SimpleFeature> features = new CopyOnWriteArrayList<>();
reader.readFeatures(features::add);
assertEquals(10.5113, features.getFirst().latLonGeometry().getCentroid().getX(), 1e-4);
assertEquals(0, features.getFirst().latLonGeometry().getCentroid().getY(), 1e-4);
Expand All @@ -105,8 +142,8 @@ private static void testReadShapefile(Path path) {

for (int i = 1; i <= 2; i++) {
assertEquals(86, reader.getFeatureCount());
List<Geometry> points = new ArrayList<>();
List<String> names = new ArrayList<>();
List<Geometry> points = new CopyOnWriteArrayList<>();
List<String> names = new CopyOnWriteArrayList<>();
WorkerPipeline.start("test", Stats.inMemory())
.fromGenerator("source", reader::readFeatures)
.addBuffer("reader_queue", 100, 1)
Expand All @@ -117,12 +154,13 @@ private static void testReadShapefile(Path path) {
points.add(elem.latLonGeometry());
names.add(elem.getTag("name").toString());
}).await();
assertEquals(numInEnv, points.stream().filter(point -> env.contains(point.getCoordinate())).count());
assertEquals(86, points.size());
assertTrue(names.contains("Van Dörn Street"));
var gc = GeoUtils.JTS_FACTORY.createGeometryCollection(points.toArray(new Geometry[0]));
var centroid = gc.getCentroid();
assertEquals(-77.0297995, centroid.getX(), 5, "iter " + i);
assertEquals(38.9119684, centroid.getY(), 5, "iter " + i);
assertEquals(-77.0297995, centroid.getX(), 1e-5, "iter " + i);
assertEquals(38.9119684, centroid.getY(), 1e-5, "iter " + i);
}
}
}
Expand Down

0 comments on commit a78e628

Please sign in to comment.