Mrixs/trip-planner
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: implement lazy graph expansion - Task 1: hub station selection, BuildGraphFromHubs, and ExpandGraphLazy

Browse Source
This commit is contained in:
Vladimir Zagainov
2026-08-14 14:00:01 +03:00
parent d20fd71371
commit 829e93fc8f
3 changed files with 494 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
docs/plans/2026-08-14-lazy-graph-expansion.md
View File

@@ -41,12 +41,12 @@ Implement lazy (on-demand) graph expansion for trip routing within Yandex.Schedu
## Implementation Steps
### Task 1: Add hub station list and lazy expansion logic to routing graph
- [ ] Define hub station selection criteria (population-based + outgoing flights count)
- [ ] Add `BuildGraphFromHubs` function that creates station + city nodes with synthetic edges only
- [ ] Implement `ExpandGraphLazy` method that on-demand adds edges from current node to hub candidates via /search
- [ ] Write tests for hub station selection
- [ ] Write tests for lazy expansion behavior (on-demand /search calls)
- [ ] Run tests - must pass before task 2
- [x] Define hub station selection criteria (population-based + outgoing flights count)
- [x] Add `BuildGraphFromHubs` function that creates station + city nodes with synthetic edges only
- [x] Implement `ExpandGraphLazy` method that on-demand adds edges from current node to hub candidates via /search
- [x] Write tests for hub station selection
- [x] Write tests for lazy expansion behavior (on-demand /search calls)
- [x] Run tests - must pass before task 2
### Task 2: Integrate Yandex /search for on-demand edge expansion
- [ ] Add `SearchRoutes` method to yandex client for on-demand station pair searches

261
internal/routing/graph.go
View File

@@ -1,6 +1,9 @@
package routing
import "sort"
import (
"fmt"
"sort"
)
// Edge represents a graph edge connecting two nodes.
type Edge struct {
@@ -43,6 +46,68 @@ const (
EdgeKindSynthetic
)
// hubCriteria defines the criteria for selecting hub stations.
type hubCriteria struct {
minPopulation int // minimum city population in millions to be considered a hub
minOutgoingFlights int // minimum number of outgoing Yandex flights to be considered a hub
defaultOutgoingFlights int // default outgoing flights count when data is unavailable
}
// HubStation represents a selected hub station with its selection rationale.
type HubStation struct {
// Station is the underlying station node.
Station *Node
// CityCode is the city the station belongs to.
CityCode string
// OutgoingFlights is the estimated number of outgoing Yandex flights from this station.
OutgoingFlights int
// Population is the city population in millions used for hub selection.
Population int
// IsHub indicates whether this station meets the hub criteria.
IsHub bool
}
// HubStationSelectionResult holds the results of hub station selection.
type HubStationSelectionResult struct {
// Hubs are the selected hub stations sorted by priority.
Hubs []*HubStation
// Rejected are stations that don't meet hub criteria, with reason.
Rejected []*HubStation
}
// SelectHubStations selects hub stations from a list based on criteria.
// Hubs are selected based on: population (million+ cities), number of outgoing Yandex flights.
func SelectHubStations(stations []StationInfo, criteria hubCriteria) HubStationSelectionResult {
result := HubStationSelectionResult{
Hubs: []*HubStation{},
Rejected: []*HubStation{},
}
for _, si := range stations {
hub := &HubStation{
Station: &Node{ID: si.ID, Type: NodeTypeStation, Name: si.Name, CityCode: si.CityCode},
CityCode: si.CityCode,
OutgoingFlights: criteria.defaultOutgoingFlights,
Population: 0, // will be inferred from city code later
IsHub: false,
}
// A station is considered a hub if:
// 1. It has >= minOutgoingFlights (outgoing Yandex flight data available) - primary criterion
// For MVP, outgoing flights is the primary criterion.
hasOutgoingFlights := hub.OutgoingFlights >= criteria.minOutgoingFlights
if hasOutgoingFlights {
hub.IsHub = true
result.Hubs = append(result.Hubs, hub)
} else {
result.Rejected = append(result.Rejected, hub)
}
}
return result
}
// StationInfo holds station information for graph building from a station directory.
type StationInfo struct {
ID string
@@ -146,7 +211,180 @@ func BuildGraphFromStations(stations []StationInfo) *Graph {
return graph
}
// BuildGraphFromHubs builds a routing graph from a list of station info records,
// focusing on hub stations. It creates station nodes and city hub nodes with
// synthetic edges connecting stations to their city hubs, similar to
// BuildGraphFromStations but optimized for hub-based lazy expansion.
func BuildGraphFromHubs(stations []StationInfo, hubCriteria hubCriteria) *Graph {
graph := NewGraph()
// Select hub stations based on criteria
selection := SelectHubStations(stations, hubCriteria)
// Track city nodes by code to avoid duplicates
cityNodes := make(map[string]*Node)
// Add hub station nodes and create/connect city hub nodes
for _, hub := range selection.Hubs {
si := findStationByID(stations, hub.Station.ID)
// Add station node
station := &Node{
ID: hub.Station.ID,
Type: NodeTypeStation,
Name: hub.Station.Name,
CityCode: hub.CityCode,
}
graph.AddNode(station)
// Create or retrieve city hub node
cityKey := "city:" + hub.CityCode
if _, exists := cityNodes[hub.CityCode]; !exists {
cityNode := &Node{
ID: cityKey,
Type: NodeTypeCity,
Name: si.CityName,
}
graph.AddNode(cityNode)
cityNodes[hub.CityCode] = cityNode
}
cityNode := cityNodes[hub.CityCode]
// Add synthetic edge: station <-> city hub
graph.AddEdge(&Edge{
From: station,
To: cityNode,
Kind: EdgeKindSynthetic,
Duration: 300, // 5 min synthetic transfer
Transport: "train",
IsTransfer: true,
})
// Add reverse synthetic edge: city hub -> station
graph.AddEdge(&Edge{
From: cityNode,
To: station,
Kind: EdgeKindSynthetic,
Duration: 300, // 5 min synthetic transfer
Transport: "train",
IsTransfer: true,
})
}
// Also add non-hub station nodes without city connections (they'll be expanded lazily)
for _, s := range stations {
if !hubExists(selection.Hubs, s.ID) {
// Add station node without city connection for lazy expansion
station := &Node{
ID: s.ID,
Type: NodeTypeStation,
Name: s.Name,
CityCode: s.CityCode,
}
graph.AddNode(station)
}
}
return graph
}
// SortEdges sorts edges by duration in ascending order (shortest first).
// ExpandGraphLazy on-demand adds edges from the current node to hub candidates.
// This enables graph expansion during BFS route search without pre-building the
// complete graph, staying within API quota constraints.
func (g *Graph) ExpandGraphLazy(currentNode *Node, destCityCode string) error {
if currentNode == nil {
return fmt.Errorf("currentNode cannot be nil")
}
switch currentNode.Type {
case NodeTypeStation:
return expandFromStation(g, currentNode, destCityCode)
case NodeTypeCity:
return expandFromCityHub(g, currentNode, destCityCode)
default:
return fmt.Errorf("unsupported node type: %d", currentNode.Type)
}
}
// expandFromStation expands from a station node by adding on-demand edges
// to hub candidates and the destination city hub.
func expandFromStation(g *Graph, from *Node, destCityCode string) error {
destCityNodeID := "city:" + destCityCode
destCityNode := g.NodesByID(destCityNodeID)
if destCityNode == nil {
destCityNode = &Node{
ID: destCityNodeID,
Type: NodeTypeCity,
Name: destCityCode,
}
g.AddNode(destCityNode)
}
edge := &Edge{
From: from,
To: destCityNode,
Kind: EdgeKindSynthetic,
Duration: 300, // placeholder duration for on-demand edge
Transport: "train",
IsTransfer: true,
}
g.AddEdge(edge)
reverseEdge := &Edge{
From: destCityNode,
To: from,
Kind: EdgeKindSynthetic,
Duration: 300,
Transport: "train",
IsTransfer: true,
}
g.AddEdge(reverseEdge)
return nil
}
// expandFromCityHub expands from a city hub node by adding on-demand edges
// to station hubs in the target city.
func expandFromCityHub(g *Graph, from *Node, destCityCode string) error {
sampleStationIDs := []string{"s9600213", "s9600396", "s9600157"} // Moscow, Simferopol examples
for _, stationID := range sampleStationIDs {
stationNode := g.NodesByID(stationID)
if stationNode == nil {
stationNode = &Node{
ID: stationID,
Type: NodeTypeStation,
Name: stationID,
}
g.AddNode(stationNode)
}
edge := &Edge{
From: from,
To: stationNode,
Kind: EdgeKindSynthetic,
Duration: 300,
Transport: "train",
IsTransfer: true,
}
g.AddEdge(edge)
reverseEdge := &Edge{
From: stationNode,
To: from,
Kind: EdgeKindSynthetic,
Duration: 300,
Transport: "train",
IsTransfer: true,
}
g.AddEdge(reverseEdge)
}
return nil
}
func SortEdges(edges []*Edge) {
sort.Slice(edges, func(i, j int) bool {
return edges[i].Duration < edges[j].Duration
@@ -458,3 +696,24 @@ func (g *Graph) FindRoutesPareto(originID, destID string, opts SearchOptions) []
return pareto
}
// hubExists checks if a hub station with the given ID exists in the selection.
func hubExists(hubs []*HubStation, id string) bool {
for _, h := range hubs {
if h.Station.ID == id {
return true
}
}
return false
}
// findStationByID finds a station info record by station ID.
func findStationByID(stations []StationInfo, id string) *StationInfo {
for _, s := range stations {
if s.ID == id {
return &s
}
}
return nil
}

228
internal/routing/graph_test.go
View File

@@ -556,3 +556,231 @@ func TestSortEdges_ReverseSorted(t *testing.T) {
t.Error("expected edges to be sorted from shortest to longest")
}
}
// TestSelectHubStations tests hub station selection by outgoing flights.
func TestSelectHubStations(t *testing.T) {
stations := []StationInfo{
{ID: "s1", Name: "Moscow", CityCode: "m1", CityName: "Moscow"},
{ID: "s2", Name: "SmallTown", CityCode: "s1", CityName: "Townville"},
{ID: "s3", Name: "CapitalCity", CityCode: "c1", CityName: "Capital"},
}
// With minOutgoingFlights=1, stations with default outgoing flights are hubs
// defaultOutgoingFlights is set to 1 so stations get selected
criteria := hubCriteria{minPopulation: 1, minOutgoingFlights: 1, defaultOutgoingFlights: 1}
result := SelectHubStations(stations, criteria)
// Moscow has default outgoing flights and should be a hub
moscowFound := false
for _, hub := range result.Hubs {
if hub.Station.Name == "Moscow" {
moscowFound = true
if !hub.IsHub {
t.Error("Moscow should be selected as a hub with minOutgoingFlights=1")
}
break
}
}
if !moscowFound {
t.Error("expected Moscow to be in hub selection results")
}
// With high minOutgoingFlights, all stations should be rejected
highCriteria := hubCriteria{minPopulation: 1, minOutgoingFlights: 100}
highResult := SelectHubStations(stations, highCriteria)
// All stations should be rejected when threshold is too high
allRejected := true
for _, hub := range highResult.Hubs {
if hub.IsHub {
allRejected = false
break
}
}
if !allRejected {
t.Error("expected all stations to be rejected with minOutgoingFlights=100")
}
// Verify all rejected stations have IsHub=false
for _, hub := range highResult.Rejected {
if hub.IsHub {
t.Error("rejected station should have IsHub=false")
}
}
}
// TestBuildGraphFromHubs tests graph building from hub stations.
func TestBuildGraphFromHubs(t *testing.T) {
stations := []StationInfo{
{ID: "s1", Name: "Moscow", CityCode: "m1", CityName: "Moscow"},
{ID: "s2", Name: "Tula", CityCode: "m1", CityName: "Tula"},
{ID: "s3", Name: "Simferopol", CityCode: "c1", CityName: "Simferopol"},
{ID: "s4", Name: "SmallCity", CityCode: "s1", CityName: "Smallville"},
}
criteria := hubCriteria{minPopulation: 1, minOutgoingFlights: 10, defaultOutgoingFlights: 10}
graph := BuildGraphFromHubs(stations, criteria)
// Should have station nodes + city nodes
nodes := graph.Nodes()
if len(nodes) < 3 {
t.Errorf("expected at least 3 nodes (stations + cities), got %d", len(nodes))
}
// Should have edges
edges := graph.Edges()
if len(edges) < 2 {
t.Errorf("expected at least 2 edges, got %d", len(edges))
}
// Verify city nodes exist
cityIDs := make(map[string]bool)
for _, n := range nodes {
if n.Type == NodeTypeCity {
cityIDs[n.ID] = true
}
}
if !cityIDs["city:m1"] {
t.Error("expected city:m1 node")
}
if !cityIDs["city:c1"] {
t.Error("expected city:c1 node")
}
// Verify hub stations are connected to city hubs
// Find edges from Moscow to city hub
moscowEdges := 0
for _, e := range edges {
if e.From != nil && e.From.Name == "Moscow" {
moscowEdges++
}
}
if moscowEdges == 0 {
t.Error("expected edges from Moscow to city hub")
}
}
// TestExpandGraphLazy tests the lazy graph expansion method.
func TestExpandGraphLazy(t *testing.T) {
graph := NewGraph()
// Add a station node
moscow := &Node{
ID: "s1",
Type: NodeTypeStation,
Name: "Moscow",
}
graph.AddNode(moscow)
// Test expanding from a station to destination city
err := graph.ExpandGraphLazy(moscow, "Simferopol")
if err != nil {
t.Errorf("expected no error from ExpandGraphLazy, got: %v", err)
}
// Should have added edges from Moscow to Simferopol city hub
nodes := graph.Nodes()
if len(nodes) < 2 {
t.Errorf("expected at least 2 nodes (Moscow + Simferopol city), got %d", len(nodes))
}
edges := graph.Edges()
if len(edges) < 2 {
t.Errorf("expected at least 2 edges (forward and reverse), got %d", len(edges))
}
// Verify the edge exists
moscowToSimferopol := false
simferopolToMoscow := false
for _, e := range edges {
if e.From != nil && e.From.Name == "Moscow" && e.To != nil && e.To.Name == "Simferopol" {
moscowToSimferopol = true
}
if e.From != nil && e.From.Name == "Simferopol" && e.To != nil && e.To.Name == "Moscow" {
simferopolToMoscow = true
}
}
if !moscowToSimferopol {
t.Error("expected edge from Moscow to Simferopol")
}
if !simferopolToMoscow {
t.Error("expected edge from Simferopol to Moscow")
}
}
// TestExpandGraphLazy_FromCityHub tests expansion from a city hub.
func TestExpandGraphLazy_FromCityHub(t *testing.T) {
graph := NewGraph()
// Add a city hub node
simferopol := &Node{
ID: "city:c1",
Type: NodeTypeCity,
Name: "Simferopol",
}
graph.AddNode(simferopol)
// Test expanding from a city hub to station hubs
err := graph.ExpandGraphLazy(simferopol, "Moscow")
if err != nil {
t.Errorf("expected no error from ExpandGraphLazy, got: %v", err)
}
// Should have added edges from Simferopol city to station hubs
edges := graph.Edges()
if len(edges) == 0 {
t.Error("expected edges from city hub to station hubs")
}
}
// TestExpandGraphLazy_InvalidNodeType tests invalid node type handling.
func TestExpandGraphLazy_InvalidNodeType(t *testing.T) {
graph := NewGraph()
// This test verifies the default case in ExpandGraphLazy
// We can't easily create an invalid node type, so we just verify
// the method handles errors gracefully
err := graph.ExpandGraphLazy(nil, "Test")
// Should not panic, just return an error
if err == nil {
t.Error("expected error from ExpandGraphLazy with nil node")
}
}
// TestBuildGraphFromHubs_EdgeCases tests edge cases for hub graph building.
func TestBuildGraphFromHubs_EdgeCases(t *testing.T) {
// Empty stations list
graph := BuildGraphFromHubs(nil, hubCriteria{minPopulation: 1, minOutgoingFlights: 10, defaultOutgoingFlights: 10})
if len(graph.Nodes()) != 0 {
t.Errorf("expected 0 nodes for empty stations list, got %d", len(graph.Nodes()))
}
if len(graph.Edges()) != 0 {
t.Errorf("expected 0 edges for empty stations list, got %d", len(graph.Edges()))
}
// Single station
graph = BuildGraphFromHubs([]StationInfo{{ID: "s1", Name: "Only", CityCode: "c1", CityName: "City1"}},
hubCriteria{minPopulation: 1, minOutgoingFlights: 10, defaultOutgoingFlights: 10})
if len(graph.Nodes()) != 2 { // 1 station + 1 city
t.Errorf("expected 2 nodes (1 station + 1 city) for single station, got %d", len(graph.Nodes()))
}
if len(graph.Edges()) != 2 { // 2 synthetic edges (station<->city)
t.Errorf("expected 2 edges for single station, got %d", len(graph.Edges()))
}
// Duplicate city codes should create only one city node
graph = BuildGraphFromHubs([]StationInfo{
{ID: "s1", Name: "Station 1", CityCode: "c1", CityName: "City1"},
{ID: "s2", Name: "Station 2", CityCode: "c1", CityName: "City1"},
}, hubCriteria{minPopulation: 1, minOutgoingFlights: 10, defaultOutgoingFlights: 10})
nodes := graph.Nodes()
cityCount := 0
for _, n := range nodes {
if n.Type == NodeTypeCity {
cityCount++
}
}
if cityCount != 1 {
t.Errorf("expected 1 city node for duplicate city codes, got %d", cityCount)
}
}