Rework indexing

ibi-group · May 22, 2024 · 8443af2 · 8443af2
1 parent 7e31651
commit 8443af2
Show file tree

Hide file tree

Showing 7 changed files with 71 additions and 40 deletions.
diff --git a/src/ext/java/org/opentripplanner/ext/geocoder/LuceneIndex.java b/src/ext/java/org/opentripplanner/ext/geocoder/LuceneIndex.java
@@ -65,10 +65,11 @@ public class LuceneIndex implements Serializable {
   private final TransitService transitService;
   private final Analyzer analyzer;
   private final SuggestIndexSearcher searcher;
+  private final StopClusterMapper stopClusterMapper;
 
   public LuceneIndex(TransitService transitService) {
     this.transitService = transitService;
-    StopClusterMapper stopClusterMapper = new StopClusterMapper(transitService);
+    this.stopClusterMapper = new StopClusterMapper(transitService);
 
     this.analyzer =
       new PerFieldAnalyzerWrapper(
@@ -184,34 +185,54 @@ public Stream<StopCluster> queryStopClusters(String query) {
   }
 
   private StopCluster toStopCluster(Document document) {
-    var clusterId = FeedScopedId.parse(document.get(ID));
-    var name = document.get(NAME);
-    var code = document.get(CODE);
-    var lat = document.getField(LAT).numericValue().doubleValue();
-    var lon = document.getField(LON).numericValue().doubleValue();
-    var modes = Arrays.asList(document.getValues(MODE));
-    var agencies = Arrays
-      .stream(document.getValues(AGENCY_IDS))
-      .map(id -> transitService.getAgencyForId(FeedScopedId.parse(id)))
-      .filter(Objects::nonNull)
-      .map(StopClusterMapper::toAgency)
-      .toList();
-    var feedPublisher = StopClusterMapper.toFeedPublisher(
-      transitService.getFeedInfo(clusterId.getFeedId())
-    );
-    var primary = new StopCluster.Location(
-      clusterId,
-      code,
-      name,
-      new Coordinate(lat, lon),
-      modes,
-      agencies,
-      feedPublisher
-    );
+    var primaryId = FeedScopedId.parse(document.get(ID));
+    var primary = toLocation(primaryId);
 
     return new StopCluster(primary, List.of());
   }
 
+  private StopCluster.Location toLocation(FeedScopedId id) {
+    var loc = transitService.getStopLocation(id);
+    if (loc != null) {
+      var feedPublisher = StopClusterMapper.toFeedPublisher(
+        transitService.getFeedInfo(id.getFeedId())
+      );
+      var agencies = stopClusterMapper
+        .agenciesForStopLocation(loc)
+        .stream()
+        .map(StopClusterMapper::toAgency)
+        .toList();
+      return new StopCluster.Location(
+        loc.getId(),
+        loc.getCode(),
+        loc.getName().toString(),
+        new Coordinate(loc.getLat(), loc.getLon()),
+        List.of(),
+        agencies,
+        feedPublisher
+      );
+    } else {
+      var group = transitService.getStopLocationsGroup(id);
+      var feedPublisher = StopClusterMapper.toFeedPublisher(
+        transitService.getFeedInfo(id.getFeedId())
+      );
+      var agencies = stopClusterMapper
+        .agenciesForStopLocationsGroup(group)
+        .stream()
+        .map(StopClusterMapper::toAgency)
+        .toList();
+      return new StopCluster.Location(
+        group.getId(),
+        group.getCode(),
+        group.getName().toString(),
+        new Coordinate(group.getLat(), group.getLon()),
+        List.of(),
+        agencies,
+        feedPublisher
+      );
+    }
+  }
+
   static IndexWriterConfig iwcWithSuggestField(Analyzer analyzer, final Set<String> suggestFields) {
     IndexWriterConfig iwc = new IndexWriterConfig(analyzer);
     Codec filterCodec = new Lucene99Codec() {
@@ -245,7 +266,7 @@ private static void addToIndex(
     Document document = new Document();
     document.add(new StoredField(ID, id));
     document.add(new TextField(TYPE, typeName, Store.YES));
-    for(var name: names) {
+    for (var name : names) {
       document.add(new TextField(NAME, Objects.toString(name), Store.YES));
       document.add(new TextField(NAME_NGRAM, Objects.toString(name), Store.YES));
       document.add(new ContextSuggestField(SUGGEST, Objects.toString(name), 1, typeName));

diff --git a/src/ext/java/org/opentripplanner/ext/geocoder/LuceneStopCluster.java b/src/ext/java/org/opentripplanner/ext/geocoder/LuceneStopCluster.java
@@ -12,5 +12,5 @@ record LuceneStopCluster(
   Collection<I18NString> names,
   Collection<String> modes,
   Collection<String> codes,
-  StopCluster.Coordinate coordinate){
-}
+  StopCluster.Coordinate coordinate
+) {}
diff --git a/src/ext/java/org/opentripplanner/ext/geocoder/StopCluster.java b/src/ext/java/org/opentripplanner/ext/geocoder/StopCluster.java
@@ -15,10 +15,7 @@
  *  - if a stop has a parent station only the parent is returned
  *  - if stops are closer than 10 meters to each and have an identical name, only one is returned
  */
-record StopCluster(
-  Location primary,
-  Collection<Location> secondaries
-) {
+record StopCluster(Location primary, Collection<Location> secondaries) {
   /**
    * Easily serializable version of a coordinate
    */
@@ -42,8 +39,8 @@ public record Location(
     Collection<String> modes,
     List<Agency> agencies,
     @Nullable FeedPublisher feedPublisher
-  ){
-    public Location{
+  ) {
+    public Location {
       Objects.requireNonNull(id);
       Objects.requireNonNull(name);
       Objects.requireNonNull(coordinate);

diff --git a/src/ext/java/org/opentripplanner/ext/geocoder/StopClusterMapper.java b/src/ext/java/org/opentripplanner/ext/geocoder/StopClusterMapper.java
@@ -65,7 +65,11 @@ LuceneStopCluster map(StopLocationsGroup g) {
 
     var childStops = g.getChildStops();
     var ids = childStops.stream().map(s -> s.getId().toString()).toList();
-    var childNames = childStops.stream().map(StopLocation::getName).filter(Objects::nonNull).toList();
+    var childNames = childStops
+      .stream()
+      .map(StopLocation::getName)
+      .filter(Objects::nonNull)
+      .toList();
     var codes = childStops.stream().map(StopLocation::getCode).filter(Objects::nonNull).toList();
 
     return new LuceneStopCluster(
@@ -94,11 +98,11 @@ Optional<LuceneStopCluster> map(StopLocation sl) {
       });
   }
 
-  private List<Agency> agenciesForStopLocation(StopLocation stop) {
+  List<Agency> agenciesForStopLocation(StopLocation stop) {
     return transitService.getRoutesForStop(stop).stream().map(Route::getAgency).distinct().toList();
   }
 
-  private List<Agency> agenciesForStopLocationsGroup(StopLocationsGroup group) {
+  List<Agency> agenciesForStopLocationsGroup(StopLocationsGroup group) {
     return group
       .getChildStops()
       .stream()

diff --git a/src/main/java/org/opentripplanner/framework/collection/ListUtils.java b/src/main/java/org/opentripplanner/framework/collection/ListUtils.java
@@ -63,10 +63,9 @@ public static <T> List<T> distinctByKey(
    * return a list with one element.
    */
   public static <T> List<T> ofNullable(T input) {
-    if(input == null){
+    if (input == null) {
       return List.of();
-    }
-    else {
+    } else {
       return List.of(input);
     }
   }

diff --git a/src/main/java/org/opentripplanner/transit/model/site/GroupOfStations.java b/src/main/java/org/opentripplanner/transit/model/site/GroupOfStations.java
@@ -51,6 +51,12 @@ public WgsCoordinate getCoordinate() {
     return coordinate;
   }
 
+  @Nullable
+  @Override
+  public String getCode() {
+    return null;
+  }
+
   @Nonnull
   public Collection<StopLocation> getChildStops() {
     return this.childStations.stream().flatMap(s -> s.getChildStops().stream()).toList();

diff --git a/src/main/java/org/opentripplanner/transit/model/site/StopLocationsGroup.java b/src/main/java/org/opentripplanner/transit/model/site/StopLocationsGroup.java
@@ -1,6 +1,7 @@
 package org.opentripplanner.transit.model.site;
 
 import java.util.Collection;
+import javax.annotation.Nullable;
 import org.opentripplanner.framework.geometry.WgsCoordinate;
 import org.opentripplanner.framework.i18n.I18NString;
 import org.opentripplanner.framework.lang.ObjectUtils;
@@ -39,4 +40,7 @@ default double getLon() {
   default String logName() {
     return ObjectUtils.ifNotNull(getName(), Object::toString, null);
   }
+
+  @Nullable
+  String getCode();
 }