shunping commented on code in PR #40142: URL: https://github.com/apache/beam/pull/40142#discussion_r4043108520
########## it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/storage/ParquetIOLT.java: ########## @@ -0,0 +1,862 @@ +/* + * 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.beam.it.gcp.storage; + +import static org.apache.beam.it.truthmatchers.PipelineAsserts.assertThatResult; +import static org.junit.Assert.assertEquals; + +import com.fasterxml.jackson.annotation.JsonProperty; +import com.fasterxml.jackson.core.JsonParser; +import com.fasterxml.jackson.core.JsonToken; +import com.fasterxml.jackson.databind.DeserializationContext; +import com.fasterxml.jackson.databind.JsonDeserializer; +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.annotation.JsonDeserialize; +import com.fasterxml.jackson.databind.node.ObjectNode; +import java.io.IOException; +import java.nio.ByteBuffer; +import java.time.Duration; +import java.time.ZoneOffset; +import java.time.format.DateTimeFormatter; +import java.util.Locale; +import java.util.Map; +import java.util.Random; +import java.util.UUID; +import java.util.regex.Matcher; +import java.util.regex.Pattern; +import org.apache.avro.Schema; +import org.apache.avro.SchemaBuilder; +import org.apache.avro.generic.GenericRecord; +import org.apache.avro.generic.GenericRecordBuilder; +import org.apache.beam.it.common.PipelineLauncher; +import org.apache.beam.it.common.PipelineOperator; +import org.apache.beam.it.common.TestProperties; +import org.apache.beam.it.common.dataflow.DefaultPipelineLauncher.PipelineMetricsType; +import org.apache.beam.it.common.storage.GcsIOLoadTestBase; +import org.apache.beam.it.common.storage.GcsResourceManager; +import org.apache.beam.it.common.utils.ResourceManagerUtils; +import org.apache.beam.sdk.extensions.avro.coders.AvroCoder; +import org.apache.beam.sdk.extensions.gcp.options.GcsOptions; +import org.apache.beam.sdk.io.FileIO; +import org.apache.beam.sdk.io.GenerateSequence; +import org.apache.beam.sdk.io.parquet.ParquetIO; +import org.apache.beam.sdk.io.synthetic.SyntheticSourceOptions; +import org.apache.beam.sdk.testing.TestPipeline; +import org.apache.beam.sdk.transforms.Create; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Strings; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; +import org.apache.parquet.hadoop.metadata.CompressionCodecName; +import org.junit.AfterClass; +import org.junit.BeforeClass; +import org.junit.FixMethodOrder; +import org.junit.Rule; +import org.junit.Test; +import org.junit.runners.MethodSorters; + +/** + * ParquetIO performance tests on Google Cloud Storage. + * + * <p>Reads and writes are kept in two separate pipelines / tests: + * + * <ul> + * <li>{@code test1ParquetWrite} generates records of the configured shape and writes them as + * Parquet files under {@code outputPrefix}. + * <li>{@code test2ParquetRead} reads all the Parquet files matching {@code inputFilePattern}, + * optionally projecting only the first {@code numFieldsToRead} fields. + * </ul> + * + * <p>The methods are ordered by name so that the write test runs first, and the dataset it produces + * is reused by the read test: running the whole class generates the dataset only once. Running the + * read test on its own still works, it generates the dataset itself. + * + * <p>Both tests enable {@code --gcsPerformanceMetrics=true} so that the {@code gcs_*} client + * metrics collected by {@link GcsIOLoadTestBase} are exported along with the runner metrics. + * + * <h3>Workload shape</h3> + * + * <p>The workload is described by two dimensions: {@code numFields} (how many columns a record has) + * and {@code maxFieldSizeBytes} (how large a single value is). Everything else is held constant so + * that runs of different shapes stay comparable: + * + * <ul> + * <li>{@code compressibility = 0.0}, i.e. incompressible payloads, so the bytes written to GCS + * match the configured dataset size. Compressible payloads would make the test measure the + * Parquet codec rather than the GCS client. + * <li>{@code compressionCodec = UNCOMPRESSED}, for the same reason. + * <li>{@code numShards} pinned, so that the number and the size of the GCS objects is identical + * across runs. + * <li>The Dataflow worker pool is pinned: autoscaling off, 3 workers, {@code e2-standard-2}. An + * autoscaled pool would give a cheap shape fewer workers than an expensive one, so the + * throughput of the two could not be compared. + * </ul> + * + * <h3>Configuration</h3> + * + * <p>{@code -Dconfiguration} takes either the name of a preset, or a json object. The json object + * may name a base preset with a {@code "preset"} property and override any of its values, so that + * one preset can be reused for several runs: + * + * <pre> + * # a preset as is + * -Dconfiguration=f100_s16 + * + * # the same shape, but a cheap local run + * -Dconfiguration='{"preset":"f100_s16","runner":"DirectRunner","totalBytes":"10MB"}' + * + * # the same shape, reading only the first field of each record + * -Dconfiguration='{"preset":"f100_s16","numFieldsToRead":1}' + * + * # no preset at all, every unset value falls back to the Configuration defaults + * -Dconfiguration='{"numFields":10,"maxFieldSizeBytes":"1KB","totalBytes":"1GB"}' + * </pre> + * + * <p>Every byte count, i.e. {@code totalBytes}, {@code maxFieldSizeBytes}, {@code + * minFieldSizeBytes} and {@code rowGroupSize}, is either a plain number of bytes or a size string + * such as {@code "10GB"}, {@code "500MB"}, {@code "64KB"} or {@code "32B"}. The units are binary, + * so {@code 1KB} is 1024 bytes, and {@code M}, {@code MB} and {@code MiB} are all accepted. + * + * <p>Example trigger command: + * + * <pre> + * ./gradlew :it:google-cloud-platform:ParquetPerformanceTest -Dconfiguration=f100_s16 \ + * -Dproject=[gcpProject] -DartifactBucket=[temp bucket] + * </pre> + * + * <p>The gradle task always passes {@code configuration} down, defaulting to {@code local}, so + * leaving the flag out runs a small local pipeline rather than a full scale one. Every run against + * a real runner has to name its preset explicitly. + */ +@FixMethodOrder(MethodSorters.NAME_ASCENDING) +public final class ParquetIOLT extends GcsIOLoadTestBase { + + private static final String READ_ELEMENT_METRIC_NAME = "read_count"; + + private static final ObjectMapper MAPPER = new ObjectMapper(); + + private static final String DATAFLOW_RUNNER = "DataflowRunner"; + + /** + * Size of the worker pool every Dataflow run gets. Frozen, see {@link #launchConfig}: the runs + * are only comparable if they all have the same amount of cpu, memory and network bandwidth. + */ + private static final int DATAFLOW_NUM_WORKERS = 3; + + /** Machine type every Dataflow worker runs on. Frozen for the same reason. */ + private static final String DATAFLOW_MACHINE_TYPE = "e2-standard-2"; + + /** + * Dataset size every shape preset generates, so that the shapes are comparable. {@link + * #parseSizeToBytes} is binary, so this is 42,949,672,960 bytes. + */ + private static final String MATRIX_TOTAL_BYTES = "40GB"; + + /** + * Wall clock budget for a matrix run. Sized from measured runs: 10GB takes roughly 5 minutes of + * worker time on the pinned pool, so 40GB needs about 20, and the read pipeline has to stay long + * enough for Cloud Monitoring to have ingested more than just its last data point. + */ + private static final int MATRIX_PIPELINE_TIMEOUT_MINUTES = 60; + + /** + * Presets, kept as json so that a caller can name one as a base and override parts of it. The + * {@code f<numFields>_s<maxFieldSize>} presets are the cells of the workload matrix. + */ + private static final Map<String, String> TEST_CONFIGS_PRESET = + ImmutableMap.<String, String>builder() + // Small run against the Configuration defaults, for local development. + .put("local", "{}") + // Legacy size presets: a single field, the shape the test used to have. + .put("medium", shape(1, "750B", "7500MB", 20)) + .put("large", shape(1, "750B", "75GB", 80)) + // Cells of the workload matrix. + .put("f1_s1k", shape(1, "1KB", MATRIX_TOTAL_BYTES, MATRIX_PIPELINE_TIMEOUT_MINUTES)) + .put("f10_s1k", shape(10, "1KB", MATRIX_TOTAL_BYTES, MATRIX_PIPELINE_TIMEOUT_MINUTES)) + .put("f100_s16", shape(100, "16B", MATRIX_TOTAL_BYTES, MATRIX_PIPELINE_TIMEOUT_MINUTES)) + .put("f1000_s16", shape(1000, "16B", MATRIX_TOTAL_BYTES, MATRIX_PIPELINE_TIMEOUT_MINUTES)) + .put("f100_s1k", shape(100, "1KB", MATRIX_TOTAL_BYTES, MATRIX_PIPELINE_TIMEOUT_MINUTES)) + .put("f1000_s1k", shape(1000, "1KB", MATRIX_TOTAL_BYTES, MATRIX_PIPELINE_TIMEOUT_MINUTES)) + .put("f10_s64k", shape(10, "64KB", MATRIX_TOTAL_BYTES, MATRIX_PIPELINE_TIMEOUT_MINUTES)) + // Blob column: a page size check every 100 rows would buffer 400 MB, so check every row. + .put( + "f1_s4m", + "{\"numFields\":1,\"maxFieldSizeBytes\":\"4MB\",\"totalBytes\":\"" + + MATRIX_TOTAL_BYTES + + "\",\"minRowCountForPageSizeCheck\":1,\"numShards\":64,\"compressibility\":0.0," + + "\"compressionCodec\":\"UNCOMPRESSED\",\"runner\":\"DataflowRunner\"," + + "\"pipelineTimeout\":" + + MATRIX_PIPELINE_TIMEOUT_MINUTES + + "}") + .build(); + + private static GcsResourceManager resourceManager; + + /** + * Prefix the write pipeline writes to, also used as read input when none is configured. Static so + * that both tests share a single dataset. + */ + private static String outputPrefix; + + private static Configuration configuration; + + /** Schema of the generated records, derived from {@code numFields}. */ + private static Schema schema; + + /** Whether {@code outputPrefix} already holds a dataset written by this class. */ + private static boolean datasetWritten; + + @Rule public TestPipeline writePipeline = TestPipeline.create(); + + @Rule public TestPipeline readPipeline = TestPipeline.create(); + + /** Returns the json of a shape preset, with all the frozen knobs pinned. */ + private static String shape( + int numFields, String maxFieldSize, String totalSize, int pipelineTimeout) { + return String.format( + "{\"numFields\":%d,\"maxFieldSizeBytes\":\"%s\",\"totalBytes\":\"%s\",\"numShards\":64," + + "\"compressibility\":0.0,\"compressionCodec\":\"UNCOMPRESSED\"," + + "\"runner\":\"DataflowRunner\",\"pipelineTimeout\":%d}", + numFields, maxFieldSize, totalSize, pipelineTimeout); + } + + /** + * Resolves the configuration and the dataset location once for the whole class, so that the write + * test and the read test operate on the same files. + */ + @BeforeClass + public static void beforeClass() { + resourceManager = + GcsResourceManager.builder(TestProperties.artifactBucket(), "parquetiolt", CREDENTIALS) + .build(); + + String testConfig = + TestProperties.getProperty("configuration", "local", TestProperties.Type.PROPERTY); + configuration = resolveConfiguration(testConfig); + validateAndDerive(configuration); + schema = buildSchema(configuration.numFields); + datasetWritten = false; + + if (!Strings.isNullOrEmpty(configuration.outputPrefix)) { + outputPrefix = configuration.outputPrefix; + } else { + String tempDirName = + "parquetiolt-" + + DateTimeFormatter.ofPattern("MMddHHmmssSSS") + .withZone(ZoneOffset.UTC) + .format(java.time.Instant.now()) + + UUID.randomUUID().toString().substring(0, 10); + resourceManager.registerTempDir(tempDirName); + outputPrefix = + String.format("gs://%s/%s/parquet", TestProperties.artifactBucket(), tempDirName); + } + printConfiguration(); + } + + @AfterClass + public static void tearDownClass() { + ResourceManagerUtils.cleanResources(resourceManager); + } + + /** Writes the configured number of records under the configured output prefix. */ + @Test + public void test1ParquetWrite() throws IOException { + PipelineLauncher.LaunchInfo writeInfo = runWritePipeline(outputPrefix); + + printMetrics( + writeInfo, + MetricsConfiguration.builder() + .setInputPCollection("Create avro records.out0") + .setInputPCollectionV2("Create avro records/ParMultiDo(CreateAvroRecord).out0") + .build()); + } + + /** Reads all the Parquet files matching the configured input file pattern. */ + @Test + public void test2ParquetRead() throws IOException { + String inputFilePattern = configuration.inputFilePattern; + long expectedRecords = configuration.numRecords; + if (Strings.isNullOrEmpty(inputFilePattern)) { + if (!datasetWritten) { + // No dataset given and the write test did not run: generate one so that the read test is + // self contained. runWritePipeline already waits for the job and asserts it succeeded. + runWritePipeline(outputPrefix); + } + inputFilePattern = outputPrefix + "*"; + } + + PCollection<FileIO.ReadableFile> files = + readPipeline + .apply("Create filepattern", Create.of(inputFilePattern)) + .apply("Match all files", FileIO.matchAll()) + .apply("Read matches", FileIO.readMatches()); + + PCollection<GenericRecord> records; + if (configuration.numFieldsToRead > 0 + && configuration.numFieldsToRead < configuration.numFields) { + // Column projection: only the leading fields are fetched from the Parquet files, which is + // what turns a sequential scan into many small ranged GETs. + Schema projection = buildSchema(configuration.numFieldsToRead); + records = + files.apply( + "Read parquet files", + ParquetIO.readFiles(schema).withProjection(projection, projection)); + } else { + records = files.apply("Read parquet files", ParquetIO.readFiles(schema)); + } + records.apply("Counting element", ParDo.of(new CountingFn<>(READ_ELEMENT_METRIC_NAME))); + + PipelineLauncher.LaunchInfo readInfo = + pipelineLauncher.launch(project, region, launchConfig("read-parquet", readPipeline)); + PipelineOperator.Result readResult = + pipelineOperator.waitUntilDone( + createConfig(readInfo, Duration.ofMinutes(configuration.pipelineTimeout))); + + // Fail the test if the pipeline failed or timed out. + assertThatResult(readResult).isLaunchFinished(); + + // Only assert the record count when we know how many records the dataset holds. + if (Strings.isNullOrEmpty(configuration.inputFilePattern)) { + double numRecords = + pipelineLauncher.getMetric( + project, + region, + readInfo.jobId(), + getBeamMetricsName(PipelineMetricsType.COUNTER, READ_ELEMENT_METRIC_NAME)); + assertEquals((double) expectedRecords, numRecords, 0.5); + } + + printMetrics( + readInfo, + MetricsConfiguration.builder() + .setOutputPCollection("Counting element.out0") + .setOutputPCollectionV2("Counting element/ParMultiDo(Counting).out0") + .build()); + } + + private PipelineLauncher.LaunchInfo runWritePipeline(String prefix) throws IOException { + ParquetIO.Sink sink = + ParquetIO.sink(schema) + .withCompressionCodec(CompressionCodecName.fromConf(configuration.compressionCodec)); + if (configuration.rowGroupSize > 0) { + sink = sink.withRowGroupSize(configuration.rowGroupSize); + } + if (configuration.minRowCountForPageSizeCheck > 0) { + // With large values the default of a page size check every 100 rows buffers far too much. + sink = sink.withMinRowCountForPageSizeCheck(configuration.minRowCountForPageSizeCheck); + } + + // FileIO.write().to(...) expects a directory, so the prefix is split into the directory the + // files are written to and the base name each file starts with. This way the written files are + // "<prefix>-0000i-of-0000n.parquet" and can be matched back with "<prefix>*" by the read test. + FileIO.Write<Void, GenericRecord> write = + FileIO.<GenericRecord>write() + .via(sink) + .to(directoryOf(prefix)) + .withNaming(FileIO.Write.defaultNaming(baseNameOf(prefix), ".parquet")); + if (configuration.numShards > 0) { + write = write.withNumShards(configuration.numShards); + } + + PCollection<GenericRecord> records = + writePipeline + .apply("Generate sequence", GenerateSequence.from(0).to(configuration.numRecords)) + .apply( + "Create avro records", + ParDo.of( + new CreateAvroRecordFn( + schema.toString(), + configuration.numFields, + (int) configuration.minFieldSizeBytes, + (int) configuration.maxFieldSizeBytes, + configuration.compressibility))) + .setCoder(AvroCoder.of(schema)); + records.apply("Write parquet files", write); + + PipelineLauncher.LaunchInfo writeInfo = + pipelineLauncher.launch(project, region, launchConfig("write-parquet", writePipeline)); + PipelineOperator.Result writeResult = + pipelineOperator.waitUntilDone( + createConfig(writeInfo, Duration.ofMinutes(configuration.pipelineTimeout))); + + // Fail the test if the pipeline failed or timed out. + assertThatResult(writeResult).isLaunchFinished(); + // The dataset now exists under `prefix`, so the read test can reuse it instead of writing a + // second copy. + datasetWritten = true; + return writeInfo; + } + + private PipelineLauncher.LaunchConfig launchConfig(String jobName, TestPipeline pipeline) { + // The launcher only turns the parameters below into pipeline options for the DataflowRunner. + // For the other runners it runs the pipeline with the options it already has, so the flag has + // to be set explicitly here, otherwise no gcs_* metric is reported. + pipeline.getOptions().as(GcsOptions.class).setGcsPerformanceMetrics(true); + + PipelineLauncher.LaunchConfig.Builder builder = + PipelineLauncher.LaunchConfig.builder(jobName) + .setSdk(PipelineLauncher.Sdk.JAVA) + .setPipeline(pipeline) + .addParameter("runner", configuration.runner) + // Required for GcsUtil to report the gcs_* client metrics. + .addParameter(GCS_PERFORMANCE_METRICS_OPTION, "true"); + + if (DATAFLOW_RUNNER.equalsIgnoreCase(configuration.runner)) { + // The worker pool is pinned so that the runs of the different workload shapes are + // comparable: with autoscaling the service would give a shape that is cheap to process + // fewer workers than an expensive one, and the throughput of the two could not be compared. + // A fixed pool also keeps the number of parallel GCS connections constant, which is what + // the gcs_* metrics measure. + // maxNumWorkers is deliberately not set, it only bounds an autoscaling pool. + builder + .addParameter("autoscalingAlgorithm", "NONE") + .addParameter("numWorkers", String.valueOf(DATAFLOW_NUM_WORKERS)) + .addParameter("workerMachineType", DATAFLOW_MACHINE_TYPE); + } + + return builder.build(); + } + + /** + * Resolves {@code -Dconfiguration} into a {@link Configuration}. + * + * <p>The value is either the name of a preset, or a json object. A json object may select a base + * preset with a {@code "preset"} property, in which case the remaining properties override the + * ones of that preset. The merge is done on the json trees rather than on the deserialized + * objects, because {@link SyntheticSourceOptions} has final properties that cannot be written + * back. + */ + private static Configuration resolveConfiguration(String spec) { + String trimmed = spec.trim(); + try { + ObjectNode overrides; + if (trimmed.startsWith("{")) { + JsonNode parsed = MAPPER.readTree(trimmed); + if (!parsed.isObject()) { + throw new IllegalArgumentException( + String.format("Configuration json must be an object, but was: [%s]", trimmed)); + } + overrides = (ObjectNode) parsed; + } else { + overrides = MAPPER.createObjectNode().put("preset", trimmed); + } + + JsonNode preset = overrides.remove("preset"); + ObjectNode merged = + preset == null + ? MAPPER.createObjectNode() + : (ObjectNode) MAPPER.readTree(presetJson(preset.asText())); + merged.setAll(overrides); + + return Configuration.fromJsonString(merged.toString(), Configuration.class); + } catch (IOException e) { + throw new IllegalArgumentException( + String.format( + "Unable to parse test configuration: [%s]. Pass a valid configuration json, or one" + + " of the presets: %s", + trimmed, TEST_CONFIGS_PRESET.keySet()), + e); + } + } + + private static String presetJson(String name) { + String preset = TEST_CONFIGS_PRESET.get(name); + if (preset == null) { + throw new IllegalArgumentException( + String.format( + "Unknown preset: [%s]. Known presets: %s", name, TEST_CONFIGS_PRESET.keySet())); + } + return preset; + } + + /** Checks the configuration and fills in the values that are derived from the others. */ + private static void validateAndDerive(Configuration configuration) { + checkConfig(configuration.numFields > 0, "numFields must be positive"); + checkConfig(configuration.maxFieldSizeBytes > 0, "maxFieldSizeBytes must be positive"); + if (configuration.minFieldSizeBytes < 0) { + configuration.minFieldSizeBytes = configuration.maxFieldSizeBytes; + } + checkConfig( + configuration.minFieldSizeBytes <= configuration.maxFieldSizeBytes, + "minFieldSizeBytes must not be greater than maxFieldSizeBytes"); + checkConfig( + configuration.compressibility >= 0.0 && configuration.compressibility <= 1.0, + "compressibility must be within [0.0, 1.0]"); + checkConfig( + configuration.numFieldsToRead >= 0 + && configuration.numFieldsToRead <= configuration.numFields, + "numFieldsToRead must be within [0, numFields]"); + + if (configuration.totalBytes > 0) { + configuration.numRecords = + Math.max(1L, configuration.totalBytes / recordBytes(configuration)); + } + checkConfig( + configuration.numRecords > 0, + "numRecords is 0. Set either numRecords or totalBytes, otherwise the write pipeline is a" + + " no-op"); + } + + /** Average number of payload bytes of a record, ignoring the Parquet overhead. */ + private static long recordBytes(Configuration configuration) { + long avgFieldSize = (configuration.minFieldSizeBytes + configuration.maxFieldSizeBytes) / 2; + return Math.max(1L, configuration.numFields * avgFieldSize); + } + + private static void checkConfig(boolean condition, String message) { + if (!condition) { + throw new IllegalArgumentException(message); + } + } + + private static final Pattern SIZE_PATTERN = + Pattern.compile("^\\s*([0-9]+(?:\\.[0-9]+)?)\\s*([a-zA-Z]*)\\s*$"); + + /** + * Parses a size string such as {@code "10GB"}, {@code "500MB"}, {@code "64KB"}, {@code "32B"} or + * {@code "1024"} into a number of bytes. The units are binary, i.e. {@code 1KB == 1024}, and both + * the short and the long spelling are accepted ({@code M}, {@code MB}, {@code MiB}). + */ + static long parseSizeToBytes(String sizeStr) { + if (sizeStr == null || sizeStr.trim().isEmpty()) { + throw new IllegalArgumentException("Size string cannot be null or empty"); + } + Matcher matcher = SIZE_PATTERN.matcher(sizeStr.trim()); + if (!matcher.matches()) { + throw new IllegalArgumentException( + "Invalid size string: '" + + sizeStr + + "'. Expected something like '10GB', '500MB', '64KB', '32B' or '1024'."); + } + double value = Double.parseDouble(matcher.group(1)); + String unit = matcher.group(2).toUpperCase(Locale.ROOT); + + long multiplier; + switch (unit) { + case "": + case "B": + case "BYTES": + multiplier = 1L; + break; + case "K": + case "KB": + case "KIB": + multiplier = 1024L; + break; + case "M": + case "MB": + case "MIB": + multiplier = 1024L * 1024L; + break; + case "G": + case "GB": + case "GIB": + multiplier = 1024L * 1024L * 1024L; + break; + case "T": + case "TB": + case "TIB": + multiplier = 1024L * 1024L * 1024L * 1024L; + break; + default: + throw new IllegalArgumentException( + "Unsupported size unit '" + unit + "' in size string: " + sizeStr); + } + return (long) (value * multiplier); + } + + /** Formats a number of bytes as e.g. {@code 9.31 GB}. */ + private static String formatBytes(long bytes) { + String[] units = {"B", "KB", "MB", "GB", "TB"}; + double value = bytes; + int unit = 0; + while (value >= 1024.0 && unit < units.length - 1) { + value /= 1024.0; + unit++; + } + return unit == 0 + ? String.format("%d B", bytes) + : String.format("%,d B (%.2f %s)", bytes, value, units[unit]); + } + + /** + * Deserializes a byte count that is either a number or a size string such as {@code "10GB"}. It + * lets the configuration json stay readable: {@code "totalBytes":"10GB"} instead of {@code + * "totalBytes":10737418240}. + */ + static final class ByteSize extends JsonDeserializer<Long> { + @Override + public Long deserialize(JsonParser parser, DeserializationContext context) throws IOException { + JsonToken token = parser.currentToken(); + if (token == JsonToken.VALUE_NUMBER_INT || token == JsonToken.VALUE_NUMBER_FLOAT) { + return parser.getLongValue(); + } + return parseSizeToBytes(parser.getText()); + } + } + + /** Same as {@link ByteSize}, for the options that the Parquet API takes as an {@code int}. */ + static final class IntByteSize extends JsonDeserializer<Integer> { + @Override + public Integer deserialize(JsonParser parser, DeserializationContext context) + throws IOException { + JsonToken token = parser.currentToken(); + if (token == JsonToken.VALUE_NUMBER_INT || token == JsonToken.VALUE_NUMBER_FLOAT) { + return parser.getIntValue(); + } + return Math.toIntExact(parseSizeToBytes(parser.getText())); + } + } + + /** Prints the effective configuration, so that a run can be matched with its metrics. */ + private static void printConfiguration() { + System.out.printf( + "%n==========================================================%n" + + " TEST CONFIGURATION%n" + + "==========================================================%n" + + " numFields: %,d%n" + + " fieldSize: %s .. %s%n" + + " recordSize: %s%n" + + " numRecords: %,d%n" + + " logicalSize: %s%n" + + " compressibility: %.2f%n" + + " compressionCodec: %s%n" + + " rowGroupSize: %s%n" + + " numShards: %d%n" + + " numFieldsToRead: %s%n" + + " runner: %s%n" + + " workerPool: %s%n" + + "==========================================================%n%n", + configuration.numFields, + formatBytes(configuration.minFieldSizeBytes), + formatBytes(configuration.maxFieldSizeBytes), + formatBytes(recordBytes(configuration)), + configuration.numRecords, + formatBytes(configuration.numRecords * recordBytes(configuration)), + configuration.compressibility, + configuration.compressionCodec, + configuration.rowGroupSize > 0 ? formatBytes(configuration.rowGroupSize) : "default", + configuration.numShards, + configuration.numFieldsToRead > 0 ? String.valueOf(configuration.numFieldsToRead) : "all", + configuration.runner, + DATAFLOW_RUNNER.equalsIgnoreCase(configuration.runner) + ? String.format("%d x %s, autoscaling off", DATAFLOW_NUM_WORKERS, DATAFLOW_MACHINE_TYPE) + : "n/a"); + } + + /** Builds a record schema of {@code numFields} byte array fields named {@code f0..fN-1}. */ + private static Schema buildSchema(int numFields) { + SchemaBuilder.FieldAssembler<Schema> fields = + SchemaBuilder.record("TestAvroLine").namespace("ioitavro").fields(); + for (int i = 0; i < numFields; i++) { + fields = fields.name(fieldName(i)).type().bytesType().noDefault(); + } + return fields.endRecord(); + } + + private static String fieldName(int index) { + return "f" + index; + } + + /** Returns the directory part of a file prefix, e.g. {@code gs://bucket/dir/} for a prefix. */ + private static String directoryOf(String prefix) { + int lastSlash = prefix.lastIndexOf('/'); + return lastSlash < 0 ? prefix : prefix.substring(0, lastSlash + 1); + } + + /** Returns the file name part of a file prefix, e.g. {@code parquet} for {@code .../parquet}. */ + private static String baseNameOf(String prefix) { + int lastSlash = prefix.lastIndexOf('/'); + String baseName = lastSlash < 0 ? prefix : prefix.substring(lastSlash + 1); + return baseName.isEmpty() ? "output" : baseName; + } + + /** Turns a sequence number into a record of the configured shape. */ + private static final class CreateAvroRecordFn extends DoFn<Long, GenericRecord> { + // Schema is not serializable, so it is carried as json and parsed on the worker. + private final String schemaJson; + private final int numFields; + private final int minFieldSizeBytes; + private final int maxFieldSizeBytes; + private final double compressibility; + + private transient Schema schema; + + CreateAvroRecordFn( + String schemaJson, + int numFields, + int minFieldSizeBytes, + int maxFieldSizeBytes, + double compressibility) { + this.schemaJson = schemaJson; + this.numFields = numFields; + this.minFieldSizeBytes = minFieldSizeBytes; + this.maxFieldSizeBytes = maxFieldSizeBytes; + this.compressibility = compressibility; + } + + @Setup + public void setup() { + schema = new Schema.Parser().parse(schemaJson); + } + + @ProcessElement + public void processElement(@Element Long element, OutputReceiver<GenericRecord> receiver) { + // Seeded with the element so that a record always holds the same content, whatever the + // runner decides to retry. + Random random = new Random(element); + GenericRecordBuilder builder = new GenericRecordBuilder(schema); + for (int i = 0; i < numFields; i++) { + int size = + minFieldSizeBytes == maxFieldSizeBytes + ? maxFieldSizeBytes + : minFieldSizeBytes + random.nextInt(maxFieldSizeBytes - minFieldSizeBytes + 1); + builder.set(fieldName(i), ByteBuffer.wrap(payload(random, size))); + } + receiver.output(builder.build()); + } + + /** + * Returns {@code size} bytes of which a {@code 1 - compressibility} fraction is random. The + * remaining bytes are left at zero, which is what the Parquet codec can collapse. + */ + private byte[] payload(Random random, int size) { Review Comment: Fixed. -- 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]
