Merge pull request #413 from PDOK/perf-prewarm-index

perf: prewarm search index

Author: rkettelerij

README.md

@@ -253,6 +253,14 @@ Set `LOG_SQL=true` environment variable to enable logging of all SQL queries to
 stdout for debug purposes. Only applies to OGC API Features with
 GeoPackage and/or PostgreSQL data source.
 
+#### PostgreSQL shared_buffers
+
+This only applies when Features Search (geocoding) is enabled: You can inspect the PostgreSQL buffercache for the search
+index by invoking `select * from gokoala_inspect_buffercache();` in the PostgreSQL database. The search index is
+automatically prewarmed after ETL is done. But in case you want to invoke it manually (one example would be after a
+restart of the database) you can invoke `select gokoala_prewarm_partitions();`. This function defaults to the current
+search index, but this can be overridden when desired.
+
 ## Develop
 
 Design principles:

internal/ogc/features_search/etl/etl.go

@@ -54,7 +54,7 @@ type Load interface {
 	PostLoad(collectionID string, index string, revision string) error
 
 	// Optimize once ETL is completed (optional)
-	Optimize() error
+	Optimize(index string) error
 
 	// Close connection to the target database
 	Close()
@@ -154,7 +154,7 @@ func ImportFile(collection config.Collection, searchIndex string, revision strin
 
 	if !skipOptimize {
 		log.Println("start optimizing")
-		if err = target.Optimize(); err != nil {
+		if err = target.Optimize(searchIndex); err != nil {
 			return err
 		}
 		log.Println("completed optimizing")

internal/ogc/features_search/etl/etl_test.go

@@ -70,6 +70,91 @@ func TestCreateSearchIndexIdempotent(t *testing.T) {
 	require.NoError(t, err)
 }
 
+func TestPreWarmPartitionFunction(t *testing.T) {
+	if testing.Short() {
+		t.Skip("Skipping integration test in short mode")
+	}
+	ctx := context.Background()
+
+	// given existing populated search index
+	dbPort, postgisContainer, err := setupPostgis(ctx, t)
+	if err != nil {
+		t.Error(err)
+	}
+	defer terminateContainer(ctx, t, postgisContainer)
+	dbConn := makeDbConnection(dbPort)
+
+	err = CreateSearchIndex(dbConn, "search_index", 28992, language.Dutch)
+	require.NoError(t, err)
+	err = insertTestData(ctx, dbConn)
+	require.NoError(t, err)
+
+	db, err := pgx.Connect(ctx, dbConn)
+	require.NoError(t, err)
+	defer db.Close(ctx)
+
+	// when/then
+	rows, err := db.Query(ctx, `SELECT * FROM gokoala_prewarm_partitions(idx_suffixes := array['ts_idx'])`)
+	require.NoError(t, err)
+	defer rows.Close()
+}
+
+func TestInspectBufferCacheFunction(t *testing.T) {
+	if testing.Short() {
+		t.Skip("Skipping integration test in short mode")
+	}
+	ctx := context.Background()
+
+	// given existing populated search index
+	dbPort, postgisContainer, err := setupPostgis(ctx, t)
+	if err != nil {
+		t.Error(err)
+	}
+	defer terminateContainer(ctx, t, postgisContainer)
+	dbConn := makeDbConnection(dbPort)
+
+	err = CreateSearchIndex(dbConn, "search_index", 28992, language.Dutch)
+	require.NoError(t, err)
+	err = insertTestData(ctx, dbConn)
+	require.NoError(t, err)
+
+	db, err := pgx.Connect(ctx, dbConn)
+	require.NoError(t, err)
+	defer db.Close(ctx)
+
+	// when
+	rows, err := db.Query(ctx, `SELECT * FROM gokoala_inspect_buffercache()`)
+	require.NoError(t, err)
+	defer rows.Close()
+
+	type result struct {
+		ObjectName       string
+		Kind             string
+		TotalSize        string
+		CachedSize       string
+		PercentageCached *float64
+	}
+
+	var got []result
+	for rows.Next() {
+		var r result
+		err := rows.Scan(
+			&r.ObjectName,
+			&r.Kind,
+			&r.TotalSize,
+			&r.CachedSize,
+			&r.PercentageCached,
+		)
+		require.NoError(t, err)
+		got = append(got, r)
+	}
+
+	// then
+	require.NoError(t, rows.Err())
+	require.NotEmpty(t, got)
+	require.NotEmpty(t, got[0].ObjectName)
+}
+
 func makeDbConnection(dbPort nat.Port) string {
 	return fmt.Sprintf("postgres://postgres:postgres@127.0.0.1:%d/%s?sslmode=disable", dbPort.Int(), "search_db")
 }

internal/ogc/features_search/etl/load/postgres.go

@@ -11,50 +11,15 @@ import (
 	t "github.com/PDOK/gokoala/internal/ogc/features_search/etl/transform"
 	"github.com/jackc/pgx/v5"
 	pgxgeom "github.com/twpayne/pgx-geom"
-	"golang.org/x/text/language"
 )
 
 const (
-	indexNameFullText = "ts_idx"
-	indexNameGeometry = "geometry_idx"
-	indexNameBbox     = "bbox_idx"
-	indexNamePreRank  = "pre_rank_idx"
-
 	alphaPartition = "_alpha"
 	betaPartition  = "_beta"
 
 	postgresDetachErr = "already pending detach in partitioned table"
 )
 
-var (
-	postgresExtensions = []string{"postgis", "unaccent"}
-
-	indexNames = []string{indexNameFullText, indexNameGeometry, indexNameBbox, indexNamePreRank}
-
-	//nolint:dupword
-	tableDefinition = `
-	create table if not exists %[1]s (
-		id 					serial,
-		feature_id 			text 					 not null,
-		external_fid        text                     null,
-		collection_id 		text					 not null,
-		collection_version 	int 					 not null,
-		display_name 		text					 not null,
-		suggest 			text					 not null,
-		geometry_type 		geometry_type			 not null,
-		bbox 				geometry(polygon, %[2]d) null,
-		geometry            geometry(point, %[2]d)   null,
-	    ts                  tsvector                 generated always as (to_tsvector('custom_dict', suggest)) stored
-	) %[3]s;`
-
-	metadataTableDefinition = `
-	create table if not exists %[1]s_metadata (
-		collection_id       text                      not null,
-		revision 			uuid                      not null,
-		revision_date       timestamptz default now() not null
-	);`
-)
-
 type Postgres struct {
 	db *pgx.Conn
 
@@ -180,10 +145,22 @@ func (p *Postgres) PostLoad(collectionID string, index string, revision string)
 	return nil
 }
 
-func (p *Postgres) Optimize() error {
-	_, err := p.db.Exec(context.Background(), `vacuum analyze;`)
+func (p *Postgres) Optimize(index string) error {
+	log.Println("perform targeted VACUUM + ANALYZE on loaded partition")
+	_, err := p.db.Exec(context.Background(), fmt.Sprintf(`vacuum analyze %s;`, p.partitionToLoad))
+	if err != nil {
+		return fmt.Errorf("failed optimizing: error performing vacuum analyze on loaded partition: %w", err)
+	}
+	_, err = p.db.Exec(context.Background(), fmt.Sprintf(`analyze %s;`, index))
+	if err != nil {
+		return fmt.Errorf("failed optimizing: error performing analyze on search index: %w", err)
+	}
+
+	// Execute pg_prewarm on all partitions and important indexes (forcing these into Postgres shared_buffers memory)
+	log.Println("prewarming partitions")
+	_, err = p.db.Exec(context.Background(), `select gokoala_prewarm_partitions()`)
 	if err != nil {
-		return fmt.Errorf("failed optimizing: error performing vacuum analyze: %w", err)
+		return fmt.Errorf("failed optimizing: prewarm function failed: %w", err)
 	}
 	return nil
 }
@@ -209,125 +186,6 @@ func (p *Postgres) GetRevision(collectionID string, index string) (string, error
 	return currentRevision, nil
 }
 
-// Init initialize search index. Should be idempotent!
-//
-// Since not all DDL in Postgres support the "if not exists" syntax we use a bit
-// of pl/pgsql to make it idempotent.
-func (p *Postgres) Init(index string, srid int, lang language.Tag) error {
-	geometryType := `
-		do $$ begin
-		    create type geometry_type as enum ('POINT', 'MULTIPOINT', 'LINESTRING', 'MULTILINESTRING', 'POLYGON', 'MULTIPOLYGON');
-		exception
-		    when duplicate_object then null;
-		end $$;`
-	_, err := p.db.Exec(context.Background(), geometryType)
-	if err != nil {
-		return fmt.Errorf("error creating geometry type: %w", err)
-	}
-
-	textSearchConfig := `
-		do $$ begin
-		    create text search configuration custom_dict (copy = simple);
-		exception
-		    when unique_violation then null;
-		end $$;`
-	_, err = p.db.Exec(context.Background(), textSearchConfig)
-	if err != nil {
-		return fmt.Errorf("error creating text search configuration: %w", err)
-	}
-
-	// This adds the 'unaccent' extension to allow searching with/without diacritics. Must happen in separate transaction.
-	alterTextSearchConfig := `
-		do $$ begin
-			alter text search configuration custom_dict
-			  alter mapping for hword, hword_part, word
-			  with unaccent, simple;
-		exception
-		    when unique_violation then null;
-		end $$;`
-	_, err = p.db.Exec(context.Background(), alterTextSearchConfig)
-	if err != nil {
-		return fmt.Errorf("error altering text search configuration: %w", err)
-	}
-
-	// create search index table
-	_, err = p.db.Exec(context.Background(), fmt.Sprintf(tableDefinition, index, srid, "partition by list(collection_id)"))
-	if err != nil {
-		return fmt.Errorf("error creating search index table: %w", err)
-	}
-
-	// create primary key when it doesn't exist yet
-	primaryKey := fmt.Sprintf(`
-		do $$
-		begin
-		    if not exists (
-		        select 1
-		        from   pg_constraint
-		        where  conrelid = '%[1]s'::regclass
-		        and    contype = 'p'
-		    )
-		    then
-		        alter table %[1]s
-		        add constraint %[1]s_pkey primary key (id, collection_id, collection_version);
-		    end if;
-		end;
-		$$;`, index)
-	_, err = p.db.Exec(context.Background(), primaryKey)
-	if err != nil {
-		return fmt.Errorf("error creating primary key: %w", err)
-	}
-
-	// create search index metadata table
-	_, err = p.db.Exec(context.Background(), fmt.Sprintf(metadataTableDefinition, index))
-	if err != nil {
-		return fmt.Errorf("error creating search index metadata table: %w", err)
-	}
-
-	// create metadata primary key when it doesn't exist yet
-	metadataPrimaryKey := fmt.Sprintf(`
-		do $$
-		begin
-		    if not exists (
-		        select 1
-		        from   pg_constraint
-		        where  conrelid = '%[1]s_metadata'::regclass
-		        and    contype = 'p'
-		    )
-		    then
-		        alter table %[1]s_metadata
-		        add constraint %[1]s_metadata_pkey primary key (collection_id);
-		    end if;
-		end;
-		$$;`, index)
-	_, err = p.db.Exec(context.Background(), metadataPrimaryKey)
-	if err != nil {
-		return fmt.Errorf("error creating metadata primary key: %w", err)
-	}
-
-	// create custom collation to correctly handle "numbers in strings" when sorting results
-	// see https://www.postgresql.org/docs/12/collation.html#id-1.6.10.4.5.7.5
-	collation := fmt.Sprintf(`create collation if not exists custom_numeric (provider = icu, locale = '%s-u-kn-true');`, lang.String())
-	_, err = p.db.Exec(context.Background(), collation)
-	if err != nil {
-		return fmt.Errorf("error creating numeric collation: %w", err)
-	}
-
-	if err = p.createIndexes(index, false); err != nil {
-		return err
-	}
-	return err
-}
-
-func createExtensions(ctx context.Context, db *pgx.Conn) error {
-	for _, ext := range postgresExtensions {
-		_, err := db.Exec(ctx, `create extension if not exists `+ext+`;`)
-		if err != nil {
-			return fmt.Errorf("error creating %s extension: %w", ext, err)
-		}
-	}
-	return nil
-}
-
 func (p *Postgres) isPartition(collectionID string, index string) (bool, error) {
 	result := false
 	err := p.db.QueryRow(context.Background(), `select exists (
@@ -341,66 +199,3 @@ func (p *Postgres) isPartition(collectionID string, index string) (bool, error)
 	    ) as is_partition_of_search_index;`, index, collectionID).Scan(&result)
 	return result, err
 }
-
-func (p *Postgres) createIndexes(table string, usePrefix bool) error {
-	// GIN indexes are best for text search
-	indexName := indexNameFullText
-	if usePrefix {
-		indexName = fmt.Sprintf("%s_%s", table, indexNameFullText)
-	}
-	_, err := p.db.Exec(context.Background(), fmt.Sprintf(`create index if not exists %[2]s on only %[1]s using gin(ts);`, table, indexName))
-	if err != nil {
-		return fmt.Errorf("error creating GIN index: %w", err)
-	}
-
-	// GIST indexes for geometry column to support search within a bounding box
-	indexName = indexNameGeometry
-	if usePrefix {
-		indexName = fmt.Sprintf("%s_%s", table, indexNameGeometry)
-	}
-	_, err = p.db.Exec(context.Background(), fmt.Sprintf(`create index if not exists %[2]s on only %[1]s using gist(geometry);`, table, indexName))
-	if err != nil {
-		return fmt.Errorf("error creating geometry GIST index: %w", err)
-	}
-
-	// GIST indexes for bbox column to support search within a bounding box
-	indexName = indexNameBbox
-	if usePrefix {
-		indexName = fmt.Sprintf("%s_%s", table, indexNameBbox)
-	}
-	_, err = p.db.Exec(context.Background(), fmt.Sprintf(`create index if not exists %[2]s on only %[1]s using gist(bbox);`, table, indexName))
-	if err != nil {
-		return fmt.Errorf("error creating bbox GIST index: %w", err)
-	}
-
-	// index used to pre-rank results when generic search terms are used
-	indexName = indexNamePreRank
-	if usePrefix {
-		indexName = fmt.Sprintf("%s_%s", table, indexNamePreRank)
-	}
-	preRankIndex := fmt.Sprintf(`create index if not exists %[2]s on only %[1]s
-		(array_length(string_to_array(suggest, ' '), 1) asc, display_name collate "custom_numeric" asc);`, table, indexName)
-	_, err = p.db.Exec(context.Background(), preRankIndex)
-	if err != nil {
-		return fmt.Errorf("error creating pre-rank index: %w", err)
-	}
-	return nil
-}
-
-// CHECK constraint is to avoid ACCESS EXCLUSIVE lock on partition as mentioned on
-// https://www.postgresql.org/docs/current/ddl-partitioning.html#DDL-PARTITIONING-DECLARATIVE-MAINTENANCE
-func (p *Postgres) createCheck(collectionID string) error {
-	dropCheck := fmt.Sprintf(`alter table %[1]s drop constraint if exists %[1]s_col_chk;`, p.partitionToLoad)
-	_, err := p.db.Exec(context.Background(), dropCheck)
-	if err != nil {
-		return fmt.Errorf("error dropping CHECK constraint: %w", err)
-	}
-
-	addCheck := fmt.Sprintf(`alter table if exists %[1]s add constraint %[1]s_col_chk check (collection_id = '%[2]s');`,
-		p.partitionToLoad, collectionID)
-	_, err = p.db.Exec(context.Background(), addCheck)
-	if err != nil {
-		return fmt.Errorf("error creating CHECK constraint: %w", err)
-	}
-	return nil
-}