speed! queries much faster now

tebben · tebben · commit e6f9ef7da491 · 2025-02-06T12:37:21.000+01:00
diff --git a/config/geocodeur.conf b/config/geocodeur.conf
@@ -18,7 +18,7 @@
         "timeoutSeconds": 30
     },
     "api": {
-        "similarityThreshold": 0.6,
+        "similarityThreshold": 0.8,
     },
     "database": {
         "name": "geocodeur",
diff --git a/src/database/setup.go b/src/database/setup.go
@@ -84,6 +84,12 @@ func CreateDB(config settings.Config) {
 		log.Fatalf("Failed to create index: %v", err)
 	}
 
+	log.Info("Creating search rank index")
+	err = createIndexRank(pool)
+	if err != nil {
+		log.Fatalf("Failed to create index: %v", err)
+	}
+
 	log.Info("Creating search trgm index")
 	err = createIndexTrgm(pool)
 	if err != nil {
@@ -171,20 +177,20 @@ func processParquet(pool *pgxpool.Pool, path string) {
 
 func processAliases(tx pgx.Tx, rec Record, id uint64) {
 	// Add name as alias
-	addAlias(tx, rec.ID, rec.Name, id)
+	addAlias(tx, rec, rec.Name, id)
 
 	// Add aliases for name aliases
 	for name, alias := range aliases {
 		if rec.Name == name {
-			addAlias(tx, rec.ID, alias, id)
+			addAlias(tx, rec, alias, id)
 		}
 	}
 
 	// Add embedding for truncated names
 	for _, truncation := range truncations {
 		if strings.Contains(rec.Name, truncation) {
 			alias := strings.Trim(strings.Replace(rec.Name, truncation, "", 1), " ")
-			addAlias(tx, rec.ID, alias, id)
+			addAlias(tx, rec, alias, id)
 		}
 	}
 
@@ -197,13 +203,13 @@ func processAliases(tx pgx.Tx, rec Record, id uint64) {
 			}
 
 			alias := rec.Name + " " + relation
-			addAlias(tx, rec.ID, alias, id)
+			addAlias(tx, rec, alias, id)
 
 			// Add entry for relation aliases
 			for name, alias := range aliases {
 				if relation == name {
 					aliasEmbedding := rec.Name + " " + alias
-					addAlias(tx, rec.ID, aliasEmbedding, id)
+					addAlias(tx, rec, aliasEmbedding, id)
 				}
 			}
 		}
@@ -217,10 +223,15 @@ func addOvertureFeature(tx pgx.Tx, rec Record, recordId uint64) error {
 	return err
 }
 
-func addAlias(tx pgx.Tx, id, alias string, recordId uint64) error {
+func addAlias(tx pgx.Tx, rec Record, alias string, recordId uint64) error {
 	alias = strings.ToLower(alias)
-	query := fmt.Sprintf(`INSERT INTO %s (feature_id, alias) VALUES ($1, $2)`, TABLE_SEARCH)
-	_, err := tx.Exec(context.Background(), query, recordId, alias)
+	classRank := getClassRank(rec.Class)
+	subclassRank := getSubclassScore(rec.Subclass)
+	wordCount := len(strings.Split(alias, " "))
+	charCount := len(alias)
+
+	query := fmt.Sprintf(`INSERT INTO %s (feature_id, alias, class_rank, subclass_rank, word_count, char_count) VALUES ($1, $2, $3, $4, $5, $6)`, TABLE_SEARCH)
+	_, err := tx.Exec(context.Background(), query, recordId, alias, classRank, subclassRank, wordCount, charCount)
 
 	return err
 }
@@ -240,6 +251,13 @@ func vacuum(pool *pgxpool.Pool) error {
 }
 
 func setupDatabase(pool *pgxpool.Pool, schema string) error {
+	if schema != "" {
+		_, err := pool.Exec(context.Background(), fmt.Sprintf("CREATE SCHEMA IF NOT EXISTS %s;", schema))
+		if err != nil {
+			return fmt.Errorf("failed to create schema: %v", err)
+		}
+	}
+
 	queryExtensions := `
 		CREATE EXTENSION IF NOT EXISTS postgis;
 		CREATE EXTENSION IF NOT EXISTS pg_trgm;
@@ -250,12 +268,7 @@ func setupDatabase(pool *pgxpool.Pool, schema string) error {
 		return fmt.Errorf("failed to create extensions: %v", err)
 	}
 
-	if schema == "" {
-		return nil
-	}
-
-	_, err = pool.Exec(context.Background(), fmt.Sprintf("CREATE SCHEMA IF NOT EXISTS %s;", schema))
-	return err
+	return nil
 }
 
 func createTableOverture(pool *pgxpool.Pool, tablespace string) error {
@@ -296,7 +309,11 @@ func createTableSearch(pool *pgxpool.Pool, tablespace string) error {
 
         CREATE TABLE %[1]s (
 			feature_id BIGINT,
-            alias TEXT
+            alias TEXT,
+			class_rank INT,
+			subclass_rank INT,
+			word_count INT,
+			char_count INT
         ) %[2]s;
     `, TABLE_SEARCH, tablespace)
 
@@ -313,6 +330,15 @@ func createForeignKey(pool *pgxpool.Pool) error {
 	return err
 }
 
+func createIndexRank(pool *pgxpool.Pool) error {
+	query := fmt.Sprintf(`
+		CREATE INDEX IF NOT EXISTS idx_%[1]s_class_subclass ON %[1]s USING btree (class_rank, subclass_rank);
+	`, TABLE_SEARCH)
+
+	_, err := pool.Exec(context.Background(), query)
+	return err
+}
+
 func createIndexGeom(pool *pgxpool.Pool) error {
 	query := fmt.Sprintf(`
 		CREATE INDEX IF NOT EXISTS idx_%[1]s_geom ON %[1]s USING GIST (geom);
@@ -341,3 +367,55 @@ func createFTSVectorColumn(pool *pgxpool.Pool) error {
 	_, err := pool.Exec(context.Background(), query)
 	return err
 }
+
+func getClassRank(class string) int {
+	switch class {
+	case "division":
+		return 1
+	case "water": // lot of division names with partly water name, maas, ijssel, etc, rank the same
+		return 1
+	case "road":
+		return 2
+	case "infra":
+		return 3
+	case "address":
+		return 4
+	case "zipcode":
+		return 5
+	case "poi":
+		return 6
+	default:
+		return 100
+	}
+}
+
+func getSubclassScore(subclass string) int {
+	switch subclass {
+	case "locality":
+		return 1
+	case "county":
+		return 2
+	case "neighboorhood":
+		return 3
+	case "microhood":
+		return 4
+	case "motorway":
+		return 1
+	case "trunk":
+		return 2
+	case "primary":
+		return 3
+	case "secondary":
+		return 4
+	case "tertiary":
+		return 5
+	case "unclassified":
+		return 6
+	case "residential":
+		return 6
+	case "living_street":
+		return 6
+	default:
+		return 100
+	}
+}
diff --git a/src/service/geocode.go b/src/service/geocode.go
@@ -148,83 +148,85 @@ func parseGeocodeResults(rows pgx.Rows) ([]GeocodeResult, error) {
 func createGeocodeQuery(options GeocodeOptions, input string) string {
 	classesIn := options.ClassesToSqlArray()
 
+	// workaround for now since we do not have class in the search table
+	classesIn = strings.Replace(classesIn, "'division'", "1", -1)
+	classesIn = strings.Replace(classesIn, "'water'", "1", -1)
+	classesIn = strings.Replace(classesIn, "'road'", "2", -1)
+	classesIn = strings.Replace(classesIn, "'infra'", "3", -1)
+	classesIn = strings.Replace(classesIn, "'address'", "4", -1)
+	classesIn = strings.Replace(classesIn, "'zipcode'", "5", -1)
+	classesIn = strings.Replace(classesIn, "'poi'", "6", -1)
+
 	// Conditional geometry column
 	geometryColumn := "'' AS geom" // Default to an empty string if geometry is not included
 	if options.IncludeGeometry {
-		geometryColumn = "ST_AsGeoJSON(a.geom) AS geom"
+		geometryColumn = "ST_AsGeoJSON(b.geom) AS geom"
 	}
 
 	return fmt.Sprintf(`
 		WITH fts AS (
-			SELECT feature_id, alias, similarity(alias, $1) AS sim, 'fts' as search
-			FROM %s AS a
-			JOIN %s AS b ON a.feature_id = b.id
-			WHERE a.vector_search @@ to_tsquery('simple',
-				replace($1, ' ', ':* & ') || ':*'
-			)
-			AND b.class IN %s
-			ORDER BY sim
+			SELECT
+				feature_id, alias, class_rank, subclass_rank, 'fts' as search
+			FROM
+				%[1]s
+			WHERE
+				ABS(word_count - array_length(string_to_array($1, ' '), 1)) < 3
+			AND
+				ABS(char_count - LENGTH($1)) < 30
+			AND
+				vector_search @@ to_tsquery('simple', replace($1, ' ', ':* & ') || ':*')
+			AND
+				class_rank IN %[3]s
+			ORDER BY
+				class_rank ASC,
+				subclass_rank ASC
+			LIMIT 100
 		),
 		trgm AS (
-			SELECT feature_id, alias, similarity(a.alias, $1) AS sim, 'trgm' as search
-			FROM %s AS a
-			JOIN %s AS b ON a.feature_id = b.id
-			WHERE a.alias %% $1
-			AND b.class IN %s
-			ORDER BY a.alias <-> $1
+			SELECT feature_id, alias, class_rank, subclass_rank, 'trgm' as search
+			FROM %[1]s
+			WHERE
+				ABS(word_count - array_length(string_to_array($1, ' '), 1)) < 3
+			AND
+				ABS(char_count - LENGTH($1)) < 30
+			AND
+				alias %% $1
+			AND
+				class_rank IN %[3]s
+			ORDER BY
+				class_rank ASC,
+				subclass_rank ASC
+			LIMIT 100
 		),
-		alias_results AS (
+		search_results AS (
 			SELECT *
 			FROM fts
 			UNION ALL
 			SELECT *
 			FROM trgm
 			WHERE NOT EXISTS (SELECT 1 FROM fts)
-		), ranked_aliases AS (
-			SELECT
-				a.id,
-				a.name,
-				a.class,
-				a.subclass,
-				array_to_string(a.divisions, ',') AS divisions,
-				b.alias,
-				b.sim,
-				b.search,
-				%s, -- Geometry column is dynamically included or excluded
-				CASE
-					WHEN a.class = 'division' THEN 1
-					WHEN a.class = 'water' THEN 2
-					WHEN a.class = 'road' THEN 3
-					WHEN a.class = 'infra' THEN 4
-					WHEN a.class = 'address' THEN 5
-					WHEN a.class = 'zipcode' THEN 6
-					WHEN a.class = 'poi' THEN 7
-					ELSE 100
-				END AS class_score,
-				CASE
-					WHEN a.subclass = 'locality' THEN 1
-					WHEN a.subclass = 'county' THEN 2
-					WHEN a.subclass = 'neighboorhood' THEN 3
-					WHEN a.subclass = 'microhood' THEN 4
-					-- roads up to living_street the rest gets a high score
-					WHEN a.subclass = 'motorway' THEN 1
-					WHEN a.subclass = 'trunk' THEN 2
-					WHEN a.subclass = 'primary' THEN 3
-					WHEN a.subclass = 'secondary' THEN 4
-					WHEN a.subclass = 'tertiary' THEN 5
-					WHEN a.subclass = 'unclassified' THEN 6
-					WHEN a.subclass = 'residential' THEN 7
-					WHEN a.subclass = 'living_street' THEN 8
-					ELSE 100
-				END AS subclass_score,
-				ROW_NUMBER() OVER (PARTITION BY a.id ORDER BY similarity(b.alias, $1) DESC) AS rnk
-			FROM %s a
-			JOIN alias_results b ON a.id = b.feature_id
+		),
+		similarity as (
+			select
+				feature_id,
+				alias,
+				class_rank,
+				subclass_rank,
+				similarity(alias, $1) AS sim,
+				search,
+				ROW_NUMBER() OVER (PARTITION BY feature_id ORDER BY similarity(alias, $1) DESC) AS rnk
+			from search_results
 		)
-		SELECT id, name, class, subclass, divisions, alias, search, sim, geom
-		FROM ranked_aliases
-		WHERE rnk = 1
-		ORDER BY sim desc, class_score asc, subclass_score asc
-		LIMIT %v;`,
-		database.TABLE_SEARCH, database.TABLE_OVERTURE, classesIn, database.TABLE_SEARCH, database.TABLE_OVERTURE, classesIn, geometryColumn, database.TABLE_OVERTURE, options.Limit)
+		SELECT
+			b.id, b.name, b.class, b.subclass, b.divisions::varchar, a.alias, a.search, a.sim, %[4]s
+		FROM similarity AS a
+		INNER JOIN
+			%[2]s AS b ON a.feature_id = b.id
+		WHERE a.rnk = 1
+		ORDER by
+			sim desc,
+			class_rank asc,
+			subclass_rank asc
+		LIMIT %[5]v;`,
+		database.TABLE_SEARCH, database.TABLE_OVERTURE, classesIn, geometryColumn, options.Limit)
 }
