Skip to content

Commit 549c197

Browse files
jbonofrejbonofre
committed
Implement source-ids to deal with multi arguments transforms
1 parent 01fe380 commit 549c197

7 files changed

Lines changed: 167 additions & 29 deletions

File tree

‎.palantir/revapi.yml‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1177,6 +1177,10 @@ acceptedBreaks:
11771177
old: "class org.apache.iceberg.Metrics"
11781178
new: "class org.apache.iceberg.Metrics"
11791179
justification: "Java serialization across versions is not guaranteed"
1180+
- code: "java.class.defaultSerializationChanged"
1181+
old: "class org.apache.iceberg.PartitionField"
1182+
new: "class org.apache.iceberg.PartitionField"
1183+
justification: "Introduce source-ids"
11801184
org.apache.iceberg:iceberg-core:
11811185
- code: "java.method.removed"
11821186
old: "method java.lang.String[] org.apache.iceberg.hadoop.Util::blockLocations(org.apache.iceberg.CombinedScanTask,\

‎api/src/main/java/org/apache/iceberg/PartitionField.java‎

Lines changed: 37 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -19,26 +19,49 @@
1919
package org.apache.iceberg;
2020

2121
import java.io.Serializable;
22+
import java.util.List;
23+
import java.util.stream.Collectors;
2224
import org.apache.iceberg.relocated.com.google.common.base.Objects;
25+
import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
2326
import org.apache.iceberg.transforms.Transform;
2427

2528
/** Represents a single field in a {@link PartitionSpec}. */
2629
public class PartitionField implements Serializable {
27-
private final int sourceId;
30+
private final List<Integer> sourceIds;
2831
private final int fieldId;
2932
private final String name;
3033
private final Transform<?, ?> transform;
3134

35+
/**
36+
* @deprecated since 1.10.0, will be removed in 2.0.0; use PartitionField(sourceIds, fieldId,
37+
* name, transform) instead.
38+
*/
39+
@Deprecated
3240
PartitionField(int sourceId, int fieldId, String name, Transform<?, ?> transform) {
33-
this.sourceId = sourceId;
41+
this(List.of(sourceId), fieldId, name, transform);
42+
}
43+
44+
PartitionField(List<Integer> sourceIds, int fieldId, String name, Transform<?, ?> transform) {
45+
this.sourceIds = sourceIds;
3446
this.fieldId = fieldId;
3547
this.name = name;
3648
this.transform = transform;
3749
}
3850

39-
/** Returns the field id of the source field in the {@link PartitionSpec spec's} table schema. */
51+
/**
52+
* Returns the source field in the {@link PartitionSpec spec's} table schema.
53+
*
54+
* @deprecated since 1.10.0, will be removed in 2.0.0; use sourceIds() instead.
55+
*/
56+
@Deprecated
4057
public int sourceId() {
41-
return sourceId;
58+
Preconditions.checkArgument(sourceIds.size() == 1, "Only one source is supported");
59+
return sourceIds.get(0);
60+
}
61+
62+
/** Returns source ids field in the {@link PartitionSpec spec's} table schema. */
63+
public List<Integer> sourceIds() {
64+
return sourceIds;
4265
}
4366

4467
/** Returns the partition field id across all the table metadata's partition specs. */
@@ -58,7 +81,14 @@ public String name() {
5881

5982
@Override
6083
public String toString() {
61-
return fieldId + ": " + name + ": " + transform + "(" + sourceId + ")";
84+
return fieldId
85+
+ ": "
86+
+ name
87+
+ ": "
88+
+ transform
89+
+ "("
90+
+ sourceIds.stream().map(String::valueOf).collect(Collectors.joining(", "))
91+
+ ")";
6292
}
6393

6494
@Override
@@ -70,14 +100,14 @@ public boolean equals(Object other) {
70100
}
71101

72102
PartitionField that = (PartitionField) other;
73-
return sourceId == that.sourceId
103+
return sourceIds.equals(that.sourceIds)
74104
&& fieldId == that.fieldId
75105
&& name.equals(that.name)
76106
&& transform.toString().equals(that.transform.toString());
77107
}
78108

79109
@Override
80110
public int hashCode() {
81-
return Objects.hashCode(sourceId, fieldId, name, transform);
111+
return Objects.hashCode(sourceIds, fieldId, name, transform);
82112
}
83113
}

‎api/src/main/java/org/apache/iceberg/PartitionSpec.java‎

Lines changed: 15 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -105,7 +105,7 @@ public UnboundPartitionSpec toUnbound() {
105105

106106
for (PartitionField field : fields) {
107107
builder.addField(
108-
field.transform().toString(), field.sourceId(), field.fieldId(), field.name());
108+
field.transform().toString(), field.sourceIds(), field.fieldId(), field.name());
109109
}
110110

111111
return builder.build();
@@ -601,17 +601,26 @@ public Builder alwaysNull(String sourceName) {
601601

602602
// add a partition field with an auto-increment partition field id starting from
603603
// PARTITION_DATA_ID_START
604+
Builder add(List<Integer> sourceIds, String name, Transform<?, ?> transform) {
605+
return add(sourceIds, nextFieldId(), name, transform);
606+
}
607+
604608
Builder add(int sourceId, String name, Transform<?, ?> transform) {
605-
return add(sourceId, nextFieldId(), name, transform);
609+
return add(List.of(sourceId), name, transform);
606610
}
607611

608-
Builder add(int sourceId, int fieldId, String name, Transform<?, ?> transform) {
609-
checkAndAddPartitionName(name, sourceId);
610-
fields.add(new PartitionField(sourceId, fieldId, name, transform));
612+
Builder add(List<Integer> sourceIds, int fieldId, String name, Transform<?, ?> transform) {
613+
// we use the first entry in the source-ids list here
614+
checkAndAddPartitionName(name, sourceIds.get(0));
615+
fields.add(new PartitionField(sourceIds, fieldId, name, transform));
611616
lastAssignedFieldId.getAndAccumulate(fieldId, Math::max);
612617
return this;
613618
}
614619

620+
Builder add(int sourceId, int fieldId, String name, Transform<?, ?> transform) {
621+
return add(List.of(sourceId), fieldId, name, transform);
622+
}
623+
615624
public PartitionSpec build() {
616625
PartitionSpec spec = buildUnchecked();
617626
checkCompatibility(spec, schema);
@@ -625,7 +634,7 @@ PartitionSpec buildUnchecked() {
625634

626635
static void checkCompatibility(PartitionSpec spec, Schema schema) {
627636
for (PartitionField field : spec.fields) {
628-
Type sourceType = schema.findType(field.sourceId());
637+
Type sourceType = schema.findType(field.sourceIds().get(0));
629638
Transform<?, ?> transform = field.transform();
630639
// In the case of a Version 1 partition-spec field gets deleted,
631640
// it is replaced with a void transform, see:

‎api/src/main/java/org/apache/iceberg/UnboundPartitionSpec.java‎

Lines changed: 25 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
package org.apache.iceberg;
2020

2121
import java.util.List;
22+
import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
2223
import org.apache.iceberg.relocated.com.google.common.collect.Lists;
2324
import org.apache.iceberg.transforms.Transform;
2425
import org.apache.iceberg.transforms.Transforms;
@@ -54,17 +55,17 @@ private PartitionSpec.Builder copyToBuilder(Schema schema) {
5455
PartitionSpec.Builder builder = PartitionSpec.builderFor(schema).withSpecId(specId);
5556

5657
for (UnboundPartitionField field : fields) {
57-
Type fieldType = schema.findType(field.sourceId);
58+
Type fieldType = schema.findType(field.sourceIds.get(0));
5859
Transform<?, ?> transform;
5960
if (fieldType != null) {
6061
transform = Transforms.fromString(fieldType, field.transform.toString());
6162
} else {
6263
transform = Transforms.fromString(field.transform.toString());
6364
}
6465
if (field.partitionId != null) {
65-
builder.add(field.sourceId, field.partitionId, field.name, transform);
66+
builder.add(field.sourceIds, field.partitionId, field.name, transform);
6667
} else {
67-
builder.add(field.sourceId, field.name, transform);
68+
builder.add(field.sourceIds, field.name, transform);
6869
}
6970
}
7071

@@ -88,14 +89,23 @@ Builder withSpecId(int newSpecId) {
8889
return this;
8990
}
9091

92+
Builder addField(
93+
String transformAsString, List<Integer> sourceIds, int partitionId, String name) {
94+
fields.add(new UnboundPartitionField(transformAsString, sourceIds, partitionId, name));
95+
return this;
96+
}
97+
9198
Builder addField(String transformAsString, int sourceId, int partitionId, String name) {
92-
fields.add(new UnboundPartitionField(transformAsString, sourceId, partitionId, name));
99+
return addField(transformAsString, List.of(sourceId), partitionId, name);
100+
}
101+
102+
Builder addField(String transformAsString, List<Integer> sourceIds, String name) {
103+
fields.add(new UnboundPartitionField(transformAsString, sourceIds, null, name));
93104
return this;
94105
}
95106

96107
Builder addField(String transformAsString, int sourceId, String name) {
97-
fields.add(new UnboundPartitionField(transformAsString, sourceId, null, name));
98-
return this;
108+
return addField(transformAsString, List.of(sourceId), name);
99109
}
100110

101111
UnboundPartitionSpec build() {
@@ -105,7 +115,7 @@ UnboundPartitionSpec build() {
105115

106116
static class UnboundPartitionField {
107117
private final Transform<?, ?> transform;
108-
private final int sourceId;
118+
private final List<Integer> sourceIds;
109119
private final Integer partitionId;
110120
private final String name;
111121

@@ -118,7 +128,12 @@ public String transformAsString() {
118128
}
119129

120130
public int sourceId() {
121-
return sourceId;
131+
Preconditions.checkArgument(sourceIds.size() == 1, "Only one source is expected");
132+
return sourceIds.get(0);
133+
}
134+
135+
public List<Integer> sourceIds() {
136+
return sourceIds;
122137
}
123138

124139
public Integer partitionId() {
@@ -130,9 +145,9 @@ public String name() {
130145
}
131146

132147
private UnboundPartitionField(
133-
String transformAsString, int sourceId, Integer partitionId, String name) {
148+
String transformAsString, List<Integer> sourceIds, Integer partitionId, String name) {
134149
this.transform = Transforms.fromString(transformAsString);
135-
this.sourceId = sourceId;
150+
this.sourceIds = sourceIds;
136151
this.partitionId = partitionId;
137152
this.name = name;
138153
}

‎api/src/test/java/org/apache/iceberg/TestHelpers.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -746,12 +746,12 @@ public ExpectedSpecBuilder withSpecId(int newSpecId) {
746746

747747
public ExpectedSpecBuilder addField(
748748
String transformAsString, int sourceId, int partitionId, String name) {
749-
unboundPartitionSpecBuilder.addField(transformAsString, sourceId, partitionId, name);
749+
unboundPartitionSpecBuilder.addField(transformAsString, List.of(sourceId), partitionId, name);
750750
return this;
751751
}
752752

753753
public ExpectedSpecBuilder addField(String transformAsString, int sourceId, String name) {
754-
unboundPartitionSpecBuilder.addField(transformAsString, sourceId, name);
754+
unboundPartitionSpecBuilder.addField(transformAsString, List.of(sourceId), name);
755755
return this;
756756
}
757757

‎core/src/main/java/org/apache/iceberg/PartitionSpecParser.java‎

Lines changed: 33 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,9 @@
2424
import com.github.benmanes.caffeine.cache.Caffeine;
2525
import java.io.IOException;
2626
import java.util.Iterator;
27+
import java.util.List;
2728
import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
29+
import org.apache.iceberg.relocated.com.google.common.collect.Lists;
2830
import org.apache.iceberg.types.Types;
2931
import org.apache.iceberg.util.JsonUtil;
3032
import org.apache.iceberg.util.Pair;
@@ -35,6 +37,7 @@ private PartitionSpecParser() {}
3537
private static final String SPEC_ID = "spec-id";
3638
private static final String FIELDS = "fields";
3739
private static final String SOURCE_ID = "source-id";
40+
private static final String SOURCE_IDS = "source-ids";
3841
private static final String FIELD_ID = "field-id";
3942
private static final String TRANSFORM = "transform";
4043
private static final String NAME = "name";
@@ -98,10 +101,22 @@ static void toJsonFields(UnboundPartitionSpec spec, JsonGenerator generator) thr
98101
generator.writeStartObject();
99102
generator.writeStringField(NAME, field.name());
100103
generator.writeStringField(TRANSFORM, field.transformAsString());
101-
generator.writeNumberField(SOURCE_ID, field.sourceId());
104+
if (field.sourceIds().size() == 1) {
105+
generator.writeNumberField(SOURCE_ID, field.sourceId());
106+
} else {
107+
generator.writeFieldName(SOURCE_IDS);
108+
generator.writeStartArray();
109+
for (Integer value : field.sourceIds()) {
110+
generator.writeNumber(value);
111+
}
112+
113+
generator.writeEndArray();
114+
}
115+
102116
generator.writeNumberField(FIELD_ID, field.partitionId());
103117
generator.writeEndObject();
104118
}
119+
105120
generator.writeEndArray();
106121
}
107122

@@ -132,15 +147,29 @@ private static void buildFromJsonFields(UnboundPartitionSpec.Builder builder, Js
132147

133148
String name = JsonUtil.getString(NAME, element);
134149
String transform = JsonUtil.getString(TRANSFORM, element);
135-
int sourceId = JsonUtil.getInt(SOURCE_ID, element);
150+
Preconditions.checkArgument(
151+
element.has(SOURCE_ID) || element.has(SOURCE_IDS),
152+
"Cannot parse partition field, either "
153+
+ SOURCE_ID
154+
+ " or "
155+
+ SOURCE_IDS
156+
+ " has to be present");
157+
List<Integer> sourceIds;
158+
Integer sourceId = JsonUtil.getIntOrNull(SOURCE_ID, element);
159+
if (sourceId != null) {
160+
sourceIds = Lists.newArrayList();
161+
sourceIds.add(sourceId);
162+
} else {
163+
sourceIds = JsonUtil.getIntegerList(SOURCE_IDS, element);
164+
}
136165

137166
// partition field ids are missing in old PartitionSpec, they always auto-increment from
138167
// PARTITION_DATA_ID_START
139168
if (element.has(FIELD_ID)) {
140-
builder.addField(transform, sourceId, JsonUtil.getInt(FIELD_ID, element), name);
169+
builder.addField(transform, sourceIds, JsonUtil.getInt(FIELD_ID, element), name);
141170
fieldIdCount++;
142171
} else {
143-
builder.addField(transform, sourceId, name);
172+
builder.addField(transform, sourceIds, name);
144173
}
145174
}
146175

‎core/src/test/java/org/apache/iceberg/TestPartitionSpecParser.java‎

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
package org.apache.iceberg;
2020

2121
import static org.assertj.core.api.Assertions.assertThat;
22+
import static org.assertj.core.api.Assertions.assertThatThrownBy;
2223

2324
import java.util.Arrays;
2425
import java.util.List;
@@ -95,6 +96,56 @@ public void testFromJsonWithFieldId() {
9596
assertThat(spec.fields().get(1).fieldId()).isEqualTo(1000);
9697
}
9798

99+
@TestTemplate
100+
public void testFromJsonWithSourceIds() {
101+
String specString =
102+
"{\n"
103+
+ " \"spec-id\" : 1,\n"
104+
+ " \"fields\" : [ {\n"
105+
+ " \"name\" : \"id_bucket\",\n"
106+
+ " \"transform\" : \"bucket[8]\",\n"
107+
+ " \"source-ids\" : [ 1, 2 ],\n"
108+
+ " \"field-id\" : 1001\n"
109+
+ " }, {\n"
110+
+ " \"name\" : \"data_bucket\",\n"
111+
+ " \"transform\" : \"bucket[16]\",\n"
112+
+ " \"source-ids\" : [ 2, 3 ],\n"
113+
+ " \"field-id\" : 1000\n"
114+
+ " } ]\n"
115+
+ "}";
116+
117+
PartitionSpec spec = PartitionSpecParser.fromJson(table.schema(), specString);
118+
119+
assertThat(spec.fields()).hasSize(2);
120+
assertThat(spec.fields().get(0).sourceIds()).hasSize(2);
121+
assertThat(spec.fields().get(1).sourceIds()).hasSize(2);
122+
}
123+
124+
@TestTemplate
125+
public void testFromJsonWithoutSourceIdAndSourceIds() {
126+
String specString =
127+
"{\n"
128+
+ " \"spec-id\" : 1,\n"
129+
+ " \"fields\" : [ {\n"
130+
+ " \"name\" : \"id_bucket\",\n"
131+
+ " \"transform\" : \"bucket[8]\",\n"
132+
+ " \"field-id\" : 1001\n"
133+
+ " }, {\n"
134+
+ " \"name\" : \"data_bucket\",\n"
135+
+ " \"transform\" : \"bucket[16]\",\n"
136+
+ " \"field-id\" : 1000\n"
137+
+ " } ]\n"
138+
+ "}";
139+
140+
assertThatThrownBy(
141+
() -> {
142+
PartitionSpecParser.fromJson(table.schema(), specString);
143+
})
144+
.isInstanceOf(IllegalArgumentException.class)
145+
.hasMessage(
146+
"Cannot parse partition field, either source-id or source-ids has to be present");
147+
}
148+
98149
@TestTemplate
99150
public void testFromJsonWithoutFieldId() {
100151
String specString =

0 commit comments

Comments
 (0)