This is an automated email from the ASF dual-hosted git repository. dockerzhang pushed a commit to branch branch-1.7 in repository https://gitbox.apache.org/repos/asf/inlong.git
commit a1ce157a122a8d3eb43d867e35e5ccb35a75fa00 Author: thexia <[email protected]> AuthorDate: Fri May 12 16:53:36 2023 +0800 [INLONG-7830][Sort] Using multi-threading to closing files ingesting data into iceberg (#8016) --- .../sink/GroupedPartitionedDeltaWriter.java | 64 +++- .../sink/GroupedPartitionedFanoutWriter.java | 69 +++- .../iceberg/sink/trick/BaseDeltaTaskWriter.java | 121 ++++++ .../sort/iceberg/sink/trick/BaseTaskWriter.java | 414 +++++++++++++++++++++ .../sort/iceberg/sink/trick/CharSequenceSet.java | 189 ++++++++++ .../iceberg/sink/trick/SortedPosDeleteWriter.java | 220 +++++++++++ .../inlong/sort/iceberg/sink/trick/StructCopy.java | 59 +++ 7 files changed, 1115 insertions(+), 21 deletions(-) diff --git a/inlong-sort/sort-connectors/iceberg/src/main/java/org/apache/inlong/sort/iceberg/sink/GroupedPartitionedDeltaWriter.java b/inlong-sort/sort-connectors/iceberg/src/main/java/org/apache/inlong/sort/iceberg/sink/GroupedPartitionedDeltaWriter.java index c72fc2d43..a27244b1a 100644 --- a/inlong-sort/sort-connectors/iceberg/src/main/java/org/apache/inlong/sort/iceberg/sink/GroupedPartitionedDeltaWriter.java +++ b/inlong-sort/sort-connectors/iceberg/src/main/java/org/apache/inlong/sort/iceberg/sink/GroupedPartitionedDeltaWriter.java @@ -17,6 +17,7 @@ package org.apache.inlong.sort.iceberg.sink; +import org.apache.flink.shaded.guava18.com.google.common.util.concurrent.ThreadFactoryBuilder; import org.apache.flink.table.data.RowData; import org.apache.flink.table.types.logical.RowType; import org.apache.iceberg.FileFormat; @@ -30,18 +31,39 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; import java.io.IOException; +import java.util.ArrayList; import java.util.List; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Future; +import java.util.concurrent.LinkedBlockingQueue; +import java.util.concurrent.ThreadPoolExecutor; +import java.util.concurrent.ThreadPoolExecutor.CallerRunsPolicy; +import java.util.concurrent.TimeUnit; -public class GroupedPartitionedDeltaWriter extends BaseDeltaTaskWriter { +public class GroupedPartitionedDeltaWriter extends org.apache.inlong.sort.iceberg.sink.trick.BaseDeltaTaskWriter { private static final Logger LOG = LoggerFactory.getLogger(GroupedPartitionedDeltaWriter.class); + private static final ExecutorService CLOSE_EXECUTOR_SERVICE = new ThreadPoolExecutor( + 10, + 20, + 100L, + TimeUnit.MILLISECONDS, + new LinkedBlockingQueue<>(100), + new ThreadFactoryBuilder().setNameFormat("iceberg-writer-close-thread-%s").build(), + new CallerRunsPolicy()); + private final PartitionKey partitionKey; private String latestPartitionPath; private RowDataDeltaWriter latestWriter; + private final List<Future> futures = new ArrayList<>(); + + private final List<Exception> failures = new ArrayList<>(); + GroupedPartitionedDeltaWriter(PartitionSpec spec, FileFormat format, FileAppenderFactory<RowData> appenderFactory, @@ -59,6 +81,7 @@ public class GroupedPartitionedDeltaWriter extends BaseDeltaTaskWriter { @Override public RowDataDeltaWriter route(RowData row) { + checkFailure(); partitionKey.partition(wrapper().wrap(row)); if (latestPartitionPath != null && partitionKey.toPath().equals(latestPartitionPath)) { return latestWriter; @@ -72,19 +95,40 @@ public class GroupedPartitionedDeltaWriter extends BaseDeltaTaskWriter { @Override public void close() { - closeCurrentWriter(); - latestWriter = null; - latestPartitionPath = null; + try { + closeCurrentWriter(); + for (Future future : futures) { + future.get(); + } + checkFailure(); + latestWriter = null; + } catch (InterruptedException | ExecutionException e) { + LOG.warn("Interrupted while waiting for tasks to finish", e); + throw new RuntimeException(e); + } } private void closeCurrentWriter() { if (latestWriter != null) { - try { - latestWriter.close(); - } catch (IOException e) { - LOG.error("Exception occur when closing file {}.", latestPartitionPath); - throw new RuntimeException(e); - } + LOG.debug("Start eliminated writer for partition {}", latestPartitionPath); + final RowDataDeltaWriter writer = latestWriter; + final String partitionPath = latestPartitionPath; + futures.add(CLOSE_EXECUTOR_SERVICE.submit(() -> { + try { + writer.close(); + LOG.debug("End eliminated writer for partition {}", partitionPath); + } catch (IOException | RuntimeException e) { + failures.add(e); + } + })); + } + } + + private void checkFailure() { + if (!failures.isEmpty()) { + throw new RuntimeException( + String.format("Failed to close equality delta writer, %d unclosed files more.", failures.size()), + failures.get(0)); } } } diff --git a/inlong-sort/sort-connectors/iceberg/src/main/java/org/apache/inlong/sort/iceberg/sink/GroupedPartitionedFanoutWriter.java b/inlong-sort/sort-connectors/iceberg/src/main/java/org/apache/inlong/sort/iceberg/sink/GroupedPartitionedFanoutWriter.java index ff3f0bf81..411831e44 100644 --- a/inlong-sort/sort-connectors/iceberg/src/main/java/org/apache/inlong/sort/iceberg/sink/GroupedPartitionedFanoutWriter.java +++ b/inlong-sort/sort-connectors/iceberg/src/main/java/org/apache/inlong/sort/iceberg/sink/GroupedPartitionedFanoutWriter.java @@ -17,10 +17,10 @@ package org.apache.inlong.sort.iceberg.sink; +import org.apache.flink.shaded.guava18.com.google.common.util.concurrent.ThreadFactoryBuilder; import org.apache.iceberg.FileFormat; import org.apache.iceberg.PartitionKey; import org.apache.iceberg.PartitionSpec; -import org.apache.iceberg.io.BaseTaskWriter; import org.apache.iceberg.io.FileAppenderFactory; import org.apache.iceberg.io.FileIO; import org.apache.iceberg.io.OutputFileFactory; @@ -28,15 +28,39 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; import java.io.IOException; +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Future; +import java.util.concurrent.LinkedBlockingQueue; +import java.util.concurrent.ThreadPoolExecutor; +import java.util.concurrent.ThreadPoolExecutor.CallerRunsPolicy; +import java.util.concurrent.TimeUnit; -public abstract class GroupedPartitionedFanoutWriter<T> extends BaseTaskWriter<T> { +public abstract class GroupedPartitionedFanoutWriter<T> + extends + org.apache.inlong.sort.iceberg.sink.trick.BaseTaskWriter<T> { private static final Logger LOG = LoggerFactory.getLogger(GroupedPartitionedFanoutWriter.class); + private static final ExecutorService CLOSE_EXECUTOR_SERVICE = new ThreadPoolExecutor( + 10, + 20, + 100L, + TimeUnit.MILLISECONDS, + new LinkedBlockingQueue<>(100), + new ThreadFactoryBuilder().setNameFormat("iceberg-writer-close-thread-%s").build(), + new CallerRunsPolicy()); + private String latestPartitionPath; private RollingFileWriter latestWriter; + private final List<Future> futures = new ArrayList<>(); + + private final List<Exception> failures = new ArrayList<>(); + protected GroupedPartitionedFanoutWriter(PartitionSpec spec, FileFormat format, FileAppenderFactory<T> appenderFactory, OutputFileFactory fileFactory, FileIO io, long targetFileSize) { @@ -54,6 +78,7 @@ public abstract class GroupedPartitionedFanoutWriter<T> extends BaseTaskWriter<T @Override public void write(T row) throws IOException { + checkFailure(); PartitionKey partitionKey = partition(row); if (latestPartitionPath == null || !partitionKey.toPath().equals(latestPartitionPath)) { // NOTICE: we need to copy a new partition key here, in case of messing up the keys in writers. @@ -67,19 +92,41 @@ public abstract class GroupedPartitionedFanoutWriter<T> extends BaseTaskWriter<T @Override public void close() throws IOException { - closeCurrentWriter(); - latestWriter = null; - latestPartitionPath = null; + try { + closeCurrentWriter(); + for (Future future : futures) { + future.get(); + } + checkFailure(); + latestWriter = null; + latestPartitionPath = null; + } catch (InterruptedException | ExecutionException e) { + LOG.warn("Interrupted while waiting for tasks to finish", e); + throw new RuntimeException(e); + } } private void closeCurrentWriter() { if (latestWriter != null) { - try { - latestWriter.close(); - } catch (IOException e) { - LOG.error("Exception occur when closing file {}.", latestPartitionPath); - throw new RuntimeException(e); - } + LOG.info("Start eliminated writer for partition {}", latestPartitionPath); + final RollingFileWriter writer = latestWriter; + final String partitionPath = latestPartitionPath; + futures.add(CLOSE_EXECUTOR_SERVICE.submit(() -> { + try { + writer.close(); + LOG.info("End eliminated writer for partition {}", partitionPath); + } catch (IOException | RuntimeException e) { + failures.add(e); + } + })); + } + } + + private void checkFailure() { + if (!failures.isEmpty()) { + throw new RuntimeException( + String.format("Failed to close equality delta writer, %d unclosed files more.", failures.size()), + failures.get(0)); } } } diff --git a/inlong-sort/sort-connectors/iceberg/src/main/java/org/apache/inlong/sort/iceberg/sink/trick/BaseDeltaTaskWriter.java b/inlong-sort/sort-connectors/iceberg/src/main/java/org/apache/inlong/sort/iceberg/sink/trick/BaseDeltaTaskWriter.java new file mode 100644 index 000000000..c070adf73 --- /dev/null +++ b/inlong-sort/sort-connectors/iceberg/src/main/java/org/apache/inlong/sort/iceberg/sink/trick/BaseDeltaTaskWriter.java @@ -0,0 +1,121 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.inlong.sort.iceberg.sink.trick; + +import org.apache.flink.table.data.RowData; +import org.apache.flink.table.types.logical.RowType; +import org.apache.iceberg.FileFormat; +import org.apache.iceberg.PartitionKey; +import org.apache.iceberg.PartitionSpec; +import org.apache.iceberg.Schema; +import org.apache.iceberg.StructLike; +import org.apache.iceberg.flink.FlinkSchemaUtil; +import org.apache.iceberg.flink.RowDataWrapper; +import org.apache.iceberg.flink.data.RowDataProjection; +import org.apache.iceberg.io.FileAppenderFactory; +import org.apache.iceberg.io.FileIO; +import org.apache.iceberg.io.OutputFileFactory; +import org.apache.iceberg.relocated.com.google.common.collect.Sets; +import org.apache.iceberg.types.TypeUtil; + +import java.io.IOException; +import java.util.List; + +public abstract class BaseDeltaTaskWriter extends BaseTaskWriter<RowData> { + + private final Schema schema; + private final Schema deleteSchema; + private final RowDataWrapper wrapper; + private final RowDataWrapper keyWrapper; + private final RowDataProjection keyProjection; + private final boolean upsert; + + public BaseDeltaTaskWriter(PartitionSpec spec, + FileFormat format, + FileAppenderFactory<RowData> appenderFactory, + OutputFileFactory fileFactory, + FileIO io, + long targetFileSize, + Schema schema, + RowType flinkSchema, + List<Integer> equalityFieldIds, + boolean upsert) { + super(spec, format, appenderFactory, fileFactory, io, targetFileSize); + this.schema = schema; + this.deleteSchema = TypeUtil.select(schema, Sets.newHashSet(equalityFieldIds)); + this.wrapper = new RowDataWrapper(flinkSchema, schema.asStruct()); + this.upsert = upsert; + this.keyWrapper = new RowDataWrapper(FlinkSchemaUtil.convert(deleteSchema), deleteSchema.asStruct()); + this.keyProjection = RowDataProjection.create(schema, deleteSchema); + } + + public abstract BaseDeltaTaskWriter.RowDataDeltaWriter route(RowData row); + + public RowDataWrapper wrapper() { + return wrapper; + } + + @Override + public void write(RowData row) throws IOException { + BaseDeltaTaskWriter.RowDataDeltaWriter writer = route(row); + + switch (row.getRowKind()) { + case INSERT: + case UPDATE_AFTER: + if (upsert) { + writer.deleteKey(keyProjection.wrap(row)); + } + writer.write(row); + break; + + case UPDATE_BEFORE: + if (upsert) { + break; // UPDATE_BEFORE is not necessary for UPDATE, we do nothing to prevent delete one row twice + } + writer.delete(row); + break; + case DELETE: + if (upsert) { // https://github.com/apache/iceberg/pull/6753/files + writer.deleteKey(keyProjection.wrap(row)); + } else { + writer.delete(row); + } + break; + + default: + throw new UnsupportedOperationException("Unknown row kind: " + row.getRowKind()); + } + } + + protected class RowDataDeltaWriter extends BaseEqualityDeltaWriter { + + public RowDataDeltaWriter(PartitionKey partition) { + super(partition, schema, deleteSchema); + } + + @Override + protected StructLike asStructLike(RowData data) { + return wrapper.wrap(data); + } + + @Override + protected StructLike asStructLikeKey(RowData data) { + return keyWrapper.wrap(data); + } + } +} diff --git a/inlong-sort/sort-connectors/iceberg/src/main/java/org/apache/inlong/sort/iceberg/sink/trick/BaseTaskWriter.java b/inlong-sort/sort-connectors/iceberg/src/main/java/org/apache/inlong/sort/iceberg/sink/trick/BaseTaskWriter.java new file mode 100644 index 000000000..9772fca0b --- /dev/null +++ b/inlong-sort/sort-connectors/iceberg/src/main/java/org/apache/inlong/sort/iceberg/sink/trick/BaseTaskWriter.java @@ -0,0 +1,414 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.inlong.sort.iceberg.sink.trick; + +import java.io.Closeable; +import java.io.IOException; +import java.io.UncheckedIOException; +import java.util.Collections; +import java.util.List; +import java.util.Map; +import org.apache.iceberg.DataFile; +import org.apache.iceberg.DeleteFile; +import org.apache.iceberg.FileFormat; +import org.apache.iceberg.PartitionSpec; +import org.apache.iceberg.Schema; +import org.apache.iceberg.StructLike; +import org.apache.iceberg.deletes.EqualityDeleteWriter; +import org.apache.iceberg.encryption.EncryptedOutputFile; +import org.apache.iceberg.io.DataWriter; +import org.apache.iceberg.io.FileAppenderFactory; +import org.apache.iceberg.io.FileIO; +import org.apache.iceberg.io.OutputFileFactory; +import org.apache.iceberg.io.TaskWriter; +import org.apache.iceberg.io.WriteResult; +import org.apache.iceberg.relocated.com.google.common.base.MoreObjects; +import org.apache.iceberg.relocated.com.google.common.base.Preconditions; +import org.apache.iceberg.relocated.com.google.common.collect.Iterables; +import org.apache.iceberg.relocated.com.google.common.collect.Lists; +import org.apache.iceberg.util.StructLikeMap; +import org.apache.iceberg.util.StructProjection; +import org.apache.iceberg.util.Tasks; +import org.apache.iceberg.util.ThreadPools; + +/** + * Copied from iceberg 1.1.x. Modify List to concurrent List. + * @param <T> + */ +public abstract class BaseTaskWriter<T> implements TaskWriter<T> { + + private final List<DataFile> completedDataFiles = Collections.synchronizedList(Lists.newArrayList()); + private final List<DeleteFile> completedDeleteFiles = Collections.synchronizedList(Lists.newArrayList()); + private final CharSequenceSet referencedDataFiles = CharSequenceSet.empty(); + + private final PartitionSpec spec; + private final FileFormat format; + private final FileAppenderFactory<T> appenderFactory; + private final OutputFileFactory fileFactory; + private final FileIO io; + private final long targetFileSize; + private Throwable failure; + + protected BaseTaskWriter( + PartitionSpec spec, + FileFormat format, + FileAppenderFactory<T> appenderFactory, + OutputFileFactory fileFactory, + FileIO io, + long targetFileSize) { + this.spec = spec; + this.format = format; + this.appenderFactory = appenderFactory; + this.fileFactory = fileFactory; + this.io = io; + this.targetFileSize = targetFileSize; + } + + protected PartitionSpec spec() { + return spec; + } + + protected void setFailure(Throwable throwable) { + if (failure == null) { + this.failure = throwable; + } + } + + @Override + public void abort() throws IOException { + close(); + + // clean up files created by this writer + Tasks.foreach(Iterables.concat(completedDataFiles, completedDeleteFiles)) + .executeWith(ThreadPools.getWorkerPool()) + .throwFailureWhenFinished() + .noRetry() + .run(file -> io.deleteFile(file.path().toString())); + } + + @Override + public WriteResult complete() throws IOException { + close(); + + Preconditions.checkState(failure == null, "Cannot return results from failed writer", failure); + + return WriteResult.builder() + .addDataFiles(completedDataFiles) + .addDeleteFiles(completedDeleteFiles) + .addReferencedDataFiles(referencedDataFiles) + .build(); + } + + /** Base equality delta writer to write both insert records and equality-deletes. */ + protected abstract class BaseEqualityDeltaWriter implements Closeable { + + private final StructProjection structProjection; + private BaseTaskWriter.RollingFileWriter dataWriter; + private BaseTaskWriter.RollingEqDeleteWriter eqDeleteWriter; + private SortedPosDeleteWriter<T> posDeleteWriter; + private Map<StructLike, BaseTaskWriter.PathOffset> insertedRowMap; + + protected BaseEqualityDeltaWriter(StructLike partition, Schema schema, Schema deleteSchema) { + Preconditions.checkNotNull(schema, "Iceberg table schema cannot be null."); + Preconditions.checkNotNull(deleteSchema, "Equality-delete schema cannot be null."); + this.structProjection = StructProjection.create(schema, deleteSchema); + + this.dataWriter = new BaseTaskWriter.RollingFileWriter(partition); + this.eqDeleteWriter = new BaseTaskWriter.RollingEqDeleteWriter(partition); + this.posDeleteWriter = + new SortedPosDeleteWriter<>(appenderFactory, fileFactory, format, partition); + this.insertedRowMap = StructLikeMap.create(deleteSchema.asStruct()); + } + + /** Wrap the data as a {@link StructLike}. */ + protected abstract StructLike asStructLike(T data); + + /** Wrap the passed in key of a row as a {@link StructLike} */ + protected abstract StructLike asStructLikeKey(T key); + + public void write(T row) throws IOException { + BaseTaskWriter.PathOffset pathOffset = + BaseTaskWriter.PathOffset.of(dataWriter.currentPath(), dataWriter.currentRows()); + + // Create a copied key from this row. + StructLike copiedKey = StructCopy.copy(structProjection.wrap(asStructLike(row))); + + // Adding a pos-delete to replace the old path-offset. + BaseTaskWriter.PathOffset previous = insertedRowMap.put(copiedKey, pathOffset); + if (previous != null) { + // TODO attach the previous row if has a positional-delete row schema in appender factory. + posDeleteWriter.delete(previous.path, previous.rowOffset, null); + } + + dataWriter.write(row); + } + + /** + * Write the pos-delete if there's an existing row matching the given key. + * + * @param key has the same columns with the equality fields. + */ + private boolean internalPosDelete(StructLike key) { + BaseTaskWriter.PathOffset previous = insertedRowMap.remove(key); + + if (previous != null) { + // TODO attach the previous row if has a positional-delete row schema in appender factory. + posDeleteWriter.delete(previous.path, previous.rowOffset, null); + return true; + } + + return false; + } + + /** + * Delete those rows whose equality fields has the same values with the given row. It will write + * the entire row into the equality-delete file. + * + * @param row the given row to delete. + */ + public void delete(T row) throws IOException { + if (!internalPosDelete(structProjection.wrap(asStructLike(row)))) { + eqDeleteWriter.write(row); + } + } + + /** + * Delete those rows with the given key. It will only write the values of equality fields into + * the equality-delete file. + * + * @param key is the projected data whose columns are the same as the equality fields. + */ + public void deleteKey(T key) throws IOException { + if (!internalPosDelete(asStructLikeKey(key))) { + eqDeleteWriter.write(key); + } + } + + @Override + public void close() throws IOException { + try { + // Close data writer and add completed data files. + if (dataWriter != null) { + try { + dataWriter.close(); + } finally { + dataWriter = null; + } + } + + // Close eq-delete writer and add completed equality-delete files. + if (eqDeleteWriter != null) { + try { + eqDeleteWriter.close(); + } finally { + eqDeleteWriter = null; + } + } + + if (insertedRowMap != null) { + insertedRowMap.clear(); + insertedRowMap = null; + } + + // Add the completed pos-delete files. + if (posDeleteWriter != null) { + try { + // complete will call close + completedDeleteFiles.addAll(posDeleteWriter.complete()); + referencedDataFiles.addAll(posDeleteWriter.referencedDataFiles()); + } finally { + posDeleteWriter = null; + } + } + } catch (IOException | RuntimeException e) { + setFailure(e); + throw e; + } + } + } + + private static class PathOffset { + + private final CharSequence path; + private final long rowOffset; + + private PathOffset(CharSequence path, long rowOffset) { + this.path = path; + this.rowOffset = rowOffset; + } + + private static BaseTaskWriter.PathOffset of(CharSequence path, long rowOffset) { + return new BaseTaskWriter.PathOffset(path, rowOffset); + } + + @Override + public String toString() { + return MoreObjects.toStringHelper(this) + .add("path", path) + .add("row_offset", rowOffset) + .toString(); + } + } + + private abstract class BaseRollingWriter<W extends Closeable> implements Closeable { + + private static final int ROWS_DIVISOR = 1000; + private final StructLike partitionKey; + + private EncryptedOutputFile currentFile = null; + private W currentWriter = null; + private long currentRows = 0; + + private BaseRollingWriter(StructLike partitionKey) { + this.partitionKey = partitionKey; + openCurrent(); + } + + abstract W newWriter(EncryptedOutputFile file, StructLike partition); + + abstract long length(W writer); + + abstract void write(W writer, T record); + + public void write(T record) throws IOException { + write(currentWriter, record); + this.currentRows++; + + if (shouldRollToNewFile()) { + closeCurrent(); + openCurrent(); + } + } + + abstract void complete(W closedWriter); + + public CharSequence currentPath() { + Preconditions.checkNotNull(currentFile, "The currentFile shouldn't be null"); + return currentFile.encryptingOutputFile().location(); + } + + public long currentRows() { + return currentRows; + } + + private void openCurrent() { + if (partitionKey == null) { + // unpartitioned + this.currentFile = fileFactory.newOutputFile(); + } else { + // partitioned + this.currentFile = fileFactory.newOutputFile(partitionKey); + } + this.currentWriter = newWriter(currentFile, partitionKey); + this.currentRows = 0; + } + + private boolean shouldRollToNewFile() { + return currentRows % ROWS_DIVISOR == 0 && length(currentWriter) >= targetFileSize; + } + + private void closeCurrent() throws IOException { + if (currentWriter != null) { + try { + currentWriter.close(); + + if (currentRows == 0L) { + try { + io.deleteFile(currentFile.encryptingOutputFile()); + } catch (UncheckedIOException e) { + // the file may not have been created, and it isn't worth failing the job to clean up, + // skip deleting + } + } else { + complete(currentWriter); + } + + } catch (IOException | RuntimeException e) { + setFailure(e); + throw e; + + } finally { + this.currentFile = null; + this.currentWriter = null; + this.currentRows = 0; + } + } + } + + @Override + public void close() throws IOException { + closeCurrent(); + } + } + + protected class RollingFileWriter extends BaseTaskWriter<T>.BaseRollingWriter<DataWriter<T>> { + + public RollingFileWriter(StructLike partitionKey) { + super(partitionKey); + } + + @Override + DataWriter<T> newWriter(EncryptedOutputFile file, StructLike partitionKey) { + return appenderFactory.newDataWriter(file, format, partitionKey); + } + + @Override + long length(DataWriter<T> writer) { + return writer.length(); + } + + @Override + void write(DataWriter<T> writer, T record) { + writer.write(record); + } + + @Override + void complete(DataWriter<T> closedWriter) { + completedDataFiles.add(closedWriter.toDataFile()); + } + } + + protected class RollingEqDeleteWriter + extends + BaseTaskWriter<T>.BaseRollingWriter<EqualityDeleteWriter<T>> { + + RollingEqDeleteWriter(StructLike partitionKey) { + super(partitionKey); + } + + @Override + EqualityDeleteWriter<T> newWriter(EncryptedOutputFile file, StructLike partitionKey) { + return appenderFactory.newEqDeleteWriter(file, format, partitionKey); + } + + @Override + long length(EqualityDeleteWriter<T> writer) { + return writer.length(); + } + + @Override + void write(EqualityDeleteWriter<T> writer, T record) { + writer.write(record); + } + + @Override + void complete(EqualityDeleteWriter<T> closedWriter) { + completedDeleteFiles.add(closedWriter.toDeleteFile()); + } + } +} diff --git a/inlong-sort/sort-connectors/iceberg/src/main/java/org/apache/inlong/sort/iceberg/sink/trick/CharSequenceSet.java b/inlong-sort/sort-connectors/iceberg/src/main/java/org/apache/inlong/sort/iceberg/sink/trick/CharSequenceSet.java new file mode 100644 index 000000000..329c714b1 --- /dev/null +++ b/inlong-sort/sort-connectors/iceberg/src/main/java/org/apache/inlong/sort/iceberg/sink/trick/CharSequenceSet.java @@ -0,0 +1,189 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.inlong.sort.iceberg.sink.trick; + +import java.io.Serializable; +import java.util.Collection; +import java.util.Collections; +import java.util.Iterator; +import java.util.Objects; +import java.util.Set; +import java.util.stream.Collectors; +import org.apache.iceberg.relocated.com.google.common.collect.ImmutableList; +import org.apache.iceberg.relocated.com.google.common.collect.Iterables; +import org.apache.iceberg.relocated.com.google.common.collect.Iterators; +import org.apache.iceberg.relocated.com.google.common.collect.Sets; +import org.apache.iceberg.relocated.com.google.common.collect.Streams; +import org.apache.iceberg.util.CharSequenceWrapper; + +/** + * Copied from iceberg 1.1.x. Modify Set to concurrent Set. + */ +public class CharSequenceSet implements Set<CharSequence>, Serializable { + + private static final ThreadLocal<CharSequenceWrapper> wrappers = + ThreadLocal.withInitial(() -> CharSequenceWrapper.wrap(null)); + + public static CharSequenceSet of(Iterable<CharSequence> charSequences) { + return new CharSequenceSet(charSequences); + } + + public static CharSequenceSet empty() { + return new CharSequenceSet(ImmutableList.of()); + } + + private final Set<CharSequenceWrapper> wrapperSet; + + private CharSequenceSet(Iterable<CharSequence> charSequences) { + this.wrapperSet = + Collections.synchronizedSet( + Sets.newHashSet(Iterables.transform(charSequences, CharSequenceWrapper::wrap))); + } + + @Override + public int size() { + return wrapperSet.size(); + } + + @Override + public boolean isEmpty() { + return wrapperSet.isEmpty(); + } + + @Override + public boolean contains(Object obj) { + if (obj instanceof CharSequence) { + CharSequenceWrapper wrapper = wrappers.get(); + boolean result = wrapperSet.contains(wrapper.set((CharSequence) obj)); + wrapper.set(null); // don't hold a reference to the value + return result; + } + return false; + } + + @Override + public Iterator<CharSequence> iterator() { + return Iterators.transform(wrapperSet.iterator(), CharSequenceWrapper::get); + } + + @Override + public Object[] toArray() { + return Iterators.toArray(iterator(), CharSequence.class); + } + + @Override + @SuppressWarnings("unchecked") + public <T> T[] toArray(T[] destArray) { + int size = wrapperSet.size(); + if (destArray.length < size) { + return (T[]) toArray(); + } + + Iterator<CharSequence> iter = iterator(); + int ind = 0; + while (iter.hasNext()) { + destArray[ind] = (T) iter.next(); + ind += 1; + } + + if (destArray.length > size) { + destArray[size] = null; + } + + return destArray; + } + + @Override + public boolean add(CharSequence charSequence) { + return wrapperSet.add(CharSequenceWrapper.wrap(charSequence)); + } + + @Override + public boolean remove(Object obj) { + if (obj instanceof CharSequence) { + CharSequenceWrapper wrapper = wrappers.get(); + boolean result = wrapperSet.remove(wrapper.set((CharSequence) obj)); + wrapper.set(null); // don't hold a reference to the value + return result; + } + return false; + } + + @Override + @SuppressWarnings("CollectionUndefinedEquality") + public boolean containsAll(Collection<?> objects) { + if (objects != null) { + return Iterables.all(objects, this::contains); + } + return false; + } + + @Override + public boolean addAll(Collection<? extends CharSequence> charSequences) { + if (charSequences != null) { + return Iterables.addAll( + wrapperSet, Iterables.transform(charSequences, CharSequenceWrapper::wrap)); + } + return false; + } + + @Override + public boolean retainAll(Collection<?> objects) { + if (objects != null) { + return Iterables.removeAll(wrapperSet, objects); + } + return false; + } + + @Override + public boolean removeAll(Collection<?> objects) { + if (objects != null) { + return Iterables.removeAll(wrapperSet, objects); + } + return false; + } + + @Override + public void clear() { + wrapperSet.clear(); + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + + if (o == null || getClass() != o.getClass()) { + return false; + } + + CharSequenceSet that = (CharSequenceSet) o; + return wrapperSet.equals(that.wrapperSet); + } + + @Override + public int hashCode() { + return Objects.hashCode(wrapperSet); + } + + @Override + public String toString() { + return Streams.stream(iterator()).collect(Collectors.joining("CharSequenceSet({", ", ", "})")); + } +} diff --git a/inlong-sort/sort-connectors/iceberg/src/main/java/org/apache/inlong/sort/iceberg/sink/trick/SortedPosDeleteWriter.java b/inlong-sort/sort-connectors/iceberg/src/main/java/org/apache/inlong/sort/iceberg/sink/trick/SortedPosDeleteWriter.java new file mode 100644 index 000000000..17a7bfad0 --- /dev/null +++ b/inlong-sort/sort-connectors/iceberg/src/main/java/org/apache/inlong/sort/iceberg/sink/trick/SortedPosDeleteWriter.java @@ -0,0 +1,220 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.inlong.sort.iceberg.sink.trick; + +import java.io.IOException; +import java.io.UncheckedIOException; +import java.util.Comparator; +import java.util.List; +import java.util.Map; +import org.apache.iceberg.DeleteFile; +import org.apache.iceberg.FileFormat; +import org.apache.iceberg.StructLike; +import org.apache.iceberg.deletes.PositionDelete; +import org.apache.iceberg.deletes.PositionDeleteWriter; +import org.apache.iceberg.encryption.EncryptedOutputFile; +import org.apache.iceberg.io.DeleteWriteResult; +import org.apache.iceberg.io.FileAppenderFactory; +import org.apache.iceberg.io.FileWriter; +import org.apache.iceberg.io.OutputFileFactory; +import org.apache.iceberg.relocated.com.google.common.base.Preconditions; +import org.apache.iceberg.relocated.com.google.common.collect.Lists; +import org.apache.iceberg.relocated.com.google.common.collect.Maps; +import org.apache.iceberg.types.Comparators; +import org.apache.iceberg.util.CharSequenceSet; +import org.apache.iceberg.util.CharSequenceWrapper; + +class SortedPosDeleteWriter<T> implements FileWriter<PositionDelete<T>, DeleteWriteResult> { + + private static final long DEFAULT_RECORDS_NUM_THRESHOLD = 100_000L; + + private final Map<CharSequenceWrapper, List<SortedPosDeleteWriter.PosRow<T>>> posDeletes = Maps.newHashMap(); + private final List<DeleteFile> completedFiles = Lists.newArrayList(); + private final CharSequenceSet referencedDataFiles = CharSequenceSet.empty(); + private final CharSequenceWrapper wrapper = CharSequenceWrapper.wrap(null); + + private final FileAppenderFactory<T> appenderFactory; + private final OutputFileFactory fileFactory; + private final FileFormat format; + private final StructLike partition; + private final long recordsNumThreshold; + + private int records = 0; + private boolean closed = false; + private Throwable failure; + + SortedPosDeleteWriter( + FileAppenderFactory<T> appenderFactory, + OutputFileFactory fileFactory, + FileFormat format, + StructLike partition, + long recordsNumThreshold) { + this.appenderFactory = appenderFactory; + this.fileFactory = fileFactory; + this.format = format; + this.partition = partition; + this.recordsNumThreshold = recordsNumThreshold; + } + + SortedPosDeleteWriter( + FileAppenderFactory<T> appenderFactory, + OutputFileFactory fileFactory, + FileFormat format, + StructLike partition) { + this(appenderFactory, fileFactory, format, partition, DEFAULT_RECORDS_NUM_THRESHOLD); + } + + protected void setFailure(Throwable throwable) { + if (failure == null) { + this.failure = throwable; + } + } + + @Override + public long length() { + throw new UnsupportedOperationException( + this.getClass().getName() + " does not implement length"); + } + + @Override + public void write(PositionDelete<T> payload) { + delete(payload.path(), payload.pos(), payload.row()); + } + + public void delete(CharSequence path, long pos) { + delete(path, pos, null); + } + + public void delete(CharSequence path, long pos, T row) { + List<SortedPosDeleteWriter.PosRow<T>> posRows = posDeletes.get(wrapper.set(path)); + if (posRows != null) { + posRows.add(SortedPosDeleteWriter.PosRow.of(pos, row)); + } else { + posDeletes.put(CharSequenceWrapper.wrap(path), Lists.newArrayList( + SortedPosDeleteWriter.PosRow.of(pos, row))); + } + + records += 1; + + // TODO Flush buffer based on the policy that checking whether whole heap memory size exceed the + // threshold. + if (records >= recordsNumThreshold) { + flushDeletes(); + } + } + + public List<DeleteFile> complete() throws IOException { + close(); + + Preconditions.checkState(failure == null, "Cannot return results from failed writer", failure); + + return completedFiles; + } + + public CharSequenceSet referencedDataFiles() { + return referencedDataFiles; + } + + @Override + public void close() throws IOException { + if (!closed) { + this.closed = true; + flushDeletes(); + } + } + + @Override + public DeleteWriteResult result() { + Preconditions.checkState(closed, "Cannot get result from unclosed writer"); + return new DeleteWriteResult(completedFiles, referencedDataFiles); + } + + private void flushDeletes() { + if (posDeletes.isEmpty()) { + return; + } + + // Create a new output file. + EncryptedOutputFile outputFile; + if (partition == null) { + outputFile = fileFactory.newOutputFile(); + } else { + outputFile = fileFactory.newOutputFile(partition); + } + + PositionDeleteWriter<T> writer = + appenderFactory.newPosDeleteWriter(outputFile, format, partition); + PositionDelete<T> posDelete = PositionDelete.create(); + try (PositionDeleteWriter<T> closeableWriter = writer) { + // Sort all the paths. + List<CharSequence> paths = Lists.newArrayListWithCapacity(posDeletes.keySet().size()); + for (CharSequenceWrapper charSequenceWrapper : posDeletes.keySet()) { + paths.add(charSequenceWrapper.get()); + } + paths.sort(Comparators.charSequences()); + + // Write all the sorted <path, pos, row> triples. + for (CharSequence path : paths) { + List<SortedPosDeleteWriter.PosRow<T>> positions = posDeletes.get(wrapper.set(path)); + positions.sort(Comparator.comparingLong(SortedPosDeleteWriter.PosRow::pos)); + + positions.forEach( + posRow -> closeableWriter.write(posDelete.set(path, posRow.pos(), posRow.row()))); + } + } catch (IOException e) { + setFailure(e); + throw new UncheckedIOException( + "Failed to write the sorted path/pos pairs to pos-delete file: " + + outputFile.encryptingOutputFile().location(), + e); + } + + // Clear the buffered pos-deletions. + posDeletes.clear(); + records = 0; + + // Add the referenced data files. + referencedDataFiles.addAll(writer.referencedDataFiles()); + + // Add the completed delete files. + completedFiles.add(writer.toDeleteFile()); + } + + private static class PosRow<R> { + + private final long pos; + private final R row; + + static <R> SortedPosDeleteWriter.PosRow<R> of(long pos, R row) { + return new SortedPosDeleteWriter.PosRow<>(pos, row); + } + + private PosRow(long pos, R row) { + this.pos = pos; + this.row = row; + } + + long pos() { + return pos; + } + + R row() { + return row; + } + } +} diff --git a/inlong-sort/sort-connectors/iceberg/src/main/java/org/apache/inlong/sort/iceberg/sink/trick/StructCopy.java b/inlong-sort/sort-connectors/iceberg/src/main/java/org/apache/inlong/sort/iceberg/sink/trick/StructCopy.java new file mode 100644 index 000000000..4f50c6b51 --- /dev/null +++ b/inlong-sort/sort-connectors/iceberg/src/main/java/org/apache/inlong/sort/iceberg/sink/trick/StructCopy.java @@ -0,0 +1,59 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.inlong.sort.iceberg.sink.trick; + +import org.apache.iceberg.StructLike; + +/** Copy the StructLike's values into a new one. It does not handle list or map values now. */ +class StructCopy implements StructLike { + + static StructLike copy(StructLike struct) { + return struct != null ? new StructCopy(struct) : null; + } + + private final Object[] values; + + private StructCopy(StructLike toCopy) { + this.values = new Object[toCopy.size()]; + + for (int i = 0; i < values.length; i += 1) { + Object value = toCopy.get(i, Object.class); + + if (value instanceof StructLike) { + values[i] = copy((StructLike) value); + } else { + values[i] = value; + } + } + } + + @Override + public int size() { + return values.length; + } + + @Override + public <T> T get(int pos, Class<T> javaClass) { + return javaClass.cast(values[pos]); + } + + @Override + public <T> void set(int pos, T value) { + throw new UnsupportedOperationException("Struct copy cannot be modified"); + } +}
