From 5fe86cbc9b05384fcb92b8e17a392aaa9e983acd Mon Sep 17 00:00:00 2001 From: Charles Huang <25107590+charleshuang119@users.noreply.github.com> Date: Tue, 21 Jul 2026 14:38:30 -0700 Subject: [PATCH 1/6] fix: coalesce executor credential refreshes across fragments --- .../internal/ExecutorNamespaceCache.java | 366 ++++++++++++++++++ .../spark/internal/LanceFragmentScanner.java | 32 +- .../internal/ExecutorNamespaceCacheTest.java | 239 ++++++++++++ .../internal/LanceFragmentScannerTest.java | 82 ++++ 4 files changed, 716 insertions(+), 3 deletions(-) create mode 100644 lance-spark-base_2.12/src/main/java/org/lance/spark/internal/ExecutorNamespaceCache.java create mode 100644 lance-spark-base_2.12/src/test/java/org/lance/spark/internal/ExecutorNamespaceCacheTest.java diff --git a/lance-spark-base_2.12/src/main/java/org/lance/spark/internal/ExecutorNamespaceCache.java b/lance-spark-base_2.12/src/main/java/org/lance/spark/internal/ExecutorNamespaceCache.java new file mode 100644 index 000000000..5c45e89b9 --- /dev/null +++ b/lance-spark-base_2.12/src/main/java/org/lance/spark/internal/ExecutorNamespaceCache.java @@ -0,0 +1,366 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.lance.spark.internal; + +import org.lance.namespace.LanceNamespace; +import org.lance.namespace.model.DescribeTableRequest; +import org.lance.namespace.model.DescribeTableResponse; +import org.lance.spark.LanceRuntime; + +import com.google.common.cache.Cache; +import com.google.common.cache.CacheBuilder; +import com.google.common.cache.RemovalNotification; +import com.google.common.util.concurrent.ExecutionError; +import com.google.common.util.concurrent.UncheckedExecutionException; +import org.apache.arrow.memory.BufferAllocator; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.Collections; +import java.util.HashMap; +import java.util.Map; +import java.util.Objects; +import java.util.concurrent.ConcurrentMap; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.FutureTask; +import java.util.concurrent.TimeUnit; +import java.util.function.LongSupplier; + +/** + * Executor-local cache for namespace clients and credential-vending table descriptions. + * + *

A Spark executor opens one Lance dataset per fragment. Opening through a namespace calls + * {@link LanceNamespace#describeTable(DescribeTableRequest)}, so reconstructing the namespace for + * every fragment turns a table with many fragments into a burst of catalog and identity-provider + * requests. This cache keeps one namespace client per Spark scan in each executor and coalesces + * identical table descriptions within that scan while preserving refresh before temporary + * credentials expire. Scan scoping prevents table locations and credentials from leaking into a + * later query that happens to reuse the same executor JVM. + */ +public final class ExecutorNamespaceCache { + private static final Logger LOG = LoggerFactory.getLogger(ExecutorNamespaceCache.class); + + static final String EXPIRES_AT_MILLIS = "expires_at_millis"; + static final long STATIC_RESPONSE_TTL_MILLIS = 5 * 60 * 1000L; + static final long UNKNOWN_CREDENTIAL_TTL_MILLIS = 60 * 1000L; + static final long MAX_EXPIRY_SAFETY_WINDOW_MILLIS = 60 * 1000L; + static final long MIN_EXPIRY_SAFETY_WINDOW_MILLIS = 1000L; + static final long MAX_DESCRIBED_TABLES_PER_NAMESPACE = 1000L; + static final long MAX_CACHED_SCANS = 1000L; + static final long CACHE_IDLE_EXPIRY_MILLIS = 60 * 60 * 1000L; + + private static final Cache NAMESPACES = + CacheBuilder.newBuilder() + .maximumSize(MAX_CACHED_SCANS) + .expireAfterAccess(CACHE_IDLE_EXPIRY_MILLIS, TimeUnit.MILLISECONDS) + .removalListener( + (RemovalNotification notification) -> { + CachedNamespace namespace = notification.getValue(); + if (namespace != null) { + namespace.evict(); + } + }) + .build(); + + private ExecutorNamespaceCache() {} + + /** Acquires a shared, credential-aware namespace client for one scan in this executor JVM. */ + public static Lease acquire( + String namespaceImpl, Map namespaceProperties, String scanId) { + NamespaceKey key = new NamespaceKey(namespaceImpl, namespaceProperties, scanId); + while (true) { + CachedNamespace cached; + try { + cached = + NAMESPACES.get( + key, + () -> + new CachedNamespace( + LanceRuntime.getOrCreateNamespace(namespaceImpl, key.properties))); + } catch (ExecutionException e) { + throw propagate(e.getCause(), "Failed to initialize executor namespace"); + } catch (UncheckedExecutionException e) { + throw propagate(e.getCause(), "Failed to initialize executor namespace"); + } catch (ExecutionError e) { + throw propagate(e.getCause(), "Failed to initialize executor namespace"); + } + Lease lease = cached.acquire(); + if (lease != null) { + return lease; + } + NAMESPACES.asMap().remove(key, cached); + } + } + + /** Clears cached namespaces. Primarily for tests. */ + static void clear() { + NAMESPACES.invalidateAll(); + NAMESPACES.cleanUp(); + } + + private static RuntimeException propagate(Throwable cause, String message) { + if (cause instanceof RuntimeException) { + return (RuntimeException) cause; + } + if (cause instanceof Error) { + throw (Error) cause; + } + return new RuntimeException(message, cause); + } + + private static final class NamespaceKey { + private final String impl; + private final Map properties; + private final String scanId; + + private NamespaceKey(String impl, Map properties, String scanId) { + this.impl = Objects.requireNonNull(impl, "namespaceImpl"); + this.properties = + properties == null + ? Collections.emptyMap() + : Collections.unmodifiableMap(new HashMap<>(properties)); + this.scanId = Objects.requireNonNull(scanId, "scanId"); + } + + @Override + public boolean equals(Object other) { + if (this == other) { + return true; + } + if (!(other instanceof NamespaceKey)) { + return false; + } + NamespaceKey that = (NamespaceKey) other; + return impl.equals(that.impl) + && properties.equals(that.properties) + && scanId.equals(that.scanId); + } + + @Override + public int hashCode() { + return Objects.hash(impl, properties, scanId); + } + } + + /** A reference-counted namespace lease held for the lifetime of one fragment scanner. */ + public static final class Lease implements AutoCloseable { + private CachedNamespace owner; + + private Lease(CachedNamespace owner) { + this.owner = owner; + } + + public LanceNamespace namespace() { + synchronized (this) { + if (owner == null) { + throw new IllegalStateException("Namespace lease is already closed"); + } + return owner.namespace; + } + } + + @Override + public void close() { + CachedNamespace toRelease; + synchronized (this) { + toRelease = owner; + owner = null; + } + if (toRelease != null) { + toRelease.release(); + } + } + } + + private static final class CachedNamespace { + private final LanceNamespace delegate; + private final CredentialCachingNamespace namespace; + private int leases; + private boolean evicted; + private boolean closed; + + private CachedNamespace(LanceNamespace delegate) { + this.delegate = Objects.requireNonNull(delegate, "delegate"); + this.namespace = new CredentialCachingNamespace(delegate, System::currentTimeMillis); + } + + private synchronized Lease acquire() { + if (evicted) { + return null; + } + leases++; + return new Lease(this); + } + + private synchronized void release() { + if (leases <= 0) { + throw new IllegalStateException("Namespace lease released too many times"); + } + leases--; + if (evicted && leases == 0) { + closeDelegate(); + } + } + + private synchronized void evict() { + evicted = true; + if (leases == 0) { + closeDelegate(); + } + } + + private void closeDelegate() { + if (closed) { + return; + } + closed = true; + if (delegate instanceof AutoCloseable) { + try { + ((AutoCloseable) delegate).close(); + } catch (Exception e) { + // Cache eviction is best effort and must not fail an unrelated Spark task. + LOG.warn("Failed to close an evicted executor namespace", e); + } + } + } + } + + static final class CredentialCachingNamespace implements LanceNamespace { + private final LanceNamespace delegate; + private final LongSupplier clock; + private final Cache> descriptionCache = + CacheBuilder.newBuilder() + .maximumSize(MAX_DESCRIBED_TABLES_PER_NAMESPACE) + .expireAfterAccess(CACHE_IDLE_EXPIRY_MILLIS, TimeUnit.MILLISECONDS) + .build(); + private final ConcurrentMap> descriptions = + descriptionCache.asMap(); + + CredentialCachingNamespace(LanceNamespace delegate, LongSupplier clock) { + this.delegate = Objects.requireNonNull(delegate, "delegate"); + this.clock = Objects.requireNonNull(clock, "clock"); + } + + @Override + public void initialize(Map properties, BufferAllocator allocator) { + throw new UnsupportedOperationException("Cached namespace is already initialized"); + } + + @Override + public String namespaceId() { + return delegate.namespaceId(); + } + + @Override + public DescribeTableResponse describeTable(DescribeTableRequest request) { + while (true) { + FutureTask candidate = + new FutureTask<>( + () -> { + long now = clock.getAsLong(); + DescribeTableResponse response = delegate.describeTable(request); + return new CachedDescription(response, refreshAtMillis(response, now)); + }); + FutureTask task = descriptions.putIfAbsent(request, candidate); + boolean loaded = task == null; + if (task == null) { + task = candidate; + task.run(); + } + + CachedDescription cached; + try { + cached = get(task); + } catch (RuntimeException | Error e) { + descriptions.remove(request, task); + throw e; + } + if (clock.getAsLong() < cached.refreshAtMillis) { + return cached.response; + } + descriptions.remove(request, task); + if (loaded) { + // Return the freshly fetched response even when it is already inside the safety window. + // The next caller will retry, but this caller must not spin forever if a server keeps + // vending expired or unusually short-lived credentials. + return cached.response; + } + } + } + + private static CachedDescription get(FutureTask task) { + try { + return task.get(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException("Interrupted while refreshing namespace credentials", e); + } catch (ExecutionException e) { + Throwable cause = e.getCause(); + if (cause instanceof RuntimeException) { + throw (RuntimeException) cause; + } + if (cause instanceof Error) { + throw (Error) cause; + } + throw new RuntimeException("Failed to refresh namespace credentials", cause); + } + } + + private static long refreshAtMillis(DescribeTableResponse response, long now) { + Map storageOptions = response.getStorageOptions(); + if (storageOptions == null || storageOptions.isEmpty()) { + return saturatedAdd(now, STATIC_RESPONSE_TTL_MILLIS); + } + + String rawExpiry = storageOptions.get(EXPIRES_AT_MILLIS); + if (rawExpiry == null) { + // The namespace spec requires temporary credentials to include expires_at_millis. Keep a + // short fallback TTL for older or non-conforming servers instead of caching indefinitely. + return saturatedAdd(now, UNKNOWN_CREDENTIAL_TTL_MILLIS); + } + + try { + long expiresAt = Long.parseLong(rawExpiry); + long remaining = expiresAt - now; + if (remaining <= 0) { + return now; + } + long safetyWindow = + Math.min( + MAX_EXPIRY_SAFETY_WINDOW_MILLIS, + Math.max(MIN_EXPIRY_SAFETY_WINDOW_MILLIS, remaining / 10)); + return Math.max(now, expiresAt - safetyWindow); + } catch (NumberFormatException ignored) { + return saturatedAdd(now, UNKNOWN_CREDENTIAL_TTL_MILLIS); + } + } + + private static long saturatedAdd(long value, long increment) { + if (value > Long.MAX_VALUE - increment) { + return Long.MAX_VALUE; + } + return value + increment; + } + } + + private static final class CachedDescription { + private final DescribeTableResponse response; + private final long refreshAtMillis; + + private CachedDescription(DescribeTableResponse response, long refreshAtMillis) { + this.response = response; + this.refreshAtMillis = refreshAtMillis; + } + } +} diff --git a/lance-spark-base_2.12/src/main/java/org/lance/spark/internal/LanceFragmentScanner.java b/lance-spark-base_2.12/src/main/java/org/lance/spark/internal/LanceFragmentScanner.java index a235a1171..83ca625c4 100644 --- a/lance-spark-base_2.12/src/main/java/org/lance/spark/internal/LanceFragmentScanner.java +++ b/lance-spark-base_2.12/src/main/java/org/lance/spark/internal/LanceFragmentScanner.java @@ -43,6 +43,7 @@ public class LanceFragmentScanner implements AutoCloseable { private final int fragmentId; private final boolean withFragmentId; private final LanceInputPartition inputPartition; + private final ExecutorNamespaceCache.Lease namespaceLease; private final long datasetOpenTimeNs; private final long scannerCreateTimeNs; @@ -61,6 +62,7 @@ private LanceFragmentScanner( int fragmentId, boolean withFragmentId, LanceInputPartition inputPartition, + ExecutorNamespaceCache.Lease namespaceLease, long datasetOpenTimeNs, long scannerCreateTimeNs, boolean withRowAddrForBlobs, @@ -70,6 +72,7 @@ private LanceFragmentScanner( this.fragmentId = fragmentId; this.withFragmentId = withFragmentId; this.inputPartition = inputPartition; + this.namespaceLease = namespaceLease; this.datasetOpenTimeNs = datasetOpenTimeNs; this.scannerCreateTimeNs = scannerCreateTimeNs; this.withRowAddrForBlobs = withRowAddrForBlobs; @@ -79,13 +82,17 @@ private LanceFragmentScanner( public static LanceFragmentScanner create(int fragmentId, LanceInputPartition inputPartition) { Dataset dataset = null; LanceScanner lanceScanner = null; + ExecutorNamespaceCache.Lease namespaceLease = null; try { LanceSparkReadOptions readOptions = inputPartition.getReadOptions(); if (inputPartition.getNamespaceImpl() != null && readOptions.isExecutorCredentialRefresh()) { if (LanceRuntime.useNamespaceOnWorkers(inputPartition.getNamespaceImpl())) { - readOptions.setNamespace( - LanceRuntime.getOrCreateNamespace( - inputPartition.getNamespaceImpl(), inputPartition.getNamespaceProperties())); + namespaceLease = + ExecutorNamespaceCache.acquire( + inputPartition.getNamespaceImpl(), + inputPartition.getNamespaceProperties(), + inputPartition.getScanId()); + readOptions.setNamespace(namespaceLease.namespace()); } else { readOptions.setNamespace(null); } @@ -156,6 +163,7 @@ public static LanceFragmentScanner create(int fragmentId, LanceInputPartition in fragmentId, withFragmentId, inputPartition, + namespaceLease, dsOpenTimeNs, scanCreateTimeNs, withRowAddrForBlobs, @@ -175,6 +183,13 @@ public static LanceFragmentScanner create(int fragmentId, LanceInputPartition in throwable.addSuppressed(closeError); } } + if (namespaceLease != null) { + try { + namespaceLease.close(); + } catch (Throwable closeError) { + throwable.addSuppressed(closeError); + } + } throw new RuntimeException(throwable); } } @@ -207,6 +222,17 @@ public void close() throws IOException { } } } + if (namespaceLease != null) { + try { + namespaceLease.close(); + } catch (Throwable t) { + if (primary != null) { + primary.addSuppressed(t); + } else { + primary = t; + } + } + } if (primary != null) { if (primary instanceof IOException) { throw (IOException) primary; diff --git a/lance-spark-base_2.12/src/test/java/org/lance/spark/internal/ExecutorNamespaceCacheTest.java b/lance-spark-base_2.12/src/test/java/org/lance/spark/internal/ExecutorNamespaceCacheTest.java new file mode 100644 index 000000000..dd709c0b6 --- /dev/null +++ b/lance-spark-base_2.12/src/test/java/org/lance/spark/internal/ExecutorNamespaceCacheTest.java @@ -0,0 +1,239 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.lance.spark.internal; + +import org.lance.namespace.LanceNamespace; +import org.lance.namespace.model.DescribeTableRequest; +import org.lance.namespace.model.DescribeTableResponse; + +import org.apache.arrow.memory.BufferAllocator; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Test; + +import java.util.Map; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicLong; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertThrows; + +public class ExecutorNamespaceCacheTest { + + @AfterEach + public void clearExecutorNamespaceCache() { + ExecutorNamespaceCache.clear(); + } + + @Test + public void defersClosingEvictedNamespaceUntilAllFragmentLeasesClose() { + CloseableNamespace.initializeCalls.set(0); + CloseableNamespace.closeCalls.set(0); + + ExecutorNamespaceCache.Lease first = + ExecutorNamespaceCache.acquire( + CloseableNamespace.class.getName(), Map.of("catalog", "test"), "scan-a"); + ExecutorNamespaceCache.Lease second = + ExecutorNamespaceCache.acquire( + CloseableNamespace.class.getName(), Map.of("catalog", "test"), "scan-a"); + + assertSame(first.namespace(), second.namespace()); + assertEquals(1, CloseableNamespace.initializeCalls.get()); + + ExecutorNamespaceCache.clear(); + assertEquals(0, CloseableNamespace.closeCalls.get()); + first.close(); + assertEquals(0, CloseableNamespace.closeCalls.get()); + second.close(); + assertEquals(1, CloseableNamespace.closeCalls.get()); + } + + @Test + public void coalescesConcurrentDescribeTableCalls() throws Exception { + AtomicLong clock = new AtomicLong(100_000L); + BlockingNamespace delegate = new BlockingNamespace(clock); + ExecutorNamespaceCache.CredentialCachingNamespace namespace = + new ExecutorNamespaceCache.CredentialCachingNamespace(delegate, clock::get); + DescribeTableRequest request = new DescribeTableRequest().addIdItem("table").version(1L); + + ExecutorService executor = Executors.newFixedThreadPool(8); + try { + @SuppressWarnings("unchecked") + Future[] futures = new Future[8]; + for (int i = 0; i < futures.length; i++) { + futures[i] = executor.submit(() -> namespace.describeTable(request)); + } + + delegate.entered.await(10, TimeUnit.SECONDS); + delegate.release.countDown(); + + DescribeTableResponse first = futures[0].get(10, TimeUnit.SECONDS); + for (Future future : futures) { + assertSame(first, future.get(10, TimeUnit.SECONDS)); + } + assertEquals(1, delegate.describeCalls.get()); + } finally { + executor.shutdownNow(); + } + } + + @Test + public void refreshesBeforeTemporaryCredentialsExpire() { + AtomicLong clock = new AtomicLong(100_000L); + RecordingNamespace delegate = new RecordingNamespace(clock); + ExecutorNamespaceCache.CredentialCachingNamespace namespace = + new ExecutorNamespaceCache.CredentialCachingNamespace(delegate, clock::get); + DescribeTableRequest request = new DescribeTableRequest().addIdItem("table").version(1L); + + DescribeTableResponse first = namespace.describeTable(request); + clock.set(189_999L); + assertSame(first, namespace.describeTable(request)); + + // Credentials expire at 200_000. With a 10% safety window, refresh starts at 190_000. + clock.set(190_000L); + DescribeTableResponse refreshed = namespace.describeTable(request); + + assertEquals(2, delegate.describeCalls.get()); + assertEquals("token-2", refreshed.getStorageOptions().get("token")); + } + + @Test + public void retriesAfterDescribeTableFailure() { + AtomicLong clock = new AtomicLong(100_000L); + RecordingNamespace delegate = new RecordingNamespace(clock); + delegate.failNext = true; + ExecutorNamespaceCache.CredentialCachingNamespace namespace = + new ExecutorNamespaceCache.CredentialCachingNamespace(delegate, clock::get); + DescribeTableRequest request = new DescribeTableRequest().addIdItem("table").version(1L); + + assertThrows(IllegalStateException.class, () -> namespace.describeTable(request)); + DescribeTableResponse response = namespace.describeTable(request); + + assertEquals(2, delegate.describeCalls.get()); + assertEquals("token-2", response.getStorageOptions().get("token")); + } + + @Test + public void doesNotSpinWhenNamespaceReturnsExpiredCredentials() { + AtomicInteger calls = new AtomicInteger(); + AtomicLong clock = new AtomicLong(100_000L); + LanceNamespace delegate = + new RecordingNamespace(clock) { + @Override + public DescribeTableResponse describeTable(DescribeTableRequest request) { + int call = calls.incrementAndGet(); + return new DescribeTableResponse() + .location("file:///tmp/table") + .storageOptions( + Map.of( + "token", + "expired-" + call, + ExecutorNamespaceCache.EXPIRES_AT_MILLIS, + Long.toString(clock.get() - 1))); + } + }; + ExecutorNamespaceCache.CredentialCachingNamespace namespace = + new ExecutorNamespaceCache.CredentialCachingNamespace(delegate, clock::get); + DescribeTableRequest request = new DescribeTableRequest().addIdItem("table").version(1L); + + assertEquals("expired-1", namespace.describeTable(request).getStorageOptions().get("token")); + assertEquals("expired-2", namespace.describeTable(request).getStorageOptions().get("token")); + assertEquals(2, calls.get()); + } + + private static class RecordingNamespace implements LanceNamespace { + final AtomicInteger describeCalls = new AtomicInteger(); + final AtomicLong clock; + volatile boolean failNext; + + RecordingNamespace(AtomicLong clock) { + this.clock = clock; + } + + @Override + public void initialize(Map properties, BufferAllocator allocator) {} + + @Override + public String namespaceId() { + return "recording"; + } + + @Override + public DescribeTableResponse describeTable(DescribeTableRequest request) { + int call = describeCalls.incrementAndGet(); + if (failNext) { + failNext = false; + throw new IllegalStateException("injected failure"); + } + return new DescribeTableResponse() + .location("file:///tmp/table") + .storageOptions( + Map.of( + "token", + "token-" + call, + ExecutorNamespaceCache.EXPIRES_AT_MILLIS, + Long.toString(clock.get() + 100_000L))); + } + } + + private static final class BlockingNamespace extends RecordingNamespace { + final CountDownLatch entered = new CountDownLatch(1); + final CountDownLatch release = new CountDownLatch(1); + + BlockingNamespace(AtomicLong clock) { + super(clock); + } + + @Override + public DescribeTableResponse describeTable(DescribeTableRequest request) { + entered.countDown(); + try { + if (!release.await(10, TimeUnit.SECONDS)) { + throw new IllegalStateException("timed out waiting to release describeTable"); + } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException(e); + } + return super.describeTable(request); + } + } + + public static final class CloseableNamespace implements LanceNamespace, AutoCloseable { + static final AtomicInteger initializeCalls = new AtomicInteger(); + static final AtomicInteger closeCalls = new AtomicInteger(); + + public CloseableNamespace() {} + + @Override + public void initialize(Map properties, BufferAllocator allocator) { + initializeCalls.incrementAndGet(); + } + + @Override + public String namespaceId() { + return "closeable"; + } + + @Override + public void close() { + closeCalls.incrementAndGet(); + } + } +} diff --git a/lance-spark-base_2.12/src/test/java/org/lance/spark/internal/LanceFragmentScannerTest.java b/lance-spark-base_2.12/src/test/java/org/lance/spark/internal/LanceFragmentScannerTest.java index e0230f994..6ec4d55a7 100644 --- a/lance-spark-base_2.12/src/test/java/org/lance/spark/internal/LanceFragmentScannerTest.java +++ b/lance-spark-base_2.12/src/test/java/org/lance/spark/internal/LanceFragmentScannerTest.java @@ -14,9 +14,13 @@ package org.lance.spark.internal; import org.lance.namespace.LanceNamespace; +import org.lance.namespace.model.DescribeTableRequest; +import org.lance.namespace.model.DescribeTableResponse; import org.lance.spark.LanceConstant; import org.lance.spark.LanceSparkReadOptions; +import org.lance.spark.TestUtils; import org.lance.spark.read.LanceInputPartition; +import org.lance.spark.read.LanceSplit; import org.lance.spark.utils.BlobUtils; import org.lance.spark.utils.Optional; @@ -25,6 +29,7 @@ import org.apache.spark.sql.types.MetadataBuilder; import org.apache.spark.sql.types.StructField; import org.apache.spark.sql.types.StructType; +import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.Test; import java.lang.reflect.InvocationTargetException; @@ -41,6 +46,11 @@ public class LanceFragmentScannerTest { + @AfterEach + public void clearExecutorNamespaceCache() { + ExecutorNamespaceCache.clear(); + } + private List callGetColumnNames(StructType schema) throws NoSuchMethodException, InvocationTargetException, IllegalAccessException { Method method = @@ -239,6 +249,7 @@ public void testGetColumnNamesWithFragmentId() throws Exception { @Test public void testCreateSkipsNamespaceRebuildWhenExecutorCredentialRefreshDisabled() { RecordingNamespace.INITIALIZE_CALLS.set(0); + RecordingNamespace.DESCRIBE_CALLS.set(0); LanceSparkReadOptions readOptions = LanceSparkReadOptions.builder() @@ -274,6 +285,69 @@ public void testCreateSkipsNamespaceRebuildWhenExecutorCredentialRefreshDisabled "executor_credential_refresh=false must not load or initialize the namespace impl"); } + @Test + public void testFragmentScansReuseExecutorNamespaceAndTableDescription() throws Exception { + ExecutorNamespaceCache.clear(); + RecordingNamespace.INITIALIZE_CALLS.set(0); + RecordingNamespace.DESCRIBE_CALLS.set(0); + RecordingNamespace.location = TestUtils.TestTable1Config.datasetUri; + + LanceInputPartition partition = createNamespacePartition("namespace-cache-test"); + + try (LanceFragmentScanner ignored = LanceFragmentScanner.create(0, partition)) { + // Opening the first fragment initializes the namespace and describes the table. + } + try (LanceFragmentScanner ignored = LanceFragmentScanner.create(1, partition)) { + // Opening another fragment must reuse both within the same executor JVM. + } + + assertEquals(1, RecordingNamespace.INITIALIZE_CALLS.get()); + assertEquals(1, RecordingNamespace.DESCRIBE_CALLS.get()); + } + + @Test + public void testDifferentScansDoNotReuseTableDescription() throws Exception { + RecordingNamespace.INITIALIZE_CALLS.set(0); + RecordingNamespace.DESCRIBE_CALLS.set(0); + RecordingNamespace.location = TestUtils.TestTable1Config.datasetUri; + + try (LanceFragmentScanner ignored = + LanceFragmentScanner.create(0, createNamespacePartition("scan-a"))) { + // First scan resolves its own table description. + } + try (LanceFragmentScanner ignored = + LanceFragmentScanner.create(1, createNamespacePartition("scan-b"))) { + // A later scan must not reuse the first scan's location or credentials. + } + + assertEquals(2, RecordingNamespace.INITIALIZE_CALLS.get()); + assertEquals(2, RecordingNamespace.DESCRIBE_CALLS.get()); + } + + private static LanceInputPartition createNamespacePartition(String scanId) { + LanceSparkReadOptions readOptions = + LanceSparkReadOptions.builder() + .datasetUri(RecordingNamespace.location) + .tableId(Collections.singletonList("test_dataset1")) + .version(6L) + .build(); + return new LanceInputPartition( + TestUtils.TestTable1Config.schema, + 0, + new LanceSplit(Arrays.asList(0, 1)), + readOptions, + Optional.empty(), + Optional.empty(), + Optional.empty(), + Optional.empty(), + Optional.empty(), + scanId, + Collections.emptyMap(), + RecordingNamespace.class.getName(), + Collections.emptyMap(), + null); + } + @Test public void getBlobColumnNamesIncludesBlobV2ReadColumns() throws Exception { Method method = @@ -306,6 +380,8 @@ public void getBlobColumnNamesIncludesBlobV2ReadColumns() throws Exception { */ public static class RecordingNamespace implements LanceNamespace { static final AtomicInteger INITIALIZE_CALLS = new AtomicInteger(); + static final AtomicInteger DESCRIBE_CALLS = new AtomicInteger(); + static volatile String location; public RecordingNamespace() {} @@ -318,5 +394,11 @@ public void initialize(Map properties, BufferAllocator allocator public String namespaceId() { return "recording"; } + + @Override + public DescribeTableResponse describeTable(DescribeTableRequest request) { + DESCRIBE_CALLS.incrementAndGet(); + return new DescribeTableResponse().location(location); + } } } From 7a6f7372be4c8b007c12c7b1f7dbe0c572bd9ef4 Mon Sep 17 00:00:00 2001 From: Charles Huang <25107590+charleshuang119@users.noreply.github.com> Date: Tue, 21 Jul 2026 15:13:31 -0700 Subject: [PATCH 2/6] fix: pin executor object stores across fragment tasks --- .../internal/ExecutorNamespaceCache.java | 109 +++++++++++++++++- .../spark/internal/LanceFragmentScanner.java | 1 + .../internal/LanceFragmentScannerTest.java | 40 ++++++- 3 files changed, 140 insertions(+), 10 deletions(-) diff --git a/lance-spark-base_2.12/src/main/java/org/lance/spark/internal/ExecutorNamespaceCache.java b/lance-spark-base_2.12/src/main/java/org/lance/spark/internal/ExecutorNamespaceCache.java index 5c45e89b9..c2fd58f04 100644 --- a/lance-spark-base_2.12/src/main/java/org/lance/spark/internal/ExecutorNamespaceCache.java +++ b/lance-spark-base_2.12/src/main/java/org/lance/spark/internal/ExecutorNamespaceCache.java @@ -13,10 +13,13 @@ */ package org.lance.spark.internal; +import org.lance.Dataset; import org.lance.namespace.LanceNamespace; import org.lance.namespace.model.DescribeTableRequest; import org.lance.namespace.model.DescribeTableResponse; import org.lance.spark.LanceRuntime; +import org.lance.spark.LanceSparkReadOptions; +import org.lance.spark.utils.Utils; import com.google.common.cache.Cache; import com.google.common.cache.CacheBuilder; @@ -38,15 +41,20 @@ import java.util.function.LongSupplier; /** - * Executor-local cache for namespace clients and credential-vending table descriptions. + * Executor-local cache for namespace clients, table descriptions, and dataset anchors. * *

A Spark executor opens one Lance dataset per fragment. Opening through a namespace calls * {@link LanceNamespace#describeTable(DescribeTableRequest)}, so reconstructing the namespace for * every fragment turns a table with many fragments into a burst of catalog and identity-provider - * requests. This cache keeps one namespace client per Spark scan in each executor and coalesces - * identical table descriptions within that scan while preserving refresh before temporary - * credentials expire. Scan scoping prevents table locations and credentials from leaking into a - * later query that happens to reuse the same executor JVM. + * requests. Closing every fragment dataset can also drop the last strong reference to Lance's + * session-cached object store, forcing the next fragment to rebuild the cloud client and acquire a + * new identity-provider token. This cache keeps one namespace client and one lightweight dataset + * anchor per Spark scan in each executor. The anchor is not used for fragment reads; it only keeps + * the shared object store and credential provider alive while fragment datasets continue to open + * and close independently. Identical table descriptions are coalesced within the scan while + * preserving refresh before temporary credentials expire. Scan scoping prevents table locations, + * credentials, and dataset metadata from leaking into a later query that happens to reuse the same + * executor JVM. */ public final class ExecutorNamespaceCache { private static final Logger LOG = LoggerFactory.getLogger(ExecutorNamespaceCache.class); @@ -153,7 +161,7 @@ public int hashCode() { } } - /** A reference-counted namespace lease held for the lifetime of one fragment scanner. */ + /** A reference-counted scan-resource lease held for the lifetime of one fragment scanner. */ public static final class Lease implements AutoCloseable { private CachedNamespace owner; @@ -170,6 +178,32 @@ public LanceNamespace namespace() { } } + /** + * Opens the scan's dataset anchor once and keeps it alive until this cache entry is evicted. + * + *

Every fragment still opens its own Dataset and scanner. The anchor deliberately never + * calls {@code getFragments()}, so its memory cost is independent of the fragment count. + */ + void pinDataset(LanceSparkReadOptions readOptions, Map initialStorageOptions) { + CachedNamespace current; + synchronized (this) { + if (owner == null) { + throw new IllegalStateException("Namespace lease is already closed"); + } + current = owner; + } + current.pinDataset(readOptions, initialStorageOptions); + } + + Dataset datasetAnchor() { + synchronized (this) { + if (owner == null) { + throw new IllegalStateException("Namespace lease is already closed"); + } + return owner.datasetAnchor(); + } + } + @Override public void close() { CachedNamespace toRelease; @@ -186,6 +220,7 @@ public void close() { private static final class CachedNamespace { private final LanceNamespace delegate; private final CredentialCachingNamespace namespace; + private FutureTask datasetAnchor; private int leases; private boolean evicted; private boolean closed; @@ -220,11 +255,73 @@ private synchronized void evict() { } } + private void pinDataset( + LanceSparkReadOptions readOptions, Map initialStorageOptions) { + FutureTask task; + boolean loaded = false; + synchronized (this) { + if (evicted) { + throw new IllegalStateException("Cannot pin a dataset on an evicted namespace"); + } + task = datasetAnchor; + if (task == null) { + task = + new FutureTask<>( + () -> + Utils.openDatasetBuilder(readOptions) + .initialStorageOptions(initialStorageOptions) + .build()); + datasetAnchor = task; + loaded = true; + } + } + + if (loaded) { + task.run(); + } + try { + getDataset(task); + } catch (RuntimeException | Error e) { + synchronized (this) { + if (datasetAnchor == task) { + datasetAnchor = null; + } + } + throw e; + } + } + + private synchronized Dataset datasetAnchor() { + if (datasetAnchor == null) { + throw new IllegalStateException("Dataset anchor has not been initialized"); + } + return getDataset(datasetAnchor); + } + + private static Dataset getDataset(FutureTask task) { + try { + return task.get(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException("Interrupted while opening executor dataset anchor", e); + } catch (ExecutionException e) { + throw propagate(e.getCause(), "Failed to open executor dataset anchor"); + } + } + private void closeDelegate() { if (closed) { return; } closed = true; + if (datasetAnchor != null) { + try { + getDataset(datasetAnchor).close(); + } catch (Exception e) { + // Cache eviction is best effort and must not fail an unrelated Spark task. + LOG.warn("Failed to close an evicted executor dataset anchor", e); + } + } if (delegate instanceof AutoCloseable) { try { ((AutoCloseable) delegate).close(); diff --git a/lance-spark-base_2.12/src/main/java/org/lance/spark/internal/LanceFragmentScanner.java b/lance-spark-base_2.12/src/main/java/org/lance/spark/internal/LanceFragmentScanner.java index 83ca625c4..946d29878 100644 --- a/lance-spark-base_2.12/src/main/java/org/lance/spark/internal/LanceFragmentScanner.java +++ b/lance-spark-base_2.12/src/main/java/org/lance/spark/internal/LanceFragmentScanner.java @@ -93,6 +93,7 @@ public static LanceFragmentScanner create(int fragmentId, LanceInputPartition in inputPartition.getNamespaceProperties(), inputPartition.getScanId()); readOptions.setNamespace(namespaceLease.namespace()); + namespaceLease.pinDataset(readOptions, inputPartition.getInitialStorageOptions()); } else { readOptions.setNamespace(null); } diff --git a/lance-spark-base_2.12/src/test/java/org/lance/spark/internal/LanceFragmentScannerTest.java b/lance-spark-base_2.12/src/test/java/org/lance/spark/internal/LanceFragmentScannerTest.java index 6ec4d55a7..c063425bf 100644 --- a/lance-spark-base_2.12/src/test/java/org/lance/spark/internal/LanceFragmentScannerTest.java +++ b/lance-spark-base_2.12/src/test/java/org/lance/spark/internal/LanceFragmentScannerTest.java @@ -13,6 +13,7 @@ */ package org.lance.spark.internal; +import org.lance.Dataset; import org.lance.namespace.LanceNamespace; import org.lance.namespace.model.DescribeTableRequest; import org.lance.namespace.model.DescribeTableResponse; @@ -41,8 +42,12 @@ import java.util.concurrent.atomic.AtomicInteger; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotSame; import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertSame; import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; public class LanceFragmentScannerTest { @@ -294,15 +299,33 @@ public void testFragmentScansReuseExecutorNamespaceAndTableDescription() throws LanceInputPartition partition = createNamespacePartition("namespace-cache-test"); - try (LanceFragmentScanner ignored = LanceFragmentScanner.create(0, partition)) { + Dataset firstAnchor; + try (LanceFragmentScanner ignored = LanceFragmentScanner.create(0, partition); + ExecutorNamespaceCache.Lease lease = + ExecutorNamespaceCache.acquire( + RecordingNamespace.class.getName(), + Collections.emptyMap(), + "namespace-cache-test")) { // Opening the first fragment initializes the namespace and describes the table. + firstAnchor = lease.datasetAnchor(); + assertFalse(firstAnchor.closed()); } - try (LanceFragmentScanner ignored = LanceFragmentScanner.create(1, partition)) { + try (LanceFragmentScanner ignored = LanceFragmentScanner.create(1, partition); + ExecutorNamespaceCache.Lease lease = + ExecutorNamespaceCache.acquire( + RecordingNamespace.class.getName(), + Collections.emptyMap(), + "namespace-cache-test")) { // Opening another fragment must reuse both within the same executor JVM. + assertSame(firstAnchor, lease.datasetAnchor()); } assertEquals(1, RecordingNamespace.INITIALIZE_CALLS.get()); assertEquals(1, RecordingNamespace.DESCRIBE_CALLS.get()); + assertFalse(firstAnchor.closed(), "the anchor must outlive individual fragment scanners"); + + ExecutorNamespaceCache.clear(); + assertTrue(firstAnchor.closed(), "cache eviction must close the dataset anchor"); } @Test @@ -311,13 +334,22 @@ public void testDifferentScansDoNotReuseTableDescription() throws Exception { RecordingNamespace.DESCRIBE_CALLS.set(0); RecordingNamespace.location = TestUtils.TestTable1Config.datasetUri; + Dataset firstAnchor; try (LanceFragmentScanner ignored = - LanceFragmentScanner.create(0, createNamespacePartition("scan-a"))) { + LanceFragmentScanner.create(0, createNamespacePartition("scan-a")); + ExecutorNamespaceCache.Lease lease = + ExecutorNamespaceCache.acquire( + RecordingNamespace.class.getName(), Collections.emptyMap(), "scan-a")) { // First scan resolves its own table description. + firstAnchor = lease.datasetAnchor(); } try (LanceFragmentScanner ignored = - LanceFragmentScanner.create(1, createNamespacePartition("scan-b"))) { + LanceFragmentScanner.create(1, createNamespacePartition("scan-b")); + ExecutorNamespaceCache.Lease lease = + ExecutorNamespaceCache.acquire( + RecordingNamespace.class.getName(), Collections.emptyMap(), "scan-b")) { // A later scan must not reuse the first scan's location or credentials. + assertNotSame(firstAnchor, lease.datasetAnchor()); } assertEquals(2, RecordingNamespace.INITIALIZE_CALLS.get()); From 4b4df9aa708d6faa2969d4900bf18f4a608161ae Mon Sep 17 00:00:00 2001 From: Charles Huang <25107590+charleshuang119@users.noreply.github.com> Date: Tue, 21 Jul 2026 16:28:37 -0700 Subject: [PATCH 3/6] fix: limit executor cache to namespace metadata --- .../internal/ExecutorNamespaceCache.java | 125 +++--------------- .../spark/internal/LanceFragmentScanner.java | 1 - .../internal/ExecutorNamespaceCacheTest.java | 39 ++++++ .../internal/LanceFragmentScannerTest.java | 48 ++----- 4 files changed, 71 insertions(+), 142 deletions(-) diff --git a/lance-spark-base_2.12/src/main/java/org/lance/spark/internal/ExecutorNamespaceCache.java b/lance-spark-base_2.12/src/main/java/org/lance/spark/internal/ExecutorNamespaceCache.java index c2fd58f04..d39d012d2 100644 --- a/lance-spark-base_2.12/src/main/java/org/lance/spark/internal/ExecutorNamespaceCache.java +++ b/lance-spark-base_2.12/src/main/java/org/lance/spark/internal/ExecutorNamespaceCache.java @@ -13,13 +13,14 @@ */ package org.lance.spark.internal; -import org.lance.Dataset; import org.lance.namespace.LanceNamespace; import org.lance.namespace.model.DescribeTableRequest; import org.lance.namespace.model.DescribeTableResponse; +import org.lance.namespace.model.DescribeTableVersionRequest; +import org.lance.namespace.model.DescribeTableVersionResponse; +import org.lance.namespace.model.ListTableVersionsRequest; +import org.lance.namespace.model.ListTableVersionsResponse; import org.lance.spark.LanceRuntime; -import org.lance.spark.LanceSparkReadOptions; -import org.lance.spark.utils.Utils; import com.google.common.cache.Cache; import com.google.common.cache.CacheBuilder; @@ -41,20 +42,15 @@ import java.util.function.LongSupplier; /** - * Executor-local cache for namespace clients, table descriptions, and dataset anchors. + * Executor-local cache for namespace clients and credential-vending table descriptions. * *

A Spark executor opens one Lance dataset per fragment. Opening through a namespace calls * {@link LanceNamespace#describeTable(DescribeTableRequest)}, so reconstructing the namespace for - * every fragment turns a table with many fragments into a burst of catalog and identity-provider - * requests. Closing every fragment dataset can also drop the last strong reference to Lance's - * session-cached object store, forcing the next fragment to rebuild the cloud client and acquire a - * new identity-provider token. This cache keeps one namespace client and one lightweight dataset - * anchor per Spark scan in each executor. The anchor is not used for fragment reads; it only keeps - * the shared object store and credential provider alive while fragment datasets continue to open - * and close independently. Identical table descriptions are coalesced within the scan while - * preserving refresh before temporary credentials expire. Scan scoping prevents table locations, - * credentials, and dataset metadata from leaking into a later query that happens to reuse the same - * executor JVM. + * every fragment turns a table with many fragments into a burst of catalog requests. This cache + * keeps one namespace client per Spark scan in each executor and coalesces identical table + * descriptions within that scan while preserving refresh before temporary credentials expire. Scan + * scoping prevents table locations and credentials from leaking into a later query that happens to + * reuse the same executor JVM. */ public final class ExecutorNamespaceCache { private static final Logger LOG = LoggerFactory.getLogger(ExecutorNamespaceCache.class); @@ -161,7 +157,7 @@ public int hashCode() { } } - /** A reference-counted scan-resource lease held for the lifetime of one fragment scanner. */ + /** A reference-counted namespace lease held for the lifetime of one fragment scanner. */ public static final class Lease implements AutoCloseable { private CachedNamespace owner; @@ -178,32 +174,6 @@ public LanceNamespace namespace() { } } - /** - * Opens the scan's dataset anchor once and keeps it alive until this cache entry is evicted. - * - *

Every fragment still opens its own Dataset and scanner. The anchor deliberately never - * calls {@code getFragments()}, so its memory cost is independent of the fragment count. - */ - void pinDataset(LanceSparkReadOptions readOptions, Map initialStorageOptions) { - CachedNamespace current; - synchronized (this) { - if (owner == null) { - throw new IllegalStateException("Namespace lease is already closed"); - } - current = owner; - } - current.pinDataset(readOptions, initialStorageOptions); - } - - Dataset datasetAnchor() { - synchronized (this) { - if (owner == null) { - throw new IllegalStateException("Namespace lease is already closed"); - } - return owner.datasetAnchor(); - } - } - @Override public void close() { CachedNamespace toRelease; @@ -220,7 +190,6 @@ public void close() { private static final class CachedNamespace { private final LanceNamespace delegate; private final CredentialCachingNamespace namespace; - private FutureTask datasetAnchor; private int leases; private boolean evicted; private boolean closed; @@ -255,73 +224,11 @@ private synchronized void evict() { } } - private void pinDataset( - LanceSparkReadOptions readOptions, Map initialStorageOptions) { - FutureTask task; - boolean loaded = false; - synchronized (this) { - if (evicted) { - throw new IllegalStateException("Cannot pin a dataset on an evicted namespace"); - } - task = datasetAnchor; - if (task == null) { - task = - new FutureTask<>( - () -> - Utils.openDatasetBuilder(readOptions) - .initialStorageOptions(initialStorageOptions) - .build()); - datasetAnchor = task; - loaded = true; - } - } - - if (loaded) { - task.run(); - } - try { - getDataset(task); - } catch (RuntimeException | Error e) { - synchronized (this) { - if (datasetAnchor == task) { - datasetAnchor = null; - } - } - throw e; - } - } - - private synchronized Dataset datasetAnchor() { - if (datasetAnchor == null) { - throw new IllegalStateException("Dataset anchor has not been initialized"); - } - return getDataset(datasetAnchor); - } - - private static Dataset getDataset(FutureTask task) { - try { - return task.get(); - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); - throw new RuntimeException("Interrupted while opening executor dataset anchor", e); - } catch (ExecutionException e) { - throw propagate(e.getCause(), "Failed to open executor dataset anchor"); - } - } - private void closeDelegate() { if (closed) { return; } closed = true; - if (datasetAnchor != null) { - try { - getDataset(datasetAnchor).close(); - } catch (Exception e) { - // Cache eviction is best effort and must not fail an unrelated Spark task. - LOG.warn("Failed to close an evicted executor dataset anchor", e); - } - } if (delegate instanceof AutoCloseable) { try { ((AutoCloseable) delegate).close(); @@ -396,6 +303,16 @@ public DescribeTableResponse describeTable(DescribeTableRequest request) { } } + @Override + public ListTableVersionsResponse listTableVersions(ListTableVersionsRequest request) { + return delegate.listTableVersions(request); + } + + @Override + public DescribeTableVersionResponse describeTableVersion(DescribeTableVersionRequest request) { + return delegate.describeTableVersion(request); + } + private static CachedDescription get(FutureTask task) { try { return task.get(); diff --git a/lance-spark-base_2.12/src/main/java/org/lance/spark/internal/LanceFragmentScanner.java b/lance-spark-base_2.12/src/main/java/org/lance/spark/internal/LanceFragmentScanner.java index 946d29878..83ca625c4 100644 --- a/lance-spark-base_2.12/src/main/java/org/lance/spark/internal/LanceFragmentScanner.java +++ b/lance-spark-base_2.12/src/main/java/org/lance/spark/internal/LanceFragmentScanner.java @@ -93,7 +93,6 @@ public static LanceFragmentScanner create(int fragmentId, LanceInputPartition in inputPartition.getNamespaceProperties(), inputPartition.getScanId()); readOptions.setNamespace(namespaceLease.namespace()); - namespaceLease.pinDataset(readOptions, inputPartition.getInitialStorageOptions()); } else { readOptions.setNamespace(null); } diff --git a/lance-spark-base_2.12/src/test/java/org/lance/spark/internal/ExecutorNamespaceCacheTest.java b/lance-spark-base_2.12/src/test/java/org/lance/spark/internal/ExecutorNamespaceCacheTest.java index dd709c0b6..96910fbfe 100644 --- a/lance-spark-base_2.12/src/test/java/org/lance/spark/internal/ExecutorNamespaceCacheTest.java +++ b/lance-spark-base_2.12/src/test/java/org/lance/spark/internal/ExecutorNamespaceCacheTest.java @@ -16,6 +16,10 @@ import org.lance.namespace.LanceNamespace; import org.lance.namespace.model.DescribeTableRequest; import org.lance.namespace.model.DescribeTableResponse; +import org.lance.namespace.model.DescribeTableVersionRequest; +import org.lance.namespace.model.DescribeTableVersionResponse; +import org.lance.namespace.model.ListTableVersionsRequest; +import org.lance.namespace.model.ListTableVersionsResponse; import org.apache.arrow.memory.BufferAllocator; import org.junit.jupiter.api.AfterEach; @@ -157,6 +161,41 @@ public DescribeTableResponse describeTable(DescribeTableRequest request) { assertEquals(2, calls.get()); } + @Test + public void delegatesManagedVersioningReads() { + ListTableVersionsResponse listResponse = new ListTableVersionsResponse(); + DescribeTableVersionResponse describeResponse = new DescribeTableVersionResponse(); + AtomicInteger listCalls = new AtomicInteger(); + AtomicInteger describeVersionCalls = new AtomicInteger(); + LanceNamespace delegate = + new RecordingNamespace(new AtomicLong()) { + @Override + public ListTableVersionsResponse listTableVersions(ListTableVersionsRequest request) { + listCalls.incrementAndGet(); + return listResponse; + } + + @Override + public DescribeTableVersionResponse describeTableVersion( + DescribeTableVersionRequest request) { + describeVersionCalls.incrementAndGet(); + return describeResponse; + } + }; + ExecutorNamespaceCache.CredentialCachingNamespace namespace = + new ExecutorNamespaceCache.CredentialCachingNamespace(delegate, () -> 0L); + + assertSame( + listResponse, + namespace.listTableVersions(new ListTableVersionsRequest().addIdItem("table"))); + assertSame( + describeResponse, + namespace.describeTableVersion( + new DescribeTableVersionRequest().addIdItem("table").version(1L))); + assertEquals(1, listCalls.get()); + assertEquals(1, describeVersionCalls.get()); + } + private static class RecordingNamespace implements LanceNamespace { final AtomicInteger describeCalls = new AtomicInteger(); final AtomicLong clock; diff --git a/lance-spark-base_2.12/src/test/java/org/lance/spark/internal/LanceFragmentScannerTest.java b/lance-spark-base_2.12/src/test/java/org/lance/spark/internal/LanceFragmentScannerTest.java index c063425bf..bbc70149e 100644 --- a/lance-spark-base_2.12/src/test/java/org/lance/spark/internal/LanceFragmentScannerTest.java +++ b/lance-spark-base_2.12/src/test/java/org/lance/spark/internal/LanceFragmentScannerTest.java @@ -13,7 +13,6 @@ */ package org.lance.spark.internal; -import org.lance.Dataset; import org.lance.namespace.LanceNamespace; import org.lance.namespace.model.DescribeTableRequest; import org.lance.namespace.model.DescribeTableResponse; @@ -42,12 +41,8 @@ import java.util.concurrent.atomic.AtomicInteger; import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.junit.jupiter.api.Assertions.assertFalse; -import static org.junit.jupiter.api.Assertions.assertNotSame; import static org.junit.jupiter.api.Assertions.assertNull; -import static org.junit.jupiter.api.Assertions.assertSame; import static org.junit.jupiter.api.Assertions.assertThrows; -import static org.junit.jupiter.api.Assertions.assertTrue; public class LanceFragmentScannerTest { @@ -299,33 +294,18 @@ public void testFragmentScansReuseExecutorNamespaceAndTableDescription() throws LanceInputPartition partition = createNamespacePartition("namespace-cache-test"); - Dataset firstAnchor; - try (LanceFragmentScanner ignored = LanceFragmentScanner.create(0, partition); - ExecutorNamespaceCache.Lease lease = - ExecutorNamespaceCache.acquire( - RecordingNamespace.class.getName(), - Collections.emptyMap(), - "namespace-cache-test")) { + try (LanceFragmentScanner ignored = LanceFragmentScanner.create(0, partition)) { // Opening the first fragment initializes the namespace and describes the table. - firstAnchor = lease.datasetAnchor(); - assertFalse(firstAnchor.closed()); } - try (LanceFragmentScanner ignored = LanceFragmentScanner.create(1, partition); - ExecutorNamespaceCache.Lease lease = - ExecutorNamespaceCache.acquire( - RecordingNamespace.class.getName(), - Collections.emptyMap(), - "namespace-cache-test")) { - // Opening another fragment must reuse both within the same executor JVM. - assertSame(firstAnchor, lease.datasetAnchor()); + assertEquals(1, RecordingNamespace.INITIALIZE_CALLS.get()); + assertEquals(1, RecordingNamespace.DESCRIBE_CALLS.get()); + + try (LanceFragmentScanner ignored = LanceFragmentScanner.create(1, partition)) { + // Opening another fragment must reuse the namespace metadata within the executor JVM. } assertEquals(1, RecordingNamespace.INITIALIZE_CALLS.get()); assertEquals(1, RecordingNamespace.DESCRIBE_CALLS.get()); - assertFalse(firstAnchor.closed(), "the anchor must outlive individual fragment scanners"); - - ExecutorNamespaceCache.clear(); - assertTrue(firstAnchor.closed(), "cache eviction must close the dataset anchor"); } @Test @@ -334,22 +314,16 @@ public void testDifferentScansDoNotReuseTableDescription() throws Exception { RecordingNamespace.DESCRIBE_CALLS.set(0); RecordingNamespace.location = TestUtils.TestTable1Config.datasetUri; - Dataset firstAnchor; try (LanceFragmentScanner ignored = - LanceFragmentScanner.create(0, createNamespacePartition("scan-a")); - ExecutorNamespaceCache.Lease lease = - ExecutorNamespaceCache.acquire( - RecordingNamespace.class.getName(), Collections.emptyMap(), "scan-a")) { + LanceFragmentScanner.create(0, createNamespacePartition("scan-a"))) { // First scan resolves its own table description. - firstAnchor = lease.datasetAnchor(); } + assertEquals(1, RecordingNamespace.INITIALIZE_CALLS.get()); + assertEquals(1, RecordingNamespace.DESCRIBE_CALLS.get()); + try (LanceFragmentScanner ignored = - LanceFragmentScanner.create(1, createNamespacePartition("scan-b")); - ExecutorNamespaceCache.Lease lease = - ExecutorNamespaceCache.acquire( - RecordingNamespace.class.getName(), Collections.emptyMap(), "scan-b")) { + LanceFragmentScanner.create(1, createNamespacePartition("scan-b"))) { // A later scan must not reuse the first scan's location or credentials. - assertNotSame(firstAnchor, lease.datasetAnchor()); } assertEquals(2, RecordingNamespace.INITIALIZE_CALLS.get()); From 40af4f87645ea943eee90e8bbb5e22827252a559 Mon Sep 17 00:00:00 2001 From: Charles Huang <25107590+charleshuang119@users.noreply.github.com> Date: Tue, 21 Jul 2026 16:33:37 -0700 Subject: [PATCH 4/6] fix: reuse namespace metadata for filtered counts --- .../read/LanceCountStarPartitionReader.java | 107 +++++++++++------- .../LanceCountStarPartitionReaderTest.java | 80 +++++++++++++ 2 files changed, 144 insertions(+), 43 deletions(-) diff --git a/lance-spark-base_2.12/src/main/java/org/lance/spark/read/LanceCountStarPartitionReader.java b/lance-spark-base_2.12/src/main/java/org/lance/spark/read/LanceCountStarPartitionReader.java index b28e6b370..fe45e1678 100644 --- a/lance-spark-base_2.12/src/main/java/org/lance/spark/read/LanceCountStarPartitionReader.java +++ b/lance-spark-base_2.12/src/main/java/org/lance/spark/read/LanceCountStarPartitionReader.java @@ -18,6 +18,7 @@ import org.lance.ipc.ScanOptions; import org.lance.spark.LanceRuntime; import org.lance.spark.LanceSparkReadOptions; +import org.lance.spark.internal.ExecutorNamespaceCache; import org.lance.spark.read.metric.LanceReadMetricsTracker; import org.lance.spark.utils.Utils; import org.lance.spark.vectorized.LanceArrowColumnVector; @@ -65,58 +66,78 @@ public boolean next() throws IOException { private long computeCount() { // This reader is only used when there are filters (metadata-based count uses LocalScan) LanceSparkReadOptions readOptions = inputPartition.getReadOptions(); + ExecutorNamespaceCache.Lease namespaceLease = null; long totalCount = 0; - long dsOpenStart = System.nanoTime(); - try (Dataset dataset = - Utils.openDatasetBuilder(readOptions) - .initialStorageOptions(inputPartition.getInitialStorageOptions()) - .build()) { - metricsTracker.addDatasetOpenTimeNs(System.nanoTime() - dsOpenStart); - - List fragmentIds = inputPartition.getLanceSplit().getFragments(); - if (fragmentIds.isEmpty()) { - return 0; + try { + if (inputPartition.getNamespaceImpl() != null && readOptions.isExecutorCredentialRefresh()) { + if (LanceRuntime.useNamespaceOnWorkers(inputPartition.getNamespaceImpl())) { + namespaceLease = + ExecutorNamespaceCache.acquire( + inputPartition.getNamespaceImpl(), + inputPartition.getNamespaceProperties(), + inputPartition.getScanId()); + readOptions.setNamespace(namespaceLease.namespace()); + } else { + readOptions.setNamespace(null); + } } - metricsTracker.addNumFragmentsScanned(fragmentIds.size()); - ScanOptions.Builder scanOptionsBuilder = new ScanOptions.Builder(); - scanOptionsBuilder.useScalarIndex(readOptions.isUseScalarIndex()); - if (inputPartition.getWhereCondition().isPresent()) { - scanOptionsBuilder.filter(inputPartition.getWhereCondition().get()); - } - scanOptionsBuilder.withRowId(true); - scanOptionsBuilder.columns(Lists.newArrayList()); - scanOptionsBuilder.fragmentIds(fragmentIds); - - // Collect scan stats - scanOptionsBuilder.collectStats(true); - - long scanCreateStart = System.nanoTime(); - try (LanceScanner scanner = dataset.newScan(scanOptionsBuilder.build())) { - metricsTracker.addScannerCreateTimeNs(System.nanoTime() - scanCreateStart); - try (ArrowReader reader = scanner.scanBatches()) { - while (true) { - long batchStart = System.nanoTime(); - boolean hasNext = reader.loadNextBatch(); - long batchTimeNs = System.nanoTime() - batchStart; - if (!hasNext) { - break; + long dsOpenStart = System.nanoTime(); + try (Dataset dataset = + Utils.openDatasetBuilder(readOptions) + .initialStorageOptions(inputPartition.getInitialStorageOptions()) + .build()) { + metricsTracker.addDatasetOpenTimeNs(System.nanoTime() - dsOpenStart); + + List fragmentIds = inputPartition.getLanceSplit().getFragments(); + if (fragmentIds.isEmpty()) { + return 0; + } + metricsTracker.addNumFragmentsScanned(fragmentIds.size()); + + ScanOptions.Builder scanOptionsBuilder = new ScanOptions.Builder(); + scanOptionsBuilder.useScalarIndex(readOptions.isUseScalarIndex()); + if (inputPartition.getWhereCondition().isPresent()) { + scanOptionsBuilder.filter(inputPartition.getWhereCondition().get()); + } + scanOptionsBuilder.withRowId(true); + scanOptionsBuilder.columns(Lists.newArrayList()); + scanOptionsBuilder.fragmentIds(fragmentIds); + + // Collect scan stats + scanOptionsBuilder.collectStats(true); + + long scanCreateStart = System.nanoTime(); + try (LanceScanner scanner = dataset.newScan(scanOptionsBuilder.build())) { + metricsTracker.addScannerCreateTimeNs(System.nanoTime() - scanCreateStart); + try (ArrowReader reader = scanner.scanBatches()) { + while (true) { + long batchStart = System.nanoTime(); + boolean hasNext = reader.loadNextBatch(); + long batchTimeNs = System.nanoTime() - batchStart; + if (!hasNext) { + break; + } + metricsTracker.addBatchLoadTimeNs(batchTimeNs); + long rowCount = reader.getVectorSchemaRoot().getRowCount(); + totalCount += rowCount; + metricsTracker.addNumBatchesLoaded(1); + metricsTracker.addNumRowsScanned(rowCount); } - metricsTracker.addBatchLoadTimeNs(batchTimeNs); - long rowCount = reader.getVectorSchemaRoot().getRowCount(); - totalCount += rowCount; - metricsTracker.addNumBatchesLoaded(1); - metricsTracker.addNumRowsScanned(rowCount); } + metricsTracker.addScanStats(scanner.getStats()); + } catch (Exception e) { + throw new RuntimeException("Failed to scan fragment " + fragmentIds, e); } - metricsTracker.addScanStats(scanner.getStats()); - } catch (Exception e) { - throw new RuntimeException("Failed to scan fragment " + fragmentIds, e); } - } - return totalCount; + return totalCount; + } finally { + if (namespaceLease != null) { + namespaceLease.close(); + } + } } private ColumnarBatch createCountResultBatch(long count, StructType resultSchema) { diff --git a/lance-spark-base_2.12/src/test/java/org/lance/spark/read/LanceCountStarPartitionReaderTest.java b/lance-spark-base_2.12/src/test/java/org/lance/spark/read/LanceCountStarPartitionReaderTest.java index 23135af59..29ad6f1b2 100644 --- a/lance-spark-base_2.12/src/test/java/org/lance/spark/read/LanceCountStarPartitionReaderTest.java +++ b/lance-spark-base_2.12/src/test/java/org/lance/spark/read/LanceCountStarPartitionReaderTest.java @@ -13,10 +13,15 @@ */ package org.lance.spark.read; +import org.lance.namespace.LanceNamespace; +import org.lance.namespace.model.DescribeTableRequest; +import org.lance.namespace.model.DescribeTableResponse; import org.lance.spark.LanceRuntime; +import org.lance.spark.LanceSparkReadOptions; import org.lance.spark.TestUtils; import org.lance.spark.utils.Optional; +import org.apache.arrow.memory.BufferAllocator; import org.apache.spark.sql.connector.expressions.Expression; import org.apache.spark.sql.connector.expressions.aggregate.AggregateFunc; import org.apache.spark.sql.connector.expressions.aggregate.Aggregation; @@ -25,12 +30,30 @@ import org.junit.jupiter.api.Test; import java.util.Arrays; +import java.util.Collections; +import java.util.Map; +import java.util.UUID; +import java.util.concurrent.atomic.AtomicInteger; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertTrue; public class LanceCountStarPartitionReaderTest { + @Test + public void testFilteredCountsReuseExecutorNamespaceMetadata() throws Exception { + RecordingNamespace.initializeCalls.set(0); + RecordingNamespace.describeCalls.set(0); + RecordingNamespace.location = TestUtils.TestTable1Config.datasetUri; + String scanId = "count-star-namespace-" + UUID.randomUUID(); + + readCount(createNamespacePartition(scanId, 0)); + readCount(createNamespacePartition(scanId, 1)); + + assertEquals(1, RecordingNamespace.initializeCalls.get()); + assertEquals(1, RecordingNamespace.describeCalls.get()); + } + @Test public void testCloseReleasesArrowMemory() throws Exception { // Build a partition with a CountStar aggregation and a filter condition @@ -82,4 +105,61 @@ public void testCloseReleasesArrowMemory() throws Exception { + ", after close: " + memAfterClose); } + + private static void readCount(LanceInputPartition partition) throws Exception { + try (LanceCountStarPartitionReader reader = new LanceCountStarPartitionReader(partition)) { + assertTrue(reader.next()); + reader.get(); + } + } + + private static LanceInputPartition createNamespacePartition(String scanId, int fragmentId) { + LanceSparkReadOptions readOptions = + LanceSparkReadOptions.builder() + .datasetUri(RecordingNamespace.location) + .tableId(Collections.singletonList("test_dataset1")) + .version(6L) + .build(); + Aggregation countStarAgg = + new Aggregation(new AggregateFunc[] {new CountStar()}, new Expression[] {}); + return new LanceInputPartition( + TestUtils.TestTable1Config.schema, + 0, + new LanceSplit(Collections.singletonList(fragmentId)), + readOptions, + Optional.of("x > 0"), + Optional.empty(), + Optional.empty(), + Optional.empty(), + Optional.of(countStarAgg), + scanId, + Collections.emptyMap(), + RecordingNamespace.class.getName(), + Collections.emptyMap(), + null); + } + + public static class RecordingNamespace implements LanceNamespace { + static final AtomicInteger initializeCalls = new AtomicInteger(); + static final AtomicInteger describeCalls = new AtomicInteger(); + static volatile String location; + + public RecordingNamespace() {} + + @Override + public void initialize(Map properties, BufferAllocator allocator) { + initializeCalls.incrementAndGet(); + } + + @Override + public String namespaceId() { + return "count-star-recording"; + } + + @Override + public DescribeTableResponse describeTable(DescribeTableRequest request) { + describeCalls.incrementAndGet(); + return new DescribeTableResponse().location(location); + } + } } From 63680683e13816ec8c785c127aa08dda28181fa1 Mon Sep 17 00:00:00 2001 From: Charles Huang <25107590+charleshuang119@users.noreply.github.com> Date: Tue, 21 Jul 2026 16:58:36 -0700 Subject: [PATCH 5/6] test: harden executor namespace cache contracts --- .../internal/ExecutorNamespaceCache.java | 8 +++ .../internal/ExecutorNamespaceCacheTest.java | 50 +++++++++++++++++++ 2 files changed, 58 insertions(+) diff --git a/lance-spark-base_2.12/src/main/java/org/lance/spark/internal/ExecutorNamespaceCache.java b/lance-spark-base_2.12/src/main/java/org/lance/spark/internal/ExecutorNamespaceCache.java index d39d012d2..4fddd15b0 100644 --- a/lance-spark-base_2.12/src/main/java/org/lance/spark/internal/ExecutorNamespaceCache.java +++ b/lance-spark-base_2.12/src/main/java/org/lance/spark/internal/ExecutorNamespaceCache.java @@ -240,6 +240,14 @@ private void closeDelegate() { } } + /** + * Read-only namespace wrapper used by executor-side Dataset opens. + * + *

Lance 9.0 executor reads use {@code namespaceId}, {@code describeTable}, and, for managed + * versioning, {@code listTableVersions} and {@code describeTableVersion}. Other namespace + * operations intentionally retain {@link LanceNamespace}'s unsupported default implementations; + * worker-side reads must not expose catalog mutation APIs through this cache. + */ static final class CredentialCachingNamespace implements LanceNamespace { private final LanceNamespace delegate; private final LongSupplier clock; diff --git a/lance-spark-base_2.12/src/test/java/org/lance/spark/internal/ExecutorNamespaceCacheTest.java b/lance-spark-base_2.12/src/test/java/org/lance/spark/internal/ExecutorNamespaceCacheTest.java index 96910fbfe..87e74ac16 100644 --- a/lance-spark-base_2.12/src/test/java/org/lance/spark/internal/ExecutorNamespaceCacheTest.java +++ b/lance-spark-base_2.12/src/test/java/org/lance/spark/internal/ExecutorNamespaceCacheTest.java @@ -20,6 +20,7 @@ import org.lance.namespace.model.DescribeTableVersionResponse; import org.lance.namespace.model.ListTableVersionsRequest; import org.lance.namespace.model.ListTableVersionsResponse; +import org.lance.namespace.model.TableExistsRequest; import org.apache.arrow.memory.BufferAllocator; import org.junit.jupiter.api.AfterEach; @@ -35,8 +36,10 @@ import java.util.concurrent.atomic.AtomicLong; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotSame; import static org.junit.jupiter.api.Assertions.assertSame; import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; public class ExecutorNamespaceCacheTest { @@ -97,6 +100,33 @@ public void coalescesConcurrentDescribeTableCalls() throws Exception { } } + @Test + public void cachesValueEquivalentDescribeTableRequests() { + AtomicLong clock = new AtomicLong(100_000L); + RecordingNamespace delegate = new RecordingNamespace(clock); + ExecutorNamespaceCache.CredentialCachingNamespace namespace = + new ExecutorNamespaceCache.CredentialCachingNamespace(delegate, clock::get); + DescribeTableRequest firstRequest = + new DescribeTableRequest() + .addIdItem("catalog") + .addIdItem("table") + .version(1L) + .branch("main") + .vendCredentials(true); + DescribeTableRequest equivalentRequest = + new DescribeTableRequest() + .addIdItem("catalog") + .addIdItem("table") + .version(1L) + .branch("main") + .vendCredentials(true); + + assertNotSame(firstRequest, equivalentRequest); + assertEquals(firstRequest, equivalentRequest); + assertSame(namespace.describeTable(firstRequest), namespace.describeTable(equivalentRequest)); + assertEquals(1, delegate.describeCalls.get()); + } + @Test public void refreshesBeforeTemporaryCredentialsExpire() { AtomicLong clock = new AtomicLong(100_000L); @@ -196,6 +226,26 @@ public DescribeTableVersionResponse describeTableVersion( assertEquals(1, describeVersionCalls.get()); } + @Test + public void keepsUnexpectedNamespaceOperationsReadOnly() { + LanceNamespace delegate = + new RecordingNamespace(new AtomicLong()) { + @Override + public void tableExists(TableExistsRequest request) { + throw new AssertionError("unexpected delegation of a non-read-path operation"); + } + }; + ExecutorNamespaceCache.CredentialCachingNamespace namespace = + new ExecutorNamespaceCache.CredentialCachingNamespace(delegate, () -> 0L); + + org.lance.namespace.errors.UnsupportedOperationException error = + assertThrows( + org.lance.namespace.errors.UnsupportedOperationException.class, + () -> namespace.tableExists(new TableExistsRequest().addIdItem("table"))); + + assertTrue(error.getMessage().contains("tableExists")); + } + private static class RecordingNamespace implements LanceNamespace { final AtomicInteger describeCalls = new AtomicInteger(); final AtomicLong clock; From e007c55f5d3779b12b6caac34539f42d6f0251ea Mon Sep 17 00:00:00 2001 From: Charles Huang <25107590+charleshuang119@users.noreply.github.com> Date: Tue, 21 Jul 2026 17:13:07 -0700 Subject: [PATCH 6/6] fix: isolate namespace providers across scans --- .../internal/ExecutorNamespaceCache.java | 21 +++++++++--- .../internal/ExecutorNamespaceCacheTest.java | 32 +++++++++++++++---- 2 files changed, 41 insertions(+), 12 deletions(-) diff --git a/lance-spark-base_2.12/src/main/java/org/lance/spark/internal/ExecutorNamespaceCache.java b/lance-spark-base_2.12/src/main/java/org/lance/spark/internal/ExecutorNamespaceCache.java index 4fddd15b0..78cb808d1 100644 --- a/lance-spark-base_2.12/src/main/java/org/lance/spark/internal/ExecutorNamespaceCache.java +++ b/lance-spark-base_2.12/src/main/java/org/lance/spark/internal/ExecutorNamespaceCache.java @@ -91,7 +91,8 @@ public static Lease acquire( key, () -> new CachedNamespace( - LanceRuntime.getOrCreateNamespace(namespaceImpl, key.properties))); + LanceRuntime.getOrCreateNamespace(namespaceImpl, key.properties), + key.scanId)); } catch (ExecutionException e) { throw propagate(e.getCause(), "Failed to initialize executor namespace"); } catch (UncheckedExecutionException e) { @@ -194,9 +195,10 @@ private static final class CachedNamespace { private boolean evicted; private boolean closed; - private CachedNamespace(LanceNamespace delegate) { + private CachedNamespace(LanceNamespace delegate, String scanId) { this.delegate = Objects.requireNonNull(delegate, "delegate"); - this.namespace = new CredentialCachingNamespace(delegate, System::currentTimeMillis); + this.namespace = + new CredentialCachingNamespace(delegate, System::currentTimeMillis, scanId); } private synchronized Lease acquire() { @@ -251,6 +253,7 @@ private void closeDelegate() { static final class CredentialCachingNamespace implements LanceNamespace { private final LanceNamespace delegate; private final LongSupplier clock; + private final String namespaceId; private final Cache> descriptionCache = CacheBuilder.newBuilder() .maximumSize(MAX_DESCRIBED_TABLES_PER_NAMESPACE) @@ -259,9 +262,17 @@ static final class CredentialCachingNamespace implements LanceNamespace { private final ConcurrentMap> descriptions = descriptionCache.asMap(); - CredentialCachingNamespace(LanceNamespace delegate, LongSupplier clock) { + CredentialCachingNamespace(LanceNamespace delegate, LongSupplier clock, String scanId) { this.delegate = Objects.requireNonNull(delegate, "delegate"); this.clock = Objects.requireNonNull(clock, "clock"); + // Lance uses namespaceId as part of its object-store/provider cache key. Keep the ID stable + // within one Spark scan, but isolate scans so a later scan cannot reuse a provider whose + // namespace lease belongs to an earlier scan and may be closed on cache eviction. + this.namespaceId = + delegate.namespaceId() + + ",sparkScan[" + + Objects.requireNonNull(scanId, "scanId") + + "]"; } @Override @@ -271,7 +282,7 @@ public void initialize(Map properties, BufferAllocator allocator @Override public String namespaceId() { - return delegate.namespaceId(); + return namespaceId; } @Override diff --git a/lance-spark-base_2.12/src/test/java/org/lance/spark/internal/ExecutorNamespaceCacheTest.java b/lance-spark-base_2.12/src/test/java/org/lance/spark/internal/ExecutorNamespaceCacheTest.java index 87e74ac16..8e3854532 100644 --- a/lance-spark-base_2.12/src/test/java/org/lance/spark/internal/ExecutorNamespaceCacheTest.java +++ b/lance-spark-base_2.12/src/test/java/org/lance/spark/internal/ExecutorNamespaceCacheTest.java @@ -36,6 +36,7 @@ import java.util.concurrent.atomic.AtomicLong; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotEquals; import static org.junit.jupiter.api.Assertions.assertNotSame; import static org.junit.jupiter.api.Assertions.assertSame; import static org.junit.jupiter.api.Assertions.assertThrows; @@ -71,12 +72,29 @@ public void defersClosingEvictedNamespaceUntilAllFragmentLeasesClose() { assertEquals(1, CloseableNamespace.closeCalls.get()); } + @Test + public void separatesObjectStoreIdentityAcrossScans() { + ExecutorNamespaceCache.Lease first = + ExecutorNamespaceCache.acquire( + CloseableNamespace.class.getName(), Map.of("catalog", "test"), "scan-a"); + ExecutorNamespaceCache.Lease second = + ExecutorNamespaceCache.acquire( + CloseableNamespace.class.getName(), Map.of("catalog", "test"), "scan-b"); + + assertNotEquals(first.namespace().namespaceId(), second.namespace().namespaceId()); + assertTrue(first.namespace().namespaceId().endsWith("sparkScan[scan-a]")); + assertTrue(second.namespace().namespaceId().endsWith("sparkScan[scan-b]")); + + first.close(); + second.close(); + } + @Test public void coalescesConcurrentDescribeTableCalls() throws Exception { AtomicLong clock = new AtomicLong(100_000L); BlockingNamespace delegate = new BlockingNamespace(clock); ExecutorNamespaceCache.CredentialCachingNamespace namespace = - new ExecutorNamespaceCache.CredentialCachingNamespace(delegate, clock::get); + new ExecutorNamespaceCache.CredentialCachingNamespace(delegate, clock::get, "test-scan"); DescribeTableRequest request = new DescribeTableRequest().addIdItem("table").version(1L); ExecutorService executor = Executors.newFixedThreadPool(8); @@ -105,7 +123,7 @@ public void cachesValueEquivalentDescribeTableRequests() { AtomicLong clock = new AtomicLong(100_000L); RecordingNamespace delegate = new RecordingNamespace(clock); ExecutorNamespaceCache.CredentialCachingNamespace namespace = - new ExecutorNamespaceCache.CredentialCachingNamespace(delegate, clock::get); + new ExecutorNamespaceCache.CredentialCachingNamespace(delegate, clock::get, "test-scan"); DescribeTableRequest firstRequest = new DescribeTableRequest() .addIdItem("catalog") @@ -132,7 +150,7 @@ public void refreshesBeforeTemporaryCredentialsExpire() { AtomicLong clock = new AtomicLong(100_000L); RecordingNamespace delegate = new RecordingNamespace(clock); ExecutorNamespaceCache.CredentialCachingNamespace namespace = - new ExecutorNamespaceCache.CredentialCachingNamespace(delegate, clock::get); + new ExecutorNamespaceCache.CredentialCachingNamespace(delegate, clock::get, "test-scan"); DescribeTableRequest request = new DescribeTableRequest().addIdItem("table").version(1L); DescribeTableResponse first = namespace.describeTable(request); @@ -153,7 +171,7 @@ public void retriesAfterDescribeTableFailure() { RecordingNamespace delegate = new RecordingNamespace(clock); delegate.failNext = true; ExecutorNamespaceCache.CredentialCachingNamespace namespace = - new ExecutorNamespaceCache.CredentialCachingNamespace(delegate, clock::get); + new ExecutorNamespaceCache.CredentialCachingNamespace(delegate, clock::get, "test-scan"); DescribeTableRequest request = new DescribeTableRequest().addIdItem("table").version(1L); assertThrows(IllegalStateException.class, () -> namespace.describeTable(request)); @@ -183,7 +201,7 @@ public DescribeTableResponse describeTable(DescribeTableRequest request) { } }; ExecutorNamespaceCache.CredentialCachingNamespace namespace = - new ExecutorNamespaceCache.CredentialCachingNamespace(delegate, clock::get); + new ExecutorNamespaceCache.CredentialCachingNamespace(delegate, clock::get, "test-scan"); DescribeTableRequest request = new DescribeTableRequest().addIdItem("table").version(1L); assertEquals("expired-1", namespace.describeTable(request).getStorageOptions().get("token")); @@ -213,7 +231,7 @@ public DescribeTableVersionResponse describeTableVersion( } }; ExecutorNamespaceCache.CredentialCachingNamespace namespace = - new ExecutorNamespaceCache.CredentialCachingNamespace(delegate, () -> 0L); + new ExecutorNamespaceCache.CredentialCachingNamespace(delegate, () -> 0L, "test-scan"); assertSame( listResponse, @@ -236,7 +254,7 @@ public void tableExists(TableExistsRequest request) { } }; ExecutorNamespaceCache.CredentialCachingNamespace namespace = - new ExecutorNamespaceCache.CredentialCachingNamespace(delegate, () -> 0L); + new ExecutorNamespaceCache.CredentialCachingNamespace(delegate, () -> 0L, "test-scan"); org.lance.namespace.errors.UnsupportedOperationException error = assertThrows(