wombatu-kun commented on code in PR #18961: URL: https://github.com/apache/hudi/pull/18961#discussion_r3835040866
########## hudi-client/hudi-spark-client/src/main/java/org/apache/hudi/io/storage/row/VariantShreddingInferenceInternalRowFileWriter.java: ########## @@ -0,0 +1,260 @@ +/* + * 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.io.storage.row; + +import org.apache.hudi.SparkAdapterSupport$; +import org.apache.hudi.common.avro.VariantShreddingSchemaInferrer; +import org.apache.hudi.common.avro.VariantShreddingSchemaInferrer.VariantSample; +import org.apache.hudi.common.schema.HoodieSchema; +import org.apache.hudi.common.util.CloseableUtils; +import org.apache.hudi.common.util.DefaultSizeEstimator; +import org.apache.hudi.core.io.storage.VariantShreddingInferenceFileWriter; + +import lombok.extern.slf4j.Slf4j; +import org.apache.spark.sql.catalyst.InternalRow; +import org.apache.spark.sql.catalyst.expressions.UnsafeRow; +import org.apache.spark.sql.types.StructField; +import org.apache.spark.sql.types.StructType; +import org.apache.spark.unsafe.types.UTF8String; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.Map; + +/** + * A {@link HoodieInternalRowFileWriter} decorator that infers a per-file variant shredding + * schema from the first rows before opening the real parquet writer; the row-writer-path + * sibling of {@link VariantShreddingInferenceFileWriter}, sharing its buffering thresholds and + * failure semantics. + * + * <p>Meta columns including the commit seqno are composed into the row by the handle BEFORE + * {@code writeRow}, so ordered replay is value-exact here by construction. Rows and keys are + * copied because Spark iterators reuse their instances.</p> + */ +@Slf4j +public class VariantShreddingInferenceInternalRowFileWriter implements HoodieInternalRowFileWriter { + + private static final int SIZE_ESTIMATE_INTERVAL = 100; + + /** Creates the real row file writer once the inferred typed_value schemas are known. */ + @FunctionalInterface + public interface InferredRowWriterFactory { + HoodieInternalRowFileWriter create(Map<String, HoodieSchema> inferredTypedValues) throws IOException; + } + + private final List<String> variantColumns; + private final int[] ordinals; + private final VariantShreddingSchemaInferrer inferrer; + private final InferredRowWriterFactory writerFactory; + private final long maxBufferedBytes; + private final DefaultSizeEstimator<InternalRow> sizeEstimator = new DefaultSizeEstimator<>(); + + private final List<BufferedRow> buffer = new ArrayList<>(); + private final List<VariantSample[]> samples = new ArrayList<>(); + private long bufferedBytes = 0; + private long estimatedRowSize = 0; + private long estimatedRowCount = 0; + private HoodieInternalRowFileWriter delegate; + private IOException fatalFailure; + private boolean closed = false; + + public VariantShreddingInferenceInternalRowFileWriter(List<String> variantColumns, + int[] ordinals, + VariantShreddingSchemaInferrer inferrer, + InferredRowWriterFactory writerFactory, + long maxFileSize) { + this.variantColumns = variantColumns; + this.ordinals = ordinals; + this.inferrer = inferrer; + this.writerFactory = writerFactory; + this.maxBufferedBytes = Math.min(VariantShreddingInferenceFileWriter.MAX_BUFFERED_BYTES, Math.max(1, maxFileSize)); + } + + /** Resolves the buffer ordinal of each variant column in {@code structType}; -1 when absent. */ + public static int[] resolveOrdinals(StructType structType, List<String> columnNames) { + int[] ordinals = new int[columnNames.size()]; + StructField[] fields = structType.fields(); + for (int i = 0; i < columnNames.size(); i++) { + ordinals[i] = -1; + for (int j = 0; j < fields.length; j++) { + if (fields[j].name().equals(columnNames.get(i))) { + ordinals[i] = j; + break; + } + } + } + return ordinals; + } + + @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 writeRow(UTF8String key, InternalRow row) throws IOException { + rethrowIfFailed(); + if (delegate != null) { + delegate.writeRow(key, row); + } else { + // copy(): Spark iterators reuse key instances. + buffer(key == null ? null : key.copy(), true, row); + } + } + + @Override + public void writeRow(InternalRow row) throws IOException { + rethrowIfFailed(); + if (delegate != null) { + delegate.writeRow(row); + } else { + buffer(null, false, row); + } + } + + @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 | Error e) { + // Error included: materialize() rethrows the Error it latches, and the delegate it + // created must still be closed. + if (delegate != null && !delegateClosed) { + // HoodieInternalRowFileWriter is not an AutoCloseable, hence the method reference. + CloseableUtils.closeSuppressing(delegate::close, e); + } + throw e; + } + } + + private void buffer(UTF8String key, boolean withKey, InternalRow row) throws IOException { + rethrowIfFailed(); + InternalRow copied = row.copy(); + samples.add(extractSamples(copied)); + buffer.add(new BufferedRow(key, withKey, copied)); + chargeSize(copied); + if (buffer.size() >= VariantShreddingInferenceFileWriter.MAX_BUFFERED_RECORDS || bufferedBytes >= maxBufferedBytes) { + materialize(); + } + } + + private VariantSample[] extractSamples(InternalRow row) { + VariantSample[] out = new VariantSample[ordinals.length]; + for (int i = 0; i < ordinals.length; i++) { + if (ordinals[i] >= 0) { + out[i] = SparkAdapterSupport$.MODULE$.sparkAdapter().extractVariantBinary(row, ordinals[i]); + } + } + return out; + } + + /** Charges {@code row} against the byte cap: an exact size for an UnsafeRow, an estimate otherwise. */ + private void chargeSize(InternalRow row) { + if (row instanceof UnsafeRow) { + bufferedBytes += ((UnsafeRow) row).getSizeInBytes(); + return; + } + // Re-estimate periodically so a small first row cannot defeat the byte cap + // (same moving-average idiom as ExternalSpillableMap). + estimatedRowCount++; + if (estimatedRowSize == 0 || estimatedRowCount % SIZE_ESTIMATE_INTERVAL == 0) { + long previous = estimatedRowSize; + long sampled = Math.max(1, sizeEstimator.sizeEstimate(row)); + estimatedRowSize = estimatedRowSize == 0 + ? sampled : (long) (estimatedRowSize * 0.9 + sampled * 0.1); + // Rescale the rows already charged at the old estimate, or the cap would trip late once the + // estimate grew. Only the estimated ones: the UnsafeRow branch above charges exact sizes. + bufferedBytes += (estimatedRowCount - 1) * (estimatedRowSize - previous); Review Comment: No test buffers an UnsafeRow and an estimated row together, so the incremental form here is unpinned - switching to the record writer's absolute assignment would pass every test while discarding the exact UnsafeRow charges. Interleaving the two row types in one buffer would pin it. ########## hudi-spark-datasource/hudi-spark/src/test/scala/org/apache/spark/sql/hudi/dml/schema/TestVariantDataType.scala: ########## @@ -1232,6 +1455,32 @@ class TestVariantDataType extends HoodieSparkSqlTestBase { } } + /** + * Sums the non-null value counts of the leaf columns under `column`.typed_value across all + * blocks of the file, from the block column statistics. + */ + private def typedValueNonNullCount(filePath: String, column: String): Long = { + val conf = spark.sparkContext.hadoopConfiguration + val inputFile = HadoopInputFile.fromPath(new HadoopPath(filePath), conf) + val reader = ParquetFileReader.open(inputFile) + try { + val prefix = s"$column.typed_value" + reader.getFooter.getBlocks.asScala.flatMap(_.getColumns.asScala) + .filter { c => + val dot = c.getPath.toDotString + dot == prefix || dot.startsWith(prefix + ".") Review Comment: This prefix match also picks up each object field's residual `value` leaf, so a file where every field fell back to its residual still sums above zero. Excluding leaves whose path ends in `.value` would keep the count on the typed leaves only. ########## hudi-common/src/main/java/org/apache/hudi/core/io/storage/VariantShreddingInferenceFileWriter.java: ########## @@ -0,0 +1,330 @@ +/* + * 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.CloseableUtils; +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}, which serves two purposes: + * an extractor that has to materialize the record to sample it (the Avro one deserializes + * payload-backed records) returns the materialized form for buffering so the replay does not + * repeat that work, and (since copy() returns the caller's own wrapper for every record type + * today) both shipped extractors return a wrapper the caller does not hold, so a handle that + * deflates its record right after the write call cannot blank a buffered one. Records with + * nothing to materialize (delete payloads) are buffered as they are, so replay still 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 Review Comment: Both `writeRow` callers - the native-log delete writer and `HoodieNativeCDCFileWriter` - pass schemas with no top-level variant, so the factory never wraps them and neither reaches this path. Saying no production caller reaches it today would be more accurate than naming the delete writer. ########## hudi-common/src/test/java/org/apache/hudi/common/schema/TestHoodieSchema.java: ########## @@ -3087,4 +3087,120 @@ public void testGetPlainTypedValueSchemaEmpty() { HoodieSchema.Variant unshreddedVariant = HoodieSchema.createVariant(); assertFalse(unshreddedVariant.getPlainTypedValueSchema().isPresent()); } + + @Test + public void testGetPlainTypedValueSchemaNestedObjectRecursion() { + // Depth-2 spec form: typed_value { a: wrapper{value, typed_value: { b: wrapper{value, typed_value: long} }} }. + // Both record levels are named "typed_value", as the schema converters produce them. + HoodieSchema innerObject = HoodieSchema.createRecord("typed_value", "inner.ns", null, + Collections.singletonList(HoodieSchemaField.of("b", + HoodieSchema.createNullable(HoodieSchema.createShreddedFieldStruct("b_wrapper", HoodieSchema.create(HoodieSchemaType.LONG)))))); + HoodieSchema topTypedValue = HoodieSchema.createRecord("typed_value", "outer.ns", null, + Collections.singletonList(HoodieSchemaField.of("a", + HoodieSchema.createNullable(HoodieSchema.createShreddedFieldStruct("a_wrapper", innerObject))))); + // Nullable typed_value, as produced by the inferred-shredding splice. + HoodieSchema.Variant variant = HoodieSchema.createVariantShredded(HoodieSchema.createNullable(topTypedValue)); + + Option<HoodieSchema> plainOpt = variant.getPlainTypedValueSchema(); + assertTrue(plainOpt.isPresent()); + HoodieSchema plain = plainOpt.get(); + assertEquals(HoodieSchemaType.RECORD, plain.getType()); + assertEquals(1, plain.getFields().size()); + + HoodieSchema aPlain = plain.getFields().get(0).schema(); + aPlain = aPlain.isNullable() ? aPlain.getNonNullType() : aPlain; + assertEquals(HoodieSchemaType.RECORD, aPlain.getType()); + assertEquals(1, aPlain.getFields().size()); + + HoodieSchema bPlain = aPlain.getFields().get(0).schema(); + bPlain = bPlain.isNullable() ? bPlain.getNonNullType() : bPlain; + assertEquals(HoodieSchemaType.LONG, bPlain.getType()); + + // Generated plain record names must be unique per nesting level: a nested record carrying + // its ancestor's fullname is an Avro self-reference, which Spark rejects as recursion. + assertNotEquals(plain.getFullName(), aPlain.getFullName()); + } + + @Test + public void testGetPlainTypedValueSchemaNamesDistinguishConcatenatingPaths() { + // Two object paths whose segments concatenate to the same string ("x_y" > "z" and "x" > "y_z") + // must still yield distinct plain record names for the objects at their leaves; the path goes + // into the namespace, where the '.' separator keeps them apart (a flat "<path>_plain" name + // gave both leaves "typed_value_x_y_z_plain"). + HoodieSchema leafObject = HoodieSchema.createRecord("typed_value", "leaf.ns", null, + Collections.singletonList(HoodieSchemaField.of("c", + HoodieSchema.createNullable(HoodieSchema.createShreddedFieldStruct("c_wrapper", HoodieSchema.create(HoodieSchemaType.LONG)))))); + HoodieSchema underXy = HoodieSchema.createRecord("typed_value", "a.ns", null, + Collections.singletonList(HoodieSchemaField.of("z", + HoodieSchema.createNullable(HoodieSchema.createShreddedFieldStruct("z_wrapper", leafObject))))); + HoodieSchema underX = HoodieSchema.createRecord("typed_value", "b.ns", null, + Collections.singletonList(HoodieSchemaField.of("y_z", + HoodieSchema.createNullable(HoodieSchema.createShreddedFieldStruct("y_z_wrapper", leafObject))))); + HoodieSchema topTypedValue = HoodieSchema.createRecord("typed_value", "outer.ns", null, Arrays.asList( + HoodieSchemaField.of("x_y", HoodieSchema.createNullable(HoodieSchema.createShreddedFieldStruct("x_y_wrapper", underXy))), + HoodieSchemaField.of("x", HoodieSchema.createNullable(HoodieSchema.createShreddedFieldStruct("x_wrapper", underX))))); + + HoodieSchema plain = HoodieSchema.createVariantShredded(topTypedValue).getPlainTypedValueSchema().get(); + HoodieSchema zLeaf = plain.getField("x_y").get().schema().getNonNullType() + .getField("z").get().schema().getNonNullType(); + HoodieSchema yzLeaf = plain.getField("x").get().schema().getNonNullType() + .getField("y_z").get().schema().getNonNullType(); + assertEquals(HoodieSchemaType.RECORD, zLeaf.getType()); + assertEquals(HoodieSchemaType.RECORD, yzLeaf.getType()); + assertNotEquals(zLeaf.getFullName(), yzLeaf.getFullName()); + // Serializing the whole tree (as the config-splice path does) must not alias the two leaves. + assertNotNull(plain.getAvroSchema().toString()); Review Comment: Avro writes a repeated name as a silent reference rather than throwing, and `toString()` never returns null, so this passes whether or not the two leaves alias. Re-parsing the serialized schema and re-asserting the two leaf full names differ would make the line do what its comment claims. ########## hudi-common/src/test/java/org/apache/hudi/core/io/storage/TestVariantShreddingInferenceFileWriter.java: ########## @@ -0,0 +1,546 @@ +/* + * 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.HoodieAvroIndexedRecord; +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.schema.HoodieSchemaField; +import org.apache.hudi.common.schema.HoodieSchemaType; +import org.apache.hudi.common.util.DefaultSizeEstimator; +import org.apache.hudi.exception.HoodieIOException; + +import org.apache.avro.generic.GenericData; +import org.apache.avro.generic.GenericRecord; +import org.junit.jupiter.api.Test; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.HashMap; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.Properties; + +import static java.util.Collections.singletonList; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +public class TestVariantShreddingInferenceFileWriter { + + private static final HoodieSchema RECORD_SCHEMA = HoodieSchema.createRecord("rec", null, null, + singletonList(HoodieSchemaField.of("id", HoodieSchema.create(HoodieSchemaType.STRING)))); + private static final Properties PROPS = new Properties(); + + private final VariantShreddingInferenceFileWriter.VariantSampleExtractor noopExtractor = + (record, schema, props) -> new VariantSample[1]; + + /** A decorator over {@link #noopExtractor} for the column {@code v}. */ + private VariantShreddingInferenceFileWriter<Object> writer( + VariantShreddingSchemaInferrer inferrer, + VariantShreddingInferenceFileWriter.InferredWriterFactory<Object> factory, + long maxFileSize) { + return new VariantShreddingInferenceFileWriter<>(singletonList("v"), noopExtractor, inferrer, factory, maxFileSize); + } + + private static HoodieRecord newRecord(String id) { + GenericRecord data = new GenericData.Record(RECORD_SCHEMA.toAvroSchema()); + data.put("id", id); + return new HoodieAvroIndexedRecord(new HoodieKey(id, "p"), data); + } + + /** Records every call so replay order and call kinds can be asserted. */ + private static class RecordingWriter implements HoodieFileWriter<Object> { + private final List<String> calls = new ArrayList<>(); + private final List<HoodieRecord> writtenRecords = new ArrayList<>(); + private final Map<String, String> footerMetadata = new LinkedHashMap<>(); + private final Object fileFormatMetadata = new Object(); + private int closeCount = 0; + /** An IOException or an Error; anything else is a misuse of the stub. */ + private Throwable failWriteWith; + private IOException failCloseWith; + + @Override + public boolean canWrite() { + return true; + } + + @Override + public void writeWithMetadata(HoodieKey key, HoodieRecord record, HoodieSchema schema, Properties props) throws IOException { + failIfConfigured(failWriteWith); + calls.add("meta:" + key.getRecordKey()); + writtenRecords.add(record); + } + + @Override + public void write(String recordKey, HoodieRecord record, HoodieSchema schema, Properties props) throws IOException { + failIfConfigured(failWriteWith); + calls.add("plain:" + recordKey); + writtenRecords.add(record); + } + + @Override + public void writeRow(String recordKey, Object record) { + calls.add("row:" + recordKey); + } + + @Override + public void addFooterMetadata(Map<String, String> footerMetadata) { + this.footerMetadata.putAll(footerMetadata); + } + + @Override + public Object getFileFormatMetadata() { + return fileFormatMetadata; + } + + @Override + public void close() throws IOException { + closeCount++; + failIfConfigured(failCloseWith); + } + + private static void failIfConfigured(Throwable failure) throws IOException { + if (failure instanceof Error) { + throw (Error) failure; + } else if (failure != null) { + throw (IOException) failure; + } + } + } + + @Test + public void testReplayPreservesOrderAndCallKinds() throws IOException { + Map<String, HoodieSchema> inferred = new HashMap<>(); + inferred.put("v", HoodieSchema.create(HoodieSchemaType.LONG)); + List<Map<String, HoodieSchema>> factoryCalls = new ArrayList<>(); + RecordingWriter delegate = new RecordingWriter(); + + VariantShreddingInferenceFileWriter<Object> writer = writer((columns, samples) -> inferred, + map -> { + factoryCalls.add(map); + return delegate; + }, Long.MAX_VALUE); + + assertTrue(writer.canWrite()); + writer.write("r1", newRecord("r1"), RECORD_SCHEMA, PROPS); + writer.writeWithMetadata(new HoodieKey("r2", "p"), newRecord("r2"), RECORD_SCHEMA, PROPS); + writer.write("r3", newRecord("r3"), RECORD_SCHEMA, PROPS); + assertTrue(delegate.calls.isEmpty()); + + writer.close(); + assertEquals(1, factoryCalls.size()); + assertSame(inferred, factoryCalls.get(0)); + assertEquals(Arrays.asList("plain:r1", "meta:r2", "plain:r3"), delegate.calls); + assertEquals(1, delegate.closeCount, "the delegate must be closed exactly once"); + + // Idempotent close + writer.close(); + assertEquals(1, factoryCalls.size()); + assertEquals(1, delegate.closeCount); + } + + @Test + public void testRecordCountThresholdTriggersMaterialization() throws IOException { + RecordingWriter delegate = new RecordingWriter(); + List<Map<String, HoodieSchema>> factoryCalls = new ArrayList<>(); + VariantShreddingInferenceFileWriter<Object> writer = writer( + (columns, samples) -> { + assertEquals(VariantShreddingInferenceFileWriter.MAX_BUFFERED_RECORDS, samples.size()); + return Collections.emptyMap(); + }, + map -> { + factoryCalls.add(map); + return delegate; + }, Long.MAX_VALUE); + + for (int i = 0; i < VariantShreddingInferenceFileWriter.MAX_BUFFERED_RECORDS; i++) { + writer.write("r" + i, newRecord("r" + i), RECORD_SCHEMA, PROPS); + } + // Threshold reached: delegate created and buffer replayed before close. + assertEquals(1, factoryCalls.size()); + assertEquals(VariantShreddingInferenceFileWriter.MAX_BUFFERED_RECORDS, delegate.calls.size()); + + // Subsequent writes stream straight through. + writer.write("tail", newRecord("tail"), RECORD_SCHEMA, PROPS); + assertEquals(VariantShreddingInferenceFileWriter.MAX_BUFFERED_RECORDS + 1, delegate.calls.size()); + writer.close(); + assertEquals(1, factoryCalls.size()); + } + + @Test + public void testByteCapTriggersEarlyMaterialization() throws IOException { + List<Map<String, HoodieSchema>> factoryCalls = new ArrayList<>(); + VariantShreddingInferenceFileWriter<Object> writer = writer((columns, samples) -> Collections.emptyMap(), + map -> { + factoryCalls.add(map); + return new RecordingWriter(); + }, 1L); + + writer.write("r1", newRecord("r1"), RECORD_SCHEMA, PROPS); + // A 1-byte cap is exceeded by any record. + assertEquals(1, factoryCalls.size()); + writer.close(); + } + + @Test + public void testInferrerFailureDeclinesAndWritesUnshredded() throws IOException { + RecordingWriter delegate = new RecordingWriter(); + List<Map<String, HoodieSchema>> factoryCalls = new ArrayList<>(); + VariantShreddingInferenceFileWriter<Object> writer = writer( + (columns, samples) -> { + throw new IllegalStateException("malformed variant"); + }, + map -> { + factoryCalls.add(map); + return delegate; + }, Long.MAX_VALUE); + + writer.write("r1", newRecord("r1"), RECORD_SCHEMA, PROPS); + writer.close(); + + assertEquals(1, factoryCalls.size()); + assertTrue(factoryCalls.get(0).isEmpty()); + assertEquals(singletonList("plain:r1"), delegate.calls); + assertEquals(1, delegate.closeCount); + } + + @Test + public void testZeroRecordCloseStillCreatesDelegate() throws IOException { + RecordingWriter delegate = new RecordingWriter(); + List<Map<String, HoodieSchema>> factoryCalls = new ArrayList<>(); + VariantShreddingInferenceFileWriter<Object> writer = writer( + (columns, samples) -> { + throw new AssertionError("inferrer must not be called with an empty buffer"); + }, + map -> { + factoryCalls.add(map); + return delegate; + }, Long.MAX_VALUE); + + writer.close(); + assertEquals(1, factoryCalls.size()); + assertTrue(factoryCalls.get(0).isEmpty()); + assertEquals(1, delegate.closeCount); + } + + @Test + public void testWriterCreationFailureIsLatchedAndRethrown() throws IOException { + IOException boom = new IOException("create failed"); + VariantShreddingInferenceFileWriter<Object> writer = writer((columns, samples) -> Collections.emptyMap(), + map -> { + throw boom; + }, Long.MAX_VALUE); + + writer.write("r1", newRecord("r1"), RECORD_SCHEMA, PROPS); + IOException fromClose = assertThrows(IOException.class, writer::close); + assertSame(boom, fromClose); + // Every subsequent call keeps failing: buffered records were never written. + IOException fromWrite = assertThrows(IOException.class, + () -> writer.write("r2", newRecord("r2"), RECORD_SCHEMA, PROPS)); + assertSame(boom, fromWrite); + } + + @Test + public void testSamplesAlignWithBufferedRecords() throws IOException { + // Snapshot: the decorator's internal list is cleared after replay. + List<List<VariantSample[]>> seenSamples = new ArrayList<>(); + VariantShreddingInferenceFileWriter.VariantSampleExtractor extractor = (record, schema, props) -> { + VariantSample[] samples = new VariantSample[1]; + samples[0] = new VariantSample(new byte[] {1}, new byte[] {2}); + return samples; + }; + VariantShreddingInferenceFileWriter<Object> writer = new VariantShreddingInferenceFileWriter<>( + singletonList("v"), extractor, (columns, samples) -> { + seenSamples.add(new ArrayList<>(samples)); + assertEquals(singletonList("v"), columns); + return Collections.emptyMap(); + }, + map -> new RecordingWriter(), Long.MAX_VALUE); + + writer.write("r1", newRecord("r1"), RECORD_SCHEMA, PROPS); + writer.write("r2", newRecord("r2"), RECORD_SCHEMA, PROPS); + writer.close(); + + assertEquals(1, seenSamples.size()); + assertEquals(2, seenSamples.get(0).size()); + assertNotNull(seenSamples.get(0).get(0)[0]); + assertEquals(1, seenSamples.get(0).get(0)[0].getValue()[0]); + } + + @Test + public void testCanWriteDelegatesAfterMaterialization() throws IOException { + VariantShreddingInferenceFileWriter<Object> writer = writer((columns, samples) -> Collections.emptyMap(), + map -> new RecordingWriter() { + @Override + public boolean canWrite() { + return false; + } + }, 1L); + + assertTrue(writer.canWrite()); + writer.write("r1", newRecord("r1"), RECORD_SCHEMA, PROPS); + assertFalse(writer.canWrite()); + writer.close(); + } + + @Test + public void testNullInferredMapTreatedAsDecline() throws IOException { + List<Map<String, HoodieSchema>> factoryCalls = new ArrayList<>(); + VariantShreddingInferenceFileWriter<Object> writer = writer((columns, samples) -> null, + map -> { + factoryCalls.add(map); + return new RecordingWriter(); + }, Long.MAX_VALUE); + + writer.write("r1", newRecord("r1"), RECORD_SCHEMA, PROPS); + writer.close(); + assertEquals(1, factoryCalls.size()); + assertNotNull(factoryCalls.get(0)); + assertTrue(factoryCalls.get(0).isEmpty()); + } + + @Test + public void testWriteRowMaterializesAndPassesThrough() throws IOException { + RecordingWriter delegate = new RecordingWriter(); + List<Map<String, HoodieSchema>> factoryCalls = new ArrayList<>(); + VariantShreddingInferenceFileWriter<Object> writer = writer((columns, samples) -> Collections.emptyMap(), + map -> { + factoryCalls.add(map); + return delegate; + }, Long.MAX_VALUE); + + writer.write("r1", newRecord("r1"), RECORD_SCHEMA, PROPS); + assertTrue(factoryCalls.isEmpty()); + // A raw row has nothing to sample from: the buffered records are replayed first, then the + // row goes straight through, preserving arrival order. + writer.writeRow("r2", new Object()); + assertEquals(1, factoryCalls.size()); + assertEquals(Arrays.asList("plain:r1", "row:r2"), delegate.calls); + writer.close(); + assertEquals(1, factoryCalls.size()); + } + + @Test + public void testFooterMetadataQueuedUntilMaterialization() throws IOException { + RecordingWriter delegate = new RecordingWriter(); + // A 1-byte cap materializes on the first write, so the forwarded leg below runs on an open writer. + VariantShreddingInferenceFileWriter<Object> writer = writer((columns, samples) -> Collections.emptyMap(), + map -> delegate, 1L); + + writer.addFooterMetadata(Collections.singletonMap("k1", "v1")); + assertTrue(delegate.footerMetadata.isEmpty(), "queued until the real writer exists"); + writer.write("r1", newRecord("r1"), RECORD_SCHEMA, PROPS); + assertEquals("v1", delegate.footerMetadata.get("k1"), "handed over at materialization"); + + // After materialization the call is forwarded directly. + writer.addFooterMetadata(Collections.singletonMap("k2", "v2")); + assertEquals("v2", delegate.footerMetadata.get("k2")); + writer.close(); + } + + @Test + public void testGetFileFormatMetadataMaterializesAndDelegates() throws IOException { + RecordingWriter delegate = new RecordingWriter(); + List<Map<String, HoodieSchema>> factoryCalls = new ArrayList<>(); + VariantShreddingInferenceFileWriter<Object> writer = writer((columns, samples) -> Collections.emptyMap(), + map -> { + factoryCalls.add(map); + return delegate; + }, Long.MAX_VALUE); + + writer.write("r1", newRecord("r1"), RECORD_SCHEMA, PROPS); + assertTrue(factoryCalls.isEmpty()); + // Footer metadata lives in the real writer, so asking for it creates that writer first. + assertSame(delegate.fileFormatMetadata, writer.getFileFormatMetadata()); + assertEquals(1, factoryCalls.size()); + assertEquals(singletonList("plain:r1"), delegate.calls); + + // The native log-format writer asks after close() (column stats): still the delegate's answer. + writer.close(); + assertSame(delegate.fileFormatMetadata, writer.getFileFormatMetadata()); + assertEquals(1, factoryCalls.size()); + } + + @Test + public void testReplayFailureIsLatchedAndRethrown() throws IOException { + RecordingWriter delegate = new RecordingWriter(); + IOException boom = new IOException("replay failed"); + delegate.failWriteWith = boom; + VariantShreddingInferenceFileWriter<Object> writer = writer((columns, samples) -> Collections.emptyMap(), + map -> delegate, Long.MAX_VALUE); + + writer.write("r1", newRecord("r1"), RECORD_SCHEMA, PROPS); + IOException fromClose = assertThrows(IOException.class, writer::close); + assertSame(boom, fromClose); + // The delegate was created but never closed by the try path, so the catch path closes it once. + assertEquals(1, delegate.closeCount); + // Latched: the buffered record was never written, so every later call keeps failing. + assertSame(boom, assertThrows(IOException.class, () -> writer.write("r2", newRecord("r2"), RECORD_SCHEMA, PROPS))); + assertSame(boom, assertThrows(HoodieIOException.class, writer::getFileFormatMetadata).getCause()); + } + + @Test + public void testReplayErrorIsLatchedTooAndCloseDoesNotFinishTheFile() throws IOException { + // An Error mid-replay latches like an exception does: inference already treats a LinkageError + // as reachable (a writer linked against another Spark than the runtime's), and an unlatched + // one would let close() finish the file without the records left in the buffer. + RecordingWriter delegate = new RecordingWriter(); + NoClassDefFoundError boom = new NoClassDefFoundError("replay failed"); + delegate.failWriteWith = boom; + // A 1-byte cap materializes on the first write, so the Error surfaces from write(), not close(). + VariantShreddingInferenceFileWriter<Object> writer = writer((columns, samples) -> Collections.emptyMap(), + map -> delegate, 1L); + + assertSame(boom, assertThrows(NoClassDefFoundError.class, + () -> writer.write("r1", newRecord("r1"), RECORD_SCHEMA, PROPS))); + assertSame(boom, assertThrows(IOException.class, writer::close).getCause()); + assertEquals(1, delegate.closeCount); + } + + @Test + public void testMaterializeErrorInsideCloseStillClosesTheDelegate() throws IOException { + // With the caps never tripped, the first materialization happens inside close(): the Error + // must still close the delegate created just above, or the file handle leaks. + RecordingWriter delegate = new RecordingWriter(); + NoClassDefFoundError boom = new NoClassDefFoundError("replay failed"); + delegate.failWriteWith = boom; + VariantShreddingInferenceFileWriter<Object> writer = writer((columns, samples) -> Collections.emptyMap(), + map -> delegate, Long.MAX_VALUE); + + writer.write("r1", newRecord("r1"), RECORD_SCHEMA, PROPS); + assertSame(boom, assertThrows(NoClassDefFoundError.class, writer::close)); + assertEquals(1, delegate.closeCount); + } + + @Test + public void testThrowingDelegateCloseSurfacesAndIsNotRetried() throws IOException { + RecordingWriter delegate = new RecordingWriter(); + IOException boom = new IOException("close failed"); + delegate.failCloseWith = boom; + VariantShreddingInferenceFileWriter<Object> writer = writer((columns, samples) -> Collections.emptyMap(), + map -> delegate, Long.MAX_VALUE); + + writer.write("r1", newRecord("r1"), RECORD_SCHEMA, PROPS); + assertSame(boom, assertThrows(IOException.class, writer::close)); + assertEquals(1, delegate.closeCount, "a throwing delegate.close() must surface, not be retried"); + } + + @Test + public void testPreparedRecordIsSampledAndReplayed() throws IOException { + // An extractor that materializes the record (the Avro one) hands the materialized form back + // via prepare(); the decorator samples that form and replays it, so the writer never redoes + // the materialization. + HoodieRecord prepared = newRecord("prepared"); + List<HoodieRecord> sampled = new ArrayList<>(); + VariantShreddingInferenceFileWriter.VariantSampleExtractor extractor = + new VariantShreddingInferenceFileWriter.VariantSampleExtractor() { + @Override + public VariantSample[] extract(HoodieRecord record, HoodieSchema schema, Properties props) { + sampled.add(record); + return new VariantSample[1]; + } + + @Override + public HoodieRecord prepare(HoodieRecord record, HoodieSchema schema, Properties props) { + return prepared; + } + }; + RecordingWriter delegate = new RecordingWriter(); + VariantShreddingInferenceFileWriter<Object> writer = new VariantShreddingInferenceFileWriter<>( + singletonList("v"), extractor, (columns, samples) -> Collections.emptyMap(), + map -> delegate, Long.MAX_VALUE); + + writer.write("r1", newRecord("r1"), RECORD_SCHEMA, PROPS); + writer.close(); + + assertEquals(singletonList(prepared), sampled); + assertEquals(singletonList(prepared), delegate.writtenRecords); + } + + @Test + public void testByteCapAccumulatesThroughTheEstimator() throws IOException { + // Same-shaped records estimate the same size, so a cap of 150 records' worth materializes on + // exactly the 150th write, after passing through the periodic re-estimation at record 100. + // That re-estimation rescales the whole buffer, so the moving average's long truncation (at + // most a byte) is charged to all 150 records at once; the slack covers that while staying + // well under one record, which is why the records are kilobyte-sized. + String padding = new String(new char[1024]).replace('\0', 'x'); + long perRecord = new DefaultSizeEstimator<HoodieRecord>().sizeEstimate(newRecord("r000" + padding)); + assertTrue(perRecord > 1000, "expected a kilobyte-sized record, got " + perRecord); + List<Map<String, HoodieSchema>> factoryCalls = new ArrayList<>(); + VariantShreddingInferenceFileWriter<Object> writer = writer((columns, samples) -> Collections.emptyMap(), + map -> { + factoryCalls.add(map); + return new RecordingWriter(); + }, 150 * perRecord - 500); + + for (int i = 0; i < 149; i++) { + writer.write("r" + i, newRecord(String.format("r%03d", i) + padding), RECORD_SCHEMA, PROPS); Review Comment: Every record here estimates the same size, so the moving average never moves and the cap trips on the 150th write with or without the `bufferedBytes` rescale; the row-writer twin has the same shape. Sizing the first 99 records well below the 100th would make both fail without the rescale. -- 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]
