Skip to content

Commit 5a5889a

Browse files
[Internal]Adding API to calculate EPK as byte-array (as alternative to HEX encoded string) (Azure#44209)
* [Internal]Adding API to calculate EPK as byte-array (as alternative to HEX encoded string) This new API is needed for Thin Client RNTBD encoding - adding this here in separate PR as preparation for ThinClient * Update PartitionKeyInternalTest.java * Update PartitionKeyInternal.java
1 parent aef0875 commit 5a5889a

File tree

3 files changed

+110
-19
lines changed

3 files changed

+110
-19
lines changed

sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/directconnectivity/PartitionKeyInternalTest.java

Lines changed: 18 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
package com.azure.cosmos.implementation.directconnectivity;
55

66
import com.azure.cosmos.implementation.Utils;
7+
import com.azure.cosmos.implementation.routing.HexConvert;
78
import com.azure.cosmos.models.PartitionKey;
89
import com.azure.cosmos.models.PartitionKeyDefinition;
910
import com.azure.cosmos.models.PartitionKeyDefinitionVersion;
@@ -288,18 +289,33 @@ public void hashEffectivePartitionKey() {
288289
assertThat(PartitionKeyInternalHelper.getEffectivePartitionKeyString(PartitionKeyInternal.InclusiveMinimum, new PartitionKeyDefinition()))
289290
.isEqualTo(PartitionKeyInternalHelper.MinimumInclusiveEffectivePartitionKey);
290291

292+
assertThat(PartitionKeyInternalHelper.getEffectivePartitionKeyBytes(PartitionKeyInternal.InclusiveMinimum, new PartitionKeyDefinition()))
293+
.isEqualTo(PartitionKeyInternalHelper.MinimumInclusiveEffectivePartitionKeyBytes);
294+
291295
assertThat(
292296
PartitionKeyInternalHelper.getEffectivePartitionKeyString(PartitionKeyInternal.ExclusiveMaximum, new PartitionKeyDefinition()))
293297
.isEqualTo(PartitionKeyInternalHelper.MaximumExclusiveEffectivePartitionKey);
294298

299+
assertThat(
300+
PartitionKeyInternalHelper.getEffectivePartitionKeyBytes(PartitionKeyInternal.ExclusiveMaximum, new PartitionKeyDefinition()))
301+
.isEqualTo(PartitionKeyInternalHelper.MaximumExclusiveEffectivePartitionKeyBytes);
302+
303+
295304
PartitionKeyDefinition partitionKeyDefinition = new PartitionKeyDefinition();
296305
partitionKeyDefinition.setPaths(Lists.newArrayList("/A", "/B", "/C", "/E", "/F", "/G"));
297306

298307
PartitionKeyInternal partitionKey = PartitionKeyInternal.fromObjectArray(
299308
new Object[]{2, true, false, null, Undefined.value(), "Привет!"}, true);
300-
String effectivePartitionKey = PartitionKeyInternalHelper.getEffectivePartitionKeyString(partitionKey, partitionKeyDefinition);
309+
String effectivePartitionKey = PartitionKeyInternalHelper
310+
.getEffectivePartitionKeyString(partitionKey, partitionKeyDefinition);
311+
312+
assertThat(effectivePartitionKey)
313+
.isEqualTo("05C1D19581B37C05C0000302010008D1A0D281D1B9D1B3D1B6D2832200");
301314

302-
assertThat(effectivePartitionKey).isEqualTo("05C1D19581B37C05C0000302010008D1A0D281D1B9D1B3D1B6D2832200");
315+
byte[] effectivePartitionKeyBytes = PartitionKeyInternalHelper
316+
.getEffectivePartitionKeyBytes(partitionKey, partitionKeyDefinition);
317+
assertThat(HexConvert.bytesToHex(effectivePartitionKeyBytes))
318+
.isEqualTo("05C1D19581B37C05C0000302010008D1A0D281D1B9D1B3D1B6D2832200");
303319
}
304320

305321
@DataProvider(name = "v2ParamProvider")

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/routing/PartitionKeyInternal.java

Lines changed: 4 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,6 @@
3030
import java.util.Arrays;
3131
import java.util.Iterator;
3232
import java.util.List;
33-
import java.util.stream.Collectors;
3433

3534
import static com.azure.cosmos.implementation.Utils.as;
3635

@@ -182,16 +181,6 @@ public boolean equals(Object obj) {
182181

183182
@Override
184183
public int hashCode() {
185-
// TODO: @kushagraThapar, @moderakh, mbhaskar to identify proper implementation.
186-
// Issue: https://github.com/Azure/azure-sdk-for-java/issues/9046
187-
// if (this.components == null || this.components.size() == 0) {
188-
// return 0;
189-
// }
190-
// int [] ordinals = new int[this.components.size()];
191-
// for (int i = 0; i < this.components.size(); i++) {
192-
// ordinals[i] = this.components.get(i).GetTypeOrdinal();
193-
// }
194-
// return Arrays.hashCode(ordinals);
195184
return super.hashCode();
196185
}
197186

@@ -250,6 +239,10 @@ public String getEffectivePartitionKeyString(PartitionKeyInternal internalPartit
250239
return PartitionKeyInternalHelper.getEffectivePartitionKeyString(internalPartitionKey, partitionKey);
251240
}
252241

242+
public byte[] getEffectivePartitionKeyBytes(PartitionKeyInternal internalPartitionKey, PartitionKeyDefinition partitionKey) {
243+
return PartitionKeyInternalHelper.getEffectivePartitionKeyBytes(internalPartitionKey, partitionKey);
244+
}
245+
253246
public Range<String> getEPKRangeForPrefixPartitionKey(PartitionKeyDefinition partitionKeyDefinition) {
254247
return PartitionKeyInternalHelper.getEPKRangeForPrefixPartitionKey(this, partitionKeyDefinition);
255248
}

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/routing/PartitionKeyInternalHelper.java

Lines changed: 88 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -17,9 +17,11 @@
1717
public class PartitionKeyInternalHelper {
1818

1919
public static final String MinimumInclusiveEffectivePartitionKey = toHexEncodedBinaryString(PartitionKeyInternal.EmptyPartitionKey.components);
20+
public static final byte[] MinimumInclusiveEffectivePartitionKeyBytes = toBinary(PartitionKeyInternal.EmptyPartitionKey.components);
2021
public static final String MaximumExclusiveEffectivePartitionKey = toHexEncodedBinaryString(PartitionKeyInternal.InfinityPartitionKey.components);
22+
public static final byte[] MaximumExclusiveEffectivePartitionKeyBytes = toBinary(PartitionKeyInternal.InfinityPartitionKey.components);
2123

22-
public static final Range<String> FullRange = new Range<String>(
24+
public static final Range<String> FullRange = new Range<>(
2325
PartitionKeyInternalHelper.MinimumInclusiveEffectivePartitionKey,
2426
PartitionKeyInternalHelper.MaximumExclusiveEffectivePartitionKey,
2527
true,
@@ -72,7 +74,29 @@ static String toHexEncodedBinaryString(List<IPartitionKeyComponent> components)
7274
return HexConvert.bytesToHex(stream.asByteBuffer());
7375
}
7476

77+
static byte[] toBinary(List<IPartitionKeyComponent> components) {
78+
ByteBufferOutputStream stream = new ByteBufferOutputStream(MaxPartitionKeyBinarySize);
79+
for (IPartitionKeyComponent component: components) {
80+
component.writeForBinaryEncoding(stream);
81+
}
82+
83+
return stream.toByteArray();
84+
}
85+
86+
static byte[] toBinary(IPartitionKeyComponent[] components) {
87+
ByteBufferOutputStream stream = new ByteBufferOutputStream(MaxPartitionKeyBinarySize);
88+
for (IPartitionKeyComponent component: components) {
89+
component.writeForBinaryEncoding(stream);
90+
}
91+
92+
return stream.toByteArray();
93+
}
94+
7595
static public String getEffectivePartitionKeyForHashPartitioningV2(PartitionKeyInternal partitionKeyInternal) {
96+
return HexConvert.bytesToHex(getEffectivePartitionKeyBytesForHashPartitioningV2(partitionKeyInternal));
97+
}
98+
99+
static public byte[] getEffectivePartitionKeyBytesForHashPartitioningV2(PartitionKeyInternal partitionKeyInternal) {
76100
try(ByteBufferOutputStream byteArrayBuffer = new ByteBufferOutputStream()) {
77101
for (int i = 0; i < partitionKeyInternal.components.size(); i++) {
78102
partitionKeyInternal.components.get(i).writeForHashingV2(byteArrayBuffer);
@@ -88,14 +112,18 @@ static public String getEffectivePartitionKeyForHashPartitioningV2(PartitionKeyI
88112
// Plus one more just in case.
89113
hash[0] &= 0x3F;
90114

91-
return HexConvert.bytesToHex(hash);
115+
return hash;
92116
} catch (IOException e) {
93117
throw new IllegalArgumentException(e);
94118
}
95119
}
96120

97121
static String getEffectivePartitionKeyForMultiHashPartitioning(PartitionKeyInternal partitionKeyInternal) {
98-
StringBuilder stringBuilder = new StringBuilder(partitionKeyInternal.components.size() * HASH_V2_EPK_LENGTH);
122+
return HexConvert.bytesToHex(getEffectivePartitionKeyBytesForMultiHashPartitioning(partitionKeyInternal));
123+
}
124+
125+
static byte[] getEffectivePartitionKeyBytesForMultiHashPartitioning(PartitionKeyInternal partitionKeyInternal) {
126+
byte[] finalHash = new byte[partitionKeyInternal.components.size() * 2 * Long.BYTES];
99127
for (int i = 0; i < partitionKeyInternal.components.size(); i++) {
100128
try(ByteBufferOutputStream byteArrayBuffer = new ByteBufferOutputStream()) {
101129
partitionKeyInternal.components.get(i).writeForHashingV2(byteArrayBuffer);
@@ -110,16 +138,22 @@ static String getEffectivePartitionKeyForMultiHashPartitioning(PartitionKeyInter
110138
// Plus one more just in case.
111139
hash[0] &= 0x3F;
112140

113-
stringBuilder.append(HexConvert.bytesToHex(hash));
141+
for (int n = 0; n < hash.length; n++) {
142+
finalHash[(i * 2 * Long.BYTES) + n] = hash[n];
143+
}
114144
} catch (IOException e) {
115145
throw new IllegalArgumentException(e);
116146
}
117147
}
118148

119-
return stringBuilder.toString();
149+
return finalHash;
120150
}
121151

122152
static String getEffectivePartitionKeyForHashPartitioning(PartitionKeyInternal partitionKeyInternal) {
153+
return HexConvert.bytesToHex(getEffectivePartitionKeyBytesForHashPartitioning(partitionKeyInternal));
154+
}
155+
156+
static byte[] getEffectivePartitionKeyBytesForHashPartitioning(PartitionKeyInternal partitionKeyInternal) {
123157
IPartitionKeyComponent[] truncatedComponents = new IPartitionKeyComponent[partitionKeyInternal.components.size()];
124158

125159
for (int i = 0; i < truncatedComponents.length; i++) {
@@ -145,7 +179,7 @@ static String getEffectivePartitionKeyForHashPartitioning(PartitionKeyInternal p
145179
partitionKeyComponents[i + 1] = truncatedComponents[i];
146180
}
147181

148-
return toHexEncodedBinaryString(partitionKeyComponents);
182+
return toBinary(partitionKeyComponents);
149183
}
150184

151185
public static String getEffectivePartitionKeyString(PartitionKeyInternal partitionKeyInternal, PartitionKeyDefinition partitionKeyDefinition) {
@@ -196,6 +230,54 @@ public static String getEffectivePartitionKeyString(PartitionKeyInternal partiti
196230
}
197231
}
198232

233+
public static byte[] getEffectivePartitionKeyBytes(PartitionKeyInternal partitionKeyInternal, PartitionKeyDefinition partitionKeyDefinition) {
234+
return getEffectivePartitionKeyBytes(partitionKeyInternal, partitionKeyDefinition, true);
235+
}
236+
237+
public static byte[] getEffectivePartitionKeyBytes(PartitionKeyInternal partitionKeyInternal, PartitionKeyDefinition partitionKeyDefinition, boolean strict) {
238+
if (partitionKeyInternal.components == null) {
239+
throw new IllegalArgumentException(RMResources.TooFewPartitionKeyComponents);
240+
}
241+
242+
if (partitionKeyInternal.equals(PartitionKeyInternal.EmptyPartitionKey)) {
243+
return MinimumInclusiveEffectivePartitionKeyBytes;
244+
}
245+
246+
if (partitionKeyInternal.equals(PartitionKeyInternal.InfinityPartitionKey)) {
247+
return MaximumExclusiveEffectivePartitionKeyBytes;
248+
}
249+
250+
if (partitionKeyInternal.components.size() < partitionKeyDefinition.getPaths().size() && partitionKeyDefinition.getKind() != PartitionKind.MULTI_HASH) {
251+
throw new IllegalArgumentException(RMResources.TooFewPartitionKeyComponents);
252+
}
253+
254+
if (partitionKeyInternal.components.size() > partitionKeyDefinition.getPaths().size() && strict) {
255+
throw new IllegalArgumentException(RMResources.TooManyPartitionKeyComponents);
256+
}
257+
258+
PartitionKind kind = partitionKeyDefinition.getKind();
259+
if (kind == null) {
260+
kind = PartitionKind.HASH;
261+
}
262+
263+
switch (kind) {
264+
case HASH:
265+
if (ModelBridgeInternal.isV2(partitionKeyDefinition)) {
266+
// V2
267+
return getEffectivePartitionKeyBytesForHashPartitioningV2(partitionKeyInternal);
268+
} else {
269+
// V1
270+
return getEffectivePartitionKeyBytesForHashPartitioning(partitionKeyInternal);
271+
}
272+
273+
case MULTI_HASH:
274+
return getEffectivePartitionKeyBytesForMultiHashPartitioning(partitionKeyInternal);
275+
276+
default:
277+
return toBinary(partitionKeyInternal.components);
278+
}
279+
}
280+
199281
static public Range<String> getEPKRangeForPrefixPartitionKey(
200282
PartitionKeyInternal internalPartitionKey,
201283
PartitionKeyDefinition partitionKeyDefinition)

0 commit comments

Comments
 (0)