voonhous commented on code in PR #18961: URL: https://github.com/apache/hudi/pull/18961#discussion_r3820934543
########## hudi-common/src/main/java/org/apache/hudi/core/io/storage/VariantShreddingInferenceFileWriter.java: ########## @@ -0,0 +1,323 @@ +/* + * 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.hudi.core.io.storage; + +import org.apache.hudi.common.avro.VariantShreddingSchemaInferrer; +import org.apache.hudi.common.avro.VariantShreddingSchemaInferrer.VariantSample; +import org.apache.hudi.common.model.HoodieKey; +import org.apache.hudi.common.model.HoodieRecord; +import org.apache.hudi.common.schema.HoodieSchema; +import org.apache.hudi.common.util.DefaultSizeEstimator; +import org.apache.hudi.common.util.SizeEstimator; +import org.apache.hudi.exception.HoodieIOException; + +import lombok.extern.slf4j.Slf4j; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.Collections; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.Properties; + +/** + * A {@link HoodieFileWriter} decorator that infers a per-file variant shredding schema from the + * first records before opening the real parquet writer. + * + * <p>Records are buffered (and their variant binaries sampled) until a threshold is reached or + * the writer closes, the sampled binaries are fed to a {@link VariantShreddingSchemaInferrer}, + * the real writer is created against the schema with the inferred typed_value spliced in, and + * the buffer is replayed in arrival order. Replay reproduces each call exactly (write vs + * writeWithMetadata), so commit seqnos, bloom filters and min/max record keys come out + * identical to a non-buffered write. Buffering thresholds mirror Spark's + * {@code ParquetOutputWriterWithVariantShredding} (4096 rows / 64MB). + * + * <p>Buffered records are {@link HoodieRecord#copy() copied} because Spark iterators reuse row + * instances, then handed to {@link VariantSampleExtractor#prepare} so an extractor that has to + * materialize the record to sample it (the Avro one deserializes payload-backed records) can + * return the materialized form for buffering and the replay does not repeat that work. For + * record types where copy() is identity (Avro), replay additionally relies on writer-level + * records being freshly allocated per record, which holds today; variant samples are extracted + * eagerly into immutable byte arrays so inference itself never depends on it. + * + * <p>Inference failures never fail the write: the file falls back to unshredded variants. This + * deliberately diverges from Spark (which propagates inference failures) because a throwing + * inference would fail compaction. Writer-creation or replay failures, however, are latched and + * rethrown from every subsequent call including {@link #close()}, so a task cannot silently + * drop buffered records that the handle already counted as written. + * + * <p>{@link #writeRow} carries neither a {@link HoodieRecord} nor a schema to sample from, so the + * first such call materializes the real writer with whatever has been sampled so far (unshredded + * when nothing has) and passes the row straight through. Footer metadata added before + * materialization is queued and handed to the real writer once it exists; parquet only consumes + * it at close, so nothing is lost. + * + * <p>Single-threaded by contract, same as the writers it wraps. + * + * <p>See https://github.com/apache/hudi/issues/18937.</p> + * + * @param <T> the engine-native record type of the wrapped writer + */ +@Slf4j +public class VariantShreddingInferenceFileWriter<T> implements HoodieFileWriter<T> { + + /** Buffer caps mirroring Spark's ParquetOutputWriterWithVariantShredding. */ + public static final int MAX_BUFFERED_RECORDS = 4096; + public static final long MAX_BUFFERED_BYTES = 64L * 1024 * 1024; + private static final int SIZE_ESTIMATE_INTERVAL = 100; + + /** + * Extracts the variant binaries of the inferable columns from a record. Bound to the writer + * schema and column set by the creating factory; must defensively copy the bytes. + */ + @FunctionalInterface + public interface VariantSampleExtractor { + VariantSample[] extract(HoodieRecord record, HoodieSchema schema, Properties props) throws IOException; + + /** + * Returns the record to buffer for replay; {@link #extract} is then called with that record. + * An extractor that must materialize the record to sample it returns the materialized form, + * so the replay does not redo the work. Defaults to the record itself. + */ + default HoodieRecord prepare(HoodieRecord record, HoodieSchema schema, Properties props) throws IOException { + return record; + } + + /** + * Bytes of state that every buffered record references but that is shared across them, so a + * deep object-size walk of one record counts it in full: the Avro {@code Schema} graph of an + * Avro record, the {@code StructType} of a Spark row. Subtracted from each record's size + * estimate so the byte cap budgets record payload rather than the schema times the record + * count (the HUDI-9499 class of over-estimate, which would shrink the inference sample to a + * fraction of the intended 4096 rows). Defaults to 0. + */ + default long sharedSizeEstimate(HoodieSchema schema) { + return 0; + } + } + + /** Creates the real file writer once the inferred typed_value schemas are known. */ + @FunctionalInterface + public interface InferredWriterFactory<T> { + HoodieFileWriter<T> create(Map<String, HoodieSchema> inferredTypedValues) throws IOException; + } + + private final List<String> variantColumns; + private final VariantSampleExtractor extractor; + private final VariantShreddingSchemaInferrer inferrer; + private final InferredWriterFactory<T> writerFactory; + private final long maxBufferedBytes; + private final SizeEstimator<HoodieRecord> sizeEstimator = new DefaultSizeEstimator<>(); + + private final List<BufferedWrite> buffer = new ArrayList<>(); + private final List<VariantSample[]> samples = new ArrayList<>(); + private final Map<String, String> pendingFooterMetadata = new LinkedHashMap<>(); + private long estimatedRecordSize = 0; + private long bufferedBytes = 0; + private HoodieFileWriter<T> delegate; + private IOException fatalFailure; + private boolean closed = false; + + public VariantShreddingInferenceFileWriter(List<String> variantColumns, + VariantSampleExtractor extractor, + VariantShreddingSchemaInferrer inferrer, + InferredWriterFactory<T> writerFactory, + long maxFileSize) { + this.variantColumns = variantColumns; + this.extractor = extractor; + this.inferrer = inferrer; + this.writerFactory = writerFactory; + this.maxBufferedBytes = Math.min(MAX_BUFFERED_BYTES, Math.max(1, maxFileSize)); + } + + @Override + public boolean canWrite() { + // Nothing has been physically written while buffering, so size-based rollover cannot apply yet. + return delegate == null || delegate.canWrite(); + } + + @Override + public void writeWithMetadata(HoodieKey key, HoodieRecord record, HoodieSchema schema, Properties props) throws IOException { + rethrowIfFailed(); + if (delegate != null) { + delegate.writeWithMetadata(key, record, schema, props); + } else { + buffer(true, key, null, record, schema, props); + } + } + + @Override + public void write(String recordKey, HoodieRecord record, HoodieSchema schema, Properties props) throws IOException { + rethrowIfFailed(); + if (delegate != null) { + delegate.write(recordKey, record, schema, props); + } else { + buffer(false, null, recordKey, record, schema, props); + } + } + + @Override + public void writeRow(String recordKey, T record) throws IOException { + rethrowIfFailed(); + // No HoodieRecord or schema to sample from: materialize with what has been sampled so far and + // pass the row through. Today only the native log-format delete writer takes this path, and its + // delete schema has no variant column to infer for. + materialize(); + delegate.writeRow(recordKey, record); + } + + @Override + public void addFooterMetadata(Map<String, String> footerMetadata) { + if (delegate != null) { + delegate.addFooterMetadata(footerMetadata); + } else { + // Footer metadata is only consumed at close, so it can wait for the real writer. + pendingFooterMetadata.putAll(footerMetadata); + } + } + + @Override + public void close() throws IOException { + if (closed) { + return; + } + closed = true; + boolean delegateClosed = false; + try { + rethrowIfFailed(); + // Materialize even with an empty buffer: handles expect the file to exist at close. + materialize(); + // Mark before close() so a throwing delegate.close() surfaces, not retried in the catch. + delegateClosed = true; + delegate.close(); + } catch (IOException | RuntimeException e) { + if (delegate != null && !delegateClosed) { + try { + delegate.close(); + } catch (Exception suppressed) { + // Best-effort cleanup; surface the original failure. + } + } + throw e; + } + } + + @Override + public Object getFileFormatMetadata() { + try { + rethrowIfFailed(); + materialize(); + } catch (IOException e) { + throw new HoodieIOException("Failed to materialize the parquet writer for format metadata", e); + } + return delegate.getFileFormatMetadata(); + } + + private void buffer(boolean withMetadata, HoodieKey key, String recordKey, HoodieRecord record, + HoodieSchema schema, Properties props) throws IOException { + rethrowIfFailed(); + HoodieRecord buffered = extractor.prepare(record.copy(), schema, props); Review Comment: Added the `prepare()` override returning a detached `HoodieSparkRecord` in 0c90b2acb1af. Scope note: the deflating combination needs the preserveMetadata compactor constructor of `HoodieWriteMergeHandle`, which has no live Spark caller (the default is `FileGroupReaderBasedMergeHandle`, which never deflates), and a deflated buffered record fails replay loudly via `getData()`'s `IllegalStateException` rather than corrupting the file. So: hardening, not a live bug. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
