diff --git a/locode_generate_test.go b/locode_generate_test.go new file mode 100644 index 0000000..e568756 --- /dev/null +++ b/locode_generate_test.go @@ -0,0 +1,30 @@ +package main + +import ( + "flag" + "testing" +) + +var ( + in = flag.String(locodeGenerateInputFlag, "", "List of paths to UN/LOCODE tables (csv)") + subdiv = flag.String(locodeGenerateSubDivFlag, "", "Path to UN/LOCODE subdivision database (csv)") + airports = flag.String(locodeGenerateAirportsFlag, "", "Path to OpenFlights airport database (csv)") + countries = flag.String(locodeGenerateCountriesFlag, "", "Path to OpenFlights country database (csv)") + continents = flag.String(locodeGenerateContinentsFlag, "", "Path to continent polygons (GeoJSON)") + out = flag.String(locodeGenerateOutputFlag, "", "Target path for generated database") +) + +func BenchmarkLocodeGenerate(b *testing.B) { + locodeGenerateInPaths = append(locodeGenerateInPaths, *in) + locodeGenerateSubDivPath = *subdiv + locodeGenerateAirportsPath = *airports + locodeGenerateCountriesPath = *countries + locodeGenerateContinentsPath = *continents + locodeGenerateOutPath = *out + + b.ResetTimer() + + for i := 0; i < b.N; i++ { + locodeGenerateCmd.Run(locodeGenerateCmd, []string{}) + } +} diff --git a/pkg/locode/db/continents/geojson/calls.go b/pkg/locode/db/continents/geojson/calls.go index b5a58d1..a766cbd 100644 --- a/pkg/locode/db/continents/geojson/calls.go +++ b/pkg/locode/db/continents/geojson/calls.go @@ -8,6 +8,7 @@ import ( "github.com/paulmach/orb" "github.com/paulmach/orb/geojson" "github.com/paulmach/orb/planar" + "github.com/paulmach/orb/quadtree" ) const continentProperty = "Continent" @@ -36,22 +37,24 @@ func (db *DB) PointContinent(point *locodedb.Point) (*locodedb.Continent, error) minDst float64 ) - for _, feature := range db.features { - if multiPolygon, ok := feature.Geometry.(orb.MultiPolygon); ok { - if planar.MultiPolygonContains(multiPolygon, planarPoint) { + pointer := db.tree.Matching(planarPoint, func(p orb.Pointer) bool { + return planar.PolygonContains( + p.(*geojson.Feature).Geometry.(orb.Polygon), + planarPoint, + ) + }) + + if pointer != nil { + continent = pointer.(*geojson.Feature).Properties.MustString(continentProperty) + } + + if continent == "" { + for _, feature := range db.features { + distance := planar.DistanceFrom(feature.Geometry, planarPoint) + if minDst == 0 || minDst > distance { + minDst = distance continent = feature.Properties.MustString(continentProperty) - break } - } else if polygon, ok := feature.Geometry.(orb.Polygon); ok { - if planar.PolygonContains(polygon, planarPoint) { - continent = feature.Properties.MustString(continentProperty) - break - } - } - distance := planar.DistanceFrom(feature.Geometry, planarPoint) - if minDst == 0 || minDst > distance { - minDst = distance - continent = feature.Properties.MustString(continentProperty) } } @@ -73,6 +76,39 @@ func (db *DB) init() error { db.features = features.Features + err = db.buildQuadtree() + if err != nil { + return fmt.Errorf("could not build quadtree: %w", err) + } + + return nil +} + +func (db *DB) buildQuadtree() error { + db.tree = quadtree.New(orb.Bound{ + Min: orb.Point{-180, -180}, + Max: orb.Point{180, 180}, + }) + + for _, feature := range db.features { + var multiPolygon orb.MultiPolygon + + if polygon, ok := feature.Geometry.(orb.Polygon); ok { + multiPolygon = append(multiPolygon, polygon) + } else { + multiPolygon = feature.Geometry.(orb.MultiPolygon) + } + + for _, polygon := range multiPolygon { + newFeature := geojson.NewFeature(polygon) + newFeature.Properties = feature.Properties.Clone() + err := db.tree.Add(newFeature) + if err != nil { + return err + } + } + } + return nil } diff --git a/pkg/locode/db/continents/geojson/db.go b/pkg/locode/db/continents/geojson/db.go index ee43bd8..244da02 100644 --- a/pkg/locode/db/continents/geojson/db.go +++ b/pkg/locode/db/continents/geojson/db.go @@ -5,6 +5,7 @@ import ( "sync" "github.com/paulmach/orb/geojson" + "github.com/paulmach/orb/quadtree" ) // Prm groups the required parameters of the DB's constructor. @@ -31,6 +32,8 @@ type DB struct { once sync.Once features []*geojson.Feature + + tree *quadtree.Quadtree } const invalidPrmValFmt = "invalid parameter %s (%T):%v" diff --git a/pkg/locode/db/db.go b/pkg/locode/db/db.go index a4a09d1..90b19ed 100644 --- a/pkg/locode/db/db.go +++ b/pkg/locode/db/db.go @@ -80,7 +80,7 @@ func FillDatabase(table SourceTable, airports AirportDB, continents ContinentsDB var errG errgroup.Group // Pick some sane default, after this the performance stopped increasing. - errG.SetLimit(runtime.NumCPU() * 4) + errG.SetLimit(runtime.NumCPU() * 16) _ = table.IterateAll(func(tableRecord locode.Record) error { errG.Go(func() error { return processTableRecord(tableRecord, airports, continents, names, db)