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..78cb808d1 --- /dev/null +++ b/lance-spark-base_2.12/src/main/java/org/lance/spark/internal/ExecutorNamespaceCache.java @@ -0,0 +1,399 @@ +/* + * 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.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 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 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), + key.scanId)); + } 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, String scanId) { + this.delegate = Objects.requireNonNull(delegate, "delegate"); + this.namespace = + new CredentialCachingNamespace(delegate, System::currentTimeMillis, scanId); + } + + 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); + } + } + } + } + + /** + * 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; + private final String namespaceId; + 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, 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 + public void initialize(Map properties, BufferAllocator allocator) { + throw new UnsupportedOperationException("Cached namespace is already initialized"); + } + + @Override + public String namespaceId() { + return 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; + } + } + } + + @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(); + } 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/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/internal/ExecutorNamespaceCacheTest.java b/lance-spark-base_2.12/src/test/java/org/lance/spark/internal/ExecutorNamespaceCacheTest.java new file mode 100644 index 000000000..8e3854532 --- /dev/null +++ b/lance-spark-base_2.12/src/test/java/org/lance/spark/internal/ExecutorNamespaceCacheTest.java @@ -0,0 +1,346 @@ +/* + * 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.namespace.model.DescribeTableVersionRequest; +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; +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.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; +import static org.junit.jupiter.api.Assertions.assertTrue; + +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 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, "test-scan"); + 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 cachesValueEquivalentDescribeTableRequests() { + AtomicLong clock = new AtomicLong(100_000L); + RecordingNamespace delegate = new RecordingNamespace(clock); + ExecutorNamespaceCache.CredentialCachingNamespace namespace = + new ExecutorNamespaceCache.CredentialCachingNamespace(delegate, clock::get, "test-scan"); + 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); + RecordingNamespace delegate = new RecordingNamespace(clock); + ExecutorNamespaceCache.CredentialCachingNamespace namespace = + new ExecutorNamespaceCache.CredentialCachingNamespace(delegate, clock::get, "test-scan"); + 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, "test-scan"); + 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, "test-scan"); + 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()); + } + + @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, "test-scan"); + + 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()); + } + + @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, "test-scan"); + + 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; + 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..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 @@ -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,75 @@ 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. + } + 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()); + } + + @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. + } + assertEquals(1, RecordingNamespace.INITIALIZE_CALLS.get()); + assertEquals(1, RecordingNamespace.DESCRIBE_CALLS.get()); + + 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 +386,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 +400,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); + } } } 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); + } + } }