shunping commented on code in PR #40142: URL: https://github.com/apache/beam/pull/40142#discussion_r4042886730
########## 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, Review Comment: Fixed. ########## it/common/src/test/java/org/apache/beam/it/common/storage/GcsIOLoadTestBase.java: ########## @@ -0,0 +1,324 @@ +/* + * 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.common.storage; + +import java.io.IOException; +import java.text.ParseException; +import java.util.Locale; +import java.util.Map; +import java.util.TreeMap; +import org.apache.beam.it.common.PipelineLauncher; +import org.apache.beam.it.common.dataflow.DefaultPipelineLauncher; +import org.apache.beam.it.common.dataflow.IOLoadTestBase; +import org.apache.beam.sdk.PipelineResult; +import org.apache.beam.sdk.metrics.DistributionResult; +import org.apache.beam.sdk.metrics.MetricQueryResults; +import org.apache.beam.sdk.metrics.MetricResult; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Base class for GCS IO load tests. + * + * <p>In addition to the runner/pipeline metrics collected by {@link IOLoadTestBase}, this class + * collects the GCS client performance metrics (all counters and distributions whose name starts + * with {@value #GCS_METRIC_PREFIX}) that are emitted by {@code GcsUtil} when the pipeline is run + * with {@code --gcsPerformanceMetrics=true}. + * + * <p>These metrics are regular Beam SDK metrics registered under the {@code GcsHttp} namespace, so + * they are collected in a runner agnostic way through {@code + * PipelineResult.metrics().allMetrics()}. Examples of collected metrics include: + * + * <ul> + * <li>{@code gcs_http_read_wire_bytes_received} / {@code gcs_http_write_wire_bytes_sent} + * <li>{@code gcs_http_read_request_count} / {@code gcs_http_write_request_count} + * <li>{@code gcs_http_read_request_count_ranged} / {@code gcs_http_read_request_count_unbounded} + * <li>{@code gcs_http_read_status_2xx} / {@code gcs_http_read_status_4xx} / {@code + * gcs_http_read_status_5xx} (and their write counterparts) + * </ul> + * + * <p>Results are currently only reported to standard output (see {@link #printMetrics}); nothing is + * persisted to BigQuery or InfluxDB. + */ +@SuppressWarnings({ Review Comment: Removed. -- 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]
