Skip to content

Commit 526a72f

Browse files
authored
feat: expose branch_identifier in python and java bindings (#6360)
### Motivation - Surface the full branch lineage (version -> uuid mapping) from Rust `BranchIdentifier` into higher-level bindings so clients can inspect a branch's lineage chain. ### Description - Update the JNI layer (`java/lance-jni/src/blocking_dataset.rs`) to construct and pass a `List<BranchVersionMapping>` when creating the Java `Branch` object. - Extend the Java model (`java/src/main/java/org/lance/Branch.java`) with an inner `BranchVersionMapping` class, add a `branchIdentifier` field and getter, update constructor/signature, and include it in `toString`, `equals`, and `hashCode`. - Update Java tests (`java/src/test/java/org/lance/DatasetTest.java`) to assert the presence and contents of `branchIdentifier` for created branches. - Add `branch_identifier: List[Tuple[int, str]]` to the Python `Branch` `TypedDict` (`python/python/lance/dataset.py`) and update the Python Rust bridge (`python/src/dataset.rs`) to include `meta.identifier.version_mapping` in the returned branch dictionaries. - Update Python tests (`python/python/tests/test_dataset.py`) to validate the `branch_identifier` contents in `branches.list()` responses. ### Testing - Ran the modified Java unit test `DatasetTest::test_branches` which exercises branch creation, listing and metadata assertions, and it passed. - Ran the modified Python test `test_branches` with `pytest` which verifies `branch_identifier` in `branches.list()` and it passed. - Exercised JNI/bridge code paths while running the above tests to validate interop and serialization of the version->uuid mappings, and no failures were observed.
1 parent 88d8faf commit 526a72f

6 files changed

Lines changed: 106 additions & 8 deletions

File tree

java/lance-jni/src/blocking_dataset.rs

Lines changed: 21 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2533,12 +2533,32 @@ fn inner_list_branches<'local>(
25332533
} else {
25342534
JObject::null()
25352535
};
2536+
let jbranch_identifier = env.new_object(
2537+
"java/util/ArrayList",
2538+
"(I)V",
2539+
&[JValue::Int(contents.identifier.version_mapping.len() as i32)],
2540+
)?;
2541+
for (version, uuid) in contents.identifier.version_mapping.iter() {
2542+
let juuid = env.new_string(uuid)?;
2543+
let jmapping = env.new_object(
2544+
"org/lance/Branch$BranchVersionMapping",
2545+
"(JLjava/lang/String;)V",
2546+
&[JValue::Long(*version as i64), JValue::Object(&juuid)],
2547+
)?;
2548+
env.call_method(
2549+
&jbranch_identifier,
2550+
"add",
2551+
"(Ljava/lang/Object;)Z",
2552+
&[JValue::Object(&jmapping)],
2553+
)?;
2554+
}
25362555
let jbranch = env.new_object(
25372556
"org/lance/Branch",
2538-
"(Ljava/lang/String;Ljava/lang/String;JJI)V",
2557+
"(Ljava/lang/String;Ljava/lang/String;Ljava/util/List;JJI)V",
25392558
&[
25402559
JValue::Object(&jname),
25412560
JValue::Object(&jparent),
2561+
JValue::Object(&jbranch_identifier),
25422562
JValue::Long(contents.parent_version as i64),
25432563
JValue::Long(contents.create_at as i64),
25442564
JValue::Int(contents.manifest_size as i32),

java/src/main/java/org/lance/Branch.java

Lines changed: 63 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -14,27 +14,78 @@
1414
package org.lance;
1515

1616
import com.google.common.base.MoreObjects;
17+
import com.google.common.collect.ImmutableList;
1718

19+
import java.util.List;
1820
import java.util.Objects;
1921
import java.util.Optional;
2022

2123
/**
2224
* Branch metadata aligned with Rust's BranchContents. name is the branch name, parentBranch may be
23-
* null (indicating main), parentVersion is the version on which the branch was created, createAt is
24-
* the unix timestamp (seconds), and manifestSize is the size of the referenced manifest file in
25-
* bytes.
25+
* null (indicating main), branchIdentifier is the lineage chain {@code [(version, uuid), ...]},
26+
* parentVersion is the version on which the branch was created, createAt is the unix timestamp
27+
* (seconds), and manifestSize is the size of the referenced manifest file in bytes.
2628
*/
2729
public class Branch {
30+
/** A single lineage hop in the Rust BranchIdentifier.version_mapping vector. */
31+
public static class BranchVersionMapping {
32+
private final long version;
33+
private final String uuid;
34+
35+
public BranchVersionMapping(long version, String uuid) {
36+
this.version = version;
37+
this.uuid = Objects.requireNonNull(uuid);
38+
}
39+
40+
public long getVersion() {
41+
return version;
42+
}
43+
44+
public String getUuid() {
45+
return uuid;
46+
}
47+
48+
@Override
49+
public String toString() {
50+
return MoreObjects.toStringHelper(this).add("version", version).add("uuid", uuid).toString();
51+
}
52+
53+
@Override
54+
public boolean equals(Object o) {
55+
if (this == o) return true;
56+
if (o == null || getClass() != o.getClass()) return false;
57+
BranchVersionMapping that = (BranchVersionMapping) o;
58+
return version == that.version && Objects.equals(uuid, that.uuid);
59+
}
60+
61+
@Override
62+
public int hashCode() {
63+
return Objects.hash(version, uuid);
64+
}
65+
}
66+
2867
private final String name;
2968
private final Optional<String> parentBranch;
69+
private final ImmutableList<BranchVersionMapping> branchIdentifier;
3070
private final long parentVersion;
3171
private final long createAt;
3272
private final int manifestSize;
3373

3474
public Branch(
3575
String name, String parentBranch, long parentVersion, long createAt, int manifestSize) {
76+
this(name, parentBranch, ImmutableList.of(), parentVersion, createAt, manifestSize);
77+
}
78+
79+
public Branch(
80+
String name,
81+
String parentBranch,
82+
List<BranchVersionMapping> branchIdentifier,
83+
long parentVersion,
84+
long createAt,
85+
int manifestSize) {
3686
this.name = name;
3787
this.parentBranch = Optional.ofNullable(parentBranch);
88+
this.branchIdentifier = ImmutableList.copyOf(Objects.requireNonNull(branchIdentifier));
3889
this.parentVersion = parentVersion;
3990
this.createAt = createAt;
4091
this.manifestSize = manifestSize;
@@ -48,6 +99,10 @@ public Optional<String> getParentBranch() {
4899
return parentBranch;
49100
}
50101

102+
public ImmutableList<BranchVersionMapping> getBranchIdentifier() {
103+
return branchIdentifier;
104+
}
105+
51106
public long getParentVersion() {
52107
return parentVersion;
53108
}
@@ -65,6 +120,7 @@ public String toString() {
65120
return MoreObjects.toStringHelper(this)
66121
.add("name", name)
67122
.add("parentBranch", parentBranch)
123+
.add("branchIdentifier", branchIdentifier)
68124
.add("parentVersion", parentVersion)
69125
.add("createAt", createAt)
70126
.add("manifestSize", manifestSize)
@@ -80,11 +136,13 @@ public boolean equals(Object o) {
80136
&& createAt == branch.createAt
81137
&& manifestSize == branch.manifestSize
82138
&& Objects.equals(name, branch.name)
83-
&& Objects.equals(parentBranch, branch.parentBranch);
139+
&& Objects.equals(parentBranch, branch.parentBranch)
140+
&& Objects.equals(branchIdentifier, branch.branchIdentifier);
84141
}
85142

86143
@Override
87144
public int hashCode() {
88-
return Objects.hash(name, parentBranch, parentVersion, createAt, manifestSize);
145+
return Objects.hash(
146+
name, parentBranch, branchIdentifier, parentVersion, createAt, manifestSize);
89147
}
90148
}

java/src/test/java/org/lance/DatasetTest.java

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1736,12 +1736,21 @@ void testBranches(@TempDir Path tempDir) {
17361736
assertEquals("branch1", branch1Meta.getName());
17371737
assertEquals(2, branch1Meta.getParentVersion());
17381738
assertFalse(branch1Meta.getParentBranch().isPresent());
1739+
assertEquals(1, branch1Meta.getBranchIdentifier().size());
1740+
assertEquals(2, branch1Meta.getBranchIdentifier().get(0).getVersion());
1741+
assertFalse(branch1Meta.getBranchIdentifier().get(0).getUuid().isEmpty());
17391742
assertTrue(branch1Meta.getCreateAt() > 0);
17401743
assertTrue(branch1Meta.getManifestSize() > 0);
17411744

17421745
assertEquals("branch2", branch2Meta.getName());
17431746
assertTrue(branch2Meta.getParentBranch().isPresent());
17441747
assertEquals("branch1", branch2Meta.getParentBranch().get());
1748+
assertEquals(2, branch2Meta.getBranchIdentifier().size());
1749+
assertEquals(
1750+
branch1Meta.getBranchIdentifier().get(0),
1751+
branch2Meta.getBranchIdentifier().get(0));
1752+
assertEquals(3, branch2Meta.getBranchIdentifier().get(1).getVersion());
1753+
assertFalse(branch2Meta.getBranchIdentifier().get(1).getUuid().isEmpty());
17451754
assertEquals(3, branch2Meta.getParentVersion());
17461755
assertTrue(branch2Meta.getCreateAt() > 0);
17471756
assertTrue(branch2Meta.getManifestSize() > 0);

python/python/lance/dataset.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4375,6 +4375,7 @@ class Tag(TypedDict):
43754375

43764376
class Branch(TypedDict):
43774377
parent_branch: Optional[str]
4378+
branch_identifier: List[Tuple[int, str]]
43784379
parent_version: int
43794380
create_at: int
43804381
manifest_size: int

python/python/tests/test_dataset.py

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5231,11 +5231,16 @@ def test_branches(tmp_path: Path):
52315231
branch1.tags.create("main_latest", (None, None))
52325232
branch1.tags.create("main_latest2", ("main", None))
52335233
branch1.create_branch("branch_from_main", ("main", None))
5234+
branches_with_main = branch1.branches.list()
52345235
assert branch1.tags.list()["branch1_latest"]["branch"] == "branch1"
52355236
assert branch1.tags.list()["main_latest"]["branch"] is None
52365237
assert branch1.tags.list()["main_latest2"]["branch"] is None
5237-
assert branch1.branches.list()["branch_from_main"]["parent_branch"] is None
5238-
assert branch1.branches.list()["branch_from_main"]["parent_version"] == 1
5238+
assert branches_with_main["branch_from_main"]["parent_branch"] is None
5239+
assert branches_with_main["branch_from_main"]["branch_identifier"][0][0] == 1
5240+
assert isinstance(
5241+
branches_with_main["branch_from_main"]["branch_identifier"][0][1], str
5242+
)
5243+
assert branches_with_main["branch_from_main"]["parent_version"] == 1
52395244
assert branch1.checkout_version("main_latest").latest_version == 1
52405245
assert branch1.checkout_version("main_latest2").latest_version == 1
52415246
assert branch1.checkout_version(("branch_from_main", None)).latest_version == 1
@@ -5261,6 +5266,9 @@ def test_branches(tmp_path: Path):
52615266
b1_meta = branches["branch1"]
52625267
assert isinstance(b1_meta["parent_version"], int)
52635268
assert b1_meta["manifest_size"] > 0
5269+
assert b1_meta["branch_identifier"][0][0] == b1_meta["parent_version"]
5270+
assert isinstance(b1_meta["branch_identifier"][0][1], str)
5271+
assert len(b1_meta["branch_identifier"][0][1]) > 0
52645272
assert "create_at" in b1_meta
52655273

52665274
try:

python/src/dataset.rs

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1818,6 +1818,7 @@ impl Dataset {
18181818
for (name, meta) in branches.iter() {
18191819
let dict = PyDict::new(py);
18201820
dict.set_item("parent_branch", meta.parent_branch.clone())?;
1821+
dict.set_item("branch_identifier", meta.identifier.version_mapping.clone())?;
18211822
dict.set_item("parent_version", meta.parent_version)?;
18221823
dict.set_item("create_at", meta.create_at)?;
18231824
dict.set_item("manifest_size", meta.manifest_size)?;
@@ -1852,6 +1853,7 @@ impl Dataset {
18521853
for (name, meta) in ordered.into_iter() {
18531854
let dict = PyDict::new(py);
18541855
dict.set_item("parent_branch", meta.parent_branch.clone())?;
1856+
dict.set_item("branch_identifier", meta.identifier.version_mapping.clone())?;
18551857
dict.set_item("parent_version", meta.parent_version)?;
18561858
dict.set_item("create_at", meta.create_at)?;
18571859
dict.set_item("manifest_size", meta.manifest_size)?;

0 commit comments

Comments
 (0)