From 4b2c593873f09b1c6ba3e541ccff1afc150d716e Mon Sep 17 00:00:00 2001
From: Andrew Kane <andrew@ankane.org>
Date: Wed, 28 Feb 2024 15:03:50 -0800
Subject: [PATCH] Updated number of parallel workers for IVFFlat index builds

---
 src/ivfbuild.c | 23 ++++++++++++++++++++++-
 1 file changed, 22 insertions(+), 1 deletion(-)

diff --git a/src/ivfbuild.c b/src/ivfbuild.c
index c4f4d4b..44a4a3a 100644
--- a/src/ivfbuild.c
+++ b/src/ivfbuild.c
@@ -904,6 +904,27 @@ IvfflatBeginParallel(IvfflatBuildState * buildstate, bool isconcurrent, int requ
 	WaitForParallelWorkersToAttach(pcxt);
 }
 
+/*
+ * Compute parallel workers
+ */
+static int
+ComputeParallelWorkers(Relation heap, Relation index)
+{
+	int			parallel_workers;
+
+	/* Make sure it's safe to use parallel workers */
+	parallel_workers = plan_create_index_workers(RelationGetRelid(heap), RelationGetRelid(index));
+	if (parallel_workers == 0)
+		return 0;
+
+	/* Use parallel_workers storage parameter on table if set */
+	parallel_workers = RelationGetParallelWorkers(heap, -1);
+	if (parallel_workers != -1)
+		return Min(parallel_workers, max_parallel_maintenance_workers);
+
+	return max_parallel_maintenance_workers;
+}
+
 /*
  * Scan table for tuples to index
  */
@@ -923,7 +944,7 @@ AssignTuples(IvfflatBuildState * buildstate)
 
 	/* Calculate parallel workers */
 	if (buildstate->heap != NULL)
-		parallel_workers = plan_create_index_workers(RelationGetRelid(buildstate->heap), RelationGetRelid(buildstate->index));
+		parallel_workers = ComputeParallelWorkers(buildstate->heap, buildstate->index);
 
 	/* Attempt to launch parallel worker scan when required */
 	if (parallel_workers > 0)