Skip to content

[improve][metadata] Add partition key resolver support for Oxia index scans #26148

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
void-ptr974 wants to merge 2 commits into
base: master
Choose a base branch
from
Open

[improve][metadata] Add partition key resolver support for Oxia index scans #26148

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
31ecc71
[improve][metadata] Add partition key resolver support for Oxia index…
void-ptr974 Jul 5, 2026
b9a9820
[fix][metadata] Harden Oxia index scan partition routing
void-ptr974 Jul 12, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 10 additions & 6 deletions ...-broker/src/main/java/org/apache/pulsar/broker/transaction/metadata/TxnMetadataStore.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -60,6 +60,10 @@ public class TxnMetadataStore {
/** Sequence-keys delta used by all append-only streams in this layout. */
private static final Option.SequenceKeysDeltas APPEND_DELTAS =
new Option.SequenceKeysDeltas(List.of(1L));
private static final Option.PartitionKeyResolver TXN_OP_PARTITION_KEY =
new Option.PartitionKeyResolver(TxnPaths::txnIdFromOpPath);
private static final Option.PartitionKeyResolver TXN_HEADER_PARTITION_KEY =
new Option.PartitionKeyResolver(TxnPaths::txnIdFromHeaderPath);

private final MetadataStore store;

Expand Down Expand Up @@ -147,7 +151,7 @@ public CompletableFuture<Void> listWritesBySegment(String segment, ScanConsumer
TxnOp op = fromJson(gr.getValue(), TxnOp.class);
return op.getKind() == TxnOpKind.WRITE && segment.equals(op.getSegment());
},
consumer);
consumer, Set.of(TXN_OP_PARTITION_KEY));
}

/** Stream all ack ops targeting {@code (segment, subscription)}. */
Expand All @@ -162,7 +166,7 @@ public CompletableFuture<Void> listAcksBySegmentSubscription(String segment, Str
&& segment.equals(op.getSegment())
&& subscription.equals(op.getSubscription());
},
consumer);
consumer, Set.of(TXN_OP_PARTITION_KEY));
}

/**
Expand All @@ -174,7 +178,7 @@ public CompletableFuture<Void> listOpsByTxn(String txnId, ScanConsumer consumer)
return store.scanByIndex(TxnPaths.TXN_OP_PREFIX, TxnPaths.IDX_OPS_BY_TXN,
txnId, txnId,
gr -> txnId.equals(TxnPaths.txnIdFromOpPath(gr.getStat().getPath())),
consumer);
consumer, Set.of(new Option.PartitionKey(txnId)));
}

/**
Expand Down Expand Up @@ -255,7 +259,7 @@ public CompletableFuture<Void> listOpenByDeadlineRange(Long fromMsInclusive, Lon
return (fromMsInclusive == null || deadline >= fromMsInclusive)
&& (toMsInclusive == null || deadline <= toMsInclusive);
},
consumer);
consumer, Set.of(TXN_HEADER_PARTITION_KEY));
}

/**
Expand Down Expand Up @@ -283,7 +287,7 @@ public CompletableFuture<Void> listFinalizedByStateAndTimeRange(TxnState state,
return (fromMsInclusive == null || finalized >= fromMsInclusive)
&& (toMsInclusive == null || finalized <= toMsInclusive);
},
consumer);
consumer, Set.of(TXN_HEADER_PARTITION_KEY));
}

// ---- Event publishing & subscription ----------------------------------
Expand Down Expand Up @@ -380,7 +384,7 @@ public CompletableFuture<Void> scanAbortedTxns(String segment,
// of TXN_SEGMENT_ABORTED_PREFIX named "<segKey>:<txnId>" — match by segKey.
return segKey.equals(TxnPaths.segmentKeyFromAbortedPath(gr.getStat().getPath()));
},
consumer);
consumer, Set.of(new Option.PartitionKey(segKey)));
}

/**
Expand Down
18 changes: 18 additions & 0 deletions pulsar-metadata/src/main/java/org/apache/pulsar/metadata/api/Option.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,8 @@
package org.apache.pulsar.metadata.api;

import java.util.List;
import java.util.Objects;
import java.util.function.Function;

/**
* An option attached to a {@link MetadataStore} operation.
Expand Down Expand Up @@ -67,6 +69,22 @@ record SecondaryIndex(String indexName, String secondaryKey) implements Option {
*/
record PartitionKey(String key) implements Option {}

/**
* Resolve the routing hint from a primary key discovered during a scan. This is useful for
* native secondary-index scans on sharded backends: the index lookup returns primary keys, and
* each primary record may need a different {@link PartitionKey} for the follow-up read.
*
* <p>Backends only consult this option when they have a primary key before fetching a record;
* regular point operations should use {@link PartitionKey}.
*
* @param resolver function mapping primary path to partition key; may return {@code null} if unknown
*/
record PartitionKeyResolver(Function<String, String> resolver) implements Option {
public PartitionKeyResolver {
Objects.requireNonNull(resolver);
}
}

/**
* Request server-assigned multi-dimensional sequence keys on {@code put}. The {@code path}
* argument to {@code put} is treated as a key prefix; the actual stored key is
Expand Down
27 changes: 27 additions & 0 deletions pulsar-metadata/src/main/java/org/apache/pulsar/metadata/api/OptionsHelper.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@

import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
Expand Down Expand Up @@ -55,6 +56,32 @@ public static String partitionKey(Set<Option> opts) {
return null;
}

/**
* Return {@code opts} with a concrete {@link Option.PartitionKey} added when a
* {@link Option.PartitionKeyResolver} can derive one from {@code path}. Existing fixed
* partition keys are preserved.
*/
public static Set<Option> withResolvedPartitionKey(Set<Option> opts, String path) {
if (opts == null || opts.isEmpty()) {
return Set.of();
}
if (partitionKey(opts) != null) {
return opts;
}
for (Option o : opts) {
if (o instanceof Option.PartitionKeyResolver resolver) {
String partitionKey = resolver.resolver().apply(path);

@Denovo1998 Denovo1998 Jul 12, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The PartitionKeyResolver explicitly allows returning null, but the current behavior then performs a lookup without a partition key. If the indexed record was originally stored with a custom partition key, this lookup can access a different shard and return an empty result—silently dropping the indexed data while the scan reports success. Could we either cause the scan to fail when routing cannot be resolved or, alternatively, clearly document and test the precise conditions under which the unpartitioned fallback is safe?

@void-ptr974 void-ptr974 Jul 12, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed. PartitionKeyResolver now requires a non-null key for every scanned primary path. Returning null fails the scan with an IllegalStateException instead of falling back to an unpartitioned read, preventing a silent cross-shard miss.

Added a test that verifies the error callback and confirms that no primary read is issued when the resolver cannot provide a key.

if (partitionKey == null) {
return opts;
}
HashSet<Option> result = new HashSet<>(opts);
result.add(new Option.PartitionKey(partitionKey));
return result;
}
}
return opts;
}

/**
* Build the {@code indexName -> secondaryKey} map from any {@link Option.SecondaryIndex} entries.
* Returns an empty map when no entries are present.
Expand Down
3 changes: 2 additions & 1 deletion pulsar-metadata/src/main/java/org/apache/pulsar/metadata/impl/oxia/OxiaMetadataStore.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -277,8 +277,9 @@ protected CompletableFuture<Void> storeScanByIndex(
// (it's the scan-and-filter compat-path predicate).
CompletableFuture<Void> chain = CompletableFuture.completedFuture(null);
for (String key : primaryKeys) {
Set<Option> getOpts = OptionsHelper.withResolvedPartitionKey(opts, key);
chain = chain
.thenCompose(__ -> storeGet(key, opts))
.thenCompose(__ -> storeGet(key, getOpts))

@Denovo1998 Denovo1998 Jul 12, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The resolver is evaluated eagerly during the chain assembly. Because the first thenCompose is attached to an already-completed future, an earlier storeGet may already be in flight when a later resolver throws. The outer stage then calls onError, but the detached get can still complete and invoke onNext afterward, which violates scanByIndex’s terminal-callback contract. Resolve the partition key inside the sequential thenCompose stage and add a resolver-failure ordering test.

@void-ptr974 void-ptr974 Jul 12, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed. The partition-key resolution now happens inside the sequential thenCompose, so the resolver for the next primary key is not evaluated until the preceding storeGet has completed.

Added a regression test that holds the first read pending and makes the second resolver fail. It verifies that the scan remains pending before the first read completes, then observes onNext followed by onError, with no read for the second key.

.thenAccept(opt -> opt.ifPresent(consumer::onNext));
}
return chain;
Expand Down
45 changes: 45 additions & 0 deletions pulsar-metadata/src/test/java/org/apache/pulsar/metadata/OxiaPartitionKeyTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,16 +21,20 @@
import static org.testng.Assert.assertEquals;
import static org.testng.Assert.assertFalse;
import static org.testng.Assert.assertTrue;
import static org.testng.Assert.fail;
import io.oxia.testcontainers.OxiaContainer;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.List;
import java.util.Optional;
import java.util.Set;
import lombok.Cleanup;
import org.apache.pulsar.metadata.api.GetResult;
import org.apache.pulsar.metadata.api.MetadataStore;
import org.apache.pulsar.metadata.api.MetadataStoreConfig;
import org.apache.pulsar.metadata.api.MetadataStoreFactory;
import org.apache.pulsar.metadata.api.Option;
import org.apache.pulsar.metadata.api.ScanConsumer;
import org.testng.annotations.AfterClass;
import org.testng.annotations.BeforeClass;
import org.testng.annotations.Test;
Expand Down Expand Up @@ -114,4 +118,45 @@ public void getChildrenWithPartitionKey() throws Exception {
assertTrue(children.containsAll(List.of("a", "b")),
"expected children a and b, got: " + children);
}

@Test
public void scanByIndexUsesPartitionKeyResolverForPrimaryReads() throws Exception {
@Cleanup
MetadataStore store = newStore();

String parent = "/partition-key-index-scan-" + System.nanoTime();
String indexName = "idx:partition-key-resolver-" + System.nanoTime();
String indexKey = "match";
RoutedKey routedKey = createIndexedKeyOnlyVisibleWithPartitionKey(store, parent, indexName, indexKey);

@Denovo1998 Denovo1998 Jul 12, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This helper may try multiple candidate records, but every unsuccessful candidate is deleted and it returns immediately after finding the first cross-shard record. Therefore only one indexed record exists when scanByIndex runs. The motivating case is a single scan containing primary records with different partition keys. Could the test retain at least two cross-shard records with distinct partition keys and assert that both are returned? This would catch an implementation that accidentally resolves or caches one partition key per scan rather than per result.

@void-ptr974 void-ptr974 Jul 12, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Updated. The integration test now retains two cross-shard indexed records with distinct partition keys and asserts that both primary paths are returned by the scan. This exercises per-result resolution rather than a single scan-level routing key.

The focused unit and multi-shard Oxia integration tests pass.


List<GetResult> results = new ArrayList<>();
Set<Option> opts = Set.of(new Option.PartitionKeyResolver(
path -> path.substring(path.lastIndexOf('/') + 1, path.lastIndexOf("-record"))));
store.scanByIndex(parent, indexName, indexKey, indexKey, __ -> true,
ScanConsumer.collectInto(results), opts).get();

assertEquals(results.stream().map(result -> result.getStat().getPath()).toList(),
List.of(routedKey.path()),
"scanByIndex should use the resolver to route follow-up primary-key reads");
}

private RoutedKey createIndexedKeyOnlyVisibleWithPartitionKey(
MetadataStore store, String parent, String indexName, String indexKey) throws Exception {
for (int i = 0; i < 100; i++) {
String partitionKey = "index-route-" + i + "-" + System.nanoTime();
String path = parent + "/" + partitionKey + "-record";
Set<Option> options = Set.of(new Option.PartitionKey(partitionKey),
new Option.SecondaryIndex(indexName, indexKey));
store.put(path, "value".getBytes(StandardCharsets.UTF_8), Optional.empty(), options).get();
if (store.get(path).get().isEmpty()) {
return new RoutedKey(path, options);
}
store.deleteIfExists(path, Optional.empty(), options).get();
}
fail("Could not find an indexed key whose PartitionKey routes to a different shard from its path");
return null;
}

private record RoutedKey(String path, Set<Option> options) {
}
}