voonhous commented on code in PR #18961: URL: https://github.com/apache/hudi/pull/18961#discussion_r3843797461
########## 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: Reshaped both twins: 99 small records, then a big 100th, with a cap of one big record. The rescale charges the earlier 99 at the grown estimate (about 90 small + 10 big) and meets the cap; without it the buffer holds about 100 small + 0.1 big and stays under. Red-checked on both writers. -- 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]
