Abacn commented on code in PR #40142:
URL: https://github.com/apache/beam/pull/40142#discussion_r4038096491
##########
sdks/java/io/parquet/src/main/java/org/apache/beam/sdk/io/parquet/ParquetIO.java:
##########
@@ -736,6 +746,34 @@ private ParquetFileReader
getParquetFileReader(ReadableFile file) throws Excepti
return ParquetFileReader.open(new
BeamParquetInputFile(file.openSeekable()), options);
}
+ /**
+ * Returns the row group metadata of {@code file}, reading the Parquet
footer only if it is
+ * not already cached.
+ *
+ * <p>{@link #getInitialRestriction}, {@link #split}, {@link
#newTracker} and {@link #getSize}
+ * all need nothing from the file but its row group metadata, and the
runner may invoke them
+ * repeatedly for the same element: {@link #getSize} and {@link
#newTracker} are called on
+ * every dynamic split attempt. Opening a {@link ParquetFileReader} each
time re-reads the
+ * footer, which costs a round trip per call on object stores such as
GCS or S3. Caching the
+ * last file's row groups reduces that to one footer read per file.
+ *
+ * <p>This is synchronized because a runner may invoke the size and
tracker callbacks from a
+ * different thread than the one processing the bundle, concurrently
with {@link
+ * #processElement}.
+ */
+ private synchronized List<BlockMetaData> getRowGroups(ReadableFile file)
throws Exception {
Review Comment:
Just curious, is this an optimization and do we know how much performance
gain from it?
##########
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:
Please avoid SuppressWarnings in new codes
##########
sdks/java/extensions/google-cloud-platform-core/src/main/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilV1.java:
##########
@@ -154,7 +159,8 @@ public GcsUtilV1 create(PipelineOptions options) {
gcsOptions.getEnableBucketWriteMetricCounter()
? gcsOptions.getGcsWriteCounterPrefix()
: null),
- gcsOptions.getGoogleCloudStorageReadOptions());
+ gcsOptions.getGoogleCloudStorageReadOptions(),
Review Comment:
Instead of one more constructor parameter, this should fit within
GcsCountersOptions framework as it's for metrics on-off switch.
##########
it/google-cloud-platform/build.gradle:
##########
@@ -79,10 +85,31 @@ dependencies {
}
tasks.register(
- "GCSPerformanceTest", IoPerformanceTestUtilities.IoPerformanceTest,
project, 'google-cloud-platform', 'FileBasedIOLT',
+ "GCSPerformanceTest", IoPerformanceTestUtilities.IoPerformanceTest,
project, 'google-cloud-platform', 'TextIOLT',
['configuration':'large','project':'apache-beam-testing',
'artifactBucket':'io-performance-temp']
+ System.properties
)
+// Note: the task always passes 'configuration' down as a system property, so
the "local" default in
+// ParquetIOLT itself is never reached through gradle. The default here is
therefore the cheap local
+// preset: an accidental `./gradlew
:it:google-cloud-platform:ParquetPerformanceTest` must not launch
+// a full scale Dataflow job. Pass an explicit preset (or a json override) for
real runs, e.g.
+// -Dconfiguration=large or
-Dconfiguration='{"preset":"f100_s16","totalBytes":"10GB"}'.
+tasks.register(
+ "ParquetPerformanceTest",
IoPerformanceTestUtilities.IoPerformanceTest, project, 'google-cloud-platform',
'ParquetIOLT',
+ ['configuration':'local','project':'apache-beam-testing',
'artifactBucket':'io-performance-temp']
Review Comment:
There are workflows running large tests:
https://github.com/apache/beam/blob/master/.github/workflows/beam_StressTests_Java_BigQueryIO.yml
https://github.com/apache/beam/blob/master/.github/workflows/beam_StressTests_Java_KafkaIO.yml
and others.
Large test should have their own workflows as it runs for long
##########
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:
in terms of memory pressure, for standard run configs (compressibility are
always 0), there is a double allocation here might superficially degrading
performance. Consider write random bytes directly to payload
##########
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({
+ "nullness" // TODO(https://github.com/apache/beam/issues/27438)
+})
+public class GcsIOLoadTestBase extends IOLoadTestBase {
+
+ private static final Logger LOG =
LoggerFactory.getLogger(GcsIOLoadTestBase.class);
+
+ /** Prefix shared by all GCS client performance metrics. */
+ public static final String GCS_METRIC_PREFIX = "gcs_";
+
+ /**
+ * Pipeline option that has to be enabled for the GCS client to report the
{@value
+ * #GCS_METRIC_PREFIX} metrics.
+ */
+ public static final String GCS_PERFORMANCE_METRICS_OPTION =
"gcsPerformanceMetrics";
+
+ /**
+ * Returns all metrics of the job, including the GCS client performance
metrics.
+ *
+ * <p>The GCS metrics are aggregated over all the steps of the pipeline so
that they can be
+ * reported as flat scalar values. The per step breakdown is printed by
{@link #printGcsMetrics}.
+ */
+ @Override
+ protected Map<String, Double> getMetrics(
+ PipelineLauncher.LaunchInfo launchInfo, MetricsConfiguration config)
+ throws IOException, InterruptedException, ParseException {
+ Map<String, Double> metrics = super.getMetrics(launchInfo, config);
+ metrics.putAll(getGcsMetrics(launchInfo.jobId()));
+ return metrics;
+ }
+
+ /**
+ * Collects all the metrics of the job and prints them to standard output.
+ *
+ * <p>This is intentionally used instead of {@code exportMetricsToBigQuery}
while these tests are
+ * still being developed: results are only reported to the console, nothing
is persisted.
+ */
+ protected void printMetrics(
+ PipelineLauncher.LaunchInfo launchInfo, MetricsConfiguration
metricsConfig) {
+ Map<String, Double> metrics;
+ try {
+ metrics = getMetrics(launchInfo, metricsConfig);
+ } catch (Exception e) {
+ LOG.warn("Unable to get metrics due to error", e);
+ return;
+ }
+
+ StringBuilder report = new StringBuilder();
+
report.append("\n==========================================================\n");
+ report.append(String.format(Locale.US, " PIPELINE METRICS (job %s)%n",
launchInfo.jobId()));
+
report.append("==========================================================\n");
+ if (metrics.isEmpty()) {
+ report.append(" No metrics found.\n");
+ } else {
+ for (Map.Entry<String, Double> entry : new
TreeMap<>(metrics).entrySet()) {
+ report.append(
+ String.format(Locale.US, " %-46s %,.3f%n", entry.getKey() + ":",
entry.getValue()));
+ }
+ }
+
report.append("==========================================================");
+ print(report.toString());
+
+ // Also print the GCS specific report, which includes the per step
breakdown.
+ printGcsMetrics(launchInfo.jobId());
+ }
+
+ /**
+ * Collects the GCS client performance metrics of the given job, aggregated
over all steps.
+ *
+ * <p>Counters are summed up across steps. Distributions are reported as
four separate scalar
+ * metrics, suffixed with {@code _COUNT}, {@code _SUM}, {@code _MIN} and
{@code _MAX}, which
+ * matches how the Dataflow launcher reports distributions.
+ *
+ * @param jobId the id of the job to query
+ * @return a map of GCS metric name to value, empty if no GCS metric was
reported
+ */
+ protected Map<String, Double> getGcsMetrics(String jobId) {
Review Comment:
getGcsMetrics and printGcsMetrics nearly identical but long codes. Consider
consolidate them.
##########
it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/storage/TextIOLT.java:
##########
@@ -54,31 +54,31 @@
import org.junit.Test;
/**
- * FileBasedIO performance tests.
+ * TextIO performance tests.
*
* <p>Example trigger command for all tests:
*
* <pre>
- * mvn test -pl it/google-cloud-platform -am -Dtest="FileBasedIOLT"
-Dproject=[gcpProject] \
+ * mvn test -pl it/google-cloud-platform -am -Dtest="TextIOLT"
-Dproject=[gcpProject] \
* -DartifactBucket=[temp bucket] -DfailIfNoTests=false
* </pre>
*
* <p>Example trigger command for specific test running on direct runner:
*
* <pre>
- * mvn test -pl it/google-cloud-platform -am
-Dtest="FileBasedIOLT#testTextIOWriteThenRead" \
+ * mvn test -pl it/google-cloud-platform -am
-Dtest="TextIOLT#testTextIOWriteThenRead" \
* -Dconfiguration=medium -Dproject=[gcpProject] -DartifactBucket=[temp
bucket] -DfailIfNoTests=false
* </pre>
*
* <p>Example trigger command for specific test and custom data configuration:
*
* <pre>mvn test -pl it/google-cloud-platform -am \
*
-Dconfiguration="{\"numRecords\":10000000,\"valueSizeBytes\":750,\"pipelineTimeout\":20,\"runner\":\"DataflowRunner\"}"
\
- * -Dtest="FileBasedIOLT#testTextIOWriteThenRead" -Dconfiguration=local
-Dproject=[gcpProject] \
+ * -Dtest="TextIOLT#testTextIOWriteThenRead" -Dconfiguration=local
-Dproject=[gcpProject] \
Review Comment:
pre-existing: duplicate -Dconfiguration=... here
##########
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({
+ "nullness" // TODO(https://github.com/apache/beam/issues/27438)
+})
+public class GcsIOLoadTestBase extends IOLoadTestBase {
+
+ private static final Logger LOG =
LoggerFactory.getLogger(GcsIOLoadTestBase.class);
+
+ /** Prefix shared by all GCS client performance metrics. */
+ public static final String GCS_METRIC_PREFIX = "gcs_";
+
+ /**
+ * Pipeline option that has to be enabled for the GCS client to report the
{@value
+ * #GCS_METRIC_PREFIX} metrics.
+ */
+ public static final String GCS_PERFORMANCE_METRICS_OPTION =
"gcsPerformanceMetrics";
+
+ /**
+ * Returns all metrics of the job, including the GCS client performance
metrics.
+ *
+ * <p>The GCS metrics are aggregated over all the steps of the pipeline so
that they can be
+ * reported as flat scalar values. The per step breakdown is printed by
{@link #printGcsMetrics}.
+ */
+ @Override
+ protected Map<String, Double> getMetrics(
+ PipelineLauncher.LaunchInfo launchInfo, MetricsConfiguration config)
+ throws IOException, InterruptedException, ParseException {
+ Map<String, Double> metrics = super.getMetrics(launchInfo, config);
+ metrics.putAll(getGcsMetrics(launchInfo.jobId()));
+ return metrics;
+ }
+
+ /**
+ * Collects all the metrics of the job and prints them to standard output.
+ *
+ * <p>This is intentionally used instead of {@code exportMetricsToBigQuery}
while these tests are
+ * still being developed: results are only reported to the console, nothing
is persisted.
+ */
+ protected void printMetrics(
+ PipelineLauncher.LaunchInfo launchInfo, MetricsConfiguration
metricsConfig) {
+ Map<String, Double> metrics;
+ try {
+ metrics = getMetrics(launchInfo, metricsConfig);
+ } catch (Exception e) {
+ LOG.warn("Unable to get metrics due to error", e);
+ return;
+ }
+
+ StringBuilder report = new StringBuilder();
+
report.append("\n==========================================================\n");
+ report.append(String.format(Locale.US, " PIPELINE METRICS (job %s)%n",
launchInfo.jobId()));
+
report.append("==========================================================\n");
+ if (metrics.isEmpty()) {
+ report.append(" No metrics found.\n");
+ } else {
+ for (Map.Entry<String, Double> entry : new
TreeMap<>(metrics).entrySet()) {
+ report.append(
+ String.format(Locale.US, " %-46s %,.3f%n", entry.getKey() + ":",
entry.getValue()));
+ }
+ }
+
report.append("==========================================================");
+ print(report.toString());
+
+ // Also print the GCS specific report, which includes the per step
breakdown.
+ printGcsMetrics(launchInfo.jobId());
+ }
+
+ /**
+ * Collects the GCS client performance metrics of the given job, aggregated
over all steps.
+ *
+ * <p>Counters are summed up across steps. Distributions are reported as
four separate scalar
+ * metrics, suffixed with {@code _COUNT}, {@code _SUM}, {@code _MIN} and
{@code _MAX}, which
+ * matches how the Dataflow launcher reports distributions.
+ *
+ * @param jobId the id of the job to query
+ * @return a map of GCS metric name to value, empty if no GCS metric was
reported
+ */
+ protected Map<String, Double> getGcsMetrics(String jobId) {
+ Map<String, Double> gcsMetrics = new TreeMap<>();
+ Map<String, Map<String, Long>> counters = getGcsCountersByStep(jobId);
+ for (Map.Entry<String, Map<String, Long>> entry : counters.entrySet()) {
+ long total =
entry.getValue().values().stream().mapToLong(Long::longValue).sum();
+ gcsMetrics.put(entry.getKey(), (double) total);
+ }
+
+ Map<String, Map<String, DistributionResult>> distributions =
getGcsDistributionsByStep(jobId);
+ for (Map.Entry<String, Map<String, DistributionResult>> entry :
distributions.entrySet()) {
+ String name = entry.getKey();
+ long count = 0;
+ long sum = 0;
+ Long min = null;
+ Long max = null;
+ for (DistributionResult distribution : entry.getValue().values()) {
+ count += distribution.getCount();
+ sum += distribution.getSum();
+ min = (min == null) ? distribution.getMin() : Math.min(min,
distribution.getMin());
+ max = (max == null) ? distribution.getMax() : Math.max(max,
distribution.getMax());
+ }
+ gcsMetrics.put(name + "_COUNT", (double) count);
+ gcsMetrics.put(name + "_SUM", (double) sum);
+ if (min != null) {
+ gcsMetrics.put(name + "_MIN", (double) min);
+ }
+ if (max != null) {
+ gcsMetrics.put(name + "_MAX", (double) max);
+ }
+ }
+
+ if (gcsMetrics.isEmpty()) {
+ LOG.warn(
+ "No {}* metrics found for job {}. Make sure the pipeline was
launched with --{}=true.",
+ GCS_METRIC_PREFIX,
+ jobId,
+ GCS_PERFORMANCE_METRICS_OPTION);
+ }
+ return gcsMetrics;
+ }
+
+ /** Prints a human readable report of the GCS metrics, including the per
step breakdown. */
+ protected void printGcsMetrics(String jobId) {
+ Map<String, Map<String, Long>> counters = getGcsCountersByStep(jobId);
+ Map<String, Map<String, DistributionResult>> distributions =
getGcsDistributionsByStep(jobId);
+
+ StringBuilder report = new StringBuilder();
+
report.append("\n==========================================================\n");
+ report.append(" GCS METRICS
\n");
+
report.append("==========================================================\n");
+ if (counters.isEmpty() && distributions.isEmpty()) {
+ report.append(" No ").append(GCS_METRIC_PREFIX).append("* metrics
found.\n");
+ } else {
+ for (Map.Entry<String, Map<String, Long>> entry : counters.entrySet()) {
+ String metricName = entry.getKey();
+ Map<String, Long> stepMap = entry.getValue();
+ long total =
stepMap.values().stream().mapToLong(Long::longValue).sum();
+ report.append(
+ String.format(Locale.US, " %-36s %s\n", metricName + ":",
format(metricName, total)));
+ if (stepMap.size() > 1 || (!stepMap.containsKey("global") &&
!stepMap.isEmpty())) {
+ for (Map.Entry<String, Long> stepEntry : stepMap.entrySet()) {
+ report.append(
+ String.format(
+ Locale.US,
+ " [%s]: %s\n",
+ stepEntry.getKey(),
+ format(metricName, stepEntry.getValue())));
+ }
+ }
+ }
+ for (Map.Entry<String, Map<String, DistributionResult>> entry :
distributions.entrySet()) {
+ String metricName = entry.getKey();
+ for (Map.Entry<String, DistributionResult> stepEntry :
entry.getValue().entrySet()) {
+ DistributionResult d = stepEntry.getValue();
+ report.append(
+ String.format(
+ Locale.US,
+ " %-36s count=%,d, sum=%,d, min=%,d, max=%,d, mean=%.2f
[%s]\n",
+ metricName + ":",
+ d.getCount(),
+ d.getSum(),
+ d.getMin(),
+ d.getMax(),
+ d.getMean(),
+ stepEntry.getKey()));
+ }
+ }
+ }
+
report.append("==========================================================");
+ print(report.toString());
+ }
+
+ /** Writes the report to standard output. */
+ private static void print(String report) {
+ System.out.println(report);
+ }
+
+ private static Map<String, Map<String, Long>> getGcsCountersByStep(String
jobId) {
Review Comment:
Similarly, getGcsCountersByStep and getGcsDistributionsByStep also have
nearly identical filtering and grouping loops.
Possible to extract a generic `Extract a single generic helper private
static <T> Map<String, Map<String, T>>
extractGcsMetricsByStep(Iterable<MetricResult<T>> results)` ?
##########
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) {
Review Comment:
It's testing code and we don't need complicated parsers (regex) and even
handling alias `switch (unit) { ... }`. Can we just use bytes or simple suffix?
##########
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) {
Review Comment:
This duplicates the same helper in GcsIOLoadTestBase
##########
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:
both minFieldSizeBytes/maxFieldSizeBytes and CreateAvroRecordFn are newly
created test codes. Please align their types so no superficial (and potential
overflowing) cast needed
--
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]
