Skip to content

Commit 7ceafb1

Browse files
rustyconoverrustyconover
committed
feat: support protocol 3.1 partition distribution
1 parent 0789d65 commit 7ceafb1

18 files changed

Lines changed: 402 additions & 23 deletions

‎README.md‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@ VGI lets [Haybarn](https://github.com/Query-farm-haybarn/haybarn) — Query Farm
1919

2020
Wire-compatible with the [Python reference implementation](https://github.com/Query-farm/vgi-python) and the [Go port](https://github.com/Query-farm/vgi-go): all three serve the same integration suite against the same C++ extension.
2121

22-
The prepared protocol uses VGI wire 3.0.0 under `vgi.v3` and requires a native VGI client speaking wire 3.0.x. Upgrade clients and workers together; published `vgi.v2` clients are incompatible. See [catalog metadata and publication](docs/catalog_publication.md).
22+
The prepared protocol uses VGI wire 3.1.0 under `vgi.v3` and requires a native VGI client speaking wire 3.1.x. Upgrade clients and workers together; published `vgi.v2` clients are incompatible. See [catalog metadata and publication](docs/catalog_publication.md).
2323

2424
## What you can serve
2525

‎vgi/src/main/java/farm/query/vgi/VgiService.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
// GENERATED by vgi.codegen.java_registry. DO NOT EDIT BY HAND.
22
//
33
// Generator: python -m vgi.codegen.java_registry v1
4-
// Content hash: e5b43b58e98f
4+
// Content hash: 06119187b8de
55
//
66
// To regenerate:
77
// uv run --project ~/Development/vgi-python python scripts/regen_generated.py
@@ -93,7 +93,7 @@ public interface VgiService {
9393
String PROTOCOL_NAME = "vgi.v3";
9494

9595
/** The protocol's semver, stamped on every request by a client of this interface. */
96-
String PROTOCOL_VERSION = "3.0.0";
96+
String PROTOCOL_VERSION = "3.1.0";
9797
/** Canonical fingerprint of the nested catalog wire schemas. */
9898
String CATALOG_CONTRACT_HASH = "1d12ce672b43ca1510721f89637bb374d27693b100d357ac2ead99fcddc27e9f";
9999

‎vgi/src/main/java/farm/query/vgi/internal/ScanSplitSerializer.java‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -59,6 +59,9 @@ public static byte[] serialize(ScanSplit split) {
5959
i64("estimated_rows", true),
6060
new Field("rows_exact", new FieldType(false, BOOL, null), null),
6161
i64("estimated_bytes", true),
62+
bin("key_value", true),
63+
i64("range_index", true),
64+
bin("column_bounds", true),
6265
bin("partition_bounds", true),
6366
bin("column_statistics", true),
6467
new Field("location_ids", new FieldType(true, new ArrowType.List(), null),
@@ -73,6 +76,9 @@ public static byte[] serialize(ScanSplit split) {
7376
setLong(root, "estimated_rows", split.estimated_rows());
7477
((BitVector) root.getVector("rows_exact")).setSafe(0, split.rows_exact() ? 1 : 0);
7578
setLong(root, "estimated_bytes", split.estimated_bytes());
79+
setBinary(root, "key_value", split.key_value());
80+
setLong(root, "range_index", split.range_index());
81+
setBinary(root, "column_bounds", split.column_bounds());
7682
setBinary(root, "partition_bounds", split.partition_bounds());
7783
setBinary(root, "column_statistics", split.column_statistics());
7884
setBinary(root, "start_position", split.start_position());

‎vgi/src/main/java/farm/query/vgi/internal/VgiServiceImpl.java‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1146,14 +1146,16 @@ public PlanResponse table_function_plan(byte[] request, CallContext ctx) {
11461146
}
11471147
PlanResponse response = PlanResponse.of(blobs);
11481148
if (!result.nextCursors().isEmpty() || result.estimatedTotalSplits() != null
1149-
|| result.estimatedTotalRows() != null || result.maxWorkers() != null) {
1149+
|| result.estimatedTotalRows() != null || result.maxWorkers() != null
1150+
|| result.keyGrouping() != null || result.rangePartitioning() != null) {
11501151
response = new PlanResponse(
11511152
response.splits(),
11521153
result.nextCursors(),
11531154
response.execution_id(), response.init_opaque_data(),
11541155
result.maxWorkers(), result.estimatedTotalSplits(), result.estimatedTotalRows(),
11551156
response.estimated_total_bytes(), response.catalog_version(), response.scope(),
1156-
response.locations(), response.partitioning(), response.sort_order(),
1157+
response.locations(), result.keyGrouping(), result.rangePartitioning(),
1158+
response.partitioning(), response.sort_order(),
11571159
response.cache_max_age_seconds(), response.start_position(),
11581160
response.end_position());
11591161
}

‎vgi/src/main/java/farm/query/vgi/protocol/FieldPath.java‎

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,37 @@
1+
// GENERATED by vgi.codegen.java_types. DO NOT EDIT BY HAND.
2+
//
3+
// Generator: python -m vgi.codegen.java_types FieldPath v1
4+
// Content hash: 7008cd58bf8c
5+
//
6+
// To regenerate:
7+
// uv run --project ~/Development/vgi-python python scripts/regen_generated.py
8+
9+
// Copyright 2025, 2026 Query Farm LLC - https://query.farm
10+
11+
package farm.query.vgi.protocol;
12+
13+
import farm.query.vgirpc.schema.ArrowSerializableRecord;
14+
15+
import java.util.List;
16+
17+
/**
18+
* One raw, component-preserving field reference used by a transform.
19+
*
20+
* <p>Protocol type {@code vgi.protocol.FieldPath}, generated from its {@code ARROW_SCHEMA}:
21+
* component order is column order and {@code @Nullable} is wire nullability.
22+
*
23+
* @param components the {&#64;code components} column
24+
*/
25+
public record FieldPath(
26+
List<String> components) implements ArrowSerializableRecord {
27+
28+
/**
29+
* Reads a {@code null} list/map component as empty: those columns are
30+
* non-nullable on the wire, and the record codec would otherwise write a
31+
* null into them, which a strict client rejects.
32+
*/
33+
public FieldPath {
34+
components = components == null ? List.of() : components;
35+
}
36+
37+
}

‎vgi/src/main/java/farm/query/vgi/protocol/KeyGrouping.java‎

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,37 @@
1+
// GENERATED by vgi.codegen.java_types. DO NOT EDIT BY HAND.
2+
//
3+
// Generator: python -m vgi.codegen.java_types KeyGrouping v1
4+
// Content hash: c2c6fdfe3599
5+
//
6+
// To regenerate:
7+
// uv run --project ~/Development/vgi-python python scripts/regen_generated.py
8+
9+
// Copyright 2025, 2026 Query Farm LLC - https://query.farm
10+
11+
package farm.query.vgi.protocol;
12+
13+
import farm.query.vgirpc.schema.ArrowSerializableRecord;
14+
15+
import java.util.List;
16+
17+
/**
18+
* Transforms whose exact result tuple is constant within every split.
19+
*
20+
* <p>Protocol type {@code vgi.protocol.KeyGrouping}, generated from its {@code ARROW_SCHEMA}:
21+
* component order is column order and {@code @Nullable} is wire nullability.
22+
*
23+
* @param keys the {&#64;code keys} column
24+
*/
25+
public record KeyGrouping(
26+
List<byte[]> keys) implements ArrowSerializableRecord {
27+
28+
/**
29+
* Reads a {@code null} list/map component as empty: those columns are
30+
* non-nullable on the wire, and the record codec would otherwise write a
31+
* null into them, which a strict client rejects.
32+
*/
33+
public KeyGrouping {
34+
keys = keys == null ? List.of() : keys;
35+
}
36+
37+
}

‎vgi/src/main/java/farm/query/vgi/protocol/PlanResponse.java‎

Lines changed: 34 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -37,10 +37,11 @@
3737
* {@code transaction}
3838
* @param locations a hoisted host list splits index into, which keeps a large
3939
* plan off the coordinator heap
40-
* @param partitioning serialised partition transforms. NOT derivable from a
41-
* split's partition values — {@code country=US} does not say whether
42-
* partitions are {@code identity(country)} or {@code bucket(16, user_id)}.
43-
* Report nothing here unless every split really is single-valued
40+
* @param key_grouping serialized {@link KeyGrouping}, when each split has one
41+
* exact transformed key tuple
42+
* @param range_partitioning serialized {@link RangePartitioning}, containing
43+
* exact ordered split points and their comparator
44+
* @param partitioning deprecated protocol-2 partition-transform list
4445
* @param sort_order ordering WITHIN each split, never a global claim across
4546
* splits: concatenating K non-contiguous sorted runs is not sorted
4647
* @param cache_max_age_seconds how long this plan stays reusable
@@ -64,6 +65,8 @@ public record PlanResponse(
6465
@Nullable Long catalog_version,
6566
String scope,
6667
@Nullable List<String> locations,
68+
@Nullable byte[] key_grouping,
69+
@Nullable byte[] range_partitioning,
6770
@Nullable List<byte[]> partitioning,
6871
@Nullable List<byte[]> sort_order,
6972
@Nullable Long cache_max_age_seconds,
@@ -79,6 +82,31 @@ public record PlanResponse(
7982
splits = splits == null ? List.of() : splits;
8083
}
8184

85+
/** Protocol-2 source compatibility while protocol 3 producers migrate. */
86+
@Deprecated
87+
public PlanResponse(
88+
List<byte[]> splits,
89+
@Nullable List<byte[]> next_cursors,
90+
@Nullable byte[] execution_id,
91+
@Nullable byte[] init_opaque_data,
92+
@Nullable Long max_workers,
93+
@Nullable Long estimated_total_splits,
94+
@Nullable Long estimated_total_rows,
95+
@Nullable Long estimated_total_bytes,
96+
@Nullable Long catalog_version,
97+
String scope,
98+
@Nullable List<String> locations,
99+
@Nullable List<byte[]> partitioning,
100+
@Nullable List<byte[]> sort_order,
101+
@Nullable Long cache_max_age_seconds,
102+
@Nullable byte[] start_position,
103+
@Nullable byte[] end_position) {
104+
this(splits, next_cursors, execution_id, init_opaque_data, max_workers,
105+
estimated_total_splits, estimated_total_rows, estimated_total_bytes,
106+
catalog_version, scope, locations, null, null, partitioning, sort_order,
107+
cache_max_age_seconds, start_position, end_position);
108+
}
109+
82110

83111
/** A plan with no splits: legal, and means the scan has no work to do.
84112
*
@@ -98,6 +126,7 @@ public static PlanResponse empty() {
98126
public static PlanResponse of(List<byte[]> splits) {
99127
return new PlanResponse(
100128
splits, List.of(), null, new byte[0], null, null, null, null, null,
101-
"catalog", null, List.of(), List.of(), null, new byte[0], new byte[0]);
129+
"catalog", null, null, null, List.of(), List.of(), null,
130+
new byte[0], new byte[0]);
102131
}
103132
}

‎vgi/src/main/java/farm/query/vgi/protocol/RangeKey.java‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,31 @@
1+
// GENERATED by vgi.codegen.java_types. DO NOT EDIT BY HAND.
2+
//
3+
// Generator: python -m vgi.codegen.java_types RangeKey v1
4+
// Content hash: 2090c2fb015e
5+
//
6+
// To regenerate:
7+
// uv run --project ~/Development/vgi-python python scripts/regen_generated.py
8+
9+
// Copyright 2025, 2026 Query Farm LLC - https://query.farm
10+
11+
package farm.query.vgi.protocol;
12+
13+
import farm.query.vgirpc.schema.ArrowField;
14+
import farm.query.vgirpc.schema.ArrowFieldType;
15+
import farm.query.vgirpc.schema.ArrowSerializableRecord;
16+
17+
/**
18+
* One transform and comparator component of a range key.
19+
*
20+
* <p>Protocol type {@code vgi.protocol.RangeKey}, generated from its {@code ARROW_SCHEMA}:
21+
* component order is column order and {@code @Nullable} is wire nullability.
22+
*
23+
* @param transform the {&#64;code transform} column
24+
* @param direction the {&#64;code direction} column
25+
* @param null_order the {&#64;code null_order} column
26+
*/
27+
public record RangeKey(
28+
byte[] transform,
29+
@ArrowField(ArrowFieldType.DICT_INT16_UTF8) String direction,
30+
@ArrowField(ArrowFieldType.DICT_INT16_UTF8) String null_order) implements ArrowSerializableRecord {
31+
}

‎vgi/src/main/java/farm/query/vgi/protocol/RangePartitioning.java‎

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,42 @@
1+
// GENERATED by vgi.codegen.java_types. DO NOT EDIT BY HAND.
2+
//
3+
// Generator: python -m vgi.codegen.java_types RangePartitioning v1
4+
// Content hash: 287a5f410fd1
5+
//
6+
// To regenerate:
7+
// uv run --project ~/Development/vgi-python python scripts/regen_generated.py
8+
9+
// Copyright 2025, 2026 Query Farm LLC - https://query.farm
10+
11+
package farm.query.vgi.protocol;
12+
13+
import farm.query.vgirpc.schema.ArrowSerializableRecord;
14+
15+
import java.util.List;
16+
17+
/**
18+
* Exact half-open range distribution shared by all splits in a plan.
19+
*
20+
* <p>{@code N} split-point rows define {@code N + 1} ranges. Values equal to point {@code i} belong
21+
* to range {@code i + 1}. Each split names its range with :attr:{@code ScanSplit.range_index}.
22+
*
23+
* <p>Protocol type {@code vgi.protocol.RangePartitioning}, generated from its {@code ARROW_SCHEMA}:
24+
* component order is column order and {@code @Nullable} is wire nullability.
25+
*
26+
* @param keys the {&#64;code keys} column
27+
* @param split_points the {&#64;code split_points} column
28+
*/
29+
public record RangePartitioning(
30+
List<byte[]> keys,
31+
byte[] split_points) implements ArrowSerializableRecord {
32+
33+
/**
34+
* Reads a {@code null} list/map component as empty: those columns are
35+
* non-nullable on the wire, and the record codec would otherwise write a
36+
* null into them, which a strict client rejects.
37+
*/
38+
public RangePartitioning {
39+
keys = keys == null ? List.of() : keys;
40+
}
41+
42+
}

‎vgi/src/main/java/farm/query/vgi/protocol/ScanSplit.java‎

Lines changed: 45 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -31,8 +31,10 @@
3131
* @param estimated_bytes byte estimate. Load-bearing for engines that bin-pack
3232
* (DataFusion weight, Trino SplitWeight); {@code null} degrades them to
3333
* round-robin by count. A greedily claiming client needs no cost model
34-
* @param partition_bounds 2-row (min, max) batch in the existing
35-
* {@code vgi_partition_values} encoding, one column per partition column
34+
* @param key_value exact one-row transformed key tuple for key grouping
35+
* @param range_index zero-based range owned by this split
36+
* @param column_bounds 2-row (min, max) containment/pruning metadata
37+
* @param partition_bounds deprecated protocol-2 spelling for column bounds
3638
* @param column_statistics per-column statistics blob for this split
3739
* @param location_ids indices into {@code PlanResponse.locations} naming hosts
3840
* where this split is cheap to read
@@ -46,19 +48,41 @@ public record ScanSplit(
4648
@Nullable Long estimated_rows,
4749
boolean rows_exact,
4850
@Nullable Long estimated_bytes,
51+
@Nullable byte[] key_value,
52+
@Nullable Long range_index,
53+
@Nullable byte[] column_bounds,
4954
@Nullable byte[] partition_bounds,
5055
@Nullable byte[] column_statistics,
5156
@Nullable List<Long> location_ids,
5257
@Nullable byte[] start_position,
5358
@Nullable byte[] end_position) implements ArrowSerializableRecord {
5459

60+
/** Protocol-2 source compatibility while protocol 3 producers migrate. */
61+
@Deprecated
62+
public ScanSplit(
63+
byte[] payload,
64+
byte[] token,
65+
@Nullable Long estimated_rows,
66+
boolean rows_exact,
67+
@Nullable Long estimated_bytes,
68+
@Nullable byte[] partition_bounds,
69+
@Nullable byte[] column_statistics,
70+
@Nullable List<Long> location_ids,
71+
@Nullable byte[] start_position,
72+
@Nullable byte[] end_position) {
73+
this(payload, token, estimated_rows, rows_exact, estimated_bytes, null, null,
74+
null, partition_bounds, column_statistics, location_ids, start_position,
75+
end_position);
76+
}
77+
5578
/** A split naming the given work, with no estimates.
5679
*
5780
* @param payload the worker's own bytes naming this unit of work
5881
* @return a split carrying that payload
5982
*/
6083
public static ScanSplit of(byte[] payload) {
61-
return new ScanSplit(payload, new byte[0], null, false, null, null, null, null, null, null);
84+
return new ScanSplit(payload, new byte[0], null, false, null, null, null,
85+
null, null, null, null, null, null);
6286
}
6387

6488
/** A split naming the given work, with an exact row count and a byte estimate.
@@ -69,7 +93,22 @@ public static ScanSplit of(byte[] payload) {
6993
* @return a split carrying that payload and those estimates
7094
*/
7195
public static ScanSplit of(byte[] payload, long rows, long bytes) {
72-
return new ScanSplit(payload, new byte[0], rows, true, bytes, null, null, null, null, null);
96+
return new ScanSplit(payload, new byte[0], rows, true, bytes, null, null,
97+
null, null, null, null, null, null);
98+
}
99+
100+
/** This split with its exact transformed key tuple. */
101+
public ScanSplit withKeyValue(byte[] value) {
102+
return new ScanSplit(payload, token, estimated_rows, rows_exact, estimated_bytes,
103+
value, range_index, column_bounds, partition_bounds, column_statistics,
104+
location_ids, start_position, end_position);
105+
}
106+
107+
/** This split assigned to one zero-based range. */
108+
public ScanSplit withRangeIndex(long value) {
109+
return new ScanSplit(payload, token, estimated_rows, rows_exact, estimated_bytes,
110+
key_value, value, column_bounds, partition_bounds, column_statistics,
111+
location_ids, start_position, end_position);
73112
}
74113

75114
/** This split with its framework-stamped token attached, and its payload cleared.
@@ -84,6 +123,7 @@ public static ScanSplit of(byte[] payload, long rows, long bytes) {
84123
*/
85124
public ScanSplit withToken(byte[] stamped) {
86125
return new ScanSplit(new byte[0], stamped, estimated_rows, rows_exact, estimated_bytes,
87-
partition_bounds, column_statistics, location_ids, start_position, end_position);
126+
key_value, range_index, column_bounds, partition_bounds, column_statistics,
127+
location_ids, start_position, end_position);
88128
}
89129
}

0 commit comments

Comments
 (0)