diff --git a/src/main/java/org/duckdb/DuckDBWritableVector.java b/src/main/java/org/duckdb/DuckDBWritableVector.java index 0fb7725fb..dec513e5d 100644 --- a/src/main/java/org/duckdb/DuckDBWritableVector.java +++ b/src/main/java/org/duckdb/DuckDBWritableVector.java @@ -32,6 +32,7 @@ public final class DuckDBWritableVector { private final DuckDBVectorTypeInfo typeInfo; private final ByteBuffer data; private final ByteBuffer validity; + private final ReentrantLock validityLock = new ReentrantLock(); private static final class StringBatchState { byte[] payload; CharsetEncoder encoder; @@ -683,13 +684,20 @@ private void setRowValidity(long row, boolean valid) { int entryOffset = Math.toIntExact(Math.multiplyExact(row / Long.SIZE, (long) Long.BYTES)); long bitIndex = row % Long.SIZE; long mask = 1L << bitIndex; - long entry = validity.getLong(entryOffset); - if (valid) { - entry |= mask; - } else { - entry &= ~mask; + // Row validity is packed 64 rows per word, so producers writing disjoint rows still share the word; + // guard only this read-modify-write so concurrent mark/clear of different bits is not lost. + validityLock.lock(); + try { + long entry = validity.getLong(entryOffset); + if (valid) { + entry |= mask; + } else { + entry &= ~mask; + } + validity.putLong(entryOffset, entry); + } finally { + validityLock.unlock(); } - validity.putLong(entryOffset, entry); } private String typeMismatchMessage(DuckDBColumnType expected) { diff --git a/src/test/java/org/duckdb/TestScalarFunctions.java b/src/test/java/org/duckdb/TestScalarFunctions.java index 732b71079..90b805477 100644 --- a/src/test/java/org/duckdb/TestScalarFunctions.java +++ b/src/test/java/org/duckdb/TestScalarFunctions.java @@ -24,7 +24,16 @@ import java.time.LocalDate; import java.time.LocalDateTime; import java.time.OffsetDateTime; +import java.util.ArrayList; import java.util.List; +import java.util.concurrent.BrokenBarrierException; +import java.util.concurrent.CyclicBarrier; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; import java.util.function.BiFunction; import java.util.function.Function; import java.util.function.Supplier; @@ -2477,18 +2486,116 @@ public static void test_scalar_function_vectorized_parallel_streams() throws Exc output.setLong(row, left + right); })) .register(conn); - try (ResultSet rs = stmt.executeQuery("SELECT java_add(x, x + 1) FROM range(0, (1<<16) + 3) r(x)")) { - for (long i = 0; i < (1 << 16) + 3; i++) { - assertTrue(rs.next()); - long expected = i + i + 1; - long actual = rs.getLong(1); - assertEquals(actual, expected); + // Repeated because a lost update on the shared validity bitmap only leaves a subset of the rows + // NULL, so the outcome depends on how the parallel scheduler interleaves the producers. + for (int repetition = 0; repetition < 3; repetition++) { + try (ResultSet rs = stmt.executeQuery("SELECT java_add(x, x + 1) FROM range(0, (1<<16) + 3) r(x)")) { + for (long i = 0; i < (1 << 16) + 3; i++) { + assertTrue(rs.next()); + long expected = i + i + 1; + long actual = rs.getLong(1); + assertFalse(rs.wasNull()); + assertEquals(actual, expected); + } + assertFalse(rs.next()); } - assertFalse(rs.next()); } } } + private static final int VALIDITY_RMW_WORKERS = 64; + + private static void awaitValidityBarrier(CyclicBarrier barrier) { + try { + barrier.await(60, TimeUnit.SECONDS); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new IllegalStateException("Validity producer was interrupted", e); + } catch (BrokenBarrierException | TimeoutException e) { + throw new IllegalStateException("Validity producer barrier failed", e); + } + } + + /** + * Lost-update regression for the packed row validity bitmap. {@code VALIDITY_RMW_WORKERS} producers are + * released by a shared barrier and each one touches its own row of the same 64-bit validity word, so any + * non-atomic read-modify-write on that word drops another producer's bit: rows that were marked valid + * stay NULL and rows that were cleared stay valid. + */ + public static void test_scalar_function_vectorized_parallel_validity_rmw() throws Exception { + ExecutorService producers = Executors.newFixedThreadPool(VALIDITY_RMW_WORKERS); + try (Connection conn = DriverManager.getConnection(JDBC_URL); Statement stmt = conn.createStatement()) { + // DuckDB must not parallelize inside one invocation here: the 64 Java producers below already + // parallelize, and DuckDB workers sharing that pool could interleave the producers' barriers. + stmt.execute("SET threads=1"); + DuckDBFunctions.scalarFunction() + .withName("java_validity_rmw") + .withParameters(long.class, long.class) + .withReturnType(long.class) + .withVectorizedFunction((input, output) -> { + long rows = input.rowCount(); + int workers = (int) Math.min(VALIDITY_RMW_WORKERS, rows); + if (workers == 0) { + return; + } + CyclicBarrier barrier = new CyclicBarrier(workers); + long rounds = rows / workers; + List> futures = new ArrayList<>(workers); + for (int worker = 0; worker < workers; worker++) { + final long index = worker; + futures.add(producers.submit(() -> { + for (long round = 0; round < rounds; round++) { + long row = round * workers + index; + output.setLong(row, input.vector(0).getLong(row, 0) + input.vector(1).getLong(row, 0)); + // Everybody marks its own row of the shared validity word valid at the same time. + awaitValidityBarrier(barrier); + // Half of the producers then clear their own row again, so a lost clear would + // leave an even row valid. + if (index % 2 == 0) { + output.setNull(row); + } + awaitValidityBarrier(barrier); + } + return null; + })); + } + for (Future future : futures) { + try { + future.get(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new IllegalStateException("Validity producer was interrupted", e); + } catch (ExecutionException e) { + throw new IllegalStateException("Validity producer failed", e.getCause()); + } + } + // Tail rows of a chunk that is not a multiple of the worker count are written by the caller, + // which keeps every barrier round aligned with the producer count. + for (long row = rounds * workers; row < rows; row++) { + output.setLong(row, input.vector(0).getLong(row, 0) + input.vector(1).getLong(row, 0)); + } + }) + .register(conn); + for (int repetition = 0; repetition < 3; repetition++) { + try (ResultSet rs = + stmt.executeQuery("SELECT java_validity_rmw(x, x + 1) FROM range(0, (1<<16)) r(x)")) { + for (long i = 0; i < (1 << 16); i++) { + assertTrue(rs.next()); + if (i % 2 == 0) { + assertNullRow(rs); + } else { + assertEquals(rs.getLong(1), i + i + 1); + assertFalse(rs.wasNull()); + } + } + assertFalse(rs.next()); + } + } + } finally { + producers.shutdownNow(); + } + } + public static void test_scalar_function_dirty_vector() throws Exception { try (DuckDBConnection conn = DriverManager.getConnection(JDBC_URL).unwrap(DuckDBConnection.class); Statement stmt = conn.createStatement()) {