From 9f859f184abd9c10aa4930ca23e2ebc624c1a25a Mon Sep 17 00:00:00 2001 From: Andre Heringer Date: Sat, 4 Jul 2026 21:30:38 -0300 Subject: [PATCH 1/3] Track index creation in sequential mode --- pgvectorscale/src/access_method/build.rs | 13 ++++++++++--- pgvectorscale/src/util/ports.rs | 2 ++ 2 files changed, 12 insertions(+), 3 deletions(-) diff --git a/pgvectorscale/src/access_method/build.rs b/pgvectorscale/src/access_method/build.rs index 0aea049c..1e59d42d 100644 --- a/pgvectorscale/src/access_method/build.rs +++ b/pgvectorscale/src/access_method/build.rs @@ -14,7 +14,7 @@ use crate::access_method::graph::Graph; use crate::access_method::options::TSVIndexOptions; use crate::access_method::pg_vector::PgVector; use crate::access_method::stats::{InsertStats, WriteStats}; -use crate::util::ports::acquire_index_lock; +use crate::util::ports::{PROGRESS_CREATE_IDX_TUPLES_DONE, PROGRESS_CREATE_IDX_TUPLES_TOTAL, acquire_index_lock}; use crate::access_method::DISKANN_DISTANCE_TYPE_PROC; use crate::util::page::PageType; @@ -441,7 +441,8 @@ pub extern "C-unwind" fn ambuild( ntuples } } else { - do_heap_scan( + pgstat_progress_update_param(PROGRESS_CREATE_IDX_TUPLES_TOTAL, heap_tuples as i64); + let ntuples = do_heap_scan( index_info, &heap_relation, &index_relation, @@ -449,7 +450,9 @@ pub extern "C-unwind" fn ambuild( write_stats, None, workers as usize, - ) + ); + pgstat_progress_update_param(PROGRESS_CREATE_IDX_TUPLES_DONE, ntuples as i64); + ntuples }; let mut result = unsafe { PgBox::::alloc0() }; @@ -1078,6 +1081,10 @@ fn build_callback_internal( check_for_interrupts!(); state.ntuples += 1; + // Only call update every few tuples + if state.ntuples % 1000 == 0 { + pgstat_progress_update_param(PROGRESS_CREATE_IDX_TUPLES_DONE, state.ntuples as i64); + } let index_pointer = storage.create_node( vector.vec().to_index_slice(), diff --git a/pgvectorscale/src/util/ports.rs b/pgvectorscale/src/util/ports.rs index 2531d471..3f059cb1 100644 --- a/pgvectorscale/src/util/ports.rs +++ b/pgvectorscale/src/util/ports.rs @@ -42,6 +42,8 @@ pub unsafe fn PageValidateSpecialPointer(page: pgrx::pg_sys::Page) { #[allow(non_upper_case_globals)] const SizeOfPageHeaderData: usize = offset_of!(pgrx::pg_sys::PageHeaderData, pd_linp); pub const PROGRESS_CREATE_IDX_SUBPHASE: c_int = 10; +pub const PROGRESS_CREATE_IDX_TUPLES_TOTAL: c_int = 11; +pub const PROGRESS_CREATE_IDX_TUPLES_DONE: c_int = 12; #[allow(non_snake_case)] pub unsafe fn PageGetContents(page: pgrx::pg_sys::Page) -> *mut std::os::raw::c_char { From 631111ca80f506c4197ab17dc94fc8294d4e4b36 Mon Sep 17 00:00:00 2001 From: Andre Heringer Date: Tue, 7 Jul 2026 00:05:30 -0300 Subject: [PATCH 2/3] Add parallel tuple progress report Parallel progress report to the main PID is done by pulling a pool of processed tuples done by each worker, the loop executes until all workers are oberved and a final barrier was kept for safety --- pgvectorscale/src/access_method/build.rs | 152 ++++++++++++++++++++--- 1 file changed, 138 insertions(+), 14 deletions(-) diff --git a/pgvectorscale/src/access_method/build.rs b/pgvectorscale/src/access_method/build.rs index 1e59d42d..440e3716 100644 --- a/pgvectorscale/src/access_method/build.rs +++ b/pgvectorscale/src/access_method/build.rs @@ -73,6 +73,8 @@ struct BuildStateParallel<'a> { local_stats: InsertStats, local_ntuples: usize, is_initializing_worker: bool, + worker_index: usize, + worker_ntuples: *mut AtomicUsize, } impl<'a> BuildState<'a> { @@ -96,6 +98,8 @@ impl<'a> BuildStateParallel<'a> { page_type: PageType, shared_state: &'a ParallelShared, is_initializing_worker: bool, + worker_index: usize, + worker_ntuples: *mut AtomicUsize, ) -> Self { let tape = unsafe { Tape::new(index_relation, page_type) }; @@ -107,6 +111,8 @@ impl<'a> BuildStateParallel<'a> { local_stats: InsertStats::default(), local_ntuples: 0, is_initializing_worker, + worker_index, + worker_ntuples, } } @@ -226,6 +232,8 @@ struct ParallelBuildState { start_nodes_initialized: AtomicBool, initializing_worker_done: AtomicBool, initialization_cv: ConditionVariable, + workers_completed: AtomicUsize, + next_worker_index: AtomicUsize, } /// Status data for parallel index builds, shared among all parallel workers. @@ -236,6 +244,16 @@ struct ParallelShared { build_state: ParallelBuildState, } +/// DSM layout wrapper: tail-allocates an `AtomicUsize` per worker +/// so the leader can sum worker-local tuple counts for progress reporting. +/// The `[AtomicUsize; 1]` field is declared but indexed past length 1. +#[repr(C, align(8))] +#[cfg_attr(not(feature = "build_parallel"), allow(dead_code))] +struct ParallelSharedLayout { + head: ParallelShared, + worker_ntuples: [AtomicUsize; 1], +} + /// Information about parallel build passed to heap scan. #[derive(Debug)] #[cfg_attr(not(feature = "build_parallel"), allow(dead_code))] @@ -243,6 +261,8 @@ struct ParallelBuildInfo { parallel_shared: *mut ParallelShared, is_initializing_worker: bool, tablescandesc: *mut pg_sys::ParallelTableScanDescData, + worker_index: usize, + worker_ntuples: *mut AtomicUsize, } fn get_meta_page( @@ -361,7 +381,9 @@ pub extern "C-unwind" fn ambuild( }; // Estimate things we put in shared memory - parallel::toc_estimate_single_chunk(pcxt, size_of::()); + let layout_size = size_of::() + + size_of::() * workers; + parallel::toc_estimate_single_chunk(pcxt, layout_size); let tablescandesc_size_estimate = pg_sys::table_parallelscan_estimate(heaprel, snapshot); parallel::toc_estimate_single_chunk(pcxt, tablescandesc_size_estimate); @@ -372,9 +394,16 @@ pub extern "C-unwind" fn ambuild( parallel::cleanup_parallel_context(pcxt, snapshot); None } else { - let parallel_shared = - pg_sys::shm_toc_allocate((*pcxt).toc, size_of::()) - .cast::(); + let shared_layout: *mut ParallelSharedLayout = pg_sys::shm_toc_allocate( + (*pcxt).toc, + layout_size, + ) + .cast::(); + let parallel_shared = &raw mut (*shared_layout).head; + let worker_ntuples_ptr = (*shared_layout).worker_ntuples.as_mut_ptr(); + for i in 0..workers { + worker_ntuples_ptr.add(i).write(AtomicUsize::new(0)); + } let shared_state = ParallelShared { params: ParallelSharedParams { heaprelid: heap_relation.rd_id, @@ -388,6 +417,8 @@ pub extern "C-unwind" fn ambuild( start_nodes_initialized: AtomicBool::new(false), initializing_worker_done: AtomicBool::new(false), initialization_cv: std::mem::zeroed(), // Will be initialized below + workers_completed: AtomicUsize::new(0), + next_worker_index: AtomicUsize::new(0), }, }; parallel_shared.write(shared_state); @@ -429,19 +460,70 @@ pub extern "C-unwind" fn ambuild( let ntuples = if let Some(ParallelData { pcxt, snapshot }) = parallel_data { unsafe { - pg_sys::WaitForParallelWorkersToFinish(pcxt); let parallel_shared: *mut ParallelShared = pg_sys::shm_toc_lookup((*pcxt).toc, parallel::SHM_TOC_SHARED_KEY, false) .cast::(); - let ntuples = (*parallel_shared) - .build_state - .ntuples - .load(Ordering::Relaxed); + let shared_layout: *mut ParallelSharedLayout = + (parallel_shared as *mut u8).cast::(); + let worker_ntuples_ptr = (*shared_layout).worker_ntuples.as_ptr(); + + pgstat_progress_update_param( + PROGRESS_CREATE_IDX_TUPLES_TOTAL, + heap_tuples as i64, + ); + + loop { + check_for_interrupts!(); + + let mut total: usize = 0; + for i in 0..workers { + total += (*worker_ntuples_ptr.add(i)).load(Ordering::Relaxed); + } + pgstat_progress_update_param( + PROGRESS_CREATE_IDX_TUPLES_DONE, + total as i64, + ); + + let completed = (*parallel_shared) + .build_state + .workers_completed + .load(Ordering::Acquire); + if completed >= workers { + break; + } + + pg_sys::WaitLatch( + pg_sys::MyLatch, + (pg_sys::WL_LATCH_SET | pg_sys::WL_TIMEOUT | pg_sys::WL_EXIT_ON_PM_DEATH) as i32, + 512, + pg_sys::PG_WAIT_EXTENSION, + ); + pg_sys::ResetLatch(pg_sys::MyLatch); + } + + pg_sys::WaitForParallelWorkersToFinish(pcxt); + + let mut final_ntuples: usize = 0; + for i in 0..workers { + final_ntuples += (*worker_ntuples_ptr.add(i)).load(Ordering::Acquire); + } + parallel::cleanup_parallel_context(pcxt, snapshot); - ntuples + + pgstat_progress_update_param( + PROGRESS_CREATE_IDX_TUPLES_DONE, + final_ntuples as i64, + ); + + final_ntuples } } else { - pgstat_progress_update_param(PROGRESS_CREATE_IDX_TUPLES_TOTAL, heap_tuples as i64); + unsafe { + pgstat_progress_update_param( + PROGRESS_CREATE_IDX_TUPLES_TOTAL, + heap_tuples as i64, + ); + } let ntuples = do_heap_scan( index_info, &heap_relation, @@ -451,7 +533,12 @@ pub extern "C-unwind" fn ambuild( None, workers as usize, ); - pgstat_progress_update_param(PROGRESS_CREATE_IDX_TUPLES_DONE, ntuples as i64); + unsafe { + pgstat_progress_update_param( + PROGRESS_CREATE_IDX_TUPLES_DONE, + ntuples as i64, + ); + } ntuples }; @@ -639,6 +726,16 @@ pub extern "C-unwind" fn _vectorscale_build_main( .cast::() }; + let shared_layout: *mut ParallelSharedLayout = + (parallel_shared as *mut u8).cast::(); + let worker_ntuples_ptr = unsafe { (*shared_layout).worker_ntuples.as_mut_ptr() }; + let worker_index = unsafe { + (*parallel_shared) + .build_state + .next_worker_index + .fetch_add(1, Ordering::Relaxed) + }; + let params = unsafe { // SAFETY: these parameters never change, so no data races (*parallel_shared).params @@ -708,10 +805,19 @@ pub extern "C-unwind" fn _vectorscale_build_main( parallel_shared, is_initializing_worker: should_initialize, tablescandesc, + worker_index, + worker_ntuples: worker_ntuples_ptr, }), params.worker_count, ); + unsafe { + (*parallel_shared) + .build_state + .workers_completed + .fetch_add(1, Ordering::Release); + } + unsafe { pg_sys::index_close(indexrel, index_lockmode); pg_sys::table_close(heaprel, heap_lockmode); @@ -762,6 +868,8 @@ fn do_heap_scan( page_type, shared_state, parallel_info.is_initializing_worker, + parallel_info.worker_index, + parallel_info.worker_ntuples, ); let mut state = StorageBuildStateParallel::Plain(&mut plain, &mut bs); @@ -797,6 +905,8 @@ fn do_heap_scan( page_type, shared_state, parallel_info.is_initializing_worker, + parallel_info.worker_index, + parallel_info.worker_ntuples, ); let mut state = StorageBuildStateParallel::SbqSpeedup(&mut bq, &mut bs); @@ -1082,8 +1192,13 @@ fn build_callback_internal( state.ntuples += 1; // Only call update every few tuples - if state.ntuples % 1000 == 0 { - pgstat_progress_update_param(PROGRESS_CREATE_IDX_TUPLES_DONE, state.ntuples as i64); + if state.ntuples % 1024 == 0 { + unsafe { + pgstat_progress_update_param( + PROGRESS_CREATE_IDX_TUPLES_DONE, + state.ntuples as i64, + ); + } } let index_pointer = storage.create_node( @@ -1155,6 +1270,15 @@ fn build_callback_parallel_internal( state .graph .maybe_flush_neighbor_cache(storage, &mut state.local_stats); + + // Lossy snapshot of worker's local count for leader progress. + // piggy-backed on the existing flush throttle. + unsafe { + state + .worker_ntuples + .add(state.worker_index) + .write(AtomicUsize::new(state.local_ntuples)); + } } } From 08926387eee5c148a3c8b74093566ddbcdeee97d Mon Sep 17 00:00:00 2001 From: Andre Heringer Date: Tue, 7 Jul 2026 02:57:18 -0300 Subject: [PATCH 3/3] Add new test cases for index creation progress report --- tests/conftest.py | 8 +- tests/test_progress_reporting.py | 159 +++++++++++++++++++++++++++++++ 2 files changed, 165 insertions(+), 2 deletions(-) create mode 100644 tests/test_progress_reporting.py diff --git a/tests/conftest.py b/tests/conftest.py index f48a7379..0d0dd5b7 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -67,10 +67,12 @@ def clean_db(db_conn): DROP TABLE IF EXISTS test_l2 CASCADE; DROP TABLE IF EXISTS test_ip CASCADE; DROP TABLE IF EXISTS documents CASCADE; + DROP TABLE IF EXISTS test_progress_inc CASCADE; + DROP TABLE IF EXISTS test_progress_done CASCADE; """) - + yield - + # Clean up after the test with db_conn.cursor() as cur: cur.execute(""" @@ -81,6 +83,8 @@ def clean_db(db_conn): DROP TABLE IF EXISTS test_l2 CASCADE; DROP TABLE IF EXISTS test_ip CASCADE; DROP TABLE IF EXISTS documents CASCADE; + DROP TABLE IF EXISTS test_progress_inc CASCADE; + DROP TABLE IF EXISTS test_progress_done CASCADE; """) diff --git a/tests/test_progress_reporting.py b/tests/test_progress_reporting.py new file mode 100644 index 00000000..f1afc563 --- /dev/null +++ b/tests/test_progress_reporting.py @@ -0,0 +1,159 @@ +""" +Tests for CREATE INDEX progress reporting via pg_stat_progress_create_index. + +Verifies that the diskann access method populates tuples_total / tuples_done / +phase during index builds. Exercises the parallel build path (forced via GUCs) +so the leader's observer loop is hit, not just the sequential path. +""" + +import threading +import time +from concurrent.futures import ThreadPoolExecutor + +import psycopg2 + +VEC_DIM = 32 +ROW_COUNT = 30_000 +POLL_INTERVAL_S = 0.05 +MAX_POLL_S = 30.0 + + +def _vec_literal(n): + """SQL expression that produces a vector of n random floats.""" + return ( + "('[' || array_to_string(" + f"ARRAY(SELECT random()::float4 FROM generate_series(1, {n}))" + ", ',', '') || ']')::vector" + ) + + +def _force_parallel_gucs(cur): + """Lower the parallel-build threshold so the 30k-row test triggers it.""" + for stmt in [ + "SET diskann.min_vectors_for_parallel_build = 1000", + "SET diskann.force_parallel_workers = 4", + "SET max_parallel_maintenance_workers = 4", + "SET max_parallel_workers = 8", + "SET parallel_tuple_cost = 0", + "SET parallel_setup_cost = 0", + ]: + cur.execute(stmt) + + +def test_progress_increases(db_setup, clean_db): + """tuples_done in pg_stat_progress_create_index rises during a parallel build.""" + table = "test_progress_inc" + index = "idx_progress_inc" + + with db_conn_for(db_setup) as setup_conn: + with setup_conn.cursor() as cur: + cur.execute( + f"CREATE TABLE {table} (id INT, v VECTOR({VEC_DIM}))" + ) + cur.execute( + f"INSERT INTO {table} " + f"SELECT i, ({_vec_literal(VEC_DIM)}) " + f"FROM generate_series(1, {ROW_COUNT}) i" + ) + cur.execute(f"ANALYZE {table}") + _force_parallel_gucs(cur) + + samples = [] + stop = threading.Event() + + def poll(poll_conn): + while not stop.is_set(): + with poll_conn.cursor() as cur: + cur.execute( + "SELECT phase, tuples_total, tuples_done " + "FROM pg_stat_progress_create_index" + ) + rows = cur.fetchall() + if rows: + samples.append((time.monotonic(), *rows[0])) + time.sleep(POLL_INTERVAL_S) + + def build(build_conn): + with build_conn.cursor() as cur: + cur.execute( + f"CREATE INDEX {index} ON {table} " + f"USING diskann (v) WITH (num_neighbors=20)" + ) + + with db_conn_for(db_setup) as poll_conn, \ + db_conn_for(db_setup) as build_conn: + with ThreadPoolExecutor(max_workers=2) as ex: + poller = ex.submit(poll, poll_conn) + builder = ex.submit(build, build_conn) + try: + builder.result(timeout=MAX_POLL_S) + finally: + stop.set() + poller.result(timeout=5) + + assert len(samples) >= 3, f"too few progress samples: {len(samples)}" + done_values = [s[3] for s in samples] + for a, b in zip(done_values, done_values[1:]): + assert b >= a, f"tuples_done went backwards: {done_values}" + assert done_values[-1] >= 1, f"final tuples_done too low: {done_values[-1]}" + phases = [s[1] for s in samples if s[1]] + assert any( + "building graph" in p + or "training quantizer" in p + or "finalizing graph" in p + for p in phases + ), f"no build phase seen in {phases}" + + +def test_progress_reaches_completion(db_conn, clean_db): + """After CREATE INDEX finishes, no stale row remains in pg_stat_progress_create_index.""" + table = "test_progress_done" + index = "idx_progress_done" + + with db_conn.cursor() as cur: + cur.execute(f"CREATE TABLE {table} (id INT, v VECTOR(16))") + cur.execute( + f"INSERT INTO {table} " + f"SELECT i, ({_vec_literal(16)}) " + f"FROM generate_series(1, {ROW_COUNT}) i" + ) + cur.execute(f"ANALYZE {table}") + _force_parallel_gucs(cur) + + cur.execute( + f"CREATE INDEX {index} ON {table} " + f"USING diskann (v) WITH (num_neighbors=20)" + ) + + cur.execute( + "SELECT indisvalid FROM pg_index " + "WHERE indexrelid = %s::regclass", + (index,), + ) + row = cur.fetchone() + assert row is not None, f"index {index} not found" + assert row[0] is True, f"index {index} not valid" + + cur.execute( + "SELECT count(*) FROM pg_stat_progress_create_index" + ) + assert cur.fetchone()[0] == 0, ( + "stale row in pg_stat_progress_create_index" + ) + + +class _ConnCtx: + """Tiny context manager so `with db_conn_for(params) as conn:` closes cleanly.""" + def __init__(self, params): + self._conn = psycopg2.connect(**params) + self._conn.autocommit = True + + def __enter__(self): + return self._conn + + def __exit__(self, exc_type, exc, tb): + self._conn.close() + + +def db_conn_for(params): + return _ConnCtx(params)