Index block_trip_entry by block_id

gtfsdb/block_trip_entry_index_test.go

@@ -0,0 +1,42 @@
+package gtfsdb
+
+import (
+	"database/sql"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+// GetBlocksForBlockTripIndexIDs joins block_trip_entry to itself on block_id. Without
+// an index there, SQLite scans every block entry for every route trip, which made
+// trips-for-route take seconds on NJ Transit's 33k-trip feed.
+func TestBlockTripEntryBlockIDJoinUsesIndex(t *testing.T) {
+	client := newTestClientWithRABA(t)
+	ctx := t.Context()
+
+	var routeID, serviceID string
+	var indexID int64
+	require.NoError(t, client.DB.QueryRowContext(ctx, `
+		SELECT trips.route_id, bte.service_id, bte.block_trip_index_id
+		FROM block_trip_entry bte JOIN trips ON trips.id = bte.trip_id
+		WHERE bte.block_id IS NOT NULL LIMIT 1`).Scan(&routeID, &serviceID, &indexID))
+
+	recorder := &queryRecorder{DBTX: client.DB}
+	_, err := New(recorder).GetBlocksForBlockTripIndexIDs(ctx, GetBlocksForBlockTripIndexIDsParams{
+		FromTime:         sql.NullInt64{Int64: 0, Valid: true},
+		ToTime:           sql.NullInt64{Int64: int64(24 * time.Hour), Valid: true},
+		RouteID:          routeID,
+		IndexIds:         []int64{indexID},
+		ActiveServiceIds: []string{serviceID},
+		RouteServiceIds:  []string{serviceID},
+	})
+	require.NoError(t, err)
+
+	// The planner may start the self-join from either side, so accept the index
+	// on either alias.
+	plan := strings.Join(explainRecordedQuery(t, client, recorder), "\n")
+	assert.Regexp(t, `SEARCH \w+ USING (COVERING )?INDEX idx_block_trip_entry_block_id \(block_id=\?`, plan)
+}

gtfsdb/query_plan_test.go

@@ -0,0 +1,41 @@
+package gtfsdb
+
+import (
+	"context"
+	"database/sql"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+// queryRecorder captures the SQL and arguments a generated query method sends,
+// so a test can explain the real query rather than a copy of it.
+type queryRecorder struct {
+	DBTX
+	query string
+	args  []any
+}
+
+func (db *queryRecorder) QueryContext(ctx context.Context, query string, args ...any) (*sql.Rows, error) {
+	db.query, db.args = query, args
+	return db.DBTX.QueryContext(ctx, query, args...)
+}
+
+// explainRecordedQuery returns the detail column of each EXPLAIN QUERY PLAN row
+// for the query the recorder last captured.
+func explainRecordedQuery(t *testing.T, client *Client, recorder *queryRecorder) []string {
+	t.Helper()
+	rows, err := client.DB.QueryContext(t.Context(), "EXPLAIN QUERY PLAN "+recorder.query, recorder.args...)
+	require.NoError(t, err)
+	defer func() { _ = rows.Close() }()
+
+	var details []string
+	for rows.Next() {
+		var id, parent, unused int
+		var detail string
+		require.NoError(t, rows.Scan(&id, &parent, &unused, &detail))
+		details = append(details, detail)
+	}
+	require.NoError(t, rows.Err())
+	return details
+}

gtfsdb/schema.sql

@@ -454,6 +454,11 @@ CREATE INDEX IF NOT EXISTS idx_block_trip_index_service_ids ON block_trip_index
 -- migrate
 CREATE INDEX IF NOT EXISTS idx_block_trip_entry_index_id ON block_trip_entry (block_trip_index_id);
 
+-- migrate
+-- GetBlocksForBlockTripIndexIDs (a route's active blocks, for trips-for-route)
+-- and GetBlockTripIndexIDsForBlocks both match block_trip_entry rows by block_id.
+CREATE INDEX IF NOT EXISTS idx_block_trip_entry_block_id ON block_trip_entry (block_id);
+
 -- migrate
 CREATE INDEX IF NOT EXISTS idx_block_trip_entry_trip_id ON block_trip_entry (trip_id);
 

gtfsdb/stops_rtree_test.go

@@ -2,7 +2,6 @@ package gtfsdb
 
 import (
 	"context"
-	"database/sql"
 	"strings"
 	"testing"
 
@@ -47,44 +46,26 @@ func TestGetActiveRoutesWithinBounds_CenteredOnStop(t *testing.T) {
 	assert.ElementsMatch(t, []string{"15", "154", "159", "160", "44X"}, stop4000RouteIDs)
 }
 
-type routeQueryRecorder struct {
-	DBTX
-	query string
-	args  []any
-}
-
-func (db *routeQueryRecorder) QueryContext(ctx context.Context, query string, args ...any) (*sql.Rows, error) {
-	db.query, db.args = query, args
-	return db.DBTX.QueryContext(ctx, query, args...)
-}
-
 func TestGetActiveRoutesWithinBounds_FilteredStopTimesLookup(t *testing.T) {
 	client := newTestClientWithRABA(t)
 	ctx := t.Context()
 	var routeID string
 	require.NoError(t, client.DB.QueryRowContext(ctx, "SELECT id FROM routes LIMIT 1").Scan(&routeID))
 
-	recorder := &routeQueryRecorder{DBTX: client.DB}
+	recorder := &queryRecorder{DBTX: client.DB}
 	_, err := New(recorder).GetActiveRoutesWithinBounds(ctx, GetActiveRoutesWithinBoundsParams{
 		Lat: 40.7347, Lon: -74.1644,
 		MinLat: -90, MaxLat: 90, MinLon: -180, MaxLon: 180,
 		RouteIDs: []string{routeID}, MaxCount: 100,
 	})
 	require.NoError(t, err)
-	plan, err := client.DB.QueryContext(ctx, "EXPLAIN QUERY PLAN "+recorder.query, recorder.args...)
-	require.NoError(t, err)
-	defer func() { _ = plan.Close() }()
 
 	var stopTimeLookups []string
-	for plan.Next() {
-		var id, parent, unused int
-		var detail string
-		require.NoError(t, plan.Scan(&id, &parent, &unused, &detail))
+	for _, detail := range explainRecordedQuery(t, client, recorder) {
 		if strings.Contains(detail, "stop_times") {
 			stopTimeLookups = append(stopTimeLookups, detail)
 		}
 	}
-	require.NoError(t, plan.Err())
 	require.NotEmpty(t, stopTimeLookups)
 	for _, lookup := range stopTimeLookups {
 		assert.Contains(t, lookup, "(trip_id=?)", "filtered searches must look up stop times by matching trip, not nearby stop")