Skip to content
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

Fix delta writer when relying on ID col config #136

Merged
merged 4 commits into from
Oct 25, 2023
Merged

Fix delta writer when relying on ID col config #136

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
05f7da2
Fix delta writer when relying on ID col config
bryanck Oct 24, 2023
d3145b7
upsert test
bryanck Oct 24, 2023
1a7cb26
upsert integration tests
bryanck Oct 24, 2023
9081837
better var name
bryanck Oct 25, 2023
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
41 changes: 41 additions & 0 deletions kafka-connect-runtime/src/test/java/io/tabular/iceberg/connect/IntegrationTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@
package io.tabular.iceberg.connect;

import static io.tabular.iceberg.connect.TestEvent.TEST_SCHEMA;
import static io.tabular.iceberg.connect.TestEvent.TEST_SCHEMA_NO_ID;
import static io.tabular.iceberg.connect.TestEvent.TEST_SPEC;
import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.fail;
Expand All @@ -29,6 +30,7 @@
import java.util.List;
import java.util.Map;
import org.apache.iceberg.DataFile;
import org.apache.iceberg.DeleteFile;
import org.apache.iceberg.PartitionSpec;
import org.apache.iceberg.Schema;
import org.apache.iceberg.Table;
Expand Down Expand Up @@ -158,6 +160,45 @@ public void testIcebergSinkAutoCreate(String branch) {
assertThat(spec.isPartitioned()).isEqualTo(useSchema);
}

@ParameterizedTest
@ValueSource(booleans = {true, false})
public void testIcebergSinkUnpartitionedUpsert(boolean hasSchemaIdCols) {
runUpsertTest(hasSchemaIdCols, PartitionSpec.unpartitioned());
}

@ParameterizedTest
@ValueSource(booleans = {true, false})
public void testIcebergSinkPartitionedUpsert(boolean hasSchemaIdCols) {
runUpsertTest(hasSchemaIdCols, TEST_SPEC);
}

private void runUpsertTest(boolean hasSchemaIdCols, PartitionSpec spec) {
Map<String, String> extraConfig = Maps.newHashMap();
extraConfig.put("iceberg.tables.upsert-mode-enabled", "true");
if (hasSchemaIdCols) {
catalog.createTable(
TABLE_IDENTIFIER, TEST_SCHEMA, spec, ImmutableMap.of("format-version", "2"));
} else {
extraConfig.put("iceberg.tables.default-id-columns", "id");
catalog.createTable(
TABLE_IDENTIFIER, TEST_SCHEMA_NO_ID, spec, ImmutableMap.of("format-version", "2"));
}

runTest(null, true, extraConfig);

List<DataFile> files = dataFiles(TABLE_IDENTIFIER, null);
// may involve 1 or 2 workers
assertThat(files).hasSizeBetween(1, 2);
assertThat(files.stream().mapToLong(DataFile::recordCount).sum()).isEqualTo(2);

// may involve 1 or 2 workers
List<DeleteFile> deleteFiles = deleteFiles(TABLE_IDENTIFIER, null);
assertThat(deleteFiles).hasSizeBetween(1, 2);
assertThat(deleteFiles.stream().mapToLong(DeleteFile::recordCount).sum()).isEqualTo(2);

assertSnapshotProps(TABLE_IDENTIFIER, null);
}

private void assertGeneratedSchema(boolean useSchema, Class<? extends Type> expectedIdType) {
Schema tableSchema = catalog.loadTable(TABLE_IDENTIFIER).schema();
assertThat(tableSchema.findField("id").type()).isInstanceOf(expectedIdType);
Expand Down
8 changes: 8 additions & 0 deletions kafka-connect-runtime/src/test/java/io/tabular/iceberg/connect/TestEvent.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,14 @@ public class TestEvent {
Types.NestedField.required(4, "payload", Types.StringType.get())),
ImmutableSet.of(1));

public static final Schema TEST_SCHEMA_NO_ID =
new Schema(
ImmutableList.of(
Types.NestedField.required(1, "id", Types.LongType.get()),
Types.NestedField.required(2, "type", Types.StringType.get()),
Types.NestedField.required(3, "ts", Types.TimestampType.withZone()),
Types.NestedField.required(4, "payload", Types.StringType.get())));

public static final org.apache.kafka.connect.data.Schema TEST_CONNECT_SCHEMA =
SchemaBuilder.struct()
.field("id", org.apache.kafka.connect.data.Schema.INT64_SCHEMA)
Expand Down
4 changes: 3 additions & 1 deletion kafka-connect/src/main/java/io/tabular/iceberg/connect/data/BaseDeltaTaskWriter.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@
package io.tabular.iceberg.connect.data;

import java.io.IOException;
import java.util.Set;
import org.apache.iceberg.FileFormat;
import org.apache.iceberg.PartitionKey;
import org.apache.iceberg.PartitionSpec;
Expand Down Expand Up @@ -50,10 +51,11 @@ abstract class BaseDeltaTaskWriter extends BaseTaskWriter<Record> {
FileIO io,
long targetFileSize,
Schema schema,
Set<Integer> identifierFieldIds,
boolean upsertMode) {
super(spec, format, appenderFactory, fileFactory, io, targetFileSize);
this.schema = schema;
this.deleteSchema = TypeUtil.select(schema, Sets.newHashSet(schema.identifierFieldIds()));
this.deleteSchema = TypeUtil.select(schema, Sets.newHashSet(identifierFieldIds));
this.wrapper = new InternalRecordWrapper(schema.asStruct());
this.keyWrapper = new InternalRecordWrapper(deleteSchema.asStruct());
this.keyProjection = RecordProjection.create(schema, deleteSchema);
Expand Down
13 changes: 12 additions & 1 deletion kafka-connect/src/main/java/io/tabular/iceberg/connect/data/PartitionedDeltaWriter.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@
import java.io.IOException;
import java.io.UncheckedIOException;
import java.util.Map;
import java.util.Set;
import org.apache.iceberg.FileFormat;
import org.apache.iceberg.PartitionKey;
import org.apache.iceberg.PartitionSpec;
Expand All @@ -45,8 +46,18 @@ public class PartitionedDeltaWriter extends BaseDeltaTaskWriter {
FileIO io,
long targetFileSize,
Schema schema,
Set<Integer> identifierFieldIds,
boolean upsertMode) {
super(spec, format, appenderFactory, fileFactory, io, targetFileSize, schema, upsertMode);
super(
spec,
format,
appenderFactory,
fileFactory,
io,
targetFileSize,
schema,
identifierFieldIds,
upsertMode);
this.partitionKey = new PartitionKey(spec, schema);
}

Expand Down
13 changes: 12 additions & 1 deletion kafka-connect/src/main/java/io/tabular/iceberg/connect/data/UnpartitionedDeltaWriter.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@
package io.tabular.iceberg.connect.data;

import java.io.IOException;
import java.util.Set;
import org.apache.iceberg.FileFormat;
import org.apache.iceberg.PartitionSpec;
import org.apache.iceberg.Schema;
Expand All @@ -38,8 +39,18 @@ public class UnpartitionedDeltaWriter extends BaseDeltaTaskWriter {
FileIO io,
long targetFileSize,
Schema schema,
Set<Integer> identifierFieldIds,
boolean upsertMode) {
super(spec, format, appenderFactory, fileFactory, io, targetFileSize, schema, upsertMode);
super(
spec,
format,
appenderFactory,
fileFactory,
io,
targetFileSize,
schema,
identifierFieldIds,
upsertMode);
this.writer = new RowDataDeltaWriter(null);
}

Expand Down
12 changes: 7 additions & 5 deletions kafka-connect/src/main/java/io/tabular/iceberg/connect/data/Utilities.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -168,19 +168,19 @@ public static TaskWriter<Record> createTableWriter(
PropertyUtil.propertyAsLong(
tableProps, WRITE_TARGET_FILE_SIZE_BYTES, WRITE_TARGET_FILE_SIZE_BYTES_DEFAULT);

Set<Integer> equalityFieldIds = table.schema().identifierFieldIds();
Set<Integer> identifierFieldIds = table.schema().identifierFieldIds();

// override the identifier fields if the config is set
List<String> idCols = config.tableConfig(tableName).idColumns();
if (!idCols.isEmpty()) {
equalityFieldIds =
identifierFieldIds =
idCols.stream()
.map(colName -> table.schema().findField(colName).fieldId())
.collect(toSet());
}

FileAppenderFactory<Record> appenderFactory;
if (equalityFieldIds == null || equalityFieldIds.isEmpty()) {
if (identifierFieldIds == null || identifierFieldIds.isEmpty()) {
appenderFactory =
new GenericAppenderFactory(table.schema(), table.spec(), null, null, null)
.setAll(tableProps);
Expand All @@ -189,8 +189,8 @@ public static TaskWriter<Record> createTableWriter(
new GenericAppenderFactory(
table.schema(),
table.spec(),
Ints.toArray(equalityFieldIds),
TypeUtil.select(table.schema(), Sets.newHashSet(equalityFieldIds)),
Ints.toArray(identifierFieldIds),
TypeUtil.select(table.schema(), Sets.newHashSet(identifierFieldIds)),
null)
.setAll(tableProps);
}
Expand Down Expand Up @@ -219,6 +219,7 @@ public static TaskWriter<Record> createTableWriter(
table.io(),
targetFileSize,
table.schema(),
identifierFieldIds,
config.upsertModeEnabled());
}
} else {
Expand All @@ -242,6 +243,7 @@ public static TaskWriter<Record> createTableWriter(
table.io(),
targetFileSize,
table.schema(),
identifierFieldIds,
config.upsertModeEnabled());
}
}
Expand Down