You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@beam.apache.org by GitBox <gi...@apache.org> on 2022/12/09 22:53:07 UTC
[GitHub] [beam] damondouglas opened a new pull request, #24630: [24469] Implement CsvIO.Write and supporting classes
damondouglas opened a new pull request, #24630:
URL: https://github.com/apache/beam/pull/24630
This PR closes #24469 and addresses #20312 by implementing CsvIO.Write and supporting classes. Please see #24469 for more details.
------------------------
Thank you for your contribution! Follow this checklist to help us incorporate your contribution quickly and easily:
- [x] Mention the appropriate issue in your description (for example: `addresses #123`), if applicable. This will automatically add a link to the pull request in the issue. If you would like the issue to automatically close on merging the pull request, comment `fixes #<ISSUE NUMBER>` instead.
- ~[ ] Update `CHANGES.md` with noteworthy changes.~
- ~[ ] If this contribution is large, please file an Apache [Individual Contributor License Agreement](https://www.apache.org/licenses/icla.pdf).~
See the [Contributor Guide](https://beam.apache.org/contribute) for more tips on [how to make review process smoother](https://beam.apache.org/contribute/get-started-contributing/#make-the-reviewers-job-easier).
To check the build health, please visit [https://github.com/apache/beam/blob/master/.test-infra/BUILD_STATUS.md](https://github.com/apache/beam/blob/master/.test-infra/BUILD_STATUS.md)
GitHub Actions Tests Status (on master branch)
------------------------------------------------------------------------------------------------
[](https://github.com/apache/beam/actions?query=workflow%3A%22Build+python+source+distribution+and+wheels%22+branch%3Amaster+event%3Aschedule)
[](https://github.com/apache/beam/actions?query=workflow%3A%22Python+Tests%22+branch%3Amaster+event%3Aschedule)
[](https://github.com/apache/beam/actions?query=workflow%3A%22Java+Tests%22+branch%3Amaster+event%3Aschedule)
[](https://github.com/apache/beam/actions?query=workflow%3A%22Go+tests%22+branch%3Amaster+event%3Aschedule)
See [CI.md](https://github.com/apache/beam/blob/master/CI.md) for more information about GitHub Actions CI.
--
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: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [beam] damondouglas commented on a diff in pull request #24630: [24469] Implement CsvIO.Write and supporting classes
Posted by "damondouglas (via GitHub)" <gi...@apache.org>.
damondouglas commented on code in PR #24630:
URL: https://github.com/apache/beam/pull/24630#discussion_r1085693708
##########
sdks/java/io/csv/src/main/java/org/apache/beam/sdk/io/csv/CsvIO.java:
##########
@@ -0,0 +1,446 @@
+/*
+ * 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.sdk.io.csv;
+
+import static java.util.Objects.requireNonNull;
+import static org.apache.beam.sdk.values.TypeDescriptors.rows;
+import static org.apache.beam.sdk.values.TypeDescriptors.strings;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+
+import com.google.auto.value.AutoValue;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Set;
+import org.apache.beam.sdk.coders.RowCoder;
+import org.apache.beam.sdk.io.Compression;
+import org.apache.beam.sdk.io.FileBasedSink;
+import org.apache.beam.sdk.io.FileIO;
+import org.apache.beam.sdk.io.ShardNameTemplate;
+import org.apache.beam.sdk.io.TextIO;
+import org.apache.beam.sdk.io.WriteFilesResult;
+import org.apache.beam.sdk.io.fs.ResourceId;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.transforms.MapElements;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.sdk.values.TypeDescriptor;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableSet;
+import org.apache.commons.csv.CSVFormat;
+
+/**
+ * {@link PTransform}s for reading and writing CSV files.
+ *
+ * <h2>Reading CSV files</h2>
+ *
+ * <p>Reading from CSV files is not yet implemented. Please see <a
+ * href="https://github.com/apache/beam/issues/24552">https://github.com/apache/beam/issues/24552</a>.
+ *
+ * <h2>Writing CSV files</h2>
+ *
+ * <p>To write a {@link PCollection} to one or more CSV files, use {@link CsvIO.Write}, using {@link
+ * CsvIO#writeRows(String, CSVFormat)} or {@link CsvIO#write(String, CSVFormat)}. {@link
+ * CsvIO.Write} supports writing {@link Row} or custom Java types using an inferred {@link Schema}.
+ * Examples below show both scenarios. See the Beam Programming Guide on <a
+ * href="https://beam.apache.org/documentation/programming-guide/#inferring-schemas">inferring
+ * schemas</a> for more information on how to enable Beam to infer a {@link Schema} from a custom
+ * Java type.
+ *
+ * <p>{@link CsvIO.Write} only supports writing {@link Schema} aware types that do not contain any
+ * nested {@link FieldType}s such a {@link org.apache.beam.sdk.schemas.Schema.TypeName#ROW} or
+ * repeated {@link org.apache.beam.sdk.schemas.Schema.TypeName#ARRAY} types. See {@link
+ * CsvIO.Write#VALID_FIELD_TYPE_SET} for valid {@link FieldType}s.
+ *
+ * <h3>Example usage:</h3>
+ *
+ * <p>Suppose we have the following <code>Transaction</code> class annotated with
+ * {@code @DefaultSchema(JavaBeanSchema.class)} so that Beam can infer its {@link Schema}:
+ *
+ * <pre>{@code @DefaultSchema(JavaBeanSchema.class)
+ * public class Transaction {
+ * public Transaction() { … }
+ * public Long getTransactionId();
+ * public void setTransactionId(Long transactionId) { … }
+ * public String getBank() { … }
+ * public void setBank(String bank) { … }
+ * public double getPurchaseAmount() { … }
+ * public void setPurchaseAmount(double purchaseAmount) { … }
+ * }
+ * }</pre>
+ *
+ * <p>From a {@code PCollection<Transaction>}, {@link CsvIO.Write} can write one or many CSV files
+ * automatically creating the header based on its inferred {@link Schema}.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions = ...
+ * transactions.apply(CsvIO.<Transaction>write("path/to/folder/prefix", CSVFormat.DEFAULT));
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header is repeated for every
+ * shard file, whereas by default, {@link CsvIO.Write} will write all fields in <b>sorted order</b>
+ * of the field names.
+ *
+ * <pre>{@code
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ *
+ * <p>To control the order and subset of fields that {@link CsvIO.Write} writes, use {@link
+ * CSVFormat#withHeader(String...)}. Note, however, the following constraints:
+ *
+ * <ol>
+ * <li>Each header column must match a field name in the {@link Schema}; matching is case
+ * sensitive.
+ * <li>Matching header columns must match {@link Schema} fields that are valid {@link FieldType}s;
+ * see {@link #VALID_FIELD_TYPE_SET}.
+ * <li>{@link CSVFormat} only allows repeated header columns when {@link
+ * CSVFormat#withAllowDuplicateHeaderNames()}
+ * </ol>
+ *
+ * <p>The following example shows the use of {@link CSVFormat#withHeader(String...)} to control the
+ * order and subset of <code>Transaction</code> fields.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions ...
+ * transactions.apply(
+ * CsvIO
+ * .<Transaction>write("path/to/folder/prefix", CSVFormat.DEFAULT.withHeader("transactionId", "purchaseAmount"))
+ * );
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header is repeated for every
+ * shard file, but only including the subset of fields in their listed order.
+ *
+ * <pre>{@code
+ * transactionId,purchaseAmount
+ * 12345,10.23
+ * 54321,54.65
+ * 98765,11.76
+ * }</pre>
+ *
+ * <p>In addition to header customization, {@link CsvIO.Write} supports {@link
+ * CSVFormat#withHeaderComments(Object...)} as shown below. Note that {@link
+ * CSVFormat#withCommentMarker(char)} is required when specifying header comments.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions = ...
+ * transactions
+ * .apply(
+ * CsvIO.<Transaction>write("path/to/folder/prefix",
+ * CSVFormat.DEFAULT.withCommentMarker('#').withHeaderComments("Bank Report", "1970-01-01", "Operator: John Doe")
+ * );
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header and header comments are
+ * repeated for every shard file.
+ *
+ * <pre>{@code
+ * # Bank Report
+ * # 1970-01-01
+ * # Operator: John Doe
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ *
+ * <p>A {@link PCollection} of {@link Row}s works just like custom Java types illustrated above,
+ * except we use {@link CsvIO#writeRows(String, CSVFormat)} as shown below for the same {@code
+ * Transaction} class. We derive {@code Transaction}'s {@link Schema} using a {@link
+ * org.apache.beam.sdk.schemas.annotations.DefaultSchema.DefaultSchemaProvider}. Note that
+ * hard-coding the {@link Row}s below is for illustration purposes. Developers are instead
+ * encouraged to take advantage of {@link
+ * org.apache.beam.sdk.schemas.annotations.DefaultSchema.DefaultSchemaProvider#toRowFunction(TypeDescriptor)}.
+ *
+ * <pre>{@code
+ * DefaultSchemaProvider defaultSchemaProvider = new DefaultSchemaProvider();
+ * Schema schema = defaultSchemaProvider.schemaFor(TypeDescriptor.of(Transaction.class));
+ * PCollection<Row> transactions = pipeline.apply(Create.of(
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "A")
+ * .withFieldValue("purchaseAmount", 10.23)
+ * .withFieldValue("transactionId", "12345")
+ * .build(),
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "B")
+ * .withFieldValue("purchaseAmount", 54.65)
+ * .withFieldValue("transactionId", "54321")
+ * .build(),
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "C")
+ * .withFieldValue("purchaseAmount", 11.76)
+ * .withFieldValue("transactionId", "98765")
+ * .build()
+ * );
+ *
+ * transactions.apply(
+ * CsvIO
+ * .writeRowsTo("gs://bucket/path/to/folder/prefix", CSVFormat.DEFAULT)
+ * );
+ * }</pre>
+ *
+ * <p>Writing the transactions {@link PCollection} of {@link Row}s would yield the following CSV
+ * file content.
+ *
+ * <pre>{@code
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ *
+ * {@link CsvIO.Write} does not support the following {@link CSVFormat} properties and will throw an
+ * {@link IllegalArgumentException}.
+ *
+ * <ul>
+ * <li>{@link CSVFormat#withSkipHeaderRecord()}
+ * <li>{@link CSVFormat#withAllowMissingColumnNames()} ()}
+ * <li>{@link CSVFormat#withAutoFlush(boolean)}
+ * <li>{@link CSVFormat#withIgnoreHeaderCase()}
+ * <li>{@link CSVFormat#withIgnoreSurroundingSpaces()}
+ * </ul>
+ */
+public class CsvIO {
+ /**
+ * The valid {@link Schema.FieldType} from which {@link CsvIO} converts CSV records to the fields.
+ *
+ * <ul>
+ * <li>{@link FieldType#BYTE}
+ * <li>{@link FieldType#BOOLEAN}
+ * <li>{@link FieldType#DATETIME}
+ * <li>{@link FieldType#DECIMAL}
+ * <li>{@link FieldType#DOUBLE}
+ * <li>{@link FieldType#INT16}
+ * <li>{@link FieldType#INT32}
+ * <li>{@link FieldType#INT64}
+ * <li>{@link FieldType#FLOAT}
+ * <li>{@link FieldType#STRING}
+ * </ul>
+ */
+ public static final Set<Schema.FieldType> VALID_FIELD_TYPE_SET =
+ ImmutableSet.of(
+ FieldType.BYTE,
+ FieldType.BOOLEAN,
+ FieldType.DATETIME,
+ FieldType.DECIMAL,
+ FieldType.DOUBLE,
+ FieldType.INT16,
+ FieldType.INT32,
+ FieldType.INT64,
+ FieldType.FLOAT,
+ FieldType.STRING);
+
+ static final String DEFAULT_FILENAME_SUFFIX = ".csv";
+
+ /** Instantiates a {@link Write} for writing user types in {@link CSVFormat} format. */
+ public static <T> Write<T> write(String to, CSVFormat csvFormat) {
+ return new AutoValue_CsvIO_Write.Builder<T>()
+ .setTextIOWrite(createDefaultTextIOWrite(to))
+ .setCSVFormat(csvFormat)
+ .build();
+ }
+
+ /** Instantiates a {@link Write} for writing {@link Row}s in {@link CSVFormat} format. */
+ public static Write<Row> writeRows(String to, CSVFormat csvFormat) {
+ return new AutoValue_CsvIO_Write.Builder<Row>()
+ .setTextIOWrite(createDefaultTextIOWrite(to))
+ .setCSVFormat(csvFormat)
+ .build();
+ }
+
+ /** {@link PTransform} for writing CSV files. */
+ @AutoValue
+ public abstract static class Write<T>
+ extends PTransform<PCollection<T>, WriteFilesResult<String>> {
+
+ /** Specifies the {@link Compression} of all generated shard files. */
+ public Write<T> withCompression(Compression compression) {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withCompression(compression)).build();
+ }
+
+ public Write<T> withNoSpilling() {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withNoSpilling()).build();
+ }
+
+ /**
+ * Specifies to use a given fixed number of shards per window. See {@link
+ * TextIO.Write#withNumShards(int)}.
+ */
+ public Write<T> withNumShards(Integer numShards) {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withNumShards(numShards)).build();
+ }
+
+ /**
+ * Forces a single file as output and empty shard name template. See {@link
+ * TextIO.Write#withoutSharding()}.
+ */
+ public Write<T> withoutSharding() {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withoutSharding()).build();
+ }
+
+ /**
+ * Uses the given {@link ShardNameTemplate} for naming output files. See {@link
+ * TextIO.Write#withShardNameTemplate(String)}.
+ */
+ public Write<T> withShardTemplate(String shardTemplate) {
+ return toBuilder()
+ .setTextIOWrite(getTextIOWrite().withShardNameTemplate(shardTemplate))
+ .build();
+ }
+
+ /**
+ * Configures the filename suffix for written files. See {@link
+ * TextIO.Write#withSuffix(String)}.
+ */
+ public Write<T> withSuffix(String suffix) {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withSuffix(suffix)).build();
+ }
+
+ /**
+ * Set the base directory used to generate temporary files. See {@link
+ * TextIO.Write#withTempDirectory(ResourceId)}.
+ */
+ public Write<T> withTempDirectory(ResourceId tempDirectory) {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withTempDirectory(tempDirectory)).build();
+ }
+
+ /**
+ * Preserves windowing of input elements and writes them to files based on the element's window.
+ * See {@link TextIO.Write#withWindowedWrites()}.
+ */
+ public Write<T> withWindowedWrites() {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withWindowedWrites()).build();
+ }
+
+ /**
+ * Returns a transform for writing to text files like this one but that has the given {@link
+ * FileBasedSink.WritableByteChannelFactory} to be used by the {@link FileBasedSink} during
+ * output. See {@link
+ * TextIO.Write#withWritableByteChannelFactory(FileBasedSink.WritableByteChannelFactory)}.
+ */
+ public Write<T> withWritableByteChannelFactory(
+ FileBasedSink.WritableByteChannelFactory writableByteChannelFactory) {
+ return toBuilder()
+ .setTextIOWrite(
+ getTextIOWrite().withWritableByteChannelFactory(writableByteChannelFactory))
+ .build();
+ }
+
+ /** The underlying {@link FileIO.Write} that writes converted input to CSV formatted output. */
+ abstract TextIO.Write getTextIOWrite();
+
+ /** The {@link CSVFormat} to inform headers, header comments, and format CSV row content. */
+ abstract CSVFormat getCSVFormat();
+
+ abstract Builder<T> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<T> {
+
+ /**
+ * The underlying {@link FileIO.Write} that writes converted input to CSV formatted output.
+ */
+ abstract Builder<T> setTextIOWrite(TextIO.Write value);
+
+ /** The {@link CSVFormat} to convert input. Defaults to {@link CSVFormat#DEFAULT}. */
+ abstract Builder<T> setCSVFormat(CSVFormat value);
+
+ abstract CSVFormat getCSVFormat();
+
+ abstract Write<T> autoBuild();
+
+ final Write<T> build() {
+ checkArgument(
+ !getCSVFormat().getSkipHeaderRecord(),
Review Comment:
CsvRowConversions performs many of the checks as they most impact RowToCsv conversion. I thought it best to do the checks closest to where they would impact the outcome.
Additionally, does the 375 to 377 comment clarify why we are checking the skipHeaderRecord property and throwing an exception?
--
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: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [beam] damondouglas commented on a diff in pull request #24630: [24469] Implement CsvIO.Write and supporting classes
Posted by "damondouglas (via GitHub)" <gi...@apache.org>.
damondouglas commented on code in PR #24630:
URL: https://github.com/apache/beam/pull/24630#discussion_r1085716045
##########
sdks/java/io/csv/src/main/java/org/apache/beam/sdk/io/csv/CsvIO.java:
##########
@@ -0,0 +1,446 @@
+/*
+ * 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.sdk.io.csv;
+
+import static java.util.Objects.requireNonNull;
+import static org.apache.beam.sdk.values.TypeDescriptors.rows;
+import static org.apache.beam.sdk.values.TypeDescriptors.strings;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+
+import com.google.auto.value.AutoValue;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Set;
+import org.apache.beam.sdk.coders.RowCoder;
+import org.apache.beam.sdk.io.Compression;
+import org.apache.beam.sdk.io.FileBasedSink;
+import org.apache.beam.sdk.io.FileIO;
+import org.apache.beam.sdk.io.ShardNameTemplate;
+import org.apache.beam.sdk.io.TextIO;
+import org.apache.beam.sdk.io.WriteFilesResult;
+import org.apache.beam.sdk.io.fs.ResourceId;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.transforms.MapElements;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.sdk.values.TypeDescriptor;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableSet;
+import org.apache.commons.csv.CSVFormat;
+
+/**
+ * {@link PTransform}s for reading and writing CSV files.
+ *
+ * <h2>Reading CSV files</h2>
+ *
+ * <p>Reading from CSV files is not yet implemented. Please see <a
+ * href="https://github.com/apache/beam/issues/24552">https://github.com/apache/beam/issues/24552</a>.
+ *
+ * <h2>Writing CSV files</h2>
+ *
+ * <p>To write a {@link PCollection} to one or more CSV files, use {@link CsvIO.Write}, using {@link
+ * CsvIO#writeRows(String, CSVFormat)} or {@link CsvIO#write(String, CSVFormat)}. {@link
+ * CsvIO.Write} supports writing {@link Row} or custom Java types using an inferred {@link Schema}.
+ * Examples below show both scenarios. See the Beam Programming Guide on <a
+ * href="https://beam.apache.org/documentation/programming-guide/#inferring-schemas">inferring
+ * schemas</a> for more information on how to enable Beam to infer a {@link Schema} from a custom
+ * Java type.
+ *
+ * <p>{@link CsvIO.Write} only supports writing {@link Schema} aware types that do not contain any
+ * nested {@link FieldType}s such a {@link org.apache.beam.sdk.schemas.Schema.TypeName#ROW} or
+ * repeated {@link org.apache.beam.sdk.schemas.Schema.TypeName#ARRAY} types. See {@link
+ * CsvIO.Write#VALID_FIELD_TYPE_SET} for valid {@link FieldType}s.
+ *
+ * <h3>Example usage:</h3>
+ *
+ * <p>Suppose we have the following <code>Transaction</code> class annotated with
+ * {@code @DefaultSchema(JavaBeanSchema.class)} so that Beam can infer its {@link Schema}:
+ *
+ * <pre>{@code @DefaultSchema(JavaBeanSchema.class)
+ * public class Transaction {
+ * public Transaction() { … }
+ * public Long getTransactionId();
+ * public void setTransactionId(Long transactionId) { … }
+ * public String getBank() { … }
+ * public void setBank(String bank) { … }
+ * public double getPurchaseAmount() { … }
+ * public void setPurchaseAmount(double purchaseAmount) { … }
+ * }
+ * }</pre>
+ *
+ * <p>From a {@code PCollection<Transaction>}, {@link CsvIO.Write} can write one or many CSV files
+ * automatically creating the header based on its inferred {@link Schema}.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions = ...
+ * transactions.apply(CsvIO.<Transaction>write("path/to/folder/prefix", CSVFormat.DEFAULT));
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header is repeated for every
+ * shard file, whereas by default, {@link CsvIO.Write} will write all fields in <b>sorted order</b>
+ * of the field names.
+ *
+ * <pre>{@code
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ *
+ * <p>To control the order and subset of fields that {@link CsvIO.Write} writes, use {@link
+ * CSVFormat#withHeader(String...)}. Note, however, the following constraints:
+ *
+ * <ol>
+ * <li>Each header column must match a field name in the {@link Schema}; matching is case
+ * sensitive.
+ * <li>Matching header columns must match {@link Schema} fields that are valid {@link FieldType}s;
+ * see {@link #VALID_FIELD_TYPE_SET}.
+ * <li>{@link CSVFormat} only allows repeated header columns when {@link
+ * CSVFormat#withAllowDuplicateHeaderNames()}
+ * </ol>
+ *
+ * <p>The following example shows the use of {@link CSVFormat#withHeader(String...)} to control the
+ * order and subset of <code>Transaction</code> fields.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions ...
+ * transactions.apply(
+ * CsvIO
+ * .<Transaction>write("path/to/folder/prefix", CSVFormat.DEFAULT.withHeader("transactionId", "purchaseAmount"))
+ * );
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header is repeated for every
+ * shard file, but only including the subset of fields in their listed order.
+ *
+ * <pre>{@code
+ * transactionId,purchaseAmount
+ * 12345,10.23
+ * 54321,54.65
+ * 98765,11.76
+ * }</pre>
+ *
+ * <p>In addition to header customization, {@link CsvIO.Write} supports {@link
+ * CSVFormat#withHeaderComments(Object...)} as shown below. Note that {@link
+ * CSVFormat#withCommentMarker(char)} is required when specifying header comments.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions = ...
+ * transactions
+ * .apply(
+ * CsvIO.<Transaction>write("path/to/folder/prefix",
+ * CSVFormat.DEFAULT.withCommentMarker('#').withHeaderComments("Bank Report", "1970-01-01", "Operator: John Doe")
+ * );
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header and header comments are
+ * repeated for every shard file.
+ *
+ * <pre>{@code
+ * # Bank Report
+ * # 1970-01-01
+ * # Operator: John Doe
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ *
+ * <p>A {@link PCollection} of {@link Row}s works just like custom Java types illustrated above,
+ * except we use {@link CsvIO#writeRows(String, CSVFormat)} as shown below for the same {@code
+ * Transaction} class. We derive {@code Transaction}'s {@link Schema} using a {@link
+ * org.apache.beam.sdk.schemas.annotations.DefaultSchema.DefaultSchemaProvider}. Note that
+ * hard-coding the {@link Row}s below is for illustration purposes. Developers are instead
+ * encouraged to take advantage of {@link
+ * org.apache.beam.sdk.schemas.annotations.DefaultSchema.DefaultSchemaProvider#toRowFunction(TypeDescriptor)}.
+ *
+ * <pre>{@code
+ * DefaultSchemaProvider defaultSchemaProvider = new DefaultSchemaProvider();
+ * Schema schema = defaultSchemaProvider.schemaFor(TypeDescriptor.of(Transaction.class));
+ * PCollection<Row> transactions = pipeline.apply(Create.of(
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "A")
+ * .withFieldValue("purchaseAmount", 10.23)
+ * .withFieldValue("transactionId", "12345")
+ * .build(),
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "B")
+ * .withFieldValue("purchaseAmount", 54.65)
+ * .withFieldValue("transactionId", "54321")
+ * .build(),
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "C")
+ * .withFieldValue("purchaseAmount", 11.76)
+ * .withFieldValue("transactionId", "98765")
+ * .build()
+ * );
+ *
+ * transactions.apply(
+ * CsvIO
+ * .writeRowsTo("gs://bucket/path/to/folder/prefix", CSVFormat.DEFAULT)
+ * );
+ * }</pre>
+ *
+ * <p>Writing the transactions {@link PCollection} of {@link Row}s would yield the following CSV
+ * file content.
+ *
+ * <pre>{@code
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ *
+ * {@link CsvIO.Write} does not support the following {@link CSVFormat} properties and will throw an
+ * {@link IllegalArgumentException}.
+ *
+ * <ul>
+ * <li>{@link CSVFormat#withSkipHeaderRecord()}
+ * <li>{@link CSVFormat#withAllowMissingColumnNames()} ()}
+ * <li>{@link CSVFormat#withAutoFlush(boolean)}
+ * <li>{@link CSVFormat#withIgnoreHeaderCase()}
+ * <li>{@link CSVFormat#withIgnoreSurroundingSpaces()}
+ * </ul>
+ */
+public class CsvIO {
+ /**
+ * The valid {@link Schema.FieldType} from which {@link CsvIO} converts CSV records to the fields.
+ *
+ * <ul>
+ * <li>{@link FieldType#BYTE}
+ * <li>{@link FieldType#BOOLEAN}
+ * <li>{@link FieldType#DATETIME}
+ * <li>{@link FieldType#DECIMAL}
+ * <li>{@link FieldType#DOUBLE}
+ * <li>{@link FieldType#INT16}
+ * <li>{@link FieldType#INT32}
+ * <li>{@link FieldType#INT64}
+ * <li>{@link FieldType#FLOAT}
+ * <li>{@link FieldType#STRING}
+ * </ul>
+ */
+ public static final Set<Schema.FieldType> VALID_FIELD_TYPE_SET =
+ ImmutableSet.of(
+ FieldType.BYTE,
+ FieldType.BOOLEAN,
+ FieldType.DATETIME,
+ FieldType.DECIMAL,
+ FieldType.DOUBLE,
+ FieldType.INT16,
+ FieldType.INT32,
+ FieldType.INT64,
+ FieldType.FLOAT,
+ FieldType.STRING);
+
+ static final String DEFAULT_FILENAME_SUFFIX = ".csv";
+
+ /** Instantiates a {@link Write} for writing user types in {@link CSVFormat} format. */
+ public static <T> Write<T> write(String to, CSVFormat csvFormat) {
+ return new AutoValue_CsvIO_Write.Builder<T>()
+ .setTextIOWrite(createDefaultTextIOWrite(to))
+ .setCSVFormat(csvFormat)
+ .build();
+ }
+
+ /** Instantiates a {@link Write} for writing {@link Row}s in {@link CSVFormat} format. */
+ public static Write<Row> writeRows(String to, CSVFormat csvFormat) {
+ return new AutoValue_CsvIO_Write.Builder<Row>()
+ .setTextIOWrite(createDefaultTextIOWrite(to))
+ .setCSVFormat(csvFormat)
+ .build();
+ }
+
+ /** {@link PTransform} for writing CSV files. */
+ @AutoValue
+ public abstract static class Write<T>
+ extends PTransform<PCollection<T>, WriteFilesResult<String>> {
+
+ /** Specifies the {@link Compression} of all generated shard files. */
+ public Write<T> withCompression(Compression compression) {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withCompression(compression)).build();
+ }
+
+ public Write<T> withNoSpilling() {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withNoSpilling()).build();
+ }
+
+ /**
+ * Specifies to use a given fixed number of shards per window. See {@link
+ * TextIO.Write#withNumShards(int)}.
+ */
+ public Write<T> withNumShards(Integer numShards) {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withNumShards(numShards)).build();
+ }
+
+ /**
+ * Forces a single file as output and empty shard name template. See {@link
+ * TextIO.Write#withoutSharding()}.
+ */
+ public Write<T> withoutSharding() {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withoutSharding()).build();
+ }
+
+ /**
+ * Uses the given {@link ShardNameTemplate} for naming output files. See {@link
+ * TextIO.Write#withShardNameTemplate(String)}.
+ */
+ public Write<T> withShardTemplate(String shardTemplate) {
+ return toBuilder()
+ .setTextIOWrite(getTextIOWrite().withShardNameTemplate(shardTemplate))
+ .build();
+ }
+
+ /**
+ * Configures the filename suffix for written files. See {@link
+ * TextIO.Write#withSuffix(String)}.
+ */
+ public Write<T> withSuffix(String suffix) {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withSuffix(suffix)).build();
+ }
+
+ /**
+ * Set the base directory used to generate temporary files. See {@link
+ * TextIO.Write#withTempDirectory(ResourceId)}.
+ */
+ public Write<T> withTempDirectory(ResourceId tempDirectory) {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withTempDirectory(tempDirectory)).build();
+ }
+
+ /**
+ * Preserves windowing of input elements and writes them to files based on the element's window.
+ * See {@link TextIO.Write#withWindowedWrites()}.
+ */
+ public Write<T> withWindowedWrites() {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withWindowedWrites()).build();
+ }
+
+ /**
+ * Returns a transform for writing to text files like this one but that has the given {@link
+ * FileBasedSink.WritableByteChannelFactory} to be used by the {@link FileBasedSink} during
+ * output. See {@link
+ * TextIO.Write#withWritableByteChannelFactory(FileBasedSink.WritableByteChannelFactory)}.
+ */
+ public Write<T> withWritableByteChannelFactory(
+ FileBasedSink.WritableByteChannelFactory writableByteChannelFactory) {
+ return toBuilder()
+ .setTextIOWrite(
+ getTextIOWrite().withWritableByteChannelFactory(writableByteChannelFactory))
+ .build();
+ }
+
+ /** The underlying {@link FileIO.Write} that writes converted input to CSV formatted output. */
+ abstract TextIO.Write getTextIOWrite();
+
+ /** The {@link CSVFormat} to inform headers, header comments, and format CSV row content. */
+ abstract CSVFormat getCSVFormat();
+
+ abstract Builder<T> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<T> {
+
+ /**
+ * The underlying {@link FileIO.Write} that writes converted input to CSV formatted output.
+ */
+ abstract Builder<T> setTextIOWrite(TextIO.Write value);
+
+ /** The {@link CSVFormat} to convert input. Defaults to {@link CSVFormat#DEFAULT}. */
+ abstract Builder<T> setCSVFormat(CSVFormat value);
+
+ abstract CSVFormat getCSVFormat();
+
+ abstract Write<T> autoBuild();
+
+ final Write<T> build() {
+ checkArgument(
+ !getCSVFormat().getSkipHeaderRecord(),
+ "withSkipHeaderRecord is an illegal CSVFormat setting");
+
+ if (getCSVFormat().getHeaderComments() != null) {
+ checkArgument(
+ getCSVFormat().isCommentMarkerSet(),
+ "CSVFormat withCommentMarker required when withHeaderComments");
+ }
+ return autoBuild();
+ }
+ }
+
+ @Override
+ public WriteFilesResult<String> expand(PCollection<T> input) {
+ if (!input.hasSchema()) {
+ throw new IllegalArgumentException(
+ String.format(
+ "%s requires an input Schema. Note that only Row or user classes are supported. Consider using TextIO or FileIO directly when writing primitive types",
+ Write.class.getName()));
+ }
+
+ Schema schema = input.getSchema();
+
+ RowCoder rowCoder = RowCoder.of(schema);
+
+ PCollection<Row> rows =
+ input
+ .apply("To Rows", MapElements.into(rows()).via(input.getToRowFunction()))
+ .setCoder(rowCoder);
+
+ CSVFormat csvFormat = applyRequiredCSVFormatSettings(schema);
+
+ String header = formatHeader(csvFormat);
+
+ SerializableFunction<Row, String> toCsvFn =
+ CsvRowConversions.RowToCsv.builder()
+ .setCSVFormat(csvFormat)
+ .setSchema(input.getSchema())
+ .build();
+
+ PCollection<String> csv = rows.apply("To CSV", MapElements.into(strings()).via(toCsvFn));
+
+ return csv.apply("Write CSV", getTextIOWrite().withHeader(header).withOutputFilenames());
+ }
+
+ CSVFormat applyRequiredCSVFormatSettings(Schema schema) {
+ CSVFormat csvFormat = getCSVFormat().withSkipHeaderRecord();
+ if (csvFormat.getHeader() == null) {
+ csvFormat = csvFormat.withHeader(schema.sorted().getFieldNames().toArray(new String[0]));
+ }
+ return csvFormat;
+ }
+
+ private static String formatHeader(CSVFormat csvFormat) {
+ String[] header = requireNonNull(csvFormat.getHeader());
+ List<String> result = new ArrayList<>();
+ if (csvFormat.getHeaderComments() != null) {
+ for (String comment : csvFormat.getHeaderComments()) {
+ result.add(csvFormat.getCommentMarker() + " " + comment);
+ }
+ }
+ CSVFormat withoutHeaderComments = csvFormat.withHeaderComments();
+ result.add(withoutHeaderComments.format((Object[]) header));
+ return String.join("\n", result);
+ }
Review Comment:
In response to this comment, I've:
1. Renamed `applyRequiredCSVFormatSettings` to `buildHeaderFromSchemaIfNeeded`
2. Refactored and renamed `formatHeader` to `writeWithCSVFormatHeaderAndComments` to hopefully provide clarity coupling this behavior with TextIO#withHeader.
Additionally, I wanted to address the questions/comments below:
1. > We are ignoring the case where the user wants to skip the header record and doesn't specify the header. We should only specify the csvFormat header if withSkipHeaderRecord is false.
CSVFormat communicates that withSkipHeaderRecord functions to not output the header. Do you agree then we should not use withSkipHeaderRecord to determine where we are building the header?
3. > Note that we must also return null from formatHeader in this case since TextIO#withHeader specifically says that a newline is automatically added after.
Are you saying that the consequence of TextIO#withHeader is:
```
col1, col2, col3
1,2,3
2,3,4
```
vs
```
col1, col2, col3
1,2,3
2,3,4
```
--
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: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [beam] damondouglas commented on a diff in pull request #24630: [24469] Implement CsvIO.Write and supporting classes
Posted by "damondouglas (via GitHub)" <gi...@apache.org>.
damondouglas commented on code in PR #24630:
URL: https://github.com/apache/beam/pull/24630#discussion_r1085716045
##########
sdks/java/io/csv/src/main/java/org/apache/beam/sdk/io/csv/CsvIO.java:
##########
@@ -0,0 +1,446 @@
+/*
+ * 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.sdk.io.csv;
+
+import static java.util.Objects.requireNonNull;
+import static org.apache.beam.sdk.values.TypeDescriptors.rows;
+import static org.apache.beam.sdk.values.TypeDescriptors.strings;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+
+import com.google.auto.value.AutoValue;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Set;
+import org.apache.beam.sdk.coders.RowCoder;
+import org.apache.beam.sdk.io.Compression;
+import org.apache.beam.sdk.io.FileBasedSink;
+import org.apache.beam.sdk.io.FileIO;
+import org.apache.beam.sdk.io.ShardNameTemplate;
+import org.apache.beam.sdk.io.TextIO;
+import org.apache.beam.sdk.io.WriteFilesResult;
+import org.apache.beam.sdk.io.fs.ResourceId;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.transforms.MapElements;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.sdk.values.TypeDescriptor;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableSet;
+import org.apache.commons.csv.CSVFormat;
+
+/**
+ * {@link PTransform}s for reading and writing CSV files.
+ *
+ * <h2>Reading CSV files</h2>
+ *
+ * <p>Reading from CSV files is not yet implemented. Please see <a
+ * href="https://github.com/apache/beam/issues/24552">https://github.com/apache/beam/issues/24552</a>.
+ *
+ * <h2>Writing CSV files</h2>
+ *
+ * <p>To write a {@link PCollection} to one or more CSV files, use {@link CsvIO.Write}, using {@link
+ * CsvIO#writeRows(String, CSVFormat)} or {@link CsvIO#write(String, CSVFormat)}. {@link
+ * CsvIO.Write} supports writing {@link Row} or custom Java types using an inferred {@link Schema}.
+ * Examples below show both scenarios. See the Beam Programming Guide on <a
+ * href="https://beam.apache.org/documentation/programming-guide/#inferring-schemas">inferring
+ * schemas</a> for more information on how to enable Beam to infer a {@link Schema} from a custom
+ * Java type.
+ *
+ * <p>{@link CsvIO.Write} only supports writing {@link Schema} aware types that do not contain any
+ * nested {@link FieldType}s such a {@link org.apache.beam.sdk.schemas.Schema.TypeName#ROW} or
+ * repeated {@link org.apache.beam.sdk.schemas.Schema.TypeName#ARRAY} types. See {@link
+ * CsvIO.Write#VALID_FIELD_TYPE_SET} for valid {@link FieldType}s.
+ *
+ * <h3>Example usage:</h3>
+ *
+ * <p>Suppose we have the following <code>Transaction</code> class annotated with
+ * {@code @DefaultSchema(JavaBeanSchema.class)} so that Beam can infer its {@link Schema}:
+ *
+ * <pre>{@code @DefaultSchema(JavaBeanSchema.class)
+ * public class Transaction {
+ * public Transaction() { … }
+ * public Long getTransactionId();
+ * public void setTransactionId(Long transactionId) { … }
+ * public String getBank() { … }
+ * public void setBank(String bank) { … }
+ * public double getPurchaseAmount() { … }
+ * public void setPurchaseAmount(double purchaseAmount) { … }
+ * }
+ * }</pre>
+ *
+ * <p>From a {@code PCollection<Transaction>}, {@link CsvIO.Write} can write one or many CSV files
+ * automatically creating the header based on its inferred {@link Schema}.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions = ...
+ * transactions.apply(CsvIO.<Transaction>write("path/to/folder/prefix", CSVFormat.DEFAULT));
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header is repeated for every
+ * shard file, whereas by default, {@link CsvIO.Write} will write all fields in <b>sorted order</b>
+ * of the field names.
+ *
+ * <pre>{@code
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ *
+ * <p>To control the order and subset of fields that {@link CsvIO.Write} writes, use {@link
+ * CSVFormat#withHeader(String...)}. Note, however, the following constraints:
+ *
+ * <ol>
+ * <li>Each header column must match a field name in the {@link Schema}; matching is case
+ * sensitive.
+ * <li>Matching header columns must match {@link Schema} fields that are valid {@link FieldType}s;
+ * see {@link #VALID_FIELD_TYPE_SET}.
+ * <li>{@link CSVFormat} only allows repeated header columns when {@link
+ * CSVFormat#withAllowDuplicateHeaderNames()}
+ * </ol>
+ *
+ * <p>The following example shows the use of {@link CSVFormat#withHeader(String...)} to control the
+ * order and subset of <code>Transaction</code> fields.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions ...
+ * transactions.apply(
+ * CsvIO
+ * .<Transaction>write("path/to/folder/prefix", CSVFormat.DEFAULT.withHeader("transactionId", "purchaseAmount"))
+ * );
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header is repeated for every
+ * shard file, but only including the subset of fields in their listed order.
+ *
+ * <pre>{@code
+ * transactionId,purchaseAmount
+ * 12345,10.23
+ * 54321,54.65
+ * 98765,11.76
+ * }</pre>
+ *
+ * <p>In addition to header customization, {@link CsvIO.Write} supports {@link
+ * CSVFormat#withHeaderComments(Object...)} as shown below. Note that {@link
+ * CSVFormat#withCommentMarker(char)} is required when specifying header comments.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions = ...
+ * transactions
+ * .apply(
+ * CsvIO.<Transaction>write("path/to/folder/prefix",
+ * CSVFormat.DEFAULT.withCommentMarker('#').withHeaderComments("Bank Report", "1970-01-01", "Operator: John Doe")
+ * );
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header and header comments are
+ * repeated for every shard file.
+ *
+ * <pre>{@code
+ * # Bank Report
+ * # 1970-01-01
+ * # Operator: John Doe
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ *
+ * <p>A {@link PCollection} of {@link Row}s works just like custom Java types illustrated above,
+ * except we use {@link CsvIO#writeRows(String, CSVFormat)} as shown below for the same {@code
+ * Transaction} class. We derive {@code Transaction}'s {@link Schema} using a {@link
+ * org.apache.beam.sdk.schemas.annotations.DefaultSchema.DefaultSchemaProvider}. Note that
+ * hard-coding the {@link Row}s below is for illustration purposes. Developers are instead
+ * encouraged to take advantage of {@link
+ * org.apache.beam.sdk.schemas.annotations.DefaultSchema.DefaultSchemaProvider#toRowFunction(TypeDescriptor)}.
+ *
+ * <pre>{@code
+ * DefaultSchemaProvider defaultSchemaProvider = new DefaultSchemaProvider();
+ * Schema schema = defaultSchemaProvider.schemaFor(TypeDescriptor.of(Transaction.class));
+ * PCollection<Row> transactions = pipeline.apply(Create.of(
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "A")
+ * .withFieldValue("purchaseAmount", 10.23)
+ * .withFieldValue("transactionId", "12345")
+ * .build(),
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "B")
+ * .withFieldValue("purchaseAmount", 54.65)
+ * .withFieldValue("transactionId", "54321")
+ * .build(),
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "C")
+ * .withFieldValue("purchaseAmount", 11.76)
+ * .withFieldValue("transactionId", "98765")
+ * .build()
+ * );
+ *
+ * transactions.apply(
+ * CsvIO
+ * .writeRowsTo("gs://bucket/path/to/folder/prefix", CSVFormat.DEFAULT)
+ * );
+ * }</pre>
+ *
+ * <p>Writing the transactions {@link PCollection} of {@link Row}s would yield the following CSV
+ * file content.
+ *
+ * <pre>{@code
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ *
+ * {@link CsvIO.Write} does not support the following {@link CSVFormat} properties and will throw an
+ * {@link IllegalArgumentException}.
+ *
+ * <ul>
+ * <li>{@link CSVFormat#withSkipHeaderRecord()}
+ * <li>{@link CSVFormat#withAllowMissingColumnNames()} ()}
+ * <li>{@link CSVFormat#withAutoFlush(boolean)}
+ * <li>{@link CSVFormat#withIgnoreHeaderCase()}
+ * <li>{@link CSVFormat#withIgnoreSurroundingSpaces()}
+ * </ul>
+ */
+public class CsvIO {
+ /**
+ * The valid {@link Schema.FieldType} from which {@link CsvIO} converts CSV records to the fields.
+ *
+ * <ul>
+ * <li>{@link FieldType#BYTE}
+ * <li>{@link FieldType#BOOLEAN}
+ * <li>{@link FieldType#DATETIME}
+ * <li>{@link FieldType#DECIMAL}
+ * <li>{@link FieldType#DOUBLE}
+ * <li>{@link FieldType#INT16}
+ * <li>{@link FieldType#INT32}
+ * <li>{@link FieldType#INT64}
+ * <li>{@link FieldType#FLOAT}
+ * <li>{@link FieldType#STRING}
+ * </ul>
+ */
+ public static final Set<Schema.FieldType> VALID_FIELD_TYPE_SET =
+ ImmutableSet.of(
+ FieldType.BYTE,
+ FieldType.BOOLEAN,
+ FieldType.DATETIME,
+ FieldType.DECIMAL,
+ FieldType.DOUBLE,
+ FieldType.INT16,
+ FieldType.INT32,
+ FieldType.INT64,
+ FieldType.FLOAT,
+ FieldType.STRING);
+
+ static final String DEFAULT_FILENAME_SUFFIX = ".csv";
+
+ /** Instantiates a {@link Write} for writing user types in {@link CSVFormat} format. */
+ public static <T> Write<T> write(String to, CSVFormat csvFormat) {
+ return new AutoValue_CsvIO_Write.Builder<T>()
+ .setTextIOWrite(createDefaultTextIOWrite(to))
+ .setCSVFormat(csvFormat)
+ .build();
+ }
+
+ /** Instantiates a {@link Write} for writing {@link Row}s in {@link CSVFormat} format. */
+ public static Write<Row> writeRows(String to, CSVFormat csvFormat) {
+ return new AutoValue_CsvIO_Write.Builder<Row>()
+ .setTextIOWrite(createDefaultTextIOWrite(to))
+ .setCSVFormat(csvFormat)
+ .build();
+ }
+
+ /** {@link PTransform} for writing CSV files. */
+ @AutoValue
+ public abstract static class Write<T>
+ extends PTransform<PCollection<T>, WriteFilesResult<String>> {
+
+ /** Specifies the {@link Compression} of all generated shard files. */
+ public Write<T> withCompression(Compression compression) {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withCompression(compression)).build();
+ }
+
+ public Write<T> withNoSpilling() {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withNoSpilling()).build();
+ }
+
+ /**
+ * Specifies to use a given fixed number of shards per window. See {@link
+ * TextIO.Write#withNumShards(int)}.
+ */
+ public Write<T> withNumShards(Integer numShards) {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withNumShards(numShards)).build();
+ }
+
+ /**
+ * Forces a single file as output and empty shard name template. See {@link
+ * TextIO.Write#withoutSharding()}.
+ */
+ public Write<T> withoutSharding() {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withoutSharding()).build();
+ }
+
+ /**
+ * Uses the given {@link ShardNameTemplate} for naming output files. See {@link
+ * TextIO.Write#withShardNameTemplate(String)}.
+ */
+ public Write<T> withShardTemplate(String shardTemplate) {
+ return toBuilder()
+ .setTextIOWrite(getTextIOWrite().withShardNameTemplate(shardTemplate))
+ .build();
+ }
+
+ /**
+ * Configures the filename suffix for written files. See {@link
+ * TextIO.Write#withSuffix(String)}.
+ */
+ public Write<T> withSuffix(String suffix) {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withSuffix(suffix)).build();
+ }
+
+ /**
+ * Set the base directory used to generate temporary files. See {@link
+ * TextIO.Write#withTempDirectory(ResourceId)}.
+ */
+ public Write<T> withTempDirectory(ResourceId tempDirectory) {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withTempDirectory(tempDirectory)).build();
+ }
+
+ /**
+ * Preserves windowing of input elements and writes them to files based on the element's window.
+ * See {@link TextIO.Write#withWindowedWrites()}.
+ */
+ public Write<T> withWindowedWrites() {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withWindowedWrites()).build();
+ }
+
+ /**
+ * Returns a transform for writing to text files like this one but that has the given {@link
+ * FileBasedSink.WritableByteChannelFactory} to be used by the {@link FileBasedSink} during
+ * output. See {@link
+ * TextIO.Write#withWritableByteChannelFactory(FileBasedSink.WritableByteChannelFactory)}.
+ */
+ public Write<T> withWritableByteChannelFactory(
+ FileBasedSink.WritableByteChannelFactory writableByteChannelFactory) {
+ return toBuilder()
+ .setTextIOWrite(
+ getTextIOWrite().withWritableByteChannelFactory(writableByteChannelFactory))
+ .build();
+ }
+
+ /** The underlying {@link FileIO.Write} that writes converted input to CSV formatted output. */
+ abstract TextIO.Write getTextIOWrite();
+
+ /** The {@link CSVFormat} to inform headers, header comments, and format CSV row content. */
+ abstract CSVFormat getCSVFormat();
+
+ abstract Builder<T> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<T> {
+
+ /**
+ * The underlying {@link FileIO.Write} that writes converted input to CSV formatted output.
+ */
+ abstract Builder<T> setTextIOWrite(TextIO.Write value);
+
+ /** The {@link CSVFormat} to convert input. Defaults to {@link CSVFormat#DEFAULT}. */
+ abstract Builder<T> setCSVFormat(CSVFormat value);
+
+ abstract CSVFormat getCSVFormat();
+
+ abstract Write<T> autoBuild();
+
+ final Write<T> build() {
+ checkArgument(
+ !getCSVFormat().getSkipHeaderRecord(),
+ "withSkipHeaderRecord is an illegal CSVFormat setting");
+
+ if (getCSVFormat().getHeaderComments() != null) {
+ checkArgument(
+ getCSVFormat().isCommentMarkerSet(),
+ "CSVFormat withCommentMarker required when withHeaderComments");
+ }
+ return autoBuild();
+ }
+ }
+
+ @Override
+ public WriteFilesResult<String> expand(PCollection<T> input) {
+ if (!input.hasSchema()) {
+ throw new IllegalArgumentException(
+ String.format(
+ "%s requires an input Schema. Note that only Row or user classes are supported. Consider using TextIO or FileIO directly when writing primitive types",
+ Write.class.getName()));
+ }
+
+ Schema schema = input.getSchema();
+
+ RowCoder rowCoder = RowCoder.of(schema);
+
+ PCollection<Row> rows =
+ input
+ .apply("To Rows", MapElements.into(rows()).via(input.getToRowFunction()))
+ .setCoder(rowCoder);
+
+ CSVFormat csvFormat = applyRequiredCSVFormatSettings(schema);
+
+ String header = formatHeader(csvFormat);
+
+ SerializableFunction<Row, String> toCsvFn =
+ CsvRowConversions.RowToCsv.builder()
+ .setCSVFormat(csvFormat)
+ .setSchema(input.getSchema())
+ .build();
+
+ PCollection<String> csv = rows.apply("To CSV", MapElements.into(strings()).via(toCsvFn));
+
+ return csv.apply("Write CSV", getTextIOWrite().withHeader(header).withOutputFilenames());
+ }
+
+ CSVFormat applyRequiredCSVFormatSettings(Schema schema) {
+ CSVFormat csvFormat = getCSVFormat().withSkipHeaderRecord();
+ if (csvFormat.getHeader() == null) {
+ csvFormat = csvFormat.withHeader(schema.sorted().getFieldNames().toArray(new String[0]));
+ }
+ return csvFormat;
+ }
+
+ private static String formatHeader(CSVFormat csvFormat) {
+ String[] header = requireNonNull(csvFormat.getHeader());
+ List<String> result = new ArrayList<>();
+ if (csvFormat.getHeaderComments() != null) {
+ for (String comment : csvFormat.getHeaderComments()) {
+ result.add(csvFormat.getCommentMarker() + " " + comment);
+ }
+ }
+ CSVFormat withoutHeaderComments = csvFormat.withHeaderComments();
+ result.add(withoutHeaderComments.format((Object[]) header));
+ return String.join("\n", result);
+ }
Review Comment:
In response to this comment, I've:
1. Renamed `applyRequiredCSVFormatSettings` to `buildHeaderFromSchemaIfNeeded`
2. Refactored and renamed `formatHeader` to `writeWithCSVFormatHeaderAndComments` to hopefully provide clarity coupling this behavior with TextIO#withHeader.
Additionally, I wanted to address the questions/comments below:
1. > We are ignoring the case where the user wants to skip the header record and doesn't specify the header. We should only specify the csvFormat header if withSkipHeaderRecord is false.
CSVFormat communicates that withSkipHeaderRecord functions to not output the header. It's also a parameter that gives an output incompatible in a beam context as CSVFormat was designed for writing a single file. Therefore, do you agree then we should not use withSkipHeaderRecord to determine where we are building the header?
3. > Note that we must also return null from formatHeader in this case since TextIO#withHeader specifically says that a newline is automatically added after.
Apologies, your comment confused me.
Are you saying that the consequence of TextIO#withHeader is:
```
col1, col2, col3
1,2,3
2,3,4
```
vs
```
col1, col2, col3
1,2,3
2,3,4
```
--
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: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [beam] damondouglas commented on a diff in pull request #24630: [24469] Implement CsvIO.Write and supporting classes
Posted by "damondouglas (via GitHub)" <gi...@apache.org>.
damondouglas commented on code in PR #24630:
URL: https://github.com/apache/beam/pull/24630#discussion_r1085716045
##########
sdks/java/io/csv/src/main/java/org/apache/beam/sdk/io/csv/CsvIO.java:
##########
@@ -0,0 +1,446 @@
+/*
+ * 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.sdk.io.csv;
+
+import static java.util.Objects.requireNonNull;
+import static org.apache.beam.sdk.values.TypeDescriptors.rows;
+import static org.apache.beam.sdk.values.TypeDescriptors.strings;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+
+import com.google.auto.value.AutoValue;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Set;
+import org.apache.beam.sdk.coders.RowCoder;
+import org.apache.beam.sdk.io.Compression;
+import org.apache.beam.sdk.io.FileBasedSink;
+import org.apache.beam.sdk.io.FileIO;
+import org.apache.beam.sdk.io.ShardNameTemplate;
+import org.apache.beam.sdk.io.TextIO;
+import org.apache.beam.sdk.io.WriteFilesResult;
+import org.apache.beam.sdk.io.fs.ResourceId;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.transforms.MapElements;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.sdk.values.TypeDescriptor;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableSet;
+import org.apache.commons.csv.CSVFormat;
+
+/**
+ * {@link PTransform}s for reading and writing CSV files.
+ *
+ * <h2>Reading CSV files</h2>
+ *
+ * <p>Reading from CSV files is not yet implemented. Please see <a
+ * href="https://github.com/apache/beam/issues/24552">https://github.com/apache/beam/issues/24552</a>.
+ *
+ * <h2>Writing CSV files</h2>
+ *
+ * <p>To write a {@link PCollection} to one or more CSV files, use {@link CsvIO.Write}, using {@link
+ * CsvIO#writeRows(String, CSVFormat)} or {@link CsvIO#write(String, CSVFormat)}. {@link
+ * CsvIO.Write} supports writing {@link Row} or custom Java types using an inferred {@link Schema}.
+ * Examples below show both scenarios. See the Beam Programming Guide on <a
+ * href="https://beam.apache.org/documentation/programming-guide/#inferring-schemas">inferring
+ * schemas</a> for more information on how to enable Beam to infer a {@link Schema} from a custom
+ * Java type.
+ *
+ * <p>{@link CsvIO.Write} only supports writing {@link Schema} aware types that do not contain any
+ * nested {@link FieldType}s such a {@link org.apache.beam.sdk.schemas.Schema.TypeName#ROW} or
+ * repeated {@link org.apache.beam.sdk.schemas.Schema.TypeName#ARRAY} types. See {@link
+ * CsvIO.Write#VALID_FIELD_TYPE_SET} for valid {@link FieldType}s.
+ *
+ * <h3>Example usage:</h3>
+ *
+ * <p>Suppose we have the following <code>Transaction</code> class annotated with
+ * {@code @DefaultSchema(JavaBeanSchema.class)} so that Beam can infer its {@link Schema}:
+ *
+ * <pre>{@code @DefaultSchema(JavaBeanSchema.class)
+ * public class Transaction {
+ * public Transaction() { … }
+ * public Long getTransactionId();
+ * public void setTransactionId(Long transactionId) { … }
+ * public String getBank() { … }
+ * public void setBank(String bank) { … }
+ * public double getPurchaseAmount() { … }
+ * public void setPurchaseAmount(double purchaseAmount) { … }
+ * }
+ * }</pre>
+ *
+ * <p>From a {@code PCollection<Transaction>}, {@link CsvIO.Write} can write one or many CSV files
+ * automatically creating the header based on its inferred {@link Schema}.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions = ...
+ * transactions.apply(CsvIO.<Transaction>write("path/to/folder/prefix", CSVFormat.DEFAULT));
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header is repeated for every
+ * shard file, whereas by default, {@link CsvIO.Write} will write all fields in <b>sorted order</b>
+ * of the field names.
+ *
+ * <pre>{@code
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ *
+ * <p>To control the order and subset of fields that {@link CsvIO.Write} writes, use {@link
+ * CSVFormat#withHeader(String...)}. Note, however, the following constraints:
+ *
+ * <ol>
+ * <li>Each header column must match a field name in the {@link Schema}; matching is case
+ * sensitive.
+ * <li>Matching header columns must match {@link Schema} fields that are valid {@link FieldType}s;
+ * see {@link #VALID_FIELD_TYPE_SET}.
+ * <li>{@link CSVFormat} only allows repeated header columns when {@link
+ * CSVFormat#withAllowDuplicateHeaderNames()}
+ * </ol>
+ *
+ * <p>The following example shows the use of {@link CSVFormat#withHeader(String...)} to control the
+ * order and subset of <code>Transaction</code> fields.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions ...
+ * transactions.apply(
+ * CsvIO
+ * .<Transaction>write("path/to/folder/prefix", CSVFormat.DEFAULT.withHeader("transactionId", "purchaseAmount"))
+ * );
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header is repeated for every
+ * shard file, but only including the subset of fields in their listed order.
+ *
+ * <pre>{@code
+ * transactionId,purchaseAmount
+ * 12345,10.23
+ * 54321,54.65
+ * 98765,11.76
+ * }</pre>
+ *
+ * <p>In addition to header customization, {@link CsvIO.Write} supports {@link
+ * CSVFormat#withHeaderComments(Object...)} as shown below. Note that {@link
+ * CSVFormat#withCommentMarker(char)} is required when specifying header comments.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions = ...
+ * transactions
+ * .apply(
+ * CsvIO.<Transaction>write("path/to/folder/prefix",
+ * CSVFormat.DEFAULT.withCommentMarker('#').withHeaderComments("Bank Report", "1970-01-01", "Operator: John Doe")
+ * );
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header and header comments are
+ * repeated for every shard file.
+ *
+ * <pre>{@code
+ * # Bank Report
+ * # 1970-01-01
+ * # Operator: John Doe
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ *
+ * <p>A {@link PCollection} of {@link Row}s works just like custom Java types illustrated above,
+ * except we use {@link CsvIO#writeRows(String, CSVFormat)} as shown below for the same {@code
+ * Transaction} class. We derive {@code Transaction}'s {@link Schema} using a {@link
+ * org.apache.beam.sdk.schemas.annotations.DefaultSchema.DefaultSchemaProvider}. Note that
+ * hard-coding the {@link Row}s below is for illustration purposes. Developers are instead
+ * encouraged to take advantage of {@link
+ * org.apache.beam.sdk.schemas.annotations.DefaultSchema.DefaultSchemaProvider#toRowFunction(TypeDescriptor)}.
+ *
+ * <pre>{@code
+ * DefaultSchemaProvider defaultSchemaProvider = new DefaultSchemaProvider();
+ * Schema schema = defaultSchemaProvider.schemaFor(TypeDescriptor.of(Transaction.class));
+ * PCollection<Row> transactions = pipeline.apply(Create.of(
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "A")
+ * .withFieldValue("purchaseAmount", 10.23)
+ * .withFieldValue("transactionId", "12345")
+ * .build(),
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "B")
+ * .withFieldValue("purchaseAmount", 54.65)
+ * .withFieldValue("transactionId", "54321")
+ * .build(),
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "C")
+ * .withFieldValue("purchaseAmount", 11.76)
+ * .withFieldValue("transactionId", "98765")
+ * .build()
+ * );
+ *
+ * transactions.apply(
+ * CsvIO
+ * .writeRowsTo("gs://bucket/path/to/folder/prefix", CSVFormat.DEFAULT)
+ * );
+ * }</pre>
+ *
+ * <p>Writing the transactions {@link PCollection} of {@link Row}s would yield the following CSV
+ * file content.
+ *
+ * <pre>{@code
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ *
+ * {@link CsvIO.Write} does not support the following {@link CSVFormat} properties and will throw an
+ * {@link IllegalArgumentException}.
+ *
+ * <ul>
+ * <li>{@link CSVFormat#withSkipHeaderRecord()}
+ * <li>{@link CSVFormat#withAllowMissingColumnNames()} ()}
+ * <li>{@link CSVFormat#withAutoFlush(boolean)}
+ * <li>{@link CSVFormat#withIgnoreHeaderCase()}
+ * <li>{@link CSVFormat#withIgnoreSurroundingSpaces()}
+ * </ul>
+ */
+public class CsvIO {
+ /**
+ * The valid {@link Schema.FieldType} from which {@link CsvIO} converts CSV records to the fields.
+ *
+ * <ul>
+ * <li>{@link FieldType#BYTE}
+ * <li>{@link FieldType#BOOLEAN}
+ * <li>{@link FieldType#DATETIME}
+ * <li>{@link FieldType#DECIMAL}
+ * <li>{@link FieldType#DOUBLE}
+ * <li>{@link FieldType#INT16}
+ * <li>{@link FieldType#INT32}
+ * <li>{@link FieldType#INT64}
+ * <li>{@link FieldType#FLOAT}
+ * <li>{@link FieldType#STRING}
+ * </ul>
+ */
+ public static final Set<Schema.FieldType> VALID_FIELD_TYPE_SET =
+ ImmutableSet.of(
+ FieldType.BYTE,
+ FieldType.BOOLEAN,
+ FieldType.DATETIME,
+ FieldType.DECIMAL,
+ FieldType.DOUBLE,
+ FieldType.INT16,
+ FieldType.INT32,
+ FieldType.INT64,
+ FieldType.FLOAT,
+ FieldType.STRING);
+
+ static final String DEFAULT_FILENAME_SUFFIX = ".csv";
+
+ /** Instantiates a {@link Write} for writing user types in {@link CSVFormat} format. */
+ public static <T> Write<T> write(String to, CSVFormat csvFormat) {
+ return new AutoValue_CsvIO_Write.Builder<T>()
+ .setTextIOWrite(createDefaultTextIOWrite(to))
+ .setCSVFormat(csvFormat)
+ .build();
+ }
+
+ /** Instantiates a {@link Write} for writing {@link Row}s in {@link CSVFormat} format. */
+ public static Write<Row> writeRows(String to, CSVFormat csvFormat) {
+ return new AutoValue_CsvIO_Write.Builder<Row>()
+ .setTextIOWrite(createDefaultTextIOWrite(to))
+ .setCSVFormat(csvFormat)
+ .build();
+ }
+
+ /** {@link PTransform} for writing CSV files. */
+ @AutoValue
+ public abstract static class Write<T>
+ extends PTransform<PCollection<T>, WriteFilesResult<String>> {
+
+ /** Specifies the {@link Compression} of all generated shard files. */
+ public Write<T> withCompression(Compression compression) {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withCompression(compression)).build();
+ }
+
+ public Write<T> withNoSpilling() {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withNoSpilling()).build();
+ }
+
+ /**
+ * Specifies to use a given fixed number of shards per window. See {@link
+ * TextIO.Write#withNumShards(int)}.
+ */
+ public Write<T> withNumShards(Integer numShards) {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withNumShards(numShards)).build();
+ }
+
+ /**
+ * Forces a single file as output and empty shard name template. See {@link
+ * TextIO.Write#withoutSharding()}.
+ */
+ public Write<T> withoutSharding() {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withoutSharding()).build();
+ }
+
+ /**
+ * Uses the given {@link ShardNameTemplate} for naming output files. See {@link
+ * TextIO.Write#withShardNameTemplate(String)}.
+ */
+ public Write<T> withShardTemplate(String shardTemplate) {
+ return toBuilder()
+ .setTextIOWrite(getTextIOWrite().withShardNameTemplate(shardTemplate))
+ .build();
+ }
+
+ /**
+ * Configures the filename suffix for written files. See {@link
+ * TextIO.Write#withSuffix(String)}.
+ */
+ public Write<T> withSuffix(String suffix) {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withSuffix(suffix)).build();
+ }
+
+ /**
+ * Set the base directory used to generate temporary files. See {@link
+ * TextIO.Write#withTempDirectory(ResourceId)}.
+ */
+ public Write<T> withTempDirectory(ResourceId tempDirectory) {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withTempDirectory(tempDirectory)).build();
+ }
+
+ /**
+ * Preserves windowing of input elements and writes them to files based on the element's window.
+ * See {@link TextIO.Write#withWindowedWrites()}.
+ */
+ public Write<T> withWindowedWrites() {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withWindowedWrites()).build();
+ }
+
+ /**
+ * Returns a transform for writing to text files like this one but that has the given {@link
+ * FileBasedSink.WritableByteChannelFactory} to be used by the {@link FileBasedSink} during
+ * output. See {@link
+ * TextIO.Write#withWritableByteChannelFactory(FileBasedSink.WritableByteChannelFactory)}.
+ */
+ public Write<T> withWritableByteChannelFactory(
+ FileBasedSink.WritableByteChannelFactory writableByteChannelFactory) {
+ return toBuilder()
+ .setTextIOWrite(
+ getTextIOWrite().withWritableByteChannelFactory(writableByteChannelFactory))
+ .build();
+ }
+
+ /** The underlying {@link FileIO.Write} that writes converted input to CSV formatted output. */
+ abstract TextIO.Write getTextIOWrite();
+
+ /** The {@link CSVFormat} to inform headers, header comments, and format CSV row content. */
+ abstract CSVFormat getCSVFormat();
+
+ abstract Builder<T> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<T> {
+
+ /**
+ * The underlying {@link FileIO.Write} that writes converted input to CSV formatted output.
+ */
+ abstract Builder<T> setTextIOWrite(TextIO.Write value);
+
+ /** The {@link CSVFormat} to convert input. Defaults to {@link CSVFormat#DEFAULT}. */
+ abstract Builder<T> setCSVFormat(CSVFormat value);
+
+ abstract CSVFormat getCSVFormat();
+
+ abstract Write<T> autoBuild();
+
+ final Write<T> build() {
+ checkArgument(
+ !getCSVFormat().getSkipHeaderRecord(),
+ "withSkipHeaderRecord is an illegal CSVFormat setting");
+
+ if (getCSVFormat().getHeaderComments() != null) {
+ checkArgument(
+ getCSVFormat().isCommentMarkerSet(),
+ "CSVFormat withCommentMarker required when withHeaderComments");
+ }
+ return autoBuild();
+ }
+ }
+
+ @Override
+ public WriteFilesResult<String> expand(PCollection<T> input) {
+ if (!input.hasSchema()) {
+ throw new IllegalArgumentException(
+ String.format(
+ "%s requires an input Schema. Note that only Row or user classes are supported. Consider using TextIO or FileIO directly when writing primitive types",
+ Write.class.getName()));
+ }
+
+ Schema schema = input.getSchema();
+
+ RowCoder rowCoder = RowCoder.of(schema);
+
+ PCollection<Row> rows =
+ input
+ .apply("To Rows", MapElements.into(rows()).via(input.getToRowFunction()))
+ .setCoder(rowCoder);
+
+ CSVFormat csvFormat = applyRequiredCSVFormatSettings(schema);
+
+ String header = formatHeader(csvFormat);
+
+ SerializableFunction<Row, String> toCsvFn =
+ CsvRowConversions.RowToCsv.builder()
+ .setCSVFormat(csvFormat)
+ .setSchema(input.getSchema())
+ .build();
+
+ PCollection<String> csv = rows.apply("To CSV", MapElements.into(strings()).via(toCsvFn));
+
+ return csv.apply("Write CSV", getTextIOWrite().withHeader(header).withOutputFilenames());
+ }
+
+ CSVFormat applyRequiredCSVFormatSettings(Schema schema) {
+ CSVFormat csvFormat = getCSVFormat().withSkipHeaderRecord();
+ if (csvFormat.getHeader() == null) {
+ csvFormat = csvFormat.withHeader(schema.sorted().getFieldNames().toArray(new String[0]));
+ }
+ return csvFormat;
+ }
+
+ private static String formatHeader(CSVFormat csvFormat) {
+ String[] header = requireNonNull(csvFormat.getHeader());
+ List<String> result = new ArrayList<>();
+ if (csvFormat.getHeaderComments() != null) {
+ for (String comment : csvFormat.getHeaderComments()) {
+ result.add(csvFormat.getCommentMarker() + " " + comment);
+ }
+ }
+ CSVFormat withoutHeaderComments = csvFormat.withHeaderComments();
+ result.add(withoutHeaderComments.format((Object[]) header));
+ return String.join("\n", result);
+ }
Review Comment:
In response to this comment, I've:
1. Renamed `applyRequiredCSVFormatSettings` to [`buildHeaderFromSchemaIfNeeded`](https://github.com/damondouglas/beam/blob/24469-csv-io-write/sdks/java/io/csv/src/main/java/org/apache/beam/sdk/io/csv/CsvIO.java#L426)
2. Refactored and renamed `formatHeader` to [`writeWithCSVFormatHeaderAndComments`](https://github.com/damondouglas/beam/blob/24469-csv-io-write/sdks/java/io/csv/src/main/java/org/apache/beam/sdk/io/csv/CsvIO.java#L436) to hopefully provide clarity coupling this behavior with TextIO#withHeader.
Additionally, I wanted to address the questions/comments below:
1. > We are ignoring the case where the user wants to skip the header record and doesn't specify the header. We should only specify the csvFormat header if withSkipHeaderRecord is false.
CSVFormat communicates that withSkipHeaderRecord functions to not output the header. It's also a parameter that gives an output incompatible in a beam context as CSVFormat was designed for writing a single file. Therefore, do you agree then we should not use withSkipHeaderRecord to drive how we are building the header?
3. > Note that we must also return null from formatHeader in this case since TextIO#withHeader specifically says that a newline is automatically added after.
Apologies, your comment confused me.
Are you saying that the consequence of TextIO#withHeader is:
```
col1, col2, col3
1,2,3
2,3,4
```
vs
```
col1, col2, col3
1,2,3
2,3,4
```
--
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: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [beam] damondouglas commented on a diff in pull request #24630: [24469] Implement CsvIO.Write and supporting classes
Posted by "damondouglas (via GitHub)" <gi...@apache.org>.
damondouglas commented on code in PR #24630:
URL: https://github.com/apache/beam/pull/24630#discussion_r1085716045
##########
sdks/java/io/csv/src/main/java/org/apache/beam/sdk/io/csv/CsvIO.java:
##########
@@ -0,0 +1,446 @@
+/*
+ * 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.sdk.io.csv;
+
+import static java.util.Objects.requireNonNull;
+import static org.apache.beam.sdk.values.TypeDescriptors.rows;
+import static org.apache.beam.sdk.values.TypeDescriptors.strings;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+
+import com.google.auto.value.AutoValue;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Set;
+import org.apache.beam.sdk.coders.RowCoder;
+import org.apache.beam.sdk.io.Compression;
+import org.apache.beam.sdk.io.FileBasedSink;
+import org.apache.beam.sdk.io.FileIO;
+import org.apache.beam.sdk.io.ShardNameTemplate;
+import org.apache.beam.sdk.io.TextIO;
+import org.apache.beam.sdk.io.WriteFilesResult;
+import org.apache.beam.sdk.io.fs.ResourceId;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.transforms.MapElements;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.sdk.values.TypeDescriptor;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableSet;
+import org.apache.commons.csv.CSVFormat;
+
+/**
+ * {@link PTransform}s for reading and writing CSV files.
+ *
+ * <h2>Reading CSV files</h2>
+ *
+ * <p>Reading from CSV files is not yet implemented. Please see <a
+ * href="https://github.com/apache/beam/issues/24552">https://github.com/apache/beam/issues/24552</a>.
+ *
+ * <h2>Writing CSV files</h2>
+ *
+ * <p>To write a {@link PCollection} to one or more CSV files, use {@link CsvIO.Write}, using {@link
+ * CsvIO#writeRows(String, CSVFormat)} or {@link CsvIO#write(String, CSVFormat)}. {@link
+ * CsvIO.Write} supports writing {@link Row} or custom Java types using an inferred {@link Schema}.
+ * Examples below show both scenarios. See the Beam Programming Guide on <a
+ * href="https://beam.apache.org/documentation/programming-guide/#inferring-schemas">inferring
+ * schemas</a> for more information on how to enable Beam to infer a {@link Schema} from a custom
+ * Java type.
+ *
+ * <p>{@link CsvIO.Write} only supports writing {@link Schema} aware types that do not contain any
+ * nested {@link FieldType}s such a {@link org.apache.beam.sdk.schemas.Schema.TypeName#ROW} or
+ * repeated {@link org.apache.beam.sdk.schemas.Schema.TypeName#ARRAY} types. See {@link
+ * CsvIO.Write#VALID_FIELD_TYPE_SET} for valid {@link FieldType}s.
+ *
+ * <h3>Example usage:</h3>
+ *
+ * <p>Suppose we have the following <code>Transaction</code> class annotated with
+ * {@code @DefaultSchema(JavaBeanSchema.class)} so that Beam can infer its {@link Schema}:
+ *
+ * <pre>{@code @DefaultSchema(JavaBeanSchema.class)
+ * public class Transaction {
+ * public Transaction() { … }
+ * public Long getTransactionId();
+ * public void setTransactionId(Long transactionId) { … }
+ * public String getBank() { … }
+ * public void setBank(String bank) { … }
+ * public double getPurchaseAmount() { … }
+ * public void setPurchaseAmount(double purchaseAmount) { … }
+ * }
+ * }</pre>
+ *
+ * <p>From a {@code PCollection<Transaction>}, {@link CsvIO.Write} can write one or many CSV files
+ * automatically creating the header based on its inferred {@link Schema}.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions = ...
+ * transactions.apply(CsvIO.<Transaction>write("path/to/folder/prefix", CSVFormat.DEFAULT));
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header is repeated for every
+ * shard file, whereas by default, {@link CsvIO.Write} will write all fields in <b>sorted order</b>
+ * of the field names.
+ *
+ * <pre>{@code
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ *
+ * <p>To control the order and subset of fields that {@link CsvIO.Write} writes, use {@link
+ * CSVFormat#withHeader(String...)}. Note, however, the following constraints:
+ *
+ * <ol>
+ * <li>Each header column must match a field name in the {@link Schema}; matching is case
+ * sensitive.
+ * <li>Matching header columns must match {@link Schema} fields that are valid {@link FieldType}s;
+ * see {@link #VALID_FIELD_TYPE_SET}.
+ * <li>{@link CSVFormat} only allows repeated header columns when {@link
+ * CSVFormat#withAllowDuplicateHeaderNames()}
+ * </ol>
+ *
+ * <p>The following example shows the use of {@link CSVFormat#withHeader(String...)} to control the
+ * order and subset of <code>Transaction</code> fields.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions ...
+ * transactions.apply(
+ * CsvIO
+ * .<Transaction>write("path/to/folder/prefix", CSVFormat.DEFAULT.withHeader("transactionId", "purchaseAmount"))
+ * );
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header is repeated for every
+ * shard file, but only including the subset of fields in their listed order.
+ *
+ * <pre>{@code
+ * transactionId,purchaseAmount
+ * 12345,10.23
+ * 54321,54.65
+ * 98765,11.76
+ * }</pre>
+ *
+ * <p>In addition to header customization, {@link CsvIO.Write} supports {@link
+ * CSVFormat#withHeaderComments(Object...)} as shown below. Note that {@link
+ * CSVFormat#withCommentMarker(char)} is required when specifying header comments.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions = ...
+ * transactions
+ * .apply(
+ * CsvIO.<Transaction>write("path/to/folder/prefix",
+ * CSVFormat.DEFAULT.withCommentMarker('#').withHeaderComments("Bank Report", "1970-01-01", "Operator: John Doe")
+ * );
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header and header comments are
+ * repeated for every shard file.
+ *
+ * <pre>{@code
+ * # Bank Report
+ * # 1970-01-01
+ * # Operator: John Doe
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ *
+ * <p>A {@link PCollection} of {@link Row}s works just like custom Java types illustrated above,
+ * except we use {@link CsvIO#writeRows(String, CSVFormat)} as shown below for the same {@code
+ * Transaction} class. We derive {@code Transaction}'s {@link Schema} using a {@link
+ * org.apache.beam.sdk.schemas.annotations.DefaultSchema.DefaultSchemaProvider}. Note that
+ * hard-coding the {@link Row}s below is for illustration purposes. Developers are instead
+ * encouraged to take advantage of {@link
+ * org.apache.beam.sdk.schemas.annotations.DefaultSchema.DefaultSchemaProvider#toRowFunction(TypeDescriptor)}.
+ *
+ * <pre>{@code
+ * DefaultSchemaProvider defaultSchemaProvider = new DefaultSchemaProvider();
+ * Schema schema = defaultSchemaProvider.schemaFor(TypeDescriptor.of(Transaction.class));
+ * PCollection<Row> transactions = pipeline.apply(Create.of(
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "A")
+ * .withFieldValue("purchaseAmount", 10.23)
+ * .withFieldValue("transactionId", "12345")
+ * .build(),
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "B")
+ * .withFieldValue("purchaseAmount", 54.65)
+ * .withFieldValue("transactionId", "54321")
+ * .build(),
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "C")
+ * .withFieldValue("purchaseAmount", 11.76)
+ * .withFieldValue("transactionId", "98765")
+ * .build()
+ * );
+ *
+ * transactions.apply(
+ * CsvIO
+ * .writeRowsTo("gs://bucket/path/to/folder/prefix", CSVFormat.DEFAULT)
+ * );
+ * }</pre>
+ *
+ * <p>Writing the transactions {@link PCollection} of {@link Row}s would yield the following CSV
+ * file content.
+ *
+ * <pre>{@code
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ *
+ * {@link CsvIO.Write} does not support the following {@link CSVFormat} properties and will throw an
+ * {@link IllegalArgumentException}.
+ *
+ * <ul>
+ * <li>{@link CSVFormat#withSkipHeaderRecord()}
+ * <li>{@link CSVFormat#withAllowMissingColumnNames()} ()}
+ * <li>{@link CSVFormat#withAutoFlush(boolean)}
+ * <li>{@link CSVFormat#withIgnoreHeaderCase()}
+ * <li>{@link CSVFormat#withIgnoreSurroundingSpaces()}
+ * </ul>
+ */
+public class CsvIO {
+ /**
+ * The valid {@link Schema.FieldType} from which {@link CsvIO} converts CSV records to the fields.
+ *
+ * <ul>
+ * <li>{@link FieldType#BYTE}
+ * <li>{@link FieldType#BOOLEAN}
+ * <li>{@link FieldType#DATETIME}
+ * <li>{@link FieldType#DECIMAL}
+ * <li>{@link FieldType#DOUBLE}
+ * <li>{@link FieldType#INT16}
+ * <li>{@link FieldType#INT32}
+ * <li>{@link FieldType#INT64}
+ * <li>{@link FieldType#FLOAT}
+ * <li>{@link FieldType#STRING}
+ * </ul>
+ */
+ public static final Set<Schema.FieldType> VALID_FIELD_TYPE_SET =
+ ImmutableSet.of(
+ FieldType.BYTE,
+ FieldType.BOOLEAN,
+ FieldType.DATETIME,
+ FieldType.DECIMAL,
+ FieldType.DOUBLE,
+ FieldType.INT16,
+ FieldType.INT32,
+ FieldType.INT64,
+ FieldType.FLOAT,
+ FieldType.STRING);
+
+ static final String DEFAULT_FILENAME_SUFFIX = ".csv";
+
+ /** Instantiates a {@link Write} for writing user types in {@link CSVFormat} format. */
+ public static <T> Write<T> write(String to, CSVFormat csvFormat) {
+ return new AutoValue_CsvIO_Write.Builder<T>()
+ .setTextIOWrite(createDefaultTextIOWrite(to))
+ .setCSVFormat(csvFormat)
+ .build();
+ }
+
+ /** Instantiates a {@link Write} for writing {@link Row}s in {@link CSVFormat} format. */
+ public static Write<Row> writeRows(String to, CSVFormat csvFormat) {
+ return new AutoValue_CsvIO_Write.Builder<Row>()
+ .setTextIOWrite(createDefaultTextIOWrite(to))
+ .setCSVFormat(csvFormat)
+ .build();
+ }
+
+ /** {@link PTransform} for writing CSV files. */
+ @AutoValue
+ public abstract static class Write<T>
+ extends PTransform<PCollection<T>, WriteFilesResult<String>> {
+
+ /** Specifies the {@link Compression} of all generated shard files. */
+ public Write<T> withCompression(Compression compression) {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withCompression(compression)).build();
+ }
+
+ public Write<T> withNoSpilling() {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withNoSpilling()).build();
+ }
+
+ /**
+ * Specifies to use a given fixed number of shards per window. See {@link
+ * TextIO.Write#withNumShards(int)}.
+ */
+ public Write<T> withNumShards(Integer numShards) {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withNumShards(numShards)).build();
+ }
+
+ /**
+ * Forces a single file as output and empty shard name template. See {@link
+ * TextIO.Write#withoutSharding()}.
+ */
+ public Write<T> withoutSharding() {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withoutSharding()).build();
+ }
+
+ /**
+ * Uses the given {@link ShardNameTemplate} for naming output files. See {@link
+ * TextIO.Write#withShardNameTemplate(String)}.
+ */
+ public Write<T> withShardTemplate(String shardTemplate) {
+ return toBuilder()
+ .setTextIOWrite(getTextIOWrite().withShardNameTemplate(shardTemplate))
+ .build();
+ }
+
+ /**
+ * Configures the filename suffix for written files. See {@link
+ * TextIO.Write#withSuffix(String)}.
+ */
+ public Write<T> withSuffix(String suffix) {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withSuffix(suffix)).build();
+ }
+
+ /**
+ * Set the base directory used to generate temporary files. See {@link
+ * TextIO.Write#withTempDirectory(ResourceId)}.
+ */
+ public Write<T> withTempDirectory(ResourceId tempDirectory) {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withTempDirectory(tempDirectory)).build();
+ }
+
+ /**
+ * Preserves windowing of input elements and writes them to files based on the element's window.
+ * See {@link TextIO.Write#withWindowedWrites()}.
+ */
+ public Write<T> withWindowedWrites() {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withWindowedWrites()).build();
+ }
+
+ /**
+ * Returns a transform for writing to text files like this one but that has the given {@link
+ * FileBasedSink.WritableByteChannelFactory} to be used by the {@link FileBasedSink} during
+ * output. See {@link
+ * TextIO.Write#withWritableByteChannelFactory(FileBasedSink.WritableByteChannelFactory)}.
+ */
+ public Write<T> withWritableByteChannelFactory(
+ FileBasedSink.WritableByteChannelFactory writableByteChannelFactory) {
+ return toBuilder()
+ .setTextIOWrite(
+ getTextIOWrite().withWritableByteChannelFactory(writableByteChannelFactory))
+ .build();
+ }
+
+ /** The underlying {@link FileIO.Write} that writes converted input to CSV formatted output. */
+ abstract TextIO.Write getTextIOWrite();
+
+ /** The {@link CSVFormat} to inform headers, header comments, and format CSV row content. */
+ abstract CSVFormat getCSVFormat();
+
+ abstract Builder<T> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<T> {
+
+ /**
+ * The underlying {@link FileIO.Write} that writes converted input to CSV formatted output.
+ */
+ abstract Builder<T> setTextIOWrite(TextIO.Write value);
+
+ /** The {@link CSVFormat} to convert input. Defaults to {@link CSVFormat#DEFAULT}. */
+ abstract Builder<T> setCSVFormat(CSVFormat value);
+
+ abstract CSVFormat getCSVFormat();
+
+ abstract Write<T> autoBuild();
+
+ final Write<T> build() {
+ checkArgument(
+ !getCSVFormat().getSkipHeaderRecord(),
+ "withSkipHeaderRecord is an illegal CSVFormat setting");
+
+ if (getCSVFormat().getHeaderComments() != null) {
+ checkArgument(
+ getCSVFormat().isCommentMarkerSet(),
+ "CSVFormat withCommentMarker required when withHeaderComments");
+ }
+ return autoBuild();
+ }
+ }
+
+ @Override
+ public WriteFilesResult<String> expand(PCollection<T> input) {
+ if (!input.hasSchema()) {
+ throw new IllegalArgumentException(
+ String.format(
+ "%s requires an input Schema. Note that only Row or user classes are supported. Consider using TextIO or FileIO directly when writing primitive types",
+ Write.class.getName()));
+ }
+
+ Schema schema = input.getSchema();
+
+ RowCoder rowCoder = RowCoder.of(schema);
+
+ PCollection<Row> rows =
+ input
+ .apply("To Rows", MapElements.into(rows()).via(input.getToRowFunction()))
+ .setCoder(rowCoder);
+
+ CSVFormat csvFormat = applyRequiredCSVFormatSettings(schema);
+
+ String header = formatHeader(csvFormat);
+
+ SerializableFunction<Row, String> toCsvFn =
+ CsvRowConversions.RowToCsv.builder()
+ .setCSVFormat(csvFormat)
+ .setSchema(input.getSchema())
+ .build();
+
+ PCollection<String> csv = rows.apply("To CSV", MapElements.into(strings()).via(toCsvFn));
+
+ return csv.apply("Write CSV", getTextIOWrite().withHeader(header).withOutputFilenames());
+ }
+
+ CSVFormat applyRequiredCSVFormatSettings(Schema schema) {
+ CSVFormat csvFormat = getCSVFormat().withSkipHeaderRecord();
+ if (csvFormat.getHeader() == null) {
+ csvFormat = csvFormat.withHeader(schema.sorted().getFieldNames().toArray(new String[0]));
+ }
+ return csvFormat;
+ }
+
+ private static String formatHeader(CSVFormat csvFormat) {
+ String[] header = requireNonNull(csvFormat.getHeader());
+ List<String> result = new ArrayList<>();
+ if (csvFormat.getHeaderComments() != null) {
+ for (String comment : csvFormat.getHeaderComments()) {
+ result.add(csvFormat.getCommentMarker() + " " + comment);
+ }
+ }
+ CSVFormat withoutHeaderComments = csvFormat.withHeaderComments();
+ result.add(withoutHeaderComments.format((Object[]) header));
+ return String.join("\n", result);
+ }
Review Comment:
In response to this comment and hopefully provide more clarity, I've:
1. Renamed `applyRequiredCSVFormatSettings` to [`buildHeaderFromSchemaIfNeeded`](https://github.com/damondouglas/beam/blob/24469-csv-io-write/sdks/java/io/csv/src/main/java/org/apache/beam/sdk/io/csv/CsvIO.java#L426)
2. Refactored and renamed `formatHeader` to [`writeWithCSVFormatHeaderAndComments`](https://github.com/damondouglas/beam/blob/24469-csv-io-write/sdks/java/io/csv/src/main/java/org/apache/beam/sdk/io/csv/CsvIO.java#L436).
Additionally, I wanted to address the questions/comments below:
1. > We are ignoring the case where the user wants to skip the header record and doesn't specify the header. We should only specify the csvFormat header if withSkipHeaderRecord is false.
CSVFormat's withSkipHeaderRecord drives output of a whole line, not what that line contains. It's also a parameter that gives an output incompatible in a Beam context as CSVFormat was designed for writing a single file. Therefore, do you agree then we should not use withSkipHeaderRecord to drive the header contents?
2. > Note that we must also return null from formatHeader in this case since TextIO#withHeader specifically says that a newline is automatically added after.
Apologies, your comment confused me.
Are you saying that the consequence of TextIO#withHeader is:
```
col1, col2, col3
1,2,3
2,3,4
```
vs
```
col1, col2, col3
1,2,3
2,3,4
```
--
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: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [beam] damondouglas commented on pull request #24630: [24469] Implement CsvIO.Write and supporting classes
Posted by GitBox <gi...@apache.org>.
damondouglas commented on PR #24630:
URL: https://github.com/apache/beam/pull/24630#issuecomment-1385953937
Run Python PreCommit
--
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: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [beam] github-actions[bot] commented on pull request #24630: [24469] Implement CsvIO.Write and supporting classes
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on PR #24630:
URL: https://github.com/apache/beam/pull/24630#issuecomment-1366607693
Reminder, please take a look at this pr: @lukecwik @Abacn @pabloem
--
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: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [beam] damondouglas commented on a diff in pull request #24630: [24469] Implement CsvIO.Write and supporting classes
Posted by "damondouglas (via GitHub)" <gi...@apache.org>.
damondouglas commented on code in PR #24630:
URL: https://github.com/apache/beam/pull/24630#discussion_r1085716045
##########
sdks/java/io/csv/src/main/java/org/apache/beam/sdk/io/csv/CsvIO.java:
##########
@@ -0,0 +1,446 @@
+/*
+ * 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.sdk.io.csv;
+
+import static java.util.Objects.requireNonNull;
+import static org.apache.beam.sdk.values.TypeDescriptors.rows;
+import static org.apache.beam.sdk.values.TypeDescriptors.strings;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+
+import com.google.auto.value.AutoValue;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Set;
+import org.apache.beam.sdk.coders.RowCoder;
+import org.apache.beam.sdk.io.Compression;
+import org.apache.beam.sdk.io.FileBasedSink;
+import org.apache.beam.sdk.io.FileIO;
+import org.apache.beam.sdk.io.ShardNameTemplate;
+import org.apache.beam.sdk.io.TextIO;
+import org.apache.beam.sdk.io.WriteFilesResult;
+import org.apache.beam.sdk.io.fs.ResourceId;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.transforms.MapElements;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.sdk.values.TypeDescriptor;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableSet;
+import org.apache.commons.csv.CSVFormat;
+
+/**
+ * {@link PTransform}s for reading and writing CSV files.
+ *
+ * <h2>Reading CSV files</h2>
+ *
+ * <p>Reading from CSV files is not yet implemented. Please see <a
+ * href="https://github.com/apache/beam/issues/24552">https://github.com/apache/beam/issues/24552</a>.
+ *
+ * <h2>Writing CSV files</h2>
+ *
+ * <p>To write a {@link PCollection} to one or more CSV files, use {@link CsvIO.Write}, using {@link
+ * CsvIO#writeRows(String, CSVFormat)} or {@link CsvIO#write(String, CSVFormat)}. {@link
+ * CsvIO.Write} supports writing {@link Row} or custom Java types using an inferred {@link Schema}.
+ * Examples below show both scenarios. See the Beam Programming Guide on <a
+ * href="https://beam.apache.org/documentation/programming-guide/#inferring-schemas">inferring
+ * schemas</a> for more information on how to enable Beam to infer a {@link Schema} from a custom
+ * Java type.
+ *
+ * <p>{@link CsvIO.Write} only supports writing {@link Schema} aware types that do not contain any
+ * nested {@link FieldType}s such a {@link org.apache.beam.sdk.schemas.Schema.TypeName#ROW} or
+ * repeated {@link org.apache.beam.sdk.schemas.Schema.TypeName#ARRAY} types. See {@link
+ * CsvIO.Write#VALID_FIELD_TYPE_SET} for valid {@link FieldType}s.
+ *
+ * <h3>Example usage:</h3>
+ *
+ * <p>Suppose we have the following <code>Transaction</code> class annotated with
+ * {@code @DefaultSchema(JavaBeanSchema.class)} so that Beam can infer its {@link Schema}:
+ *
+ * <pre>{@code @DefaultSchema(JavaBeanSchema.class)
+ * public class Transaction {
+ * public Transaction() { … }
+ * public Long getTransactionId();
+ * public void setTransactionId(Long transactionId) { … }
+ * public String getBank() { … }
+ * public void setBank(String bank) { … }
+ * public double getPurchaseAmount() { … }
+ * public void setPurchaseAmount(double purchaseAmount) { … }
+ * }
+ * }</pre>
+ *
+ * <p>From a {@code PCollection<Transaction>}, {@link CsvIO.Write} can write one or many CSV files
+ * automatically creating the header based on its inferred {@link Schema}.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions = ...
+ * transactions.apply(CsvIO.<Transaction>write("path/to/folder/prefix", CSVFormat.DEFAULT));
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header is repeated for every
+ * shard file, whereas by default, {@link CsvIO.Write} will write all fields in <b>sorted order</b>
+ * of the field names.
+ *
+ * <pre>{@code
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ *
+ * <p>To control the order and subset of fields that {@link CsvIO.Write} writes, use {@link
+ * CSVFormat#withHeader(String...)}. Note, however, the following constraints:
+ *
+ * <ol>
+ * <li>Each header column must match a field name in the {@link Schema}; matching is case
+ * sensitive.
+ * <li>Matching header columns must match {@link Schema} fields that are valid {@link FieldType}s;
+ * see {@link #VALID_FIELD_TYPE_SET}.
+ * <li>{@link CSVFormat} only allows repeated header columns when {@link
+ * CSVFormat#withAllowDuplicateHeaderNames()}
+ * </ol>
+ *
+ * <p>The following example shows the use of {@link CSVFormat#withHeader(String...)} to control the
+ * order and subset of <code>Transaction</code> fields.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions ...
+ * transactions.apply(
+ * CsvIO
+ * .<Transaction>write("path/to/folder/prefix", CSVFormat.DEFAULT.withHeader("transactionId", "purchaseAmount"))
+ * );
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header is repeated for every
+ * shard file, but only including the subset of fields in their listed order.
+ *
+ * <pre>{@code
+ * transactionId,purchaseAmount
+ * 12345,10.23
+ * 54321,54.65
+ * 98765,11.76
+ * }</pre>
+ *
+ * <p>In addition to header customization, {@link CsvIO.Write} supports {@link
+ * CSVFormat#withHeaderComments(Object...)} as shown below. Note that {@link
+ * CSVFormat#withCommentMarker(char)} is required when specifying header comments.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions = ...
+ * transactions
+ * .apply(
+ * CsvIO.<Transaction>write("path/to/folder/prefix",
+ * CSVFormat.DEFAULT.withCommentMarker('#').withHeaderComments("Bank Report", "1970-01-01", "Operator: John Doe")
+ * );
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header and header comments are
+ * repeated for every shard file.
+ *
+ * <pre>{@code
+ * # Bank Report
+ * # 1970-01-01
+ * # Operator: John Doe
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ *
+ * <p>A {@link PCollection} of {@link Row}s works just like custom Java types illustrated above,
+ * except we use {@link CsvIO#writeRows(String, CSVFormat)} as shown below for the same {@code
+ * Transaction} class. We derive {@code Transaction}'s {@link Schema} using a {@link
+ * org.apache.beam.sdk.schemas.annotations.DefaultSchema.DefaultSchemaProvider}. Note that
+ * hard-coding the {@link Row}s below is for illustration purposes. Developers are instead
+ * encouraged to take advantage of {@link
+ * org.apache.beam.sdk.schemas.annotations.DefaultSchema.DefaultSchemaProvider#toRowFunction(TypeDescriptor)}.
+ *
+ * <pre>{@code
+ * DefaultSchemaProvider defaultSchemaProvider = new DefaultSchemaProvider();
+ * Schema schema = defaultSchemaProvider.schemaFor(TypeDescriptor.of(Transaction.class));
+ * PCollection<Row> transactions = pipeline.apply(Create.of(
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "A")
+ * .withFieldValue("purchaseAmount", 10.23)
+ * .withFieldValue("transactionId", "12345")
+ * .build(),
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "B")
+ * .withFieldValue("purchaseAmount", 54.65)
+ * .withFieldValue("transactionId", "54321")
+ * .build(),
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "C")
+ * .withFieldValue("purchaseAmount", 11.76)
+ * .withFieldValue("transactionId", "98765")
+ * .build()
+ * );
+ *
+ * transactions.apply(
+ * CsvIO
+ * .writeRowsTo("gs://bucket/path/to/folder/prefix", CSVFormat.DEFAULT)
+ * );
+ * }</pre>
+ *
+ * <p>Writing the transactions {@link PCollection} of {@link Row}s would yield the following CSV
+ * file content.
+ *
+ * <pre>{@code
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ *
+ * {@link CsvIO.Write} does not support the following {@link CSVFormat} properties and will throw an
+ * {@link IllegalArgumentException}.
+ *
+ * <ul>
+ * <li>{@link CSVFormat#withSkipHeaderRecord()}
+ * <li>{@link CSVFormat#withAllowMissingColumnNames()} ()}
+ * <li>{@link CSVFormat#withAutoFlush(boolean)}
+ * <li>{@link CSVFormat#withIgnoreHeaderCase()}
+ * <li>{@link CSVFormat#withIgnoreSurroundingSpaces()}
+ * </ul>
+ */
+public class CsvIO {
+ /**
+ * The valid {@link Schema.FieldType} from which {@link CsvIO} converts CSV records to the fields.
+ *
+ * <ul>
+ * <li>{@link FieldType#BYTE}
+ * <li>{@link FieldType#BOOLEAN}
+ * <li>{@link FieldType#DATETIME}
+ * <li>{@link FieldType#DECIMAL}
+ * <li>{@link FieldType#DOUBLE}
+ * <li>{@link FieldType#INT16}
+ * <li>{@link FieldType#INT32}
+ * <li>{@link FieldType#INT64}
+ * <li>{@link FieldType#FLOAT}
+ * <li>{@link FieldType#STRING}
+ * </ul>
+ */
+ public static final Set<Schema.FieldType> VALID_FIELD_TYPE_SET =
+ ImmutableSet.of(
+ FieldType.BYTE,
+ FieldType.BOOLEAN,
+ FieldType.DATETIME,
+ FieldType.DECIMAL,
+ FieldType.DOUBLE,
+ FieldType.INT16,
+ FieldType.INT32,
+ FieldType.INT64,
+ FieldType.FLOAT,
+ FieldType.STRING);
+
+ static final String DEFAULT_FILENAME_SUFFIX = ".csv";
+
+ /** Instantiates a {@link Write} for writing user types in {@link CSVFormat} format. */
+ public static <T> Write<T> write(String to, CSVFormat csvFormat) {
+ return new AutoValue_CsvIO_Write.Builder<T>()
+ .setTextIOWrite(createDefaultTextIOWrite(to))
+ .setCSVFormat(csvFormat)
+ .build();
+ }
+
+ /** Instantiates a {@link Write} for writing {@link Row}s in {@link CSVFormat} format. */
+ public static Write<Row> writeRows(String to, CSVFormat csvFormat) {
+ return new AutoValue_CsvIO_Write.Builder<Row>()
+ .setTextIOWrite(createDefaultTextIOWrite(to))
+ .setCSVFormat(csvFormat)
+ .build();
+ }
+
+ /** {@link PTransform} for writing CSV files. */
+ @AutoValue
+ public abstract static class Write<T>
+ extends PTransform<PCollection<T>, WriteFilesResult<String>> {
+
+ /** Specifies the {@link Compression} of all generated shard files. */
+ public Write<T> withCompression(Compression compression) {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withCompression(compression)).build();
+ }
+
+ public Write<T> withNoSpilling() {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withNoSpilling()).build();
+ }
+
+ /**
+ * Specifies to use a given fixed number of shards per window. See {@link
+ * TextIO.Write#withNumShards(int)}.
+ */
+ public Write<T> withNumShards(Integer numShards) {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withNumShards(numShards)).build();
+ }
+
+ /**
+ * Forces a single file as output and empty shard name template. See {@link
+ * TextIO.Write#withoutSharding()}.
+ */
+ public Write<T> withoutSharding() {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withoutSharding()).build();
+ }
+
+ /**
+ * Uses the given {@link ShardNameTemplate} for naming output files. See {@link
+ * TextIO.Write#withShardNameTemplate(String)}.
+ */
+ public Write<T> withShardTemplate(String shardTemplate) {
+ return toBuilder()
+ .setTextIOWrite(getTextIOWrite().withShardNameTemplate(shardTemplate))
+ .build();
+ }
+
+ /**
+ * Configures the filename suffix for written files. See {@link
+ * TextIO.Write#withSuffix(String)}.
+ */
+ public Write<T> withSuffix(String suffix) {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withSuffix(suffix)).build();
+ }
+
+ /**
+ * Set the base directory used to generate temporary files. See {@link
+ * TextIO.Write#withTempDirectory(ResourceId)}.
+ */
+ public Write<T> withTempDirectory(ResourceId tempDirectory) {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withTempDirectory(tempDirectory)).build();
+ }
+
+ /**
+ * Preserves windowing of input elements and writes them to files based on the element's window.
+ * See {@link TextIO.Write#withWindowedWrites()}.
+ */
+ public Write<T> withWindowedWrites() {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withWindowedWrites()).build();
+ }
+
+ /**
+ * Returns a transform for writing to text files like this one but that has the given {@link
+ * FileBasedSink.WritableByteChannelFactory} to be used by the {@link FileBasedSink} during
+ * output. See {@link
+ * TextIO.Write#withWritableByteChannelFactory(FileBasedSink.WritableByteChannelFactory)}.
+ */
+ public Write<T> withWritableByteChannelFactory(
+ FileBasedSink.WritableByteChannelFactory writableByteChannelFactory) {
+ return toBuilder()
+ .setTextIOWrite(
+ getTextIOWrite().withWritableByteChannelFactory(writableByteChannelFactory))
+ .build();
+ }
+
+ /** The underlying {@link FileIO.Write} that writes converted input to CSV formatted output. */
+ abstract TextIO.Write getTextIOWrite();
+
+ /** The {@link CSVFormat} to inform headers, header comments, and format CSV row content. */
+ abstract CSVFormat getCSVFormat();
+
+ abstract Builder<T> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<T> {
+
+ /**
+ * The underlying {@link FileIO.Write} that writes converted input to CSV formatted output.
+ */
+ abstract Builder<T> setTextIOWrite(TextIO.Write value);
+
+ /** The {@link CSVFormat} to convert input. Defaults to {@link CSVFormat#DEFAULT}. */
+ abstract Builder<T> setCSVFormat(CSVFormat value);
+
+ abstract CSVFormat getCSVFormat();
+
+ abstract Write<T> autoBuild();
+
+ final Write<T> build() {
+ checkArgument(
+ !getCSVFormat().getSkipHeaderRecord(),
+ "withSkipHeaderRecord is an illegal CSVFormat setting");
+
+ if (getCSVFormat().getHeaderComments() != null) {
+ checkArgument(
+ getCSVFormat().isCommentMarkerSet(),
+ "CSVFormat withCommentMarker required when withHeaderComments");
+ }
+ return autoBuild();
+ }
+ }
+
+ @Override
+ public WriteFilesResult<String> expand(PCollection<T> input) {
+ if (!input.hasSchema()) {
+ throw new IllegalArgumentException(
+ String.format(
+ "%s requires an input Schema. Note that only Row or user classes are supported. Consider using TextIO or FileIO directly when writing primitive types",
+ Write.class.getName()));
+ }
+
+ Schema schema = input.getSchema();
+
+ RowCoder rowCoder = RowCoder.of(schema);
+
+ PCollection<Row> rows =
+ input
+ .apply("To Rows", MapElements.into(rows()).via(input.getToRowFunction()))
+ .setCoder(rowCoder);
+
+ CSVFormat csvFormat = applyRequiredCSVFormatSettings(schema);
+
+ String header = formatHeader(csvFormat);
+
+ SerializableFunction<Row, String> toCsvFn =
+ CsvRowConversions.RowToCsv.builder()
+ .setCSVFormat(csvFormat)
+ .setSchema(input.getSchema())
+ .build();
+
+ PCollection<String> csv = rows.apply("To CSV", MapElements.into(strings()).via(toCsvFn));
+
+ return csv.apply("Write CSV", getTextIOWrite().withHeader(header).withOutputFilenames());
+ }
+
+ CSVFormat applyRequiredCSVFormatSettings(Schema schema) {
+ CSVFormat csvFormat = getCSVFormat().withSkipHeaderRecord();
+ if (csvFormat.getHeader() == null) {
+ csvFormat = csvFormat.withHeader(schema.sorted().getFieldNames().toArray(new String[0]));
+ }
+ return csvFormat;
+ }
+
+ private static String formatHeader(CSVFormat csvFormat) {
+ String[] header = requireNonNull(csvFormat.getHeader());
+ List<String> result = new ArrayList<>();
+ if (csvFormat.getHeaderComments() != null) {
+ for (String comment : csvFormat.getHeaderComments()) {
+ result.add(csvFormat.getCommentMarker() + " " + comment);
+ }
+ }
+ CSVFormat withoutHeaderComments = csvFormat.withHeaderComments();
+ result.add(withoutHeaderComments.format((Object[]) header));
+ return String.join("\n", result);
+ }
Review Comment:
In response to this comment, I've:
1. Renamed `applyRequiredCSVFormatSettings` to [`buildHeaderFromSchemaIfNeeded`](https://github.com/damondouglas/beam/blob/24469-csv-io-write/sdks/java/io/csv/src/main/java/org/apache/beam/sdk/io/csv/CsvIO.java#L426)
2. Refactored and renamed `formatHeader` to `writeWithCSVFormatHeaderAndComments` to hopefully provide clarity coupling this behavior with TextIO#withHeader.
Additionally, I wanted to address the questions/comments below:
1. > We are ignoring the case where the user wants to skip the header record and doesn't specify the header. We should only specify the csvFormat header if withSkipHeaderRecord is false.
CSVFormat communicates that withSkipHeaderRecord functions to not output the header. It's also a parameter that gives an output incompatible in a beam context as CSVFormat was designed for writing a single file. Therefore, do you agree then we should not use withSkipHeaderRecord to drive how we are building the header?
3. > Note that we must also return null from formatHeader in this case since TextIO#withHeader specifically says that a newline is automatically added after.
Apologies, your comment confused me.
Are you saying that the consequence of TextIO#withHeader is:
```
col1, col2, col3
1,2,3
2,3,4
```
vs
```
col1, col2, col3
1,2,3
2,3,4
```
--
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: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [beam] damondouglas commented on a diff in pull request #24630: [24469] Implement CsvIO.Write and supporting classes
Posted by "damondouglas (via GitHub)" <gi...@apache.org>.
damondouglas commented on code in PR #24630:
URL: https://github.com/apache/beam/pull/24630#discussion_r1085768571
##########
sdks/java/io/csv/src/main/java/org/apache/beam/sdk/io/csv/CsvIO.java:
##########
@@ -0,0 +1,446 @@
+/*
+ * 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.sdk.io.csv;
+
+import static java.util.Objects.requireNonNull;
+import static org.apache.beam.sdk.values.TypeDescriptors.rows;
+import static org.apache.beam.sdk.values.TypeDescriptors.strings;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+
+import com.google.auto.value.AutoValue;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Set;
+import org.apache.beam.sdk.coders.RowCoder;
+import org.apache.beam.sdk.io.Compression;
+import org.apache.beam.sdk.io.FileBasedSink;
+import org.apache.beam.sdk.io.FileIO;
+import org.apache.beam.sdk.io.ShardNameTemplate;
+import org.apache.beam.sdk.io.TextIO;
+import org.apache.beam.sdk.io.WriteFilesResult;
+import org.apache.beam.sdk.io.fs.ResourceId;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.transforms.MapElements;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.sdk.values.TypeDescriptor;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableSet;
+import org.apache.commons.csv.CSVFormat;
+
+/**
+ * {@link PTransform}s for reading and writing CSV files.
+ *
+ * <h2>Reading CSV files</h2>
+ *
+ * <p>Reading from CSV files is not yet implemented. Please see <a
+ * href="https://github.com/apache/beam/issues/24552">https://github.com/apache/beam/issues/24552</a>.
+ *
+ * <h2>Writing CSV files</h2>
+ *
+ * <p>To write a {@link PCollection} to one or more CSV files, use {@link CsvIO.Write}, using {@link
+ * CsvIO#writeRows(String, CSVFormat)} or {@link CsvIO#write(String, CSVFormat)}. {@link
+ * CsvIO.Write} supports writing {@link Row} or custom Java types using an inferred {@link Schema}.
+ * Examples below show both scenarios. See the Beam Programming Guide on <a
+ * href="https://beam.apache.org/documentation/programming-guide/#inferring-schemas">inferring
+ * schemas</a> for more information on how to enable Beam to infer a {@link Schema} from a custom
+ * Java type.
+ *
+ * <p>{@link CsvIO.Write} only supports writing {@link Schema} aware types that do not contain any
+ * nested {@link FieldType}s such a {@link org.apache.beam.sdk.schemas.Schema.TypeName#ROW} or
+ * repeated {@link org.apache.beam.sdk.schemas.Schema.TypeName#ARRAY} types. See {@link
+ * CsvIO.Write#VALID_FIELD_TYPE_SET} for valid {@link FieldType}s.
+ *
+ * <h3>Example usage:</h3>
+ *
+ * <p>Suppose we have the following <code>Transaction</code> class annotated with
+ * {@code @DefaultSchema(JavaBeanSchema.class)} so that Beam can infer its {@link Schema}:
+ *
+ * <pre>{@code @DefaultSchema(JavaBeanSchema.class)
+ * public class Transaction {
+ * public Transaction() { … }
+ * public Long getTransactionId();
+ * public void setTransactionId(Long transactionId) { … }
+ * public String getBank() { … }
+ * public void setBank(String bank) { … }
+ * public double getPurchaseAmount() { … }
+ * public void setPurchaseAmount(double purchaseAmount) { … }
+ * }
+ * }</pre>
+ *
+ * <p>From a {@code PCollection<Transaction>}, {@link CsvIO.Write} can write one or many CSV files
+ * automatically creating the header based on its inferred {@link Schema}.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions = ...
+ * transactions.apply(CsvIO.<Transaction>write("path/to/folder/prefix", CSVFormat.DEFAULT));
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header is repeated for every
+ * shard file, whereas by default, {@link CsvIO.Write} will write all fields in <b>sorted order</b>
+ * of the field names.
+ *
+ * <pre>{@code
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ *
+ * <p>To control the order and subset of fields that {@link CsvIO.Write} writes, use {@link
+ * CSVFormat#withHeader(String...)}. Note, however, the following constraints:
+ *
+ * <ol>
+ * <li>Each header column must match a field name in the {@link Schema}; matching is case
+ * sensitive.
+ * <li>Matching header columns must match {@link Schema} fields that are valid {@link FieldType}s;
+ * see {@link #VALID_FIELD_TYPE_SET}.
+ * <li>{@link CSVFormat} only allows repeated header columns when {@link
+ * CSVFormat#withAllowDuplicateHeaderNames()}
+ * </ol>
+ *
+ * <p>The following example shows the use of {@link CSVFormat#withHeader(String...)} to control the
+ * order and subset of <code>Transaction</code> fields.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions ...
+ * transactions.apply(
+ * CsvIO
+ * .<Transaction>write("path/to/folder/prefix", CSVFormat.DEFAULT.withHeader("transactionId", "purchaseAmount"))
+ * );
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header is repeated for every
+ * shard file, but only including the subset of fields in their listed order.
+ *
+ * <pre>{@code
+ * transactionId,purchaseAmount
+ * 12345,10.23
+ * 54321,54.65
+ * 98765,11.76
+ * }</pre>
+ *
+ * <p>In addition to header customization, {@link CsvIO.Write} supports {@link
+ * CSVFormat#withHeaderComments(Object...)} as shown below. Note that {@link
+ * CSVFormat#withCommentMarker(char)} is required when specifying header comments.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions = ...
+ * transactions
+ * .apply(
+ * CsvIO.<Transaction>write("path/to/folder/prefix",
+ * CSVFormat.DEFAULT.withCommentMarker('#').withHeaderComments("Bank Report", "1970-01-01", "Operator: John Doe")
+ * );
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header and header comments are
+ * repeated for every shard file.
+ *
+ * <pre>{@code
+ * # Bank Report
+ * # 1970-01-01
+ * # Operator: John Doe
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ *
+ * <p>A {@link PCollection} of {@link Row}s works just like custom Java types illustrated above,
+ * except we use {@link CsvIO#writeRows(String, CSVFormat)} as shown below for the same {@code
+ * Transaction} class. We derive {@code Transaction}'s {@link Schema} using a {@link
+ * org.apache.beam.sdk.schemas.annotations.DefaultSchema.DefaultSchemaProvider}. Note that
+ * hard-coding the {@link Row}s below is for illustration purposes. Developers are instead
+ * encouraged to take advantage of {@link
+ * org.apache.beam.sdk.schemas.annotations.DefaultSchema.DefaultSchemaProvider#toRowFunction(TypeDescriptor)}.
+ *
+ * <pre>{@code
+ * DefaultSchemaProvider defaultSchemaProvider = new DefaultSchemaProvider();
+ * Schema schema = defaultSchemaProvider.schemaFor(TypeDescriptor.of(Transaction.class));
+ * PCollection<Row> transactions = pipeline.apply(Create.of(
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "A")
+ * .withFieldValue("purchaseAmount", 10.23)
+ * .withFieldValue("transactionId", "12345")
+ * .build(),
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "B")
+ * .withFieldValue("purchaseAmount", 54.65)
+ * .withFieldValue("transactionId", "54321")
+ * .build(),
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "C")
+ * .withFieldValue("purchaseAmount", 11.76)
+ * .withFieldValue("transactionId", "98765")
+ * .build()
+ * );
+ *
+ * transactions.apply(
+ * CsvIO
+ * .writeRowsTo("gs://bucket/path/to/folder/prefix", CSVFormat.DEFAULT)
+ * );
+ * }</pre>
+ *
+ * <p>Writing the transactions {@link PCollection} of {@link Row}s would yield the following CSV
+ * file content.
+ *
+ * <pre>{@code
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ *
+ * {@link CsvIO.Write} does not support the following {@link CSVFormat} properties and will throw an
+ * {@link IllegalArgumentException}.
+ *
+ * <ul>
+ * <li>{@link CSVFormat#withSkipHeaderRecord()}
+ * <li>{@link CSVFormat#withAllowMissingColumnNames()} ()}
+ * <li>{@link CSVFormat#withAutoFlush(boolean)}
+ * <li>{@link CSVFormat#withIgnoreHeaderCase()}
+ * <li>{@link CSVFormat#withIgnoreSurroundingSpaces()}
+ * </ul>
+ */
+public class CsvIO {
+ /**
+ * The valid {@link Schema.FieldType} from which {@link CsvIO} converts CSV records to the fields.
+ *
+ * <ul>
+ * <li>{@link FieldType#BYTE}
+ * <li>{@link FieldType#BOOLEAN}
+ * <li>{@link FieldType#DATETIME}
+ * <li>{@link FieldType#DECIMAL}
+ * <li>{@link FieldType#DOUBLE}
+ * <li>{@link FieldType#INT16}
+ * <li>{@link FieldType#INT32}
+ * <li>{@link FieldType#INT64}
+ * <li>{@link FieldType#FLOAT}
+ * <li>{@link FieldType#STRING}
+ * </ul>
+ */
+ public static final Set<Schema.FieldType> VALID_FIELD_TYPE_SET =
+ ImmutableSet.of(
+ FieldType.BYTE,
+ FieldType.BOOLEAN,
+ FieldType.DATETIME,
+ FieldType.DECIMAL,
+ FieldType.DOUBLE,
+ FieldType.INT16,
+ FieldType.INT32,
+ FieldType.INT64,
+ FieldType.FLOAT,
+ FieldType.STRING);
+
+ static final String DEFAULT_FILENAME_SUFFIX = ".csv";
+
+ /** Instantiates a {@link Write} for writing user types in {@link CSVFormat} format. */
+ public static <T> Write<T> write(String to, CSVFormat csvFormat) {
+ return new AutoValue_CsvIO_Write.Builder<T>()
+ .setTextIOWrite(createDefaultTextIOWrite(to))
+ .setCSVFormat(csvFormat)
+ .build();
+ }
+
+ /** Instantiates a {@link Write} for writing {@link Row}s in {@link CSVFormat} format. */
+ public static Write<Row> writeRows(String to, CSVFormat csvFormat) {
+ return new AutoValue_CsvIO_Write.Builder<Row>()
+ .setTextIOWrite(createDefaultTextIOWrite(to))
+ .setCSVFormat(csvFormat)
+ .build();
+ }
+
+ /** {@link PTransform} for writing CSV files. */
+ @AutoValue
+ public abstract static class Write<T>
+ extends PTransform<PCollection<T>, WriteFilesResult<String>> {
+
+ /** Specifies the {@link Compression} of all generated shard files. */
+ public Write<T> withCompression(Compression compression) {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withCompression(compression)).build();
+ }
+
+ public Write<T> withNoSpilling() {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withNoSpilling()).build();
+ }
+
+ /**
+ * Specifies to use a given fixed number of shards per window. See {@link
+ * TextIO.Write#withNumShards(int)}.
+ */
+ public Write<T> withNumShards(Integer numShards) {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withNumShards(numShards)).build();
+ }
+
+ /**
+ * Forces a single file as output and empty shard name template. See {@link
+ * TextIO.Write#withoutSharding()}.
+ */
+ public Write<T> withoutSharding() {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withoutSharding()).build();
+ }
+
+ /**
+ * Uses the given {@link ShardNameTemplate} for naming output files. See {@link
+ * TextIO.Write#withShardNameTemplate(String)}.
+ */
+ public Write<T> withShardTemplate(String shardTemplate) {
+ return toBuilder()
+ .setTextIOWrite(getTextIOWrite().withShardNameTemplate(shardTemplate))
+ .build();
+ }
+
+ /**
+ * Configures the filename suffix for written files. See {@link
+ * TextIO.Write#withSuffix(String)}.
+ */
+ public Write<T> withSuffix(String suffix) {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withSuffix(suffix)).build();
+ }
+
+ /**
+ * Set the base directory used to generate temporary files. See {@link
+ * TextIO.Write#withTempDirectory(ResourceId)}.
+ */
+ public Write<T> withTempDirectory(ResourceId tempDirectory) {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withTempDirectory(tempDirectory)).build();
+ }
+
+ /**
+ * Preserves windowing of input elements and writes them to files based on the element's window.
+ * See {@link TextIO.Write#withWindowedWrites()}.
+ */
+ public Write<T> withWindowedWrites() {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withWindowedWrites()).build();
+ }
+
+ /**
+ * Returns a transform for writing to text files like this one but that has the given {@link
+ * FileBasedSink.WritableByteChannelFactory} to be used by the {@link FileBasedSink} during
+ * output. See {@link
+ * TextIO.Write#withWritableByteChannelFactory(FileBasedSink.WritableByteChannelFactory)}.
+ */
+ public Write<T> withWritableByteChannelFactory(
+ FileBasedSink.WritableByteChannelFactory writableByteChannelFactory) {
+ return toBuilder()
+ .setTextIOWrite(
+ getTextIOWrite().withWritableByteChannelFactory(writableByteChannelFactory))
+ .build();
+ }
+
+ /** The underlying {@link FileIO.Write} that writes converted input to CSV formatted output. */
+ abstract TextIO.Write getTextIOWrite();
+
+ /** The {@link CSVFormat} to inform headers, header comments, and format CSV row content. */
+ abstract CSVFormat getCSVFormat();
+
+ abstract Builder<T> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<T> {
+
+ /**
+ * The underlying {@link FileIO.Write} that writes converted input to CSV formatted output.
+ */
+ abstract Builder<T> setTextIOWrite(TextIO.Write value);
+
+ /** The {@link CSVFormat} to convert input. Defaults to {@link CSVFormat#DEFAULT}. */
+ abstract Builder<T> setCSVFormat(CSVFormat value);
+
+ abstract CSVFormat getCSVFormat();
+
+ abstract Write<T> autoBuild();
+
+ final Write<T> build() {
+ checkArgument(
+ !getCSVFormat().getSkipHeaderRecord(),
+ "withSkipHeaderRecord is an illegal CSVFormat setting");
+
+ if (getCSVFormat().getHeaderComments() != null) {
+ checkArgument(
+ getCSVFormat().isCommentMarkerSet(),
+ "CSVFormat withCommentMarker required when withHeaderComments");
+ }
+ return autoBuild();
+ }
+ }
+
+ @Override
+ public WriteFilesResult<String> expand(PCollection<T> input) {
+ if (!input.hasSchema()) {
+ throw new IllegalArgumentException(
+ String.format(
+ "%s requires an input Schema. Note that only Row or user classes are supported. Consider using TextIO or FileIO directly when writing primitive types",
+ Write.class.getName()));
+ }
+
+ Schema schema = input.getSchema();
+
+ RowCoder rowCoder = RowCoder.of(schema);
+
+ PCollection<Row> rows =
+ input
+ .apply("To Rows", MapElements.into(rows()).via(input.getToRowFunction()))
+ .setCoder(rowCoder);
+
+ CSVFormat csvFormat = applyRequiredCSVFormatSettings(schema);
+
+ String header = formatHeader(csvFormat);
+
+ SerializableFunction<Row, String> toCsvFn =
+ CsvRowConversions.RowToCsv.builder()
+ .setCSVFormat(csvFormat)
+ .setSchema(input.getSchema())
+ .build();
+
+ PCollection<String> csv = rows.apply("To CSV", MapElements.into(strings()).via(toCsvFn));
+
+ return csv.apply("Write CSV", getTextIOWrite().withHeader(header).withOutputFilenames());
+ }
+
+ CSVFormat applyRequiredCSVFormatSettings(Schema schema) {
+ CSVFormat csvFormat = getCSVFormat().withSkipHeaderRecord();
+ if (csvFormat.getHeader() == null) {
+ csvFormat = csvFormat.withHeader(schema.sorted().getFieldNames().toArray(new String[0]));
+ }
+ return csvFormat;
+ }
+
+ private static String formatHeader(CSVFormat csvFormat) {
+ String[] header = requireNonNull(csvFormat.getHeader());
+ List<String> result = new ArrayList<>();
+ if (csvFormat.getHeaderComments() != null) {
+ for (String comment : csvFormat.getHeaderComments()) {
+ result.add(csvFormat.getCommentMarker() + " " + comment);
+ }
+ }
+ CSVFormat withoutHeaderComments = csvFormat.withHeaderComments();
+ result.add(withoutHeaderComments.format((Object[]) header));
+ return String.join("\n", result);
+ }
Review Comment:
I would like to anchor our discussion on CSV header-related software behavior using the following:
```gherkin
Feature: write a CSV header
Scenario: User provides CSVFormat.withHeader
Given CSVFormat is not null
When CsvIO writes PCollection<Row> to CSV formatted file shards
Then for each shard the CSVFormat withHeader appears as the first line
Scenario: User does not provide CSVFormat.withHeader
Given CSVFormat is null
When CsvIO writes PCollection<Row> to CSV formatted file shards
Then for each shard the Schema field names in sorted order appears as the first line
Scenario: User provides CSVFormat.withSkipHeaderRecord == true
Given CSVFormat.withSkipHeaderRecord is true
When the user applies the CsvIO.Write PTransform
Then they see an error "withSkipHeaderRecord is an illegal CSVFormat setting; CsvIO.Write always writes the header"
```
--
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: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [beam] damondouglas commented on a diff in pull request #24630: [24469] Implement CsvIO.Write and supporting classes
Posted by "damondouglas (via GitHub)" <gi...@apache.org>.
damondouglas commented on code in PR #24630:
URL: https://github.com/apache/beam/pull/24630#discussion_r1085704532
##########
sdks/java/io/csv/src/test/java/org/apache/beam/sdk/io/csv/RowToCsvCSVFormatTest.java:
##########
@@ -0,0 +1,353 @@
+/*
+ * 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.sdk.io.csv;
+
+import static org.apache.beam.sdk.io.csv.CsvIOTestData.DATA;
+import static org.apache.beam.sdk.io.csv.CsvIOTestJavaBeans.ALL_PRIMITIVE_DATA_TYPES_SCHEMA;
+import static org.apache.beam.sdk.io.csv.CsvIOTestJavaBeans.NULLABLE_ALL_PRIMITIVE_DATA_TYPES_SCHEMA;
+import static org.apache.beam.sdk.io.csv.CsvIOTestJavaBeans.TIME_CONTAINING_SCHEMA;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkNotNull;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertThrows;
+
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.Field;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.values.Row;
+import org.apache.commons.csv.CSVFormat;
+import org.apache.commons.csv.QuoteMode;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/**
+ * Tests {@link org.apache.commons.csv.CSVFormat} settings in the context of {@link
+ * CsvRowConversions.RowToCsv}.
+ */
+@RunWith(JUnit4.class)
+public class RowToCsvCSVFormatTest {
+ @Test
+ public void invalidCSVFormatHeader() {
+ NullPointerException nullHeaderError =
+ assertThrows(
+ NullPointerException.class,
+ () ->
+ CsvRowConversions.RowToCsv.builder()
+ .setCSVFormat(CSVFormat.DEFAULT)
+ .setSchema(ALL_PRIMITIVE_DATA_TYPES_SCHEMA)
+ .build());
+ assertEquals("CSVFormat withHeader is required", nullHeaderError.getMessage());
+
+ IllegalArgumentException emptyHeaderError =
+ assertThrows(
+ IllegalArgumentException.class,
+ () ->
+ CsvRowConversions.RowToCsv.builder()
+ .setCSVFormat(CSVFormat.DEFAULT.withHeader())
+ .setSchema(ALL_PRIMITIVE_DATA_TYPES_SCHEMA)
+ .build());
+
+ assertEquals(
+ "CSVFormat withHeader requires at least one column", emptyHeaderError.getMessage());
+
+ IllegalArgumentException mismatchHeaderError =
+ assertThrows(
+ IllegalArgumentException.class,
+ () ->
+ CsvRowConversions.RowToCsv.builder()
+ .setCSVFormat(
+ CSVFormat.DEFAULT.withHeader("aString", "idontexist1", "idontexist2"))
+ .setSchema(ALL_PRIMITIVE_DATA_TYPES_SCHEMA)
+ .build());
+
+ assertEquals(
+ "columns in CSVFormat header do not exist in Schema: idontexist2,idontexist1",
+ mismatchHeaderError.getMessage());
+ }
+
+ @Test
+ public void invalidSchema() {
+ IllegalArgumentException emptySchemaError =
+ assertThrows(
+ IllegalArgumentException.class,
+ () ->
+ CsvRowConversions.RowToCsv.builder()
+ .setCSVFormat(CSVFormat.DEFAULT.withHeader())
+ .setSchema(Schema.of())
+ .build());
+
+ assertEquals("Schema has no fields", emptySchemaError.getMessage());
+
+ IllegalArgumentException invalidArrayFieldsError =
+ assertThrows(
+ IllegalArgumentException.class,
+ () ->
+ CsvRowConversions.RowToCsv.builder()
+ .setCSVFormat(CSVFormat.DEFAULT.withHeader("instant", "instantList"))
+ .setSchema(TIME_CONTAINING_SCHEMA)
+ .build());
+
+ assertEquals(
+ "columns in header match fields in Schema with invalid types: instantList. See CsvIO#VALID_FIELD_TYPE_SET for a list of valid field types.",
+ invalidArrayFieldsError.getMessage());
+
+ // Should not throw Exception when limited to valid fields.
+ CsvRowConversions.RowToCsv.builder()
+ .setCSVFormat(CSVFormat.DEFAULT.withHeader("instant"))
+ .setSchema(TIME_CONTAINING_SCHEMA)
+ .build();
+ }
+
+ @Test
+ public void withAllowDuplicateHeaderNamesDuplicatesRowFieldOutput() {
+ assertEquals(
+ "allowDuplicateHeaderNames=true",
+ "a,a,a",
+ rowToCsv(
+ DATA.allPrimitiveDataTypesRow,
+ csvFormat(ALL_PRIMITIVE_DATA_TYPES_SCHEMA)
+ .withAllowDuplicateHeaderNames(true)
+ .withHeader("aString", "aString", "aString")));
+ }
+
+ @Test
+ public void withAllowMissingColumnNamesSettingThrowsException() {
+ IllegalArgumentException exception =
+ assertThrows(
+ "allowMissingColumnNames=true",
+ IllegalArgumentException.class,
+ () ->
+ rowToCsv(
+ DATA.allPrimitiveDataTypesRow,
+ csvFormat(ALL_PRIMITIVE_DATA_TYPES_SCHEMA).withAllowMissingColumnNames(true)));
+
+ assertEquals(
+ "withAllowMissingColumnNames is an illegal CSVFormat setting", exception.getMessage());
+ }
+
+ @Test
+ public void withAutoFlushThrowsException() {
+ IllegalArgumentException exception =
+ assertThrows(
+ "autoFlush=true",
+ IllegalArgumentException.class,
+ () ->
+ rowToCsv(
+ DATA.allPrimitiveDataTypesRow,
+ csvFormat(ALL_PRIMITIVE_DATA_TYPES_SCHEMA).withAutoFlush(true)));
+
+ assertEquals("withAutoFlush is an illegal CSVFormat setting", exception.getMessage());
+ }
+
+ @Test
+ public void withCommentMarkerDoesNotEffectConversion() {
+ Schema schema = Schema.of(Field.of("aString", FieldType.STRING));
+ Row row = Row.withSchema(schema).attachValues("$abc");
+ assertEquals("$abc", rowToCsv(row, csvFormat(schema).withCommentMarker('$')));
+ assertEquals("$abc", rowToCsv(row, csvFormat(schema).withCommentMarker(null)));
+ }
+
+ @Test
+ public void withDelimiterDrivesCellBorders() {
+ assertEquals(
+ "false~1~10~1.0~1.0~1~1~a~1",
+ rowToCsv(
+ DATA.allPrimitiveDataTypesRow,
+ csvFormat(ALL_PRIMITIVE_DATA_TYPES_SCHEMA).withDelimiter('~')));
+ }
+
+ @Test
+ public void withEscapeDrivesOutput() {
+ Schema schema =
+ Schema.of(Field.of("aString", FieldType.STRING), Field.of("anInt", FieldType.INT32));
+ Row row = Row.withSchema(schema).attachValues(",a", 1);
+ String[] header = new String[] {"anInt", "aString"};
+ assertEquals(
+ "1,#,a",
+ rowToCsv(
+ row,
+ csvFormat(ALL_PRIMITIVE_DATA_TYPES_SCHEMA)
+ .withHeader(header)
+ .withEscape('#')
+ .withQuoteMode(QuoteMode.NONE)));
+ assertEquals(
+ "1,\",a\"", rowToCsv(row, csvFormat(ALL_PRIMITIVE_DATA_TYPES_SCHEMA).withHeader(header)));
+ }
+
+ @Test
+ public void withHeaderDrivesFieldOrderSubsetOutput() {
+ assertEquals(
+ "1,false,a",
+ rowToCsv(
+ DATA.allPrimitiveDataTypesRow,
+ csvFormat(ALL_PRIMITIVE_DATA_TYPES_SCHEMA)
+ .withHeader("anInteger", "aBoolean", "aString")));
+ }
+
+ @Test
+ public void withHeaderCommentsDoesNotEffectConversion() {
+ assertEquals(
+ "false,1,10,1.0,1.0,1,1,a,1",
+ rowToCsv(
+ DATA.allPrimitiveDataTypesRow,
+ csvFormat(ALL_PRIMITIVE_DATA_TYPES_SCHEMA)
+ .withHeaderComments("some", "header", "comments")));
+ }
+
+ @Test
+ public void withIgnoreEmptyLinesDoesNotEffectOutput() {
+ assertEquals(
+ "false,1,10,1.0,1.0,1,1,a,1",
+ rowToCsv(
+ DATA.allPrimitiveDataTypesRow,
+ csvFormat(ALL_PRIMITIVE_DATA_TYPES_SCHEMA).withIgnoreEmptyLines(true)));
+ assertEquals(
+ "false,1,10,1.0,1.0,1,1,a,1",
+ rowToCsv(
+ DATA.allPrimitiveDataTypesRow,
+ csvFormat(ALL_PRIMITIVE_DATA_TYPES_SCHEMA).withIgnoreEmptyLines(false)));
+ }
+
+ @Test
+ public void withIgnoreHeaderCaseThrowsException() {
+ IllegalArgumentException exception =
+ assertThrows(
+ IllegalArgumentException.class,
+ () ->
+ rowToCsv(
+ DATA.allPrimitiveDataTypesRow,
+ csvFormat(ALL_PRIMITIVE_DATA_TYPES_SCHEMA).withIgnoreHeaderCase(true)));
+ assertEquals("withIgnoreHeaderCase is an illegal CSVFormat setting", exception.getMessage());
+ }
+
+ @Test
+ public void withIgnoreSurroundingSpacesThrowsException() {
+ IllegalArgumentException exception =
+ assertThrows(
+ IllegalArgumentException.class,
+ () ->
+ rowToCsv(
+ DATA.allPrimitiveDataTypesRow,
+ csvFormat(ALL_PRIMITIVE_DATA_TYPES_SCHEMA).withIgnoreSurroundingSpaces(true)));
+ assertEquals(
+ "withIgnoreSurroundingSpaces is an illegal CSVFormat setting", exception.getMessage());
+ }
+
+ @Test
+ public void withNullStringReplacesNullValues() {
+ assertEquals(
+ "🦄,🦄,🦄,🦄,🦄,🦄",
Review Comment:
🦄🦄🦄🦄
--
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: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [beam] damondouglas commented on a diff in pull request #24630: [24469] Implement CsvIO.Write and supporting classes
Posted by "damondouglas (via GitHub)" <gi...@apache.org>.
damondouglas commented on code in PR #24630:
URL: https://github.com/apache/beam/pull/24630#discussion_r1085716045
##########
sdks/java/io/csv/src/main/java/org/apache/beam/sdk/io/csv/CsvIO.java:
##########
@@ -0,0 +1,446 @@
+/*
+ * 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.sdk.io.csv;
+
+import static java.util.Objects.requireNonNull;
+import static org.apache.beam.sdk.values.TypeDescriptors.rows;
+import static org.apache.beam.sdk.values.TypeDescriptors.strings;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+
+import com.google.auto.value.AutoValue;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Set;
+import org.apache.beam.sdk.coders.RowCoder;
+import org.apache.beam.sdk.io.Compression;
+import org.apache.beam.sdk.io.FileBasedSink;
+import org.apache.beam.sdk.io.FileIO;
+import org.apache.beam.sdk.io.ShardNameTemplate;
+import org.apache.beam.sdk.io.TextIO;
+import org.apache.beam.sdk.io.WriteFilesResult;
+import org.apache.beam.sdk.io.fs.ResourceId;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.transforms.MapElements;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.sdk.values.TypeDescriptor;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableSet;
+import org.apache.commons.csv.CSVFormat;
+
+/**
+ * {@link PTransform}s for reading and writing CSV files.
+ *
+ * <h2>Reading CSV files</h2>
+ *
+ * <p>Reading from CSV files is not yet implemented. Please see <a
+ * href="https://github.com/apache/beam/issues/24552">https://github.com/apache/beam/issues/24552</a>.
+ *
+ * <h2>Writing CSV files</h2>
+ *
+ * <p>To write a {@link PCollection} to one or more CSV files, use {@link CsvIO.Write}, using {@link
+ * CsvIO#writeRows(String, CSVFormat)} or {@link CsvIO#write(String, CSVFormat)}. {@link
+ * CsvIO.Write} supports writing {@link Row} or custom Java types using an inferred {@link Schema}.
+ * Examples below show both scenarios. See the Beam Programming Guide on <a
+ * href="https://beam.apache.org/documentation/programming-guide/#inferring-schemas">inferring
+ * schemas</a> for more information on how to enable Beam to infer a {@link Schema} from a custom
+ * Java type.
+ *
+ * <p>{@link CsvIO.Write} only supports writing {@link Schema} aware types that do not contain any
+ * nested {@link FieldType}s such a {@link org.apache.beam.sdk.schemas.Schema.TypeName#ROW} or
+ * repeated {@link org.apache.beam.sdk.schemas.Schema.TypeName#ARRAY} types. See {@link
+ * CsvIO.Write#VALID_FIELD_TYPE_SET} for valid {@link FieldType}s.
+ *
+ * <h3>Example usage:</h3>
+ *
+ * <p>Suppose we have the following <code>Transaction</code> class annotated with
+ * {@code @DefaultSchema(JavaBeanSchema.class)} so that Beam can infer its {@link Schema}:
+ *
+ * <pre>{@code @DefaultSchema(JavaBeanSchema.class)
+ * public class Transaction {
+ * public Transaction() { … }
+ * public Long getTransactionId();
+ * public void setTransactionId(Long transactionId) { … }
+ * public String getBank() { … }
+ * public void setBank(String bank) { … }
+ * public double getPurchaseAmount() { … }
+ * public void setPurchaseAmount(double purchaseAmount) { … }
+ * }
+ * }</pre>
+ *
+ * <p>From a {@code PCollection<Transaction>}, {@link CsvIO.Write} can write one or many CSV files
+ * automatically creating the header based on its inferred {@link Schema}.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions = ...
+ * transactions.apply(CsvIO.<Transaction>write("path/to/folder/prefix", CSVFormat.DEFAULT));
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header is repeated for every
+ * shard file, whereas by default, {@link CsvIO.Write} will write all fields in <b>sorted order</b>
+ * of the field names.
+ *
+ * <pre>{@code
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ *
+ * <p>To control the order and subset of fields that {@link CsvIO.Write} writes, use {@link
+ * CSVFormat#withHeader(String...)}. Note, however, the following constraints:
+ *
+ * <ol>
+ * <li>Each header column must match a field name in the {@link Schema}; matching is case
+ * sensitive.
+ * <li>Matching header columns must match {@link Schema} fields that are valid {@link FieldType}s;
+ * see {@link #VALID_FIELD_TYPE_SET}.
+ * <li>{@link CSVFormat} only allows repeated header columns when {@link
+ * CSVFormat#withAllowDuplicateHeaderNames()}
+ * </ol>
+ *
+ * <p>The following example shows the use of {@link CSVFormat#withHeader(String...)} to control the
+ * order and subset of <code>Transaction</code> fields.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions ...
+ * transactions.apply(
+ * CsvIO
+ * .<Transaction>write("path/to/folder/prefix", CSVFormat.DEFAULT.withHeader("transactionId", "purchaseAmount"))
+ * );
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header is repeated for every
+ * shard file, but only including the subset of fields in their listed order.
+ *
+ * <pre>{@code
+ * transactionId,purchaseAmount
+ * 12345,10.23
+ * 54321,54.65
+ * 98765,11.76
+ * }</pre>
+ *
+ * <p>In addition to header customization, {@link CsvIO.Write} supports {@link
+ * CSVFormat#withHeaderComments(Object...)} as shown below. Note that {@link
+ * CSVFormat#withCommentMarker(char)} is required when specifying header comments.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions = ...
+ * transactions
+ * .apply(
+ * CsvIO.<Transaction>write("path/to/folder/prefix",
+ * CSVFormat.DEFAULT.withCommentMarker('#').withHeaderComments("Bank Report", "1970-01-01", "Operator: John Doe")
+ * );
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header and header comments are
+ * repeated for every shard file.
+ *
+ * <pre>{@code
+ * # Bank Report
+ * # 1970-01-01
+ * # Operator: John Doe
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ *
+ * <p>A {@link PCollection} of {@link Row}s works just like custom Java types illustrated above,
+ * except we use {@link CsvIO#writeRows(String, CSVFormat)} as shown below for the same {@code
+ * Transaction} class. We derive {@code Transaction}'s {@link Schema} using a {@link
+ * org.apache.beam.sdk.schemas.annotations.DefaultSchema.DefaultSchemaProvider}. Note that
+ * hard-coding the {@link Row}s below is for illustration purposes. Developers are instead
+ * encouraged to take advantage of {@link
+ * org.apache.beam.sdk.schemas.annotations.DefaultSchema.DefaultSchemaProvider#toRowFunction(TypeDescriptor)}.
+ *
+ * <pre>{@code
+ * DefaultSchemaProvider defaultSchemaProvider = new DefaultSchemaProvider();
+ * Schema schema = defaultSchemaProvider.schemaFor(TypeDescriptor.of(Transaction.class));
+ * PCollection<Row> transactions = pipeline.apply(Create.of(
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "A")
+ * .withFieldValue("purchaseAmount", 10.23)
+ * .withFieldValue("transactionId", "12345")
+ * .build(),
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "B")
+ * .withFieldValue("purchaseAmount", 54.65)
+ * .withFieldValue("transactionId", "54321")
+ * .build(),
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "C")
+ * .withFieldValue("purchaseAmount", 11.76)
+ * .withFieldValue("transactionId", "98765")
+ * .build()
+ * );
+ *
+ * transactions.apply(
+ * CsvIO
+ * .writeRowsTo("gs://bucket/path/to/folder/prefix", CSVFormat.DEFAULT)
+ * );
+ * }</pre>
+ *
+ * <p>Writing the transactions {@link PCollection} of {@link Row}s would yield the following CSV
+ * file content.
+ *
+ * <pre>{@code
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ *
+ * {@link CsvIO.Write} does not support the following {@link CSVFormat} properties and will throw an
+ * {@link IllegalArgumentException}.
+ *
+ * <ul>
+ * <li>{@link CSVFormat#withSkipHeaderRecord()}
+ * <li>{@link CSVFormat#withAllowMissingColumnNames()} ()}
+ * <li>{@link CSVFormat#withAutoFlush(boolean)}
+ * <li>{@link CSVFormat#withIgnoreHeaderCase()}
+ * <li>{@link CSVFormat#withIgnoreSurroundingSpaces()}
+ * </ul>
+ */
+public class CsvIO {
+ /**
+ * The valid {@link Schema.FieldType} from which {@link CsvIO} converts CSV records to the fields.
+ *
+ * <ul>
+ * <li>{@link FieldType#BYTE}
+ * <li>{@link FieldType#BOOLEAN}
+ * <li>{@link FieldType#DATETIME}
+ * <li>{@link FieldType#DECIMAL}
+ * <li>{@link FieldType#DOUBLE}
+ * <li>{@link FieldType#INT16}
+ * <li>{@link FieldType#INT32}
+ * <li>{@link FieldType#INT64}
+ * <li>{@link FieldType#FLOAT}
+ * <li>{@link FieldType#STRING}
+ * </ul>
+ */
+ public static final Set<Schema.FieldType> VALID_FIELD_TYPE_SET =
+ ImmutableSet.of(
+ FieldType.BYTE,
+ FieldType.BOOLEAN,
+ FieldType.DATETIME,
+ FieldType.DECIMAL,
+ FieldType.DOUBLE,
+ FieldType.INT16,
+ FieldType.INT32,
+ FieldType.INT64,
+ FieldType.FLOAT,
+ FieldType.STRING);
+
+ static final String DEFAULT_FILENAME_SUFFIX = ".csv";
+
+ /** Instantiates a {@link Write} for writing user types in {@link CSVFormat} format. */
+ public static <T> Write<T> write(String to, CSVFormat csvFormat) {
+ return new AutoValue_CsvIO_Write.Builder<T>()
+ .setTextIOWrite(createDefaultTextIOWrite(to))
+ .setCSVFormat(csvFormat)
+ .build();
+ }
+
+ /** Instantiates a {@link Write} for writing {@link Row}s in {@link CSVFormat} format. */
+ public static Write<Row> writeRows(String to, CSVFormat csvFormat) {
+ return new AutoValue_CsvIO_Write.Builder<Row>()
+ .setTextIOWrite(createDefaultTextIOWrite(to))
+ .setCSVFormat(csvFormat)
+ .build();
+ }
+
+ /** {@link PTransform} for writing CSV files. */
+ @AutoValue
+ public abstract static class Write<T>
+ extends PTransform<PCollection<T>, WriteFilesResult<String>> {
+
+ /** Specifies the {@link Compression} of all generated shard files. */
+ public Write<T> withCompression(Compression compression) {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withCompression(compression)).build();
+ }
+
+ public Write<T> withNoSpilling() {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withNoSpilling()).build();
+ }
+
+ /**
+ * Specifies to use a given fixed number of shards per window. See {@link
+ * TextIO.Write#withNumShards(int)}.
+ */
+ public Write<T> withNumShards(Integer numShards) {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withNumShards(numShards)).build();
+ }
+
+ /**
+ * Forces a single file as output and empty shard name template. See {@link
+ * TextIO.Write#withoutSharding()}.
+ */
+ public Write<T> withoutSharding() {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withoutSharding()).build();
+ }
+
+ /**
+ * Uses the given {@link ShardNameTemplate} for naming output files. See {@link
+ * TextIO.Write#withShardNameTemplate(String)}.
+ */
+ public Write<T> withShardTemplate(String shardTemplate) {
+ return toBuilder()
+ .setTextIOWrite(getTextIOWrite().withShardNameTemplate(shardTemplate))
+ .build();
+ }
+
+ /**
+ * Configures the filename suffix for written files. See {@link
+ * TextIO.Write#withSuffix(String)}.
+ */
+ public Write<T> withSuffix(String suffix) {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withSuffix(suffix)).build();
+ }
+
+ /**
+ * Set the base directory used to generate temporary files. See {@link
+ * TextIO.Write#withTempDirectory(ResourceId)}.
+ */
+ public Write<T> withTempDirectory(ResourceId tempDirectory) {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withTempDirectory(tempDirectory)).build();
+ }
+
+ /**
+ * Preserves windowing of input elements and writes them to files based on the element's window.
+ * See {@link TextIO.Write#withWindowedWrites()}.
+ */
+ public Write<T> withWindowedWrites() {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withWindowedWrites()).build();
+ }
+
+ /**
+ * Returns a transform for writing to text files like this one but that has the given {@link
+ * FileBasedSink.WritableByteChannelFactory} to be used by the {@link FileBasedSink} during
+ * output. See {@link
+ * TextIO.Write#withWritableByteChannelFactory(FileBasedSink.WritableByteChannelFactory)}.
+ */
+ public Write<T> withWritableByteChannelFactory(
+ FileBasedSink.WritableByteChannelFactory writableByteChannelFactory) {
+ return toBuilder()
+ .setTextIOWrite(
+ getTextIOWrite().withWritableByteChannelFactory(writableByteChannelFactory))
+ .build();
+ }
+
+ /** The underlying {@link FileIO.Write} that writes converted input to CSV formatted output. */
+ abstract TextIO.Write getTextIOWrite();
+
+ /** The {@link CSVFormat} to inform headers, header comments, and format CSV row content. */
+ abstract CSVFormat getCSVFormat();
+
+ abstract Builder<T> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<T> {
+
+ /**
+ * The underlying {@link FileIO.Write} that writes converted input to CSV formatted output.
+ */
+ abstract Builder<T> setTextIOWrite(TextIO.Write value);
+
+ /** The {@link CSVFormat} to convert input. Defaults to {@link CSVFormat#DEFAULT}. */
+ abstract Builder<T> setCSVFormat(CSVFormat value);
+
+ abstract CSVFormat getCSVFormat();
+
+ abstract Write<T> autoBuild();
+
+ final Write<T> build() {
+ checkArgument(
+ !getCSVFormat().getSkipHeaderRecord(),
+ "withSkipHeaderRecord is an illegal CSVFormat setting");
+
+ if (getCSVFormat().getHeaderComments() != null) {
+ checkArgument(
+ getCSVFormat().isCommentMarkerSet(),
+ "CSVFormat withCommentMarker required when withHeaderComments");
+ }
+ return autoBuild();
+ }
+ }
+
+ @Override
+ public WriteFilesResult<String> expand(PCollection<T> input) {
+ if (!input.hasSchema()) {
+ throw new IllegalArgumentException(
+ String.format(
+ "%s requires an input Schema. Note that only Row or user classes are supported. Consider using TextIO or FileIO directly when writing primitive types",
+ Write.class.getName()));
+ }
+
+ Schema schema = input.getSchema();
+
+ RowCoder rowCoder = RowCoder.of(schema);
+
+ PCollection<Row> rows =
+ input
+ .apply("To Rows", MapElements.into(rows()).via(input.getToRowFunction()))
+ .setCoder(rowCoder);
+
+ CSVFormat csvFormat = applyRequiredCSVFormatSettings(schema);
+
+ String header = formatHeader(csvFormat);
+
+ SerializableFunction<Row, String> toCsvFn =
+ CsvRowConversions.RowToCsv.builder()
+ .setCSVFormat(csvFormat)
+ .setSchema(input.getSchema())
+ .build();
+
+ PCollection<String> csv = rows.apply("To CSV", MapElements.into(strings()).via(toCsvFn));
+
+ return csv.apply("Write CSV", getTextIOWrite().withHeader(header).withOutputFilenames());
+ }
+
+ CSVFormat applyRequiredCSVFormatSettings(Schema schema) {
+ CSVFormat csvFormat = getCSVFormat().withSkipHeaderRecord();
+ if (csvFormat.getHeader() == null) {
+ csvFormat = csvFormat.withHeader(schema.sorted().getFieldNames().toArray(new String[0]));
+ }
+ return csvFormat;
+ }
+
+ private static String formatHeader(CSVFormat csvFormat) {
+ String[] header = requireNonNull(csvFormat.getHeader());
+ List<String> result = new ArrayList<>();
+ if (csvFormat.getHeaderComments() != null) {
+ for (String comment : csvFormat.getHeaderComments()) {
+ result.add(csvFormat.getCommentMarker() + " " + comment);
+ }
+ }
+ CSVFormat withoutHeaderComments = csvFormat.withHeaderComments();
+ result.add(withoutHeaderComments.format((Object[]) header));
+ return String.join("\n", result);
+ }
Review Comment:
In response to this comment, I've:
1. Renamed `applyRequiredCSVFormatSettings` to `buildHeaderFromSchemaIfNeeded`
2. Refactored and renamed `formatHeader` to `writeWithCSVFormatHeaderAndComments` to hopefully provide clarity coupling this behavior with TextIO#withHeader.
Additionally, I wanted to address the questions/comments below:
1. > We are ignoring the case where the user wants to skip the header record and doesn't specify the header. We should only specify the csvFormat header if withSkipHeaderRecord is false.
CSVFormat communicates that withSkipHeaderRecord functions to not output the header. Do you agree then we should keep withSkipHeaderRecord separate from where we are deriving the header?
3. > Note that we must also return null from formatHeader in this case since TextIO#withHeader specifically says that a newline is automatically added after.
Are you saying that the consequence of TextIO#withHeader is:
```
col1, col2, col3
1,2,3
2,3,4
```
vs
```
col1, col2, col3
1,2,3
2,3,4
```
--
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: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [beam] damondouglas commented on a diff in pull request #24630: [24469] Implement CsvIO.Write and supporting classes
Posted by GitBox <gi...@apache.org>.
damondouglas commented on code in PR #24630:
URL: https://github.com/apache/beam/pull/24630#discussion_r1047881763
##########
sdks/java/io/csv/src/main/java/org/apache/beam/sdk/io/csv/CsvIO.java:
##########
@@ -0,0 +1,384 @@
+/*
+ * 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.sdk.io.csv;
+
+import static java.nio.charset.StandardCharsets.UTF_8;
+import static org.apache.beam.sdk.values.TypeDescriptors.rows;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+
+import com.google.auto.value.AutoValue;
+import java.io.BufferedWriter;
+import java.io.IOException;
+import java.io.OutputStreamWriter;
+import java.io.PrintWriter;
+import java.nio.channels.Channels;
+import java.nio.channels.WritableByteChannel;
+import java.util.List;
+import java.util.Objects;
+import java.util.Optional;
+import org.apache.beam.sdk.io.Compression;
+import org.apache.beam.sdk.io.FileIO;
+import org.apache.beam.sdk.io.FileSystems;
+import org.apache.beam.sdk.io.fs.ResourceId;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.transforms.MapElements;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.SimpleFunction;
+import org.apache.beam.sdk.transforms.display.DisplayData;
+import org.apache.beam.sdk.transforms.display.HasDisplayData;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PDone;
+import org.apache.beam.sdk.values.Row;
+import org.apache.commons.csv.CSVFormat;
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+/** Transforms for reading and writing CSV files. */
+@SuppressWarnings({
+ "unused",
+ "nullness" // TODO(https://github.com/apache/beam/issues/20497)
Review Comment:
I removed the "unused". Without "nullness", we get multiple `[dereference.of.nullable] dereference of possibly-null reference` errors despite proceeding null checks.
--
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: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [beam] damondouglas commented on a diff in pull request #24630: [24469] Implement CsvIO.Write and supporting classes
Posted by "damondouglas (via GitHub)" <gi...@apache.org>.
damondouglas commented on code in PR #24630:
URL: https://github.com/apache/beam/pull/24630#discussion_r1085698752
##########
sdks/java/io/csv/src/main/java/org/apache/beam/sdk/io/csv/CsvRowConversions.java:
##########
@@ -0,0 +1,148 @@
+/*
+ * 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.sdk.io.csv;
+
+import static org.apache.beam.sdk.io.csv.CsvIO.VALID_FIELD_TYPE_SET;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkNotNull;
+
+import com.google.auto.value.AutoValue;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Set;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.schemas.Schema.TypeName;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.values.Row;
+import org.apache.commons.csv.CSVFormat;
+import org.checkerframework.checker.nullness.qual.NonNull;
+
+/** Contains classes and methods to help with converting between {@link Row} and CSV strings. */
+class CsvRowConversions {
+
+ /** Converts between {@link Row} and CSV string using a {@link CSVFormat}. */
+ @AutoValue
+ abstract static class RowToCsv implements SerializableFunction<Row, String> {
+
+ static Builder builder() {
+ return new AutoValue_CsvRowConversions_RowToCsv.Builder();
+ }
+
+ /** The expected {@link Schema} of the {@link Row} input. */
+ abstract Schema getSchema();
+
+ /** The {@link CSVFormat} of the converted {@link Row} input. */
+ abstract CSVFormat getCSVFormat();
+
+ /** Converts a {@link Row} to a CSV string formatted using {@link #getCSVFormat()}. */
+ @Override
+ public String apply(Row input) {
+ Row safeInput = checkNotNull(input);
+ String[] header = getHeader();
+ Object[] values = new Object[header.length];
+ for (int i = 0; i < header.length; i++) {
+ values[i] = safeInput.getValue(header[i]);
+ }
+ return getCSVFormat().format(values);
+ }
+
+ @NonNull
+ String[] getHeader() {
+ return checkNotNull(getCSVFormat().getHeader());
+ }
+
+ @AutoValue.Builder
+ abstract static class Builder {
+
+ /** The expected {@link Schema} of the {@link Row} input. */
+ abstract Builder setSchema(Schema schema);
+
+ abstract Schema getSchema();
+
+ /** The {@link CSVFormat} of the converted {@link Row} input. */
+ abstract Builder setCSVFormat(CSVFormat format);
+
+ abstract CSVFormat getCSVFormat();
+
+ abstract RowToCsv autoBuild();
+
+ final RowToCsv build() {
+ checkArgument(getSchema().getFieldCount() > 0, "Schema has no fields");
+ setCSVFormat(
+ getCSVFormat()
+ .withSkipHeaderRecord()
Review Comment:
I wanted to keep CsvRowConversions and CsvIO independent. In the setting of CsvRowConversions.RowToCSV didn't want to assume anyone would remember to set withSkipHeaderRecord.
--
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: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [beam] damondouglas commented on a diff in pull request #24630: [24469] Implement CsvIO.Write and supporting classes
Posted by "damondouglas (via GitHub)" <gi...@apache.org>.
damondouglas commented on code in PR #24630:
URL: https://github.com/apache/beam/pull/24630#discussion_r1085716045
##########
sdks/java/io/csv/src/main/java/org/apache/beam/sdk/io/csv/CsvIO.java:
##########
@@ -0,0 +1,446 @@
+/*
+ * 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.sdk.io.csv;
+
+import static java.util.Objects.requireNonNull;
+import static org.apache.beam.sdk.values.TypeDescriptors.rows;
+import static org.apache.beam.sdk.values.TypeDescriptors.strings;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+
+import com.google.auto.value.AutoValue;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Set;
+import org.apache.beam.sdk.coders.RowCoder;
+import org.apache.beam.sdk.io.Compression;
+import org.apache.beam.sdk.io.FileBasedSink;
+import org.apache.beam.sdk.io.FileIO;
+import org.apache.beam.sdk.io.ShardNameTemplate;
+import org.apache.beam.sdk.io.TextIO;
+import org.apache.beam.sdk.io.WriteFilesResult;
+import org.apache.beam.sdk.io.fs.ResourceId;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.transforms.MapElements;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.sdk.values.TypeDescriptor;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableSet;
+import org.apache.commons.csv.CSVFormat;
+
+/**
+ * {@link PTransform}s for reading and writing CSV files.
+ *
+ * <h2>Reading CSV files</h2>
+ *
+ * <p>Reading from CSV files is not yet implemented. Please see <a
+ * href="https://github.com/apache/beam/issues/24552">https://github.com/apache/beam/issues/24552</a>.
+ *
+ * <h2>Writing CSV files</h2>
+ *
+ * <p>To write a {@link PCollection} to one or more CSV files, use {@link CsvIO.Write}, using {@link
+ * CsvIO#writeRows(String, CSVFormat)} or {@link CsvIO#write(String, CSVFormat)}. {@link
+ * CsvIO.Write} supports writing {@link Row} or custom Java types using an inferred {@link Schema}.
+ * Examples below show both scenarios. See the Beam Programming Guide on <a
+ * href="https://beam.apache.org/documentation/programming-guide/#inferring-schemas">inferring
+ * schemas</a> for more information on how to enable Beam to infer a {@link Schema} from a custom
+ * Java type.
+ *
+ * <p>{@link CsvIO.Write} only supports writing {@link Schema} aware types that do not contain any
+ * nested {@link FieldType}s such a {@link org.apache.beam.sdk.schemas.Schema.TypeName#ROW} or
+ * repeated {@link org.apache.beam.sdk.schemas.Schema.TypeName#ARRAY} types. See {@link
+ * CsvIO.Write#VALID_FIELD_TYPE_SET} for valid {@link FieldType}s.
+ *
+ * <h3>Example usage:</h3>
+ *
+ * <p>Suppose we have the following <code>Transaction</code> class annotated with
+ * {@code @DefaultSchema(JavaBeanSchema.class)} so that Beam can infer its {@link Schema}:
+ *
+ * <pre>{@code @DefaultSchema(JavaBeanSchema.class)
+ * public class Transaction {
+ * public Transaction() { … }
+ * public Long getTransactionId();
+ * public void setTransactionId(Long transactionId) { … }
+ * public String getBank() { … }
+ * public void setBank(String bank) { … }
+ * public double getPurchaseAmount() { … }
+ * public void setPurchaseAmount(double purchaseAmount) { … }
+ * }
+ * }</pre>
+ *
+ * <p>From a {@code PCollection<Transaction>}, {@link CsvIO.Write} can write one or many CSV files
+ * automatically creating the header based on its inferred {@link Schema}.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions = ...
+ * transactions.apply(CsvIO.<Transaction>write("path/to/folder/prefix", CSVFormat.DEFAULT));
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header is repeated for every
+ * shard file, whereas by default, {@link CsvIO.Write} will write all fields in <b>sorted order</b>
+ * of the field names.
+ *
+ * <pre>{@code
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ *
+ * <p>To control the order and subset of fields that {@link CsvIO.Write} writes, use {@link
+ * CSVFormat#withHeader(String...)}. Note, however, the following constraints:
+ *
+ * <ol>
+ * <li>Each header column must match a field name in the {@link Schema}; matching is case
+ * sensitive.
+ * <li>Matching header columns must match {@link Schema} fields that are valid {@link FieldType}s;
+ * see {@link #VALID_FIELD_TYPE_SET}.
+ * <li>{@link CSVFormat} only allows repeated header columns when {@link
+ * CSVFormat#withAllowDuplicateHeaderNames()}
+ * </ol>
+ *
+ * <p>The following example shows the use of {@link CSVFormat#withHeader(String...)} to control the
+ * order and subset of <code>Transaction</code> fields.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions ...
+ * transactions.apply(
+ * CsvIO
+ * .<Transaction>write("path/to/folder/prefix", CSVFormat.DEFAULT.withHeader("transactionId", "purchaseAmount"))
+ * );
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header is repeated for every
+ * shard file, but only including the subset of fields in their listed order.
+ *
+ * <pre>{@code
+ * transactionId,purchaseAmount
+ * 12345,10.23
+ * 54321,54.65
+ * 98765,11.76
+ * }</pre>
+ *
+ * <p>In addition to header customization, {@link CsvIO.Write} supports {@link
+ * CSVFormat#withHeaderComments(Object...)} as shown below. Note that {@link
+ * CSVFormat#withCommentMarker(char)} is required when specifying header comments.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions = ...
+ * transactions
+ * .apply(
+ * CsvIO.<Transaction>write("path/to/folder/prefix",
+ * CSVFormat.DEFAULT.withCommentMarker('#').withHeaderComments("Bank Report", "1970-01-01", "Operator: John Doe")
+ * );
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header and header comments are
+ * repeated for every shard file.
+ *
+ * <pre>{@code
+ * # Bank Report
+ * # 1970-01-01
+ * # Operator: John Doe
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ *
+ * <p>A {@link PCollection} of {@link Row}s works just like custom Java types illustrated above,
+ * except we use {@link CsvIO#writeRows(String, CSVFormat)} as shown below for the same {@code
+ * Transaction} class. We derive {@code Transaction}'s {@link Schema} using a {@link
+ * org.apache.beam.sdk.schemas.annotations.DefaultSchema.DefaultSchemaProvider}. Note that
+ * hard-coding the {@link Row}s below is for illustration purposes. Developers are instead
+ * encouraged to take advantage of {@link
+ * org.apache.beam.sdk.schemas.annotations.DefaultSchema.DefaultSchemaProvider#toRowFunction(TypeDescriptor)}.
+ *
+ * <pre>{@code
+ * DefaultSchemaProvider defaultSchemaProvider = new DefaultSchemaProvider();
+ * Schema schema = defaultSchemaProvider.schemaFor(TypeDescriptor.of(Transaction.class));
+ * PCollection<Row> transactions = pipeline.apply(Create.of(
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "A")
+ * .withFieldValue("purchaseAmount", 10.23)
+ * .withFieldValue("transactionId", "12345")
+ * .build(),
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "B")
+ * .withFieldValue("purchaseAmount", 54.65)
+ * .withFieldValue("transactionId", "54321")
+ * .build(),
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "C")
+ * .withFieldValue("purchaseAmount", 11.76)
+ * .withFieldValue("transactionId", "98765")
+ * .build()
+ * );
+ *
+ * transactions.apply(
+ * CsvIO
+ * .writeRowsTo("gs://bucket/path/to/folder/prefix", CSVFormat.DEFAULT)
+ * );
+ * }</pre>
+ *
+ * <p>Writing the transactions {@link PCollection} of {@link Row}s would yield the following CSV
+ * file content.
+ *
+ * <pre>{@code
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ *
+ * {@link CsvIO.Write} does not support the following {@link CSVFormat} properties and will throw an
+ * {@link IllegalArgumentException}.
+ *
+ * <ul>
+ * <li>{@link CSVFormat#withSkipHeaderRecord()}
+ * <li>{@link CSVFormat#withAllowMissingColumnNames()} ()}
+ * <li>{@link CSVFormat#withAutoFlush(boolean)}
+ * <li>{@link CSVFormat#withIgnoreHeaderCase()}
+ * <li>{@link CSVFormat#withIgnoreSurroundingSpaces()}
+ * </ul>
+ */
+public class CsvIO {
+ /**
+ * The valid {@link Schema.FieldType} from which {@link CsvIO} converts CSV records to the fields.
+ *
+ * <ul>
+ * <li>{@link FieldType#BYTE}
+ * <li>{@link FieldType#BOOLEAN}
+ * <li>{@link FieldType#DATETIME}
+ * <li>{@link FieldType#DECIMAL}
+ * <li>{@link FieldType#DOUBLE}
+ * <li>{@link FieldType#INT16}
+ * <li>{@link FieldType#INT32}
+ * <li>{@link FieldType#INT64}
+ * <li>{@link FieldType#FLOAT}
+ * <li>{@link FieldType#STRING}
+ * </ul>
+ */
+ public static final Set<Schema.FieldType> VALID_FIELD_TYPE_SET =
+ ImmutableSet.of(
+ FieldType.BYTE,
+ FieldType.BOOLEAN,
+ FieldType.DATETIME,
+ FieldType.DECIMAL,
+ FieldType.DOUBLE,
+ FieldType.INT16,
+ FieldType.INT32,
+ FieldType.INT64,
+ FieldType.FLOAT,
+ FieldType.STRING);
+
+ static final String DEFAULT_FILENAME_SUFFIX = ".csv";
+
+ /** Instantiates a {@link Write} for writing user types in {@link CSVFormat} format. */
+ public static <T> Write<T> write(String to, CSVFormat csvFormat) {
+ return new AutoValue_CsvIO_Write.Builder<T>()
+ .setTextIOWrite(createDefaultTextIOWrite(to))
+ .setCSVFormat(csvFormat)
+ .build();
+ }
+
+ /** Instantiates a {@link Write} for writing {@link Row}s in {@link CSVFormat} format. */
+ public static Write<Row> writeRows(String to, CSVFormat csvFormat) {
+ return new AutoValue_CsvIO_Write.Builder<Row>()
+ .setTextIOWrite(createDefaultTextIOWrite(to))
+ .setCSVFormat(csvFormat)
+ .build();
+ }
+
+ /** {@link PTransform} for writing CSV files. */
+ @AutoValue
+ public abstract static class Write<T>
+ extends PTransform<PCollection<T>, WriteFilesResult<String>> {
+
+ /** Specifies the {@link Compression} of all generated shard files. */
+ public Write<T> withCompression(Compression compression) {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withCompression(compression)).build();
+ }
+
+ public Write<T> withNoSpilling() {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withNoSpilling()).build();
+ }
+
+ /**
+ * Specifies to use a given fixed number of shards per window. See {@link
+ * TextIO.Write#withNumShards(int)}.
+ */
+ public Write<T> withNumShards(Integer numShards) {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withNumShards(numShards)).build();
+ }
+
+ /**
+ * Forces a single file as output and empty shard name template. See {@link
+ * TextIO.Write#withoutSharding()}.
+ */
+ public Write<T> withoutSharding() {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withoutSharding()).build();
+ }
+
+ /**
+ * Uses the given {@link ShardNameTemplate} for naming output files. See {@link
+ * TextIO.Write#withShardNameTemplate(String)}.
+ */
+ public Write<T> withShardTemplate(String shardTemplate) {
+ return toBuilder()
+ .setTextIOWrite(getTextIOWrite().withShardNameTemplate(shardTemplate))
+ .build();
+ }
+
+ /**
+ * Configures the filename suffix for written files. See {@link
+ * TextIO.Write#withSuffix(String)}.
+ */
+ public Write<T> withSuffix(String suffix) {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withSuffix(suffix)).build();
+ }
+
+ /**
+ * Set the base directory used to generate temporary files. See {@link
+ * TextIO.Write#withTempDirectory(ResourceId)}.
+ */
+ public Write<T> withTempDirectory(ResourceId tempDirectory) {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withTempDirectory(tempDirectory)).build();
+ }
+
+ /**
+ * Preserves windowing of input elements and writes them to files based on the element's window.
+ * See {@link TextIO.Write#withWindowedWrites()}.
+ */
+ public Write<T> withWindowedWrites() {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withWindowedWrites()).build();
+ }
+
+ /**
+ * Returns a transform for writing to text files like this one but that has the given {@link
+ * FileBasedSink.WritableByteChannelFactory} to be used by the {@link FileBasedSink} during
+ * output. See {@link
+ * TextIO.Write#withWritableByteChannelFactory(FileBasedSink.WritableByteChannelFactory)}.
+ */
+ public Write<T> withWritableByteChannelFactory(
+ FileBasedSink.WritableByteChannelFactory writableByteChannelFactory) {
+ return toBuilder()
+ .setTextIOWrite(
+ getTextIOWrite().withWritableByteChannelFactory(writableByteChannelFactory))
+ .build();
+ }
+
+ /** The underlying {@link FileIO.Write} that writes converted input to CSV formatted output. */
+ abstract TextIO.Write getTextIOWrite();
+
+ /** The {@link CSVFormat} to inform headers, header comments, and format CSV row content. */
+ abstract CSVFormat getCSVFormat();
+
+ abstract Builder<T> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<T> {
+
+ /**
+ * The underlying {@link FileIO.Write} that writes converted input to CSV formatted output.
+ */
+ abstract Builder<T> setTextIOWrite(TextIO.Write value);
+
+ /** The {@link CSVFormat} to convert input. Defaults to {@link CSVFormat#DEFAULT}. */
+ abstract Builder<T> setCSVFormat(CSVFormat value);
+
+ abstract CSVFormat getCSVFormat();
+
+ abstract Write<T> autoBuild();
+
+ final Write<T> build() {
+ checkArgument(
+ !getCSVFormat().getSkipHeaderRecord(),
+ "withSkipHeaderRecord is an illegal CSVFormat setting");
+
+ if (getCSVFormat().getHeaderComments() != null) {
+ checkArgument(
+ getCSVFormat().isCommentMarkerSet(),
+ "CSVFormat withCommentMarker required when withHeaderComments");
+ }
+ return autoBuild();
+ }
+ }
+
+ @Override
+ public WriteFilesResult<String> expand(PCollection<T> input) {
+ if (!input.hasSchema()) {
+ throw new IllegalArgumentException(
+ String.format(
+ "%s requires an input Schema. Note that only Row or user classes are supported. Consider using TextIO or FileIO directly when writing primitive types",
+ Write.class.getName()));
+ }
+
+ Schema schema = input.getSchema();
+
+ RowCoder rowCoder = RowCoder.of(schema);
+
+ PCollection<Row> rows =
+ input
+ .apply("To Rows", MapElements.into(rows()).via(input.getToRowFunction()))
+ .setCoder(rowCoder);
+
+ CSVFormat csvFormat = applyRequiredCSVFormatSettings(schema);
+
+ String header = formatHeader(csvFormat);
+
+ SerializableFunction<Row, String> toCsvFn =
+ CsvRowConversions.RowToCsv.builder()
+ .setCSVFormat(csvFormat)
+ .setSchema(input.getSchema())
+ .build();
+
+ PCollection<String> csv = rows.apply("To CSV", MapElements.into(strings()).via(toCsvFn));
+
+ return csv.apply("Write CSV", getTextIOWrite().withHeader(header).withOutputFilenames());
+ }
+
+ CSVFormat applyRequiredCSVFormatSettings(Schema schema) {
+ CSVFormat csvFormat = getCSVFormat().withSkipHeaderRecord();
+ if (csvFormat.getHeader() == null) {
+ csvFormat = csvFormat.withHeader(schema.sorted().getFieldNames().toArray(new String[0]));
+ }
+ return csvFormat;
+ }
+
+ private static String formatHeader(CSVFormat csvFormat) {
+ String[] header = requireNonNull(csvFormat.getHeader());
+ List<String> result = new ArrayList<>();
+ if (csvFormat.getHeaderComments() != null) {
+ for (String comment : csvFormat.getHeaderComments()) {
+ result.add(csvFormat.getCommentMarker() + " " + comment);
+ }
+ }
+ CSVFormat withoutHeaderComments = csvFormat.withHeaderComments();
+ result.add(withoutHeaderComments.format((Object[]) header));
+ return String.join("\n", result);
+ }
Review Comment:
In response to this comment, I've:
1. Renamed `applyRequiredCSVFormatSettings` to `buildHeaderFromSchemaIfNeeded`
2. Refactored and renamed `formatHeader` to `writeWithCSVFormatHeaderAndComments` to hopefully provide clarity coupling this behavior with TextIO#withHeader.
Additionally, I wanted to address the questions/comments below:
1. > We are ignoring the case where the user wants to skip the header record and doesn't specify the header. We should only specify the csvFormat header if withSkipHeaderRecord is false.
CSVFormat communicates that withSkipHeaderRecord functions to not output the header. It's also a parameter that gives an output incompatible in a beam context as CSVFormat was designed for writing a single file. Therefore, do you agree then we should not use withSkipHeaderRecord to drive how we are building the header?
3. > Note that we must also return null from formatHeader in this case since TextIO#withHeader specifically says that a newline is automatically added after.
Apologies, your comment confused me.
Are you saying that the consequence of TextIO#withHeader is:
```
col1, col2, col3
1,2,3
2,3,4
```
vs
```
col1, col2, col3
1,2,3
2,3,4
```
--
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: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [beam] lukecwik commented on a diff in pull request #24630: [24469] Implement CsvIO.Write and supporting classes
Posted by "lukecwik (via GitHub)" <gi...@apache.org>.
lukecwik commented on code in PR #24630:
URL: https://github.com/apache/beam/pull/24630#discussion_r1085735320
##########
sdks/java/io/csv/src/main/java/org/apache/beam/sdk/io/csv/CsvIO.java:
##########
@@ -0,0 +1,446 @@
+/*
+ * 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.sdk.io.csv;
+
+import static java.util.Objects.requireNonNull;
+import static org.apache.beam.sdk.values.TypeDescriptors.rows;
+import static org.apache.beam.sdk.values.TypeDescriptors.strings;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+
+import com.google.auto.value.AutoValue;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Set;
+import org.apache.beam.sdk.coders.RowCoder;
+import org.apache.beam.sdk.io.Compression;
+import org.apache.beam.sdk.io.FileBasedSink;
+import org.apache.beam.sdk.io.FileIO;
+import org.apache.beam.sdk.io.ShardNameTemplate;
+import org.apache.beam.sdk.io.TextIO;
+import org.apache.beam.sdk.io.WriteFilesResult;
+import org.apache.beam.sdk.io.fs.ResourceId;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.transforms.MapElements;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.sdk.values.TypeDescriptor;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableSet;
+import org.apache.commons.csv.CSVFormat;
+
+/**
+ * {@link PTransform}s for reading and writing CSV files.
+ *
+ * <h2>Reading CSV files</h2>
+ *
+ * <p>Reading from CSV files is not yet implemented. Please see <a
+ * href="https://github.com/apache/beam/issues/24552">https://github.com/apache/beam/issues/24552</a>.
+ *
+ * <h2>Writing CSV files</h2>
+ *
+ * <p>To write a {@link PCollection} to one or more CSV files, use {@link CsvIO.Write}, using {@link
+ * CsvIO#writeRows(String, CSVFormat)} or {@link CsvIO#write(String, CSVFormat)}. {@link
+ * CsvIO.Write} supports writing {@link Row} or custom Java types using an inferred {@link Schema}.
+ * Examples below show both scenarios. See the Beam Programming Guide on <a
+ * href="https://beam.apache.org/documentation/programming-guide/#inferring-schemas">inferring
+ * schemas</a> for more information on how to enable Beam to infer a {@link Schema} from a custom
+ * Java type.
+ *
+ * <p>{@link CsvIO.Write} only supports writing {@link Schema} aware types that do not contain any
+ * nested {@link FieldType}s such a {@link org.apache.beam.sdk.schemas.Schema.TypeName#ROW} or
+ * repeated {@link org.apache.beam.sdk.schemas.Schema.TypeName#ARRAY} types. See {@link
+ * CsvIO.Write#VALID_FIELD_TYPE_SET} for valid {@link FieldType}s.
+ *
+ * <h3>Example usage:</h3>
+ *
+ * <p>Suppose we have the following <code>Transaction</code> class annotated with
+ * {@code @DefaultSchema(JavaBeanSchema.class)} so that Beam can infer its {@link Schema}:
+ *
+ * <pre>{@code @DefaultSchema(JavaBeanSchema.class)
+ * public class Transaction {
+ * public Transaction() { … }
+ * public Long getTransactionId();
+ * public void setTransactionId(Long transactionId) { … }
+ * public String getBank() { … }
+ * public void setBank(String bank) { … }
+ * public double getPurchaseAmount() { … }
+ * public void setPurchaseAmount(double purchaseAmount) { … }
+ * }
+ * }</pre>
+ *
+ * <p>From a {@code PCollection<Transaction>}, {@link CsvIO.Write} can write one or many CSV files
+ * automatically creating the header based on its inferred {@link Schema}.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions = ...
+ * transactions.apply(CsvIO.<Transaction>write("path/to/folder/prefix", CSVFormat.DEFAULT));
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header is repeated for every
+ * shard file, whereas by default, {@link CsvIO.Write} will write all fields in <b>sorted order</b>
+ * of the field names.
+ *
+ * <pre>{@code
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ *
+ * <p>To control the order and subset of fields that {@link CsvIO.Write} writes, use {@link
+ * CSVFormat#withHeader(String...)}. Note, however, the following constraints:
+ *
+ * <ol>
+ * <li>Each header column must match a field name in the {@link Schema}; matching is case
+ * sensitive.
+ * <li>Matching header columns must match {@link Schema} fields that are valid {@link FieldType}s;
+ * see {@link #VALID_FIELD_TYPE_SET}.
+ * <li>{@link CSVFormat} only allows repeated header columns when {@link
+ * CSVFormat#withAllowDuplicateHeaderNames()}
+ * </ol>
+ *
+ * <p>The following example shows the use of {@link CSVFormat#withHeader(String...)} to control the
+ * order and subset of <code>Transaction</code> fields.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions ...
+ * transactions.apply(
+ * CsvIO
+ * .<Transaction>write("path/to/folder/prefix", CSVFormat.DEFAULT.withHeader("transactionId", "purchaseAmount"))
+ * );
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header is repeated for every
+ * shard file, but only including the subset of fields in their listed order.
+ *
+ * <pre>{@code
+ * transactionId,purchaseAmount
+ * 12345,10.23
+ * 54321,54.65
+ * 98765,11.76
+ * }</pre>
+ *
+ * <p>In addition to header customization, {@link CsvIO.Write} supports {@link
+ * CSVFormat#withHeaderComments(Object...)} as shown below. Note that {@link
+ * CSVFormat#withCommentMarker(char)} is required when specifying header comments.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions = ...
+ * transactions
+ * .apply(
+ * CsvIO.<Transaction>write("path/to/folder/prefix",
+ * CSVFormat.DEFAULT.withCommentMarker('#').withHeaderComments("Bank Report", "1970-01-01", "Operator: John Doe")
+ * );
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header and header comments are
+ * repeated for every shard file.
+ *
+ * <pre>{@code
+ * # Bank Report
+ * # 1970-01-01
+ * # Operator: John Doe
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ *
+ * <p>A {@link PCollection} of {@link Row}s works just like custom Java types illustrated above,
+ * except we use {@link CsvIO#writeRows(String, CSVFormat)} as shown below for the same {@code
+ * Transaction} class. We derive {@code Transaction}'s {@link Schema} using a {@link
+ * org.apache.beam.sdk.schemas.annotations.DefaultSchema.DefaultSchemaProvider}. Note that
+ * hard-coding the {@link Row}s below is for illustration purposes. Developers are instead
+ * encouraged to take advantage of {@link
+ * org.apache.beam.sdk.schemas.annotations.DefaultSchema.DefaultSchemaProvider#toRowFunction(TypeDescriptor)}.
+ *
+ * <pre>{@code
+ * DefaultSchemaProvider defaultSchemaProvider = new DefaultSchemaProvider();
+ * Schema schema = defaultSchemaProvider.schemaFor(TypeDescriptor.of(Transaction.class));
+ * PCollection<Row> transactions = pipeline.apply(Create.of(
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "A")
+ * .withFieldValue("purchaseAmount", 10.23)
+ * .withFieldValue("transactionId", "12345")
+ * .build(),
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "B")
+ * .withFieldValue("purchaseAmount", 54.65)
+ * .withFieldValue("transactionId", "54321")
+ * .build(),
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "C")
+ * .withFieldValue("purchaseAmount", 11.76)
+ * .withFieldValue("transactionId", "98765")
+ * .build()
+ * );
+ *
+ * transactions.apply(
+ * CsvIO
+ * .writeRowsTo("gs://bucket/path/to/folder/prefix", CSVFormat.DEFAULT)
+ * );
+ * }</pre>
+ *
+ * <p>Writing the transactions {@link PCollection} of {@link Row}s would yield the following CSV
+ * file content.
+ *
+ * <pre>{@code
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ *
+ * {@link CsvIO.Write} does not support the following {@link CSVFormat} properties and will throw an
+ * {@link IllegalArgumentException}.
+ *
+ * <ul>
+ * <li>{@link CSVFormat#withSkipHeaderRecord()}
+ * <li>{@link CSVFormat#withAllowMissingColumnNames()} ()}
+ * <li>{@link CSVFormat#withAutoFlush(boolean)}
+ * <li>{@link CSVFormat#withIgnoreHeaderCase()}
+ * <li>{@link CSVFormat#withIgnoreSurroundingSpaces()}
+ * </ul>
+ */
+public class CsvIO {
+ /**
+ * The valid {@link Schema.FieldType} from which {@link CsvIO} converts CSV records to the fields.
+ *
+ * <ul>
+ * <li>{@link FieldType#BYTE}
+ * <li>{@link FieldType#BOOLEAN}
+ * <li>{@link FieldType#DATETIME}
+ * <li>{@link FieldType#DECIMAL}
+ * <li>{@link FieldType#DOUBLE}
+ * <li>{@link FieldType#INT16}
+ * <li>{@link FieldType#INT32}
+ * <li>{@link FieldType#INT64}
+ * <li>{@link FieldType#FLOAT}
+ * <li>{@link FieldType#STRING}
+ * </ul>
+ */
+ public static final Set<Schema.FieldType> VALID_FIELD_TYPE_SET =
+ ImmutableSet.of(
+ FieldType.BYTE,
+ FieldType.BOOLEAN,
+ FieldType.DATETIME,
+ FieldType.DECIMAL,
+ FieldType.DOUBLE,
+ FieldType.INT16,
+ FieldType.INT32,
+ FieldType.INT64,
+ FieldType.FLOAT,
+ FieldType.STRING);
+
+ static final String DEFAULT_FILENAME_SUFFIX = ".csv";
+
+ /** Instantiates a {@link Write} for writing user types in {@link CSVFormat} format. */
+ public static <T> Write<T> write(String to, CSVFormat csvFormat) {
+ return new AutoValue_CsvIO_Write.Builder<T>()
+ .setTextIOWrite(createDefaultTextIOWrite(to))
+ .setCSVFormat(csvFormat)
+ .build();
+ }
+
+ /** Instantiates a {@link Write} for writing {@link Row}s in {@link CSVFormat} format. */
+ public static Write<Row> writeRows(String to, CSVFormat csvFormat) {
+ return new AutoValue_CsvIO_Write.Builder<Row>()
+ .setTextIOWrite(createDefaultTextIOWrite(to))
+ .setCSVFormat(csvFormat)
+ .build();
+ }
+
+ /** {@link PTransform} for writing CSV files. */
+ @AutoValue
+ public abstract static class Write<T>
+ extends PTransform<PCollection<T>, WriteFilesResult<String>> {
+
+ /** Specifies the {@link Compression} of all generated shard files. */
+ public Write<T> withCompression(Compression compression) {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withCompression(compression)).build();
+ }
+
+ public Write<T> withNoSpilling() {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withNoSpilling()).build();
+ }
+
+ /**
+ * Specifies to use a given fixed number of shards per window. See {@link
+ * TextIO.Write#withNumShards(int)}.
+ */
+ public Write<T> withNumShards(Integer numShards) {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withNumShards(numShards)).build();
+ }
+
+ /**
+ * Forces a single file as output and empty shard name template. See {@link
+ * TextIO.Write#withoutSharding()}.
+ */
+ public Write<T> withoutSharding() {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withoutSharding()).build();
+ }
+
+ /**
+ * Uses the given {@link ShardNameTemplate} for naming output files. See {@link
+ * TextIO.Write#withShardNameTemplate(String)}.
+ */
+ public Write<T> withShardTemplate(String shardTemplate) {
+ return toBuilder()
+ .setTextIOWrite(getTextIOWrite().withShardNameTemplate(shardTemplate))
+ .build();
+ }
+
+ /**
+ * Configures the filename suffix for written files. See {@link
+ * TextIO.Write#withSuffix(String)}.
+ */
+ public Write<T> withSuffix(String suffix) {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withSuffix(suffix)).build();
+ }
+
+ /**
+ * Set the base directory used to generate temporary files. See {@link
+ * TextIO.Write#withTempDirectory(ResourceId)}.
+ */
+ public Write<T> withTempDirectory(ResourceId tempDirectory) {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withTempDirectory(tempDirectory)).build();
+ }
+
+ /**
+ * Preserves windowing of input elements and writes them to files based on the element's window.
+ * See {@link TextIO.Write#withWindowedWrites()}.
+ */
+ public Write<T> withWindowedWrites() {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withWindowedWrites()).build();
+ }
+
+ /**
+ * Returns a transform for writing to text files like this one but that has the given {@link
+ * FileBasedSink.WritableByteChannelFactory} to be used by the {@link FileBasedSink} during
+ * output. See {@link
+ * TextIO.Write#withWritableByteChannelFactory(FileBasedSink.WritableByteChannelFactory)}.
+ */
+ public Write<T> withWritableByteChannelFactory(
+ FileBasedSink.WritableByteChannelFactory writableByteChannelFactory) {
+ return toBuilder()
+ .setTextIOWrite(
+ getTextIOWrite().withWritableByteChannelFactory(writableByteChannelFactory))
+ .build();
+ }
+
+ /** The underlying {@link FileIO.Write} that writes converted input to CSV formatted output. */
+ abstract TextIO.Write getTextIOWrite();
+
+ /** The {@link CSVFormat} to inform headers, header comments, and format CSV row content. */
+ abstract CSVFormat getCSVFormat();
+
+ abstract Builder<T> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<T> {
+
+ /**
+ * The underlying {@link FileIO.Write} that writes converted input to CSV formatted output.
+ */
+ abstract Builder<T> setTextIOWrite(TextIO.Write value);
+
+ /** The {@link CSVFormat} to convert input. Defaults to {@link CSVFormat#DEFAULT}. */
+ abstract Builder<T> setCSVFormat(CSVFormat value);
+
+ abstract CSVFormat getCSVFormat();
+
+ abstract Write<T> autoBuild();
+
+ final Write<T> build() {
+ checkArgument(
+ !getCSVFormat().getSkipHeaderRecord(),
+ "withSkipHeaderRecord is an illegal CSVFormat setting");
+
+ if (getCSVFormat().getHeaderComments() != null) {
+ checkArgument(
+ getCSVFormat().isCommentMarkerSet(),
+ "CSVFormat withCommentMarker required when withHeaderComments");
+ }
+ return autoBuild();
+ }
+ }
+
+ @Override
+ public WriteFilesResult<String> expand(PCollection<T> input) {
+ if (!input.hasSchema()) {
+ throw new IllegalArgumentException(
+ String.format(
+ "%s requires an input Schema. Note that only Row or user classes are supported. Consider using TextIO or FileIO directly when writing primitive types",
+ Write.class.getName()));
+ }
+
+ Schema schema = input.getSchema();
+
+ RowCoder rowCoder = RowCoder.of(schema);
+
+ PCollection<Row> rows =
+ input
+ .apply("To Rows", MapElements.into(rows()).via(input.getToRowFunction()))
+ .setCoder(rowCoder);
+
+ CSVFormat csvFormat = applyRequiredCSVFormatSettings(schema);
+
+ String header = formatHeader(csvFormat);
+
+ SerializableFunction<Row, String> toCsvFn =
+ CsvRowConversions.RowToCsv.builder()
+ .setCSVFormat(csvFormat)
+ .setSchema(input.getSchema())
+ .build();
+
+ PCollection<String> csv = rows.apply("To CSV", MapElements.into(strings()).via(toCsvFn));
+
+ return csv.apply("Write CSV", getTextIOWrite().withHeader(header).withOutputFilenames());
+ }
+
+ CSVFormat applyRequiredCSVFormatSettings(Schema schema) {
+ CSVFormat csvFormat = getCSVFormat().withSkipHeaderRecord();
+ if (csvFormat.getHeader() == null) {
+ csvFormat = csvFormat.withHeader(schema.sorted().getFieldNames().toArray(new String[0]));
+ }
+ return csvFormat;
+ }
+
+ private static String formatHeader(CSVFormat csvFormat) {
+ String[] header = requireNonNull(csvFormat.getHeader());
+ List<String> result = new ArrayList<>();
+ if (csvFormat.getHeaderComments() != null) {
+ for (String comment : csvFormat.getHeaderComments()) {
+ result.add(csvFormat.getCommentMarker() + " " + comment);
+ }
+ }
+ CSVFormat withoutHeaderComments = csvFormat.withHeaderComments();
+ result.add(withoutHeaderComments.format((Object[]) header));
+ return String.join("\n", result);
+ }
Review Comment:
I'm saying that when header is `null` the output is:
```
1,2,3
2,3,4
```
while when header is `""` the output is:
```
1,2,3
2,3,4
```
--
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: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [beam] github-actions[bot] commented on pull request #24630: [24469] Implement CsvIO.Write and supporting classes
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on PR #24630:
URL: https://github.com/apache/beam/pull/24630#issuecomment-1344888010
Assigning reviewers. If you would like to opt out of this review, comment `assign to next reviewer`:
R: @lukecwik for label java.
R: @Abacn for label build.
R: @pabloem for label io.
Available commands:
- `stop reviewer notifications` - opt out of the automated review tooling
- `remind me after tests pass` - tag the comment author after tests pass
- `waiting on author` - shift the attention set back to the author (any comment or push by the author will return the attention set to the reviewers)
The PR bot will only process comments in the main thread (not review comments).
--
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: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [beam] damondouglas commented on a diff in pull request #24630: [24469] Implement CsvIO.Write and supporting classes
Posted by "damondouglas (via GitHub)" <gi...@apache.org>.
damondouglas commented on code in PR #24630:
URL: https://github.com/apache/beam/pull/24630#discussion_r1085716045
##########
sdks/java/io/csv/src/main/java/org/apache/beam/sdk/io/csv/CsvIO.java:
##########
@@ -0,0 +1,446 @@
+/*
+ * 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.sdk.io.csv;
+
+import static java.util.Objects.requireNonNull;
+import static org.apache.beam.sdk.values.TypeDescriptors.rows;
+import static org.apache.beam.sdk.values.TypeDescriptors.strings;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+
+import com.google.auto.value.AutoValue;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Set;
+import org.apache.beam.sdk.coders.RowCoder;
+import org.apache.beam.sdk.io.Compression;
+import org.apache.beam.sdk.io.FileBasedSink;
+import org.apache.beam.sdk.io.FileIO;
+import org.apache.beam.sdk.io.ShardNameTemplate;
+import org.apache.beam.sdk.io.TextIO;
+import org.apache.beam.sdk.io.WriteFilesResult;
+import org.apache.beam.sdk.io.fs.ResourceId;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.transforms.MapElements;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.sdk.values.TypeDescriptor;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableSet;
+import org.apache.commons.csv.CSVFormat;
+
+/**
+ * {@link PTransform}s for reading and writing CSV files.
+ *
+ * <h2>Reading CSV files</h2>
+ *
+ * <p>Reading from CSV files is not yet implemented. Please see <a
+ * href="https://github.com/apache/beam/issues/24552">https://github.com/apache/beam/issues/24552</a>.
+ *
+ * <h2>Writing CSV files</h2>
+ *
+ * <p>To write a {@link PCollection} to one or more CSV files, use {@link CsvIO.Write}, using {@link
+ * CsvIO#writeRows(String, CSVFormat)} or {@link CsvIO#write(String, CSVFormat)}. {@link
+ * CsvIO.Write} supports writing {@link Row} or custom Java types using an inferred {@link Schema}.
+ * Examples below show both scenarios. See the Beam Programming Guide on <a
+ * href="https://beam.apache.org/documentation/programming-guide/#inferring-schemas">inferring
+ * schemas</a> for more information on how to enable Beam to infer a {@link Schema} from a custom
+ * Java type.
+ *
+ * <p>{@link CsvIO.Write} only supports writing {@link Schema} aware types that do not contain any
+ * nested {@link FieldType}s such a {@link org.apache.beam.sdk.schemas.Schema.TypeName#ROW} or
+ * repeated {@link org.apache.beam.sdk.schemas.Schema.TypeName#ARRAY} types. See {@link
+ * CsvIO.Write#VALID_FIELD_TYPE_SET} for valid {@link FieldType}s.
+ *
+ * <h3>Example usage:</h3>
+ *
+ * <p>Suppose we have the following <code>Transaction</code> class annotated with
+ * {@code @DefaultSchema(JavaBeanSchema.class)} so that Beam can infer its {@link Schema}:
+ *
+ * <pre>{@code @DefaultSchema(JavaBeanSchema.class)
+ * public class Transaction {
+ * public Transaction() { … }
+ * public Long getTransactionId();
+ * public void setTransactionId(Long transactionId) { … }
+ * public String getBank() { … }
+ * public void setBank(String bank) { … }
+ * public double getPurchaseAmount() { … }
+ * public void setPurchaseAmount(double purchaseAmount) { … }
+ * }
+ * }</pre>
+ *
+ * <p>From a {@code PCollection<Transaction>}, {@link CsvIO.Write} can write one or many CSV files
+ * automatically creating the header based on its inferred {@link Schema}.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions = ...
+ * transactions.apply(CsvIO.<Transaction>write("path/to/folder/prefix", CSVFormat.DEFAULT));
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header is repeated for every
+ * shard file, whereas by default, {@link CsvIO.Write} will write all fields in <b>sorted order</b>
+ * of the field names.
+ *
+ * <pre>{@code
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ *
+ * <p>To control the order and subset of fields that {@link CsvIO.Write} writes, use {@link
+ * CSVFormat#withHeader(String...)}. Note, however, the following constraints:
+ *
+ * <ol>
+ * <li>Each header column must match a field name in the {@link Schema}; matching is case
+ * sensitive.
+ * <li>Matching header columns must match {@link Schema} fields that are valid {@link FieldType}s;
+ * see {@link #VALID_FIELD_TYPE_SET}.
+ * <li>{@link CSVFormat} only allows repeated header columns when {@link
+ * CSVFormat#withAllowDuplicateHeaderNames()}
+ * </ol>
+ *
+ * <p>The following example shows the use of {@link CSVFormat#withHeader(String...)} to control the
+ * order and subset of <code>Transaction</code> fields.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions ...
+ * transactions.apply(
+ * CsvIO
+ * .<Transaction>write("path/to/folder/prefix", CSVFormat.DEFAULT.withHeader("transactionId", "purchaseAmount"))
+ * );
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header is repeated for every
+ * shard file, but only including the subset of fields in their listed order.
+ *
+ * <pre>{@code
+ * transactionId,purchaseAmount
+ * 12345,10.23
+ * 54321,54.65
+ * 98765,11.76
+ * }</pre>
+ *
+ * <p>In addition to header customization, {@link CsvIO.Write} supports {@link
+ * CSVFormat#withHeaderComments(Object...)} as shown below. Note that {@link
+ * CSVFormat#withCommentMarker(char)} is required when specifying header comments.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions = ...
+ * transactions
+ * .apply(
+ * CsvIO.<Transaction>write("path/to/folder/prefix",
+ * CSVFormat.DEFAULT.withCommentMarker('#').withHeaderComments("Bank Report", "1970-01-01", "Operator: John Doe")
+ * );
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header and header comments are
+ * repeated for every shard file.
+ *
+ * <pre>{@code
+ * # Bank Report
+ * # 1970-01-01
+ * # Operator: John Doe
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ *
+ * <p>A {@link PCollection} of {@link Row}s works just like custom Java types illustrated above,
+ * except we use {@link CsvIO#writeRows(String, CSVFormat)} as shown below for the same {@code
+ * Transaction} class. We derive {@code Transaction}'s {@link Schema} using a {@link
+ * org.apache.beam.sdk.schemas.annotations.DefaultSchema.DefaultSchemaProvider}. Note that
+ * hard-coding the {@link Row}s below is for illustration purposes. Developers are instead
+ * encouraged to take advantage of {@link
+ * org.apache.beam.sdk.schemas.annotations.DefaultSchema.DefaultSchemaProvider#toRowFunction(TypeDescriptor)}.
+ *
+ * <pre>{@code
+ * DefaultSchemaProvider defaultSchemaProvider = new DefaultSchemaProvider();
+ * Schema schema = defaultSchemaProvider.schemaFor(TypeDescriptor.of(Transaction.class));
+ * PCollection<Row> transactions = pipeline.apply(Create.of(
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "A")
+ * .withFieldValue("purchaseAmount", 10.23)
+ * .withFieldValue("transactionId", "12345")
+ * .build(),
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "B")
+ * .withFieldValue("purchaseAmount", 54.65)
+ * .withFieldValue("transactionId", "54321")
+ * .build(),
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "C")
+ * .withFieldValue("purchaseAmount", 11.76)
+ * .withFieldValue("transactionId", "98765")
+ * .build()
+ * );
+ *
+ * transactions.apply(
+ * CsvIO
+ * .writeRowsTo("gs://bucket/path/to/folder/prefix", CSVFormat.DEFAULT)
+ * );
+ * }</pre>
+ *
+ * <p>Writing the transactions {@link PCollection} of {@link Row}s would yield the following CSV
+ * file content.
+ *
+ * <pre>{@code
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ *
+ * {@link CsvIO.Write} does not support the following {@link CSVFormat} properties and will throw an
+ * {@link IllegalArgumentException}.
+ *
+ * <ul>
+ * <li>{@link CSVFormat#withSkipHeaderRecord()}
+ * <li>{@link CSVFormat#withAllowMissingColumnNames()} ()}
+ * <li>{@link CSVFormat#withAutoFlush(boolean)}
+ * <li>{@link CSVFormat#withIgnoreHeaderCase()}
+ * <li>{@link CSVFormat#withIgnoreSurroundingSpaces()}
+ * </ul>
+ */
+public class CsvIO {
+ /**
+ * The valid {@link Schema.FieldType} from which {@link CsvIO} converts CSV records to the fields.
+ *
+ * <ul>
+ * <li>{@link FieldType#BYTE}
+ * <li>{@link FieldType#BOOLEAN}
+ * <li>{@link FieldType#DATETIME}
+ * <li>{@link FieldType#DECIMAL}
+ * <li>{@link FieldType#DOUBLE}
+ * <li>{@link FieldType#INT16}
+ * <li>{@link FieldType#INT32}
+ * <li>{@link FieldType#INT64}
+ * <li>{@link FieldType#FLOAT}
+ * <li>{@link FieldType#STRING}
+ * </ul>
+ */
+ public static final Set<Schema.FieldType> VALID_FIELD_TYPE_SET =
+ ImmutableSet.of(
+ FieldType.BYTE,
+ FieldType.BOOLEAN,
+ FieldType.DATETIME,
+ FieldType.DECIMAL,
+ FieldType.DOUBLE,
+ FieldType.INT16,
+ FieldType.INT32,
+ FieldType.INT64,
+ FieldType.FLOAT,
+ FieldType.STRING);
+
+ static final String DEFAULT_FILENAME_SUFFIX = ".csv";
+
+ /** Instantiates a {@link Write} for writing user types in {@link CSVFormat} format. */
+ public static <T> Write<T> write(String to, CSVFormat csvFormat) {
+ return new AutoValue_CsvIO_Write.Builder<T>()
+ .setTextIOWrite(createDefaultTextIOWrite(to))
+ .setCSVFormat(csvFormat)
+ .build();
+ }
+
+ /** Instantiates a {@link Write} for writing {@link Row}s in {@link CSVFormat} format. */
+ public static Write<Row> writeRows(String to, CSVFormat csvFormat) {
+ return new AutoValue_CsvIO_Write.Builder<Row>()
+ .setTextIOWrite(createDefaultTextIOWrite(to))
+ .setCSVFormat(csvFormat)
+ .build();
+ }
+
+ /** {@link PTransform} for writing CSV files. */
+ @AutoValue
+ public abstract static class Write<T>
+ extends PTransform<PCollection<T>, WriteFilesResult<String>> {
+
+ /** Specifies the {@link Compression} of all generated shard files. */
+ public Write<T> withCompression(Compression compression) {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withCompression(compression)).build();
+ }
+
+ public Write<T> withNoSpilling() {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withNoSpilling()).build();
+ }
+
+ /**
+ * Specifies to use a given fixed number of shards per window. See {@link
+ * TextIO.Write#withNumShards(int)}.
+ */
+ public Write<T> withNumShards(Integer numShards) {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withNumShards(numShards)).build();
+ }
+
+ /**
+ * Forces a single file as output and empty shard name template. See {@link
+ * TextIO.Write#withoutSharding()}.
+ */
+ public Write<T> withoutSharding() {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withoutSharding()).build();
+ }
+
+ /**
+ * Uses the given {@link ShardNameTemplate} for naming output files. See {@link
+ * TextIO.Write#withShardNameTemplate(String)}.
+ */
+ public Write<T> withShardTemplate(String shardTemplate) {
+ return toBuilder()
+ .setTextIOWrite(getTextIOWrite().withShardNameTemplate(shardTemplate))
+ .build();
+ }
+
+ /**
+ * Configures the filename suffix for written files. See {@link
+ * TextIO.Write#withSuffix(String)}.
+ */
+ public Write<T> withSuffix(String suffix) {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withSuffix(suffix)).build();
+ }
+
+ /**
+ * Set the base directory used to generate temporary files. See {@link
+ * TextIO.Write#withTempDirectory(ResourceId)}.
+ */
+ public Write<T> withTempDirectory(ResourceId tempDirectory) {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withTempDirectory(tempDirectory)).build();
+ }
+
+ /**
+ * Preserves windowing of input elements and writes them to files based on the element's window.
+ * See {@link TextIO.Write#withWindowedWrites()}.
+ */
+ public Write<T> withWindowedWrites() {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withWindowedWrites()).build();
+ }
+
+ /**
+ * Returns a transform for writing to text files like this one but that has the given {@link
+ * FileBasedSink.WritableByteChannelFactory} to be used by the {@link FileBasedSink} during
+ * output. See {@link
+ * TextIO.Write#withWritableByteChannelFactory(FileBasedSink.WritableByteChannelFactory)}.
+ */
+ public Write<T> withWritableByteChannelFactory(
+ FileBasedSink.WritableByteChannelFactory writableByteChannelFactory) {
+ return toBuilder()
+ .setTextIOWrite(
+ getTextIOWrite().withWritableByteChannelFactory(writableByteChannelFactory))
+ .build();
+ }
+
+ /** The underlying {@link FileIO.Write} that writes converted input to CSV formatted output. */
+ abstract TextIO.Write getTextIOWrite();
+
+ /** The {@link CSVFormat} to inform headers, header comments, and format CSV row content. */
+ abstract CSVFormat getCSVFormat();
+
+ abstract Builder<T> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<T> {
+
+ /**
+ * The underlying {@link FileIO.Write} that writes converted input to CSV formatted output.
+ */
+ abstract Builder<T> setTextIOWrite(TextIO.Write value);
+
+ /** The {@link CSVFormat} to convert input. Defaults to {@link CSVFormat#DEFAULT}. */
+ abstract Builder<T> setCSVFormat(CSVFormat value);
+
+ abstract CSVFormat getCSVFormat();
+
+ abstract Write<T> autoBuild();
+
+ final Write<T> build() {
+ checkArgument(
+ !getCSVFormat().getSkipHeaderRecord(),
+ "withSkipHeaderRecord is an illegal CSVFormat setting");
+
+ if (getCSVFormat().getHeaderComments() != null) {
+ checkArgument(
+ getCSVFormat().isCommentMarkerSet(),
+ "CSVFormat withCommentMarker required when withHeaderComments");
+ }
+ return autoBuild();
+ }
+ }
+
+ @Override
+ public WriteFilesResult<String> expand(PCollection<T> input) {
+ if (!input.hasSchema()) {
+ throw new IllegalArgumentException(
+ String.format(
+ "%s requires an input Schema. Note that only Row or user classes are supported. Consider using TextIO or FileIO directly when writing primitive types",
+ Write.class.getName()));
+ }
+
+ Schema schema = input.getSchema();
+
+ RowCoder rowCoder = RowCoder.of(schema);
+
+ PCollection<Row> rows =
+ input
+ .apply("To Rows", MapElements.into(rows()).via(input.getToRowFunction()))
+ .setCoder(rowCoder);
+
+ CSVFormat csvFormat = applyRequiredCSVFormatSettings(schema);
+
+ String header = formatHeader(csvFormat);
+
+ SerializableFunction<Row, String> toCsvFn =
+ CsvRowConversions.RowToCsv.builder()
+ .setCSVFormat(csvFormat)
+ .setSchema(input.getSchema())
+ .build();
+
+ PCollection<String> csv = rows.apply("To CSV", MapElements.into(strings()).via(toCsvFn));
+
+ return csv.apply("Write CSV", getTextIOWrite().withHeader(header).withOutputFilenames());
+ }
+
+ CSVFormat applyRequiredCSVFormatSettings(Schema schema) {
+ CSVFormat csvFormat = getCSVFormat().withSkipHeaderRecord();
+ if (csvFormat.getHeader() == null) {
+ csvFormat = csvFormat.withHeader(schema.sorted().getFieldNames().toArray(new String[0]));
+ }
+ return csvFormat;
+ }
+
+ private static String formatHeader(CSVFormat csvFormat) {
+ String[] header = requireNonNull(csvFormat.getHeader());
+ List<String> result = new ArrayList<>();
+ if (csvFormat.getHeaderComments() != null) {
+ for (String comment : csvFormat.getHeaderComments()) {
+ result.add(csvFormat.getCommentMarker() + " " + comment);
+ }
+ }
+ CSVFormat withoutHeaderComments = csvFormat.withHeaderComments();
+ result.add(withoutHeaderComments.format((Object[]) header));
+ return String.join("\n", result);
+ }
Review Comment:
In response to this comment and hopefully provide more clarity, I've:
1. Renamed `applyRequiredCSVFormatSettings` to [`buildHeaderFromSchemaIfNeeded`](https://github.com/damondouglas/beam/blob/24469-csv-io-write/sdks/java/io/csv/src/main/java/org/apache/beam/sdk/io/csv/CsvIO.java#L426)
2. Refactored and renamed `formatHeader` to [`writeWithCSVFormatHeaderAndComments`](https://github.com/damondouglas/beam/blob/24469-csv-io-write/sdks/java/io/csv/src/main/java/org/apache/beam/sdk/io/csv/CsvIO.java#L436).
Additionally, I wanted to address the questions/comments below:
1. > We are ignoring the case where the user wants to skip the header record and doesn't specify the header. We should only specify the csvFormat header if withSkipHeaderRecord is false.
Responding to the comment and the refactored code suggestion. CSVFormat's withSkipHeaderRecord drives output of a whole line, not what that line contains. It's also a parameter that gives an output incompatible in a Beam context as CSVFormat was designed for writing a single file. Therefore, do you agree then we should not use withSkipHeaderRecord to drive the header contents?
2. > Note that we must also return null from formatHeader in this case since TextIO#withHeader specifically says that a newline is automatically added after.
Apologies, your comment confused me.
Are you saying that the consequence of TextIO#withHeader is:
```
col1, col2, col3
1,2,3
2,3,4
```
vs
```
col1, col2, col3
1,2,3
2,3,4
```
--
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: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [beam] damondouglas commented on a diff in pull request #24630: [24469] Implement CsvIO.Write and supporting classes
Posted by "damondouglas (via GitHub)" <gi...@apache.org>.
damondouglas commented on code in PR #24630:
URL: https://github.com/apache/beam/pull/24630#discussion_r1085768571
##########
sdks/java/io/csv/src/main/java/org/apache/beam/sdk/io/csv/CsvIO.java:
##########
@@ -0,0 +1,446 @@
+/*
+ * 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.sdk.io.csv;
+
+import static java.util.Objects.requireNonNull;
+import static org.apache.beam.sdk.values.TypeDescriptors.rows;
+import static org.apache.beam.sdk.values.TypeDescriptors.strings;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+
+import com.google.auto.value.AutoValue;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Set;
+import org.apache.beam.sdk.coders.RowCoder;
+import org.apache.beam.sdk.io.Compression;
+import org.apache.beam.sdk.io.FileBasedSink;
+import org.apache.beam.sdk.io.FileIO;
+import org.apache.beam.sdk.io.ShardNameTemplate;
+import org.apache.beam.sdk.io.TextIO;
+import org.apache.beam.sdk.io.WriteFilesResult;
+import org.apache.beam.sdk.io.fs.ResourceId;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.transforms.MapElements;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.sdk.values.TypeDescriptor;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableSet;
+import org.apache.commons.csv.CSVFormat;
+
+/**
+ * {@link PTransform}s for reading and writing CSV files.
+ *
+ * <h2>Reading CSV files</h2>
+ *
+ * <p>Reading from CSV files is not yet implemented. Please see <a
+ * href="https://github.com/apache/beam/issues/24552">https://github.com/apache/beam/issues/24552</a>.
+ *
+ * <h2>Writing CSV files</h2>
+ *
+ * <p>To write a {@link PCollection} to one or more CSV files, use {@link CsvIO.Write}, using {@link
+ * CsvIO#writeRows(String, CSVFormat)} or {@link CsvIO#write(String, CSVFormat)}. {@link
+ * CsvIO.Write} supports writing {@link Row} or custom Java types using an inferred {@link Schema}.
+ * Examples below show both scenarios. See the Beam Programming Guide on <a
+ * href="https://beam.apache.org/documentation/programming-guide/#inferring-schemas">inferring
+ * schemas</a> for more information on how to enable Beam to infer a {@link Schema} from a custom
+ * Java type.
+ *
+ * <p>{@link CsvIO.Write} only supports writing {@link Schema} aware types that do not contain any
+ * nested {@link FieldType}s such a {@link org.apache.beam.sdk.schemas.Schema.TypeName#ROW} or
+ * repeated {@link org.apache.beam.sdk.schemas.Schema.TypeName#ARRAY} types. See {@link
+ * CsvIO.Write#VALID_FIELD_TYPE_SET} for valid {@link FieldType}s.
+ *
+ * <h3>Example usage:</h3>
+ *
+ * <p>Suppose we have the following <code>Transaction</code> class annotated with
+ * {@code @DefaultSchema(JavaBeanSchema.class)} so that Beam can infer its {@link Schema}:
+ *
+ * <pre>{@code @DefaultSchema(JavaBeanSchema.class)
+ * public class Transaction {
+ * public Transaction() { … }
+ * public Long getTransactionId();
+ * public void setTransactionId(Long transactionId) { … }
+ * public String getBank() { … }
+ * public void setBank(String bank) { … }
+ * public double getPurchaseAmount() { … }
+ * public void setPurchaseAmount(double purchaseAmount) { … }
+ * }
+ * }</pre>
+ *
+ * <p>From a {@code PCollection<Transaction>}, {@link CsvIO.Write} can write one or many CSV files
+ * automatically creating the header based on its inferred {@link Schema}.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions = ...
+ * transactions.apply(CsvIO.<Transaction>write("path/to/folder/prefix", CSVFormat.DEFAULT));
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header is repeated for every
+ * shard file, whereas by default, {@link CsvIO.Write} will write all fields in <b>sorted order</b>
+ * of the field names.
+ *
+ * <pre>{@code
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ *
+ * <p>To control the order and subset of fields that {@link CsvIO.Write} writes, use {@link
+ * CSVFormat#withHeader(String...)}. Note, however, the following constraints:
+ *
+ * <ol>
+ * <li>Each header column must match a field name in the {@link Schema}; matching is case
+ * sensitive.
+ * <li>Matching header columns must match {@link Schema} fields that are valid {@link FieldType}s;
+ * see {@link #VALID_FIELD_TYPE_SET}.
+ * <li>{@link CSVFormat} only allows repeated header columns when {@link
+ * CSVFormat#withAllowDuplicateHeaderNames()}
+ * </ol>
+ *
+ * <p>The following example shows the use of {@link CSVFormat#withHeader(String...)} to control the
+ * order and subset of <code>Transaction</code> fields.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions ...
+ * transactions.apply(
+ * CsvIO
+ * .<Transaction>write("path/to/folder/prefix", CSVFormat.DEFAULT.withHeader("transactionId", "purchaseAmount"))
+ * );
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header is repeated for every
+ * shard file, but only including the subset of fields in their listed order.
+ *
+ * <pre>{@code
+ * transactionId,purchaseAmount
+ * 12345,10.23
+ * 54321,54.65
+ * 98765,11.76
+ * }</pre>
+ *
+ * <p>In addition to header customization, {@link CsvIO.Write} supports {@link
+ * CSVFormat#withHeaderComments(Object...)} as shown below. Note that {@link
+ * CSVFormat#withCommentMarker(char)} is required when specifying header comments.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions = ...
+ * transactions
+ * .apply(
+ * CsvIO.<Transaction>write("path/to/folder/prefix",
+ * CSVFormat.DEFAULT.withCommentMarker('#').withHeaderComments("Bank Report", "1970-01-01", "Operator: John Doe")
+ * );
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header and header comments are
+ * repeated for every shard file.
+ *
+ * <pre>{@code
+ * # Bank Report
+ * # 1970-01-01
+ * # Operator: John Doe
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ *
+ * <p>A {@link PCollection} of {@link Row}s works just like custom Java types illustrated above,
+ * except we use {@link CsvIO#writeRows(String, CSVFormat)} as shown below for the same {@code
+ * Transaction} class. We derive {@code Transaction}'s {@link Schema} using a {@link
+ * org.apache.beam.sdk.schemas.annotations.DefaultSchema.DefaultSchemaProvider}. Note that
+ * hard-coding the {@link Row}s below is for illustration purposes. Developers are instead
+ * encouraged to take advantage of {@link
+ * org.apache.beam.sdk.schemas.annotations.DefaultSchema.DefaultSchemaProvider#toRowFunction(TypeDescriptor)}.
+ *
+ * <pre>{@code
+ * DefaultSchemaProvider defaultSchemaProvider = new DefaultSchemaProvider();
+ * Schema schema = defaultSchemaProvider.schemaFor(TypeDescriptor.of(Transaction.class));
+ * PCollection<Row> transactions = pipeline.apply(Create.of(
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "A")
+ * .withFieldValue("purchaseAmount", 10.23)
+ * .withFieldValue("transactionId", "12345")
+ * .build(),
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "B")
+ * .withFieldValue("purchaseAmount", 54.65)
+ * .withFieldValue("transactionId", "54321")
+ * .build(),
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "C")
+ * .withFieldValue("purchaseAmount", 11.76)
+ * .withFieldValue("transactionId", "98765")
+ * .build()
+ * );
+ *
+ * transactions.apply(
+ * CsvIO
+ * .writeRowsTo("gs://bucket/path/to/folder/prefix", CSVFormat.DEFAULT)
+ * );
+ * }</pre>
+ *
+ * <p>Writing the transactions {@link PCollection} of {@link Row}s would yield the following CSV
+ * file content.
+ *
+ * <pre>{@code
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ *
+ * {@link CsvIO.Write} does not support the following {@link CSVFormat} properties and will throw an
+ * {@link IllegalArgumentException}.
+ *
+ * <ul>
+ * <li>{@link CSVFormat#withSkipHeaderRecord()}
+ * <li>{@link CSVFormat#withAllowMissingColumnNames()} ()}
+ * <li>{@link CSVFormat#withAutoFlush(boolean)}
+ * <li>{@link CSVFormat#withIgnoreHeaderCase()}
+ * <li>{@link CSVFormat#withIgnoreSurroundingSpaces()}
+ * </ul>
+ */
+public class CsvIO {
+ /**
+ * The valid {@link Schema.FieldType} from which {@link CsvIO} converts CSV records to the fields.
+ *
+ * <ul>
+ * <li>{@link FieldType#BYTE}
+ * <li>{@link FieldType#BOOLEAN}
+ * <li>{@link FieldType#DATETIME}
+ * <li>{@link FieldType#DECIMAL}
+ * <li>{@link FieldType#DOUBLE}
+ * <li>{@link FieldType#INT16}
+ * <li>{@link FieldType#INT32}
+ * <li>{@link FieldType#INT64}
+ * <li>{@link FieldType#FLOAT}
+ * <li>{@link FieldType#STRING}
+ * </ul>
+ */
+ public static final Set<Schema.FieldType> VALID_FIELD_TYPE_SET =
+ ImmutableSet.of(
+ FieldType.BYTE,
+ FieldType.BOOLEAN,
+ FieldType.DATETIME,
+ FieldType.DECIMAL,
+ FieldType.DOUBLE,
+ FieldType.INT16,
+ FieldType.INT32,
+ FieldType.INT64,
+ FieldType.FLOAT,
+ FieldType.STRING);
+
+ static final String DEFAULT_FILENAME_SUFFIX = ".csv";
+
+ /** Instantiates a {@link Write} for writing user types in {@link CSVFormat} format. */
+ public static <T> Write<T> write(String to, CSVFormat csvFormat) {
+ return new AutoValue_CsvIO_Write.Builder<T>()
+ .setTextIOWrite(createDefaultTextIOWrite(to))
+ .setCSVFormat(csvFormat)
+ .build();
+ }
+
+ /** Instantiates a {@link Write} for writing {@link Row}s in {@link CSVFormat} format. */
+ public static Write<Row> writeRows(String to, CSVFormat csvFormat) {
+ return new AutoValue_CsvIO_Write.Builder<Row>()
+ .setTextIOWrite(createDefaultTextIOWrite(to))
+ .setCSVFormat(csvFormat)
+ .build();
+ }
+
+ /** {@link PTransform} for writing CSV files. */
+ @AutoValue
+ public abstract static class Write<T>
+ extends PTransform<PCollection<T>, WriteFilesResult<String>> {
+
+ /** Specifies the {@link Compression} of all generated shard files. */
+ public Write<T> withCompression(Compression compression) {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withCompression(compression)).build();
+ }
+
+ public Write<T> withNoSpilling() {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withNoSpilling()).build();
+ }
+
+ /**
+ * Specifies to use a given fixed number of shards per window. See {@link
+ * TextIO.Write#withNumShards(int)}.
+ */
+ public Write<T> withNumShards(Integer numShards) {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withNumShards(numShards)).build();
+ }
+
+ /**
+ * Forces a single file as output and empty shard name template. See {@link
+ * TextIO.Write#withoutSharding()}.
+ */
+ public Write<T> withoutSharding() {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withoutSharding()).build();
+ }
+
+ /**
+ * Uses the given {@link ShardNameTemplate} for naming output files. See {@link
+ * TextIO.Write#withShardNameTemplate(String)}.
+ */
+ public Write<T> withShardTemplate(String shardTemplate) {
+ return toBuilder()
+ .setTextIOWrite(getTextIOWrite().withShardNameTemplate(shardTemplate))
+ .build();
+ }
+
+ /**
+ * Configures the filename suffix for written files. See {@link
+ * TextIO.Write#withSuffix(String)}.
+ */
+ public Write<T> withSuffix(String suffix) {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withSuffix(suffix)).build();
+ }
+
+ /**
+ * Set the base directory used to generate temporary files. See {@link
+ * TextIO.Write#withTempDirectory(ResourceId)}.
+ */
+ public Write<T> withTempDirectory(ResourceId tempDirectory) {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withTempDirectory(tempDirectory)).build();
+ }
+
+ /**
+ * Preserves windowing of input elements and writes them to files based on the element's window.
+ * See {@link TextIO.Write#withWindowedWrites()}.
+ */
+ public Write<T> withWindowedWrites() {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withWindowedWrites()).build();
+ }
+
+ /**
+ * Returns a transform for writing to text files like this one but that has the given {@link
+ * FileBasedSink.WritableByteChannelFactory} to be used by the {@link FileBasedSink} during
+ * output. See {@link
+ * TextIO.Write#withWritableByteChannelFactory(FileBasedSink.WritableByteChannelFactory)}.
+ */
+ public Write<T> withWritableByteChannelFactory(
+ FileBasedSink.WritableByteChannelFactory writableByteChannelFactory) {
+ return toBuilder()
+ .setTextIOWrite(
+ getTextIOWrite().withWritableByteChannelFactory(writableByteChannelFactory))
+ .build();
+ }
+
+ /** The underlying {@link FileIO.Write} that writes converted input to CSV formatted output. */
+ abstract TextIO.Write getTextIOWrite();
+
+ /** The {@link CSVFormat} to inform headers, header comments, and format CSV row content. */
+ abstract CSVFormat getCSVFormat();
+
+ abstract Builder<T> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<T> {
+
+ /**
+ * The underlying {@link FileIO.Write} that writes converted input to CSV formatted output.
+ */
+ abstract Builder<T> setTextIOWrite(TextIO.Write value);
+
+ /** The {@link CSVFormat} to convert input. Defaults to {@link CSVFormat#DEFAULT}. */
+ abstract Builder<T> setCSVFormat(CSVFormat value);
+
+ abstract CSVFormat getCSVFormat();
+
+ abstract Write<T> autoBuild();
+
+ final Write<T> build() {
+ checkArgument(
+ !getCSVFormat().getSkipHeaderRecord(),
+ "withSkipHeaderRecord is an illegal CSVFormat setting");
+
+ if (getCSVFormat().getHeaderComments() != null) {
+ checkArgument(
+ getCSVFormat().isCommentMarkerSet(),
+ "CSVFormat withCommentMarker required when withHeaderComments");
+ }
+ return autoBuild();
+ }
+ }
+
+ @Override
+ public WriteFilesResult<String> expand(PCollection<T> input) {
+ if (!input.hasSchema()) {
+ throw new IllegalArgumentException(
+ String.format(
+ "%s requires an input Schema. Note that only Row or user classes are supported. Consider using TextIO or FileIO directly when writing primitive types",
+ Write.class.getName()));
+ }
+
+ Schema schema = input.getSchema();
+
+ RowCoder rowCoder = RowCoder.of(schema);
+
+ PCollection<Row> rows =
+ input
+ .apply("To Rows", MapElements.into(rows()).via(input.getToRowFunction()))
+ .setCoder(rowCoder);
+
+ CSVFormat csvFormat = applyRequiredCSVFormatSettings(schema);
+
+ String header = formatHeader(csvFormat);
+
+ SerializableFunction<Row, String> toCsvFn =
+ CsvRowConversions.RowToCsv.builder()
+ .setCSVFormat(csvFormat)
+ .setSchema(input.getSchema())
+ .build();
+
+ PCollection<String> csv = rows.apply("To CSV", MapElements.into(strings()).via(toCsvFn));
+
+ return csv.apply("Write CSV", getTextIOWrite().withHeader(header).withOutputFilenames());
+ }
+
+ CSVFormat applyRequiredCSVFormatSettings(Schema schema) {
+ CSVFormat csvFormat = getCSVFormat().withSkipHeaderRecord();
+ if (csvFormat.getHeader() == null) {
+ csvFormat = csvFormat.withHeader(schema.sorted().getFieldNames().toArray(new String[0]));
+ }
+ return csvFormat;
+ }
+
+ private static String formatHeader(CSVFormat csvFormat) {
+ String[] header = requireNonNull(csvFormat.getHeader());
+ List<String> result = new ArrayList<>();
+ if (csvFormat.getHeaderComments() != null) {
+ for (String comment : csvFormat.getHeaderComments()) {
+ result.add(csvFormat.getCommentMarker() + " " + comment);
+ }
+ }
+ CSVFormat withoutHeaderComments = csvFormat.withHeaderComments();
+ result.add(withoutHeaderComments.format((Object[]) header));
+ return String.join("\n", result);
+ }
Review Comment:
I would like to anchor our discussion on CSV header-related software behavior using the following:
```gherkin
Feature: write a CSV header
Scenario: User provides CSVFormat.withHeader
Given CSVFormat is not null
When CsvIO writes PCollection<Row> to CSV formatted file shards
Then for each shard the CSVFormat withHeader appears as the first line
Scenario: User does not provide CSVFormat.withHeader
Given CSVFormat is null
When CsvIO writes PCollection<Row> to CSV formatted file shards
Then for each shard the Schema field names in sorted order appears as the first line
Scenario: User provides CSVFormat.withSkipHeaderRecord == true
Given CSVFormat.withSkipHeaderRecord is true
When the user applies the CsvIO.Write PTransform
Then they see an error that this is an illegal setting
```
--
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: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [beam] damondouglas commented on a diff in pull request #24630: [24469] Implement CsvIO.Write and supporting classes
Posted by "damondouglas (via GitHub)" <gi...@apache.org>.
damondouglas commented on code in PR #24630:
URL: https://github.com/apache/beam/pull/24630#discussion_r1085698752
##########
sdks/java/io/csv/src/main/java/org/apache/beam/sdk/io/csv/CsvRowConversions.java:
##########
@@ -0,0 +1,148 @@
+/*
+ * 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.sdk.io.csv;
+
+import static org.apache.beam.sdk.io.csv.CsvIO.VALID_FIELD_TYPE_SET;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkNotNull;
+
+import com.google.auto.value.AutoValue;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Set;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.schemas.Schema.TypeName;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.values.Row;
+import org.apache.commons.csv.CSVFormat;
+import org.checkerframework.checker.nullness.qual.NonNull;
+
+/** Contains classes and methods to help with converting between {@link Row} and CSV strings. */
+class CsvRowConversions {
+
+ /** Converts between {@link Row} and CSV string using a {@link CSVFormat}. */
+ @AutoValue
+ abstract static class RowToCsv implements SerializableFunction<Row, String> {
+
+ static Builder builder() {
+ return new AutoValue_CsvRowConversions_RowToCsv.Builder();
+ }
+
+ /** The expected {@link Schema} of the {@link Row} input. */
+ abstract Schema getSchema();
+
+ /** The {@link CSVFormat} of the converted {@link Row} input. */
+ abstract CSVFormat getCSVFormat();
+
+ /** Converts a {@link Row} to a CSV string formatted using {@link #getCSVFormat()}. */
+ @Override
+ public String apply(Row input) {
+ Row safeInput = checkNotNull(input);
+ String[] header = getHeader();
+ Object[] values = new Object[header.length];
+ for (int i = 0; i < header.length; i++) {
+ values[i] = safeInput.getValue(header[i]);
+ }
+ return getCSVFormat().format(values);
+ }
+
+ @NonNull
+ String[] getHeader() {
+ return checkNotNull(getCSVFormat().getHeader());
+ }
+
+ @AutoValue.Builder
+ abstract static class Builder {
+
+ /** The expected {@link Schema} of the {@link Row} input. */
+ abstract Builder setSchema(Schema schema);
+
+ abstract Schema getSchema();
+
+ /** The {@link CSVFormat} of the converted {@link Row} input. */
+ abstract Builder setCSVFormat(CSVFormat format);
+
+ abstract CSVFormat getCSVFormat();
+
+ abstract RowToCsv autoBuild();
+
+ final RowToCsv build() {
+ checkArgument(getSchema().getFieldCount() > 0, "Schema has no fields");
+ setCSVFormat(
+ getCSVFormat()
+ .withSkipHeaderRecord()
Review Comment:
I wanted to keep CsvRowConversions and CsvIO independent. In the setting of CsvRowConversions.RowToCSV didn't want to assume anyone would remember to set withSkipHeaderRecord. Also, for a package-private class, I didn't want to throw an error either for a class that has nothing to do with the header other than lookups against the Schema and Row values. The withSkipHeaderRecord is an output parameter and a hack to still use CSVFormat which was designed to write single files.
--
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: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [beam] github-actions[bot] commented on pull request #24630: [24469] Implement CsvIO.Write and supporting classes
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on PR #24630:
URL: https://github.com/apache/beam/pull/24630#issuecomment-1378660518
Reminder, please take a look at this pr: @kennknowles @damccorm @johnjcasey
--
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: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [beam] damondouglas commented on a diff in pull request #24630: [24469] Implement CsvIO.Write and supporting classes
Posted by "damondouglas (via GitHub)" <gi...@apache.org>.
damondouglas commented on code in PR #24630:
URL: https://github.com/apache/beam/pull/24630#discussion_r1085716045
##########
sdks/java/io/csv/src/main/java/org/apache/beam/sdk/io/csv/CsvIO.java:
##########
@@ -0,0 +1,446 @@
+/*
+ * 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.sdk.io.csv;
+
+import static java.util.Objects.requireNonNull;
+import static org.apache.beam.sdk.values.TypeDescriptors.rows;
+import static org.apache.beam.sdk.values.TypeDescriptors.strings;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+
+import com.google.auto.value.AutoValue;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Set;
+import org.apache.beam.sdk.coders.RowCoder;
+import org.apache.beam.sdk.io.Compression;
+import org.apache.beam.sdk.io.FileBasedSink;
+import org.apache.beam.sdk.io.FileIO;
+import org.apache.beam.sdk.io.ShardNameTemplate;
+import org.apache.beam.sdk.io.TextIO;
+import org.apache.beam.sdk.io.WriteFilesResult;
+import org.apache.beam.sdk.io.fs.ResourceId;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.transforms.MapElements;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.sdk.values.TypeDescriptor;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableSet;
+import org.apache.commons.csv.CSVFormat;
+
+/**
+ * {@link PTransform}s for reading and writing CSV files.
+ *
+ * <h2>Reading CSV files</h2>
+ *
+ * <p>Reading from CSV files is not yet implemented. Please see <a
+ * href="https://github.com/apache/beam/issues/24552">https://github.com/apache/beam/issues/24552</a>.
+ *
+ * <h2>Writing CSV files</h2>
+ *
+ * <p>To write a {@link PCollection} to one or more CSV files, use {@link CsvIO.Write}, using {@link
+ * CsvIO#writeRows(String, CSVFormat)} or {@link CsvIO#write(String, CSVFormat)}. {@link
+ * CsvIO.Write} supports writing {@link Row} or custom Java types using an inferred {@link Schema}.
+ * Examples below show both scenarios. See the Beam Programming Guide on <a
+ * href="https://beam.apache.org/documentation/programming-guide/#inferring-schemas">inferring
+ * schemas</a> for more information on how to enable Beam to infer a {@link Schema} from a custom
+ * Java type.
+ *
+ * <p>{@link CsvIO.Write} only supports writing {@link Schema} aware types that do not contain any
+ * nested {@link FieldType}s such a {@link org.apache.beam.sdk.schemas.Schema.TypeName#ROW} or
+ * repeated {@link org.apache.beam.sdk.schemas.Schema.TypeName#ARRAY} types. See {@link
+ * CsvIO.Write#VALID_FIELD_TYPE_SET} for valid {@link FieldType}s.
+ *
+ * <h3>Example usage:</h3>
+ *
+ * <p>Suppose we have the following <code>Transaction</code> class annotated with
+ * {@code @DefaultSchema(JavaBeanSchema.class)} so that Beam can infer its {@link Schema}:
+ *
+ * <pre>{@code @DefaultSchema(JavaBeanSchema.class)
+ * public class Transaction {
+ * public Transaction() { … }
+ * public Long getTransactionId();
+ * public void setTransactionId(Long transactionId) { … }
+ * public String getBank() { … }
+ * public void setBank(String bank) { … }
+ * public double getPurchaseAmount() { … }
+ * public void setPurchaseAmount(double purchaseAmount) { … }
+ * }
+ * }</pre>
+ *
+ * <p>From a {@code PCollection<Transaction>}, {@link CsvIO.Write} can write one or many CSV files
+ * automatically creating the header based on its inferred {@link Schema}.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions = ...
+ * transactions.apply(CsvIO.<Transaction>write("path/to/folder/prefix", CSVFormat.DEFAULT));
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header is repeated for every
+ * shard file, whereas by default, {@link CsvIO.Write} will write all fields in <b>sorted order</b>
+ * of the field names.
+ *
+ * <pre>{@code
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ *
+ * <p>To control the order and subset of fields that {@link CsvIO.Write} writes, use {@link
+ * CSVFormat#withHeader(String...)}. Note, however, the following constraints:
+ *
+ * <ol>
+ * <li>Each header column must match a field name in the {@link Schema}; matching is case
+ * sensitive.
+ * <li>Matching header columns must match {@link Schema} fields that are valid {@link FieldType}s;
+ * see {@link #VALID_FIELD_TYPE_SET}.
+ * <li>{@link CSVFormat} only allows repeated header columns when {@link
+ * CSVFormat#withAllowDuplicateHeaderNames()}
+ * </ol>
+ *
+ * <p>The following example shows the use of {@link CSVFormat#withHeader(String...)} to control the
+ * order and subset of <code>Transaction</code> fields.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions ...
+ * transactions.apply(
+ * CsvIO
+ * .<Transaction>write("path/to/folder/prefix", CSVFormat.DEFAULT.withHeader("transactionId", "purchaseAmount"))
+ * );
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header is repeated for every
+ * shard file, but only including the subset of fields in their listed order.
+ *
+ * <pre>{@code
+ * transactionId,purchaseAmount
+ * 12345,10.23
+ * 54321,54.65
+ * 98765,11.76
+ * }</pre>
+ *
+ * <p>In addition to header customization, {@link CsvIO.Write} supports {@link
+ * CSVFormat#withHeaderComments(Object...)} as shown below. Note that {@link
+ * CSVFormat#withCommentMarker(char)} is required when specifying header comments.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions = ...
+ * transactions
+ * .apply(
+ * CsvIO.<Transaction>write("path/to/folder/prefix",
+ * CSVFormat.DEFAULT.withCommentMarker('#').withHeaderComments("Bank Report", "1970-01-01", "Operator: John Doe")
+ * );
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header and header comments are
+ * repeated for every shard file.
+ *
+ * <pre>{@code
+ * # Bank Report
+ * # 1970-01-01
+ * # Operator: John Doe
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ *
+ * <p>A {@link PCollection} of {@link Row}s works just like custom Java types illustrated above,
+ * except we use {@link CsvIO#writeRows(String, CSVFormat)} as shown below for the same {@code
+ * Transaction} class. We derive {@code Transaction}'s {@link Schema} using a {@link
+ * org.apache.beam.sdk.schemas.annotations.DefaultSchema.DefaultSchemaProvider}. Note that
+ * hard-coding the {@link Row}s below is for illustration purposes. Developers are instead
+ * encouraged to take advantage of {@link
+ * org.apache.beam.sdk.schemas.annotations.DefaultSchema.DefaultSchemaProvider#toRowFunction(TypeDescriptor)}.
+ *
+ * <pre>{@code
+ * DefaultSchemaProvider defaultSchemaProvider = new DefaultSchemaProvider();
+ * Schema schema = defaultSchemaProvider.schemaFor(TypeDescriptor.of(Transaction.class));
+ * PCollection<Row> transactions = pipeline.apply(Create.of(
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "A")
+ * .withFieldValue("purchaseAmount", 10.23)
+ * .withFieldValue("transactionId", "12345")
+ * .build(),
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "B")
+ * .withFieldValue("purchaseAmount", 54.65)
+ * .withFieldValue("transactionId", "54321")
+ * .build(),
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "C")
+ * .withFieldValue("purchaseAmount", 11.76)
+ * .withFieldValue("transactionId", "98765")
+ * .build()
+ * );
+ *
+ * transactions.apply(
+ * CsvIO
+ * .writeRowsTo("gs://bucket/path/to/folder/prefix", CSVFormat.DEFAULT)
+ * );
+ * }</pre>
+ *
+ * <p>Writing the transactions {@link PCollection} of {@link Row}s would yield the following CSV
+ * file content.
+ *
+ * <pre>{@code
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ *
+ * {@link CsvIO.Write} does not support the following {@link CSVFormat} properties and will throw an
+ * {@link IllegalArgumentException}.
+ *
+ * <ul>
+ * <li>{@link CSVFormat#withSkipHeaderRecord()}
+ * <li>{@link CSVFormat#withAllowMissingColumnNames()} ()}
+ * <li>{@link CSVFormat#withAutoFlush(boolean)}
+ * <li>{@link CSVFormat#withIgnoreHeaderCase()}
+ * <li>{@link CSVFormat#withIgnoreSurroundingSpaces()}
+ * </ul>
+ */
+public class CsvIO {
+ /**
+ * The valid {@link Schema.FieldType} from which {@link CsvIO} converts CSV records to the fields.
+ *
+ * <ul>
+ * <li>{@link FieldType#BYTE}
+ * <li>{@link FieldType#BOOLEAN}
+ * <li>{@link FieldType#DATETIME}
+ * <li>{@link FieldType#DECIMAL}
+ * <li>{@link FieldType#DOUBLE}
+ * <li>{@link FieldType#INT16}
+ * <li>{@link FieldType#INT32}
+ * <li>{@link FieldType#INT64}
+ * <li>{@link FieldType#FLOAT}
+ * <li>{@link FieldType#STRING}
+ * </ul>
+ */
+ public static final Set<Schema.FieldType> VALID_FIELD_TYPE_SET =
+ ImmutableSet.of(
+ FieldType.BYTE,
+ FieldType.BOOLEAN,
+ FieldType.DATETIME,
+ FieldType.DECIMAL,
+ FieldType.DOUBLE,
+ FieldType.INT16,
+ FieldType.INT32,
+ FieldType.INT64,
+ FieldType.FLOAT,
+ FieldType.STRING);
+
+ static final String DEFAULT_FILENAME_SUFFIX = ".csv";
+
+ /** Instantiates a {@link Write} for writing user types in {@link CSVFormat} format. */
+ public static <T> Write<T> write(String to, CSVFormat csvFormat) {
+ return new AutoValue_CsvIO_Write.Builder<T>()
+ .setTextIOWrite(createDefaultTextIOWrite(to))
+ .setCSVFormat(csvFormat)
+ .build();
+ }
+
+ /** Instantiates a {@link Write} for writing {@link Row}s in {@link CSVFormat} format. */
+ public static Write<Row> writeRows(String to, CSVFormat csvFormat) {
+ return new AutoValue_CsvIO_Write.Builder<Row>()
+ .setTextIOWrite(createDefaultTextIOWrite(to))
+ .setCSVFormat(csvFormat)
+ .build();
+ }
+
+ /** {@link PTransform} for writing CSV files. */
+ @AutoValue
+ public abstract static class Write<T>
+ extends PTransform<PCollection<T>, WriteFilesResult<String>> {
+
+ /** Specifies the {@link Compression} of all generated shard files. */
+ public Write<T> withCompression(Compression compression) {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withCompression(compression)).build();
+ }
+
+ public Write<T> withNoSpilling() {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withNoSpilling()).build();
+ }
+
+ /**
+ * Specifies to use a given fixed number of shards per window. See {@link
+ * TextIO.Write#withNumShards(int)}.
+ */
+ public Write<T> withNumShards(Integer numShards) {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withNumShards(numShards)).build();
+ }
+
+ /**
+ * Forces a single file as output and empty shard name template. See {@link
+ * TextIO.Write#withoutSharding()}.
+ */
+ public Write<T> withoutSharding() {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withoutSharding()).build();
+ }
+
+ /**
+ * Uses the given {@link ShardNameTemplate} for naming output files. See {@link
+ * TextIO.Write#withShardNameTemplate(String)}.
+ */
+ public Write<T> withShardTemplate(String shardTemplate) {
+ return toBuilder()
+ .setTextIOWrite(getTextIOWrite().withShardNameTemplate(shardTemplate))
+ .build();
+ }
+
+ /**
+ * Configures the filename suffix for written files. See {@link
+ * TextIO.Write#withSuffix(String)}.
+ */
+ public Write<T> withSuffix(String suffix) {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withSuffix(suffix)).build();
+ }
+
+ /**
+ * Set the base directory used to generate temporary files. See {@link
+ * TextIO.Write#withTempDirectory(ResourceId)}.
+ */
+ public Write<T> withTempDirectory(ResourceId tempDirectory) {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withTempDirectory(tempDirectory)).build();
+ }
+
+ /**
+ * Preserves windowing of input elements and writes them to files based on the element's window.
+ * See {@link TextIO.Write#withWindowedWrites()}.
+ */
+ public Write<T> withWindowedWrites() {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withWindowedWrites()).build();
+ }
+
+ /**
+ * Returns a transform for writing to text files like this one but that has the given {@link
+ * FileBasedSink.WritableByteChannelFactory} to be used by the {@link FileBasedSink} during
+ * output. See {@link
+ * TextIO.Write#withWritableByteChannelFactory(FileBasedSink.WritableByteChannelFactory)}.
+ */
+ public Write<T> withWritableByteChannelFactory(
+ FileBasedSink.WritableByteChannelFactory writableByteChannelFactory) {
+ return toBuilder()
+ .setTextIOWrite(
+ getTextIOWrite().withWritableByteChannelFactory(writableByteChannelFactory))
+ .build();
+ }
+
+ /** The underlying {@link FileIO.Write} that writes converted input to CSV formatted output. */
+ abstract TextIO.Write getTextIOWrite();
+
+ /** The {@link CSVFormat} to inform headers, header comments, and format CSV row content. */
+ abstract CSVFormat getCSVFormat();
+
+ abstract Builder<T> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<T> {
+
+ /**
+ * The underlying {@link FileIO.Write} that writes converted input to CSV formatted output.
+ */
+ abstract Builder<T> setTextIOWrite(TextIO.Write value);
+
+ /** The {@link CSVFormat} to convert input. Defaults to {@link CSVFormat#DEFAULT}. */
+ abstract Builder<T> setCSVFormat(CSVFormat value);
+
+ abstract CSVFormat getCSVFormat();
+
+ abstract Write<T> autoBuild();
+
+ final Write<T> build() {
+ checkArgument(
+ !getCSVFormat().getSkipHeaderRecord(),
+ "withSkipHeaderRecord is an illegal CSVFormat setting");
+
+ if (getCSVFormat().getHeaderComments() != null) {
+ checkArgument(
+ getCSVFormat().isCommentMarkerSet(),
+ "CSVFormat withCommentMarker required when withHeaderComments");
+ }
+ return autoBuild();
+ }
+ }
+
+ @Override
+ public WriteFilesResult<String> expand(PCollection<T> input) {
+ if (!input.hasSchema()) {
+ throw new IllegalArgumentException(
+ String.format(
+ "%s requires an input Schema. Note that only Row or user classes are supported. Consider using TextIO or FileIO directly when writing primitive types",
+ Write.class.getName()));
+ }
+
+ Schema schema = input.getSchema();
+
+ RowCoder rowCoder = RowCoder.of(schema);
+
+ PCollection<Row> rows =
+ input
+ .apply("To Rows", MapElements.into(rows()).via(input.getToRowFunction()))
+ .setCoder(rowCoder);
+
+ CSVFormat csvFormat = applyRequiredCSVFormatSettings(schema);
+
+ String header = formatHeader(csvFormat);
+
+ SerializableFunction<Row, String> toCsvFn =
+ CsvRowConversions.RowToCsv.builder()
+ .setCSVFormat(csvFormat)
+ .setSchema(input.getSchema())
+ .build();
+
+ PCollection<String> csv = rows.apply("To CSV", MapElements.into(strings()).via(toCsvFn));
+
+ return csv.apply("Write CSV", getTextIOWrite().withHeader(header).withOutputFilenames());
+ }
+
+ CSVFormat applyRequiredCSVFormatSettings(Schema schema) {
+ CSVFormat csvFormat = getCSVFormat().withSkipHeaderRecord();
+ if (csvFormat.getHeader() == null) {
+ csvFormat = csvFormat.withHeader(schema.sorted().getFieldNames().toArray(new String[0]));
+ }
+ return csvFormat;
+ }
+
+ private static String formatHeader(CSVFormat csvFormat) {
+ String[] header = requireNonNull(csvFormat.getHeader());
+ List<String> result = new ArrayList<>();
+ if (csvFormat.getHeaderComments() != null) {
+ for (String comment : csvFormat.getHeaderComments()) {
+ result.add(csvFormat.getCommentMarker() + " " + comment);
+ }
+ }
+ CSVFormat withoutHeaderComments = csvFormat.withHeaderComments();
+ result.add(withoutHeaderComments.format((Object[]) header));
+ return String.join("\n", result);
+ }
Review Comment:
Overall, I changed this method to incorporate applying to TextIO.Write.withHeader to provide more clarity.
> We are ignoring the case where the user wants to skip the header record and doesn't specify the header. We should only specify the csvFormat header if withSkipHeaderRecord is false.
CSVFormat communicates that withSkipHeaderRecord functions to not output the header. Do you agree then we should keep withSkipHeaderRecord separate from where we are deriving the header?
> Note that we must also return null from formatHeader in this case since TextIO#withHeader specifically says that a newline is automatically added after.
Are you saying that the consequence of TextIO#withHeader is:
```
col1, col2, col3
1,2,3
2,3,4
```
vs
```
col1, col2, col3
1,2,3
2,3,4
```
--
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: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [beam] damondouglas commented on a diff in pull request #24630: [24469] Implement CsvIO.Write and supporting classes
Posted by "damondouglas (via GitHub)" <gi...@apache.org>.
damondouglas commented on code in PR #24630:
URL: https://github.com/apache/beam/pull/24630#discussion_r1085693708
##########
sdks/java/io/csv/src/main/java/org/apache/beam/sdk/io/csv/CsvIO.java:
##########
@@ -0,0 +1,446 @@
+/*
+ * 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.sdk.io.csv;
+
+import static java.util.Objects.requireNonNull;
+import static org.apache.beam.sdk.values.TypeDescriptors.rows;
+import static org.apache.beam.sdk.values.TypeDescriptors.strings;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+
+import com.google.auto.value.AutoValue;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Set;
+import org.apache.beam.sdk.coders.RowCoder;
+import org.apache.beam.sdk.io.Compression;
+import org.apache.beam.sdk.io.FileBasedSink;
+import org.apache.beam.sdk.io.FileIO;
+import org.apache.beam.sdk.io.ShardNameTemplate;
+import org.apache.beam.sdk.io.TextIO;
+import org.apache.beam.sdk.io.WriteFilesResult;
+import org.apache.beam.sdk.io.fs.ResourceId;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.transforms.MapElements;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.sdk.values.TypeDescriptor;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableSet;
+import org.apache.commons.csv.CSVFormat;
+
+/**
+ * {@link PTransform}s for reading and writing CSV files.
+ *
+ * <h2>Reading CSV files</h2>
+ *
+ * <p>Reading from CSV files is not yet implemented. Please see <a
+ * href="https://github.com/apache/beam/issues/24552">https://github.com/apache/beam/issues/24552</a>.
+ *
+ * <h2>Writing CSV files</h2>
+ *
+ * <p>To write a {@link PCollection} to one or more CSV files, use {@link CsvIO.Write}, using {@link
+ * CsvIO#writeRows(String, CSVFormat)} or {@link CsvIO#write(String, CSVFormat)}. {@link
+ * CsvIO.Write} supports writing {@link Row} or custom Java types using an inferred {@link Schema}.
+ * Examples below show both scenarios. See the Beam Programming Guide on <a
+ * href="https://beam.apache.org/documentation/programming-guide/#inferring-schemas">inferring
+ * schemas</a> for more information on how to enable Beam to infer a {@link Schema} from a custom
+ * Java type.
+ *
+ * <p>{@link CsvIO.Write} only supports writing {@link Schema} aware types that do not contain any
+ * nested {@link FieldType}s such a {@link org.apache.beam.sdk.schemas.Schema.TypeName#ROW} or
+ * repeated {@link org.apache.beam.sdk.schemas.Schema.TypeName#ARRAY} types. See {@link
+ * CsvIO.Write#VALID_FIELD_TYPE_SET} for valid {@link FieldType}s.
+ *
+ * <h3>Example usage:</h3>
+ *
+ * <p>Suppose we have the following <code>Transaction</code> class annotated with
+ * {@code @DefaultSchema(JavaBeanSchema.class)} so that Beam can infer its {@link Schema}:
+ *
+ * <pre>{@code @DefaultSchema(JavaBeanSchema.class)
+ * public class Transaction {
+ * public Transaction() { … }
+ * public Long getTransactionId();
+ * public void setTransactionId(Long transactionId) { … }
+ * public String getBank() { … }
+ * public void setBank(String bank) { … }
+ * public double getPurchaseAmount() { … }
+ * public void setPurchaseAmount(double purchaseAmount) { … }
+ * }
+ * }</pre>
+ *
+ * <p>From a {@code PCollection<Transaction>}, {@link CsvIO.Write} can write one or many CSV files
+ * automatically creating the header based on its inferred {@link Schema}.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions = ...
+ * transactions.apply(CsvIO.<Transaction>write("path/to/folder/prefix", CSVFormat.DEFAULT));
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header is repeated for every
+ * shard file, whereas by default, {@link CsvIO.Write} will write all fields in <b>sorted order</b>
+ * of the field names.
+ *
+ * <pre>{@code
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ *
+ * <p>To control the order and subset of fields that {@link CsvIO.Write} writes, use {@link
+ * CSVFormat#withHeader(String...)}. Note, however, the following constraints:
+ *
+ * <ol>
+ * <li>Each header column must match a field name in the {@link Schema}; matching is case
+ * sensitive.
+ * <li>Matching header columns must match {@link Schema} fields that are valid {@link FieldType}s;
+ * see {@link #VALID_FIELD_TYPE_SET}.
+ * <li>{@link CSVFormat} only allows repeated header columns when {@link
+ * CSVFormat#withAllowDuplicateHeaderNames()}
+ * </ol>
+ *
+ * <p>The following example shows the use of {@link CSVFormat#withHeader(String...)} to control the
+ * order and subset of <code>Transaction</code> fields.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions ...
+ * transactions.apply(
+ * CsvIO
+ * .<Transaction>write("path/to/folder/prefix", CSVFormat.DEFAULT.withHeader("transactionId", "purchaseAmount"))
+ * );
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header is repeated for every
+ * shard file, but only including the subset of fields in their listed order.
+ *
+ * <pre>{@code
+ * transactionId,purchaseAmount
+ * 12345,10.23
+ * 54321,54.65
+ * 98765,11.76
+ * }</pre>
+ *
+ * <p>In addition to header customization, {@link CsvIO.Write} supports {@link
+ * CSVFormat#withHeaderComments(Object...)} as shown below. Note that {@link
+ * CSVFormat#withCommentMarker(char)} is required when specifying header comments.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions = ...
+ * transactions
+ * .apply(
+ * CsvIO.<Transaction>write("path/to/folder/prefix",
+ * CSVFormat.DEFAULT.withCommentMarker('#').withHeaderComments("Bank Report", "1970-01-01", "Operator: John Doe")
+ * );
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header and header comments are
+ * repeated for every shard file.
+ *
+ * <pre>{@code
+ * # Bank Report
+ * # 1970-01-01
+ * # Operator: John Doe
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ *
+ * <p>A {@link PCollection} of {@link Row}s works just like custom Java types illustrated above,
+ * except we use {@link CsvIO#writeRows(String, CSVFormat)} as shown below for the same {@code
+ * Transaction} class. We derive {@code Transaction}'s {@link Schema} using a {@link
+ * org.apache.beam.sdk.schemas.annotations.DefaultSchema.DefaultSchemaProvider}. Note that
+ * hard-coding the {@link Row}s below is for illustration purposes. Developers are instead
+ * encouraged to take advantage of {@link
+ * org.apache.beam.sdk.schemas.annotations.DefaultSchema.DefaultSchemaProvider#toRowFunction(TypeDescriptor)}.
+ *
+ * <pre>{@code
+ * DefaultSchemaProvider defaultSchemaProvider = new DefaultSchemaProvider();
+ * Schema schema = defaultSchemaProvider.schemaFor(TypeDescriptor.of(Transaction.class));
+ * PCollection<Row> transactions = pipeline.apply(Create.of(
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "A")
+ * .withFieldValue("purchaseAmount", 10.23)
+ * .withFieldValue("transactionId", "12345")
+ * .build(),
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "B")
+ * .withFieldValue("purchaseAmount", 54.65)
+ * .withFieldValue("transactionId", "54321")
+ * .build(),
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "C")
+ * .withFieldValue("purchaseAmount", 11.76)
+ * .withFieldValue("transactionId", "98765")
+ * .build()
+ * );
+ *
+ * transactions.apply(
+ * CsvIO
+ * .writeRowsTo("gs://bucket/path/to/folder/prefix", CSVFormat.DEFAULT)
+ * );
+ * }</pre>
+ *
+ * <p>Writing the transactions {@link PCollection} of {@link Row}s would yield the following CSV
+ * file content.
+ *
+ * <pre>{@code
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ *
+ * {@link CsvIO.Write} does not support the following {@link CSVFormat} properties and will throw an
+ * {@link IllegalArgumentException}.
+ *
+ * <ul>
+ * <li>{@link CSVFormat#withSkipHeaderRecord()}
+ * <li>{@link CSVFormat#withAllowMissingColumnNames()} ()}
+ * <li>{@link CSVFormat#withAutoFlush(boolean)}
+ * <li>{@link CSVFormat#withIgnoreHeaderCase()}
+ * <li>{@link CSVFormat#withIgnoreSurroundingSpaces()}
+ * </ul>
+ */
+public class CsvIO {
+ /**
+ * The valid {@link Schema.FieldType} from which {@link CsvIO} converts CSV records to the fields.
+ *
+ * <ul>
+ * <li>{@link FieldType#BYTE}
+ * <li>{@link FieldType#BOOLEAN}
+ * <li>{@link FieldType#DATETIME}
+ * <li>{@link FieldType#DECIMAL}
+ * <li>{@link FieldType#DOUBLE}
+ * <li>{@link FieldType#INT16}
+ * <li>{@link FieldType#INT32}
+ * <li>{@link FieldType#INT64}
+ * <li>{@link FieldType#FLOAT}
+ * <li>{@link FieldType#STRING}
+ * </ul>
+ */
+ public static final Set<Schema.FieldType> VALID_FIELD_TYPE_SET =
+ ImmutableSet.of(
+ FieldType.BYTE,
+ FieldType.BOOLEAN,
+ FieldType.DATETIME,
+ FieldType.DECIMAL,
+ FieldType.DOUBLE,
+ FieldType.INT16,
+ FieldType.INT32,
+ FieldType.INT64,
+ FieldType.FLOAT,
+ FieldType.STRING);
+
+ static final String DEFAULT_FILENAME_SUFFIX = ".csv";
+
+ /** Instantiates a {@link Write} for writing user types in {@link CSVFormat} format. */
+ public static <T> Write<T> write(String to, CSVFormat csvFormat) {
+ return new AutoValue_CsvIO_Write.Builder<T>()
+ .setTextIOWrite(createDefaultTextIOWrite(to))
+ .setCSVFormat(csvFormat)
+ .build();
+ }
+
+ /** Instantiates a {@link Write} for writing {@link Row}s in {@link CSVFormat} format. */
+ public static Write<Row> writeRows(String to, CSVFormat csvFormat) {
+ return new AutoValue_CsvIO_Write.Builder<Row>()
+ .setTextIOWrite(createDefaultTextIOWrite(to))
+ .setCSVFormat(csvFormat)
+ .build();
+ }
+
+ /** {@link PTransform} for writing CSV files. */
+ @AutoValue
+ public abstract static class Write<T>
+ extends PTransform<PCollection<T>, WriteFilesResult<String>> {
+
+ /** Specifies the {@link Compression} of all generated shard files. */
+ public Write<T> withCompression(Compression compression) {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withCompression(compression)).build();
+ }
+
+ public Write<T> withNoSpilling() {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withNoSpilling()).build();
+ }
+
+ /**
+ * Specifies to use a given fixed number of shards per window. See {@link
+ * TextIO.Write#withNumShards(int)}.
+ */
+ public Write<T> withNumShards(Integer numShards) {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withNumShards(numShards)).build();
+ }
+
+ /**
+ * Forces a single file as output and empty shard name template. See {@link
+ * TextIO.Write#withoutSharding()}.
+ */
+ public Write<T> withoutSharding() {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withoutSharding()).build();
+ }
+
+ /**
+ * Uses the given {@link ShardNameTemplate} for naming output files. See {@link
+ * TextIO.Write#withShardNameTemplate(String)}.
+ */
+ public Write<T> withShardTemplate(String shardTemplate) {
+ return toBuilder()
+ .setTextIOWrite(getTextIOWrite().withShardNameTemplate(shardTemplate))
+ .build();
+ }
+
+ /**
+ * Configures the filename suffix for written files. See {@link
+ * TextIO.Write#withSuffix(String)}.
+ */
+ public Write<T> withSuffix(String suffix) {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withSuffix(suffix)).build();
+ }
+
+ /**
+ * Set the base directory used to generate temporary files. See {@link
+ * TextIO.Write#withTempDirectory(ResourceId)}.
+ */
+ public Write<T> withTempDirectory(ResourceId tempDirectory) {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withTempDirectory(tempDirectory)).build();
+ }
+
+ /**
+ * Preserves windowing of input elements and writes them to files based on the element's window.
+ * See {@link TextIO.Write#withWindowedWrites()}.
+ */
+ public Write<T> withWindowedWrites() {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withWindowedWrites()).build();
+ }
+
+ /**
+ * Returns a transform for writing to text files like this one but that has the given {@link
+ * FileBasedSink.WritableByteChannelFactory} to be used by the {@link FileBasedSink} during
+ * output. See {@link
+ * TextIO.Write#withWritableByteChannelFactory(FileBasedSink.WritableByteChannelFactory)}.
+ */
+ public Write<T> withWritableByteChannelFactory(
+ FileBasedSink.WritableByteChannelFactory writableByteChannelFactory) {
+ return toBuilder()
+ .setTextIOWrite(
+ getTextIOWrite().withWritableByteChannelFactory(writableByteChannelFactory))
+ .build();
+ }
+
+ /** The underlying {@link FileIO.Write} that writes converted input to CSV formatted output. */
+ abstract TextIO.Write getTextIOWrite();
+
+ /** The {@link CSVFormat} to inform headers, header comments, and format CSV row content. */
+ abstract CSVFormat getCSVFormat();
+
+ abstract Builder<T> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<T> {
+
+ /**
+ * The underlying {@link FileIO.Write} that writes converted input to CSV formatted output.
+ */
+ abstract Builder<T> setTextIOWrite(TextIO.Write value);
+
+ /** The {@link CSVFormat} to convert input. Defaults to {@link CSVFormat#DEFAULT}. */
+ abstract Builder<T> setCSVFormat(CSVFormat value);
+
+ abstract CSVFormat getCSVFormat();
+
+ abstract Write<T> autoBuild();
+
+ final Write<T> build() {
+ checkArgument(
+ !getCSVFormat().getSkipHeaderRecord(),
Review Comment:
CsvRowConversions performs many of the checks as they most impact RowToCsv conversion. While package-private, I thought it best to do the checks closest to where they would impact the outcome.
Additionally, does the 375 to 377 comment clarify why we are checking the skipHeaderRecord property and throwing an exception?
--
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: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [beam] lukecwik merged pull request #24630: [24469] Implement CsvIO.Write and supporting classes
Posted by "lukecwik (via GitHub)" <gi...@apache.org>.
lukecwik merged PR #24630:
URL: https://github.com/apache/beam/pull/24630
--
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: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [beam] damondouglas commented on pull request #24630: [24469] Implement CsvIO.Write and supporting classes
Posted by GitBox <gi...@apache.org>.
damondouglas commented on PR #24630:
URL: https://github.com/apache/beam/pull/24630#issuecomment-1359553823
Run Java PreCommit
--
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: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [beam] lukecwik commented on a diff in pull request #24630: [24469] Implement CsvIO.Write and supporting classes
Posted by "lukecwik (via GitHub)" <gi...@apache.org>.
lukecwik commented on code in PR #24630:
URL: https://github.com/apache/beam/pull/24630#discussion_r1084642090
##########
sdks/java/io/csv/src/main/java/org/apache/beam/sdk/io/csv/CsvIO.java:
##########
@@ -0,0 +1,446 @@
+/*
+ * 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.sdk.io.csv;
+
+import static java.util.Objects.requireNonNull;
+import static org.apache.beam.sdk.values.TypeDescriptors.rows;
+import static org.apache.beam.sdk.values.TypeDescriptors.strings;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+
+import com.google.auto.value.AutoValue;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Set;
+import org.apache.beam.sdk.coders.RowCoder;
+import org.apache.beam.sdk.io.Compression;
+import org.apache.beam.sdk.io.FileBasedSink;
+import org.apache.beam.sdk.io.FileIO;
+import org.apache.beam.sdk.io.ShardNameTemplate;
+import org.apache.beam.sdk.io.TextIO;
+import org.apache.beam.sdk.io.WriteFilesResult;
+import org.apache.beam.sdk.io.fs.ResourceId;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.transforms.MapElements;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.sdk.values.TypeDescriptor;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableSet;
+import org.apache.commons.csv.CSVFormat;
+
+/**
+ * {@link PTransform}s for reading and writing CSV files.
+ *
+ * <h2>Reading CSV files</h2>
+ *
+ * <p>Reading from CSV files is not yet implemented. Please see <a
+ * href="https://github.com/apache/beam/issues/24552">https://github.com/apache/beam/issues/24552</a>.
+ *
+ * <h2>Writing CSV files</h2>
+ *
+ * <p>To write a {@link PCollection} to one or more CSV files, use {@link CsvIO.Write}, using {@link
+ * CsvIO#writeRows(String, CSVFormat)} or {@link CsvIO#write(String, CSVFormat)}. {@link
+ * CsvIO.Write} supports writing {@link Row} or custom Java types using an inferred {@link Schema}.
+ * Examples below show both scenarios. See the Beam Programming Guide on <a
+ * href="https://beam.apache.org/documentation/programming-guide/#inferring-schemas">inferring
+ * schemas</a> for more information on how to enable Beam to infer a {@link Schema} from a custom
+ * Java type.
+ *
+ * <p>{@link CsvIO.Write} only supports writing {@link Schema} aware types that do not contain any
+ * nested {@link FieldType}s such a {@link org.apache.beam.sdk.schemas.Schema.TypeName#ROW} or
+ * repeated {@link org.apache.beam.sdk.schemas.Schema.TypeName#ARRAY} types. See {@link
+ * CsvIO.Write#VALID_FIELD_TYPE_SET} for valid {@link FieldType}s.
+ *
+ * <h3>Example usage:</h3>
+ *
+ * <p>Suppose we have the following <code>Transaction</code> class annotated with
+ * {@code @DefaultSchema(JavaBeanSchema.class)} so that Beam can infer its {@link Schema}:
+ *
+ * <pre>{@code @DefaultSchema(JavaBeanSchema.class)
+ * public class Transaction {
+ * public Transaction() { … }
+ * public Long getTransactionId();
+ * public void setTransactionId(Long transactionId) { … }
+ * public String getBank() { … }
+ * public void setBank(String bank) { … }
+ * public double getPurchaseAmount() { … }
+ * public void setPurchaseAmount(double purchaseAmount) { … }
+ * }
+ * }</pre>
+ *
+ * <p>From a {@code PCollection<Transaction>}, {@link CsvIO.Write} can write one or many CSV files
+ * automatically creating the header based on its inferred {@link Schema}.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions = ...
+ * transactions.apply(CsvIO.<Transaction>write("path/to/folder/prefix", CSVFormat.DEFAULT));
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header is repeated for every
+ * shard file, whereas by default, {@link CsvIO.Write} will write all fields in <b>sorted order</b>
+ * of the field names.
+ *
+ * <pre>{@code
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ *
+ * <p>To control the order and subset of fields that {@link CsvIO.Write} writes, use {@link
+ * CSVFormat#withHeader(String...)}. Note, however, the following constraints:
+ *
+ * <ol>
+ * <li>Each header column must match a field name in the {@link Schema}; matching is case
+ * sensitive.
+ * <li>Matching header columns must match {@link Schema} fields that are valid {@link FieldType}s;
+ * see {@link #VALID_FIELD_TYPE_SET}.
+ * <li>{@link CSVFormat} only allows repeated header columns when {@link
+ * CSVFormat#withAllowDuplicateHeaderNames()}
+ * </ol>
+ *
+ * <p>The following example shows the use of {@link CSVFormat#withHeader(String...)} to control the
+ * order and subset of <code>Transaction</code> fields.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions ...
+ * transactions.apply(
+ * CsvIO
+ * .<Transaction>write("path/to/folder/prefix", CSVFormat.DEFAULT.withHeader("transactionId", "purchaseAmount"))
+ * );
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header is repeated for every
+ * shard file, but only including the subset of fields in their listed order.
Review Comment:
```suggestion
* file, but will only include the subset of fields in their listed order.
```
##########
sdks/java/io/csv/src/main/java/org/apache/beam/sdk/io/csv/CsvIO.java:
##########
@@ -0,0 +1,446 @@
+/*
+ * 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.sdk.io.csv;
+
+import static java.util.Objects.requireNonNull;
+import static org.apache.beam.sdk.values.TypeDescriptors.rows;
+import static org.apache.beam.sdk.values.TypeDescriptors.strings;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+
+import com.google.auto.value.AutoValue;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Set;
+import org.apache.beam.sdk.coders.RowCoder;
+import org.apache.beam.sdk.io.Compression;
+import org.apache.beam.sdk.io.FileBasedSink;
+import org.apache.beam.sdk.io.FileIO;
+import org.apache.beam.sdk.io.ShardNameTemplate;
+import org.apache.beam.sdk.io.TextIO;
+import org.apache.beam.sdk.io.WriteFilesResult;
+import org.apache.beam.sdk.io.fs.ResourceId;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.transforms.MapElements;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.sdk.values.TypeDescriptor;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableSet;
+import org.apache.commons.csv.CSVFormat;
+
+/**
+ * {@link PTransform}s for reading and writing CSV files.
+ *
+ * <h2>Reading CSV files</h2>
+ *
+ * <p>Reading from CSV files is not yet implemented. Please see <a
+ * href="https://github.com/apache/beam/issues/24552">https://github.com/apache/beam/issues/24552</a>.
+ *
+ * <h2>Writing CSV files</h2>
+ *
+ * <p>To write a {@link PCollection} to one or more CSV files, use {@link CsvIO.Write}, using {@link
+ * CsvIO#writeRows(String, CSVFormat)} or {@link CsvIO#write(String, CSVFormat)}. {@link
+ * CsvIO.Write} supports writing {@link Row} or custom Java types using an inferred {@link Schema}.
+ * Examples below show both scenarios. See the Beam Programming Guide on <a
+ * href="https://beam.apache.org/documentation/programming-guide/#inferring-schemas">inferring
+ * schemas</a> for more information on how to enable Beam to infer a {@link Schema} from a custom
+ * Java type.
+ *
+ * <p>{@link CsvIO.Write} only supports writing {@link Schema} aware types that do not contain any
+ * nested {@link FieldType}s such a {@link org.apache.beam.sdk.schemas.Schema.TypeName#ROW} or
+ * repeated {@link org.apache.beam.sdk.schemas.Schema.TypeName#ARRAY} types. See {@link
+ * CsvIO.Write#VALID_FIELD_TYPE_SET} for valid {@link FieldType}s.
Review Comment:
```suggestion
* <p>{@link CsvIO.Write} only supports writing the parts of {@link Schema} aware types that do not contain any
* nested {@link FieldType}s such a {@link org.apache.beam.sdk.schemas.Schema.TypeName#ROW} or
* repeated {@link org.apache.beam.sdk.schemas.Schema.TypeName#ARRAY} types. See {@link
* CsvIO.Write#VALID_FIELD_TYPE_SET} for valid {@link FieldType}s.
```
##########
sdks/java/io/csv/src/main/java/org/apache/beam/sdk/io/csv/CsvRowConversions.java:
##########
@@ -0,0 +1,148 @@
+/*
+ * 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.sdk.io.csv;
+
+import static org.apache.beam.sdk.io.csv.CsvIO.VALID_FIELD_TYPE_SET;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkNotNull;
+
+import com.google.auto.value.AutoValue;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Set;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.schemas.Schema.TypeName;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.values.Row;
+import org.apache.commons.csv.CSVFormat;
+import org.checkerframework.checker.nullness.qual.NonNull;
+
+/** Contains classes and methods to help with converting between {@link Row} and CSV strings. */
+class CsvRowConversions {
+
+ /** Converts between {@link Row} and CSV string using a {@link CSVFormat}. */
+ @AutoValue
+ abstract static class RowToCsv implements SerializableFunction<Row, String> {
+
+ static Builder builder() {
+ return new AutoValue_CsvRowConversions_RowToCsv.Builder();
+ }
+
+ /** The expected {@link Schema} of the {@link Row} input. */
+ abstract Schema getSchema();
+
+ /** The {@link CSVFormat} of the converted {@link Row} input. */
+ abstract CSVFormat getCSVFormat();
+
+ /** Converts a {@link Row} to a CSV string formatted using {@link #getCSVFormat()}. */
+ @Override
+ public String apply(Row input) {
+ Row safeInput = checkNotNull(input);
+ String[] header = getHeader();
+ Object[] values = new Object[header.length];
+ for (int i = 0; i < header.length; i++) {
+ values[i] = safeInput.getValue(header[i]);
+ }
+ return getCSVFormat().format(values);
+ }
+
+ @NonNull
+ String[] getHeader() {
+ return checkNotNull(getCSVFormat().getHeader());
+ }
+
+ @AutoValue.Builder
+ abstract static class Builder {
+
+ /** The expected {@link Schema} of the {@link Row} input. */
+ abstract Builder setSchema(Schema schema);
+
+ abstract Schema getSchema();
+
+ /** The {@link CSVFormat} of the converted {@link Row} input. */
+ abstract Builder setCSVFormat(CSVFormat format);
+
+ abstract CSVFormat getCSVFormat();
+
+ abstract RowToCsv autoBuild();
+
+ final RowToCsv build() {
+ checkArgument(getSchema().getFieldCount() > 0, "Schema has no fields");
+ setCSVFormat(
+ getCSVFormat()
+ .withSkipHeaderRecord()
+ // Delegate to TextIO.Write.withDelimiter instead.
+ .withRecordSeparator((char) 32)
Review Comment:
Why is this necessary?
```suggestion
.withRecordSeparator(' ')
```
##########
sdks/java/io/csv/src/main/java/org/apache/beam/sdk/io/csv/CsvIO.java:
##########
@@ -0,0 +1,446 @@
+/*
+ * 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.sdk.io.csv;
+
+import static java.util.Objects.requireNonNull;
+import static org.apache.beam.sdk.values.TypeDescriptors.rows;
+import static org.apache.beam.sdk.values.TypeDescriptors.strings;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+
+import com.google.auto.value.AutoValue;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Set;
+import org.apache.beam.sdk.coders.RowCoder;
+import org.apache.beam.sdk.io.Compression;
+import org.apache.beam.sdk.io.FileBasedSink;
+import org.apache.beam.sdk.io.FileIO;
+import org.apache.beam.sdk.io.ShardNameTemplate;
+import org.apache.beam.sdk.io.TextIO;
+import org.apache.beam.sdk.io.WriteFilesResult;
+import org.apache.beam.sdk.io.fs.ResourceId;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.transforms.MapElements;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.sdk.values.TypeDescriptor;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableSet;
+import org.apache.commons.csv.CSVFormat;
+
+/**
+ * {@link PTransform}s for reading and writing CSV files.
+ *
+ * <h2>Reading CSV files</h2>
+ *
+ * <p>Reading from CSV files is not yet implemented. Please see <a
+ * href="https://github.com/apache/beam/issues/24552">https://github.com/apache/beam/issues/24552</a>.
+ *
+ * <h2>Writing CSV files</h2>
+ *
+ * <p>To write a {@link PCollection} to one or more CSV files, use {@link CsvIO.Write}, using {@link
+ * CsvIO#writeRows(String, CSVFormat)} or {@link CsvIO#write(String, CSVFormat)}. {@link
+ * CsvIO.Write} supports writing {@link Row} or custom Java types using an inferred {@link Schema}.
+ * Examples below show both scenarios. See the Beam Programming Guide on <a
+ * href="https://beam.apache.org/documentation/programming-guide/#inferring-schemas">inferring
+ * schemas</a> for more information on how to enable Beam to infer a {@link Schema} from a custom
+ * Java type.
+ *
+ * <p>{@link CsvIO.Write} only supports writing {@link Schema} aware types that do not contain any
+ * nested {@link FieldType}s such a {@link org.apache.beam.sdk.schemas.Schema.TypeName#ROW} or
+ * repeated {@link org.apache.beam.sdk.schemas.Schema.TypeName#ARRAY} types. See {@link
+ * CsvIO.Write#VALID_FIELD_TYPE_SET} for valid {@link FieldType}s.
+ *
+ * <h3>Example usage:</h3>
+ *
+ * <p>Suppose we have the following <code>Transaction</code> class annotated with
+ * {@code @DefaultSchema(JavaBeanSchema.class)} so that Beam can infer its {@link Schema}:
+ *
+ * <pre>{@code @DefaultSchema(JavaBeanSchema.class)
+ * public class Transaction {
+ * public Transaction() { … }
+ * public Long getTransactionId();
+ * public void setTransactionId(Long transactionId) { … }
+ * public String getBank() { … }
+ * public void setBank(String bank) { … }
+ * public double getPurchaseAmount() { … }
+ * public void setPurchaseAmount(double purchaseAmount) { … }
+ * }
+ * }</pre>
+ *
+ * <p>From a {@code PCollection<Transaction>}, {@link CsvIO.Write} can write one or many CSV files
+ * automatically creating the header based on its inferred {@link Schema}.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions = ...
+ * transactions.apply(CsvIO.<Transaction>write("path/to/folder/prefix", CSVFormat.DEFAULT));
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header is repeated for every
+ * shard file, whereas by default, {@link CsvIO.Write} will write all fields in <b>sorted order</b>
+ * of the field names.
+ *
+ * <pre>{@code
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ *
+ * <p>To control the order and subset of fields that {@link CsvIO.Write} writes, use {@link
+ * CSVFormat#withHeader(String...)}. Note, however, the following constraints:
+ *
+ * <ol>
+ * <li>Each header column must match a field name in the {@link Schema}; matching is case
+ * sensitive.
+ * <li>Matching header columns must match {@link Schema} fields that are valid {@link FieldType}s;
+ * see {@link #VALID_FIELD_TYPE_SET}.
+ * <li>{@link CSVFormat} only allows repeated header columns when {@link
+ * CSVFormat#withAllowDuplicateHeaderNames()}
+ * </ol>
+ *
+ * <p>The following example shows the use of {@link CSVFormat#withHeader(String...)} to control the
+ * order and subset of <code>Transaction</code> fields.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions ...
+ * transactions.apply(
+ * CsvIO
+ * .<Transaction>write("path/to/folder/prefix", CSVFormat.DEFAULT.withHeader("transactionId", "purchaseAmount"))
+ * );
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header is repeated for every
+ * shard file, but only including the subset of fields in their listed order.
+ *
+ * <pre>{@code
+ * transactionId,purchaseAmount
+ * 12345,10.23
+ * 54321,54.65
+ * 98765,11.76
+ * }</pre>
+ *
+ * <p>In addition to header customization, {@link CsvIO.Write} supports {@link
+ * CSVFormat#withHeaderComments(Object...)} as shown below. Note that {@link
+ * CSVFormat#withCommentMarker(char)} is required when specifying header comments.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions = ...
+ * transactions
+ * .apply(
+ * CsvIO.<Transaction>write("path/to/folder/prefix",
+ * CSVFormat.DEFAULT.withCommentMarker('#').withHeaderComments("Bank Report", "1970-01-01", "Operator: John Doe")
+ * );
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header and header comments are
+ * repeated for every shard file.
+ *
+ * <pre>{@code
+ * # Bank Report
+ * # 1970-01-01
+ * # Operator: John Doe
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ *
+ * <p>A {@link PCollection} of {@link Row}s works just like custom Java types illustrated above,
+ * except we use {@link CsvIO#writeRows(String, CSVFormat)} as shown below for the same {@code
+ * Transaction} class. We derive {@code Transaction}'s {@link Schema} using a {@link
+ * org.apache.beam.sdk.schemas.annotations.DefaultSchema.DefaultSchemaProvider}. Note that
+ * hard-coding the {@link Row}s below is for illustration purposes. Developers are instead
+ * encouraged to take advantage of {@link
+ * org.apache.beam.sdk.schemas.annotations.DefaultSchema.DefaultSchemaProvider#toRowFunction(TypeDescriptor)}.
+ *
+ * <pre>{@code
+ * DefaultSchemaProvider defaultSchemaProvider = new DefaultSchemaProvider();
+ * Schema schema = defaultSchemaProvider.schemaFor(TypeDescriptor.of(Transaction.class));
+ * PCollection<Row> transactions = pipeline.apply(Create.of(
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "A")
+ * .withFieldValue("purchaseAmount", 10.23)
+ * .withFieldValue("transactionId", "12345")
+ * .build(),
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "B")
+ * .withFieldValue("purchaseAmount", 54.65)
+ * .withFieldValue("transactionId", "54321")
+ * .build(),
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "C")
+ * .withFieldValue("purchaseAmount", 11.76)
+ * .withFieldValue("transactionId", "98765")
+ * .build()
+ * );
+ *
+ * transactions.apply(
+ * CsvIO
+ * .writeRowsTo("gs://bucket/path/to/folder/prefix", CSVFormat.DEFAULT)
+ * );
+ * }</pre>
+ *
+ * <p>Writing the transactions {@link PCollection} of {@link Row}s would yield the following CSV
+ * file content.
+ *
+ * <pre>{@code
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ *
+ * {@link CsvIO.Write} does not support the following {@link CSVFormat} properties and will throw an
+ * {@link IllegalArgumentException}.
+ *
+ * <ul>
+ * <li>{@link CSVFormat#withSkipHeaderRecord()}
+ * <li>{@link CSVFormat#withAllowMissingColumnNames()} ()}
+ * <li>{@link CSVFormat#withAutoFlush(boolean)}
+ * <li>{@link CSVFormat#withIgnoreHeaderCase()}
+ * <li>{@link CSVFormat#withIgnoreSurroundingSpaces()}
+ * </ul>
+ */
+public class CsvIO {
+ /**
+ * The valid {@link Schema.FieldType} from which {@link CsvIO} converts CSV records to the fields.
+ *
+ * <ul>
+ * <li>{@link FieldType#BYTE}
+ * <li>{@link FieldType#BOOLEAN}
+ * <li>{@link FieldType#DATETIME}
+ * <li>{@link FieldType#DECIMAL}
+ * <li>{@link FieldType#DOUBLE}
+ * <li>{@link FieldType#INT16}
+ * <li>{@link FieldType#INT32}
+ * <li>{@link FieldType#INT64}
+ * <li>{@link FieldType#FLOAT}
+ * <li>{@link FieldType#STRING}
+ * </ul>
+ */
+ public static final Set<Schema.FieldType> VALID_FIELD_TYPE_SET =
+ ImmutableSet.of(
+ FieldType.BYTE,
+ FieldType.BOOLEAN,
+ FieldType.DATETIME,
+ FieldType.DECIMAL,
+ FieldType.DOUBLE,
+ FieldType.INT16,
+ FieldType.INT32,
+ FieldType.INT64,
+ FieldType.FLOAT,
+ FieldType.STRING);
+
+ static final String DEFAULT_FILENAME_SUFFIX = ".csv";
+
+ /** Instantiates a {@link Write} for writing user types in {@link CSVFormat} format. */
+ public static <T> Write<T> write(String to, CSVFormat csvFormat) {
+ return new AutoValue_CsvIO_Write.Builder<T>()
+ .setTextIOWrite(createDefaultTextIOWrite(to))
+ .setCSVFormat(csvFormat)
+ .build();
+ }
+
+ /** Instantiates a {@link Write} for writing {@link Row}s in {@link CSVFormat} format. */
+ public static Write<Row> writeRows(String to, CSVFormat csvFormat) {
+ return new AutoValue_CsvIO_Write.Builder<Row>()
+ .setTextIOWrite(createDefaultTextIOWrite(to))
+ .setCSVFormat(csvFormat)
+ .build();
+ }
+
+ /** {@link PTransform} for writing CSV files. */
+ @AutoValue
+ public abstract static class Write<T>
+ extends PTransform<PCollection<T>, WriteFilesResult<String>> {
+
+ /** Specifies the {@link Compression} of all generated shard files. */
+ public Write<T> withCompression(Compression compression) {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withCompression(compression)).build();
+ }
+
+ public Write<T> withNoSpilling() {
Review Comment:
javadoc?
##########
sdks/java/io/csv/src/test/java/org/apache/beam/sdk/io/csv/RowToCsvCSVFormatTest.java:
##########
@@ -0,0 +1,353 @@
+/*
+ * 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.sdk.io.csv;
+
+import static org.apache.beam.sdk.io.csv.CsvIOTestData.DATA;
+import static org.apache.beam.sdk.io.csv.CsvIOTestJavaBeans.ALL_PRIMITIVE_DATA_TYPES_SCHEMA;
+import static org.apache.beam.sdk.io.csv.CsvIOTestJavaBeans.NULLABLE_ALL_PRIMITIVE_DATA_TYPES_SCHEMA;
+import static org.apache.beam.sdk.io.csv.CsvIOTestJavaBeans.TIME_CONTAINING_SCHEMA;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkNotNull;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertThrows;
+
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.Field;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.values.Row;
+import org.apache.commons.csv.CSVFormat;
+import org.apache.commons.csv.QuoteMode;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/**
+ * Tests {@link org.apache.commons.csv.CSVFormat} settings in the context of {@link
+ * CsvRowConversions.RowToCsv}.
+ */
+@RunWith(JUnit4.class)
+public class RowToCsvCSVFormatTest {
+ @Test
+ public void invalidCSVFormatHeader() {
+ NullPointerException nullHeaderError =
+ assertThrows(
+ NullPointerException.class,
+ () ->
+ CsvRowConversions.RowToCsv.builder()
+ .setCSVFormat(CSVFormat.DEFAULT)
+ .setSchema(ALL_PRIMITIVE_DATA_TYPES_SCHEMA)
+ .build());
+ assertEquals("CSVFormat withHeader is required", nullHeaderError.getMessage());
+
+ IllegalArgumentException emptyHeaderError =
+ assertThrows(
+ IllegalArgumentException.class,
+ () ->
+ CsvRowConversions.RowToCsv.builder()
+ .setCSVFormat(CSVFormat.DEFAULT.withHeader())
+ .setSchema(ALL_PRIMITIVE_DATA_TYPES_SCHEMA)
+ .build());
+
+ assertEquals(
+ "CSVFormat withHeader requires at least one column", emptyHeaderError.getMessage());
+
+ IllegalArgumentException mismatchHeaderError =
+ assertThrows(
+ IllegalArgumentException.class,
+ () ->
+ CsvRowConversions.RowToCsv.builder()
+ .setCSVFormat(
+ CSVFormat.DEFAULT.withHeader("aString", "idontexist1", "idontexist2"))
+ .setSchema(ALL_PRIMITIVE_DATA_TYPES_SCHEMA)
+ .build());
+
+ assertEquals(
+ "columns in CSVFormat header do not exist in Schema: idontexist2,idontexist1",
+ mismatchHeaderError.getMessage());
+ }
+
+ @Test
+ public void invalidSchema() {
+ IllegalArgumentException emptySchemaError =
+ assertThrows(
+ IllegalArgumentException.class,
+ () ->
+ CsvRowConversions.RowToCsv.builder()
+ .setCSVFormat(CSVFormat.DEFAULT.withHeader())
+ .setSchema(Schema.of())
+ .build());
+
+ assertEquals("Schema has no fields", emptySchemaError.getMessage());
+
+ IllegalArgumentException invalidArrayFieldsError =
+ assertThrows(
+ IllegalArgumentException.class,
+ () ->
+ CsvRowConversions.RowToCsv.builder()
+ .setCSVFormat(CSVFormat.DEFAULT.withHeader("instant", "instantList"))
+ .setSchema(TIME_CONTAINING_SCHEMA)
+ .build());
+
+ assertEquals(
+ "columns in header match fields in Schema with invalid types: instantList. See CsvIO#VALID_FIELD_TYPE_SET for a list of valid field types.",
+ invalidArrayFieldsError.getMessage());
+
+ // Should not throw Exception when limited to valid fields.
+ CsvRowConversions.RowToCsv.builder()
+ .setCSVFormat(CSVFormat.DEFAULT.withHeader("instant"))
+ .setSchema(TIME_CONTAINING_SCHEMA)
+ .build();
+ }
+
+ @Test
+ public void withAllowDuplicateHeaderNamesDuplicatesRowFieldOutput() {
+ assertEquals(
+ "allowDuplicateHeaderNames=true",
+ "a,a,a",
+ rowToCsv(
+ DATA.allPrimitiveDataTypesRow,
+ csvFormat(ALL_PRIMITIVE_DATA_TYPES_SCHEMA)
+ .withAllowDuplicateHeaderNames(true)
+ .withHeader("aString", "aString", "aString")));
+ }
+
+ @Test
+ public void withAllowMissingColumnNamesSettingThrowsException() {
+ IllegalArgumentException exception =
+ assertThrows(
+ "allowMissingColumnNames=true",
+ IllegalArgumentException.class,
+ () ->
+ rowToCsv(
+ DATA.allPrimitiveDataTypesRow,
+ csvFormat(ALL_PRIMITIVE_DATA_TYPES_SCHEMA).withAllowMissingColumnNames(true)));
+
+ assertEquals(
+ "withAllowMissingColumnNames is an illegal CSVFormat setting", exception.getMessage());
+ }
+
+ @Test
+ public void withAutoFlushThrowsException() {
+ IllegalArgumentException exception =
+ assertThrows(
+ "autoFlush=true",
+ IllegalArgumentException.class,
+ () ->
+ rowToCsv(
+ DATA.allPrimitiveDataTypesRow,
+ csvFormat(ALL_PRIMITIVE_DATA_TYPES_SCHEMA).withAutoFlush(true)));
+
+ assertEquals("withAutoFlush is an illegal CSVFormat setting", exception.getMessage());
+ }
+
+ @Test
+ public void withCommentMarkerDoesNotEffectConversion() {
+ Schema schema = Schema.of(Field.of("aString", FieldType.STRING));
+ Row row = Row.withSchema(schema).attachValues("$abc");
+ assertEquals("$abc", rowToCsv(row, csvFormat(schema).withCommentMarker('$')));
+ assertEquals("$abc", rowToCsv(row, csvFormat(schema).withCommentMarker(null)));
+ }
+
+ @Test
+ public void withDelimiterDrivesCellBorders() {
+ assertEquals(
+ "false~1~10~1.0~1.0~1~1~a~1",
+ rowToCsv(
+ DATA.allPrimitiveDataTypesRow,
+ csvFormat(ALL_PRIMITIVE_DATA_TYPES_SCHEMA).withDelimiter('~')));
+ }
+
+ @Test
+ public void withEscapeDrivesOutput() {
+ Schema schema =
+ Schema.of(Field.of("aString", FieldType.STRING), Field.of("anInt", FieldType.INT32));
+ Row row = Row.withSchema(schema).attachValues(",a", 1);
+ String[] header = new String[] {"anInt", "aString"};
+ assertEquals(
+ "1,#,a",
+ rowToCsv(
+ row,
+ csvFormat(ALL_PRIMITIVE_DATA_TYPES_SCHEMA)
+ .withHeader(header)
+ .withEscape('#')
+ .withQuoteMode(QuoteMode.NONE)));
+ assertEquals(
+ "1,\",a\"", rowToCsv(row, csvFormat(ALL_PRIMITIVE_DATA_TYPES_SCHEMA).withHeader(header)));
+ }
+
+ @Test
+ public void withHeaderDrivesFieldOrderSubsetOutput() {
+ assertEquals(
+ "1,false,a",
+ rowToCsv(
+ DATA.allPrimitiveDataTypesRow,
+ csvFormat(ALL_PRIMITIVE_DATA_TYPES_SCHEMA)
+ .withHeader("anInteger", "aBoolean", "aString")));
+ }
+
+ @Test
+ public void withHeaderCommentsDoesNotEffectConversion() {
+ assertEquals(
+ "false,1,10,1.0,1.0,1,1,a,1",
+ rowToCsv(
+ DATA.allPrimitiveDataTypesRow,
+ csvFormat(ALL_PRIMITIVE_DATA_TYPES_SCHEMA)
+ .withHeaderComments("some", "header", "comments")));
+ }
+
+ @Test
+ public void withIgnoreEmptyLinesDoesNotEffectOutput() {
+ assertEquals(
+ "false,1,10,1.0,1.0,1,1,a,1",
+ rowToCsv(
+ DATA.allPrimitiveDataTypesRow,
+ csvFormat(ALL_PRIMITIVE_DATA_TYPES_SCHEMA).withIgnoreEmptyLines(true)));
+ assertEquals(
+ "false,1,10,1.0,1.0,1,1,a,1",
+ rowToCsv(
+ DATA.allPrimitiveDataTypesRow,
+ csvFormat(ALL_PRIMITIVE_DATA_TYPES_SCHEMA).withIgnoreEmptyLines(false)));
+ }
+
+ @Test
+ public void withIgnoreHeaderCaseThrowsException() {
+ IllegalArgumentException exception =
+ assertThrows(
+ IllegalArgumentException.class,
+ () ->
+ rowToCsv(
+ DATA.allPrimitiveDataTypesRow,
+ csvFormat(ALL_PRIMITIVE_DATA_TYPES_SCHEMA).withIgnoreHeaderCase(true)));
+ assertEquals("withIgnoreHeaderCase is an illegal CSVFormat setting", exception.getMessage());
+ }
+
+ @Test
+ public void withIgnoreSurroundingSpacesThrowsException() {
+ IllegalArgumentException exception =
+ assertThrows(
+ IllegalArgumentException.class,
+ () ->
+ rowToCsv(
+ DATA.allPrimitiveDataTypesRow,
+ csvFormat(ALL_PRIMITIVE_DATA_TYPES_SCHEMA).withIgnoreSurroundingSpaces(true)));
+ assertEquals(
+ "withIgnoreSurroundingSpaces is an illegal CSVFormat setting", exception.getMessage());
+ }
+
+ @Test
+ public void withNullStringReplacesNullValues() {
+ assertEquals(
+ "🦄,🦄,🦄,🦄,🦄,🦄",
Review Comment:
nice
##########
sdks/java/io/csv/src/main/java/org/apache/beam/sdk/io/csv/CsvIO.java:
##########
@@ -0,0 +1,446 @@
+/*
+ * 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.sdk.io.csv;
+
+import static java.util.Objects.requireNonNull;
+import static org.apache.beam.sdk.values.TypeDescriptors.rows;
+import static org.apache.beam.sdk.values.TypeDescriptors.strings;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+
+import com.google.auto.value.AutoValue;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Set;
+import org.apache.beam.sdk.coders.RowCoder;
+import org.apache.beam.sdk.io.Compression;
+import org.apache.beam.sdk.io.FileBasedSink;
+import org.apache.beam.sdk.io.FileIO;
+import org.apache.beam.sdk.io.ShardNameTemplate;
+import org.apache.beam.sdk.io.TextIO;
+import org.apache.beam.sdk.io.WriteFilesResult;
+import org.apache.beam.sdk.io.fs.ResourceId;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.transforms.MapElements;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.sdk.values.TypeDescriptor;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableSet;
+import org.apache.commons.csv.CSVFormat;
+
+/**
+ * {@link PTransform}s for reading and writing CSV files.
+ *
+ * <h2>Reading CSV files</h2>
+ *
+ * <p>Reading from CSV files is not yet implemented. Please see <a
+ * href="https://github.com/apache/beam/issues/24552">https://github.com/apache/beam/issues/24552</a>.
+ *
+ * <h2>Writing CSV files</h2>
+ *
+ * <p>To write a {@link PCollection} to one or more CSV files, use {@link CsvIO.Write}, using {@link
+ * CsvIO#writeRows(String, CSVFormat)} or {@link CsvIO#write(String, CSVFormat)}. {@link
+ * CsvIO.Write} supports writing {@link Row} or custom Java types using an inferred {@link Schema}.
+ * Examples below show both scenarios. See the Beam Programming Guide on <a
+ * href="https://beam.apache.org/documentation/programming-guide/#inferring-schemas">inferring
+ * schemas</a> for more information on how to enable Beam to infer a {@link Schema} from a custom
+ * Java type.
+ *
+ * <p>{@link CsvIO.Write} only supports writing {@link Schema} aware types that do not contain any
+ * nested {@link FieldType}s such a {@link org.apache.beam.sdk.schemas.Schema.TypeName#ROW} or
+ * repeated {@link org.apache.beam.sdk.schemas.Schema.TypeName#ARRAY} types. See {@link
+ * CsvIO.Write#VALID_FIELD_TYPE_SET} for valid {@link FieldType}s.
+ *
+ * <h3>Example usage:</h3>
+ *
+ * <p>Suppose we have the following <code>Transaction</code> class annotated with
+ * {@code @DefaultSchema(JavaBeanSchema.class)} so that Beam can infer its {@link Schema}:
+ *
+ * <pre>{@code @DefaultSchema(JavaBeanSchema.class)
+ * public class Transaction {
+ * public Transaction() { … }
+ * public Long getTransactionId();
+ * public void setTransactionId(Long transactionId) { … }
+ * public String getBank() { … }
+ * public void setBank(String bank) { … }
+ * public double getPurchaseAmount() { … }
+ * public void setPurchaseAmount(double purchaseAmount) { … }
+ * }
+ * }</pre>
+ *
+ * <p>From a {@code PCollection<Transaction>}, {@link CsvIO.Write} can write one or many CSV files
+ * automatically creating the header based on its inferred {@link Schema}.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions = ...
+ * transactions.apply(CsvIO.<Transaction>write("path/to/folder/prefix", CSVFormat.DEFAULT));
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header is repeated for every
+ * shard file, whereas by default, {@link CsvIO.Write} will write all fields in <b>sorted order</b>
+ * of the field names.
+ *
+ * <pre>{@code
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ *
+ * <p>To control the order and subset of fields that {@link CsvIO.Write} writes, use {@link
+ * CSVFormat#withHeader(String...)}. Note, however, the following constraints:
+ *
+ * <ol>
+ * <li>Each header column must match a field name in the {@link Schema}; matching is case
+ * sensitive.
+ * <li>Matching header columns must match {@link Schema} fields that are valid {@link FieldType}s;
+ * see {@link #VALID_FIELD_TYPE_SET}.
+ * <li>{@link CSVFormat} only allows repeated header columns when {@link
+ * CSVFormat#withAllowDuplicateHeaderNames()}
+ * </ol>
+ *
+ * <p>The following example shows the use of {@link CSVFormat#withHeader(String...)} to control the
+ * order and subset of <code>Transaction</code> fields.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions ...
+ * transactions.apply(
+ * CsvIO
+ * .<Transaction>write("path/to/folder/prefix", CSVFormat.DEFAULT.withHeader("transactionId", "purchaseAmount"))
+ * );
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header is repeated for every
+ * shard file, but only including the subset of fields in their listed order.
+ *
+ * <pre>{@code
+ * transactionId,purchaseAmount
+ * 12345,10.23
+ * 54321,54.65
+ * 98765,11.76
+ * }</pre>
+ *
+ * <p>In addition to header customization, {@link CsvIO.Write} supports {@link
+ * CSVFormat#withHeaderComments(Object...)} as shown below. Note that {@link
+ * CSVFormat#withCommentMarker(char)} is required when specifying header comments.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions = ...
+ * transactions
+ * .apply(
+ * CsvIO.<Transaction>write("path/to/folder/prefix",
+ * CSVFormat.DEFAULT.withCommentMarker('#').withHeaderComments("Bank Report", "1970-01-01", "Operator: John Doe")
+ * );
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header and header comments are
+ * repeated for every shard file.
+ *
+ * <pre>{@code
+ * # Bank Report
+ * # 1970-01-01
+ * # Operator: John Doe
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ *
+ * <p>A {@link PCollection} of {@link Row}s works just like custom Java types illustrated above,
+ * except we use {@link CsvIO#writeRows(String, CSVFormat)} as shown below for the same {@code
+ * Transaction} class. We derive {@code Transaction}'s {@link Schema} using a {@link
+ * org.apache.beam.sdk.schemas.annotations.DefaultSchema.DefaultSchemaProvider}. Note that
+ * hard-coding the {@link Row}s below is for illustration purposes. Developers are instead
+ * encouraged to take advantage of {@link
+ * org.apache.beam.sdk.schemas.annotations.DefaultSchema.DefaultSchemaProvider#toRowFunction(TypeDescriptor)}.
+ *
+ * <pre>{@code
+ * DefaultSchemaProvider defaultSchemaProvider = new DefaultSchemaProvider();
+ * Schema schema = defaultSchemaProvider.schemaFor(TypeDescriptor.of(Transaction.class));
+ * PCollection<Row> transactions = pipeline.apply(Create.of(
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "A")
+ * .withFieldValue("purchaseAmount", 10.23)
+ * .withFieldValue("transactionId", "12345")
+ * .build(),
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "B")
+ * .withFieldValue("purchaseAmount", 54.65)
+ * .withFieldValue("transactionId", "54321")
+ * .build(),
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "C")
+ * .withFieldValue("purchaseAmount", 11.76)
+ * .withFieldValue("transactionId", "98765")
+ * .build()
+ * );
+ *
+ * transactions.apply(
+ * CsvIO
+ * .writeRowsTo("gs://bucket/path/to/folder/prefix", CSVFormat.DEFAULT)
+ * );
+ * }</pre>
+ *
+ * <p>Writing the transactions {@link PCollection} of {@link Row}s would yield the following CSV
+ * file content.
+ *
+ * <pre>{@code
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ *
+ * {@link CsvIO.Write} does not support the following {@link CSVFormat} properties and will throw an
+ * {@link IllegalArgumentException}.
+ *
+ * <ul>
+ * <li>{@link CSVFormat#withSkipHeaderRecord()}
+ * <li>{@link CSVFormat#withAllowMissingColumnNames()} ()}
+ * <li>{@link CSVFormat#withAutoFlush(boolean)}
+ * <li>{@link CSVFormat#withIgnoreHeaderCase()}
+ * <li>{@link CSVFormat#withIgnoreSurroundingSpaces()}
Review Comment:
```suggestion
* <li>{@link CSVFormat#withSkipHeaderRecord}
* <li>{@link CSVFormat#withAllowMissingColumnNames}
* <li>{@link CSVFormat#withAutoFlush}
* <li>{@link CSVFormat#withIgnoreHeaderCase}
* <li>{@link CSVFormat#withIgnoreSurroundingSpaces}
```
##########
sdks/java/io/csv/src/main/java/org/apache/beam/sdk/io/csv/CsvIO.java:
##########
@@ -0,0 +1,446 @@
+/*
+ * 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.sdk.io.csv;
+
+import static java.util.Objects.requireNonNull;
+import static org.apache.beam.sdk.values.TypeDescriptors.rows;
+import static org.apache.beam.sdk.values.TypeDescriptors.strings;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+
+import com.google.auto.value.AutoValue;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Set;
+import org.apache.beam.sdk.coders.RowCoder;
+import org.apache.beam.sdk.io.Compression;
+import org.apache.beam.sdk.io.FileBasedSink;
+import org.apache.beam.sdk.io.FileIO;
+import org.apache.beam.sdk.io.ShardNameTemplate;
+import org.apache.beam.sdk.io.TextIO;
+import org.apache.beam.sdk.io.WriteFilesResult;
+import org.apache.beam.sdk.io.fs.ResourceId;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.transforms.MapElements;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.sdk.values.TypeDescriptor;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableSet;
+import org.apache.commons.csv.CSVFormat;
+
+/**
+ * {@link PTransform}s for reading and writing CSV files.
+ *
+ * <h2>Reading CSV files</h2>
+ *
+ * <p>Reading from CSV files is not yet implemented. Please see <a
+ * href="https://github.com/apache/beam/issues/24552">https://github.com/apache/beam/issues/24552</a>.
+ *
+ * <h2>Writing CSV files</h2>
+ *
+ * <p>To write a {@link PCollection} to one or more CSV files, use {@link CsvIO.Write}, using {@link
+ * CsvIO#writeRows(String, CSVFormat)} or {@link CsvIO#write(String, CSVFormat)}. {@link
+ * CsvIO.Write} supports writing {@link Row} or custom Java types using an inferred {@link Schema}.
+ * Examples below show both scenarios. See the Beam Programming Guide on <a
+ * href="https://beam.apache.org/documentation/programming-guide/#inferring-schemas">inferring
+ * schemas</a> for more information on how to enable Beam to infer a {@link Schema} from a custom
+ * Java type.
+ *
+ * <p>{@link CsvIO.Write} only supports writing {@link Schema} aware types that do not contain any
+ * nested {@link FieldType}s such a {@link org.apache.beam.sdk.schemas.Schema.TypeName#ROW} or
+ * repeated {@link org.apache.beam.sdk.schemas.Schema.TypeName#ARRAY} types. See {@link
+ * CsvIO.Write#VALID_FIELD_TYPE_SET} for valid {@link FieldType}s.
+ *
+ * <h3>Example usage:</h3>
+ *
+ * <p>Suppose we have the following <code>Transaction</code> class annotated with
+ * {@code @DefaultSchema(JavaBeanSchema.class)} so that Beam can infer its {@link Schema}:
+ *
+ * <pre>{@code @DefaultSchema(JavaBeanSchema.class)
+ * public class Transaction {
+ * public Transaction() { … }
+ * public Long getTransactionId();
+ * public void setTransactionId(Long transactionId) { … }
+ * public String getBank() { … }
+ * public void setBank(String bank) { … }
+ * public double getPurchaseAmount() { … }
+ * public void setPurchaseAmount(double purchaseAmount) { … }
+ * }
+ * }</pre>
+ *
+ * <p>From a {@code PCollection<Transaction>}, {@link CsvIO.Write} can write one or many CSV files
+ * automatically creating the header based on its inferred {@link Schema}.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions = ...
+ * transactions.apply(CsvIO.<Transaction>write("path/to/folder/prefix", CSVFormat.DEFAULT));
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header is repeated for every
+ * shard file, whereas by default, {@link CsvIO.Write} will write all fields in <b>sorted order</b>
+ * of the field names.
+ *
+ * <pre>{@code
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ *
+ * <p>To control the order and subset of fields that {@link CsvIO.Write} writes, use {@link
+ * CSVFormat#withHeader(String...)}. Note, however, the following constraints:
Review Comment:
Note that you can drop the parameter types like (String...) in the above. It helps keep the javadoc links working longer and it is only useful if you want to specify an exact version of an overloaded method.
##########
sdks/java/io/csv/src/main/java/org/apache/beam/sdk/io/csv/CsvIO.java:
##########
@@ -0,0 +1,446 @@
+/*
+ * 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.sdk.io.csv;
+
+import static java.util.Objects.requireNonNull;
+import static org.apache.beam.sdk.values.TypeDescriptors.rows;
+import static org.apache.beam.sdk.values.TypeDescriptors.strings;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+
+import com.google.auto.value.AutoValue;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Set;
+import org.apache.beam.sdk.coders.RowCoder;
+import org.apache.beam.sdk.io.Compression;
+import org.apache.beam.sdk.io.FileBasedSink;
+import org.apache.beam.sdk.io.FileIO;
+import org.apache.beam.sdk.io.ShardNameTemplate;
+import org.apache.beam.sdk.io.TextIO;
+import org.apache.beam.sdk.io.WriteFilesResult;
+import org.apache.beam.sdk.io.fs.ResourceId;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.transforms.MapElements;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.sdk.values.TypeDescriptor;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableSet;
+import org.apache.commons.csv.CSVFormat;
+
+/**
+ * {@link PTransform}s for reading and writing CSV files.
+ *
+ * <h2>Reading CSV files</h2>
+ *
+ * <p>Reading from CSV files is not yet implemented. Please see <a
+ * href="https://github.com/apache/beam/issues/24552">https://github.com/apache/beam/issues/24552</a>.
+ *
+ * <h2>Writing CSV files</h2>
+ *
+ * <p>To write a {@link PCollection} to one or more CSV files, use {@link CsvIO.Write}, using {@link
+ * CsvIO#writeRows(String, CSVFormat)} or {@link CsvIO#write(String, CSVFormat)}. {@link
+ * CsvIO.Write} supports writing {@link Row} or custom Java types using an inferred {@link Schema}.
+ * Examples below show both scenarios. See the Beam Programming Guide on <a
+ * href="https://beam.apache.org/documentation/programming-guide/#inferring-schemas">inferring
+ * schemas</a> for more information on how to enable Beam to infer a {@link Schema} from a custom
+ * Java type.
+ *
+ * <p>{@link CsvIO.Write} only supports writing {@link Schema} aware types that do not contain any
+ * nested {@link FieldType}s such a {@link org.apache.beam.sdk.schemas.Schema.TypeName#ROW} or
+ * repeated {@link org.apache.beam.sdk.schemas.Schema.TypeName#ARRAY} types. See {@link
+ * CsvIO.Write#VALID_FIELD_TYPE_SET} for valid {@link FieldType}s.
+ *
+ * <h3>Example usage:</h3>
+ *
+ * <p>Suppose we have the following <code>Transaction</code> class annotated with
+ * {@code @DefaultSchema(JavaBeanSchema.class)} so that Beam can infer its {@link Schema}:
+ *
+ * <pre>{@code @DefaultSchema(JavaBeanSchema.class)
+ * public class Transaction {
+ * public Transaction() { … }
+ * public Long getTransactionId();
+ * public void setTransactionId(Long transactionId) { … }
+ * public String getBank() { … }
+ * public void setBank(String bank) { … }
+ * public double getPurchaseAmount() { … }
+ * public void setPurchaseAmount(double purchaseAmount) { … }
+ * }
+ * }</pre>
+ *
+ * <p>From a {@code PCollection<Transaction>}, {@link CsvIO.Write} can write one or many CSV files
+ * automatically creating the header based on its inferred {@link Schema}.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions = ...
+ * transactions.apply(CsvIO.<Transaction>write("path/to/folder/prefix", CSVFormat.DEFAULT));
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header is repeated for every
+ * shard file, whereas by default, {@link CsvIO.Write} will write all fields in <b>sorted order</b>
+ * of the field names.
+ *
+ * <pre>{@code
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ *
+ * <p>To control the order and subset of fields that {@link CsvIO.Write} writes, use {@link
+ * CSVFormat#withHeader(String...)}. Note, however, the following constraints:
+ *
+ * <ol>
+ * <li>Each header column must match a field name in the {@link Schema}; matching is case
+ * sensitive.
+ * <li>Matching header columns must match {@link Schema} fields that are valid {@link FieldType}s;
+ * see {@link #VALID_FIELD_TYPE_SET}.
+ * <li>{@link CSVFormat} only allows repeated header columns when {@link
+ * CSVFormat#withAllowDuplicateHeaderNames()}
+ * </ol>
+ *
+ * <p>The following example shows the use of {@link CSVFormat#withHeader(String...)} to control the
+ * order and subset of <code>Transaction</code> fields.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions ...
+ * transactions.apply(
+ * CsvIO
+ * .<Transaction>write("path/to/folder/prefix", CSVFormat.DEFAULT.withHeader("transactionId", "purchaseAmount"))
+ * );
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header is repeated for every
+ * shard file, but only including the subset of fields in their listed order.
+ *
+ * <pre>{@code
+ * transactionId,purchaseAmount
+ * 12345,10.23
+ * 54321,54.65
+ * 98765,11.76
+ * }</pre>
+ *
+ * <p>In addition to header customization, {@link CsvIO.Write} supports {@link
+ * CSVFormat#withHeaderComments(Object...)} as shown below. Note that {@link
+ * CSVFormat#withCommentMarker(char)} is required when specifying header comments.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions = ...
+ * transactions
+ * .apply(
+ * CsvIO.<Transaction>write("path/to/folder/prefix",
+ * CSVFormat.DEFAULT.withCommentMarker('#').withHeaderComments("Bank Report", "1970-01-01", "Operator: John Doe")
+ * );
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header and header comments are
+ * repeated for every shard file.
+ *
+ * <pre>{@code
+ * # Bank Report
+ * # 1970-01-01
+ * # Operator: John Doe
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ *
+ * <p>A {@link PCollection} of {@link Row}s works just like custom Java types illustrated above,
+ * except we use {@link CsvIO#writeRows(String, CSVFormat)} as shown below for the same {@code
+ * Transaction} class. We derive {@code Transaction}'s {@link Schema} using a {@link
+ * org.apache.beam.sdk.schemas.annotations.DefaultSchema.DefaultSchemaProvider}. Note that
+ * hard-coding the {@link Row}s below is for illustration purposes. Developers are instead
+ * encouraged to take advantage of {@link
+ * org.apache.beam.sdk.schemas.annotations.DefaultSchema.DefaultSchemaProvider#toRowFunction(TypeDescriptor)}.
+ *
+ * <pre>{@code
+ * DefaultSchemaProvider defaultSchemaProvider = new DefaultSchemaProvider();
+ * Schema schema = defaultSchemaProvider.schemaFor(TypeDescriptor.of(Transaction.class));
+ * PCollection<Row> transactions = pipeline.apply(Create.of(
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "A")
+ * .withFieldValue("purchaseAmount", 10.23)
+ * .withFieldValue("transactionId", "12345")
+ * .build(),
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "B")
+ * .withFieldValue("purchaseAmount", 54.65)
+ * .withFieldValue("transactionId", "54321")
+ * .build(),
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "C")
+ * .withFieldValue("purchaseAmount", 11.76)
+ * .withFieldValue("transactionId", "98765")
+ * .build()
+ * );
+ *
+ * transactions.apply(
+ * CsvIO
+ * .writeRowsTo("gs://bucket/path/to/folder/prefix", CSVFormat.DEFAULT)
+ * );
+ * }</pre>
+ *
+ * <p>Writing the transactions {@link PCollection} of {@link Row}s would yield the following CSV
+ * file content.
+ *
+ * <pre>{@code
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ *
+ * {@link CsvIO.Write} does not support the following {@link CSVFormat} properties and will throw an
+ * {@link IllegalArgumentException}.
+ *
+ * <ul>
+ * <li>{@link CSVFormat#withSkipHeaderRecord()}
+ * <li>{@link CSVFormat#withAllowMissingColumnNames()} ()}
+ * <li>{@link CSVFormat#withAutoFlush(boolean)}
+ * <li>{@link CSVFormat#withIgnoreHeaderCase()}
+ * <li>{@link CSVFormat#withIgnoreSurroundingSpaces()}
+ * </ul>
+ */
+public class CsvIO {
+ /**
+ * The valid {@link Schema.FieldType} from which {@link CsvIO} converts CSV records to the fields.
+ *
+ * <ul>
+ * <li>{@link FieldType#BYTE}
+ * <li>{@link FieldType#BOOLEAN}
+ * <li>{@link FieldType#DATETIME}
+ * <li>{@link FieldType#DECIMAL}
+ * <li>{@link FieldType#DOUBLE}
+ * <li>{@link FieldType#INT16}
+ * <li>{@link FieldType#INT32}
+ * <li>{@link FieldType#INT64}
+ * <li>{@link FieldType#FLOAT}
+ * <li>{@link FieldType#STRING}
+ * </ul>
+ */
+ public static final Set<Schema.FieldType> VALID_FIELD_TYPE_SET =
+ ImmutableSet.of(
+ FieldType.BYTE,
+ FieldType.BOOLEAN,
+ FieldType.DATETIME,
+ FieldType.DECIMAL,
+ FieldType.DOUBLE,
+ FieldType.INT16,
+ FieldType.INT32,
+ FieldType.INT64,
+ FieldType.FLOAT,
+ FieldType.STRING);
+
+ static final String DEFAULT_FILENAME_SUFFIX = ".csv";
+
+ /** Instantiates a {@link Write} for writing user types in {@link CSVFormat} format. */
+ public static <T> Write<T> write(String to, CSVFormat csvFormat) {
+ return new AutoValue_CsvIO_Write.Builder<T>()
+ .setTextIOWrite(createDefaultTextIOWrite(to))
+ .setCSVFormat(csvFormat)
+ .build();
+ }
+
+ /** Instantiates a {@link Write} for writing {@link Row}s in {@link CSVFormat} format. */
+ public static Write<Row> writeRows(String to, CSVFormat csvFormat) {
+ return new AutoValue_CsvIO_Write.Builder<Row>()
+ .setTextIOWrite(createDefaultTextIOWrite(to))
+ .setCSVFormat(csvFormat)
+ .build();
+ }
+
+ /** {@link PTransform} for writing CSV files. */
+ @AutoValue
+ public abstract static class Write<T>
+ extends PTransform<PCollection<T>, WriteFilesResult<String>> {
+
+ /** Specifies the {@link Compression} of all generated shard files. */
+ public Write<T> withCompression(Compression compression) {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withCompression(compression)).build();
+ }
+
+ public Write<T> withNoSpilling() {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withNoSpilling()).build();
+ }
+
+ /**
+ * Specifies to use a given fixed number of shards per window. See {@link
+ * TextIO.Write#withNumShards(int)}.
+ */
+ public Write<T> withNumShards(Integer numShards) {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withNumShards(numShards)).build();
+ }
+
+ /**
+ * Forces a single file as output and empty shard name template. See {@link
+ * TextIO.Write#withoutSharding()}.
+ */
+ public Write<T> withoutSharding() {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withoutSharding()).build();
+ }
+
+ /**
+ * Uses the given {@link ShardNameTemplate} for naming output files. See {@link
+ * TextIO.Write#withShardNameTemplate(String)}.
+ */
+ public Write<T> withShardTemplate(String shardTemplate) {
+ return toBuilder()
+ .setTextIOWrite(getTextIOWrite().withShardNameTemplate(shardTemplate))
+ .build();
+ }
+
+ /**
+ * Configures the filename suffix for written files. See {@link
+ * TextIO.Write#withSuffix(String)}.
+ */
+ public Write<T> withSuffix(String suffix) {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withSuffix(suffix)).build();
+ }
+
+ /**
+ * Set the base directory used to generate temporary files. See {@link
+ * TextIO.Write#withTempDirectory(ResourceId)}.
+ */
+ public Write<T> withTempDirectory(ResourceId tempDirectory) {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withTempDirectory(tempDirectory)).build();
+ }
+
+ /**
+ * Preserves windowing of input elements and writes them to files based on the element's window.
+ * See {@link TextIO.Write#withWindowedWrites()}.
+ */
+ public Write<T> withWindowedWrites() {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withWindowedWrites()).build();
+ }
+
+ /**
+ * Returns a transform for writing to text files like this one but that has the given {@link
+ * FileBasedSink.WritableByteChannelFactory} to be used by the {@link FileBasedSink} during
+ * output. See {@link
+ * TextIO.Write#withWritableByteChannelFactory(FileBasedSink.WritableByteChannelFactory)}.
+ */
+ public Write<T> withWritableByteChannelFactory(
+ FileBasedSink.WritableByteChannelFactory writableByteChannelFactory) {
+ return toBuilder()
+ .setTextIOWrite(
+ getTextIOWrite().withWritableByteChannelFactory(writableByteChannelFactory))
+ .build();
+ }
+
+ /** The underlying {@link FileIO.Write} that writes converted input to CSV formatted output. */
+ abstract TextIO.Write getTextIOWrite();
+
+ /** The {@link CSVFormat} to inform headers, header comments, and format CSV row content. */
+ abstract CSVFormat getCSVFormat();
+
+ abstract Builder<T> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<T> {
+
+ /**
+ * The underlying {@link FileIO.Write} that writes converted input to CSV formatted output.
+ */
+ abstract Builder<T> setTextIOWrite(TextIO.Write value);
+
+ /** The {@link CSVFormat} to convert input. Defaults to {@link CSVFormat#DEFAULT}. */
+ abstract Builder<T> setCSVFormat(CSVFormat value);
+
+ abstract CSVFormat getCSVFormat();
+
+ abstract Write<T> autoBuild();
+
+ final Write<T> build() {
+ checkArgument(
+ !getCSVFormat().getSkipHeaderRecord(),
+ "withSkipHeaderRecord is an illegal CSVFormat setting");
+
+ if (getCSVFormat().getHeaderComments() != null) {
+ checkArgument(
+ getCSVFormat().isCommentMarkerSet(),
+ "CSVFormat withCommentMarker required when withHeaderComments");
+ }
+ return autoBuild();
+ }
+ }
+
+ @Override
+ public WriteFilesResult<String> expand(PCollection<T> input) {
+ if (!input.hasSchema()) {
+ throw new IllegalArgumentException(
+ String.format(
+ "%s requires an input Schema. Note that only Row or user classes are supported. Consider using TextIO or FileIO directly when writing primitive types",
+ Write.class.getName()));
+ }
+
+ Schema schema = input.getSchema();
+
+ RowCoder rowCoder = RowCoder.of(schema);
+
+ PCollection<Row> rows =
+ input
+ .apply("To Rows", MapElements.into(rows()).via(input.getToRowFunction()))
+ .setCoder(rowCoder);
+
+ CSVFormat csvFormat = applyRequiredCSVFormatSettings(schema);
+
+ String header = formatHeader(csvFormat);
+
+ SerializableFunction<Row, String> toCsvFn =
+ CsvRowConversions.RowToCsv.builder()
+ .setCSVFormat(csvFormat)
+ .setSchema(input.getSchema())
+ .build();
+
+ PCollection<String> csv = rows.apply("To CSV", MapElements.into(strings()).via(toCsvFn));
+
+ return csv.apply("Write CSV", getTextIOWrite().withHeader(header).withOutputFilenames());
+ }
+
+ CSVFormat applyRequiredCSVFormatSettings(Schema schema) {
+ CSVFormat csvFormat = getCSVFormat().withSkipHeaderRecord();
+ if (csvFormat.getHeader() == null) {
+ csvFormat = csvFormat.withHeader(schema.sorted().getFieldNames().toArray(new String[0]));
+ }
+ return csvFormat;
+ }
+
+ private static String formatHeader(CSVFormat csvFormat) {
+ String[] header = requireNonNull(csvFormat.getHeader());
+ List<String> result = new ArrayList<>();
+ if (csvFormat.getHeaderComments() != null) {
+ for (String comment : csvFormat.getHeaderComments()) {
+ result.add(csvFormat.getCommentMarker() + " " + comment);
+ }
+ }
+ CSVFormat withoutHeaderComments = csvFormat.withHeaderComments();
+ result.add(withoutHeaderComments.format((Object[]) header));
+ return String.join("\n", result);
+ }
Review Comment:
We are ignoring the case where the user wants to skip the header record and doesn't specify the header. We should only specify the `csvFormat` header if `withSkipHeaderRecord` is `false`.
Note that we must also return `null` from formatHeader in this case since `TextIO#withHeader` specifically says that a newline is automatically added after.
Add a test case for this scenario.
```suggestion
CSVFormat applyRequiredCSVFormatSettings(Schema schema) {
CSVFormat csvFormat = getCSVFormat();
if (!csvFormat.getSkipHeaderRecord() && csvFormat.getHeader() == null) {
return csvFormat.withSkipHeaderRecord().withHeader(schema.sorted().getFieldNames().toArray(new String[0]));
}
return csvFormat;
}
private static @Nullable String formatHeader(CSVFormat csvFormat) {
String[] header = requireNonNull(csvFormat.getHeader());
if (header == null) {
return null;
}
List<String> result = new ArrayList<>();
if (csvFormat.getHeaderComments() != null) {
for (String comment : csvFormat.getHeaderComments()) {
result.add(csvFormat.getCommentMarker() + " " + comment);
}
}
CSVFormat withoutHeaderComments = csvFormat.withHeaderComments();
result.add(withoutHeaderComments.format((Object[]) header));
return String.join("\n", result);
}
```
##########
sdks/java/io/csv/src/test/java/org/apache/beam/sdk/io/csv/CsvIOWriteTest.java:
##########
@@ -0,0 +1,299 @@
+/*
+ * 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.sdk.io.csv;
+
+import static java.util.Objects.requireNonNull;
+import static org.apache.beam.sdk.io.csv.CsvIOTestData.DATA;
+import static org.apache.beam.sdk.io.csv.CsvIOTestJavaBeans.ALL_PRIMITIVE_DATA_TYPES_SCHEMA;
+import static org.apache.beam.sdk.io.csv.CsvIOTestJavaBeans.AllPrimitiveDataTypes;
+import static org.apache.beam.sdk.io.csv.CsvIOTestJavaBeans.allPrimitiveDataTypes;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertThrows;
+import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.fail;
+
+import java.io.File;
+import java.io.IOException;
+import java.math.BigDecimal;
+import java.util.List;
+import org.apache.beam.sdk.io.FileIO;
+import org.apache.beam.sdk.io.TextIO;
+import org.apache.beam.sdk.testing.PAssert;
+import org.apache.beam.sdk.testing.TestPipeline;
+import org.apache.beam.sdk.transforms.Create;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Splitter;
+import org.apache.commons.csv.CSVFormat;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.TemporaryFolder;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/** Tests for {@link CsvIO.Write}. */
+@RunWith(JUnit4.class)
+public class CsvIOWriteTest {
+ @Rule public TestPipeline writePipeline = TestPipeline.create();
+
+ @Rule public TestPipeline readPipeline = TestPipeline.create();
+
+ @Rule
+ public TestPipeline errorPipeline = TestPipeline.create().enableAbandonedNodeEnforcement(false);
+
+ @Rule public TemporaryFolder tempFolder = new TemporaryFolder();
+
+ @Test
+ public void headersWithCommentsWrittenFirstOnEachShard() {
+ File folder =
+ createFolder(
+ AllPrimitiveDataTypes.class.getSimpleName(),
+ "headersWithCommentsWrittenFirstOnEachShard");
+
+ PCollection<Row> input =
+ writePipeline.apply(
+ Create.of(DATA.allPrimitiveDataTypeRows)
+ .withRowSchema(ALL_PRIMITIVE_DATA_TYPES_SCHEMA));
+ String expectedHeader = "aBoolean,aByte,aDecimal,aDouble,aFloat,aLong,aShort,aString,anInteger";
+ CSVFormat csvFormat =
+ CSVFormat.DEFAULT.withHeaderComments("foo", "bar", "baz").withCommentMarker('#');
+
+ input.apply(CsvIO.writeRows(toFilenamePrefix(folder), csvFormat).withNumShards(3));
+ writePipeline.run().waitUntilFinish();
+
+ PCollection<FileIO.ReadableFile> files =
+ readPipeline
+ .apply(FileIO.match().filepattern(toFilenamePrefix(folder) + "*"))
+ .apply(FileIO.readMatches());
+ PAssert.that(files)
+ .satisfies(
+ (Iterable<FileIO.ReadableFile> itr) -> {
+ Iterable<FileIO.ReadableFile> safeItr = requireNonNull(itr);
+ for (FileIO.ReadableFile file : safeItr) {
+ try {
+ List<String> lines = Splitter.on('\n').splitToList(file.readFullyAsUTF8String());
+ assertFalse(lines.isEmpty());
+ assertEquals("# foo", lines.get(0));
+ assertEquals("# bar", lines.get(1));
+ assertEquals("# baz", lines.get(2));
+ assertEquals(expectedHeader, lines.get(3));
+
+ assertTrue(
+ lines.subList(4, lines.size()).stream().noneMatch(expectedHeader::equals));
+
Review Comment:
```suggestion
```
##########
sdks/java/io/csv/src/main/java/org/apache/beam/sdk/io/csv/CsvIO.java:
##########
@@ -0,0 +1,446 @@
+/*
+ * 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.sdk.io.csv;
+
+import static java.util.Objects.requireNonNull;
+import static org.apache.beam.sdk.values.TypeDescriptors.rows;
+import static org.apache.beam.sdk.values.TypeDescriptors.strings;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+
+import com.google.auto.value.AutoValue;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Set;
+import org.apache.beam.sdk.coders.RowCoder;
+import org.apache.beam.sdk.io.Compression;
+import org.apache.beam.sdk.io.FileBasedSink;
+import org.apache.beam.sdk.io.FileIO;
+import org.apache.beam.sdk.io.ShardNameTemplate;
+import org.apache.beam.sdk.io.TextIO;
+import org.apache.beam.sdk.io.WriteFilesResult;
+import org.apache.beam.sdk.io.fs.ResourceId;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.transforms.MapElements;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.sdk.values.TypeDescriptor;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableSet;
+import org.apache.commons.csv.CSVFormat;
+
+/**
+ * {@link PTransform}s for reading and writing CSV files.
+ *
+ * <h2>Reading CSV files</h2>
+ *
+ * <p>Reading from CSV files is not yet implemented. Please see <a
+ * href="https://github.com/apache/beam/issues/24552">https://github.com/apache/beam/issues/24552</a>.
+ *
+ * <h2>Writing CSV files</h2>
+ *
+ * <p>To write a {@link PCollection} to one or more CSV files, use {@link CsvIO.Write}, using {@link
+ * CsvIO#writeRows(String, CSVFormat)} or {@link CsvIO#write(String, CSVFormat)}. {@link
+ * CsvIO.Write} supports writing {@link Row} or custom Java types using an inferred {@link Schema}.
+ * Examples below show both scenarios. See the Beam Programming Guide on <a
+ * href="https://beam.apache.org/documentation/programming-guide/#inferring-schemas">inferring
+ * schemas</a> for more information on how to enable Beam to infer a {@link Schema} from a custom
+ * Java type.
+ *
+ * <p>{@link CsvIO.Write} only supports writing {@link Schema} aware types that do not contain any
+ * nested {@link FieldType}s such a {@link org.apache.beam.sdk.schemas.Schema.TypeName#ROW} or
+ * repeated {@link org.apache.beam.sdk.schemas.Schema.TypeName#ARRAY} types. See {@link
+ * CsvIO.Write#VALID_FIELD_TYPE_SET} for valid {@link FieldType}s.
+ *
+ * <h3>Example usage:</h3>
+ *
+ * <p>Suppose we have the following <code>Transaction</code> class annotated with
+ * {@code @DefaultSchema(JavaBeanSchema.class)} so that Beam can infer its {@link Schema}:
+ *
+ * <pre>{@code @DefaultSchema(JavaBeanSchema.class)
+ * public class Transaction {
+ * public Transaction() { … }
+ * public Long getTransactionId();
+ * public void setTransactionId(Long transactionId) { … }
+ * public String getBank() { … }
+ * public void setBank(String bank) { … }
+ * public double getPurchaseAmount() { … }
+ * public void setPurchaseAmount(double purchaseAmount) { … }
+ * }
+ * }</pre>
+ *
+ * <p>From a {@code PCollection<Transaction>}, {@link CsvIO.Write} can write one or many CSV files
+ * automatically creating the header based on its inferred {@link Schema}.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions = ...
+ * transactions.apply(CsvIO.<Transaction>write("path/to/folder/prefix", CSVFormat.DEFAULT));
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header is repeated for every
+ * shard file, whereas by default, {@link CsvIO.Write} will write all fields in <b>sorted order</b>
+ * of the field names.
+ *
+ * <pre>{@code
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ *
+ * <p>To control the order and subset of fields that {@link CsvIO.Write} writes, use {@link
+ * CSVFormat#withHeader(String...)}. Note, however, the following constraints:
+ *
+ * <ol>
+ * <li>Each header column must match a field name in the {@link Schema}; matching is case
+ * sensitive.
+ * <li>Matching header columns must match {@link Schema} fields that are valid {@link FieldType}s;
+ * see {@link #VALID_FIELD_TYPE_SET}.
+ * <li>{@link CSVFormat} only allows repeated header columns when {@link
+ * CSVFormat#withAllowDuplicateHeaderNames()}
+ * </ol>
+ *
+ * <p>The following example shows the use of {@link CSVFormat#withHeader(String...)} to control the
+ * order and subset of <code>Transaction</code> fields.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions ...
+ * transactions.apply(
+ * CsvIO
+ * .<Transaction>write("path/to/folder/prefix", CSVFormat.DEFAULT.withHeader("transactionId", "purchaseAmount"))
+ * );
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header is repeated for every
+ * shard file, but only including the subset of fields in their listed order.
+ *
+ * <pre>{@code
+ * transactionId,purchaseAmount
+ * 12345,10.23
+ * 54321,54.65
+ * 98765,11.76
+ * }</pre>
+ *
+ * <p>In addition to header customization, {@link CsvIO.Write} supports {@link
+ * CSVFormat#withHeaderComments(Object...)} as shown below. Note that {@link
+ * CSVFormat#withCommentMarker(char)} is required when specifying header comments.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions = ...
+ * transactions
+ * .apply(
+ * CsvIO.<Transaction>write("path/to/folder/prefix",
+ * CSVFormat.DEFAULT.withCommentMarker('#').withHeaderComments("Bank Report", "1970-01-01", "Operator: John Doe")
+ * );
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header and header comments are
+ * repeated for every shard file.
+ *
+ * <pre>{@code
+ * # Bank Report
+ * # 1970-01-01
+ * # Operator: John Doe
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ *
+ * <p>A {@link PCollection} of {@link Row}s works just like custom Java types illustrated above,
+ * except we use {@link CsvIO#writeRows(String, CSVFormat)} as shown below for the same {@code
+ * Transaction} class. We derive {@code Transaction}'s {@link Schema} using a {@link
+ * org.apache.beam.sdk.schemas.annotations.DefaultSchema.DefaultSchemaProvider}. Note that
+ * hard-coding the {@link Row}s below is for illustration purposes. Developers are instead
+ * encouraged to take advantage of {@link
+ * org.apache.beam.sdk.schemas.annotations.DefaultSchema.DefaultSchemaProvider#toRowFunction(TypeDescriptor)}.
+ *
+ * <pre>{@code
+ * DefaultSchemaProvider defaultSchemaProvider = new DefaultSchemaProvider();
+ * Schema schema = defaultSchemaProvider.schemaFor(TypeDescriptor.of(Transaction.class));
+ * PCollection<Row> transactions = pipeline.apply(Create.of(
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "A")
+ * .withFieldValue("purchaseAmount", 10.23)
+ * .withFieldValue("transactionId", "12345")
+ * .build(),
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "B")
+ * .withFieldValue("purchaseAmount", 54.65)
+ * .withFieldValue("transactionId", "54321")
+ * .build(),
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "C")
+ * .withFieldValue("purchaseAmount", 11.76)
+ * .withFieldValue("transactionId", "98765")
+ * .build()
+ * );
+ *
+ * transactions.apply(
+ * CsvIO
+ * .writeRowsTo("gs://bucket/path/to/folder/prefix", CSVFormat.DEFAULT)
+ * );
+ * }</pre>
+ *
+ * <p>Writing the transactions {@link PCollection} of {@link Row}s would yield the following CSV
+ * file content.
+ *
+ * <pre>{@code
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ *
+ * {@link CsvIO.Write} does not support the following {@link CSVFormat} properties and will throw an
+ * {@link IllegalArgumentException}.
+ *
+ * <ul>
+ * <li>{@link CSVFormat#withSkipHeaderRecord()}
+ * <li>{@link CSVFormat#withAllowMissingColumnNames()} ()}
+ * <li>{@link CSVFormat#withAutoFlush(boolean)}
+ * <li>{@link CSVFormat#withIgnoreHeaderCase()}
+ * <li>{@link CSVFormat#withIgnoreSurroundingSpaces()}
+ * </ul>
+ */
+public class CsvIO {
+ /**
+ * The valid {@link Schema.FieldType} from which {@link CsvIO} converts CSV records to the fields.
+ *
+ * <ul>
+ * <li>{@link FieldType#BYTE}
+ * <li>{@link FieldType#BOOLEAN}
+ * <li>{@link FieldType#DATETIME}
+ * <li>{@link FieldType#DECIMAL}
+ * <li>{@link FieldType#DOUBLE}
+ * <li>{@link FieldType#INT16}
+ * <li>{@link FieldType#INT32}
+ * <li>{@link FieldType#INT64}
+ * <li>{@link FieldType#FLOAT}
+ * <li>{@link FieldType#STRING}
+ * </ul>
+ */
+ public static final Set<Schema.FieldType> VALID_FIELD_TYPE_SET =
+ ImmutableSet.of(
+ FieldType.BYTE,
+ FieldType.BOOLEAN,
+ FieldType.DATETIME,
+ FieldType.DECIMAL,
+ FieldType.DOUBLE,
+ FieldType.INT16,
+ FieldType.INT32,
+ FieldType.INT64,
+ FieldType.FLOAT,
+ FieldType.STRING);
+
+ static final String DEFAULT_FILENAME_SUFFIX = ".csv";
+
+ /** Instantiates a {@link Write} for writing user types in {@link CSVFormat} format. */
+ public static <T> Write<T> write(String to, CSVFormat csvFormat) {
+ return new AutoValue_CsvIO_Write.Builder<T>()
+ .setTextIOWrite(createDefaultTextIOWrite(to))
+ .setCSVFormat(csvFormat)
+ .build();
+ }
+
+ /** Instantiates a {@link Write} for writing {@link Row}s in {@link CSVFormat} format. */
+ public static Write<Row> writeRows(String to, CSVFormat csvFormat) {
+ return new AutoValue_CsvIO_Write.Builder<Row>()
+ .setTextIOWrite(createDefaultTextIOWrite(to))
+ .setCSVFormat(csvFormat)
+ .build();
+ }
+
+ /** {@link PTransform} for writing CSV files. */
+ @AutoValue
+ public abstract static class Write<T>
+ extends PTransform<PCollection<T>, WriteFilesResult<String>> {
+
+ /** Specifies the {@link Compression} of all generated shard files. */
+ public Write<T> withCompression(Compression compression) {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withCompression(compression)).build();
+ }
+
+ public Write<T> withNoSpilling() {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withNoSpilling()).build();
+ }
+
+ /**
+ * Specifies to use a given fixed number of shards per window. See {@link
+ * TextIO.Write#withNumShards(int)}.
+ */
+ public Write<T> withNumShards(Integer numShards) {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withNumShards(numShards)).build();
+ }
+
+ /**
+ * Forces a single file as output and empty shard name template. See {@link
+ * TextIO.Write#withoutSharding()}.
+ */
+ public Write<T> withoutSharding() {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withoutSharding()).build();
+ }
+
+ /**
+ * Uses the given {@link ShardNameTemplate} for naming output files. See {@link
+ * TextIO.Write#withShardNameTemplate(String)}.
+ */
+ public Write<T> withShardTemplate(String shardTemplate) {
+ return toBuilder()
+ .setTextIOWrite(getTextIOWrite().withShardNameTemplate(shardTemplate))
+ .build();
+ }
+
+ /**
+ * Configures the filename suffix for written files. See {@link
+ * TextIO.Write#withSuffix(String)}.
+ */
+ public Write<T> withSuffix(String suffix) {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withSuffix(suffix)).build();
+ }
+
+ /**
+ * Set the base directory used to generate temporary files. See {@link
+ * TextIO.Write#withTempDirectory(ResourceId)}.
+ */
+ public Write<T> withTempDirectory(ResourceId tempDirectory) {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withTempDirectory(tempDirectory)).build();
+ }
+
+ /**
+ * Preserves windowing of input elements and writes them to files based on the element's window.
+ * See {@link TextIO.Write#withWindowedWrites()}.
+ */
+ public Write<T> withWindowedWrites() {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withWindowedWrites()).build();
+ }
+
+ /**
+ * Returns a transform for writing to text files like this one but that has the given {@link
+ * FileBasedSink.WritableByteChannelFactory} to be used by the {@link FileBasedSink} during
+ * output. See {@link
+ * TextIO.Write#withWritableByteChannelFactory(FileBasedSink.WritableByteChannelFactory)}.
+ */
+ public Write<T> withWritableByteChannelFactory(
+ FileBasedSink.WritableByteChannelFactory writableByteChannelFactory) {
+ return toBuilder()
+ .setTextIOWrite(
+ getTextIOWrite().withWritableByteChannelFactory(writableByteChannelFactory))
+ .build();
+ }
+
+ /** The underlying {@link FileIO.Write} that writes converted input to CSV formatted output. */
+ abstract TextIO.Write getTextIOWrite();
+
+ /** The {@link CSVFormat} to inform headers, header comments, and format CSV row content. */
+ abstract CSVFormat getCSVFormat();
+
+ abstract Builder<T> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<T> {
+
+ /**
+ * The underlying {@link FileIO.Write} that writes converted input to CSV formatted output.
+ */
+ abstract Builder<T> setTextIOWrite(TextIO.Write value);
+
+ /** The {@link CSVFormat} to convert input. Defaults to {@link CSVFormat#DEFAULT}. */
+ abstract Builder<T> setCSVFormat(CSVFormat value);
+
+ abstract CSVFormat getCSVFormat();
+
+ abstract Write<T> autoBuild();
+
+ final Write<T> build() {
+ checkArgument(
+ !getCSVFormat().getSkipHeaderRecord(),
Review Comment:
You specify a few other properties in the javadoc for things that aren't supported, should we check for them here as well OR it seems as though CSVFormat has a mix of properties used during reading, writing or both and instead of checking here we could instead just list properties that are ignored during writing in the javadoc (this way a person can use the same CSVFormat object when reading and writing)?
##########
sdks/java/io/csv/src/main/java/org/apache/beam/sdk/io/csv/CsvRowConversions.java:
##########
@@ -0,0 +1,148 @@
+/*
+ * 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.sdk.io.csv;
+
+import static org.apache.beam.sdk.io.csv.CsvIO.VALID_FIELD_TYPE_SET;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkNotNull;
+
+import com.google.auto.value.AutoValue;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Set;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.schemas.Schema.TypeName;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.values.Row;
+import org.apache.commons.csv.CSVFormat;
+import org.checkerframework.checker.nullness.qual.NonNull;
+
+/** Contains classes and methods to help with converting between {@link Row} and CSV strings. */
+class CsvRowConversions {
+
+ /** Converts between {@link Row} and CSV string using a {@link CSVFormat}. */
+ @AutoValue
+ abstract static class RowToCsv implements SerializableFunction<Row, String> {
+
+ static Builder builder() {
+ return new AutoValue_CsvRowConversions_RowToCsv.Builder();
+ }
+
+ /** The expected {@link Schema} of the {@link Row} input. */
+ abstract Schema getSchema();
+
+ /** The {@link CSVFormat} of the converted {@link Row} input. */
+ abstract CSVFormat getCSVFormat();
+
+ /** Converts a {@link Row} to a CSV string formatted using {@link #getCSVFormat()}. */
+ @Override
+ public String apply(Row input) {
+ Row safeInput = checkNotNull(input);
+ String[] header = getHeader();
+ Object[] values = new Object[header.length];
+ for (int i = 0; i < header.length; i++) {
+ values[i] = safeInput.getValue(header[i]);
+ }
+ return getCSVFormat().format(values);
+ }
+
+ @NonNull
+ String[] getHeader() {
+ return checkNotNull(getCSVFormat().getHeader());
+ }
+
+ @AutoValue.Builder
+ abstract static class Builder {
+
+ /** The expected {@link Schema} of the {@link Row} input. */
+ abstract Builder setSchema(Schema schema);
+
+ abstract Schema getSchema();
+
+ /** The {@link CSVFormat} of the converted {@link Row} input. */
+ abstract Builder setCSVFormat(CSVFormat format);
+
+ abstract CSVFormat getCSVFormat();
+
+ abstract RowToCsv autoBuild();
+
+ final RowToCsv build() {
+ checkArgument(getSchema().getFieldCount() > 0, "Schema has no fields");
+ setCSVFormat(
+ getCSVFormat()
+ .withSkipHeaderRecord()
+ // Delegate to TextIO.Write.withDelimiter instead.
+ .withRecordSeparator((char) 32)
+ .withHeaderComments());
+ validateCSVFormat(getCSVFormat());
+ validateHeaderAgainstSchema(getCSVFormat().getHeader(), getSchema());
+
+ return autoBuild();
+ }
+ }
+ }
+
+ private static void validateCSVFormat(CSVFormat csvFormat) {
+ String[] header = checkNotNull(csvFormat.getHeader(), "CSVFormat withHeader is required");
+
+ checkArgument(header.length > 0, "CSVFormat withHeader requires at least one column");
+
+ checkArgument(!csvFormat.getAutoFlush(), "withAutoFlush is an illegal CSVFormat setting");
+
+ checkArgument(
+ !csvFormat.getIgnoreHeaderCase(), "withIgnoreHeaderCase is an illegal CSVFormat setting");
+
+ checkArgument(
+ !csvFormat.getAllowMissingColumnNames(),
+ "withAllowMissingColumnNames is an illegal CSVFormat setting");
+
+ checkArgument(
+ !csvFormat.getIgnoreSurroundingSpaces(),
+ "withIgnoreSurroundingSpaces is an illegal CSVFormat setting");
+ }
+
+ private static void validateHeaderAgainstSchema(String[] csvHeader, Schema schema) {
+ Set<String> distinctColumns = Stream.of(csvHeader).collect(Collectors.toSet());
Review Comment:
```suggestion
Set<String> distinctColumns = new HashSet<>(Arrays.asList(csvHeader));
```
##########
sdks/java/io/csv/src/main/java/org/apache/beam/sdk/io/csv/CsvIO.java:
##########
@@ -0,0 +1,446 @@
+/*
+ * 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.sdk.io.csv;
+
+import static java.util.Objects.requireNonNull;
+import static org.apache.beam.sdk.values.TypeDescriptors.rows;
+import static org.apache.beam.sdk.values.TypeDescriptors.strings;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+
+import com.google.auto.value.AutoValue;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Set;
+import org.apache.beam.sdk.coders.RowCoder;
+import org.apache.beam.sdk.io.Compression;
+import org.apache.beam.sdk.io.FileBasedSink;
+import org.apache.beam.sdk.io.FileIO;
+import org.apache.beam.sdk.io.ShardNameTemplate;
+import org.apache.beam.sdk.io.TextIO;
+import org.apache.beam.sdk.io.WriteFilesResult;
+import org.apache.beam.sdk.io.fs.ResourceId;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.transforms.MapElements;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.sdk.values.TypeDescriptor;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableSet;
+import org.apache.commons.csv.CSVFormat;
+
+/**
+ * {@link PTransform}s for reading and writing CSV files.
+ *
+ * <h2>Reading CSV files</h2>
+ *
+ * <p>Reading from CSV files is not yet implemented. Please see <a
+ * href="https://github.com/apache/beam/issues/24552">https://github.com/apache/beam/issues/24552</a>.
+ *
+ * <h2>Writing CSV files</h2>
+ *
+ * <p>To write a {@link PCollection} to one or more CSV files, use {@link CsvIO.Write}, using {@link
+ * CsvIO#writeRows(String, CSVFormat)} or {@link CsvIO#write(String, CSVFormat)}. {@link
+ * CsvIO.Write} supports writing {@link Row} or custom Java types using an inferred {@link Schema}.
+ * Examples below show both scenarios. See the Beam Programming Guide on <a
+ * href="https://beam.apache.org/documentation/programming-guide/#inferring-schemas">inferring
+ * schemas</a> for more information on how to enable Beam to infer a {@link Schema} from a custom
+ * Java type.
+ *
+ * <p>{@link CsvIO.Write} only supports writing {@link Schema} aware types that do not contain any
+ * nested {@link FieldType}s such a {@link org.apache.beam.sdk.schemas.Schema.TypeName#ROW} or
+ * repeated {@link org.apache.beam.sdk.schemas.Schema.TypeName#ARRAY} types. See {@link
+ * CsvIO.Write#VALID_FIELD_TYPE_SET} for valid {@link FieldType}s.
+ *
+ * <h3>Example usage:</h3>
+ *
+ * <p>Suppose we have the following <code>Transaction</code> class annotated with
+ * {@code @DefaultSchema(JavaBeanSchema.class)} so that Beam can infer its {@link Schema}:
+ *
+ * <pre>{@code @DefaultSchema(JavaBeanSchema.class)
+ * public class Transaction {
+ * public Transaction() { … }
+ * public Long getTransactionId();
+ * public void setTransactionId(Long transactionId) { … }
+ * public String getBank() { … }
+ * public void setBank(String bank) { … }
+ * public double getPurchaseAmount() { … }
+ * public void setPurchaseAmount(double purchaseAmount) { … }
+ * }
+ * }</pre>
+ *
+ * <p>From a {@code PCollection<Transaction>}, {@link CsvIO.Write} can write one or many CSV files
+ * automatically creating the header based on its inferred {@link Schema}.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions = ...
+ * transactions.apply(CsvIO.<Transaction>write("path/to/folder/prefix", CSVFormat.DEFAULT));
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header is repeated for every
+ * shard file, whereas by default, {@link CsvIO.Write} will write all fields in <b>sorted order</b>
Review Comment:
```suggestion
* file, whereas by default, {@link CsvIO.Write} will write all fields in <b>sorted order</b>
```
##########
sdks/java/io/csv/src/main/java/org/apache/beam/sdk/io/csv/CsvRowConversions.java:
##########
@@ -0,0 +1,148 @@
+/*
+ * 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.sdk.io.csv;
+
+import static org.apache.beam.sdk.io.csv.CsvIO.VALID_FIELD_TYPE_SET;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkNotNull;
+
+import com.google.auto.value.AutoValue;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Set;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.schemas.Schema.TypeName;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.values.Row;
+import org.apache.commons.csv.CSVFormat;
+import org.checkerframework.checker.nullness.qual.NonNull;
+
+/** Contains classes and methods to help with converting between {@link Row} and CSV strings. */
+class CsvRowConversions {
+
+ /** Converts between {@link Row} and CSV string using a {@link CSVFormat}. */
+ @AutoValue
+ abstract static class RowToCsv implements SerializableFunction<Row, String> {
+
+ static Builder builder() {
+ return new AutoValue_CsvRowConversions_RowToCsv.Builder();
+ }
+
+ /** The expected {@link Schema} of the {@link Row} input. */
+ abstract Schema getSchema();
+
+ /** The {@link CSVFormat} of the converted {@link Row} input. */
+ abstract CSVFormat getCSVFormat();
+
+ /** Converts a {@link Row} to a CSV string formatted using {@link #getCSVFormat()}. */
+ @Override
+ public String apply(Row input) {
+ Row safeInput = checkNotNull(input);
+ String[] header = getHeader();
+ Object[] values = new Object[header.length];
+ for (int i = 0; i < header.length; i++) {
+ values[i] = safeInput.getValue(header[i]);
+ }
+ return getCSVFormat().format(values);
+ }
+
+ @NonNull
+ String[] getHeader() {
+ return checkNotNull(getCSVFormat().getHeader());
+ }
+
+ @AutoValue.Builder
+ abstract static class Builder {
+
+ /** The expected {@link Schema} of the {@link Row} input. */
+ abstract Builder setSchema(Schema schema);
+
+ abstract Schema getSchema();
+
+ /** The {@link CSVFormat} of the converted {@link Row} input. */
+ abstract Builder setCSVFormat(CSVFormat format);
+
+ abstract CSVFormat getCSVFormat();
+
+ abstract RowToCsv autoBuild();
+
+ final RowToCsv build() {
+ checkArgument(getSchema().getFieldCount() > 0, "Schema has no fields");
+ setCSVFormat(
+ getCSVFormat()
+ .withSkipHeaderRecord()
Review Comment:
Don't we do this in CsvIO already?, why repeat here?
--
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: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [beam] damondouglas commented on a diff in pull request #24630: [24469] Implement CsvIO.Write and supporting classes
Posted by "damondouglas (via GitHub)" <gi...@apache.org>.
damondouglas commented on code in PR #24630:
URL: https://github.com/apache/beam/pull/24630#discussion_r1085716045
##########
sdks/java/io/csv/src/main/java/org/apache/beam/sdk/io/csv/CsvIO.java:
##########
@@ -0,0 +1,446 @@
+/*
+ * 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.sdk.io.csv;
+
+import static java.util.Objects.requireNonNull;
+import static org.apache.beam.sdk.values.TypeDescriptors.rows;
+import static org.apache.beam.sdk.values.TypeDescriptors.strings;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+
+import com.google.auto.value.AutoValue;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Set;
+import org.apache.beam.sdk.coders.RowCoder;
+import org.apache.beam.sdk.io.Compression;
+import org.apache.beam.sdk.io.FileBasedSink;
+import org.apache.beam.sdk.io.FileIO;
+import org.apache.beam.sdk.io.ShardNameTemplate;
+import org.apache.beam.sdk.io.TextIO;
+import org.apache.beam.sdk.io.WriteFilesResult;
+import org.apache.beam.sdk.io.fs.ResourceId;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.transforms.MapElements;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.sdk.values.TypeDescriptor;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableSet;
+import org.apache.commons.csv.CSVFormat;
+
+/**
+ * {@link PTransform}s for reading and writing CSV files.
+ *
+ * <h2>Reading CSV files</h2>
+ *
+ * <p>Reading from CSV files is not yet implemented. Please see <a
+ * href="https://github.com/apache/beam/issues/24552">https://github.com/apache/beam/issues/24552</a>.
+ *
+ * <h2>Writing CSV files</h2>
+ *
+ * <p>To write a {@link PCollection} to one or more CSV files, use {@link CsvIO.Write}, using {@link
+ * CsvIO#writeRows(String, CSVFormat)} or {@link CsvIO#write(String, CSVFormat)}. {@link
+ * CsvIO.Write} supports writing {@link Row} or custom Java types using an inferred {@link Schema}.
+ * Examples below show both scenarios. See the Beam Programming Guide on <a
+ * href="https://beam.apache.org/documentation/programming-guide/#inferring-schemas">inferring
+ * schemas</a> for more information on how to enable Beam to infer a {@link Schema} from a custom
+ * Java type.
+ *
+ * <p>{@link CsvIO.Write} only supports writing {@link Schema} aware types that do not contain any
+ * nested {@link FieldType}s such a {@link org.apache.beam.sdk.schemas.Schema.TypeName#ROW} or
+ * repeated {@link org.apache.beam.sdk.schemas.Schema.TypeName#ARRAY} types. See {@link
+ * CsvIO.Write#VALID_FIELD_TYPE_SET} for valid {@link FieldType}s.
+ *
+ * <h3>Example usage:</h3>
+ *
+ * <p>Suppose we have the following <code>Transaction</code> class annotated with
+ * {@code @DefaultSchema(JavaBeanSchema.class)} so that Beam can infer its {@link Schema}:
+ *
+ * <pre>{@code @DefaultSchema(JavaBeanSchema.class)
+ * public class Transaction {
+ * public Transaction() { … }
+ * public Long getTransactionId();
+ * public void setTransactionId(Long transactionId) { … }
+ * public String getBank() { … }
+ * public void setBank(String bank) { … }
+ * public double getPurchaseAmount() { … }
+ * public void setPurchaseAmount(double purchaseAmount) { … }
+ * }
+ * }</pre>
+ *
+ * <p>From a {@code PCollection<Transaction>}, {@link CsvIO.Write} can write one or many CSV files
+ * automatically creating the header based on its inferred {@link Schema}.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions = ...
+ * transactions.apply(CsvIO.<Transaction>write("path/to/folder/prefix", CSVFormat.DEFAULT));
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header is repeated for every
+ * shard file, whereas by default, {@link CsvIO.Write} will write all fields in <b>sorted order</b>
+ * of the field names.
+ *
+ * <pre>{@code
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ *
+ * <p>To control the order and subset of fields that {@link CsvIO.Write} writes, use {@link
+ * CSVFormat#withHeader(String...)}. Note, however, the following constraints:
+ *
+ * <ol>
+ * <li>Each header column must match a field name in the {@link Schema}; matching is case
+ * sensitive.
+ * <li>Matching header columns must match {@link Schema} fields that are valid {@link FieldType}s;
+ * see {@link #VALID_FIELD_TYPE_SET}.
+ * <li>{@link CSVFormat} only allows repeated header columns when {@link
+ * CSVFormat#withAllowDuplicateHeaderNames()}
+ * </ol>
+ *
+ * <p>The following example shows the use of {@link CSVFormat#withHeader(String...)} to control the
+ * order and subset of <code>Transaction</code> fields.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions ...
+ * transactions.apply(
+ * CsvIO
+ * .<Transaction>write("path/to/folder/prefix", CSVFormat.DEFAULT.withHeader("transactionId", "purchaseAmount"))
+ * );
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header is repeated for every
+ * shard file, but only including the subset of fields in their listed order.
+ *
+ * <pre>{@code
+ * transactionId,purchaseAmount
+ * 12345,10.23
+ * 54321,54.65
+ * 98765,11.76
+ * }</pre>
+ *
+ * <p>In addition to header customization, {@link CsvIO.Write} supports {@link
+ * CSVFormat#withHeaderComments(Object...)} as shown below. Note that {@link
+ * CSVFormat#withCommentMarker(char)} is required when specifying header comments.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions = ...
+ * transactions
+ * .apply(
+ * CsvIO.<Transaction>write("path/to/folder/prefix",
+ * CSVFormat.DEFAULT.withCommentMarker('#').withHeaderComments("Bank Report", "1970-01-01", "Operator: John Doe")
+ * );
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header and header comments are
+ * repeated for every shard file.
+ *
+ * <pre>{@code
+ * # Bank Report
+ * # 1970-01-01
+ * # Operator: John Doe
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ *
+ * <p>A {@link PCollection} of {@link Row}s works just like custom Java types illustrated above,
+ * except we use {@link CsvIO#writeRows(String, CSVFormat)} as shown below for the same {@code
+ * Transaction} class. We derive {@code Transaction}'s {@link Schema} using a {@link
+ * org.apache.beam.sdk.schemas.annotations.DefaultSchema.DefaultSchemaProvider}. Note that
+ * hard-coding the {@link Row}s below is for illustration purposes. Developers are instead
+ * encouraged to take advantage of {@link
+ * org.apache.beam.sdk.schemas.annotations.DefaultSchema.DefaultSchemaProvider#toRowFunction(TypeDescriptor)}.
+ *
+ * <pre>{@code
+ * DefaultSchemaProvider defaultSchemaProvider = new DefaultSchemaProvider();
+ * Schema schema = defaultSchemaProvider.schemaFor(TypeDescriptor.of(Transaction.class));
+ * PCollection<Row> transactions = pipeline.apply(Create.of(
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "A")
+ * .withFieldValue("purchaseAmount", 10.23)
+ * .withFieldValue("transactionId", "12345")
+ * .build(),
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "B")
+ * .withFieldValue("purchaseAmount", 54.65)
+ * .withFieldValue("transactionId", "54321")
+ * .build(),
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "C")
+ * .withFieldValue("purchaseAmount", 11.76)
+ * .withFieldValue("transactionId", "98765")
+ * .build()
+ * );
+ *
+ * transactions.apply(
+ * CsvIO
+ * .writeRowsTo("gs://bucket/path/to/folder/prefix", CSVFormat.DEFAULT)
+ * );
+ * }</pre>
+ *
+ * <p>Writing the transactions {@link PCollection} of {@link Row}s would yield the following CSV
+ * file content.
+ *
+ * <pre>{@code
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ *
+ * {@link CsvIO.Write} does not support the following {@link CSVFormat} properties and will throw an
+ * {@link IllegalArgumentException}.
+ *
+ * <ul>
+ * <li>{@link CSVFormat#withSkipHeaderRecord()}
+ * <li>{@link CSVFormat#withAllowMissingColumnNames()} ()}
+ * <li>{@link CSVFormat#withAutoFlush(boolean)}
+ * <li>{@link CSVFormat#withIgnoreHeaderCase()}
+ * <li>{@link CSVFormat#withIgnoreSurroundingSpaces()}
+ * </ul>
+ */
+public class CsvIO {
+ /**
+ * The valid {@link Schema.FieldType} from which {@link CsvIO} converts CSV records to the fields.
+ *
+ * <ul>
+ * <li>{@link FieldType#BYTE}
+ * <li>{@link FieldType#BOOLEAN}
+ * <li>{@link FieldType#DATETIME}
+ * <li>{@link FieldType#DECIMAL}
+ * <li>{@link FieldType#DOUBLE}
+ * <li>{@link FieldType#INT16}
+ * <li>{@link FieldType#INT32}
+ * <li>{@link FieldType#INT64}
+ * <li>{@link FieldType#FLOAT}
+ * <li>{@link FieldType#STRING}
+ * </ul>
+ */
+ public static final Set<Schema.FieldType> VALID_FIELD_TYPE_SET =
+ ImmutableSet.of(
+ FieldType.BYTE,
+ FieldType.BOOLEAN,
+ FieldType.DATETIME,
+ FieldType.DECIMAL,
+ FieldType.DOUBLE,
+ FieldType.INT16,
+ FieldType.INT32,
+ FieldType.INT64,
+ FieldType.FLOAT,
+ FieldType.STRING);
+
+ static final String DEFAULT_FILENAME_SUFFIX = ".csv";
+
+ /** Instantiates a {@link Write} for writing user types in {@link CSVFormat} format. */
+ public static <T> Write<T> write(String to, CSVFormat csvFormat) {
+ return new AutoValue_CsvIO_Write.Builder<T>()
+ .setTextIOWrite(createDefaultTextIOWrite(to))
+ .setCSVFormat(csvFormat)
+ .build();
+ }
+
+ /** Instantiates a {@link Write} for writing {@link Row}s in {@link CSVFormat} format. */
+ public static Write<Row> writeRows(String to, CSVFormat csvFormat) {
+ return new AutoValue_CsvIO_Write.Builder<Row>()
+ .setTextIOWrite(createDefaultTextIOWrite(to))
+ .setCSVFormat(csvFormat)
+ .build();
+ }
+
+ /** {@link PTransform} for writing CSV files. */
+ @AutoValue
+ public abstract static class Write<T>
+ extends PTransform<PCollection<T>, WriteFilesResult<String>> {
+
+ /** Specifies the {@link Compression} of all generated shard files. */
+ public Write<T> withCompression(Compression compression) {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withCompression(compression)).build();
+ }
+
+ public Write<T> withNoSpilling() {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withNoSpilling()).build();
+ }
+
+ /**
+ * Specifies to use a given fixed number of shards per window. See {@link
+ * TextIO.Write#withNumShards(int)}.
+ */
+ public Write<T> withNumShards(Integer numShards) {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withNumShards(numShards)).build();
+ }
+
+ /**
+ * Forces a single file as output and empty shard name template. See {@link
+ * TextIO.Write#withoutSharding()}.
+ */
+ public Write<T> withoutSharding() {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withoutSharding()).build();
+ }
+
+ /**
+ * Uses the given {@link ShardNameTemplate} for naming output files. See {@link
+ * TextIO.Write#withShardNameTemplate(String)}.
+ */
+ public Write<T> withShardTemplate(String shardTemplate) {
+ return toBuilder()
+ .setTextIOWrite(getTextIOWrite().withShardNameTemplate(shardTemplate))
+ .build();
+ }
+
+ /**
+ * Configures the filename suffix for written files. See {@link
+ * TextIO.Write#withSuffix(String)}.
+ */
+ public Write<T> withSuffix(String suffix) {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withSuffix(suffix)).build();
+ }
+
+ /**
+ * Set the base directory used to generate temporary files. See {@link
+ * TextIO.Write#withTempDirectory(ResourceId)}.
+ */
+ public Write<T> withTempDirectory(ResourceId tempDirectory) {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withTempDirectory(tempDirectory)).build();
+ }
+
+ /**
+ * Preserves windowing of input elements and writes them to files based on the element's window.
+ * See {@link TextIO.Write#withWindowedWrites()}.
+ */
+ public Write<T> withWindowedWrites() {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withWindowedWrites()).build();
+ }
+
+ /**
+ * Returns a transform for writing to text files like this one but that has the given {@link
+ * FileBasedSink.WritableByteChannelFactory} to be used by the {@link FileBasedSink} during
+ * output. See {@link
+ * TextIO.Write#withWritableByteChannelFactory(FileBasedSink.WritableByteChannelFactory)}.
+ */
+ public Write<T> withWritableByteChannelFactory(
+ FileBasedSink.WritableByteChannelFactory writableByteChannelFactory) {
+ return toBuilder()
+ .setTextIOWrite(
+ getTextIOWrite().withWritableByteChannelFactory(writableByteChannelFactory))
+ .build();
+ }
+
+ /** The underlying {@link FileIO.Write} that writes converted input to CSV formatted output. */
+ abstract TextIO.Write getTextIOWrite();
+
+ /** The {@link CSVFormat} to inform headers, header comments, and format CSV row content. */
+ abstract CSVFormat getCSVFormat();
+
+ abstract Builder<T> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<T> {
+
+ /**
+ * The underlying {@link FileIO.Write} that writes converted input to CSV formatted output.
+ */
+ abstract Builder<T> setTextIOWrite(TextIO.Write value);
+
+ /** The {@link CSVFormat} to convert input. Defaults to {@link CSVFormat#DEFAULT}. */
+ abstract Builder<T> setCSVFormat(CSVFormat value);
+
+ abstract CSVFormat getCSVFormat();
+
+ abstract Write<T> autoBuild();
+
+ final Write<T> build() {
+ checkArgument(
+ !getCSVFormat().getSkipHeaderRecord(),
+ "withSkipHeaderRecord is an illegal CSVFormat setting");
+
+ if (getCSVFormat().getHeaderComments() != null) {
+ checkArgument(
+ getCSVFormat().isCommentMarkerSet(),
+ "CSVFormat withCommentMarker required when withHeaderComments");
+ }
+ return autoBuild();
+ }
+ }
+
+ @Override
+ public WriteFilesResult<String> expand(PCollection<T> input) {
+ if (!input.hasSchema()) {
+ throw new IllegalArgumentException(
+ String.format(
+ "%s requires an input Schema. Note that only Row or user classes are supported. Consider using TextIO or FileIO directly when writing primitive types",
+ Write.class.getName()));
+ }
+
+ Schema schema = input.getSchema();
+
+ RowCoder rowCoder = RowCoder.of(schema);
+
+ PCollection<Row> rows =
+ input
+ .apply("To Rows", MapElements.into(rows()).via(input.getToRowFunction()))
+ .setCoder(rowCoder);
+
+ CSVFormat csvFormat = applyRequiredCSVFormatSettings(schema);
+
+ String header = formatHeader(csvFormat);
+
+ SerializableFunction<Row, String> toCsvFn =
+ CsvRowConversions.RowToCsv.builder()
+ .setCSVFormat(csvFormat)
+ .setSchema(input.getSchema())
+ .build();
+
+ PCollection<String> csv = rows.apply("To CSV", MapElements.into(strings()).via(toCsvFn));
+
+ return csv.apply("Write CSV", getTextIOWrite().withHeader(header).withOutputFilenames());
+ }
+
+ CSVFormat applyRequiredCSVFormatSettings(Schema schema) {
+ CSVFormat csvFormat = getCSVFormat().withSkipHeaderRecord();
+ if (csvFormat.getHeader() == null) {
+ csvFormat = csvFormat.withHeader(schema.sorted().getFieldNames().toArray(new String[0]));
+ }
+ return csvFormat;
+ }
+
+ private static String formatHeader(CSVFormat csvFormat) {
+ String[] header = requireNonNull(csvFormat.getHeader());
+ List<String> result = new ArrayList<>();
+ if (csvFormat.getHeaderComments() != null) {
+ for (String comment : csvFormat.getHeaderComments()) {
+ result.add(csvFormat.getCommentMarker() + " " + comment);
+ }
+ }
+ CSVFormat withoutHeaderComments = csvFormat.withHeaderComments();
+ result.add(withoutHeaderComments.format((Object[]) header));
+ return String.join("\n", result);
+ }
Review Comment:
Overall, I changed this method to incorporate applying to TextIO.Write.withHeader to provide more clarity.
1. > We are ignoring the case where the user wants to skip the header record and doesn't specify the header. We should only specify the csvFormat header if withSkipHeaderRecord is false.
CSVFormat communicates that withSkipHeaderRecord functions to not output the header. Do you agree then we should keep withSkipHeaderRecord separate from where we are deriving the header?
2. > Note that we must also return null from formatHeader in this case since TextIO#withHeader specifically says that a newline is automatically added after.
Are you saying that the consequence of TextIO#withHeader is:
```
col1, col2, col3
1,2,3
2,3,4
```
vs
```
col1, col2, col3
1,2,3
2,3,4
```
--
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: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [beam] damondouglas commented on a diff in pull request #24630: [24469] Implement CsvIO.Write and supporting classes
Posted by "damondouglas (via GitHub)" <gi...@apache.org>.
damondouglas commented on code in PR #24630:
URL: https://github.com/apache/beam/pull/24630#discussion_r1085716045
##########
sdks/java/io/csv/src/main/java/org/apache/beam/sdk/io/csv/CsvIO.java:
##########
@@ -0,0 +1,446 @@
+/*
+ * 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.sdk.io.csv;
+
+import static java.util.Objects.requireNonNull;
+import static org.apache.beam.sdk.values.TypeDescriptors.rows;
+import static org.apache.beam.sdk.values.TypeDescriptors.strings;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+
+import com.google.auto.value.AutoValue;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Set;
+import org.apache.beam.sdk.coders.RowCoder;
+import org.apache.beam.sdk.io.Compression;
+import org.apache.beam.sdk.io.FileBasedSink;
+import org.apache.beam.sdk.io.FileIO;
+import org.apache.beam.sdk.io.ShardNameTemplate;
+import org.apache.beam.sdk.io.TextIO;
+import org.apache.beam.sdk.io.WriteFilesResult;
+import org.apache.beam.sdk.io.fs.ResourceId;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.transforms.MapElements;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.sdk.values.TypeDescriptor;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableSet;
+import org.apache.commons.csv.CSVFormat;
+
+/**
+ * {@link PTransform}s for reading and writing CSV files.
+ *
+ * <h2>Reading CSV files</h2>
+ *
+ * <p>Reading from CSV files is not yet implemented. Please see <a
+ * href="https://github.com/apache/beam/issues/24552">https://github.com/apache/beam/issues/24552</a>.
+ *
+ * <h2>Writing CSV files</h2>
+ *
+ * <p>To write a {@link PCollection} to one or more CSV files, use {@link CsvIO.Write}, using {@link
+ * CsvIO#writeRows(String, CSVFormat)} or {@link CsvIO#write(String, CSVFormat)}. {@link
+ * CsvIO.Write} supports writing {@link Row} or custom Java types using an inferred {@link Schema}.
+ * Examples below show both scenarios. See the Beam Programming Guide on <a
+ * href="https://beam.apache.org/documentation/programming-guide/#inferring-schemas">inferring
+ * schemas</a> for more information on how to enable Beam to infer a {@link Schema} from a custom
+ * Java type.
+ *
+ * <p>{@link CsvIO.Write} only supports writing {@link Schema} aware types that do not contain any
+ * nested {@link FieldType}s such a {@link org.apache.beam.sdk.schemas.Schema.TypeName#ROW} or
+ * repeated {@link org.apache.beam.sdk.schemas.Schema.TypeName#ARRAY} types. See {@link
+ * CsvIO.Write#VALID_FIELD_TYPE_SET} for valid {@link FieldType}s.
+ *
+ * <h3>Example usage:</h3>
+ *
+ * <p>Suppose we have the following <code>Transaction</code> class annotated with
+ * {@code @DefaultSchema(JavaBeanSchema.class)} so that Beam can infer its {@link Schema}:
+ *
+ * <pre>{@code @DefaultSchema(JavaBeanSchema.class)
+ * public class Transaction {
+ * public Transaction() { … }
+ * public Long getTransactionId();
+ * public void setTransactionId(Long transactionId) { … }
+ * public String getBank() { … }
+ * public void setBank(String bank) { … }
+ * public double getPurchaseAmount() { … }
+ * public void setPurchaseAmount(double purchaseAmount) { … }
+ * }
+ * }</pre>
+ *
+ * <p>From a {@code PCollection<Transaction>}, {@link CsvIO.Write} can write one or many CSV files
+ * automatically creating the header based on its inferred {@link Schema}.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions = ...
+ * transactions.apply(CsvIO.<Transaction>write("path/to/folder/prefix", CSVFormat.DEFAULT));
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header is repeated for every
+ * shard file, whereas by default, {@link CsvIO.Write} will write all fields in <b>sorted order</b>
+ * of the field names.
+ *
+ * <pre>{@code
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ *
+ * <p>To control the order and subset of fields that {@link CsvIO.Write} writes, use {@link
+ * CSVFormat#withHeader(String...)}. Note, however, the following constraints:
+ *
+ * <ol>
+ * <li>Each header column must match a field name in the {@link Schema}; matching is case
+ * sensitive.
+ * <li>Matching header columns must match {@link Schema} fields that are valid {@link FieldType}s;
+ * see {@link #VALID_FIELD_TYPE_SET}.
+ * <li>{@link CSVFormat} only allows repeated header columns when {@link
+ * CSVFormat#withAllowDuplicateHeaderNames()}
+ * </ol>
+ *
+ * <p>The following example shows the use of {@link CSVFormat#withHeader(String...)} to control the
+ * order and subset of <code>Transaction</code> fields.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions ...
+ * transactions.apply(
+ * CsvIO
+ * .<Transaction>write("path/to/folder/prefix", CSVFormat.DEFAULT.withHeader("transactionId", "purchaseAmount"))
+ * );
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header is repeated for every
+ * shard file, but only including the subset of fields in their listed order.
+ *
+ * <pre>{@code
+ * transactionId,purchaseAmount
+ * 12345,10.23
+ * 54321,54.65
+ * 98765,11.76
+ * }</pre>
+ *
+ * <p>In addition to header customization, {@link CsvIO.Write} supports {@link
+ * CSVFormat#withHeaderComments(Object...)} as shown below. Note that {@link
+ * CSVFormat#withCommentMarker(char)} is required when specifying header comments.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions = ...
+ * transactions
+ * .apply(
+ * CsvIO.<Transaction>write("path/to/folder/prefix",
+ * CSVFormat.DEFAULT.withCommentMarker('#').withHeaderComments("Bank Report", "1970-01-01", "Operator: John Doe")
+ * );
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header and header comments are
+ * repeated for every shard file.
+ *
+ * <pre>{@code
+ * # Bank Report
+ * # 1970-01-01
+ * # Operator: John Doe
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ *
+ * <p>A {@link PCollection} of {@link Row}s works just like custom Java types illustrated above,
+ * except we use {@link CsvIO#writeRows(String, CSVFormat)} as shown below for the same {@code
+ * Transaction} class. We derive {@code Transaction}'s {@link Schema} using a {@link
+ * org.apache.beam.sdk.schemas.annotations.DefaultSchema.DefaultSchemaProvider}. Note that
+ * hard-coding the {@link Row}s below is for illustration purposes. Developers are instead
+ * encouraged to take advantage of {@link
+ * org.apache.beam.sdk.schemas.annotations.DefaultSchema.DefaultSchemaProvider#toRowFunction(TypeDescriptor)}.
+ *
+ * <pre>{@code
+ * DefaultSchemaProvider defaultSchemaProvider = new DefaultSchemaProvider();
+ * Schema schema = defaultSchemaProvider.schemaFor(TypeDescriptor.of(Transaction.class));
+ * PCollection<Row> transactions = pipeline.apply(Create.of(
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "A")
+ * .withFieldValue("purchaseAmount", 10.23)
+ * .withFieldValue("transactionId", "12345")
+ * .build(),
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "B")
+ * .withFieldValue("purchaseAmount", 54.65)
+ * .withFieldValue("transactionId", "54321")
+ * .build(),
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "C")
+ * .withFieldValue("purchaseAmount", 11.76)
+ * .withFieldValue("transactionId", "98765")
+ * .build()
+ * );
+ *
+ * transactions.apply(
+ * CsvIO
+ * .writeRowsTo("gs://bucket/path/to/folder/prefix", CSVFormat.DEFAULT)
+ * );
+ * }</pre>
+ *
+ * <p>Writing the transactions {@link PCollection} of {@link Row}s would yield the following CSV
+ * file content.
+ *
+ * <pre>{@code
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ *
+ * {@link CsvIO.Write} does not support the following {@link CSVFormat} properties and will throw an
+ * {@link IllegalArgumentException}.
+ *
+ * <ul>
+ * <li>{@link CSVFormat#withSkipHeaderRecord()}
+ * <li>{@link CSVFormat#withAllowMissingColumnNames()} ()}
+ * <li>{@link CSVFormat#withAutoFlush(boolean)}
+ * <li>{@link CSVFormat#withIgnoreHeaderCase()}
+ * <li>{@link CSVFormat#withIgnoreSurroundingSpaces()}
+ * </ul>
+ */
+public class CsvIO {
+ /**
+ * The valid {@link Schema.FieldType} from which {@link CsvIO} converts CSV records to the fields.
+ *
+ * <ul>
+ * <li>{@link FieldType#BYTE}
+ * <li>{@link FieldType#BOOLEAN}
+ * <li>{@link FieldType#DATETIME}
+ * <li>{@link FieldType#DECIMAL}
+ * <li>{@link FieldType#DOUBLE}
+ * <li>{@link FieldType#INT16}
+ * <li>{@link FieldType#INT32}
+ * <li>{@link FieldType#INT64}
+ * <li>{@link FieldType#FLOAT}
+ * <li>{@link FieldType#STRING}
+ * </ul>
+ */
+ public static final Set<Schema.FieldType> VALID_FIELD_TYPE_SET =
+ ImmutableSet.of(
+ FieldType.BYTE,
+ FieldType.BOOLEAN,
+ FieldType.DATETIME,
+ FieldType.DECIMAL,
+ FieldType.DOUBLE,
+ FieldType.INT16,
+ FieldType.INT32,
+ FieldType.INT64,
+ FieldType.FLOAT,
+ FieldType.STRING);
+
+ static final String DEFAULT_FILENAME_SUFFIX = ".csv";
+
+ /** Instantiates a {@link Write} for writing user types in {@link CSVFormat} format. */
+ public static <T> Write<T> write(String to, CSVFormat csvFormat) {
+ return new AutoValue_CsvIO_Write.Builder<T>()
+ .setTextIOWrite(createDefaultTextIOWrite(to))
+ .setCSVFormat(csvFormat)
+ .build();
+ }
+
+ /** Instantiates a {@link Write} for writing {@link Row}s in {@link CSVFormat} format. */
+ public static Write<Row> writeRows(String to, CSVFormat csvFormat) {
+ return new AutoValue_CsvIO_Write.Builder<Row>()
+ .setTextIOWrite(createDefaultTextIOWrite(to))
+ .setCSVFormat(csvFormat)
+ .build();
+ }
+
+ /** {@link PTransform} for writing CSV files. */
+ @AutoValue
+ public abstract static class Write<T>
+ extends PTransform<PCollection<T>, WriteFilesResult<String>> {
+
+ /** Specifies the {@link Compression} of all generated shard files. */
+ public Write<T> withCompression(Compression compression) {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withCompression(compression)).build();
+ }
+
+ public Write<T> withNoSpilling() {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withNoSpilling()).build();
+ }
+
+ /**
+ * Specifies to use a given fixed number of shards per window. See {@link
+ * TextIO.Write#withNumShards(int)}.
+ */
+ public Write<T> withNumShards(Integer numShards) {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withNumShards(numShards)).build();
+ }
+
+ /**
+ * Forces a single file as output and empty shard name template. See {@link
+ * TextIO.Write#withoutSharding()}.
+ */
+ public Write<T> withoutSharding() {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withoutSharding()).build();
+ }
+
+ /**
+ * Uses the given {@link ShardNameTemplate} for naming output files. See {@link
+ * TextIO.Write#withShardNameTemplate(String)}.
+ */
+ public Write<T> withShardTemplate(String shardTemplate) {
+ return toBuilder()
+ .setTextIOWrite(getTextIOWrite().withShardNameTemplate(shardTemplate))
+ .build();
+ }
+
+ /**
+ * Configures the filename suffix for written files. See {@link
+ * TextIO.Write#withSuffix(String)}.
+ */
+ public Write<T> withSuffix(String suffix) {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withSuffix(suffix)).build();
+ }
+
+ /**
+ * Set the base directory used to generate temporary files. See {@link
+ * TextIO.Write#withTempDirectory(ResourceId)}.
+ */
+ public Write<T> withTempDirectory(ResourceId tempDirectory) {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withTempDirectory(tempDirectory)).build();
+ }
+
+ /**
+ * Preserves windowing of input elements and writes them to files based on the element's window.
+ * See {@link TextIO.Write#withWindowedWrites()}.
+ */
+ public Write<T> withWindowedWrites() {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withWindowedWrites()).build();
+ }
+
+ /**
+ * Returns a transform for writing to text files like this one but that has the given {@link
+ * FileBasedSink.WritableByteChannelFactory} to be used by the {@link FileBasedSink} during
+ * output. See {@link
+ * TextIO.Write#withWritableByteChannelFactory(FileBasedSink.WritableByteChannelFactory)}.
+ */
+ public Write<T> withWritableByteChannelFactory(
+ FileBasedSink.WritableByteChannelFactory writableByteChannelFactory) {
+ return toBuilder()
+ .setTextIOWrite(
+ getTextIOWrite().withWritableByteChannelFactory(writableByteChannelFactory))
+ .build();
+ }
+
+ /** The underlying {@link FileIO.Write} that writes converted input to CSV formatted output. */
+ abstract TextIO.Write getTextIOWrite();
+
+ /** The {@link CSVFormat} to inform headers, header comments, and format CSV row content. */
+ abstract CSVFormat getCSVFormat();
+
+ abstract Builder<T> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<T> {
+
+ /**
+ * The underlying {@link FileIO.Write} that writes converted input to CSV formatted output.
+ */
+ abstract Builder<T> setTextIOWrite(TextIO.Write value);
+
+ /** The {@link CSVFormat} to convert input. Defaults to {@link CSVFormat#DEFAULT}. */
+ abstract Builder<T> setCSVFormat(CSVFormat value);
+
+ abstract CSVFormat getCSVFormat();
+
+ abstract Write<T> autoBuild();
+
+ final Write<T> build() {
+ checkArgument(
+ !getCSVFormat().getSkipHeaderRecord(),
+ "withSkipHeaderRecord is an illegal CSVFormat setting");
+
+ if (getCSVFormat().getHeaderComments() != null) {
+ checkArgument(
+ getCSVFormat().isCommentMarkerSet(),
+ "CSVFormat withCommentMarker required when withHeaderComments");
+ }
+ return autoBuild();
+ }
+ }
+
+ @Override
+ public WriteFilesResult<String> expand(PCollection<T> input) {
+ if (!input.hasSchema()) {
+ throw new IllegalArgumentException(
+ String.format(
+ "%s requires an input Schema. Note that only Row or user classes are supported. Consider using TextIO or FileIO directly when writing primitive types",
+ Write.class.getName()));
+ }
+
+ Schema schema = input.getSchema();
+
+ RowCoder rowCoder = RowCoder.of(schema);
+
+ PCollection<Row> rows =
+ input
+ .apply("To Rows", MapElements.into(rows()).via(input.getToRowFunction()))
+ .setCoder(rowCoder);
+
+ CSVFormat csvFormat = applyRequiredCSVFormatSettings(schema);
+
+ String header = formatHeader(csvFormat);
+
+ SerializableFunction<Row, String> toCsvFn =
+ CsvRowConversions.RowToCsv.builder()
+ .setCSVFormat(csvFormat)
+ .setSchema(input.getSchema())
+ .build();
+
+ PCollection<String> csv = rows.apply("To CSV", MapElements.into(strings()).via(toCsvFn));
+
+ return csv.apply("Write CSV", getTextIOWrite().withHeader(header).withOutputFilenames());
+ }
+
+ CSVFormat applyRequiredCSVFormatSettings(Schema schema) {
+ CSVFormat csvFormat = getCSVFormat().withSkipHeaderRecord();
+ if (csvFormat.getHeader() == null) {
+ csvFormat = csvFormat.withHeader(schema.sorted().getFieldNames().toArray(new String[0]));
+ }
+ return csvFormat;
+ }
+
+ private static String formatHeader(CSVFormat csvFormat) {
+ String[] header = requireNonNull(csvFormat.getHeader());
+ List<String> result = new ArrayList<>();
+ if (csvFormat.getHeaderComments() != null) {
+ for (String comment : csvFormat.getHeaderComments()) {
+ result.add(csvFormat.getCommentMarker() + " " + comment);
+ }
+ }
+ CSVFormat withoutHeaderComments = csvFormat.withHeaderComments();
+ result.add(withoutHeaderComments.format((Object[]) header));
+ return String.join("\n", result);
+ }
Review Comment:
Overall, I changed this method to incorporate applying to TextIO.Write.withHeader to provide more clarity.
> We are ignoring the case where the user wants to skip the header record and doesn't specify the header. We should only specify the csvFormat header if withSkipHeaderRecord is false.
CSVFormat communicates that withSkipHeaderRecord functions to not output the header. Do you agree then we should keep withSkipHeaderRecord separate from where we are deriving the header?
> Note that we must also return null from formatHeader in this case since TextIO#withHeader specifically says that a newline is automatically added after.
Are you saying that the consequence of TextIO#withHeader is:
```
col1, col2, col3
1,2,3
2,3,4
```
vs
```
col1, col2, col3
1,2,3
2,3,4
```
--
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: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [beam] codecov[bot] commented on pull request #24630: [24469] Implement CsvIO.Write and supporting classes
Posted by GitBox <gi...@apache.org>.
codecov[bot] commented on PR #24630:
URL: https://github.com/apache/beam/pull/24630#issuecomment-1344891653
# [Codecov](https://codecov.io/gh/apache/beam/pull/24630?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=The+Apache+Software+Foundation) Report
> Merging [#24630](https://codecov.io/gh/apache/beam/pull/24630?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=The+Apache+Software+Foundation) (e374c0f) into [master](https://codecov.io/gh/apache/beam/commit/cb329e410e9975ceeac751cb7ac28cd27cf2903d?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=The+Apache+Software+Foundation) (cb329e4) will **decrease** coverage by `0.00%`.
> The diff coverage is `n/a`.
```diff
@@ Coverage Diff @@
## master #24630 +/- ##
==========================================
- Coverage 73.37% 73.36% -0.01%
==========================================
Files 719 719
Lines 97209 97126 -83
==========================================
- Hits 71323 71256 -67
+ Misses 24538 24522 -16
Partials 1348 1348
```
| Flag | Coverage Δ | |
|---|---|---|
| python | `82.98% <ø> (+<0.01%)` | :arrow_up: |
Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=The+Apache+Software+Foundation#carryforward-flags-in-the-pull-request-comment) to find out more.
| [Impacted Files](https://codecov.io/gh/apache/beam/pull/24630?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=The+Apache+Software+Foundation) | Coverage Δ | |
|---|---|---|
| [...on/apache\_beam/runners/dataflow/dataflow\_runner.py](https://codecov.io/gh/apache/beam/pull/24630/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=The+Apache+Software+Foundation#diff-c2Rrcy9weXRob24vYXBhY2hlX2JlYW0vcnVubmVycy9kYXRhZmxvdy9kYXRhZmxvd19ydW5uZXIucHk=) | `81.88% <0.00%> (-0.81%)` | :arrow_down: |
| [...apache\_beam/runners/dataflow/internal/apiclient.py](https://codecov.io/gh/apache/beam/pull/24630/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=The+Apache+Software+Foundation#diff-c2Rrcy9weXRob24vYXBhY2hlX2JlYW0vcnVubmVycy9kYXRhZmxvdy9pbnRlcm5hbC9hcGljbGllbnQucHk=) | `76.72% <0.00%> (-0.63%)` | :arrow_down: |
| [...s/python/apache\_beam/runners/portability/stager.py](https://codecov.io/gh/apache/beam/pull/24630/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=The+Apache+Software+Foundation#diff-c2Rrcy9weXRob24vYXBhY2hlX2JlYW0vcnVubmVycy9wb3J0YWJpbGl0eS9zdGFnZXIucHk=) | `87.04% <0.00%> (-0.18%)` | :arrow_down: |
| [...dks/python/apache\_beam/options/pipeline\_options.py](https://codecov.io/gh/apache/beam/pull/24630/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=The+Apache+Software+Foundation#diff-c2Rrcy9weXRob24vYXBhY2hlX2JlYW0vb3B0aW9ucy9waXBlbGluZV9vcHRpb25zLnB5) | `93.95% <0.00%> (-0.02%)` | :arrow_down: |
| [...che\_beam/runners/interactive/interactive\_runner.py](https://codecov.io/gh/apache/beam/pull/24630/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=The+Apache+Software+Foundation#diff-c2Rrcy9weXRob24vYXBhY2hlX2JlYW0vcnVubmVycy9pbnRlcmFjdGl2ZS9pbnRlcmFjdGl2ZV9ydW5uZXIucHk=) | `91.77% <0.00%> (+1.26%)` | :arrow_up: |
| [...python/apache\_beam/runners/worker/worker\_status.py](https://codecov.io/gh/apache/beam/pull/24630/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=The+Apache+Software+Foundation#diff-c2Rrcy9weXRob24vYXBhY2hlX2JlYW0vcnVubmVycy93b3JrZXIvd29ya2VyX3N0YXR1cy5weQ==) | `76.66% <0.00%> (+1.33%)` | :arrow_up: |
| [sdks/python/apache\_beam/utils/interactive\_utils.py](https://codecov.io/gh/apache/beam/pull/24630/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=The+Apache+Software+Foundation#diff-c2Rrcy9weXRob24vYXBhY2hlX2JlYW0vdXRpbHMvaW50ZXJhY3RpdmVfdXRpbHMucHk=) | `97.56% <0.00%> (+2.43%)` | :arrow_up: |
| [.../python/apache\_beam/testing/test\_stream\_service.py](https://codecov.io/gh/apache/beam/pull/24630/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=The+Apache+Software+Foundation#diff-c2Rrcy9weXRob24vYXBhY2hlX2JlYW0vdGVzdGluZy90ZXN0X3N0cmVhbV9zZXJ2aWNlLnB5) | `92.85% <0.00%> (+4.76%)` | :arrow_up: |
| [...ache\_beam/runners/dataflow/ptransform\_overrides.py](https://codecov.io/gh/apache/beam/pull/24630/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=The+Apache+Software+Foundation#diff-c2Rrcy9weXRob24vYXBhY2hlX2JlYW0vcnVubmVycy9kYXRhZmxvdy9wdHJhbnNmb3JtX292ZXJyaWRlcy5weQ==) | `98.66% <0.00%> (+7.75%)` | :arrow_up: |
:mega: We’re building smart automated test selection to slash your CI/CD build times. [Learn more](https://about.codecov.io/iterative-testing/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=The+Apache+Software+Foundation)
--
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: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [beam] damondouglas commented on a diff in pull request #24630: [24469] Implement CsvIO.Write and supporting classes
Posted by "damondouglas (via GitHub)" <gi...@apache.org>.
damondouglas commented on code in PR #24630:
URL: https://github.com/apache/beam/pull/24630#discussion_r1085716045
##########
sdks/java/io/csv/src/main/java/org/apache/beam/sdk/io/csv/CsvIO.java:
##########
@@ -0,0 +1,446 @@
+/*
+ * 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.sdk.io.csv;
+
+import static java.util.Objects.requireNonNull;
+import static org.apache.beam.sdk.values.TypeDescriptors.rows;
+import static org.apache.beam.sdk.values.TypeDescriptors.strings;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+
+import com.google.auto.value.AutoValue;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Set;
+import org.apache.beam.sdk.coders.RowCoder;
+import org.apache.beam.sdk.io.Compression;
+import org.apache.beam.sdk.io.FileBasedSink;
+import org.apache.beam.sdk.io.FileIO;
+import org.apache.beam.sdk.io.ShardNameTemplate;
+import org.apache.beam.sdk.io.TextIO;
+import org.apache.beam.sdk.io.WriteFilesResult;
+import org.apache.beam.sdk.io.fs.ResourceId;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.transforms.MapElements;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.sdk.values.TypeDescriptor;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableSet;
+import org.apache.commons.csv.CSVFormat;
+
+/**
+ * {@link PTransform}s for reading and writing CSV files.
+ *
+ * <h2>Reading CSV files</h2>
+ *
+ * <p>Reading from CSV files is not yet implemented. Please see <a
+ * href="https://github.com/apache/beam/issues/24552">https://github.com/apache/beam/issues/24552</a>.
+ *
+ * <h2>Writing CSV files</h2>
+ *
+ * <p>To write a {@link PCollection} to one or more CSV files, use {@link CsvIO.Write}, using {@link
+ * CsvIO#writeRows(String, CSVFormat)} or {@link CsvIO#write(String, CSVFormat)}. {@link
+ * CsvIO.Write} supports writing {@link Row} or custom Java types using an inferred {@link Schema}.
+ * Examples below show both scenarios. See the Beam Programming Guide on <a
+ * href="https://beam.apache.org/documentation/programming-guide/#inferring-schemas">inferring
+ * schemas</a> for more information on how to enable Beam to infer a {@link Schema} from a custom
+ * Java type.
+ *
+ * <p>{@link CsvIO.Write} only supports writing {@link Schema} aware types that do not contain any
+ * nested {@link FieldType}s such a {@link org.apache.beam.sdk.schemas.Schema.TypeName#ROW} or
+ * repeated {@link org.apache.beam.sdk.schemas.Schema.TypeName#ARRAY} types. See {@link
+ * CsvIO.Write#VALID_FIELD_TYPE_SET} for valid {@link FieldType}s.
+ *
+ * <h3>Example usage:</h3>
+ *
+ * <p>Suppose we have the following <code>Transaction</code> class annotated with
+ * {@code @DefaultSchema(JavaBeanSchema.class)} so that Beam can infer its {@link Schema}:
+ *
+ * <pre>{@code @DefaultSchema(JavaBeanSchema.class)
+ * public class Transaction {
+ * public Transaction() { … }
+ * public Long getTransactionId();
+ * public void setTransactionId(Long transactionId) { … }
+ * public String getBank() { … }
+ * public void setBank(String bank) { … }
+ * public double getPurchaseAmount() { … }
+ * public void setPurchaseAmount(double purchaseAmount) { … }
+ * }
+ * }</pre>
+ *
+ * <p>From a {@code PCollection<Transaction>}, {@link CsvIO.Write} can write one or many CSV files
+ * automatically creating the header based on its inferred {@link Schema}.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions = ...
+ * transactions.apply(CsvIO.<Transaction>write("path/to/folder/prefix", CSVFormat.DEFAULT));
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header is repeated for every
+ * shard file, whereas by default, {@link CsvIO.Write} will write all fields in <b>sorted order</b>
+ * of the field names.
+ *
+ * <pre>{@code
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ *
+ * <p>To control the order and subset of fields that {@link CsvIO.Write} writes, use {@link
+ * CSVFormat#withHeader(String...)}. Note, however, the following constraints:
+ *
+ * <ol>
+ * <li>Each header column must match a field name in the {@link Schema}; matching is case
+ * sensitive.
+ * <li>Matching header columns must match {@link Schema} fields that are valid {@link FieldType}s;
+ * see {@link #VALID_FIELD_TYPE_SET}.
+ * <li>{@link CSVFormat} only allows repeated header columns when {@link
+ * CSVFormat#withAllowDuplicateHeaderNames()}
+ * </ol>
+ *
+ * <p>The following example shows the use of {@link CSVFormat#withHeader(String...)} to control the
+ * order and subset of <code>Transaction</code> fields.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions ...
+ * transactions.apply(
+ * CsvIO
+ * .<Transaction>write("path/to/folder/prefix", CSVFormat.DEFAULT.withHeader("transactionId", "purchaseAmount"))
+ * );
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header is repeated for every
+ * shard file, but only including the subset of fields in their listed order.
+ *
+ * <pre>{@code
+ * transactionId,purchaseAmount
+ * 12345,10.23
+ * 54321,54.65
+ * 98765,11.76
+ * }</pre>
+ *
+ * <p>In addition to header customization, {@link CsvIO.Write} supports {@link
+ * CSVFormat#withHeaderComments(Object...)} as shown below. Note that {@link
+ * CSVFormat#withCommentMarker(char)} is required when specifying header comments.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions = ...
+ * transactions
+ * .apply(
+ * CsvIO.<Transaction>write("path/to/folder/prefix",
+ * CSVFormat.DEFAULT.withCommentMarker('#').withHeaderComments("Bank Report", "1970-01-01", "Operator: John Doe")
+ * );
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header and header comments are
+ * repeated for every shard file.
+ *
+ * <pre>{@code
+ * # Bank Report
+ * # 1970-01-01
+ * # Operator: John Doe
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ *
+ * <p>A {@link PCollection} of {@link Row}s works just like custom Java types illustrated above,
+ * except we use {@link CsvIO#writeRows(String, CSVFormat)} as shown below for the same {@code
+ * Transaction} class. We derive {@code Transaction}'s {@link Schema} using a {@link
+ * org.apache.beam.sdk.schemas.annotations.DefaultSchema.DefaultSchemaProvider}. Note that
+ * hard-coding the {@link Row}s below is for illustration purposes. Developers are instead
+ * encouraged to take advantage of {@link
+ * org.apache.beam.sdk.schemas.annotations.DefaultSchema.DefaultSchemaProvider#toRowFunction(TypeDescriptor)}.
+ *
+ * <pre>{@code
+ * DefaultSchemaProvider defaultSchemaProvider = new DefaultSchemaProvider();
+ * Schema schema = defaultSchemaProvider.schemaFor(TypeDescriptor.of(Transaction.class));
+ * PCollection<Row> transactions = pipeline.apply(Create.of(
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "A")
+ * .withFieldValue("purchaseAmount", 10.23)
+ * .withFieldValue("transactionId", "12345")
+ * .build(),
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "B")
+ * .withFieldValue("purchaseAmount", 54.65)
+ * .withFieldValue("transactionId", "54321")
+ * .build(),
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "C")
+ * .withFieldValue("purchaseAmount", 11.76)
+ * .withFieldValue("transactionId", "98765")
+ * .build()
+ * );
+ *
+ * transactions.apply(
+ * CsvIO
+ * .writeRowsTo("gs://bucket/path/to/folder/prefix", CSVFormat.DEFAULT)
+ * );
+ * }</pre>
+ *
+ * <p>Writing the transactions {@link PCollection} of {@link Row}s would yield the following CSV
+ * file content.
+ *
+ * <pre>{@code
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ *
+ * {@link CsvIO.Write} does not support the following {@link CSVFormat} properties and will throw an
+ * {@link IllegalArgumentException}.
+ *
+ * <ul>
+ * <li>{@link CSVFormat#withSkipHeaderRecord()}
+ * <li>{@link CSVFormat#withAllowMissingColumnNames()} ()}
+ * <li>{@link CSVFormat#withAutoFlush(boolean)}
+ * <li>{@link CSVFormat#withIgnoreHeaderCase()}
+ * <li>{@link CSVFormat#withIgnoreSurroundingSpaces()}
+ * </ul>
+ */
+public class CsvIO {
+ /**
+ * The valid {@link Schema.FieldType} from which {@link CsvIO} converts CSV records to the fields.
+ *
+ * <ul>
+ * <li>{@link FieldType#BYTE}
+ * <li>{@link FieldType#BOOLEAN}
+ * <li>{@link FieldType#DATETIME}
+ * <li>{@link FieldType#DECIMAL}
+ * <li>{@link FieldType#DOUBLE}
+ * <li>{@link FieldType#INT16}
+ * <li>{@link FieldType#INT32}
+ * <li>{@link FieldType#INT64}
+ * <li>{@link FieldType#FLOAT}
+ * <li>{@link FieldType#STRING}
+ * </ul>
+ */
+ public static final Set<Schema.FieldType> VALID_FIELD_TYPE_SET =
+ ImmutableSet.of(
+ FieldType.BYTE,
+ FieldType.BOOLEAN,
+ FieldType.DATETIME,
+ FieldType.DECIMAL,
+ FieldType.DOUBLE,
+ FieldType.INT16,
+ FieldType.INT32,
+ FieldType.INT64,
+ FieldType.FLOAT,
+ FieldType.STRING);
+
+ static final String DEFAULT_FILENAME_SUFFIX = ".csv";
+
+ /** Instantiates a {@link Write} for writing user types in {@link CSVFormat} format. */
+ public static <T> Write<T> write(String to, CSVFormat csvFormat) {
+ return new AutoValue_CsvIO_Write.Builder<T>()
+ .setTextIOWrite(createDefaultTextIOWrite(to))
+ .setCSVFormat(csvFormat)
+ .build();
+ }
+
+ /** Instantiates a {@link Write} for writing {@link Row}s in {@link CSVFormat} format. */
+ public static Write<Row> writeRows(String to, CSVFormat csvFormat) {
+ return new AutoValue_CsvIO_Write.Builder<Row>()
+ .setTextIOWrite(createDefaultTextIOWrite(to))
+ .setCSVFormat(csvFormat)
+ .build();
+ }
+
+ /** {@link PTransform} for writing CSV files. */
+ @AutoValue
+ public abstract static class Write<T>
+ extends PTransform<PCollection<T>, WriteFilesResult<String>> {
+
+ /** Specifies the {@link Compression} of all generated shard files. */
+ public Write<T> withCompression(Compression compression) {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withCompression(compression)).build();
+ }
+
+ public Write<T> withNoSpilling() {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withNoSpilling()).build();
+ }
+
+ /**
+ * Specifies to use a given fixed number of shards per window. See {@link
+ * TextIO.Write#withNumShards(int)}.
+ */
+ public Write<T> withNumShards(Integer numShards) {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withNumShards(numShards)).build();
+ }
+
+ /**
+ * Forces a single file as output and empty shard name template. See {@link
+ * TextIO.Write#withoutSharding()}.
+ */
+ public Write<T> withoutSharding() {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withoutSharding()).build();
+ }
+
+ /**
+ * Uses the given {@link ShardNameTemplate} for naming output files. See {@link
+ * TextIO.Write#withShardNameTemplate(String)}.
+ */
+ public Write<T> withShardTemplate(String shardTemplate) {
+ return toBuilder()
+ .setTextIOWrite(getTextIOWrite().withShardNameTemplate(shardTemplate))
+ .build();
+ }
+
+ /**
+ * Configures the filename suffix for written files. See {@link
+ * TextIO.Write#withSuffix(String)}.
+ */
+ public Write<T> withSuffix(String suffix) {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withSuffix(suffix)).build();
+ }
+
+ /**
+ * Set the base directory used to generate temporary files. See {@link
+ * TextIO.Write#withTempDirectory(ResourceId)}.
+ */
+ public Write<T> withTempDirectory(ResourceId tempDirectory) {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withTempDirectory(tempDirectory)).build();
+ }
+
+ /**
+ * Preserves windowing of input elements and writes them to files based on the element's window.
+ * See {@link TextIO.Write#withWindowedWrites()}.
+ */
+ public Write<T> withWindowedWrites() {
+ return toBuilder().setTextIOWrite(getTextIOWrite().withWindowedWrites()).build();
+ }
+
+ /**
+ * Returns a transform for writing to text files like this one but that has the given {@link
+ * FileBasedSink.WritableByteChannelFactory} to be used by the {@link FileBasedSink} during
+ * output. See {@link
+ * TextIO.Write#withWritableByteChannelFactory(FileBasedSink.WritableByteChannelFactory)}.
+ */
+ public Write<T> withWritableByteChannelFactory(
+ FileBasedSink.WritableByteChannelFactory writableByteChannelFactory) {
+ return toBuilder()
+ .setTextIOWrite(
+ getTextIOWrite().withWritableByteChannelFactory(writableByteChannelFactory))
+ .build();
+ }
+
+ /** The underlying {@link FileIO.Write} that writes converted input to CSV formatted output. */
+ abstract TextIO.Write getTextIOWrite();
+
+ /** The {@link CSVFormat} to inform headers, header comments, and format CSV row content. */
+ abstract CSVFormat getCSVFormat();
+
+ abstract Builder<T> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<T> {
+
+ /**
+ * The underlying {@link FileIO.Write} that writes converted input to CSV formatted output.
+ */
+ abstract Builder<T> setTextIOWrite(TextIO.Write value);
+
+ /** The {@link CSVFormat} to convert input. Defaults to {@link CSVFormat#DEFAULT}. */
+ abstract Builder<T> setCSVFormat(CSVFormat value);
+
+ abstract CSVFormat getCSVFormat();
+
+ abstract Write<T> autoBuild();
+
+ final Write<T> build() {
+ checkArgument(
+ !getCSVFormat().getSkipHeaderRecord(),
+ "withSkipHeaderRecord is an illegal CSVFormat setting");
+
+ if (getCSVFormat().getHeaderComments() != null) {
+ checkArgument(
+ getCSVFormat().isCommentMarkerSet(),
+ "CSVFormat withCommentMarker required when withHeaderComments");
+ }
+ return autoBuild();
+ }
+ }
+
+ @Override
+ public WriteFilesResult<String> expand(PCollection<T> input) {
+ if (!input.hasSchema()) {
+ throw new IllegalArgumentException(
+ String.format(
+ "%s requires an input Schema. Note that only Row or user classes are supported. Consider using TextIO or FileIO directly when writing primitive types",
+ Write.class.getName()));
+ }
+
+ Schema schema = input.getSchema();
+
+ RowCoder rowCoder = RowCoder.of(schema);
+
+ PCollection<Row> rows =
+ input
+ .apply("To Rows", MapElements.into(rows()).via(input.getToRowFunction()))
+ .setCoder(rowCoder);
+
+ CSVFormat csvFormat = applyRequiredCSVFormatSettings(schema);
+
+ String header = formatHeader(csvFormat);
+
+ SerializableFunction<Row, String> toCsvFn =
+ CsvRowConversions.RowToCsv.builder()
+ .setCSVFormat(csvFormat)
+ .setSchema(input.getSchema())
+ .build();
+
+ PCollection<String> csv = rows.apply("To CSV", MapElements.into(strings()).via(toCsvFn));
+
+ return csv.apply("Write CSV", getTextIOWrite().withHeader(header).withOutputFilenames());
+ }
+
+ CSVFormat applyRequiredCSVFormatSettings(Schema schema) {
+ CSVFormat csvFormat = getCSVFormat().withSkipHeaderRecord();
+ if (csvFormat.getHeader() == null) {
+ csvFormat = csvFormat.withHeader(schema.sorted().getFieldNames().toArray(new String[0]));
+ }
+ return csvFormat;
+ }
+
+ private static String formatHeader(CSVFormat csvFormat) {
+ String[] header = requireNonNull(csvFormat.getHeader());
+ List<String> result = new ArrayList<>();
+ if (csvFormat.getHeaderComments() != null) {
+ for (String comment : csvFormat.getHeaderComments()) {
+ result.add(csvFormat.getCommentMarker() + " " + comment);
+ }
+ }
+ CSVFormat withoutHeaderComments = csvFormat.withHeaderComments();
+ result.add(withoutHeaderComments.format((Object[]) header));
+ return String.join("\n", result);
+ }
Review Comment:
In response to this comment and hopefully provide more clarity, I've:
1. Renamed `applyRequiredCSVFormatSettings` to [`buildHeaderFromSchemaIfNeeded`](https://github.com/damondouglas/beam/blob/24469-csv-io-write/sdks/java/io/csv/src/main/java/org/apache/beam/sdk/io/csv/CsvIO.java#L426)
2. Refactored and renamed `formatHeader` to [`writeWithCSVFormatHeaderAndComments`](https://github.com/damondouglas/beam/blob/24469-csv-io-write/sdks/java/io/csv/src/main/java/org/apache/beam/sdk/io/csv/CsvIO.java#L436).
Additionally, I wanted to address the questions/comments below:
1. > We are ignoring the case where the user wants to skip the header record and doesn't specify the header. We should only specify the csvFormat header if withSkipHeaderRecord is false.
CSVFormat communicates that withSkipHeaderRecord functions to not output the header. It's also a parameter that gives an output incompatible in a beam context as CSVFormat was designed for writing a single file. Therefore, do you agree then we should not use withSkipHeaderRecord to drive how we are building the header?
3. > Note that we must also return null from formatHeader in this case since TextIO#withHeader specifically says that a newline is automatically added after.
Apologies, your comment confused me.
Are you saying that the consequence of TextIO#withHeader is:
```
col1, col2, col3
1,2,3
2,3,4
```
vs
```
col1, col2, col3
1,2,3
2,3,4
```
--
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: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [beam] lukecwik commented on a diff in pull request #24630: [24469] Implement CsvIO.Write and supporting classes
Posted by GitBox <gi...@apache.org>.
lukecwik commented on code in PR #24630:
URL: https://github.com/apache/beam/pull/24630#discussion_r1052700616
##########
sdks/java/io/csv/src/main/java/org/apache/beam/sdk/io/csv/CsvIO.java:
##########
@@ -0,0 +1,554 @@
+/*
+ * 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.sdk.io.csv;
+
+import static java.nio.charset.StandardCharsets.UTF_8;
+import static org.apache.beam.sdk.values.TypeDescriptors.rows;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
+
+import com.google.auto.value.AutoValue;
+import edu.umd.cs.findbugs.annotations.NonNull;
+import java.io.BufferedWriter;
+import java.io.IOException;
+import java.io.OutputStreamWriter;
+import java.io.PrintWriter;
+import java.nio.channels.Channels;
+import java.nio.channels.WritableByteChannel;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+import java.util.List;
+import java.util.Optional;
+import java.util.Set;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.io.Compression;
+import org.apache.beam.sdk.io.FileIO;
+import org.apache.beam.sdk.io.FileIO.Write.FileNaming;
+import org.apache.beam.sdk.io.csv.CsvRowConversions.RowToCsv;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.Field;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.transforms.Contextful;
+import org.apache.beam.sdk.transforms.MapElements;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PCollectionView;
+import org.apache.beam.sdk.values.PDone;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableSet;
+import org.apache.commons.csv.CSVFormat;
+
+/**
+ * {@link PTransform}s for reading and writing CSV files.
+ *
+ * <h2>Reading CSV files</h2>
+ *
+ * <p>Reading from CSV files is not yet implemented. Please see <a
+ * href="https://github.com/apache/beam/issues/24552">https://github.com/apache/beam/issues/24552</a>.
+ *
+ * <h2>Writing CSV files</h2>
+ *
+ * <p>To write a {@link PCollection} to one or more CSV files, use {@link CsvIO.Write}, using {@code
+ * CsvIO.write().to(String)}. {@link CsvIO.Write} supports writing {@link Row} or custom Java types
+ * using an inferred {@link Schema}. See the Beam Programming Guide on <a
+ * href="https://beam.apache.org/documentation/programming-guide/#inferring-schemas">Inferring
Review Comment:
```suggestion
* href="https://beam.apache.org/documentation/programming-guide/#inferring-schemas">inferring
```
##########
sdks/java/io/csv/src/main/java/org/apache/beam/sdk/io/csv/CsvIO.java:
##########
@@ -0,0 +1,554 @@
+/*
+ * 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.sdk.io.csv;
+
+import static java.nio.charset.StandardCharsets.UTF_8;
+import static org.apache.beam.sdk.values.TypeDescriptors.rows;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
+
+import com.google.auto.value.AutoValue;
+import edu.umd.cs.findbugs.annotations.NonNull;
+import java.io.BufferedWriter;
+import java.io.IOException;
+import java.io.OutputStreamWriter;
+import java.io.PrintWriter;
+import java.nio.channels.Channels;
+import java.nio.channels.WritableByteChannel;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+import java.util.List;
+import java.util.Optional;
+import java.util.Set;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.io.Compression;
+import org.apache.beam.sdk.io.FileIO;
+import org.apache.beam.sdk.io.FileIO.Write.FileNaming;
+import org.apache.beam.sdk.io.csv.CsvRowConversions.RowToCsv;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.Field;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.transforms.Contextful;
+import org.apache.beam.sdk.transforms.MapElements;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PCollectionView;
+import org.apache.beam.sdk.values.PDone;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableSet;
+import org.apache.commons.csv.CSVFormat;
+
+/**
+ * {@link PTransform}s for reading and writing CSV files.
+ *
+ * <h2>Reading CSV files</h2>
+ *
+ * <p>Reading from CSV files is not yet implemented. Please see <a
+ * href="https://github.com/apache/beam/issues/24552">https://github.com/apache/beam/issues/24552</a>.
+ *
+ * <h2>Writing CSV files</h2>
+ *
+ * <p>To write a {@link PCollection} to one or more CSV files, use {@link CsvIO.Write}, using {@code
+ * CsvIO.write().to(String)}. {@link CsvIO.Write} supports writing {@link Row} or custom Java types
+ * using an inferred {@link Schema}. See the Beam Programming Guide on <a
+ * href="https://beam.apache.org/documentation/programming-guide/#inferring-schemas">Inferring
+ * schemas</a> for more information on how to enable Beam to infer a {@link Schema} from a custom
+ * Java type.
+ *
+ * <p>For each {@link Schema.Field} included in the {@link Schema}, according to {@link
+ * CsvIO.Write#withSchemaFields(List)}, {@link CsvIO.Write} restricts the use of flat {@link
+ * Schema}s and checks whether the {@link Schema.FieldType}s are one of {@link
+ * #VALID_FIELD_TYPE_SET}.
Review Comment:
```suggestion
* <p>{@link CsvIO.Write} only supports writing schemas that do not contain any nested types. {@link CsvIO.Write} also only supports a {@link #VALID_FIELD_TYPE_SET limited set} of {@link Schema.FieldType schema field types}.
```
##########
sdks/java/io/csv/src/main/java/org/apache/beam/sdk/io/csv/CsvIO.java:
##########
@@ -0,0 +1,553 @@
+/*
+ * 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.sdk.io.csv;
+
+import static java.nio.charset.StandardCharsets.UTF_8;
+import static org.apache.beam.sdk.values.TypeDescriptors.rows;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
+
+import com.google.auto.value.AutoValue;
+import java.io.BufferedWriter;
+import java.io.IOException;
+import java.io.OutputStreamWriter;
+import java.io.PrintWriter;
+import java.nio.channels.Channels;
+import java.nio.channels.WritableByteChannel;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+import java.util.List;
+import java.util.Optional;
+import java.util.Set;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.io.Compression;
+import org.apache.beam.sdk.io.FileIO;
+import org.apache.beam.sdk.io.FileIO.Write.FileNaming;
+import org.apache.beam.sdk.io.csv.CsvRowConversions.RowToCsv;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.Field;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.transforms.Contextful;
+import org.apache.beam.sdk.transforms.MapElements;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PCollectionView;
+import org.apache.beam.sdk.values.PDone;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableSet;
+import org.apache.commons.csv.CSVFormat;
+
+/**
+ * {@link PTransform}s for reading and writing CSV files.
+ *
+ * <h2>Reading CSV files</h2>
+ *
+ * <p>Reading from CSV files is not yet implemented. Please see <a
+ * href="https://github.com/apache/beam/issues/24552">https://github.com/apache/beam/issues/24552</a>.
+ *
+ * <h2>Writing CSV files</h2>
+ *
+ * <p>To write a {@link PCollection} to one or more CSV files, use {@link CsvIO.Write}, using {@code
+ * CsvIO.write().to(String)}. {@link CsvIO.Write} supports writing {@link Row} or custom Java types
+ * using an inferred {@link Schema}. See the Beam Programming Guide on <a
+ * href="https://beam.apache.org/documentation/programming-guide/#inferring-schemas">Inferring
+ * schemas</a> for more information on how to enable Beam to infer a {@link Schema} from a custom
+ * Java type.
+ *
+ * <p>For each {@link Schema.Field} included in the {@link Schema}, according to {@link
+ * CsvIO.Write#withSchemaFields(List)}, {@link CsvIO.Write} restricts the use of flat {@link
+ * Schema}s and checks whether the {@link Schema.FieldType}s are one of {@link
+ * #VALID_FIELD_TYPE_SET}.
+ *
+ * <h3>Example usage:</h3>
+ *
+ * <p>Suppose we have a <code>Transaction</code> class annotated with
+ * {@code @DefaultSchema(JavaBeanSchema.class)} so that Beam can infer its {@link Schema}:
+ *
+ * <pre>{@code @DefaultSchema(JavaBeanSchema.class)
+ * public class Transaction {
+ * public Transaction() { … }
+ * public Long getTransactionId();
+ * public void setTransactionId(Long transactionId) { … }
+ * public String getBank() { … }
+ * public void setBank(String bank) { … }
+ * public double getPurchaseAmount() { … }
+ * public void setPurchaseAmount(double purchaseAmount) { … }
+ * }
+ * }</pre>
+ *
+ * <p>From a {@code PCollection<Transaction>}, we can write one or many CSV files and {@link
+ * CsvIO.Write} will automatically create the header based on its inferred {@link Schema}
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions ...
+ * transactions.apply(CsvIO.<Transaction>write().to("gs://bucket/path/to/folder/prefix"));
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header is repeated for every
+ * file, whereas by default, {@link CsvIO.Write} will write all fields in <b>sorted order</b> of the
+ * field names.
+ *
+ * <pre>{@code
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ *
+ * <p>{@link CsvIO.Write} allows the use a subset of the available fields or control the order of
+ * their output using {@link CsvIO.Write#withSchemaFields(List)} listing of matching {@link
+ * Schema#getField(String)} names.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions ...
+ * transactions.apply(
+ * CsvIO
+ * .<Transaction>write()
+ * .to("gs://bucket/path/to/folder/prefix")
+ * .withSchemaFields(Arrays.asList("transactionId", "purchaseAmount"))
+ * );
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header is repeated for every
+ * file, but only including the subset of fields in their listed order.
+ *
+ * <pre>{@code
+ * transactionId,purchaseAmount
+ * 12345,10.23
+ * 54321,54.65
+ * 98765,11.76
+ * }</pre>
+ *
+ * <p>Additionally, we may override the default {@link CSVFormat#DEFAULT}.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions ...
+ * transactions.apply(
+ * CsvIO
+ * .<Transaction>write()
+ * .to("gs://bucket/path/to/folder/prefix")
+ * .withCSVFormat(CSVFormat.POSTGRESQL_CSV)
+ * );
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following formatted according to {@link
+ * CSVFormat#POSTGRESQL_CSV}.
+ *
+ * <pre>{@code
+ * "bank","purchaseAmount","transactionId"
+ * "A","10.23","12345"
+ * "B","54.65","54321"
+ * "C","11.76","98765"
+ * }</pre>
+ *
+ * <p>A {@link PCollection} of {@link Row}s works just like custom Java types illustrated above,
+ * except we use {@link CsvIO#writeRows()} as shown below for the same {@code Transaction} class. We
+ * derive {@code Transaction}'s {@link Schema} using {@link
+ * org.apache.beam.sdk.schemas.annotations.DefaultSchema.DefaultSchemaProvider}.
+ *
+ * <pre>{@code
+ * DefaultSchemaProvider defaultSchemaProvider = new DefaultSchemaProvider();
+ * Schema schema = defaultSchemaProvider.schemaFor(TypeDescriptor.of(Transaction.class));
+ * PCollection<Row> transactions = pipeline.apply(Create.of(
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "A")
+ * .withFieldValue("purchaseAmount", 10.23)
+ * .withFieldValue("transactionId", "12345")
+ * .build(),
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "B")
+ * .withFieldValue("purchaseAmount", 54.65)
+ * .withFieldValue("transactionId", "54321")
+ * .build(),
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "C")
+ * .withFieldValue("purchaseAmount", 11.76)
+ * .withFieldValue("transactionId", "98765")
+ * .build()
+ * );
+ *
+ * transactions.apply(
+ * CsvIO
+ * .writeRows()
+ * .to("gs://bucket/path/to/folder/prefix")
+ * );
+ * }</pre>
+ *
+ * <p>Writing the transactions {@link PCollection} of {@link Row}s would yield the following CSV
+ * file content.
+ *
+ * <pre>{@code
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ */
+public class CsvIO {
+ /**
+ * The valid {@link Schema.FieldType} from which {@link CsvIO} converts CSV records.
+ *
+ * <p>{@link FieldType#BYTE}
+ *
+ * <p>{@link FieldType#BOOLEAN}
+ *
+ * <p>{@link FieldType#DATETIME}
+ *
+ * <p>{@link FieldType#DECIMAL}
+ *
+ * <p>{@link FieldType#DOUBLE}
+ *
+ * <p>{@link FieldType#INT16}
+ *
+ * <p>{@link FieldType#INT32}
+ *
+ * <p>{@link FieldType#INT64}
+ *
+ * <p>{@link FieldType#FLOAT}
+ *
+ * <p>{@link FieldType#STRING}
+ */
+ public static final Set<Schema.FieldType> VALID_FIELD_TYPE_SET =
+ ImmutableSet.of(
+ FieldType.BYTE,
+ FieldType.BOOLEAN,
+ FieldType.DATETIME,
+ FieldType.DECIMAL,
+ FieldType.DOUBLE,
+ FieldType.INT16,
+ FieldType.INT32,
+ FieldType.INT64,
+ FieldType.FLOAT,
+ FieldType.STRING);
+
+ static final String DEFAULT_FILENAME_SUFFIX = ".csv";
+
+ /** Instantiates a {@link Write} for writing user types in {@link CSVFormat#DEFAULT} format. */
+ public static <T> Write<T> write() {
+ return defaultWrite();
+ }
+
+ /** Instantiates a {@link Write} for {@link Row}s in {@link CSVFormat#DEFAULT} format. */
+ public static Write<Row> writeRows() {
+ return defaultWrite();
+ }
+
+ private static <T> Write<T> defaultWrite() {
+ return new AutoValue_CsvIO_Write.Builder<T>()
+ .setCSVFormat(CSVFormat.DEFAULT)
+ .setFileWrite(FileIO.write())
+ .build();
+ }
+
+ /** Instantiates a {@link Sink.Builder} for writing {@link Row} elements to CSV file sinks. */
+ static Sink.Builder<Row> sinkBuilder() {
+ return new AutoValue_CsvIO_Sink.Builder<>();
+ }
+
+ /** Implementation of {@link FileIO.Sink}. */
+ @AutoValue
+ abstract static class Sink<T> implements FileIO.Sink<T> {
+
+ private transient @Nullable PrintWriter writer;
+
+ /**
+ * Opens a {@link WritableByteChannel} for writing CSV files. Writes the {@link #getPreamble()}
+ * if available followed by the {@link #getHeader()}.
+ */
+ @Override
+ public void open(WritableByteChannel channel) throws IOException {
+ writer =
+ new PrintWriter(
+ new BufferedWriter(new OutputStreamWriter(Channels.newOutputStream(channel), UTF_8)));
+
+ if (getPreamble() != null) {
+ getWriter().println(getPreamble());
+ }
+ getWriter().println(getHeader());
+ }
+
+ /** Serializes and writes the element to a file. */
+ @Override
+ public void write(T element) throws IOException {
+ String line = getFormatFunction().apply(element);
+ getWriter().println(line);
+ }
+
+ @Override
+ public void flush() throws IOException {
+ getWriter().flush();
+ }
+
+ private PrintWriter getWriter() {
+ // resolves [dereference.of.nullable] error
+ Optional<PrintWriter> printWriter = Optional.ofNullable(writer);
+ checkState(printWriter.isPresent());
+ return printWriter.get();
+ }
+
+ /**
+ * Not to be confused with the CSV header, it is content written to the top of every sharded
+ * file prior to the header. See {@link Write#withPreamble(String)}.
+ */
+ @Nullable
+ abstract String getPreamble();
+
+ /**
+ * The column names of the CSV file written at the top line of each shard after the preamble, if
+ * available.
+ */
+ abstract String getHeader();
+
+ /** A {@link SerializableFunction} for converting to a CSV formatted string. */
+ abstract SerializableFunction<T, String> getFormatFunction();
+
+ @AutoValue.Builder
+ abstract static class Builder<T> {
+
+ /**
+ * Not to be confused with the CSV header, it is content written to the top of every sharded
+ * file prior to the header. See {@link Write#withPreamble(String)}.
+ */
+ abstract Builder<T> setPreamble(String value);
+
+ /**
+ * The column names of the CSV file written at the top line of each shard after the preamble,
+ * if available.
+ */
+ abstract Builder<T> setHeader(String value);
+
+ /** A {@link SerializableFunction} for converting a to a CSV formatted string. */
+ abstract Builder<T> setFormatFunction(SerializableFunction<T, String> value);
+
+ abstract Sink<T> build();
+ }
+ }
+
+ /** {@link PTransform} for writing CSV files. */
+ @AutoValue
+ public abstract static class Write<T> extends PTransform<PCollection<T>, PDone> {
+
+ /** Specifies a common prefix for all generated files. */
+ public Write<T> to(String filenamePrefix) {
+ Path path = Paths.get(filenamePrefix);
+ String directory = "";
+ String name = "";
+ if (path.getParent() != null) {
+ directory = path.getParent().toString();
+ }
+ if (path.getFileName() != null) {
+ name = path.getFileName().toString();
+ }
+ return toBuilder()
+ .setFileWrite(
+ FileIO.<Row>write()
+ .to(directory)
+ .withPrefix(name)
+ .withSuffix(DEFAULT_FILENAME_SUFFIX))
+ .build();
+ }
+
+ /**
+ * Overrides the default {@link CSVFormat#DEFAULT} with format of the destination CSV file data.
+ */
+ public Write<T> withCSVFormat(CSVFormat format) {
+ return toBuilder().setCSVFormat(format).build();
+ }
+
+ /**
+ * Controls the subset and order of included fields when writing records. Defaults to {@link
+ * Schema#getFieldNames()} of {@link Schema#sorted()} from {@link PCollection} input. {@link
+ * CsvIO.Write} will only validate whether {@link Field#getType()} of the {@link Schema}'s
+ * {@link Field#getName()}, included in schemaFields, matches any {@link #VALID_FIELD_TYPE_SET}.
+ */
+ public Write<T> withSchemaFields(List<String> schemaFields) {
+ return toBuilder().setSchemaFields(schemaFields).build();
+ }
+
+ /**
+ * Not to be confused with the CSV header, it is content written to the top of every sharded
+ * file prior to the header. In the example below, all the text proceeding the header
+ * 'column1,column2,column3' is the preamble.
+ *
+ * <pre>{@code
+ * Fake company, Inc.
+ * Lab experiment: abcdefg123456
+ * Experiment date: 2022-12-05
+ * Operator: John Doe
+ * column1,column2,colum3
+ * 1,2,3
+ * 4,5,6
+ * }</pre>
+ */
+ public Write<T> withPreamble(String preamble) {
+ return toBuilder().setPreamble(preamble).build();
+ }
+
+ /** Specifies the {@link Compression} of all generated shard files. */
+ public Write<T> withCompression(Compression compression) {
+ return toBuilder().setFileWrite(getFileWrite().withCompression(compression)).build();
+ }
+
+ /**
+ * Specifies to use a given fixed number of shards per window. See {@link
+ * FileIO.Write#withNumShards(int)} for details.
+ */
+ public Write<T> withNumShards(Integer numShards) {
+ return toBuilder().setFileWrite(getFileWrite().withNumShards(numShards)).build();
+ }
+
+ /**
+ * The sharding {@link PTransform} to use with the underlying {@link FileIO.Write}. Since {@link
+ * CsvIO.Write} converts user types to {@link Row}, the sharding transform is limited to {@link
+ * PCollection} of {@link Row}s.
+ */
+ public CsvIO.Write<T> withSharding(
+ PTransform<PCollection<Row>, PCollectionView<Integer>> sharding) {
+ return toBuilder().setFileWrite(getFileWrite().withSharding(sharding)).build();
+ }
+
+ /**
+ * The {@link Contextful} destinationFn to use with the underlying {@link FileIO.Write}. See
+ * {@link FileIO.Write#by(Contextful)} for more details.
+ */
+ public CsvIO.Write<T> by(Contextful<Contextful.Fn<Row, Void>> destinationFn) {
+ return toBuilder().setFileWrite(getFileWrite().by(destinationFn)).build();
+ }
+
+ /**
+ * The {@link Contextful} namingFn to use with the underlying {@link FileIO.Write}. See {@link
+ * FileIO.Write#withNaming(Contextful)} for more details.
+ */
+ public CsvIO.Write<T> withNaming(
+ Contextful<Contextful.Fn<Void, FileIO.Write.FileNaming>> namingFn) {
+ return toBuilder().setFileWrite(getFileWrite().withNaming(namingFn)).build();
+ }
+
+ /**
+ * The {@link FileIO.Write.FileNaming} naming to use with the underlying {@link FileIO.Write}.
+ * See {@link FileIO.Write#withNaming(FileNaming)} for more details.
+ */
+ public CsvIO.Write<T> withNaming(FileNaming naming) {
+ return toBuilder().setFileWrite(getFileWrite().withNaming(naming)).build();
+ }
+
+ /**
+ * The String prefix to use with the underlying {@link FileIO.Write}. See {@link
+ * FileIO.Write#withPrefix(String)} for more details.
+ */
+ public CsvIO.Write<T> withPrefix(String prefix) {
+ return toBuilder().setFileWrite(getFileWrite().withPrefix(prefix)).build();
+ }
+
+ /**
+ * Determines whether to use {@link FileIO.Write#withNoSpilling()} with the underlying {@link
+ * FileIO.Write}.
+ */
+ public CsvIO.Write<T> withNoSpilling() {
+ return toBuilder().setFileWrite(getFileWrite().withNoSpilling()).build();
+ }
+
+ /** The underlying {@link FileIO.Write} that writes converted input to CSV formatted output. */
+ abstract FileIO.Write<Void, Row> getFileWrite();
+
+ /** The {@link CSVFormat} to convert input. Defaults to {@link CSVFormat#DEFAULT}. */
+ abstract CSVFormat getCSVFormat();
+
+ /**
+ * The order and set of {@link Schema} fields driving the CSV conversion. See {@link
+ * #withSchemaFields(List)}.
+ */
+ @Nullable
+ abstract List<String> getSchemaFields();
+
+ /** The text preceding the CSV header. See {@link #withPreamble(String)}. */
+ @Nullable
+ abstract String getPreamble();
+
+ abstract Builder<T> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<T> {
+
+ /**
+ * The underlying {@link FileIO.Write} that writes converted input to CSV formatted output.
+ */
+ abstract Builder<T> setFileWrite(FileIO.Write<Void, Row> value);
+
+ /** The {@link CSVFormat} to convert input. Defaults to {@link CSVFormat#DEFAULT}. */
+ abstract Builder<T> setCSVFormat(CSVFormat value);
+
+ /**
+ * The order and set of {@link Schema} fields driving the CSV conversion. See {@link
+ * #withSchemaFields(List)}.
+ */
+ abstract Builder<T> setSchemaFields(List<String> value);
+
+ /** The text preceding the CSV header. See {@link #withPreamble(String)}. */
+ abstract Builder<T> setPreamble(String value);
+
+ abstract Write<T> build();
+ }
+
+ @Override
+ public PDone expand(PCollection<T> input) {
+ if (!input.hasSchema()) {
+ throw new IllegalArgumentException(
+ String.format(
+ "%s requires an input Schema. Note that only Row or user classes are supported. Consider using TextIO or FileIO directly when writing primitive types",
+ Write.class.getName()));
+ }
+
+ Schema schema = input.getSchema();
+ SerializableFunction<T, Row> toRowFn = input.getToRowFunction();
+ PCollection<Row> rows =
+ input.apply("To Rows", MapElements.into(rows()).via(toRowFn)).setRowSchema(schema);
+ rows.apply("Write Rows To CSV", getFileWrite().via(buildSink(schema)));
+ return PDone.in(input.getPipeline());
+ }
+
+ private Sink<Row> buildSink(Schema schema) {
+
+ List<String> schemaFields = getSchemaFields();
+ if (schemaFields == null) {
+ schemaFields = schema.sorted().getFieldNames();
+ }
+
+ RowToCsv rowToCsv =
+ RowToCsv.builder()
+ .setSchema(schema)
+ .setCSVFormat(getCSVFormat())
+ .setSchemaFields(schemaFields)
+ .build();
+
+ Sink.Builder<Row> builder = sinkBuilder();
+
+ // resolves [dereference.of.nullable] error
+ Optional<String> preamble = Optional.ofNullable(getPreamble());
+ if (preamble.isPresent()) {
+ builder = builder.setPreamble(preamble.get());
+ }
Review Comment:
Can this be simplified to:
```suggestion
String preamble = getPreamble();
if (preamble != null) {
builder = builder.setPreamble(preamble);
}
```
##########
sdks/java/io/csv/src/main/java/org/apache/beam/sdk/io/csv/CsvIO.java:
##########
@@ -0,0 +1,554 @@
+/*
+ * 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.sdk.io.csv;
+
+import static java.nio.charset.StandardCharsets.UTF_8;
+import static org.apache.beam.sdk.values.TypeDescriptors.rows;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
+
+import com.google.auto.value.AutoValue;
+import edu.umd.cs.findbugs.annotations.NonNull;
+import java.io.BufferedWriter;
+import java.io.IOException;
+import java.io.OutputStreamWriter;
+import java.io.PrintWriter;
+import java.nio.channels.Channels;
+import java.nio.channels.WritableByteChannel;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+import java.util.List;
+import java.util.Optional;
+import java.util.Set;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.io.Compression;
+import org.apache.beam.sdk.io.FileIO;
+import org.apache.beam.sdk.io.FileIO.Write.FileNaming;
+import org.apache.beam.sdk.io.csv.CsvRowConversions.RowToCsv;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.Field;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.transforms.Contextful;
+import org.apache.beam.sdk.transforms.MapElements;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PCollectionView;
+import org.apache.beam.sdk.values.PDone;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableSet;
+import org.apache.commons.csv.CSVFormat;
+
+/**
+ * {@link PTransform}s for reading and writing CSV files.
+ *
+ * <h2>Reading CSV files</h2>
+ *
+ * <p>Reading from CSV files is not yet implemented. Please see <a
+ * href="https://github.com/apache/beam/issues/24552">https://github.com/apache/beam/issues/24552</a>.
+ *
+ * <h2>Writing CSV files</h2>
+ *
+ * <p>To write a {@link PCollection} to one or more CSV files, use {@link CsvIO.Write}, using {@code
+ * CsvIO.write().to(String)}. {@link CsvIO.Write} supports writing {@link Row} or custom Java types
+ * using an inferred {@link Schema}. See the Beam Programming Guide on <a
+ * href="https://beam.apache.org/documentation/programming-guide/#inferring-schemas">Inferring
+ * schemas</a> for more information on how to enable Beam to infer a {@link Schema} from a custom
+ * Java type.
+ *
+ * <p>For each {@link Schema.Field} included in the {@link Schema}, according to {@link
+ * CsvIO.Write#withSchemaFields(List)}, {@link CsvIO.Write} restricts the use of flat {@link
+ * Schema}s and checks whether the {@link Schema.FieldType}s are one of {@link
+ * #VALID_FIELD_TYPE_SET}.
+ *
+ * <h3>Example usage:</h3>
+ *
+ * <p>Suppose we have a <code>Transaction</code> class annotated with
+ * {@code @DefaultSchema(JavaBeanSchema.class)} so that Beam can infer its {@link Schema}:
+ *
+ * <pre>{@code @DefaultSchema(JavaBeanSchema.class)
+ * public class Transaction {
+ * public Transaction() { … }
+ * public Long getTransactionId();
+ * public void setTransactionId(Long transactionId) { … }
+ * public String getBank() { … }
+ * public void setBank(String bank) { … }
+ * public double getPurchaseAmount() { … }
+ * public void setPurchaseAmount(double purchaseAmount) { … }
+ * }
+ * }</pre>
+ *
+ * <p>From a {@code PCollection<Transaction>}, we can write one or many CSV files and {@link
+ * CsvIO.Write} will automatically create the header based on its inferred {@link Schema}
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions ...
+ * transactions.apply(CsvIO.<Transaction>write().to("gs://bucket/path/to/folder/prefix"));
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header is repeated for every
+ * file, whereas by default, {@link CsvIO.Write} will write all fields in <b>sorted order</b> of the
+ * field names.
+ *
+ * <pre>{@code
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ *
+ * <p>{@link CsvIO.Write} allows the use a subset of the available fields or control the order of
+ * their output using {@link CsvIO.Write#withSchemaFields(List)} listing of matching {@link
+ * Schema#getField(String)} names.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions ...
+ * transactions.apply(
+ * CsvIO
+ * .<Transaction>write()
+ * .to("gs://bucket/path/to/folder/prefix")
+ * .withSchemaFields(Arrays.asList("transactionId", "purchaseAmount"))
+ * );
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header is repeated for every
+ * file, but only including the subset of fields in their listed order.
+ *
+ * <pre>{@code
+ * transactionId,purchaseAmount
+ * 12345,10.23
+ * 54321,54.65
+ * 98765,11.76
+ * }</pre>
+ *
+ * <p>Additionally, we may override the default {@link CSVFormat#DEFAULT}.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions ...
+ * transactions.apply(
+ * CsvIO
+ * .<Transaction>write()
+ * .to("gs://bucket/path/to/folder/prefix")
+ * .withCSVFormat(CSVFormat.POSTGRESQL_CSV)
+ * );
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following formatted according to {@link
+ * CSVFormat#POSTGRESQL_CSV}.
+ *
+ * <pre>{@code
+ * "bank","purchaseAmount","transactionId"
+ * "A","10.23","12345"
+ * "B","54.65","54321"
+ * "C","11.76","98765"
+ * }</pre>
+ *
+ * <p>A {@link PCollection} of {@link Row}s works just like custom Java types illustrated above,
+ * except we use {@link CsvIO#writeRows()} as shown below for the same {@code Transaction} class. We
+ * derive {@code Transaction}'s {@link Schema} using {@link
+ * org.apache.beam.sdk.schemas.annotations.DefaultSchema.DefaultSchemaProvider}.
+ *
+ * <pre>{@code
+ * DefaultSchemaProvider defaultSchemaProvider = new DefaultSchemaProvider();
+ * Schema schema = defaultSchemaProvider.schemaFor(TypeDescriptor.of(Transaction.class));
+ * PCollection<Row> transactions = pipeline.apply(Create.of(
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "A")
+ * .withFieldValue("purchaseAmount", 10.23)
+ * .withFieldValue("transactionId", "12345")
+ * .build(),
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "B")
+ * .withFieldValue("purchaseAmount", 54.65)
+ * .withFieldValue("transactionId", "54321")
+ * .build(),
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "C")
+ * .withFieldValue("purchaseAmount", 11.76)
+ * .withFieldValue("transactionId", "98765")
+ * .build()
+ * );
+ *
+ * transactions.apply(
+ * CsvIO
+ * .writeRows()
+ * .to("gs://bucket/path/to/folder/prefix")
+ * );
+ * }</pre>
+ *
+ * <p>Writing the transactions {@link PCollection} of {@link Row}s would yield the following CSV
+ * file content.
+ *
+ * <pre>{@code
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ */
+public class CsvIO {
+ /**
+ * The valid {@link Schema.FieldType} from which {@link CsvIO} converts CSV records.
+ *
+ * <p>{@link FieldType#BYTE}
+ *
+ * <p>{@link FieldType#BOOLEAN}
+ *
+ * <p>{@link FieldType#DATETIME}
+ *
+ * <p>{@link FieldType#DECIMAL}
+ *
+ * <p>{@link FieldType#DOUBLE}
+ *
+ * <p>{@link FieldType#INT16}
+ *
+ * <p>{@link FieldType#INT32}
+ *
+ * <p>{@link FieldType#INT64}
+ *
+ * <p>{@link FieldType#FLOAT}
+ *
+ * <p>{@link FieldType#STRING}
+ */
+ public static final Set<Schema.FieldType> VALID_FIELD_TYPE_SET =
+ ImmutableSet.of(
+ FieldType.BYTE,
+ FieldType.BOOLEAN,
+ FieldType.DATETIME,
+ FieldType.DECIMAL,
+ FieldType.DOUBLE,
+ FieldType.INT16,
+ FieldType.INT32,
+ FieldType.INT64,
+ FieldType.FLOAT,
+ FieldType.STRING);
+
+ static final String DEFAULT_FILENAME_SUFFIX = ".csv";
+
+ /** Instantiates a {@link Write} for writing user types in {@link CSVFormat#DEFAULT} format. */
+ public static <T> Write<T> write() {
+ return defaultWrite();
+ }
+
+ /** Instantiates a {@link Write} for {@link Row}s in {@link CSVFormat#DEFAULT} format. */
+ public static Write<Row> writeRows() {
+ return defaultWrite();
+ }
+
+ private static <T> Write<T> defaultWrite() {
+ return new AutoValue_CsvIO_Write.Builder<T>()
+ .setCSVFormat(CSVFormat.DEFAULT)
+ .setFileWrite(FileIO.write())
+ .build();
+ }
+
+ /** Instantiates a {@link Sink.Builder} for writing {@link Row} elements to CSV file sinks. */
+ static Sink.Builder<Row> sinkBuilder() {
+ return new AutoValue_CsvIO_Sink.Builder<>();
+ }
+
+ /** Implementation of {@link FileIO.Sink}. */
+ @AutoValue
+ abstract static class Sink<T> implements FileIO.Sink<T> {
+
+ private transient @Nullable PrintWriter writer;
+
+ /**
+ * Opens a {@link WritableByteChannel} for writing CSV files. Writes the {@link #getPreamble()}
+ * if available followed by the {@link #getHeader()}.
+ */
+ @Override
+ public void open(@NonNull WritableByteChannel channel) throws IOException {
+ writer =
+ new PrintWriter(
+ new BufferedWriter(new OutputStreamWriter(Channels.newOutputStream(channel), UTF_8)));
+
+ if (getPreamble() != null) {
+ getWriter().println(getPreamble());
+ }
+ getWriter().println(getHeader());
+ }
+
+ /** Serializes and writes the element to a file. */
+ @Override
+ public void write(T element) throws IOException {
+ String line = getFormatFunction().apply(element);
+ getWriter().println(line);
+ }
+
+ @Override
+ public void flush() throws IOException {
+ getWriter().flush();
+ }
+
+ private PrintWriter getWriter() {
+ // resolves [dereference.of.nullable] error
+ Optional<PrintWriter> printWriter = Optional.ofNullable(writer);
+ checkState(printWriter.isPresent());
+ return printWriter.get();
+ }
+
+ /**
+ * Not to be confused with the CSV header, it is content written to the top of every sharded
+ * file prior to the header. See {@link Write#withPreamble(String)}.
Review Comment:
```suggestion
* Content written to the top of every file prior to the header. See {@link Write#withPreamble}.
```
##########
sdks/java/io/csv/src/main/java/org/apache/beam/sdk/io/csv/CsvIO.java:
##########
@@ -0,0 +1,553 @@
+/*
+ * 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.sdk.io.csv;
+
+import static java.nio.charset.StandardCharsets.UTF_8;
+import static org.apache.beam.sdk.values.TypeDescriptors.rows;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
+
+import com.google.auto.value.AutoValue;
+import java.io.BufferedWriter;
+import java.io.IOException;
+import java.io.OutputStreamWriter;
+import java.io.PrintWriter;
+import java.nio.channels.Channels;
+import java.nio.channels.WritableByteChannel;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+import java.util.List;
+import java.util.Optional;
+import java.util.Set;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.io.Compression;
+import org.apache.beam.sdk.io.FileIO;
+import org.apache.beam.sdk.io.FileIO.Write.FileNaming;
+import org.apache.beam.sdk.io.csv.CsvRowConversions.RowToCsv;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.Field;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.transforms.Contextful;
+import org.apache.beam.sdk.transforms.MapElements;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PCollectionView;
+import org.apache.beam.sdk.values.PDone;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableSet;
+import org.apache.commons.csv.CSVFormat;
+
+/**
+ * {@link PTransform}s for reading and writing CSV files.
+ *
+ * <h2>Reading CSV files</h2>
+ *
+ * <p>Reading from CSV files is not yet implemented. Please see <a
+ * href="https://github.com/apache/beam/issues/24552">https://github.com/apache/beam/issues/24552</a>.
+ *
+ * <h2>Writing CSV files</h2>
+ *
+ * <p>To write a {@link PCollection} to one or more CSV files, use {@link CsvIO.Write}, using {@code
+ * CsvIO.write().to(String)}. {@link CsvIO.Write} supports writing {@link Row} or custom Java types
+ * using an inferred {@link Schema}. See the Beam Programming Guide on <a
+ * href="https://beam.apache.org/documentation/programming-guide/#inferring-schemas">Inferring
+ * schemas</a> for more information on how to enable Beam to infer a {@link Schema} from a custom
+ * Java type.
+ *
+ * <p>For each {@link Schema.Field} included in the {@link Schema}, according to {@link
+ * CsvIO.Write#withSchemaFields(List)}, {@link CsvIO.Write} restricts the use of flat {@link
+ * Schema}s and checks whether the {@link Schema.FieldType}s are one of {@link
+ * #VALID_FIELD_TYPE_SET}.
+ *
+ * <h3>Example usage:</h3>
+ *
+ * <p>Suppose we have a <code>Transaction</code> class annotated with
+ * {@code @DefaultSchema(JavaBeanSchema.class)} so that Beam can infer its {@link Schema}:
+ *
+ * <pre>{@code @DefaultSchema(JavaBeanSchema.class)
+ * public class Transaction {
+ * public Transaction() { … }
+ * public Long getTransactionId();
+ * public void setTransactionId(Long transactionId) { … }
+ * public String getBank() { … }
+ * public void setBank(String bank) { … }
+ * public double getPurchaseAmount() { … }
+ * public void setPurchaseAmount(double purchaseAmount) { … }
+ * }
+ * }</pre>
+ *
+ * <p>From a {@code PCollection<Transaction>}, we can write one or many CSV files and {@link
+ * CsvIO.Write} will automatically create the header based on its inferred {@link Schema}
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions ...
+ * transactions.apply(CsvIO.<Transaction>write().to("gs://bucket/path/to/folder/prefix"));
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header is repeated for every
+ * file, whereas by default, {@link CsvIO.Write} will write all fields in <b>sorted order</b> of the
+ * field names.
+ *
+ * <pre>{@code
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ *
+ * <p>{@link CsvIO.Write} allows the use a subset of the available fields or control the order of
+ * their output using {@link CsvIO.Write#withSchemaFields(List)} listing of matching {@link
+ * Schema#getField(String)} names.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions ...
+ * transactions.apply(
+ * CsvIO
+ * .<Transaction>write()
+ * .to("gs://bucket/path/to/folder/prefix")
+ * .withSchemaFields(Arrays.asList("transactionId", "purchaseAmount"))
+ * );
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header is repeated for every
+ * file, but only including the subset of fields in their listed order.
+ *
+ * <pre>{@code
+ * transactionId,purchaseAmount
+ * 12345,10.23
+ * 54321,54.65
+ * 98765,11.76
+ * }</pre>
+ *
+ * <p>Additionally, we may override the default {@link CSVFormat#DEFAULT}.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions ...
+ * transactions.apply(
+ * CsvIO
+ * .<Transaction>write()
+ * .to("gs://bucket/path/to/folder/prefix")
+ * .withCSVFormat(CSVFormat.POSTGRESQL_CSV)
+ * );
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following formatted according to {@link
+ * CSVFormat#POSTGRESQL_CSV}.
+ *
+ * <pre>{@code
+ * "bank","purchaseAmount","transactionId"
+ * "A","10.23","12345"
+ * "B","54.65","54321"
+ * "C","11.76","98765"
+ * }</pre>
+ *
+ * <p>A {@link PCollection} of {@link Row}s works just like custom Java types illustrated above,
+ * except we use {@link CsvIO#writeRows()} as shown below for the same {@code Transaction} class. We
+ * derive {@code Transaction}'s {@link Schema} using {@link
+ * org.apache.beam.sdk.schemas.annotations.DefaultSchema.DefaultSchemaProvider}.
+ *
+ * <pre>{@code
+ * DefaultSchemaProvider defaultSchemaProvider = new DefaultSchemaProvider();
+ * Schema schema = defaultSchemaProvider.schemaFor(TypeDescriptor.of(Transaction.class));
+ * PCollection<Row> transactions = pipeline.apply(Create.of(
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "A")
+ * .withFieldValue("purchaseAmount", 10.23)
+ * .withFieldValue("transactionId", "12345")
+ * .build(),
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "B")
+ * .withFieldValue("purchaseAmount", 54.65)
+ * .withFieldValue("transactionId", "54321")
+ * .build(),
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "C")
+ * .withFieldValue("purchaseAmount", 11.76)
+ * .withFieldValue("transactionId", "98765")
+ * .build()
+ * );
+ *
+ * transactions.apply(
+ * CsvIO
+ * .writeRows()
+ * .to("gs://bucket/path/to/folder/prefix")
+ * );
+ * }</pre>
+ *
+ * <p>Writing the transactions {@link PCollection} of {@link Row}s would yield the following CSV
+ * file content.
+ *
+ * <pre>{@code
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ */
+public class CsvIO {
+ /**
+ * The valid {@link Schema.FieldType} from which {@link CsvIO} converts CSV records.
+ *
+ * <p>{@link FieldType#BYTE}
+ *
+ * <p>{@link FieldType#BOOLEAN}
+ *
+ * <p>{@link FieldType#DATETIME}
+ *
+ * <p>{@link FieldType#DECIMAL}
+ *
+ * <p>{@link FieldType#DOUBLE}
+ *
+ * <p>{@link FieldType#INT16}
+ *
+ * <p>{@link FieldType#INT32}
+ *
+ * <p>{@link FieldType#INT64}
+ *
+ * <p>{@link FieldType#FLOAT}
+ *
+ * <p>{@link FieldType#STRING}
+ */
+ public static final Set<Schema.FieldType> VALID_FIELD_TYPE_SET =
+ ImmutableSet.of(
+ FieldType.BYTE,
+ FieldType.BOOLEAN,
+ FieldType.DATETIME,
+ FieldType.DECIMAL,
+ FieldType.DOUBLE,
+ FieldType.INT16,
+ FieldType.INT32,
+ FieldType.INT64,
+ FieldType.FLOAT,
+ FieldType.STRING);
+
+ static final String DEFAULT_FILENAME_SUFFIX = ".csv";
+
+ /** Instantiates a {@link Write} for writing user types in {@link CSVFormat#DEFAULT} format. */
+ public static <T> Write<T> write() {
+ return defaultWrite();
+ }
+
+ /** Instantiates a {@link Write} for {@link Row}s in {@link CSVFormat#DEFAULT} format. */
+ public static Write<Row> writeRows() {
+ return defaultWrite();
+ }
+
+ private static <T> Write<T> defaultWrite() {
+ return new AutoValue_CsvIO_Write.Builder<T>()
+ .setCSVFormat(CSVFormat.DEFAULT)
+ .setFileWrite(FileIO.write())
+ .build();
+ }
+
+ /** Instantiates a {@link Sink.Builder} for writing {@link Row} elements to CSV file sinks. */
+ static Sink.Builder<Row> sinkBuilder() {
+ return new AutoValue_CsvIO_Sink.Builder<>();
+ }
+
+ /** Implementation of {@link FileIO.Sink}. */
+ @AutoValue
+ abstract static class Sink<T> implements FileIO.Sink<T> {
+
+ private transient @Nullable PrintWriter writer;
+
+ /**
+ * Opens a {@link WritableByteChannel} for writing CSV files. Writes the {@link #getPreamble()}
+ * if available followed by the {@link #getHeader()}.
+ */
+ @Override
+ public void open(WritableByteChannel channel) throws IOException {
+ writer =
+ new PrintWriter(
+ new BufferedWriter(new OutputStreamWriter(Channels.newOutputStream(channel), UTF_8)));
+
+ if (getPreamble() != null) {
+ getWriter().println(getPreamble());
+ }
+ getWriter().println(getHeader());
+ }
+
+ /** Serializes and writes the element to a file. */
+ @Override
+ public void write(T element) throws IOException {
+ String line = getFormatFunction().apply(element);
+ getWriter().println(line);
+ }
+
+ @Override
+ public void flush() throws IOException {
+ getWriter().flush();
+ }
+
+ private PrintWriter getWriter() {
+ // resolves [dereference.of.nullable] error
+ Optional<PrintWriter> printWriter = Optional.ofNullable(writer);
+ checkState(printWriter.isPresent());
+ return printWriter.get();
+ }
+
+ /**
+ * Not to be confused with the CSV header, it is content written to the top of every sharded
+ * file prior to the header. See {@link Write#withPreamble(String)}.
+ */
+ @Nullable
+ abstract String getPreamble();
+
+ /**
+ * The column names of the CSV file written at the top line of each shard after the preamble, if
+ * available.
+ */
+ abstract String getHeader();
+
+ /** A {@link SerializableFunction} for converting to a CSV formatted string. */
+ abstract SerializableFunction<T, String> getFormatFunction();
+
+ @AutoValue.Builder
+ abstract static class Builder<T> {
+
+ /**
+ * Not to be confused with the CSV header, it is content written to the top of every sharded
+ * file prior to the header. See {@link Write#withPreamble(String)}.
+ */
+ abstract Builder<T> setPreamble(String value);
+
+ /**
+ * The column names of the CSV file written at the top line of each shard after the preamble,
+ * if available.
+ */
+ abstract Builder<T> setHeader(String value);
+
+ /** A {@link SerializableFunction} for converting a to a CSV formatted string. */
+ abstract Builder<T> setFormatFunction(SerializableFunction<T, String> value);
+
+ abstract Sink<T> build();
+ }
+ }
+
+ /** {@link PTransform} for writing CSV files. */
+ @AutoValue
+ public abstract static class Write<T> extends PTransform<PCollection<T>, PDone> {
+
+ /** Specifies a common prefix for all generated files. */
+ public Write<T> to(String filenamePrefix) {
+ Path path = Paths.get(filenamePrefix);
+ String directory = "";
+ String name = "";
+ if (path.getParent() != null) {
+ directory = path.getParent().toString();
+ }
+ if (path.getFileName() != null) {
+ name = path.getFileName().toString();
+ }
+ return toBuilder()
+ .setFileWrite(
+ FileIO.<Row>write()
+ .to(directory)
+ .withPrefix(name)
+ .withSuffix(DEFAULT_FILENAME_SUFFIX))
+ .build();
Review Comment:
Can this be simplified to:
```suggestion
return toBuilder()
.setFileWrite(
FileIO.<Row>write()
.to(filenamePrefix)
.withSuffix(DEFAULT_FILENAME_SUFFIX))
.build();
```
##########
sdks/java/io/csv/src/main/java/org/apache/beam/sdk/io/csv/CsvIO.java:
##########
@@ -0,0 +1,554 @@
+/*
+ * 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.sdk.io.csv;
+
+import static java.nio.charset.StandardCharsets.UTF_8;
+import static org.apache.beam.sdk.values.TypeDescriptors.rows;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
+
+import com.google.auto.value.AutoValue;
+import edu.umd.cs.findbugs.annotations.NonNull;
+import java.io.BufferedWriter;
+import java.io.IOException;
+import java.io.OutputStreamWriter;
+import java.io.PrintWriter;
+import java.nio.channels.Channels;
+import java.nio.channels.WritableByteChannel;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+import java.util.List;
+import java.util.Optional;
+import java.util.Set;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.io.Compression;
+import org.apache.beam.sdk.io.FileIO;
+import org.apache.beam.sdk.io.FileIO.Write.FileNaming;
+import org.apache.beam.sdk.io.csv.CsvRowConversions.RowToCsv;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.Field;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.transforms.Contextful;
+import org.apache.beam.sdk.transforms.MapElements;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PCollectionView;
+import org.apache.beam.sdk.values.PDone;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableSet;
+import org.apache.commons.csv.CSVFormat;
+
+/**
+ * {@link PTransform}s for reading and writing CSV files.
+ *
+ * <h2>Reading CSV files</h2>
+ *
+ * <p>Reading from CSV files is not yet implemented. Please see <a
+ * href="https://github.com/apache/beam/issues/24552">https://github.com/apache/beam/issues/24552</a>.
+ *
+ * <h2>Writing CSV files</h2>
+ *
+ * <p>To write a {@link PCollection} to one or more CSV files, use {@link CsvIO.Write}, using {@code
+ * CsvIO.write().to(String)}. {@link CsvIO.Write} supports writing {@link Row} or custom Java types
+ * using an inferred {@link Schema}. See the Beam Programming Guide on <a
+ * href="https://beam.apache.org/documentation/programming-guide/#inferring-schemas">Inferring
+ * schemas</a> for more information on how to enable Beam to infer a {@link Schema} from a custom
+ * Java type.
+ *
+ * <p>For each {@link Schema.Field} included in the {@link Schema}, according to {@link
+ * CsvIO.Write#withSchemaFields(List)}, {@link CsvIO.Write} restricts the use of flat {@link
+ * Schema}s and checks whether the {@link Schema.FieldType}s are one of {@link
+ * #VALID_FIELD_TYPE_SET}.
+ *
+ * <h3>Example usage:</h3>
+ *
+ * <p>Suppose we have a <code>Transaction</code> class annotated with
Review Comment:
```suggestion
* <p>Suppose we have the following <code>Transaction</code> class annotated with
```
##########
sdks/java/io/csv/src/test/java/org/apache/beam/sdk/io/csv/CsvIOWriteTest.java:
##########
@@ -0,0 +1,534 @@
+/*
+ * 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.sdk.io.csv;
+
+import static org.apache.beam.sdk.io.csv.CsvIOTestHelpers.ALL_DATA_TYPES_SCHEMA;
+import static org.apache.beam.sdk.io.csv.CsvIOTestHelpers.rowOf;
+import static org.junit.Assert.assertArrayEquals;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertThrows;
+import static org.junit.Assert.assertTrue;
+
+import java.io.BufferedReader;
+import java.io.File;
+import java.io.IOException;
+import java.math.BigDecimal;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.regex.Pattern;
+import java.util.stream.Collectors;
+import org.apache.beam.sdk.io.Compression;
+import org.apache.beam.sdk.io.csv.CsvIOTestHelpers.AllDataTypes;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.Field;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.testing.TestPipeline;
+import org.apache.beam.sdk.transforms.Create;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Charsets;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.commons.csv.CSVFormat;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Instant;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.TemporaryFolder;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/** Tests for {@link CsvIO.Write}. */
+@RunWith(JUnit4.class)
+public class CsvIOWriteTest {
+
+ @Rule public TemporaryFolder tmpFolder = TemporaryFolder.builder().build();
+
+ @Rule public TestPipeline pipeline = TestPipeline.create();
+
+ @Rule
+ public TestPipeline errorPipeline = TestPipeline.create().enableAbandonedNodeEnforcement(false);
+
+ private final List<AllDataTypes> userTypes =
+ Arrays.asList(
+ AllDataTypes.of(
+ true,
+ (byte) 0,
+ Instant.ofEpochMilli(0L).toDateTime(),
+ BigDecimal.valueOf(0L),
+ 0.1,
+ 0.2f,
+ (short) 3,
+ 4,
+ 5L,
+ "a"),
+ AllDataTypes.of(
+ true,
+ (byte) 1,
+ Instant.ofEpochMilli(1L).toDateTime(),
+ BigDecimal.valueOf(1L),
+ 1.1,
+ 1.2f,
+ (short) 4,
+ 5,
+ 6L,
+ "b"),
+ AllDataTypes.of(
+ false,
+ (byte) 2,
+ Instant.ofEpochMilli(2L).toDateTime(),
+ BigDecimal.valueOf(2L),
+ 2.0,
+ 2.1f,
+ (short) 5,
+ 6,
+ 7L,
+ "c"));
+
+ private final List<Row> rows =
+ Arrays.asList(
+ rowOf(
+ true,
+ (byte) 0,
+ Instant.ofEpochMilli(0L).toDateTime(),
+ BigDecimal.valueOf(0L),
+ 0.1,
+ 0.2f,
+ (short) 3,
+ 4,
+ 5L,
+ "a"),
+ rowOf(
+ true,
+ (byte) 1,
+ Instant.ofEpochMilli(1L).toDateTime(),
+ BigDecimal.valueOf(1L),
+ 1.1,
+ 1.2f,
+ (short) 4,
+ 5,
+ 6L,
+ "b"),
+ rowOf(
+ false,
+ (byte) 2,
+ Instant.ofEpochMilli(2L).toDateTime(),
+ BigDecimal.valueOf(2L),
+ 2.0,
+ 2.1f,
+ (short) 5,
+ 6,
+ 7L,
+ "c"));
+
+ @Test
+ public void validateSchema() {
+ String to = tmpFolder.getRoot().getAbsolutePath() + "/somewhere";
+ Schema emptySchema = Schema.of();
+ Schema bytesSchema =
+ Schema.of(Field.of("badfield", FieldType.BYTES), Field.of("ok", FieldType.STRING));
+ Schema mapSchema =
+ Schema.of(
+ Field.of("badfield", FieldType.map(FieldType.STRING, FieldType.STRING)),
+ Field.of("ok", FieldType.STRING));
+ Schema arraySchema =
+ Schema.of(
+ Field.of("badfield", FieldType.array(FieldType.STRING)),
+ Field.of("ok", FieldType.STRING));
+ Schema rowSchema =
+ Schema.of(
+ Field.of("badfield", FieldType.row(ALL_DATA_TYPES_SCHEMA)),
+ Field.of("ok", FieldType.STRING));
+
+ assertThrows(
+ "schema should not be empty",
+ IllegalArgumentException.class,
+ () ->
+ errorPipeline
+ .apply(Create.of(Row.nullRow(emptySchema)).withRowSchema(emptySchema))
+ .apply(CsvIO.writeRows().to(to)));
+
+ assertThrows(
+ "schema should not contain bytes field",
+ IllegalArgumentException.class,
+ () ->
+ errorPipeline
+ .apply(
+ Create.of(
+ Row.withSchema(bytesSchema)
+ .withFieldValue("badfield", new byte[] {1, 2, 3})
+ .withFieldValue("ok", "abcdefg")
+ .build()))
+ .setRowSchema(bytesSchema)
+ .apply(CsvIO.writeRows().withCSVFormat(CSVFormat.DEFAULT).to(to)));
+
+ assertThrows(
+ "schema should not contain map field",
+ IllegalArgumentException.class,
+ () ->
+ errorPipeline
+ .apply(
+ Create.of(
+ Row.withSchema(mapSchema)
+ .withFieldValue("badfield", ImmutableMap.of("a", "b"))
+ .withFieldValue("ok", "abcdefg")
+ .build()))
+ .setRowSchema(mapSchema)
+ .apply(CsvIO.writeRows().to(to)));
+
+ assertThrows(
+ "schema should not contain array field",
+ IllegalArgumentException.class,
+ () ->
+ errorPipeline
+ .apply(
+ Create.of(
+ Row.withSchema(arraySchema)
+ .withFieldValue("badfield", Arrays.asList("a", "b"))
+ .withFieldValue("ok", "abcdefg")
+ .build()))
+ .setRowSchema(arraySchema)
+ .apply(CsvIO.writeRows().withCSVFormat(CSVFormat.DEFAULT).to(to)));
+
+ assertThrows(
+ "schema should not contain row field",
+ IllegalArgumentException.class,
+ () ->
+ errorPipeline
+ .apply(
+ Create.of(
+ Row.withSchema(rowSchema)
+ .withFieldValue("badfield", rows.get(0))
+ .withFieldValue("ok", "abcdefg")
+ .build()))
+ .setRowSchema(rowSchema)
+ .apply(CsvIO.writeRows().to(to)));
+ }
+
+ @Test
+ public void invalidSchemaFields() {
+ String to = tmpFolder.getRoot().getAbsolutePath() + "/somewhere";
+
+ assertThrows(
+ "schema field names should not be empty",
+ IllegalArgumentException.class,
+ () ->
+ errorPipeline
+ .apply(Create.of(rows))
+ .setRowSchema(ALL_DATA_TYPES_SCHEMA)
+ .apply(
+ CsvIO.writeRows()
+ .to(to)
+ .withSchemaFields(Arrays.asList("", "string", "anInt"))));
+
+ assertThrows(
+ "schema fields should not be an empty list",
+ IllegalArgumentException.class,
+ () ->
+ errorPipeline
+ .apply(Create.of(rows))
+ .setRowSchema(ALL_DATA_TYPES_SCHEMA)
+ .apply(CsvIO.writeRows().to(to).withSchemaFields(Collections.emptyList())));
+
+ assertThrows(
+ "field names should match schema",
+ IllegalArgumentException.class,
+ () ->
+ errorPipeline
+ .apply(Create.of(rows))
+ .setRowSchema(ALL_DATA_TYPES_SCHEMA)
+ .apply(
+ CsvIO.writeRows()
+ .to(to)
+ .withSchemaFields(Arrays.asList("doesnotexist", "string", "anInt"))));
+ }
+
+ @Test
+ public void withCompression() {
+ String to = "with_compression";
+ String regex = String.format("^%s.*$", to);
+ pipeline
+ .apply(Create.of(rows).withRowSchema(ALL_DATA_TYPES_SCHEMA))
+ .apply(
+ CsvIO.writeRows()
+ .to(tmpFolder.getRoot().getAbsolutePath() + "/" + to)
+ .withCompression(Compression.GZIP));
+ pipeline.run().waitUntilFinish();
+ String[] files = filesMatching(regex);
+ for (String name : files) {
+ Path p = Paths.get(tmpFolder.getRoot().getAbsolutePath(), name);
+ File f = new File(p.toString());
+ assertTrue(Compression.GZIP.isCompressed(f.getAbsolutePath()));
+ }
+ }
+
+ @Test
+ public void withInvalidType() {
+ assertThrows(
+ IllegalArgumentException.class,
+ () ->
+ errorPipeline
+ .apply(Create.of("1,2,3", "4.5.6"))
+ .apply(
+ CsvIO.<String>write().to(tmpFolder.getRoot().getAbsolutePath() + "/badtype")));
+ }
+
+ @Test
+ public void withNumShards() {
+ String to = "with_num_shards";
+ String regex = String.format("^%s.*$", to);
+ pipeline
+ .apply(Create.of(rows).withRowSchema(ALL_DATA_TYPES_SCHEMA))
+ .apply(
+ CsvIO.writeRows()
+ .to(tmpFolder.getRoot().getAbsolutePath() + "/" + to)
+ .withNumShards(1));
+ pipeline.run().waitUntilFinish();
+ String[] files = filesMatching(regex);
+ assertEquals(1, files.length);
+ }
+
+ @Test
+ public void withPreamble() throws IOException {
+ String preamble =
+ "id=44ebe8ee-dc0f-4c93-8b47-2e627e1eb988\ndate=2022-01-02\ntest=e3bbfc47-9bcb-46f0-8944-829827ee5222";
+ String to = "rows_preamble";
+ String regex = String.format("^%s.*$", to);
+ pipeline
+ .apply(Create.of(rows).withRowSchema(ALL_DATA_TYPES_SCHEMA))
+ .apply(
+ CsvIO.writeRows()
+ .to(tmpFolder.getRoot().getAbsolutePath() + "/" + to)
+ .withPreamble(preamble));
+ pipeline.run().waitUntilFinish();
+ String[] files = filesMatching(regex);
+ assertTrue(
+ "CsvIO.writeRows should write to " + tmpFolder.getRoot().getAbsolutePath() + "/" + to,
+ files.length > 0);
+ String expectedHeader =
+ "aBoolean,aByte,aDouble,aFloat,aLong,aShort,anInt,dateTime,decimal,string";
+ for (String name : files) {
+ Path p = Paths.get(tmpFolder.getRoot().getAbsolutePath(), name);
+ File f = new File(p.toString());
+ assertFileContentsMatchInAnyOrder(
+ preamble,
+ expectedHeader,
+ f,
+ "true,0,0.1,0.2,5,3,4,1970-01-01T00:00:00.000Z,0,a",
+ "true,1,1.1,1.2,6,4,5,1970-01-01T00:00:00.001Z,1,b",
+ "false,2,2.0,2.1,7,5,6,1970-01-01T00:00:00.002Z,2,c");
+ }
+ }
+
+ @Test
+ public void rowsWithCsvFormatNonDefault() throws IOException {
+ String to = "rows_non_default";
+ String regex = String.format("^%s.*$", to);
+ pipeline
+ .apply(Create.of(rows).withRowSchema(ALL_DATA_TYPES_SCHEMA))
+ .apply(
+ CsvIO.writeRows()
+ .to(tmpFolder.getRoot().getAbsolutePath() + "/" + to)
+ .withCSVFormat(CSVFormat.MYSQL));
+ pipeline.run().waitUntilFinish();
+ String[] files = filesMatching(regex);
+ assertTrue(
+ "CsvIO.writeRows should write to " + tmpFolder.getRoot().getAbsolutePath() + "/" + to,
+ files.length > 0);
+ String expectedHeader =
+ "aBoolean\taByte\taDouble\taFloat\taLong\taShort\tanInt\tdateTime\tdecimal\tstring";
+ for (String name : files) {
+ Path p = Paths.get(tmpFolder.getRoot().getAbsolutePath(), name);
+ File f = new File(p.toString());
+ assertFileContentsMatchInAnyOrder(
+ null,
+ expectedHeader,
+ f,
+ "true\t0\t0.1\t0.2\t5\t3\t4\t1970-01-01T00:00:00.000Z\t0\ta",
+ "true\t1\t1.1\t1.2\t6\t4\t5\t1970-01-01T00:00:00.001Z\t1\tb",
+ "false\t2\t2.0\t2.1\t7\t5\t6\t1970-01-01T00:00:00.002Z\t2\tc");
+ }
+ }
+
+ @Test
+ public void withSchemaFields() throws IOException {
+ String to = "with_schema_fields";
+ String regex = String.format("^%s.*$", to);
+ pipeline
+ .apply(Create.of(rows).withRowSchema(ALL_DATA_TYPES_SCHEMA))
+ .apply(
+ CsvIO.writeRows()
+ .to(tmpFolder.getRoot().getAbsolutePath() + "/" + to)
+ .withSchemaFields(Arrays.asList("aShort", "aBoolean", "aLong")));
+ pipeline.run().waitUntilFinish();
+ String[] files = filesMatching(regex);
+ assertTrue(
+ "CsvIO.writeRows should write to " + tmpFolder.getRoot().getAbsolutePath() + "/" + to,
+ files.length > 0);
+ String expectedHeader = "aShort,aBoolean,aLong";
+ for (String name : files) {
+ Path p = Paths.get(tmpFolder.getRoot().getAbsolutePath(), name);
+ File f = new File(p.toString());
+ assertFileContentsMatchInAnyOrder(
+ null, expectedHeader, f, "3,true,5", "4,true,6", "5,false,7");
+ }
+ }
+
+ @Test
+ public void rowsWithDefaults() throws IOException {
+ String to = "rows";
+ String regex = String.format("^%s.*$", to);
+ pipeline
+ .apply(Create.of(rows).withRowSchema(ALL_DATA_TYPES_SCHEMA))
+ .apply(CsvIO.writeRows().to(tmpFolder.getRoot().getAbsolutePath() + "/" + to));
+ pipeline.run().waitUntilFinish();
+ String[] files = filesMatching(regex);
+ assertTrue(
+ "CsvIO.writeRows should write to " + tmpFolder.getRoot().getAbsolutePath() + "/" + to,
+ files.length > 0);
+ String expectedHeader =
+ "aBoolean,aByte,aDouble,aFloat,aLong,aShort,anInt,dateTime,decimal,string";
+ for (String name : files) {
+ Pattern matchCSVSuffix = Pattern.compile("^.*csv$");
+ assertTrue(
+ String.format("%s should match %s", name, matchCSVSuffix.pattern()),
+ matchCSVSuffix.matcher(name).matches());
+ Path p = Paths.get(tmpFolder.getRoot().getAbsolutePath(), name);
+ File f = new File(p.toString());
+ assertFileContentsMatchInAnyOrder(
+ null,
+ expectedHeader,
+ f,
+ "true,0,0.1,0.2,5,3,4,1970-01-01T00:00:00.000Z,0,a",
+ "true,1,1.1,1.2,6,4,5,1970-01-01T00:00:00.001Z,1,b",
+ "false,2,2.0,2.1,7,5,6,1970-01-01T00:00:00.002Z,2,c");
+ }
+ }
+
+ @Test
+ public void userTypesWithDefaults() throws IOException {
+ String to = "user_types";
+ String regex = String.format("^%s.*$", to);
+ pipeline
+ .apply(Create.of(userTypes))
+ .apply(CsvIO.<AllDataTypes>write().to(tmpFolder.getRoot().getAbsolutePath() + "/" + to));
+ pipeline.run().waitUntilFinish();
+ String[] files = filesMatching(regex);
+ assertTrue(
+ "CsvIO.writeRows should write to " + tmpFolder.getRoot().getAbsolutePath() + "/" + to,
+ files.length > 0);
+ String expectedHeader =
+ "aBoolean,aByte,aDouble,aFloat,aLong,aShort,anInt,dateTime,decimal,string";
+ for (String name : files) {
+ Path p = Paths.get(tmpFolder.getRoot().getAbsolutePath(), name);
+ File f = new File(p.toString());
+ assertFileContentsMatchInAnyOrder(
+ null,
+ expectedHeader,
+ f,
+ "true,0,0.1,0.2,5,3,4,1970-01-01T00:00:00.000Z,0,a",
+ "true,1,1.1,1.2,6,4,5,1970-01-01T00:00:00.001Z,1,b",
+ "false,2,2.0,2.1,7,5,6,1970-01-01T00:00:00.002Z,2,c");
+ }
+ }
+
+ @Test
+ public void rowsWithNulls() throws IOException {
+ String to = "rows_with_nulls";
+ String regex = String.format("^%s.*$", to);
+ Schema schema =
+ Schema.of(
+ Field.of("string", FieldType.STRING),
+ Field.nullable("aDouble", FieldType.DOUBLE),
+ Field.nullable("anInt", FieldType.INT32));
+
+ List<Row> nullableRows =
+ Arrays.asList(
+ Row.withSchema(schema).withFieldValue("string", "a").build(),
+ Row.withSchema(schema)
+ .withFieldValue("string", "b")
+ .withFieldValue("aDouble", 1.2345)
+ .build(),
+ Row.withSchema(schema).withFieldValue("string", "c").withFieldValue("anInt", 2).build(),
+ Row.withSchema(schema)
+ .withFieldValue("string", "d")
+ .withFieldValue("aDouble", 3.1234)
+ .withFieldValue("anInt", 3)
+ .build());
+
+ pipeline
+ .apply(Create.of(nullableRows).withRowSchema(schema))
+ .apply(CsvIO.writeRows().to(tmpFolder.getRoot().getAbsolutePath() + "/" + to));
+
+ pipeline.run().waitUntilFinish();
+ String[] files = filesMatching(regex);
+ assertTrue(
+ "CsvIO.writeRows should write to " + tmpFolder.getRoot().getAbsolutePath() + "/" + to,
+ files.length > 0);
+ String expectedHeader = "aDouble,anInt,string";
+ for (String name : files) {
+ Path p = Paths.get(tmpFolder.getRoot().getAbsolutePath(), name);
+ File f = new File(p.toString());
+ assertFileContentsMatchInAnyOrder(
+ null, expectedHeader, f, ",,a", "1.2345,,b", ",2,c", "3.1234,3,d");
+ }
+ }
+
+ private static List<String> readLinesFromFile(File f) throws IOException {
+ List<String> currentFile = new ArrayList<>();
+ try (BufferedReader reader = Files.newBufferedReader(f.toPath(), Charsets.UTF_8)) {
+ while (true) {
+ String line = reader.readLine();
+ if (line == null) {
+ break;
+ }
+ currentFile.add(line);
+ }
+ }
+ return currentFile;
+ }
+
+ private void assertFileContentsMatchInAnyOrder(
+ @Nullable String preamble, String header, File f, String... expected) throws IOException {
+ List<String> actual = readLinesFromFile(f);
+ int headerIndex = 0;
+ if (preamble != null) {
+ String[] preambleLines = preamble.split("\n");
+ assertArrayEquals(preambleLines, actual.subList(0, preambleLines.length).toArray());
+ headerIndex = preambleLines.length;
+ }
+
+ assertEquals("header should match", header, actual.get(headerIndex));
+
+ String[] rest = actual.subList(headerIndex + 1, actual.size()).toArray(new String[0]);
+ List<String> expectedList = Arrays.stream(expected).collect(Collectors.toList());
+ String expectedString = String.join(", ", expectedList);
+ for (String line : rest) {
+ String message = String.format("[%s]\nshould contain %s", expectedString, line);
+ assertTrue(message, expectedList.contains(line));
+ }
Review Comment:
```suggestion
List<String> rest = actual.subList(headerIndex + 1, actual.size());
assertThat(rest, containsInAnyOrder(expected));
```
##########
sdks/java/io/csv/src/main/java/org/apache/beam/sdk/io/csv/CsvIO.java:
##########
@@ -0,0 +1,553 @@
+/*
+ * 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.sdk.io.csv;
+
+import static java.nio.charset.StandardCharsets.UTF_8;
+import static org.apache.beam.sdk.values.TypeDescriptors.rows;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
+
+import com.google.auto.value.AutoValue;
+import java.io.BufferedWriter;
+import java.io.IOException;
+import java.io.OutputStreamWriter;
+import java.io.PrintWriter;
+import java.nio.channels.Channels;
+import java.nio.channels.WritableByteChannel;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+import java.util.List;
+import java.util.Optional;
+import java.util.Set;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.io.Compression;
+import org.apache.beam.sdk.io.FileIO;
+import org.apache.beam.sdk.io.FileIO.Write.FileNaming;
+import org.apache.beam.sdk.io.csv.CsvRowConversions.RowToCsv;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.Field;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.transforms.Contextful;
+import org.apache.beam.sdk.transforms.MapElements;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PCollectionView;
+import org.apache.beam.sdk.values.PDone;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableSet;
+import org.apache.commons.csv.CSVFormat;
+
+/**
+ * {@link PTransform}s for reading and writing CSV files.
+ *
+ * <h2>Reading CSV files</h2>
+ *
+ * <p>Reading from CSV files is not yet implemented. Please see <a
+ * href="https://github.com/apache/beam/issues/24552">https://github.com/apache/beam/issues/24552</a>.
+ *
+ * <h2>Writing CSV files</h2>
+ *
+ * <p>To write a {@link PCollection} to one or more CSV files, use {@link CsvIO.Write}, using {@code
+ * CsvIO.write().to(String)}. {@link CsvIO.Write} supports writing {@link Row} or custom Java types
+ * using an inferred {@link Schema}. See the Beam Programming Guide on <a
+ * href="https://beam.apache.org/documentation/programming-guide/#inferring-schemas">Inferring
+ * schemas</a> for more information on how to enable Beam to infer a {@link Schema} from a custom
+ * Java type.
+ *
+ * <p>For each {@link Schema.Field} included in the {@link Schema}, according to {@link
+ * CsvIO.Write#withSchemaFields(List)}, {@link CsvIO.Write} restricts the use of flat {@link
+ * Schema}s and checks whether the {@link Schema.FieldType}s are one of {@link
+ * #VALID_FIELD_TYPE_SET}.
+ *
+ * <h3>Example usage:</h3>
+ *
+ * <p>Suppose we have a <code>Transaction</code> class annotated with
+ * {@code @DefaultSchema(JavaBeanSchema.class)} so that Beam can infer its {@link Schema}:
+ *
+ * <pre>{@code @DefaultSchema(JavaBeanSchema.class)
+ * public class Transaction {
+ * public Transaction() { … }
+ * public Long getTransactionId();
+ * public void setTransactionId(Long transactionId) { … }
+ * public String getBank() { … }
+ * public void setBank(String bank) { … }
+ * public double getPurchaseAmount() { … }
+ * public void setPurchaseAmount(double purchaseAmount) { … }
+ * }
+ * }</pre>
+ *
+ * <p>From a {@code PCollection<Transaction>}, we can write one or many CSV files and {@link
+ * CsvIO.Write} will automatically create the header based on its inferred {@link Schema}
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions ...
+ * transactions.apply(CsvIO.<Transaction>write().to("gs://bucket/path/to/folder/prefix"));
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header is repeated for every
+ * file, whereas by default, {@link CsvIO.Write} will write all fields in <b>sorted order</b> of the
+ * field names.
+ *
+ * <pre>{@code
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ *
+ * <p>{@link CsvIO.Write} allows the use a subset of the available fields or control the order of
+ * their output using {@link CsvIO.Write#withSchemaFields(List)} listing of matching {@link
+ * Schema#getField(String)} names.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions ...
+ * transactions.apply(
+ * CsvIO
+ * .<Transaction>write()
+ * .to("gs://bucket/path/to/folder/prefix")
+ * .withSchemaFields(Arrays.asList("transactionId", "purchaseAmount"))
+ * );
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header is repeated for every
+ * file, but only including the subset of fields in their listed order.
+ *
+ * <pre>{@code
+ * transactionId,purchaseAmount
+ * 12345,10.23
+ * 54321,54.65
+ * 98765,11.76
+ * }</pre>
+ *
+ * <p>Additionally, we may override the default {@link CSVFormat#DEFAULT}.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions ...
+ * transactions.apply(
+ * CsvIO
+ * .<Transaction>write()
+ * .to("gs://bucket/path/to/folder/prefix")
+ * .withCSVFormat(CSVFormat.POSTGRESQL_CSV)
+ * );
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following formatted according to {@link
+ * CSVFormat#POSTGRESQL_CSV}.
+ *
+ * <pre>{@code
+ * "bank","purchaseAmount","transactionId"
+ * "A","10.23","12345"
+ * "B","54.65","54321"
+ * "C","11.76","98765"
+ * }</pre>
+ *
+ * <p>A {@link PCollection} of {@link Row}s works just like custom Java types illustrated above,
+ * except we use {@link CsvIO#writeRows()} as shown below for the same {@code Transaction} class. We
+ * derive {@code Transaction}'s {@link Schema} using {@link
+ * org.apache.beam.sdk.schemas.annotations.DefaultSchema.DefaultSchemaProvider}.
+ *
+ * <pre>{@code
+ * DefaultSchemaProvider defaultSchemaProvider = new DefaultSchemaProvider();
+ * Schema schema = defaultSchemaProvider.schemaFor(TypeDescriptor.of(Transaction.class));
+ * PCollection<Row> transactions = pipeline.apply(Create.of(
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "A")
+ * .withFieldValue("purchaseAmount", 10.23)
+ * .withFieldValue("transactionId", "12345")
+ * .build(),
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "B")
+ * .withFieldValue("purchaseAmount", 54.65)
+ * .withFieldValue("transactionId", "54321")
+ * .build(),
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "C")
+ * .withFieldValue("purchaseAmount", 11.76)
+ * .withFieldValue("transactionId", "98765")
+ * .build()
+ * );
+ *
+ * transactions.apply(
+ * CsvIO
+ * .writeRows()
+ * .to("gs://bucket/path/to/folder/prefix")
+ * );
+ * }</pre>
+ *
+ * <p>Writing the transactions {@link PCollection} of {@link Row}s would yield the following CSV
+ * file content.
+ *
+ * <pre>{@code
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ */
+public class CsvIO {
+ /**
+ * The valid {@link Schema.FieldType} from which {@link CsvIO} converts CSV records.
+ *
+ * <p>{@link FieldType#BYTE}
+ *
+ * <p>{@link FieldType#BOOLEAN}
+ *
+ * <p>{@link FieldType#DATETIME}
+ *
+ * <p>{@link FieldType#DECIMAL}
+ *
+ * <p>{@link FieldType#DOUBLE}
+ *
+ * <p>{@link FieldType#INT16}
+ *
+ * <p>{@link FieldType#INT32}
+ *
+ * <p>{@link FieldType#INT64}
+ *
+ * <p>{@link FieldType#FLOAT}
+ *
+ * <p>{@link FieldType#STRING}
+ */
+ public static final Set<Schema.FieldType> VALID_FIELD_TYPE_SET =
+ ImmutableSet.of(
+ FieldType.BYTE,
+ FieldType.BOOLEAN,
+ FieldType.DATETIME,
+ FieldType.DECIMAL,
+ FieldType.DOUBLE,
+ FieldType.INT16,
+ FieldType.INT32,
+ FieldType.INT64,
+ FieldType.FLOAT,
+ FieldType.STRING);
+
+ static final String DEFAULT_FILENAME_SUFFIX = ".csv";
+
+ /** Instantiates a {@link Write} for writing user types in {@link CSVFormat#DEFAULT} format. */
+ public static <T> Write<T> write() {
+ return defaultWrite();
+ }
+
+ /** Instantiates a {@link Write} for {@link Row}s in {@link CSVFormat#DEFAULT} format. */
+ public static Write<Row> writeRows() {
+ return defaultWrite();
+ }
+
+ private static <T> Write<T> defaultWrite() {
+ return new AutoValue_CsvIO_Write.Builder<T>()
+ .setCSVFormat(CSVFormat.DEFAULT)
+ .setFileWrite(FileIO.write())
+ .build();
+ }
+
+ /** Instantiates a {@link Sink.Builder} for writing {@link Row} elements to CSV file sinks. */
+ static Sink.Builder<Row> sinkBuilder() {
+ return new AutoValue_CsvIO_Sink.Builder<>();
+ }
+
+ /** Implementation of {@link FileIO.Sink}. */
+ @AutoValue
+ abstract static class Sink<T> implements FileIO.Sink<T> {
+
+ private transient @Nullable PrintWriter writer;
+
+ /**
+ * Opens a {@link WritableByteChannel} for writing CSV files. Writes the {@link #getPreamble()}
+ * if available followed by the {@link #getHeader()}.
+ */
+ @Override
+ public void open(WritableByteChannel channel) throws IOException {
+ writer =
+ new PrintWriter(
+ new BufferedWriter(new OutputStreamWriter(Channels.newOutputStream(channel), UTF_8)));
+
+ if (getPreamble() != null) {
+ getWriter().println(getPreamble());
+ }
+ getWriter().println(getHeader());
+ }
+
+ /** Serializes and writes the element to a file. */
+ @Override
+ public void write(T element) throws IOException {
+ String line = getFormatFunction().apply(element);
+ getWriter().println(line);
+ }
+
+ @Override
+ public void flush() throws IOException {
+ getWriter().flush();
+ }
+
+ private PrintWriter getWriter() {
+ // resolves [dereference.of.nullable] error
+ Optional<PrintWriter> printWriter = Optional.ofNullable(writer);
+ checkState(printWriter.isPresent());
+ return printWriter.get();
+ }
+
+ /**
+ * Not to be confused with the CSV header, it is content written to the top of every sharded
+ * file prior to the header. See {@link Write#withPreamble(String)}.
+ */
+ @Nullable
+ abstract String getPreamble();
+
+ /**
+ * The column names of the CSV file written at the top line of each shard after the preamble, if
+ * available.
+ */
+ abstract String getHeader();
+
+ /** A {@link SerializableFunction} for converting to a CSV formatted string. */
+ abstract SerializableFunction<T, String> getFormatFunction();
+
+ @AutoValue.Builder
+ abstract static class Builder<T> {
+
+ /**
+ * Not to be confused with the CSV header, it is content written to the top of every sharded
+ * file prior to the header. See {@link Write#withPreamble(String)}.
+ */
+ abstract Builder<T> setPreamble(String value);
+
+ /**
+ * The column names of the CSV file written at the top line of each shard after the preamble,
+ * if available.
+ */
+ abstract Builder<T> setHeader(String value);
+
+ /** A {@link SerializableFunction} for converting a to a CSV formatted string. */
+ abstract Builder<T> setFormatFunction(SerializableFunction<T, String> value);
+
+ abstract Sink<T> build();
+ }
+ }
+
+ /** {@link PTransform} for writing CSV files. */
+ @AutoValue
+ public abstract static class Write<T> extends PTransform<PCollection<T>, PDone> {
+
+ /** Specifies a common prefix for all generated files. */
+ public Write<T> to(String filenamePrefix) {
+ Path path = Paths.get(filenamePrefix);
+ String directory = "";
+ String name = "";
+ if (path.getParent() != null) {
+ directory = path.getParent().toString();
+ }
+ if (path.getFileName() != null) {
+ name = path.getFileName().toString();
+ }
+ return toBuilder()
+ .setFileWrite(
+ FileIO.<Row>write()
+ .to(directory)
+ .withPrefix(name)
+ .withSuffix(DEFAULT_FILENAME_SUFFIX))
+ .build();
+ }
+
+ /**
+ * Overrides the default {@link CSVFormat#DEFAULT} with format of the destination CSV file data.
+ */
+ public Write<T> withCSVFormat(CSVFormat format) {
+ return toBuilder().setCSVFormat(format).build();
+ }
+
+ /**
+ * Controls the subset and order of included fields when writing records. Defaults to {@link
+ * Schema#getFieldNames()} of {@link Schema#sorted()} from {@link PCollection} input. {@link
+ * CsvIO.Write} will only validate whether {@link Field#getType()} of the {@link Schema}'s
+ * {@link Field#getName()}, included in schemaFields, matches any {@link #VALID_FIELD_TYPE_SET}.
+ */
+ public Write<T> withSchemaFields(List<String> schemaFields) {
Review Comment:
I think we should drop this field and instead rely on CSVFormat.withHeader and point people to that. This way we don't have a point of confusion where the CSVFormat specifies a header and the user specifies withSchemaFields.
We can explicitly tell people that we will filter down to the set of fields in CSVFormat.withHeader if specified.
##########
sdks/java/io/csv/src/main/java/org/apache/beam/sdk/io/csv/CsvIO.java:
##########
@@ -0,0 +1,553 @@
+/*
+ * 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.sdk.io.csv;
+
+import static java.nio.charset.StandardCharsets.UTF_8;
+import static org.apache.beam.sdk.values.TypeDescriptors.rows;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
+
+import com.google.auto.value.AutoValue;
+import java.io.BufferedWriter;
+import java.io.IOException;
+import java.io.OutputStreamWriter;
+import java.io.PrintWriter;
+import java.nio.channels.Channels;
+import java.nio.channels.WritableByteChannel;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+import java.util.List;
+import java.util.Optional;
+import java.util.Set;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.io.Compression;
+import org.apache.beam.sdk.io.FileIO;
+import org.apache.beam.sdk.io.FileIO.Write.FileNaming;
+import org.apache.beam.sdk.io.csv.CsvRowConversions.RowToCsv;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.Field;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.transforms.Contextful;
+import org.apache.beam.sdk.transforms.MapElements;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PCollectionView;
+import org.apache.beam.sdk.values.PDone;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableSet;
+import org.apache.commons.csv.CSVFormat;
+
+/**
+ * {@link PTransform}s for reading and writing CSV files.
+ *
+ * <h2>Reading CSV files</h2>
+ *
+ * <p>Reading from CSV files is not yet implemented. Please see <a
+ * href="https://github.com/apache/beam/issues/24552">https://github.com/apache/beam/issues/24552</a>.
+ *
+ * <h2>Writing CSV files</h2>
+ *
+ * <p>To write a {@link PCollection} to one or more CSV files, use {@link CsvIO.Write}, using {@code
+ * CsvIO.write().to(String)}. {@link CsvIO.Write} supports writing {@link Row} or custom Java types
+ * using an inferred {@link Schema}. See the Beam Programming Guide on <a
+ * href="https://beam.apache.org/documentation/programming-guide/#inferring-schemas">Inferring
+ * schemas</a> for more information on how to enable Beam to infer a {@link Schema} from a custom
+ * Java type.
+ *
+ * <p>For each {@link Schema.Field} included in the {@link Schema}, according to {@link
+ * CsvIO.Write#withSchemaFields(List)}, {@link CsvIO.Write} restricts the use of flat {@link
+ * Schema}s and checks whether the {@link Schema.FieldType}s are one of {@link
+ * #VALID_FIELD_TYPE_SET}.
+ *
+ * <h3>Example usage:</h3>
+ *
+ * <p>Suppose we have a <code>Transaction</code> class annotated with
+ * {@code @DefaultSchema(JavaBeanSchema.class)} so that Beam can infer its {@link Schema}:
+ *
+ * <pre>{@code @DefaultSchema(JavaBeanSchema.class)
+ * public class Transaction {
+ * public Transaction() { … }
+ * public Long getTransactionId();
+ * public void setTransactionId(Long transactionId) { … }
+ * public String getBank() { … }
+ * public void setBank(String bank) { … }
+ * public double getPurchaseAmount() { … }
+ * public void setPurchaseAmount(double purchaseAmount) { … }
+ * }
+ * }</pre>
+ *
+ * <p>From a {@code PCollection<Transaction>}, we can write one or many CSV files and {@link
+ * CsvIO.Write} will automatically create the header based on its inferred {@link Schema}
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions ...
+ * transactions.apply(CsvIO.<Transaction>write().to("gs://bucket/path/to/folder/prefix"));
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header is repeated for every
+ * file, whereas by default, {@link CsvIO.Write} will write all fields in <b>sorted order</b> of the
+ * field names.
+ *
+ * <pre>{@code
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ *
+ * <p>{@link CsvIO.Write} allows the use a subset of the available fields or control the order of
+ * their output using {@link CsvIO.Write#withSchemaFields(List)} listing of matching {@link
+ * Schema#getField(String)} names.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions ...
+ * transactions.apply(
+ * CsvIO
+ * .<Transaction>write()
+ * .to("gs://bucket/path/to/folder/prefix")
+ * .withSchemaFields(Arrays.asList("transactionId", "purchaseAmount"))
+ * );
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header is repeated for every
+ * file, but only including the subset of fields in their listed order.
+ *
+ * <pre>{@code
+ * transactionId,purchaseAmount
+ * 12345,10.23
+ * 54321,54.65
+ * 98765,11.76
+ * }</pre>
+ *
+ * <p>Additionally, we may override the default {@link CSVFormat#DEFAULT}.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions ...
+ * transactions.apply(
+ * CsvIO
+ * .<Transaction>write()
+ * .to("gs://bucket/path/to/folder/prefix")
+ * .withCSVFormat(CSVFormat.POSTGRESQL_CSV)
+ * );
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following formatted according to {@link
+ * CSVFormat#POSTGRESQL_CSV}.
+ *
+ * <pre>{@code
+ * "bank","purchaseAmount","transactionId"
+ * "A","10.23","12345"
+ * "B","54.65","54321"
+ * "C","11.76","98765"
+ * }</pre>
+ *
+ * <p>A {@link PCollection} of {@link Row}s works just like custom Java types illustrated above,
+ * except we use {@link CsvIO#writeRows()} as shown below for the same {@code Transaction} class. We
+ * derive {@code Transaction}'s {@link Schema} using {@link
+ * org.apache.beam.sdk.schemas.annotations.DefaultSchema.DefaultSchemaProvider}.
+ *
+ * <pre>{@code
+ * DefaultSchemaProvider defaultSchemaProvider = new DefaultSchemaProvider();
+ * Schema schema = defaultSchemaProvider.schemaFor(TypeDescriptor.of(Transaction.class));
+ * PCollection<Row> transactions = pipeline.apply(Create.of(
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "A")
+ * .withFieldValue("purchaseAmount", 10.23)
+ * .withFieldValue("transactionId", "12345")
+ * .build(),
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "B")
+ * .withFieldValue("purchaseAmount", 54.65)
+ * .withFieldValue("transactionId", "54321")
+ * .build(),
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "C")
+ * .withFieldValue("purchaseAmount", 11.76)
+ * .withFieldValue("transactionId", "98765")
+ * .build()
+ * );
+ *
+ * transactions.apply(
+ * CsvIO
+ * .writeRows()
+ * .to("gs://bucket/path/to/folder/prefix")
+ * );
+ * }</pre>
+ *
+ * <p>Writing the transactions {@link PCollection} of {@link Row}s would yield the following CSV
+ * file content.
+ *
+ * <pre>{@code
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ */
+public class CsvIO {
+ /**
+ * The valid {@link Schema.FieldType} from which {@link CsvIO} converts CSV records.
+ *
+ * <p>{@link FieldType#BYTE}
+ *
+ * <p>{@link FieldType#BOOLEAN}
+ *
+ * <p>{@link FieldType#DATETIME}
+ *
+ * <p>{@link FieldType#DECIMAL}
+ *
+ * <p>{@link FieldType#DOUBLE}
+ *
+ * <p>{@link FieldType#INT16}
+ *
+ * <p>{@link FieldType#INT32}
+ *
+ * <p>{@link FieldType#INT64}
+ *
+ * <p>{@link FieldType#FLOAT}
+ *
+ * <p>{@link FieldType#STRING}
+ */
+ public static final Set<Schema.FieldType> VALID_FIELD_TYPE_SET =
+ ImmutableSet.of(
+ FieldType.BYTE,
+ FieldType.BOOLEAN,
+ FieldType.DATETIME,
+ FieldType.DECIMAL,
+ FieldType.DOUBLE,
+ FieldType.INT16,
+ FieldType.INT32,
+ FieldType.INT64,
+ FieldType.FLOAT,
+ FieldType.STRING);
+
+ static final String DEFAULT_FILENAME_SUFFIX = ".csv";
+
+ /** Instantiates a {@link Write} for writing user types in {@link CSVFormat#DEFAULT} format. */
+ public static <T> Write<T> write() {
+ return defaultWrite();
+ }
+
+ /** Instantiates a {@link Write} for {@link Row}s in {@link CSVFormat#DEFAULT} format. */
+ public static Write<Row> writeRows() {
+ return defaultWrite();
+ }
+
+ private static <T> Write<T> defaultWrite() {
+ return new AutoValue_CsvIO_Write.Builder<T>()
+ .setCSVFormat(CSVFormat.DEFAULT)
+ .setFileWrite(FileIO.write())
+ .build();
+ }
+
+ /** Instantiates a {@link Sink.Builder} for writing {@link Row} elements to CSV file sinks. */
+ static Sink.Builder<Row> sinkBuilder() {
+ return new AutoValue_CsvIO_Sink.Builder<>();
+ }
+
+ /** Implementation of {@link FileIO.Sink}. */
+ @AutoValue
+ abstract static class Sink<T> implements FileIO.Sink<T> {
+
+ private transient @Nullable PrintWriter writer;
+
+ /**
+ * Opens a {@link WritableByteChannel} for writing CSV files. Writes the {@link #getPreamble()}
+ * if available followed by the {@link #getHeader()}.
+ */
+ @Override
+ public void open(WritableByteChannel channel) throws IOException {
+ writer =
+ new PrintWriter(
+ new BufferedWriter(new OutputStreamWriter(Channels.newOutputStream(channel), UTF_8)));
+
+ if (getPreamble() != null) {
+ getWriter().println(getPreamble());
+ }
+ getWriter().println(getHeader());
+ }
+
+ /** Serializes and writes the element to a file. */
+ @Override
+ public void write(T element) throws IOException {
+ String line = getFormatFunction().apply(element);
+ getWriter().println(line);
+ }
+
+ @Override
+ public void flush() throws IOException {
+ getWriter().flush();
+ }
+
+ private PrintWriter getWriter() {
+ // resolves [dereference.of.nullable] error
+ Optional<PrintWriter> printWriter = Optional.ofNullable(writer);
+ checkState(printWriter.isPresent());
+ return printWriter.get();
+ }
+
+ /**
+ * Not to be confused with the CSV header, it is content written to the top of every sharded
+ * file prior to the header. See {@link Write#withPreamble(String)}.
+ */
+ @Nullable
+ abstract String getPreamble();
+
+ /**
+ * The column names of the CSV file written at the top line of each shard after the preamble, if
+ * available.
+ */
+ abstract String getHeader();
+
+ /** A {@link SerializableFunction} for converting to a CSV formatted string. */
+ abstract SerializableFunction<T, String> getFormatFunction();
+
+ @AutoValue.Builder
+ abstract static class Builder<T> {
+
+ /**
+ * Not to be confused with the CSV header, it is content written to the top of every sharded
+ * file prior to the header. See {@link Write#withPreamble(String)}.
+ */
+ abstract Builder<T> setPreamble(String value);
+
+ /**
+ * The column names of the CSV file written at the top line of each shard after the preamble,
+ * if available.
+ */
+ abstract Builder<T> setHeader(String value);
+
+ /** A {@link SerializableFunction} for converting a to a CSV formatted string. */
+ abstract Builder<T> setFormatFunction(SerializableFunction<T, String> value);
+
+ abstract Sink<T> build();
+ }
+ }
+
+ /** {@link PTransform} for writing CSV files. */
+ @AutoValue
+ public abstract static class Write<T> extends PTransform<PCollection<T>, PDone> {
+
+ /** Specifies a common prefix for all generated files. */
+ public Write<T> to(String filenamePrefix) {
+ Path path = Paths.get(filenamePrefix);
+ String directory = "";
+ String name = "";
+ if (path.getParent() != null) {
+ directory = path.getParent().toString();
+ }
+ if (path.getFileName() != null) {
+ name = path.getFileName().toString();
+ }
+ return toBuilder()
+ .setFileWrite(
+ FileIO.<Row>write()
+ .to(directory)
+ .withPrefix(name)
+ .withSuffix(DEFAULT_FILENAME_SUFFIX))
+ .build();
+ }
+
+ /**
+ * Overrides the default {@link CSVFormat#DEFAULT} with format of the destination CSV file data.
+ */
+ public Write<T> withCSVFormat(CSVFormat format) {
+ return toBuilder().setCSVFormat(format).build();
+ }
+
+ /**
+ * Controls the subset and order of included fields when writing records. Defaults to {@link
+ * Schema#getFieldNames()} of {@link Schema#sorted()} from {@link PCollection} input. {@link
+ * CsvIO.Write} will only validate whether {@link Field#getType()} of the {@link Schema}'s
+ * {@link Field#getName()}, included in schemaFields, matches any {@link #VALID_FIELD_TYPE_SET}.
+ */
+ public Write<T> withSchemaFields(List<String> schemaFields) {
+ return toBuilder().setSchemaFields(schemaFields).build();
+ }
+
+ /**
+ * Not to be confused with the CSV header, it is content written to the top of every sharded
+ * file prior to the header. In the example below, all the text proceeding the header
+ * 'column1,column2,column3' is the preamble.
+ *
+ * <pre>{@code
+ * Fake company, Inc.
Review Comment:
Should the preamble be replaced with header comments from CSVFormat?
##########
sdks/java/io/csv/src/main/java/org/apache/beam/sdk/io/csv/CsvIO.java:
##########
@@ -0,0 +1,554 @@
+/*
+ * 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.sdk.io.csv;
+
+import static java.nio.charset.StandardCharsets.UTF_8;
+import static org.apache.beam.sdk.values.TypeDescriptors.rows;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
+
+import com.google.auto.value.AutoValue;
+import edu.umd.cs.findbugs.annotations.NonNull;
+import java.io.BufferedWriter;
+import java.io.IOException;
+import java.io.OutputStreamWriter;
+import java.io.PrintWriter;
+import java.nio.channels.Channels;
+import java.nio.channels.WritableByteChannel;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+import java.util.List;
+import java.util.Optional;
+import java.util.Set;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.io.Compression;
+import org.apache.beam.sdk.io.FileIO;
+import org.apache.beam.sdk.io.FileIO.Write.FileNaming;
+import org.apache.beam.sdk.io.csv.CsvRowConversions.RowToCsv;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.Field;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.transforms.Contextful;
+import org.apache.beam.sdk.transforms.MapElements;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PCollectionView;
+import org.apache.beam.sdk.values.PDone;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableSet;
+import org.apache.commons.csv.CSVFormat;
+
+/**
+ * {@link PTransform}s for reading and writing CSV files.
+ *
+ * <h2>Reading CSV files</h2>
+ *
+ * <p>Reading from CSV files is not yet implemented. Please see <a
+ * href="https://github.com/apache/beam/issues/24552">https://github.com/apache/beam/issues/24552</a>.
+ *
+ * <h2>Writing CSV files</h2>
+ *
+ * <p>To write a {@link PCollection} to one or more CSV files, use {@link CsvIO.Write}, using {@code
+ * CsvIO.write().to(String)}. {@link CsvIO.Write} supports writing {@link Row} or custom Java types
+ * using an inferred {@link Schema}. See the Beam Programming Guide on <a
+ * href="https://beam.apache.org/documentation/programming-guide/#inferring-schemas">Inferring
+ * schemas</a> for more information on how to enable Beam to infer a {@link Schema} from a custom
+ * Java type.
+ *
+ * <p>For each {@link Schema.Field} included in the {@link Schema}, according to {@link
+ * CsvIO.Write#withSchemaFields(List)}, {@link CsvIO.Write} restricts the use of flat {@link
+ * Schema}s and checks whether the {@link Schema.FieldType}s are one of {@link
+ * #VALID_FIELD_TYPE_SET}.
+ *
+ * <h3>Example usage:</h3>
+ *
+ * <p>Suppose we have a <code>Transaction</code> class annotated with
+ * {@code @DefaultSchema(JavaBeanSchema.class)} so that Beam can infer its {@link Schema}:
+ *
+ * <pre>{@code @DefaultSchema(JavaBeanSchema.class)
+ * public class Transaction {
+ * public Transaction() { … }
+ * public Long getTransactionId();
+ * public void setTransactionId(Long transactionId) { … }
+ * public String getBank() { … }
+ * public void setBank(String bank) { … }
+ * public double getPurchaseAmount() { … }
+ * public void setPurchaseAmount(double purchaseAmount) { … }
+ * }
+ * }</pre>
+ *
+ * <p>From a {@code PCollection<Transaction>}, we can write one or many CSV files and {@link
+ * CsvIO.Write} will automatically create the header based on its inferred {@link Schema}
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions ...
+ * transactions.apply(CsvIO.<Transaction>write().to("gs://bucket/path/to/folder/prefix"));
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header is repeated for every
+ * file, whereas by default, {@link CsvIO.Write} will write all fields in <b>sorted order</b> of the
+ * field names.
+ *
+ * <pre>{@code
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ *
+ * <p>{@link CsvIO.Write} allows the use a subset of the available fields or control the order of
+ * their output using {@link CsvIO.Write#withSchemaFields(List)} listing of matching {@link
+ * Schema#getField(String)} names.
Review Comment:
```suggestion
* <p>{@link CsvIO.Write} allows outputting a subset of the available fields or control the order of
* their output using {@link CsvIO.Write#withSchemaFields} listing of matching {@link
* Schema schema field names}.
```
##########
sdks/java/io/csv/src/main/java/org/apache/beam/sdk/io/csv/CsvIO.java:
##########
@@ -0,0 +1,554 @@
+/*
+ * 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.sdk.io.csv;
+
+import static java.nio.charset.StandardCharsets.UTF_8;
+import static org.apache.beam.sdk.values.TypeDescriptors.rows;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
+
+import com.google.auto.value.AutoValue;
+import edu.umd.cs.findbugs.annotations.NonNull;
+import java.io.BufferedWriter;
+import java.io.IOException;
+import java.io.OutputStreamWriter;
+import java.io.PrintWriter;
+import java.nio.channels.Channels;
+import java.nio.channels.WritableByteChannel;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+import java.util.List;
+import java.util.Optional;
+import java.util.Set;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.io.Compression;
+import org.apache.beam.sdk.io.FileIO;
+import org.apache.beam.sdk.io.FileIO.Write.FileNaming;
+import org.apache.beam.sdk.io.csv.CsvRowConversions.RowToCsv;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.Field;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.transforms.Contextful;
+import org.apache.beam.sdk.transforms.MapElements;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PCollectionView;
+import org.apache.beam.sdk.values.PDone;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableSet;
+import org.apache.commons.csv.CSVFormat;
+
+/**
+ * {@link PTransform}s for reading and writing CSV files.
+ *
+ * <h2>Reading CSV files</h2>
+ *
+ * <p>Reading from CSV files is not yet implemented. Please see <a
+ * href="https://github.com/apache/beam/issues/24552">https://github.com/apache/beam/issues/24552</a>.
+ *
+ * <h2>Writing CSV files</h2>
+ *
+ * <p>To write a {@link PCollection} to one or more CSV files, use {@link CsvIO.Write}, using {@code
+ * CsvIO.write().to(String)}. {@link CsvIO.Write} supports writing {@link Row} or custom Java types
+ * using an inferred {@link Schema}. See the Beam Programming Guide on <a
+ * href="https://beam.apache.org/documentation/programming-guide/#inferring-schemas">Inferring
+ * schemas</a> for more information on how to enable Beam to infer a {@link Schema} from a custom
+ * Java type.
+ *
+ * <p>For each {@link Schema.Field} included in the {@link Schema}, according to {@link
+ * CsvIO.Write#withSchemaFields(List)}, {@link CsvIO.Write} restricts the use of flat {@link
+ * Schema}s and checks whether the {@link Schema.FieldType}s are one of {@link
+ * #VALID_FIELD_TYPE_SET}.
+ *
+ * <h3>Example usage:</h3>
+ *
+ * <p>Suppose we have a <code>Transaction</code> class annotated with
+ * {@code @DefaultSchema(JavaBeanSchema.class)} so that Beam can infer its {@link Schema}:
+ *
+ * <pre>{@code @DefaultSchema(JavaBeanSchema.class)
+ * public class Transaction {
+ * public Transaction() { … }
+ * public Long getTransactionId();
+ * public void setTransactionId(Long transactionId) { … }
+ * public String getBank() { … }
+ * public void setBank(String bank) { … }
+ * public double getPurchaseAmount() { … }
+ * public void setPurchaseAmount(double purchaseAmount) { … }
+ * }
+ * }</pre>
+ *
+ * <p>From a {@code PCollection<Transaction>}, we can write one or many CSV files and {@link
+ * CsvIO.Write} will automatically create the header based on its inferred {@link Schema}
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions ...
+ * transactions.apply(CsvIO.<Transaction>write().to("gs://bucket/path/to/folder/prefix"));
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header is repeated for every
+ * file, whereas by default, {@link CsvIO.Write} will write all fields in <b>sorted order</b> of the
+ * field names.
+ *
+ * <pre>{@code
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ *
+ * <p>{@link CsvIO.Write} allows the use a subset of the available fields or control the order of
+ * their output using {@link CsvIO.Write#withSchemaFields(List)} listing of matching {@link
+ * Schema#getField(String)} names.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions ...
+ * transactions.apply(
+ * CsvIO
+ * .<Transaction>write()
+ * .to("gs://bucket/path/to/folder/prefix")
+ * .withSchemaFields(Arrays.asList("transactionId", "purchaseAmount"))
+ * );
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header is repeated for every
+ * file, but only including the subset of fields in their listed order.
+ *
+ * <pre>{@code
+ * transactionId,purchaseAmount
+ * 12345,10.23
+ * 54321,54.65
+ * 98765,11.76
+ * }</pre>
+ *
+ * <p>Additionally, we may override the default {@link CSVFormat#DEFAULT}.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions ...
+ * transactions.apply(
+ * CsvIO
+ * .<Transaction>write()
+ * .to("gs://bucket/path/to/folder/prefix")
+ * .withCSVFormat(CSVFormat.POSTGRESQL_CSV)
+ * );
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following formatted according to {@link
+ * CSVFormat#POSTGRESQL_CSV}.
+ *
+ * <pre>{@code
+ * "bank","purchaseAmount","transactionId"
+ * "A","10.23","12345"
+ * "B","54.65","54321"
+ * "C","11.76","98765"
+ * }</pre>
+ *
+ * <p>A {@link PCollection} of {@link Row}s works just like custom Java types illustrated above,
+ * except we use {@link CsvIO#writeRows()} as shown below for the same {@code Transaction} class. We
+ * derive {@code Transaction}'s {@link Schema} using {@link
+ * org.apache.beam.sdk.schemas.annotations.DefaultSchema.DefaultSchemaProvider}.
+ *
+ * <pre>{@code
+ * DefaultSchemaProvider defaultSchemaProvider = new DefaultSchemaProvider();
+ * Schema schema = defaultSchemaProvider.schemaFor(TypeDescriptor.of(Transaction.class));
+ * PCollection<Row> transactions = pipeline.apply(Create.of(
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "A")
+ * .withFieldValue("purchaseAmount", 10.23)
+ * .withFieldValue("transactionId", "12345")
+ * .build(),
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "B")
+ * .withFieldValue("purchaseAmount", 54.65)
+ * .withFieldValue("transactionId", "54321")
+ * .build(),
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "C")
+ * .withFieldValue("purchaseAmount", 11.76)
+ * .withFieldValue("transactionId", "98765")
+ * .build()
+ * );
+ *
+ * transactions.apply(
+ * CsvIO
+ * .writeRows()
+ * .to("gs://bucket/path/to/folder/prefix")
+ * );
+ * }</pre>
+ *
+ * <p>Writing the transactions {@link PCollection} of {@link Row}s would yield the following CSV
+ * file content.
+ *
+ * <pre>{@code
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ */
+public class CsvIO {
+ /**
+ * The valid {@link Schema.FieldType} from which {@link CsvIO} converts CSV records.
+ *
+ * <p>{@link FieldType#BYTE}
+ *
+ * <p>{@link FieldType#BOOLEAN}
+ *
+ * <p>{@link FieldType#DATETIME}
+ *
+ * <p>{@link FieldType#DECIMAL}
+ *
+ * <p>{@link FieldType#DOUBLE}
+ *
+ * <p>{@link FieldType#INT16}
+ *
+ * <p>{@link FieldType#INT32}
+ *
+ * <p>{@link FieldType#INT64}
+ *
+ * <p>{@link FieldType#FLOAT}
+ *
+ * <p>{@link FieldType#STRING}
+ */
Review Comment:
Swap to use `<ul><li>...</ul>` instead of `<p>` to get an unordered list.
##########
sdks/java/io/csv/src/main/java/org/apache/beam/sdk/io/csv/CsvIO.java:
##########
@@ -0,0 +1,554 @@
+/*
+ * 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.sdk.io.csv;
+
+import static java.nio.charset.StandardCharsets.UTF_8;
+import static org.apache.beam.sdk.values.TypeDescriptors.rows;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
+
+import com.google.auto.value.AutoValue;
+import edu.umd.cs.findbugs.annotations.NonNull;
+import java.io.BufferedWriter;
+import java.io.IOException;
+import java.io.OutputStreamWriter;
+import java.io.PrintWriter;
+import java.nio.channels.Channels;
+import java.nio.channels.WritableByteChannel;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+import java.util.List;
+import java.util.Optional;
+import java.util.Set;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.io.Compression;
+import org.apache.beam.sdk.io.FileIO;
+import org.apache.beam.sdk.io.FileIO.Write.FileNaming;
+import org.apache.beam.sdk.io.csv.CsvRowConversions.RowToCsv;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.Field;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.transforms.Contextful;
+import org.apache.beam.sdk.transforms.MapElements;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PCollectionView;
+import org.apache.beam.sdk.values.PDone;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableSet;
+import org.apache.commons.csv.CSVFormat;
+
+/**
+ * {@link PTransform}s for reading and writing CSV files.
+ *
+ * <h2>Reading CSV files</h2>
+ *
+ * <p>Reading from CSV files is not yet implemented. Please see <a
+ * href="https://github.com/apache/beam/issues/24552">https://github.com/apache/beam/issues/24552</a>.
+ *
+ * <h2>Writing CSV files</h2>
+ *
+ * <p>To write a {@link PCollection} to one or more CSV files, use {@link CsvIO.Write}, using {@code
+ * CsvIO.write().to(String)}. {@link CsvIO.Write} supports writing {@link Row} or custom Java types
+ * using an inferred {@link Schema}. See the Beam Programming Guide on <a
+ * href="https://beam.apache.org/documentation/programming-guide/#inferring-schemas">Inferring
+ * schemas</a> for more information on how to enable Beam to infer a {@link Schema} from a custom
+ * Java type.
+ *
+ * <p>For each {@link Schema.Field} included in the {@link Schema}, according to {@link
+ * CsvIO.Write#withSchemaFields(List)}, {@link CsvIO.Write} restricts the use of flat {@link
+ * Schema}s and checks whether the {@link Schema.FieldType}s are one of {@link
+ * #VALID_FIELD_TYPE_SET}.
+ *
+ * <h3>Example usage:</h3>
+ *
+ * <p>Suppose we have a <code>Transaction</code> class annotated with
+ * {@code @DefaultSchema(JavaBeanSchema.class)} so that Beam can infer its {@link Schema}:
+ *
+ * <pre>{@code @DefaultSchema(JavaBeanSchema.class)
+ * public class Transaction {
+ * public Transaction() { … }
+ * public Long getTransactionId();
+ * public void setTransactionId(Long transactionId) { … }
+ * public String getBank() { … }
+ * public void setBank(String bank) { … }
+ * public double getPurchaseAmount() { … }
+ * public void setPurchaseAmount(double purchaseAmount) { … }
+ * }
+ * }</pre>
+ *
+ * <p>From a {@code PCollection<Transaction>}, we can write one or many CSV files and {@link
+ * CsvIO.Write} will automatically create the header based on its inferred {@link Schema}
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions ...
+ * transactions.apply(CsvIO.<Transaction>write().to("gs://bucket/path/to/folder/prefix"));
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header is repeated for every
+ * file, whereas by default, {@link CsvIO.Write} will write all fields in <b>sorted order</b> of the
+ * field names.
+ *
+ * <pre>{@code
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ *
+ * <p>{@link CsvIO.Write} allows the use a subset of the available fields or control the order of
+ * their output using {@link CsvIO.Write#withSchemaFields(List)} listing of matching {@link
+ * Schema#getField(String)} names.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions ...
+ * transactions.apply(
+ * CsvIO
+ * .<Transaction>write()
+ * .to("gs://bucket/path/to/folder/prefix")
+ * .withSchemaFields(Arrays.asList("transactionId", "purchaseAmount"))
+ * );
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header is repeated for every
+ * file, but only including the subset of fields in their listed order.
+ *
+ * <pre>{@code
+ * transactionId,purchaseAmount
+ * 12345,10.23
+ * 54321,54.65
+ * 98765,11.76
+ * }</pre>
+ *
+ * <p>Additionally, we may override the default {@link CSVFormat#DEFAULT}.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions ...
+ * transactions.apply(
+ * CsvIO
+ * .<Transaction>write()
+ * .to("gs://bucket/path/to/folder/prefix")
+ * .withCSVFormat(CSVFormat.POSTGRESQL_CSV)
+ * );
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following formatted according to {@link
+ * CSVFormat#POSTGRESQL_CSV}.
+ *
+ * <pre>{@code
+ * "bank","purchaseAmount","transactionId"
+ * "A","10.23","12345"
+ * "B","54.65","54321"
+ * "C","11.76","98765"
+ * }</pre>
+ *
+ * <p>A {@link PCollection} of {@link Row}s works just like custom Java types illustrated above,
+ * except we use {@link CsvIO#writeRows()} as shown below for the same {@code Transaction} class. We
+ * derive {@code Transaction}'s {@link Schema} using {@link
+ * org.apache.beam.sdk.schemas.annotations.DefaultSchema.DefaultSchemaProvider}.
+ *
+ * <pre>{@code
+ * DefaultSchemaProvider defaultSchemaProvider = new DefaultSchemaProvider();
+ * Schema schema = defaultSchemaProvider.schemaFor(TypeDescriptor.of(Transaction.class));
+ * PCollection<Row> transactions = pipeline.apply(Create.of(
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "A")
+ * .withFieldValue("purchaseAmount", 10.23)
+ * .withFieldValue("transactionId", "12345")
+ * .build(),
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "B")
+ * .withFieldValue("purchaseAmount", 54.65)
+ * .withFieldValue("transactionId", "54321")
+ * .build(),
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "C")
+ * .withFieldValue("purchaseAmount", 11.76)
+ * .withFieldValue("transactionId", "98765")
+ * .build()
+ * );
+ *
+ * transactions.apply(
+ * CsvIO
+ * .writeRows()
+ * .to("gs://bucket/path/to/folder/prefix")
+ * );
+ * }</pre>
+ *
+ * <p>Writing the transactions {@link PCollection} of {@link Row}s would yield the following CSV
+ * file content.
+ *
+ * <pre>{@code
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ */
+public class CsvIO {
+ /**
+ * The valid {@link Schema.FieldType} from which {@link CsvIO} converts CSV records.
+ *
+ * <p>{@link FieldType#BYTE}
+ *
+ * <p>{@link FieldType#BOOLEAN}
+ *
+ * <p>{@link FieldType#DATETIME}
+ *
+ * <p>{@link FieldType#DECIMAL}
+ *
+ * <p>{@link FieldType#DOUBLE}
+ *
+ * <p>{@link FieldType#INT16}
+ *
+ * <p>{@link FieldType#INT32}
+ *
+ * <p>{@link FieldType#INT64}
+ *
+ * <p>{@link FieldType#FLOAT}
+ *
+ * <p>{@link FieldType#STRING}
+ */
+ public static final Set<Schema.FieldType> VALID_FIELD_TYPE_SET =
+ ImmutableSet.of(
+ FieldType.BYTE,
+ FieldType.BOOLEAN,
+ FieldType.DATETIME,
+ FieldType.DECIMAL,
+ FieldType.DOUBLE,
+ FieldType.INT16,
+ FieldType.INT32,
+ FieldType.INT64,
+ FieldType.FLOAT,
+ FieldType.STRING);
+
+ static final String DEFAULT_FILENAME_SUFFIX = ".csv";
+
+ /** Instantiates a {@link Write} for writing user types in {@link CSVFormat#DEFAULT} format. */
+ public static <T> Write<T> write() {
+ return defaultWrite();
+ }
+
+ /** Instantiates a {@link Write} for {@link Row}s in {@link CSVFormat#DEFAULT} format. */
+ public static Write<Row> writeRows() {
+ return defaultWrite();
+ }
+
+ private static <T> Write<T> defaultWrite() {
+ return new AutoValue_CsvIO_Write.Builder<T>()
+ .setCSVFormat(CSVFormat.DEFAULT)
+ .setFileWrite(FileIO.write())
+ .build();
+ }
+
+ /** Instantiates a {@link Sink.Builder} for writing {@link Row} elements to CSV file sinks. */
+ static Sink.Builder<Row> sinkBuilder() {
+ return new AutoValue_CsvIO_Sink.Builder<>();
+ }
+
+ /** Implementation of {@link FileIO.Sink}. */
+ @AutoValue
+ abstract static class Sink<T> implements FileIO.Sink<T> {
+
+ private transient @Nullable PrintWriter writer;
+
+ /**
+ * Opens a {@link WritableByteChannel} for writing CSV files. Writes the {@link #getPreamble()}
+ * if available followed by the {@link #getHeader()}.
+ */
+ @Override
+ public void open(@NonNull WritableByteChannel channel) throws IOException {
+ writer =
+ new PrintWriter(
+ new BufferedWriter(new OutputStreamWriter(Channels.newOutputStream(channel), UTF_8)));
+
+ if (getPreamble() != null) {
+ getWriter().println(getPreamble());
+ }
+ getWriter().println(getHeader());
+ }
+
+ /** Serializes and writes the element to a file. */
+ @Override
+ public void write(T element) throws IOException {
+ String line = getFormatFunction().apply(element);
+ getWriter().println(line);
+ }
+
+ @Override
+ public void flush() throws IOException {
+ getWriter().flush();
+ }
+
+ private PrintWriter getWriter() {
+ // resolves [dereference.of.nullable] error
+ Optional<PrintWriter> printWriter = Optional.ofNullable(writer);
+ checkState(printWriter.isPresent());
+ return printWriter.get();
+ }
+
+ /**
+ * Not to be confused with the CSV header, it is content written to the top of every sharded
+ * file prior to the header. See {@link Write#withPreamble(String)}.
+ */
+ @Nullable
+ abstract String getPreamble();
+
+ /**
+ * The column names of the CSV file written at the top line of each shard after the preamble, if
Review Comment:
why would this not be available?
##########
sdks/java/io/csv/src/main/java/org/apache/beam/sdk/io/csv/CsvIO.java:
##########
@@ -0,0 +1,554 @@
+/*
+ * 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.sdk.io.csv;
+
+import static java.nio.charset.StandardCharsets.UTF_8;
+import static org.apache.beam.sdk.values.TypeDescriptors.rows;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
+
+import com.google.auto.value.AutoValue;
+import edu.umd.cs.findbugs.annotations.NonNull;
+import java.io.BufferedWriter;
+import java.io.IOException;
+import java.io.OutputStreamWriter;
+import java.io.PrintWriter;
+import java.nio.channels.Channels;
+import java.nio.channels.WritableByteChannel;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+import java.util.List;
+import java.util.Optional;
+import java.util.Set;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.io.Compression;
+import org.apache.beam.sdk.io.FileIO;
+import org.apache.beam.sdk.io.FileIO.Write.FileNaming;
+import org.apache.beam.sdk.io.csv.CsvRowConversions.RowToCsv;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.Field;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.transforms.Contextful;
+import org.apache.beam.sdk.transforms.MapElements;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PCollectionView;
+import org.apache.beam.sdk.values.PDone;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableSet;
+import org.apache.commons.csv.CSVFormat;
+
+/**
+ * {@link PTransform}s for reading and writing CSV files.
+ *
+ * <h2>Reading CSV files</h2>
+ *
+ * <p>Reading from CSV files is not yet implemented. Please see <a
+ * href="https://github.com/apache/beam/issues/24552">https://github.com/apache/beam/issues/24552</a>.
+ *
+ * <h2>Writing CSV files</h2>
+ *
+ * <p>To write a {@link PCollection} to one or more CSV files, use {@link CsvIO.Write}, using {@code
+ * CsvIO.write().to(String)}. {@link CsvIO.Write} supports writing {@link Row} or custom Java types
+ * using an inferred {@link Schema}. See the Beam Programming Guide on <a
+ * href="https://beam.apache.org/documentation/programming-guide/#inferring-schemas">Inferring
+ * schemas</a> for more information on how to enable Beam to infer a {@link Schema} from a custom
+ * Java type.
+ *
+ * <p>For each {@link Schema.Field} included in the {@link Schema}, according to {@link
+ * CsvIO.Write#withSchemaFields(List)}, {@link CsvIO.Write} restricts the use of flat {@link
+ * Schema}s and checks whether the {@link Schema.FieldType}s are one of {@link
+ * #VALID_FIELD_TYPE_SET}.
+ *
+ * <h3>Example usage:</h3>
+ *
+ * <p>Suppose we have a <code>Transaction</code> class annotated with
+ * {@code @DefaultSchema(JavaBeanSchema.class)} so that Beam can infer its {@link Schema}:
+ *
+ * <pre>{@code @DefaultSchema(JavaBeanSchema.class)
+ * public class Transaction {
+ * public Transaction() { … }
+ * public Long getTransactionId();
+ * public void setTransactionId(Long transactionId) { … }
+ * public String getBank() { … }
+ * public void setBank(String bank) { … }
+ * public double getPurchaseAmount() { … }
+ * public void setPurchaseAmount(double purchaseAmount) { … }
+ * }
+ * }</pre>
+ *
+ * <p>From a {@code PCollection<Transaction>}, we can write one or many CSV files and {@link
+ * CsvIO.Write} will automatically create the header based on its inferred {@link Schema}
Review Comment:
```suggestion
* <p>From a {@code PCollection<Transaction>}, {@link CsvIO.Write} can write one or many CSV files
* automatically creating the header based on its inferred {@link Schema}.
```
##########
sdks/java/io/csv/src/main/java/org/apache/beam/sdk/io/csv/CsvIO.java:
##########
@@ -0,0 +1,554 @@
+/*
+ * 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.sdk.io.csv;
+
+import static java.nio.charset.StandardCharsets.UTF_8;
+import static org.apache.beam.sdk.values.TypeDescriptors.rows;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
+
+import com.google.auto.value.AutoValue;
+import edu.umd.cs.findbugs.annotations.NonNull;
+import java.io.BufferedWriter;
+import java.io.IOException;
+import java.io.OutputStreamWriter;
+import java.io.PrintWriter;
+import java.nio.channels.Channels;
+import java.nio.channels.WritableByteChannel;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+import java.util.List;
+import java.util.Optional;
+import java.util.Set;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.io.Compression;
+import org.apache.beam.sdk.io.FileIO;
+import org.apache.beam.sdk.io.FileIO.Write.FileNaming;
+import org.apache.beam.sdk.io.csv.CsvRowConversions.RowToCsv;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.Field;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.transforms.Contextful;
+import org.apache.beam.sdk.transforms.MapElements;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PCollectionView;
+import org.apache.beam.sdk.values.PDone;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableSet;
+import org.apache.commons.csv.CSVFormat;
+
+/**
+ * {@link PTransform}s for reading and writing CSV files.
+ *
+ * <h2>Reading CSV files</h2>
+ *
+ * <p>Reading from CSV files is not yet implemented. Please see <a
+ * href="https://github.com/apache/beam/issues/24552">https://github.com/apache/beam/issues/24552</a>.
+ *
+ * <h2>Writing CSV files</h2>
+ *
+ * <p>To write a {@link PCollection} to one or more CSV files, use {@link CsvIO.Write}, using {@code
+ * CsvIO.write().to(String)}. {@link CsvIO.Write} supports writing {@link Row} or custom Java types
+ * using an inferred {@link Schema}. See the Beam Programming Guide on <a
+ * href="https://beam.apache.org/documentation/programming-guide/#inferring-schemas">Inferring
+ * schemas</a> for more information on how to enable Beam to infer a {@link Schema} from a custom
+ * Java type.
+ *
+ * <p>For each {@link Schema.Field} included in the {@link Schema}, according to {@link
+ * CsvIO.Write#withSchemaFields(List)}, {@link CsvIO.Write} restricts the use of flat {@link
+ * Schema}s and checks whether the {@link Schema.FieldType}s are one of {@link
+ * #VALID_FIELD_TYPE_SET}.
+ *
+ * <h3>Example usage:</h3>
+ *
+ * <p>Suppose we have a <code>Transaction</code> class annotated with
+ * {@code @DefaultSchema(JavaBeanSchema.class)} so that Beam can infer its {@link Schema}:
+ *
+ * <pre>{@code @DefaultSchema(JavaBeanSchema.class)
+ * public class Transaction {
+ * public Transaction() { … }
+ * public Long getTransactionId();
+ * public void setTransactionId(Long transactionId) { … }
+ * public String getBank() { … }
+ * public void setBank(String bank) { … }
+ * public double getPurchaseAmount() { … }
+ * public void setPurchaseAmount(double purchaseAmount) { … }
+ * }
+ * }</pre>
+ *
+ * <p>From a {@code PCollection<Transaction>}, we can write one or many CSV files and {@link
+ * CsvIO.Write} will automatically create the header based on its inferred {@link Schema}
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions ...
+ * transactions.apply(CsvIO.<Transaction>write().to("gs://bucket/path/to/folder/prefix"));
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header is repeated for every
+ * file, whereas by default, {@link CsvIO.Write} will write all fields in <b>sorted order</b> of the
+ * field names.
+ *
+ * <pre>{@code
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ *
+ * <p>{@link CsvIO.Write} allows the use a subset of the available fields or control the order of
+ * their output using {@link CsvIO.Write#withSchemaFields(List)} listing of matching {@link
+ * Schema#getField(String)} names.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions ...
+ * transactions.apply(
+ * CsvIO
+ * .<Transaction>write()
+ * .to("gs://bucket/path/to/folder/prefix")
+ * .withSchemaFields(Arrays.asList("transactionId", "purchaseAmount"))
+ * );
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header is repeated for every
+ * file, but only including the subset of fields in their listed order.
+ *
+ * <pre>{@code
+ * transactionId,purchaseAmount
+ * 12345,10.23
+ * 54321,54.65
+ * 98765,11.76
+ * }</pre>
+ *
+ * <p>Additionally, we may override the default {@link CSVFormat#DEFAULT}.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions ...
+ * transactions.apply(
+ * CsvIO
+ * .<Transaction>write()
+ * .to("gs://bucket/path/to/folder/prefix")
+ * .withCSVFormat(CSVFormat.POSTGRESQL_CSV)
+ * );
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following formatted according to {@link
+ * CSVFormat#POSTGRESQL_CSV}.
+ *
+ * <pre>{@code
+ * "bank","purchaseAmount","transactionId"
+ * "A","10.23","12345"
+ * "B","54.65","54321"
+ * "C","11.76","98765"
+ * }</pre>
+ *
+ * <p>A {@link PCollection} of {@link Row}s works just like custom Java types illustrated above,
+ * except we use {@link CsvIO#writeRows()} as shown below for the same {@code Transaction} class. We
+ * derive {@code Transaction}'s {@link Schema} using {@link
Review Comment:
```suggestion
* derive {@code Transaction}'s {@link Schema} using a {@link
```
##########
sdks/java/io/csv/src/main/java/org/apache/beam/sdk/io/csv/CsvIO.java:
##########
@@ -0,0 +1,554 @@
+/*
+ * 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.sdk.io.csv;
+
+import static java.nio.charset.StandardCharsets.UTF_8;
+import static org.apache.beam.sdk.values.TypeDescriptors.rows;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
+
+import com.google.auto.value.AutoValue;
+import edu.umd.cs.findbugs.annotations.NonNull;
+import java.io.BufferedWriter;
+import java.io.IOException;
+import java.io.OutputStreamWriter;
+import java.io.PrintWriter;
+import java.nio.channels.Channels;
+import java.nio.channels.WritableByteChannel;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+import java.util.List;
+import java.util.Optional;
+import java.util.Set;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.io.Compression;
+import org.apache.beam.sdk.io.FileIO;
+import org.apache.beam.sdk.io.FileIO.Write.FileNaming;
+import org.apache.beam.sdk.io.csv.CsvRowConversions.RowToCsv;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.Field;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.transforms.Contextful;
+import org.apache.beam.sdk.transforms.MapElements;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PCollectionView;
+import org.apache.beam.sdk.values.PDone;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableSet;
+import org.apache.commons.csv.CSVFormat;
+
+/**
+ * {@link PTransform}s for reading and writing CSV files.
+ *
+ * <h2>Reading CSV files</h2>
+ *
+ * <p>Reading from CSV files is not yet implemented. Please see <a
+ * href="https://github.com/apache/beam/issues/24552">https://github.com/apache/beam/issues/24552</a>.
+ *
+ * <h2>Writing CSV files</h2>
+ *
+ * <p>To write a {@link PCollection} to one or more CSV files, use {@link CsvIO.Write}, using {@code
+ * CsvIO.write().to(String)}. {@link CsvIO.Write} supports writing {@link Row} or custom Java types
+ * using an inferred {@link Schema}. See the Beam Programming Guide on <a
+ * href="https://beam.apache.org/documentation/programming-guide/#inferring-schemas">Inferring
+ * schemas</a> for more information on how to enable Beam to infer a {@link Schema} from a custom
+ * Java type.
+ *
+ * <p>For each {@link Schema.Field} included in the {@link Schema}, according to {@link
+ * CsvIO.Write#withSchemaFields(List)}, {@link CsvIO.Write} restricts the use of flat {@link
+ * Schema}s and checks whether the {@link Schema.FieldType}s are one of {@link
+ * #VALID_FIELD_TYPE_SET}.
+ *
+ * <h3>Example usage:</h3>
+ *
+ * <p>Suppose we have a <code>Transaction</code> class annotated with
+ * {@code @DefaultSchema(JavaBeanSchema.class)} so that Beam can infer its {@link Schema}:
+ *
+ * <pre>{@code @DefaultSchema(JavaBeanSchema.class)
+ * public class Transaction {
+ * public Transaction() { … }
+ * public Long getTransactionId();
+ * public void setTransactionId(Long transactionId) { … }
+ * public String getBank() { … }
+ * public void setBank(String bank) { … }
+ * public double getPurchaseAmount() { … }
+ * public void setPurchaseAmount(double purchaseAmount) { … }
+ * }
+ * }</pre>
+ *
+ * <p>From a {@code PCollection<Transaction>}, we can write one or many CSV files and {@link
+ * CsvIO.Write} will automatically create the header based on its inferred {@link Schema}
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions ...
+ * transactions.apply(CsvIO.<Transaction>write().to("gs://bucket/path/to/folder/prefix"));
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header is repeated for every
+ * file, whereas by default, {@link CsvIO.Write} will write all fields in <b>sorted order</b> of the
+ * field names.
+ *
+ * <pre>{@code
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ *
+ * <p>{@link CsvIO.Write} allows the use a subset of the available fields or control the order of
+ * their output using {@link CsvIO.Write#withSchemaFields(List)} listing of matching {@link
+ * Schema#getField(String)} names.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions ...
+ * transactions.apply(
+ * CsvIO
+ * .<Transaction>write()
+ * .to("gs://bucket/path/to/folder/prefix")
+ * .withSchemaFields(Arrays.asList("transactionId", "purchaseAmount"))
+ * );
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header is repeated for every
+ * file, but only including the subset of fields in their listed order.
+ *
+ * <pre>{@code
+ * transactionId,purchaseAmount
+ * 12345,10.23
+ * 54321,54.65
+ * 98765,11.76
+ * }</pre>
+ *
+ * <p>Additionally, we may override the default {@link CSVFormat#DEFAULT}.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions ...
+ * transactions.apply(
+ * CsvIO
+ * .<Transaction>write()
+ * .to("gs://bucket/path/to/folder/prefix")
+ * .withCSVFormat(CSVFormat.POSTGRESQL_CSV)
+ * );
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following formatted according to {@link
+ * CSVFormat#POSTGRESQL_CSV}.
+ *
+ * <pre>{@code
+ * "bank","purchaseAmount","transactionId"
+ * "A","10.23","12345"
+ * "B","54.65","54321"
+ * "C","11.76","98765"
+ * }</pre>
+ *
+ * <p>A {@link PCollection} of {@link Row}s works just like custom Java types illustrated above,
+ * except we use {@link CsvIO#writeRows()} as shown below for the same {@code Transaction} class. We
+ * derive {@code Transaction}'s {@link Schema} using {@link
+ * org.apache.beam.sdk.schemas.annotations.DefaultSchema.DefaultSchemaProvider}.
+ *
+ * <pre>{@code
+ * DefaultSchemaProvider defaultSchemaProvider = new DefaultSchemaProvider();
+ * Schema schema = defaultSchemaProvider.schemaFor(TypeDescriptor.of(Transaction.class));
+ * PCollection<Row> transactions = pipeline.apply(Create.of(
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "A")
+ * .withFieldValue("purchaseAmount", 10.23)
+ * .withFieldValue("transactionId", "12345")
+ * .build(),
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "B")
+ * .withFieldValue("purchaseAmount", 54.65)
+ * .withFieldValue("transactionId", "54321")
+ * .build(),
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "C")
+ * .withFieldValue("purchaseAmount", 11.76)
+ * .withFieldValue("transactionId", "98765")
+ * .build()
+ * );
+ *
+ * transactions.apply(
+ * CsvIO
+ * .writeRows()
+ * .to("gs://bucket/path/to/folder/prefix")
+ * );
+ * }</pre>
+ *
+ * <p>Writing the transactions {@link PCollection} of {@link Row}s would yield the following CSV
+ * file content.
+ *
+ * <pre>{@code
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ */
+public class CsvIO {
+ /**
+ * The valid {@link Schema.FieldType} from which {@link CsvIO} converts CSV records.
+ *
+ * <p>{@link FieldType#BYTE}
+ *
+ * <p>{@link FieldType#BOOLEAN}
+ *
+ * <p>{@link FieldType#DATETIME}
+ *
+ * <p>{@link FieldType#DECIMAL}
+ *
+ * <p>{@link FieldType#DOUBLE}
+ *
+ * <p>{@link FieldType#INT16}
+ *
+ * <p>{@link FieldType#INT32}
+ *
+ * <p>{@link FieldType#INT64}
+ *
+ * <p>{@link FieldType#FLOAT}
+ *
+ * <p>{@link FieldType#STRING}
+ */
+ public static final Set<Schema.FieldType> VALID_FIELD_TYPE_SET =
+ ImmutableSet.of(
+ FieldType.BYTE,
+ FieldType.BOOLEAN,
+ FieldType.DATETIME,
+ FieldType.DECIMAL,
+ FieldType.DOUBLE,
+ FieldType.INT16,
+ FieldType.INT32,
+ FieldType.INT64,
+ FieldType.FLOAT,
+ FieldType.STRING);
+
+ static final String DEFAULT_FILENAME_SUFFIX = ".csv";
+
+ /** Instantiates a {@link Write} for writing user types in {@link CSVFormat#DEFAULT} format. */
+ public static <T> Write<T> write() {
+ return defaultWrite();
+ }
+
+ /** Instantiates a {@link Write} for {@link Row}s in {@link CSVFormat#DEFAULT} format. */
+ public static Write<Row> writeRows() {
+ return defaultWrite();
+ }
+
+ private static <T> Write<T> defaultWrite() {
+ return new AutoValue_CsvIO_Write.Builder<T>()
+ .setCSVFormat(CSVFormat.DEFAULT)
+ .setFileWrite(FileIO.write())
+ .build();
+ }
+
+ /** Instantiates a {@link Sink.Builder} for writing {@link Row} elements to CSV file sinks. */
+ static Sink.Builder<Row> sinkBuilder() {
+ return new AutoValue_CsvIO_Sink.Builder<>();
+ }
+
+ /** Implementation of {@link FileIO.Sink}. */
+ @AutoValue
+ abstract static class Sink<T> implements FileIO.Sink<T> {
+
+ private transient @Nullable PrintWriter writer;
+
+ /**
+ * Opens a {@link WritableByteChannel} for writing CSV files. Writes the {@link #getPreamble()}
+ * if available followed by the {@link #getHeader()}.
+ */
+ @Override
+ public void open(@NonNull WritableByteChannel channel) throws IOException {
+ writer =
+ new PrintWriter(
+ new BufferedWriter(new OutputStreamWriter(Channels.newOutputStream(channel), UTF_8)));
+
+ if (getPreamble() != null) {
+ getWriter().println(getPreamble());
+ }
+ getWriter().println(getHeader());
+ }
+
+ /** Serializes and writes the element to a file. */
+ @Override
+ public void write(T element) throws IOException {
+ String line = getFormatFunction().apply(element);
+ getWriter().println(line);
+ }
+
+ @Override
+ public void flush() throws IOException {
+ getWriter().flush();
+ }
+
+ private PrintWriter getWriter() {
+ // resolves [dereference.of.nullable] error
+ Optional<PrintWriter> printWriter = Optional.ofNullable(writer);
+ checkState(printWriter.isPresent());
+ return printWriter.get();
Review Comment:
```suggestion
if (writer == null) {
throw new IllegalStateException("writer unexpected to be null");
}
return writer;
```
##########
sdks/java/io/csv/src/main/java/org/apache/beam/sdk/io/csv/CsvIO.java:
##########
@@ -0,0 +1,553 @@
+/*
+ * 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.sdk.io.csv;
+
+import static java.nio.charset.StandardCharsets.UTF_8;
+import static org.apache.beam.sdk.values.TypeDescriptors.rows;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
+
+import com.google.auto.value.AutoValue;
+import java.io.BufferedWriter;
+import java.io.IOException;
+import java.io.OutputStreamWriter;
+import java.io.PrintWriter;
+import java.nio.channels.Channels;
+import java.nio.channels.WritableByteChannel;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+import java.util.List;
+import java.util.Optional;
+import java.util.Set;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.io.Compression;
+import org.apache.beam.sdk.io.FileIO;
+import org.apache.beam.sdk.io.FileIO.Write.FileNaming;
+import org.apache.beam.sdk.io.csv.CsvRowConversions.RowToCsv;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.Field;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.transforms.Contextful;
+import org.apache.beam.sdk.transforms.MapElements;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PCollectionView;
+import org.apache.beam.sdk.values.PDone;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableSet;
+import org.apache.commons.csv.CSVFormat;
+
+/**
+ * {@link PTransform}s for reading and writing CSV files.
+ *
+ * <h2>Reading CSV files</h2>
+ *
+ * <p>Reading from CSV files is not yet implemented. Please see <a
+ * href="https://github.com/apache/beam/issues/24552">https://github.com/apache/beam/issues/24552</a>.
+ *
+ * <h2>Writing CSV files</h2>
+ *
+ * <p>To write a {@link PCollection} to one or more CSV files, use {@link CsvIO.Write}, using {@code
+ * CsvIO.write().to(String)}. {@link CsvIO.Write} supports writing {@link Row} or custom Java types
+ * using an inferred {@link Schema}. See the Beam Programming Guide on <a
+ * href="https://beam.apache.org/documentation/programming-guide/#inferring-schemas">Inferring
+ * schemas</a> for more information on how to enable Beam to infer a {@link Schema} from a custom
+ * Java type.
+ *
+ * <p>For each {@link Schema.Field} included in the {@link Schema}, according to {@link
+ * CsvIO.Write#withSchemaFields(List)}, {@link CsvIO.Write} restricts the use of flat {@link
+ * Schema}s and checks whether the {@link Schema.FieldType}s are one of {@link
+ * #VALID_FIELD_TYPE_SET}.
+ *
+ * <h3>Example usage:</h3>
+ *
+ * <p>Suppose we have a <code>Transaction</code> class annotated with
+ * {@code @DefaultSchema(JavaBeanSchema.class)} so that Beam can infer its {@link Schema}:
+ *
+ * <pre>{@code @DefaultSchema(JavaBeanSchema.class)
+ * public class Transaction {
+ * public Transaction() { … }
+ * public Long getTransactionId();
+ * public void setTransactionId(Long transactionId) { … }
+ * public String getBank() { … }
+ * public void setBank(String bank) { … }
+ * public double getPurchaseAmount() { … }
+ * public void setPurchaseAmount(double purchaseAmount) { … }
+ * }
+ * }</pre>
+ *
+ * <p>From a {@code PCollection<Transaction>}, we can write one or many CSV files and {@link
+ * CsvIO.Write} will automatically create the header based on its inferred {@link Schema}
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions ...
+ * transactions.apply(CsvIO.<Transaction>write().to("gs://bucket/path/to/folder/prefix"));
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header is repeated for every
+ * file, whereas by default, {@link CsvIO.Write} will write all fields in <b>sorted order</b> of the
+ * field names.
+ *
+ * <pre>{@code
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ *
+ * <p>{@link CsvIO.Write} allows the use a subset of the available fields or control the order of
+ * their output using {@link CsvIO.Write#withSchemaFields(List)} listing of matching {@link
+ * Schema#getField(String)} names.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions ...
+ * transactions.apply(
+ * CsvIO
+ * .<Transaction>write()
+ * .to("gs://bucket/path/to/folder/prefix")
+ * .withSchemaFields(Arrays.asList("transactionId", "purchaseAmount"))
+ * );
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following where the header is repeated for every
+ * file, but only including the subset of fields in their listed order.
+ *
+ * <pre>{@code
+ * transactionId,purchaseAmount
+ * 12345,10.23
+ * 54321,54.65
+ * 98765,11.76
+ * }</pre>
+ *
+ * <p>Additionally, we may override the default {@link CSVFormat#DEFAULT}.
+ *
+ * <pre>{@code
+ * PCollection<Transaction> transactions ...
+ * transactions.apply(
+ * CsvIO
+ * .<Transaction>write()
+ * .to("gs://bucket/path/to/folder/prefix")
+ * .withCSVFormat(CSVFormat.POSTGRESQL_CSV)
+ * );
+ * }</pre>
+ *
+ * <p>The resulting CSV files will look like the following formatted according to {@link
+ * CSVFormat#POSTGRESQL_CSV}.
+ *
+ * <pre>{@code
+ * "bank","purchaseAmount","transactionId"
+ * "A","10.23","12345"
+ * "B","54.65","54321"
+ * "C","11.76","98765"
+ * }</pre>
+ *
+ * <p>A {@link PCollection} of {@link Row}s works just like custom Java types illustrated above,
+ * except we use {@link CsvIO#writeRows()} as shown below for the same {@code Transaction} class. We
+ * derive {@code Transaction}'s {@link Schema} using {@link
+ * org.apache.beam.sdk.schemas.annotations.DefaultSchema.DefaultSchemaProvider}.
+ *
+ * <pre>{@code
+ * DefaultSchemaProvider defaultSchemaProvider = new DefaultSchemaProvider();
+ * Schema schema = defaultSchemaProvider.schemaFor(TypeDescriptor.of(Transaction.class));
+ * PCollection<Row> transactions = pipeline.apply(Create.of(
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "A")
+ * .withFieldValue("purchaseAmount", 10.23)
+ * .withFieldValue("transactionId", "12345")
+ * .build(),
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "B")
+ * .withFieldValue("purchaseAmount", 54.65)
+ * .withFieldValue("transactionId", "54321")
+ * .build(),
+ * Row
+ * .withSchema(schema)
+ * .withFieldValue("bank", "C")
+ * .withFieldValue("purchaseAmount", 11.76)
+ * .withFieldValue("transactionId", "98765")
+ * .build()
+ * );
+ *
+ * transactions.apply(
+ * CsvIO
+ * .writeRows()
+ * .to("gs://bucket/path/to/folder/prefix")
+ * );
+ * }</pre>
+ *
+ * <p>Writing the transactions {@link PCollection} of {@link Row}s would yield the following CSV
+ * file content.
+ *
+ * <pre>{@code
+ * bank,purchaseAmount,transactionId
+ * A,10.23,12345
+ * B,54.65,54321
+ * C,11.76,98765
+ * }</pre>
+ */
+public class CsvIO {
+ /**
+ * The valid {@link Schema.FieldType} from which {@link CsvIO} converts CSV records.
+ *
+ * <p>{@link FieldType#BYTE}
+ *
+ * <p>{@link FieldType#BOOLEAN}
+ *
+ * <p>{@link FieldType#DATETIME}
+ *
+ * <p>{@link FieldType#DECIMAL}
+ *
+ * <p>{@link FieldType#DOUBLE}
+ *
+ * <p>{@link FieldType#INT16}
+ *
+ * <p>{@link FieldType#INT32}
+ *
+ * <p>{@link FieldType#INT64}
+ *
+ * <p>{@link FieldType#FLOAT}
+ *
+ * <p>{@link FieldType#STRING}
+ */
+ public static final Set<Schema.FieldType> VALID_FIELD_TYPE_SET =
+ ImmutableSet.of(
+ FieldType.BYTE,
+ FieldType.BOOLEAN,
+ FieldType.DATETIME,
+ FieldType.DECIMAL,
+ FieldType.DOUBLE,
+ FieldType.INT16,
+ FieldType.INT32,
+ FieldType.INT64,
+ FieldType.FLOAT,
+ FieldType.STRING);
+
+ static final String DEFAULT_FILENAME_SUFFIX = ".csv";
+
+ /** Instantiates a {@link Write} for writing user types in {@link CSVFormat#DEFAULT} format. */
+ public static <T> Write<T> write() {
+ return defaultWrite();
+ }
+
+ /** Instantiates a {@link Write} for {@link Row}s in {@link CSVFormat#DEFAULT} format. */
+ public static Write<Row> writeRows() {
+ return defaultWrite();
+ }
+
+ private static <T> Write<T> defaultWrite() {
+ return new AutoValue_CsvIO_Write.Builder<T>()
+ .setCSVFormat(CSVFormat.DEFAULT)
+ .setFileWrite(FileIO.write())
+ .build();
+ }
+
+ /** Instantiates a {@link Sink.Builder} for writing {@link Row} elements to CSV file sinks. */
+ static Sink.Builder<Row> sinkBuilder() {
+ return new AutoValue_CsvIO_Sink.Builder<>();
+ }
+
+ /** Implementation of {@link FileIO.Sink}. */
+ @AutoValue
+ abstract static class Sink<T> implements FileIO.Sink<T> {
+
+ private transient @Nullable PrintWriter writer;
+
+ /**
+ * Opens a {@link WritableByteChannel} for writing CSV files. Writes the {@link #getPreamble()}
+ * if available followed by the {@link #getHeader()}.
+ */
+ @Override
+ public void open(WritableByteChannel channel) throws IOException {
+ writer =
+ new PrintWriter(
+ new BufferedWriter(new OutputStreamWriter(Channels.newOutputStream(channel), UTF_8)));
+
+ if (getPreamble() != null) {
+ getWriter().println(getPreamble());
+ }
+ getWriter().println(getHeader());
+ }
+
+ /** Serializes and writes the element to a file. */
+ @Override
+ public void write(T element) throws IOException {
+ String line = getFormatFunction().apply(element);
+ getWriter().println(line);
+ }
+
+ @Override
+ public void flush() throws IOException {
+ getWriter().flush();
+ }
+
+ private PrintWriter getWriter() {
+ // resolves [dereference.of.nullable] error
+ Optional<PrintWriter> printWriter = Optional.ofNullable(writer);
+ checkState(printWriter.isPresent());
+ return printWriter.get();
+ }
+
+ /**
+ * Not to be confused with the CSV header, it is content written to the top of every sharded
+ * file prior to the header. See {@link Write#withPreamble(String)}.
+ */
+ @Nullable
+ abstract String getPreamble();
+
+ /**
+ * The column names of the CSV file written at the top line of each shard after the preamble, if
+ * available.
+ */
+ abstract String getHeader();
+
+ /** A {@link SerializableFunction} for converting to a CSV formatted string. */
+ abstract SerializableFunction<T, String> getFormatFunction();
+
+ @AutoValue.Builder
+ abstract static class Builder<T> {
+
+ /**
+ * Not to be confused with the CSV header, it is content written to the top of every sharded
+ * file prior to the header. See {@link Write#withPreamble(String)}.
+ */
+ abstract Builder<T> setPreamble(String value);
+
+ /**
+ * The column names of the CSV file written at the top line of each shard after the preamble,
+ * if available.
Review Comment:
ditto on why not available?
##########
sdks/java/io/csv/src/main/java/org/apache/beam/sdk/io/csv/CsvRowConversions.java:
##########
@@ -0,0 +1,166 @@
+/*
+ * 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.sdk.io.csv;
+
+import com.google.auto.value.AutoValue;
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.stream.Collectors;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.Field;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.values.Row;
+import org.apache.commons.csv.CSVFormat;
+
+/** Contains classes and methods to help with converting between {@link Row} and CSV strings. */
+class CsvRowConversions {
+
+ /** Converts between {@link Row} and CSV string using a {@link CSVFormat}. */
+ @AutoValue
+ abstract static class RowToCsv implements SerializableFunction<Row, String> {
+
+ static Builder builder() {
+ return new AutoValue_CsvRowConversions_RowToCsv.Builder();
+ }
+
+ /** The expected {@link Schema} of the {@link Row} input. */
+ abstract Schema getSchema();
+
+ /** The {@link CSVFormat} of the converted {@link Row} input. */
+ abstract CSVFormat getCSVFormat();
+
+ /** The order and subset of {@link Schema} field names to drive the {@link Row} conversion. */
+ abstract List<String> getSchemaFields();
+
+ /** Builds a header from {@link #getSchemaFields()} using the {@link #getCSVFormat()}. */
+ String buildHeader() {
+ return buildHeaderFrom(getSchemaFields(), getCSVFormat());
+ }
+
+ /**
+ * Converts a {@link Row} to a CSV string formatted using {@link #getCSVFormat()}. Null values
+ * are converted to empty strings.
+ */
+ @Override
+ public String apply(Row input) {
+ List<String> schemaFields = getSchemaFields();
+ Object[] values = new Object[schemaFields.size()];
+ for (int i = 0; i < schemaFields.size(); i++) {
+ String name = schemaFields.get(i);
+ values[i] = input.getValue(name);
+ }
+ return getCSVFormat().format(values);
+ }
+
+ @AutoValue.Builder
+ abstract static class Builder {
+
+ /** The expected {@link Schema} of the {@link Row} input. */
+ abstract Builder setSchema(Schema schema);
+
+ abstract Schema getSchema();
+
+ /** The {@link CSVFormat} of the converted {@link Row} input. */
+ abstract Builder setCSVFormat(CSVFormat format);
+
+ /** The order and subset of {@link Schema} field names to drive the {@link Row} conversion. */
+ abstract Builder setSchemaFields(List<String> fields);
+
+ abstract List<String> getSchemaFields();
+
+ abstract RowToCsv autoBuild();
+
+ final RowToCsv build() {
+
+ validateHeaderAgainstSchema(getSchemaFields(), getSchema());
+
+ return autoBuild();
+ }
+ }
+ }
+
+ /** Formats columns into a header String based on {@link CSVFormat}. */
Review Comment:
```suggestion
/** Formats columns into a header String based upon {@link CSVFormat}. */
```
##########
sdks/java/io/csv/src/main/java/org/apache/beam/sdk/io/csv/CsvRowConversions.java:
##########
@@ -0,0 +1,166 @@
+/*
+ * 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.sdk.io.csv;
+
+import com.google.auto.value.AutoValue;
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.stream.Collectors;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.Field;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.values.Row;
+import org.apache.commons.csv.CSVFormat;
+
+/** Contains classes and methods to help with converting between {@link Row} and CSV strings. */
+class CsvRowConversions {
+
+ /** Converts between {@link Row} and CSV string using a {@link CSVFormat}. */
+ @AutoValue
+ abstract static class RowToCsv implements SerializableFunction<Row, String> {
+
+ static Builder builder() {
+ return new AutoValue_CsvRowConversions_RowToCsv.Builder();
+ }
+
+ /** The expected {@link Schema} of the {@link Row} input. */
+ abstract Schema getSchema();
+
+ /** The {@link CSVFormat} of the converted {@link Row} input. */
+ abstract CSVFormat getCSVFormat();
+
+ /** The order and subset of {@link Schema} field names to drive the {@link Row} conversion. */
+ abstract List<String> getSchemaFields();
+
+ /** Builds a header from {@link #getSchemaFields()} using the {@link #getCSVFormat()}. */
+ String buildHeader() {
+ return buildHeaderFrom(getSchemaFields(), getCSVFormat());
+ }
+
+ /**
+ * Converts a {@link Row} to a CSV string formatted using {@link #getCSVFormat()}. Null values
+ * are converted to empty strings.
Review Comment:
Aren't null values converted based upon the CSV format (based upon the null string)?
```suggestion
* Converts a {@link Row} to a CSV string formatted using {@link #getCSVFormat()}.
```
##########
sdks/java/io/csv/src/test/java/org/apache/beam/sdk/io/csv/CsvIOWriteTest.java:
##########
@@ -0,0 +1,534 @@
+/*
+ * 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.sdk.io.csv;
+
+import static org.apache.beam.sdk.io.csv.CsvIOTestHelpers.ALL_DATA_TYPES_SCHEMA;
+import static org.apache.beam.sdk.io.csv.CsvIOTestHelpers.rowOf;
+import static org.junit.Assert.assertArrayEquals;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertThrows;
+import static org.junit.Assert.assertTrue;
+
+import java.io.BufferedReader;
+import java.io.File;
+import java.io.IOException;
+import java.math.BigDecimal;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.regex.Pattern;
+import java.util.stream.Collectors;
+import org.apache.beam.sdk.io.Compression;
+import org.apache.beam.sdk.io.csv.CsvIOTestHelpers.AllDataTypes;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.Field;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.testing.TestPipeline;
+import org.apache.beam.sdk.transforms.Create;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Charsets;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.commons.csv.CSVFormat;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Instant;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.TemporaryFolder;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/** Tests for {@link CsvIO.Write}. */
+@RunWith(JUnit4.class)
+public class CsvIOWriteTest {
+
+ @Rule public TemporaryFolder tmpFolder = TemporaryFolder.builder().build();
+
+ @Rule public TestPipeline pipeline = TestPipeline.create();
+
+ @Rule
+ public TestPipeline errorPipeline = TestPipeline.create().enableAbandonedNodeEnforcement(false);
+
+ private final List<AllDataTypes> userTypes =
+ Arrays.asList(
+ AllDataTypes.of(
+ true,
+ (byte) 0,
+ Instant.ofEpochMilli(0L).toDateTime(),
+ BigDecimal.valueOf(0L),
+ 0.1,
+ 0.2f,
+ (short) 3,
+ 4,
+ 5L,
+ "a"),
+ AllDataTypes.of(
+ true,
+ (byte) 1,
+ Instant.ofEpochMilli(1L).toDateTime(),
+ BigDecimal.valueOf(1L),
+ 1.1,
+ 1.2f,
+ (short) 4,
+ 5,
+ 6L,
+ "b"),
+ AllDataTypes.of(
+ false,
+ (byte) 2,
+ Instant.ofEpochMilli(2L).toDateTime(),
+ BigDecimal.valueOf(2L),
+ 2.0,
+ 2.1f,
+ (short) 5,
+ 6,
+ 7L,
+ "c"));
+
+ private final List<Row> rows =
+ Arrays.asList(
+ rowOf(
+ true,
+ (byte) 0,
+ Instant.ofEpochMilli(0L).toDateTime(),
+ BigDecimal.valueOf(0L),
+ 0.1,
+ 0.2f,
+ (short) 3,
+ 4,
+ 5L,
+ "a"),
+ rowOf(
+ true,
+ (byte) 1,
+ Instant.ofEpochMilli(1L).toDateTime(),
+ BigDecimal.valueOf(1L),
+ 1.1,
+ 1.2f,
+ (short) 4,
+ 5,
+ 6L,
+ "b"),
+ rowOf(
+ false,
+ (byte) 2,
+ Instant.ofEpochMilli(2L).toDateTime(),
+ BigDecimal.valueOf(2L),
+ 2.0,
+ 2.1f,
+ (short) 5,
+ 6,
+ 7L,
+ "c"));
+
+ @Test
+ public void validateSchema() {
+ String to = tmpFolder.getRoot().getAbsolutePath() + "/somewhere";
+ Schema emptySchema = Schema.of();
+ Schema bytesSchema =
+ Schema.of(Field.of("badfield", FieldType.BYTES), Field.of("ok", FieldType.STRING));
+ Schema mapSchema =
+ Schema.of(
+ Field.of("badfield", FieldType.map(FieldType.STRING, FieldType.STRING)),
+ Field.of("ok", FieldType.STRING));
+ Schema arraySchema =
+ Schema.of(
+ Field.of("badfield", FieldType.array(FieldType.STRING)),
+ Field.of("ok", FieldType.STRING));
+ Schema rowSchema =
+ Schema.of(
+ Field.of("badfield", FieldType.row(ALL_DATA_TYPES_SCHEMA)),
+ Field.of("ok", FieldType.STRING));
+
+ assertThrows(
+ "schema should not be empty",
+ IllegalArgumentException.class,
+ () ->
+ errorPipeline
+ .apply(Create.of(Row.nullRow(emptySchema)).withRowSchema(emptySchema))
+ .apply(CsvIO.writeRows().to(to)));
+
+ assertThrows(
+ "schema should not contain bytes field",
+ IllegalArgumentException.class,
+ () ->
+ errorPipeline
+ .apply(
+ Create.of(
+ Row.withSchema(bytesSchema)
+ .withFieldValue("badfield", new byte[] {1, 2, 3})
+ .withFieldValue("ok", "abcdefg")
+ .build()))
+ .setRowSchema(bytesSchema)
+ .apply(CsvIO.writeRows().withCSVFormat(CSVFormat.DEFAULT).to(to)));
+
+ assertThrows(
+ "schema should not contain map field",
+ IllegalArgumentException.class,
+ () ->
+ errorPipeline
+ .apply(
+ Create.of(
+ Row.withSchema(mapSchema)
+ .withFieldValue("badfield", ImmutableMap.of("a", "b"))
+ .withFieldValue("ok", "abcdefg")
+ .build()))
+ .setRowSchema(mapSchema)
+ .apply(CsvIO.writeRows().to(to)));
+
+ assertThrows(
+ "schema should not contain array field",
+ IllegalArgumentException.class,
+ () ->
+ errorPipeline
+ .apply(
+ Create.of(
+ Row.withSchema(arraySchema)
+ .withFieldValue("badfield", Arrays.asList("a", "b"))
+ .withFieldValue("ok", "abcdefg")
+ .build()))
+ .setRowSchema(arraySchema)
+ .apply(CsvIO.writeRows().withCSVFormat(CSVFormat.DEFAULT).to(to)));
+
+ assertThrows(
+ "schema should not contain row field",
+ IllegalArgumentException.class,
+ () ->
+ errorPipeline
+ .apply(
+ Create.of(
+ Row.withSchema(rowSchema)
+ .withFieldValue("badfield", rows.get(0))
+ .withFieldValue("ok", "abcdefg")
+ .build()))
+ .setRowSchema(rowSchema)
+ .apply(CsvIO.writeRows().to(to)));
+ }
+
+ @Test
+ public void invalidSchemaFields() {
+ String to = tmpFolder.getRoot().getAbsolutePath() + "/somewhere";
+
+ assertThrows(
+ "schema field names should not be empty",
+ IllegalArgumentException.class,
+ () ->
+ errorPipeline
+ .apply(Create.of(rows))
+ .setRowSchema(ALL_DATA_TYPES_SCHEMA)
+ .apply(
+ CsvIO.writeRows()
+ .to(to)
+ .withSchemaFields(Arrays.asList("", "string", "anInt"))));
+
+ assertThrows(
+ "schema fields should not be an empty list",
+ IllegalArgumentException.class,
+ () ->
+ errorPipeline
+ .apply(Create.of(rows))
+ .setRowSchema(ALL_DATA_TYPES_SCHEMA)
+ .apply(CsvIO.writeRows().to(to).withSchemaFields(Collections.emptyList())));
+
+ assertThrows(
+ "field names should match schema",
+ IllegalArgumentException.class,
+ () ->
+ errorPipeline
+ .apply(Create.of(rows))
+ .setRowSchema(ALL_DATA_TYPES_SCHEMA)
+ .apply(
+ CsvIO.writeRows()
+ .to(to)
+ .withSchemaFields(Arrays.asList("doesnotexist", "string", "anInt"))));
+ }
+
+ @Test
+ public void withCompression() {
+ String to = "with_compression";
Review Comment:
nit: It is simpler to reason about the tests if you use a new temporary folder instead of using a special prefix. It also makes file matching easier because then it is just all files in the directory. See https://junit.org/junit4/javadoc/4.12/org/junit/rules/TemporaryFolder.html#newFolder(java.lang.String)
##########
sdks/java/io/csv/src/main/java/org/apache/beam/sdk/io/csv/CsvRowConversions.java:
##########
@@ -0,0 +1,166 @@
+/*
+ * 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.sdk.io.csv;
+
+import com.google.auto.value.AutoValue;
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.stream.Collectors;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.Field;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.values.Row;
+import org.apache.commons.csv.CSVFormat;
+
+/** Contains classes and methods to help with converting between {@link Row} and CSV strings. */
+class CsvRowConversions {
+
+ /** Converts between {@link Row} and CSV string using a {@link CSVFormat}. */
+ @AutoValue
+ abstract static class RowToCsv implements SerializableFunction<Row, String> {
+
+ static Builder builder() {
+ return new AutoValue_CsvRowConversions_RowToCsv.Builder();
+ }
+
+ /** The expected {@link Schema} of the {@link Row} input. */
+ abstract Schema getSchema();
+
+ /** The {@link CSVFormat} of the converted {@link Row} input. */
+ abstract CSVFormat getCSVFormat();
+
+ /** The order and subset of {@link Schema} field names to drive the {@link Row} conversion. */
+ abstract List<String> getSchemaFields();
+
+ /** Builds a header from {@link #getSchemaFields()} using the {@link #getCSVFormat()}. */
+ String buildHeader() {
+ return buildHeaderFrom(getSchemaFields(), getCSVFormat());
+ }
+
+ /**
+ * Converts a {@link Row} to a CSV string formatted using {@link #getCSVFormat()}. Null values
+ * are converted to empty strings.
+ */
+ @Override
+ public String apply(Row input) {
+ List<String> schemaFields = getSchemaFields();
+ Object[] values = new Object[schemaFields.size()];
+ for (int i = 0; i < schemaFields.size(); i++) {
+ String name = schemaFields.get(i);
+ values[i] = input.getValue(name);
+ }
+ return getCSVFormat().format(values);
+ }
+
+ @AutoValue.Builder
+ abstract static class Builder {
+
+ /** The expected {@link Schema} of the {@link Row} input. */
+ abstract Builder setSchema(Schema schema);
+
+ abstract Schema getSchema();
+
+ /** The {@link CSVFormat} of the converted {@link Row} input. */
+ abstract Builder setCSVFormat(CSVFormat format);
+
+ /** The order and subset of {@link Schema} field names to drive the {@link Row} conversion. */
+ abstract Builder setSchemaFields(List<String> fields);
+
+ abstract List<String> getSchemaFields();
+
+ abstract RowToCsv autoBuild();
+
+ final RowToCsv build() {
+
+ validateHeaderAgainstSchema(getSchemaFields(), getSchema());
+
+ return autoBuild();
+ }
+ }
+ }
+
+ /** Formats columns into a header String based on {@link CSVFormat}. */
+ static String buildHeaderFrom(List<String> columns, CSVFormat csvFormat) {
+ StringBuilder builder = new StringBuilder();
+ try {
+ boolean newRecord = true;
+ for (String name : columns) {
+ csvFormat.print(name, builder, newRecord);
+ newRecord = false;
+ }
+ } catch (IOException e) {
+ throw new IllegalStateException(e);
+ }
+ return builder.toString();
+ }
+
+ private static void validateHeaderAgainstSchema(List<String> schemaFields, Schema schema) {
+ if (schema.getFieldCount() == 0) {
+ throw new IllegalArgumentException("schema is empty");
+ }
+ if (schemaFields.isEmpty()) {
+ throw new IllegalArgumentException(
+ "Columns is empty. An intent to not override columns, should assign null to the columns parameter.");
+ }
+ List<String> missing = new ArrayList<>();
+ List<Field> invalid = new ArrayList<>();
+
+ for (String name : schemaFields) {
+ if (name.isEmpty()) {
Review Comment:
A quick look seems to show that we can have schemas and CSVs with empty field names
I still think it would make sense to ensure that `name != null`
##########
sdks/java/io/csv/src/test/java/org/apache/beam/sdk/io/csv/CsvIOWriteTest.java:
##########
@@ -0,0 +1,534 @@
+/*
+ * 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.sdk.io.csv;
+
+import static org.apache.beam.sdk.io.csv.CsvIOTestHelpers.ALL_DATA_TYPES_SCHEMA;
+import static org.apache.beam.sdk.io.csv.CsvIOTestHelpers.rowOf;
+import static org.junit.Assert.assertArrayEquals;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertThrows;
+import static org.junit.Assert.assertTrue;
+
+import java.io.BufferedReader;
+import java.io.File;
+import java.io.IOException;
+import java.math.BigDecimal;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.regex.Pattern;
+import java.util.stream.Collectors;
+import org.apache.beam.sdk.io.Compression;
+import org.apache.beam.sdk.io.csv.CsvIOTestHelpers.AllDataTypes;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.Field;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.testing.TestPipeline;
+import org.apache.beam.sdk.transforms.Create;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Charsets;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.commons.csv.CSVFormat;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Instant;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.TemporaryFolder;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/** Tests for {@link CsvIO.Write}. */
+@RunWith(JUnit4.class)
+public class CsvIOWriteTest {
+
+ @Rule public TemporaryFolder tmpFolder = TemporaryFolder.builder().build();
+
+ @Rule public TestPipeline pipeline = TestPipeline.create();
+
+ @Rule
+ public TestPipeline errorPipeline = TestPipeline.create().enableAbandonedNodeEnforcement(false);
+
+ private final List<AllDataTypes> userTypes =
+ Arrays.asList(
+ AllDataTypes.of(
+ true,
+ (byte) 0,
+ Instant.ofEpochMilli(0L).toDateTime(),
+ BigDecimal.valueOf(0L),
+ 0.1,
+ 0.2f,
+ (short) 3,
+ 4,
+ 5L,
+ "a"),
+ AllDataTypes.of(
+ true,
+ (byte) 1,
+ Instant.ofEpochMilli(1L).toDateTime(),
+ BigDecimal.valueOf(1L),
+ 1.1,
+ 1.2f,
+ (short) 4,
+ 5,
+ 6L,
+ "b"),
+ AllDataTypes.of(
+ false,
+ (byte) 2,
+ Instant.ofEpochMilli(2L).toDateTime(),
+ BigDecimal.valueOf(2L),
+ 2.0,
+ 2.1f,
+ (short) 5,
+ 6,
+ 7L,
+ "c"));
+
+ private final List<Row> rows =
+ Arrays.asList(
+ rowOf(
+ true,
+ (byte) 0,
+ Instant.ofEpochMilli(0L).toDateTime(),
+ BigDecimal.valueOf(0L),
+ 0.1,
+ 0.2f,
+ (short) 3,
+ 4,
+ 5L,
+ "a"),
+ rowOf(
+ true,
+ (byte) 1,
+ Instant.ofEpochMilli(1L).toDateTime(),
+ BigDecimal.valueOf(1L),
+ 1.1,
+ 1.2f,
+ (short) 4,
+ 5,
+ 6L,
+ "b"),
+ rowOf(
+ false,
+ (byte) 2,
+ Instant.ofEpochMilli(2L).toDateTime(),
+ BigDecimal.valueOf(2L),
+ 2.0,
+ 2.1f,
+ (short) 5,
+ 6,
+ 7L,
+ "c"));
+
+ @Test
+ public void validateSchema() {
+ String to = tmpFolder.getRoot().getAbsolutePath() + "/somewhere";
+ Schema emptySchema = Schema.of();
+ Schema bytesSchema =
+ Schema.of(Field.of("badfield", FieldType.BYTES), Field.of("ok", FieldType.STRING));
+ Schema mapSchema =
+ Schema.of(
+ Field.of("badfield", FieldType.map(FieldType.STRING, FieldType.STRING)),
+ Field.of("ok", FieldType.STRING));
+ Schema arraySchema =
+ Schema.of(
+ Field.of("badfield", FieldType.array(FieldType.STRING)),
+ Field.of("ok", FieldType.STRING));
+ Schema rowSchema =
+ Schema.of(
+ Field.of("badfield", FieldType.row(ALL_DATA_TYPES_SCHEMA)),
+ Field.of("ok", FieldType.STRING));
+
+ assertThrows(
+ "schema should not be empty",
+ IllegalArgumentException.class,
+ () ->
+ errorPipeline
+ .apply(Create.of(Row.nullRow(emptySchema)).withRowSchema(emptySchema))
+ .apply(CsvIO.writeRows().to(to)));
+
+ assertThrows(
+ "schema should not contain bytes field",
+ IllegalArgumentException.class,
+ () ->
+ errorPipeline
+ .apply(
+ Create.of(
+ Row.withSchema(bytesSchema)
+ .withFieldValue("badfield", new byte[] {1, 2, 3})
+ .withFieldValue("ok", "abcdefg")
+ .build()))
+ .setRowSchema(bytesSchema)
+ .apply(CsvIO.writeRows().withCSVFormat(CSVFormat.DEFAULT).to(to)));
+
+ assertThrows(
+ "schema should not contain map field",
+ IllegalArgumentException.class,
+ () ->
+ errorPipeline
+ .apply(
+ Create.of(
+ Row.withSchema(mapSchema)
+ .withFieldValue("badfield", ImmutableMap.of("a", "b"))
+ .withFieldValue("ok", "abcdefg")
+ .build()))
+ .setRowSchema(mapSchema)
+ .apply(CsvIO.writeRows().to(to)));
+
+ assertThrows(
+ "schema should not contain array field",
+ IllegalArgumentException.class,
+ () ->
+ errorPipeline
+ .apply(
+ Create.of(
+ Row.withSchema(arraySchema)
+ .withFieldValue("badfield", Arrays.asList("a", "b"))
+ .withFieldValue("ok", "abcdefg")
+ .build()))
+ .setRowSchema(arraySchema)
+ .apply(CsvIO.writeRows().withCSVFormat(CSVFormat.DEFAULT).to(to)));
+
+ assertThrows(
+ "schema should not contain row field",
+ IllegalArgumentException.class,
+ () ->
+ errorPipeline
+ .apply(
+ Create.of(
+ Row.withSchema(rowSchema)
+ .withFieldValue("badfield", rows.get(0))
+ .withFieldValue("ok", "abcdefg")
+ .build()))
+ .setRowSchema(rowSchema)
+ .apply(CsvIO.writeRows().to(to)));
+ }
+
+ @Test
+ public void invalidSchemaFields() {
+ String to = tmpFolder.getRoot().getAbsolutePath() + "/somewhere";
+
+ assertThrows(
+ "schema field names should not be empty",
+ IllegalArgumentException.class,
+ () ->
+ errorPipeline
+ .apply(Create.of(rows))
+ .setRowSchema(ALL_DATA_TYPES_SCHEMA)
+ .apply(
+ CsvIO.writeRows()
+ .to(to)
+ .withSchemaFields(Arrays.asList("", "string", "anInt"))));
+
+ assertThrows(
+ "schema fields should not be an empty list",
+ IllegalArgumentException.class,
+ () ->
+ errorPipeline
+ .apply(Create.of(rows))
+ .setRowSchema(ALL_DATA_TYPES_SCHEMA)
+ .apply(CsvIO.writeRows().to(to).withSchemaFields(Collections.emptyList())));
+
+ assertThrows(
+ "field names should match schema",
+ IllegalArgumentException.class,
+ () ->
+ errorPipeline
+ .apply(Create.of(rows))
+ .setRowSchema(ALL_DATA_TYPES_SCHEMA)
+ .apply(
+ CsvIO.writeRows()
+ .to(to)
+ .withSchemaFields(Arrays.asList("doesnotexist", "string", "anInt"))));
+ }
+
+ @Test
+ public void withCompression() {
+ String to = "with_compression";
+ String regex = String.format("^%s.*$", to);
+ pipeline
+ .apply(Create.of(rows).withRowSchema(ALL_DATA_TYPES_SCHEMA))
+ .apply(
+ CsvIO.writeRows()
+ .to(tmpFolder.getRoot().getAbsolutePath() + "/" + to)
+ .withCompression(Compression.GZIP));
+ pipeline.run().waitUntilFinish();
+ String[] files = filesMatching(regex);
+ for (String name : files) {
+ Path p = Paths.get(tmpFolder.getRoot().getAbsolutePath(), name);
+ File f = new File(p.toString());
+ assertTrue(Compression.GZIP.isCompressed(f.getAbsolutePath()));
+ }
+ }
+
+ @Test
+ public void withInvalidType() {
+ assertThrows(
+ IllegalArgumentException.class,
+ () ->
+ errorPipeline
+ .apply(Create.of("1,2,3", "4.5.6"))
+ .apply(
+ CsvIO.<String>write().to(tmpFolder.getRoot().getAbsolutePath() + "/badtype")));
+ }
+
+ @Test
+ public void withNumShards() {
+ String to = "with_num_shards";
+ String regex = String.format("^%s.*$", to);
+ pipeline
+ .apply(Create.of(rows).withRowSchema(ALL_DATA_TYPES_SCHEMA))
+ .apply(
+ CsvIO.writeRows()
+ .to(tmpFolder.getRoot().getAbsolutePath() + "/" + to)
+ .withNumShards(1));
+ pipeline.run().waitUntilFinish();
+ String[] files = filesMatching(regex);
+ assertEquals(1, files.length);
+ }
+
+ @Test
+ public void withPreamble() throws IOException {
+ String preamble =
+ "id=44ebe8ee-dc0f-4c93-8b47-2e627e1eb988\ndate=2022-01-02\ntest=e3bbfc47-9bcb-46f0-8944-829827ee5222";
+ String to = "rows_preamble";
+ String regex = String.format("^%s.*$", to);
+ pipeline
+ .apply(Create.of(rows).withRowSchema(ALL_DATA_TYPES_SCHEMA))
+ .apply(
+ CsvIO.writeRows()
+ .to(tmpFolder.getRoot().getAbsolutePath() + "/" + to)
+ .withPreamble(preamble));
+ pipeline.run().waitUntilFinish();
+ String[] files = filesMatching(regex);
+ assertTrue(
+ "CsvIO.writeRows should write to " + tmpFolder.getRoot().getAbsolutePath() + "/" + to,
+ files.length > 0);
+ String expectedHeader =
+ "aBoolean,aByte,aDouble,aFloat,aLong,aShort,anInt,dateTime,decimal,string";
+ for (String name : files) {
Review Comment:
why would we expect all files to have the same contents?
I think what you want to do is check that the header/preamble are correct and return the rest that is found in each file and then combine all that across all the files and check that it matches expected
here and other places where you use `assertFileContentsMatchInAnyOrder`
##########
sdks/java/io/csv/src/test/java/org/apache/beam/sdk/io/csv/CsvIOWriteTest.java:
##########
@@ -0,0 +1,534 @@
+/*
+ * 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.sdk.io.csv;
+
+import static org.apache.beam.sdk.io.csv.CsvIOTestHelpers.ALL_DATA_TYPES_SCHEMA;
+import static org.apache.beam.sdk.io.csv.CsvIOTestHelpers.rowOf;
+import static org.junit.Assert.assertArrayEquals;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertThrows;
+import static org.junit.Assert.assertTrue;
+
+import java.io.BufferedReader;
+import java.io.File;
+import java.io.IOException;
+import java.math.BigDecimal;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.regex.Pattern;
+import java.util.stream.Collectors;
+import org.apache.beam.sdk.io.Compression;
+import org.apache.beam.sdk.io.csv.CsvIOTestHelpers.AllDataTypes;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.Field;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.testing.TestPipeline;
+import org.apache.beam.sdk.transforms.Create;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Charsets;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.commons.csv.CSVFormat;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Instant;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.TemporaryFolder;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/** Tests for {@link CsvIO.Write}. */
+@RunWith(JUnit4.class)
+public class CsvIOWriteTest {
+
+ @Rule public TemporaryFolder tmpFolder = TemporaryFolder.builder().build();
+
+ @Rule public TestPipeline pipeline = TestPipeline.create();
+
+ @Rule
+ public TestPipeline errorPipeline = TestPipeline.create().enableAbandonedNodeEnforcement(false);
+
+ private final List<AllDataTypes> userTypes =
+ Arrays.asList(
+ AllDataTypes.of(
+ true,
+ (byte) 0,
+ Instant.ofEpochMilli(0L).toDateTime(),
+ BigDecimal.valueOf(0L),
+ 0.1,
+ 0.2f,
+ (short) 3,
+ 4,
+ 5L,
+ "a"),
+ AllDataTypes.of(
+ true,
+ (byte) 1,
+ Instant.ofEpochMilli(1L).toDateTime(),
+ BigDecimal.valueOf(1L),
+ 1.1,
+ 1.2f,
+ (short) 4,
+ 5,
+ 6L,
+ "b"),
+ AllDataTypes.of(
+ false,
+ (byte) 2,
+ Instant.ofEpochMilli(2L).toDateTime(),
+ BigDecimal.valueOf(2L),
+ 2.0,
+ 2.1f,
+ (short) 5,
+ 6,
+ 7L,
+ "c"));
+
+ private final List<Row> rows =
+ Arrays.asList(
+ rowOf(
+ true,
+ (byte) 0,
+ Instant.ofEpochMilli(0L).toDateTime(),
+ BigDecimal.valueOf(0L),
+ 0.1,
+ 0.2f,
+ (short) 3,
+ 4,
+ 5L,
+ "a"),
+ rowOf(
+ true,
+ (byte) 1,
+ Instant.ofEpochMilli(1L).toDateTime(),
+ BigDecimal.valueOf(1L),
+ 1.1,
+ 1.2f,
+ (short) 4,
+ 5,
+ 6L,
+ "b"),
+ rowOf(
+ false,
+ (byte) 2,
+ Instant.ofEpochMilli(2L).toDateTime(),
+ BigDecimal.valueOf(2L),
+ 2.0,
+ 2.1f,
+ (short) 5,
+ 6,
+ 7L,
+ "c"));
+
+ @Test
+ public void validateSchema() {
+ String to = tmpFolder.getRoot().getAbsolutePath() + "/somewhere";
+ Schema emptySchema = Schema.of();
+ Schema bytesSchema =
+ Schema.of(Field.of("badfield", FieldType.BYTES), Field.of("ok", FieldType.STRING));
+ Schema mapSchema =
+ Schema.of(
+ Field.of("badfield", FieldType.map(FieldType.STRING, FieldType.STRING)),
+ Field.of("ok", FieldType.STRING));
+ Schema arraySchema =
+ Schema.of(
+ Field.of("badfield", FieldType.array(FieldType.STRING)),
+ Field.of("ok", FieldType.STRING));
+ Schema rowSchema =
+ Schema.of(
+ Field.of("badfield", FieldType.row(ALL_DATA_TYPES_SCHEMA)),
+ Field.of("ok", FieldType.STRING));
+
+ assertThrows(
+ "schema should not be empty",
+ IllegalArgumentException.class,
+ () ->
+ errorPipeline
+ .apply(Create.of(Row.nullRow(emptySchema)).withRowSchema(emptySchema))
+ .apply(CsvIO.writeRows().to(to)));
+
+ assertThrows(
+ "schema should not contain bytes field",
+ IllegalArgumentException.class,
+ () ->
+ errorPipeline
+ .apply(
+ Create.of(
+ Row.withSchema(bytesSchema)
+ .withFieldValue("badfield", new byte[] {1, 2, 3})
+ .withFieldValue("ok", "abcdefg")
+ .build()))
+ .setRowSchema(bytesSchema)
+ .apply(CsvIO.writeRows().withCSVFormat(CSVFormat.DEFAULT).to(to)));
+
+ assertThrows(
+ "schema should not contain map field",
+ IllegalArgumentException.class,
+ () ->
+ errorPipeline
+ .apply(
+ Create.of(
+ Row.withSchema(mapSchema)
+ .withFieldValue("badfield", ImmutableMap.of("a", "b"))
+ .withFieldValue("ok", "abcdefg")
+ .build()))
+ .setRowSchema(mapSchema)
+ .apply(CsvIO.writeRows().to(to)));
+
+ assertThrows(
+ "schema should not contain array field",
+ IllegalArgumentException.class,
+ () ->
+ errorPipeline
+ .apply(
+ Create.of(
+ Row.withSchema(arraySchema)
+ .withFieldValue("badfield", Arrays.asList("a", "b"))
+ .withFieldValue("ok", "abcdefg")
+ .build()))
+ .setRowSchema(arraySchema)
+ .apply(CsvIO.writeRows().withCSVFormat(CSVFormat.DEFAULT).to(to)));
+
+ assertThrows(
+ "schema should not contain row field",
+ IllegalArgumentException.class,
+ () ->
+ errorPipeline
+ .apply(
+ Create.of(
+ Row.withSchema(rowSchema)
+ .withFieldValue("badfield", rows.get(0))
+ .withFieldValue("ok", "abcdefg")
+ .build()))
+ .setRowSchema(rowSchema)
+ .apply(CsvIO.writeRows().to(to)));
+ }
+
+ @Test
+ public void invalidSchemaFields() {
+ String to = tmpFolder.getRoot().getAbsolutePath() + "/somewhere";
+
+ assertThrows(
+ "schema field names should not be empty",
+ IllegalArgumentException.class,
+ () ->
+ errorPipeline
+ .apply(Create.of(rows))
+ .setRowSchema(ALL_DATA_TYPES_SCHEMA)
+ .apply(
+ CsvIO.writeRows()
+ .to(to)
+ .withSchemaFields(Arrays.asList("", "string", "anInt"))));
+
+ assertThrows(
+ "schema fields should not be an empty list",
+ IllegalArgumentException.class,
+ () ->
+ errorPipeline
+ .apply(Create.of(rows))
+ .setRowSchema(ALL_DATA_TYPES_SCHEMA)
+ .apply(CsvIO.writeRows().to(to).withSchemaFields(Collections.emptyList())));
+
+ assertThrows(
+ "field names should match schema",
+ IllegalArgumentException.class,
+ () ->
+ errorPipeline
+ .apply(Create.of(rows))
+ .setRowSchema(ALL_DATA_TYPES_SCHEMA)
+ .apply(
+ CsvIO.writeRows()
+ .to(to)
+ .withSchemaFields(Arrays.asList("doesnotexist", "string", "anInt"))));
+ }
+
+ @Test
+ public void withCompression() {
+ String to = "with_compression";
+ String regex = String.format("^%s.*$", to);
+ pipeline
+ .apply(Create.of(rows).withRowSchema(ALL_DATA_TYPES_SCHEMA))
+ .apply(
+ CsvIO.writeRows()
+ .to(tmpFolder.getRoot().getAbsolutePath() + "/" + to)
+ .withCompression(Compression.GZIP));
+ pipeline.run().waitUntilFinish();
+ String[] files = filesMatching(regex);
+ for (String name : files) {
+ Path p = Paths.get(tmpFolder.getRoot().getAbsolutePath(), name);
+ File f = new File(p.toString());
+ assertTrue(Compression.GZIP.isCompressed(f.getAbsolutePath()));
+ }
+ }
+
+ @Test
+ public void withInvalidType() {
+ assertThrows(
+ IllegalArgumentException.class,
Review Comment:
Can you match against the error string as well to ensure we are catching the exception from the right place and not just any `IllegalArgumentException`
##########
sdks/java/io/csv/src/test/java/org/apache/beam/sdk/io/csv/CsvIOWriteTest.java:
##########
@@ -0,0 +1,534 @@
+/*
+ * 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.sdk.io.csv;
+
+import static org.apache.beam.sdk.io.csv.CsvIOTestHelpers.ALL_DATA_TYPES_SCHEMA;
+import static org.apache.beam.sdk.io.csv.CsvIOTestHelpers.rowOf;
+import static org.junit.Assert.assertArrayEquals;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertThrows;
+import static org.junit.Assert.assertTrue;
+
+import java.io.BufferedReader;
+import java.io.File;
+import java.io.IOException;
+import java.math.BigDecimal;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.regex.Pattern;
+import java.util.stream.Collectors;
+import org.apache.beam.sdk.io.Compression;
+import org.apache.beam.sdk.io.csv.CsvIOTestHelpers.AllDataTypes;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.Field;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.testing.TestPipeline;
+import org.apache.beam.sdk.transforms.Create;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Charsets;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.commons.csv.CSVFormat;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Instant;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.TemporaryFolder;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/** Tests for {@link CsvIO.Write}. */
+@RunWith(JUnit4.class)
+public class CsvIOWriteTest {
+
+ @Rule public TemporaryFolder tmpFolder = TemporaryFolder.builder().build();
+
+ @Rule public TestPipeline pipeline = TestPipeline.create();
+
+ @Rule
+ public TestPipeline errorPipeline = TestPipeline.create().enableAbandonedNodeEnforcement(false);
+
+ private final List<AllDataTypes> userTypes =
+ Arrays.asList(
+ AllDataTypes.of(
+ true,
+ (byte) 0,
+ Instant.ofEpochMilli(0L).toDateTime(),
+ BigDecimal.valueOf(0L),
+ 0.1,
+ 0.2f,
+ (short) 3,
+ 4,
+ 5L,
+ "a"),
+ AllDataTypes.of(
+ true,
+ (byte) 1,
+ Instant.ofEpochMilli(1L).toDateTime(),
+ BigDecimal.valueOf(1L),
+ 1.1,
+ 1.2f,
+ (short) 4,
+ 5,
+ 6L,
+ "b"),
+ AllDataTypes.of(
+ false,
+ (byte) 2,
+ Instant.ofEpochMilli(2L).toDateTime(),
+ BigDecimal.valueOf(2L),
+ 2.0,
+ 2.1f,
+ (short) 5,
+ 6,
+ 7L,
+ "c"));
+
+ private final List<Row> rows =
+ Arrays.asList(
+ rowOf(
+ true,
+ (byte) 0,
+ Instant.ofEpochMilli(0L).toDateTime(),
+ BigDecimal.valueOf(0L),
+ 0.1,
+ 0.2f,
+ (short) 3,
+ 4,
+ 5L,
+ "a"),
+ rowOf(
+ true,
+ (byte) 1,
+ Instant.ofEpochMilli(1L).toDateTime(),
+ BigDecimal.valueOf(1L),
+ 1.1,
+ 1.2f,
+ (short) 4,
+ 5,
+ 6L,
+ "b"),
+ rowOf(
+ false,
+ (byte) 2,
+ Instant.ofEpochMilli(2L).toDateTime(),
+ BigDecimal.valueOf(2L),
+ 2.0,
+ 2.1f,
+ (short) 5,
+ 6,
+ 7L,
+ "c"));
+
+ @Test
+ public void validateSchema() {
+ String to = tmpFolder.getRoot().getAbsolutePath() + "/somewhere";
+ Schema emptySchema = Schema.of();
+ Schema bytesSchema =
+ Schema.of(Field.of("badfield", FieldType.BYTES), Field.of("ok", FieldType.STRING));
+ Schema mapSchema =
+ Schema.of(
+ Field.of("badfield", FieldType.map(FieldType.STRING, FieldType.STRING)),
+ Field.of("ok", FieldType.STRING));
+ Schema arraySchema =
+ Schema.of(
+ Field.of("badfield", FieldType.array(FieldType.STRING)),
+ Field.of("ok", FieldType.STRING));
+ Schema rowSchema =
+ Schema.of(
+ Field.of("badfield", FieldType.row(ALL_DATA_TYPES_SCHEMA)),
+ Field.of("ok", FieldType.STRING));
+
+ assertThrows(
+ "schema should not be empty",
+ IllegalArgumentException.class,
+ () ->
+ errorPipeline
+ .apply(Create.of(Row.nullRow(emptySchema)).withRowSchema(emptySchema))
+ .apply(CsvIO.writeRows().to(to)));
+
+ assertThrows(
+ "schema should not contain bytes field",
+ IllegalArgumentException.class,
+ () ->
+ errorPipeline
+ .apply(
+ Create.of(
+ Row.withSchema(bytesSchema)
+ .withFieldValue("badfield", new byte[] {1, 2, 3})
+ .withFieldValue("ok", "abcdefg")
+ .build()))
+ .setRowSchema(bytesSchema)
+ .apply(CsvIO.writeRows().withCSVFormat(CSVFormat.DEFAULT).to(to)));
+
+ assertThrows(
+ "schema should not contain map field",
+ IllegalArgumentException.class,
+ () ->
+ errorPipeline
+ .apply(
+ Create.of(
+ Row.withSchema(mapSchema)
+ .withFieldValue("badfield", ImmutableMap.of("a", "b"))
+ .withFieldValue("ok", "abcdefg")
+ .build()))
+ .setRowSchema(mapSchema)
+ .apply(CsvIO.writeRows().to(to)));
+
+ assertThrows(
+ "schema should not contain array field",
+ IllegalArgumentException.class,
+ () ->
+ errorPipeline
+ .apply(
+ Create.of(
+ Row.withSchema(arraySchema)
+ .withFieldValue("badfield", Arrays.asList("a", "b"))
+ .withFieldValue("ok", "abcdefg")
+ .build()))
+ .setRowSchema(arraySchema)
+ .apply(CsvIO.writeRows().withCSVFormat(CSVFormat.DEFAULT).to(to)));
+
+ assertThrows(
+ "schema should not contain row field",
+ IllegalArgumentException.class,
+ () ->
+ errorPipeline
+ .apply(
+ Create.of(
+ Row.withSchema(rowSchema)
+ .withFieldValue("badfield", rows.get(0))
+ .withFieldValue("ok", "abcdefg")
+ .build()))
+ .setRowSchema(rowSchema)
+ .apply(CsvIO.writeRows().to(to)));
+ }
+
+ @Test
+ public void invalidSchemaFields() {
+ String to = tmpFolder.getRoot().getAbsolutePath() + "/somewhere";
+
+ assertThrows(
+ "schema field names should not be empty",
+ IllegalArgumentException.class,
+ () ->
+ errorPipeline
+ .apply(Create.of(rows))
+ .setRowSchema(ALL_DATA_TYPES_SCHEMA)
+ .apply(
+ CsvIO.writeRows()
+ .to(to)
+ .withSchemaFields(Arrays.asList("", "string", "anInt"))));
+
+ assertThrows(
+ "schema fields should not be an empty list",
+ IllegalArgumentException.class,
+ () ->
+ errorPipeline
+ .apply(Create.of(rows))
+ .setRowSchema(ALL_DATA_TYPES_SCHEMA)
+ .apply(CsvIO.writeRows().to(to).withSchemaFields(Collections.emptyList())));
+
+ assertThrows(
+ "field names should match schema",
+ IllegalArgumentException.class,
+ () ->
+ errorPipeline
+ .apply(Create.of(rows))
+ .setRowSchema(ALL_DATA_TYPES_SCHEMA)
+ .apply(
+ CsvIO.writeRows()
+ .to(to)
+ .withSchemaFields(Arrays.asList("doesnotexist", "string", "anInt"))));
+ }
+
+ @Test
+ public void withCompression() {
+ String to = "with_compression";
+ String regex = String.format("^%s.*$", to);
+ pipeline
+ .apply(Create.of(rows).withRowSchema(ALL_DATA_TYPES_SCHEMA))
+ .apply(
+ CsvIO.writeRows()
+ .to(tmpFolder.getRoot().getAbsolutePath() + "/" + to)
+ .withCompression(Compression.GZIP));
+ pipeline.run().waitUntilFinish();
+ String[] files = filesMatching(regex);
+ for (String name : files) {
+ Path p = Paths.get(tmpFolder.getRoot().getAbsolutePath(), name);
+ File f = new File(p.toString());
+ assertTrue(Compression.GZIP.isCompressed(f.getAbsolutePath()));
+ }
+ }
+
+ @Test
+ public void withInvalidType() {
+ assertThrows(
+ IllegalArgumentException.class,
+ () ->
+ errorPipeline
+ .apply(Create.of("1,2,3", "4.5.6"))
+ .apply(
+ CsvIO.<String>write().to(tmpFolder.getRoot().getAbsolutePath() + "/badtype")));
+ }
+
+ @Test
+ public void withNumShards() {
+ String to = "with_num_shards";
+ String regex = String.format("^%s.*$", to);
+ pipeline
+ .apply(Create.of(rows).withRowSchema(ALL_DATA_TYPES_SCHEMA))
+ .apply(
+ CsvIO.writeRows()
+ .to(tmpFolder.getRoot().getAbsolutePath() + "/" + to)
+ .withNumShards(1));
+ pipeline.run().waitUntilFinish();
+ String[] files = filesMatching(regex);
+ assertEquals(1, files.length);
+ }
+
+ @Test
+ public void withPreamble() throws IOException {
+ String preamble =
+ "id=44ebe8ee-dc0f-4c93-8b47-2e627e1eb988\ndate=2022-01-02\ntest=e3bbfc47-9bcb-46f0-8944-829827ee5222";
+ String to = "rows_preamble";
+ String regex = String.format("^%s.*$", to);
+ pipeline
+ .apply(Create.of(rows).withRowSchema(ALL_DATA_TYPES_SCHEMA))
+ .apply(
+ CsvIO.writeRows()
+ .to(tmpFolder.getRoot().getAbsolutePath() + "/" + to)
+ .withPreamble(preamble));
+ pipeline.run().waitUntilFinish();
+ String[] files = filesMatching(regex);
+ assertTrue(
+ "CsvIO.writeRows should write to " + tmpFolder.getRoot().getAbsolutePath() + "/" + to,
+ files.length > 0);
+ String expectedHeader =
+ "aBoolean,aByte,aDouble,aFloat,aLong,aShort,anInt,dateTime,decimal,string";
+ for (String name : files) {
+ Path p = Paths.get(tmpFolder.getRoot().getAbsolutePath(), name);
+ File f = new File(p.toString());
+ assertFileContentsMatchInAnyOrder(
+ preamble,
+ expectedHeader,
+ f,
+ "true,0,0.1,0.2,5,3,4,1970-01-01T00:00:00.000Z,0,a",
+ "true,1,1.1,1.2,6,4,5,1970-01-01T00:00:00.001Z,1,b",
+ "false,2,2.0,2.1,7,5,6,1970-01-01T00:00:00.002Z,2,c");
+ }
+ }
+
+ @Test
+ public void rowsWithCsvFormatNonDefault() throws IOException {
+ String to = "rows_non_default";
+ String regex = String.format("^%s.*$", to);
+ pipeline
+ .apply(Create.of(rows).withRowSchema(ALL_DATA_TYPES_SCHEMA))
+ .apply(
+ CsvIO.writeRows()
+ .to(tmpFolder.getRoot().getAbsolutePath() + "/" + to)
+ .withCSVFormat(CSVFormat.MYSQL));
+ pipeline.run().waitUntilFinish();
+ String[] files = filesMatching(regex);
+ assertTrue(
+ "CsvIO.writeRows should write to " + tmpFolder.getRoot().getAbsolutePath() + "/" + to,
+ files.length > 0);
+ String expectedHeader =
+ "aBoolean\taByte\taDouble\taFloat\taLong\taShort\tanInt\tdateTime\tdecimal\tstring";
+ for (String name : files) {
+ Path p = Paths.get(tmpFolder.getRoot().getAbsolutePath(), name);
+ File f = new File(p.toString());
+ assertFileContentsMatchInAnyOrder(
+ null,
+ expectedHeader,
+ f,
+ "true\t0\t0.1\t0.2\t5\t3\t4\t1970-01-01T00:00:00.000Z\t0\ta",
+ "true\t1\t1.1\t1.2\t6\t4\t5\t1970-01-01T00:00:00.001Z\t1\tb",
+ "false\t2\t2.0\t2.1\t7\t5\t6\t1970-01-01T00:00:00.002Z\t2\tc");
+ }
+ }
+
+ @Test
+ public void withSchemaFields() throws IOException {
+ String to = "with_schema_fields";
+ String regex = String.format("^%s.*$", to);
+ pipeline
+ .apply(Create.of(rows).withRowSchema(ALL_DATA_TYPES_SCHEMA))
+ .apply(
+ CsvIO.writeRows()
+ .to(tmpFolder.getRoot().getAbsolutePath() + "/" + to)
+ .withSchemaFields(Arrays.asList("aShort", "aBoolean", "aLong")));
+ pipeline.run().waitUntilFinish();
+ String[] files = filesMatching(regex);
+ assertTrue(
+ "CsvIO.writeRows should write to " + tmpFolder.getRoot().getAbsolutePath() + "/" + to,
+ files.length > 0);
+ String expectedHeader = "aShort,aBoolean,aLong";
+ for (String name : files) {
+ Path p = Paths.get(tmpFolder.getRoot().getAbsolutePath(), name);
+ File f = new File(p.toString());
+ assertFileContentsMatchInAnyOrder(
+ null, expectedHeader, f, "3,true,5", "4,true,6", "5,false,7");
+ }
+ }
+
+ @Test
+ public void rowsWithDefaults() throws IOException {
+ String to = "rows";
+ String regex = String.format("^%s.*$", to);
+ pipeline
+ .apply(Create.of(rows).withRowSchema(ALL_DATA_TYPES_SCHEMA))
+ .apply(CsvIO.writeRows().to(tmpFolder.getRoot().getAbsolutePath() + "/" + to));
+ pipeline.run().waitUntilFinish();
+ String[] files = filesMatching(regex);
+ assertTrue(
+ "CsvIO.writeRows should write to " + tmpFolder.getRoot().getAbsolutePath() + "/" + to,
+ files.length > 0);
+ String expectedHeader =
+ "aBoolean,aByte,aDouble,aFloat,aLong,aShort,anInt,dateTime,decimal,string";
+ for (String name : files) {
+ Pattern matchCSVSuffix = Pattern.compile("^.*csv$");
+ assertTrue(
+ String.format("%s should match %s", name, matchCSVSuffix.pattern()),
+ matchCSVSuffix.matcher(name).matches());
+ Path p = Paths.get(tmpFolder.getRoot().getAbsolutePath(), name);
+ File f = new File(p.toString());
+ assertFileContentsMatchInAnyOrder(
+ null,
+ expectedHeader,
+ f,
+ "true,0,0.1,0.2,5,3,4,1970-01-01T00:00:00.000Z,0,a",
+ "true,1,1.1,1.2,6,4,5,1970-01-01T00:00:00.001Z,1,b",
+ "false,2,2.0,2.1,7,5,6,1970-01-01T00:00:00.002Z,2,c");
+ }
+ }
+
+ @Test
+ public void userTypesWithDefaults() throws IOException {
+ String to = "user_types";
+ String regex = String.format("^%s.*$", to);
+ pipeline
+ .apply(Create.of(userTypes))
+ .apply(CsvIO.<AllDataTypes>write().to(tmpFolder.getRoot().getAbsolutePath() + "/" + to));
+ pipeline.run().waitUntilFinish();
+ String[] files = filesMatching(regex);
+ assertTrue(
+ "CsvIO.writeRows should write to " + tmpFolder.getRoot().getAbsolutePath() + "/" + to,
+ files.length > 0);
+ String expectedHeader =
+ "aBoolean,aByte,aDouble,aFloat,aLong,aShort,anInt,dateTime,decimal,string";
+ for (String name : files) {
+ Path p = Paths.get(tmpFolder.getRoot().getAbsolutePath(), name);
+ File f = new File(p.toString());
+ assertFileContentsMatchInAnyOrder(
+ null,
+ expectedHeader,
+ f,
+ "true,0,0.1,0.2,5,3,4,1970-01-01T00:00:00.000Z,0,a",
+ "true,1,1.1,1.2,6,4,5,1970-01-01T00:00:00.001Z,1,b",
+ "false,2,2.0,2.1,7,5,6,1970-01-01T00:00:00.002Z,2,c");
+ }
+ }
+
+ @Test
+ public void rowsWithNulls() throws IOException {
+ String to = "rows_with_nulls";
+ String regex = String.format("^%s.*$", to);
+ Schema schema =
+ Schema.of(
+ Field.of("string", FieldType.STRING),
+ Field.nullable("aDouble", FieldType.DOUBLE),
+ Field.nullable("anInt", FieldType.INT32));
+
+ List<Row> nullableRows =
+ Arrays.asList(
+ Row.withSchema(schema).withFieldValue("string", "a").build(),
+ Row.withSchema(schema)
+ .withFieldValue("string", "b")
+ .withFieldValue("aDouble", 1.2345)
+ .build(),
+ Row.withSchema(schema).withFieldValue("string", "c").withFieldValue("anInt", 2).build(),
+ Row.withSchema(schema)
+ .withFieldValue("string", "d")
+ .withFieldValue("aDouble", 3.1234)
+ .withFieldValue("anInt", 3)
+ .build());
+
+ pipeline
+ .apply(Create.of(nullableRows).withRowSchema(schema))
+ .apply(CsvIO.writeRows().to(tmpFolder.getRoot().getAbsolutePath() + "/" + to));
+
+ pipeline.run().waitUntilFinish();
+ String[] files = filesMatching(regex);
+ assertTrue(
+ "CsvIO.writeRows should write to " + tmpFolder.getRoot().getAbsolutePath() + "/" + to,
+ files.length > 0);
+ String expectedHeader = "aDouble,anInt,string";
+ for (String name : files) {
+ Path p = Paths.get(tmpFolder.getRoot().getAbsolutePath(), name);
+ File f = new File(p.toString());
+ assertFileContentsMatchInAnyOrder(
+ null, expectedHeader, f, ",,a", "1.2345,,b", ",2,c", "3.1234,3,d");
+ }
+ }
+
+ private static List<String> readLinesFromFile(File f) throws IOException {
+ List<String> currentFile = new ArrayList<>();
+ try (BufferedReader reader = Files.newBufferedReader(f.toPath(), Charsets.UTF_8)) {
+ while (true) {
+ String line = reader.readLine();
+ if (line == null) {
+ break;
+ }
+ currentFile.add(line);
+ }
+ }
+ return currentFile;
+ }
Review Comment:
Replace usage with https://guava.dev/releases/21.0/api/docs/com/google/common/io/Files.html#readLines-java.io.File-java.nio.charset.Charset-
--
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: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [beam] damondouglas commented on pull request #24630: [24469] Implement CsvIO.Write and supporting classes
Posted by GitBox <gi...@apache.org>.
damondouglas commented on PR #24630:
URL: https://github.com/apache/beam/pull/24630#issuecomment-1385654685
Run Python PreCommit
--
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: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [beam] damondouglas commented on pull request #24630: [24469] Implement CsvIO.Write and supporting classes
Posted by "damondouglas (via GitHub)" <gi...@apache.org>.
damondouglas commented on PR #24630:
URL: https://github.com/apache/beam/pull/24630#issuecomment-1402805586
Run Java PreCommit
--
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: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [beam] damondouglas commented on a diff in pull request #24630: [24469] Implement CsvIO.Write and supporting classes
Posted by "damondouglas (via GitHub)" <gi...@apache.org>.
damondouglas commented on code in PR #24630:
URL: https://github.com/apache/beam/pull/24630#discussion_r1085698752
##########
sdks/java/io/csv/src/main/java/org/apache/beam/sdk/io/csv/CsvRowConversions.java:
##########
@@ -0,0 +1,148 @@
+/*
+ * 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.sdk.io.csv;
+
+import static org.apache.beam.sdk.io.csv.CsvIO.VALID_FIELD_TYPE_SET;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkNotNull;
+
+import com.google.auto.value.AutoValue;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Set;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.schemas.Schema.TypeName;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.values.Row;
+import org.apache.commons.csv.CSVFormat;
+import org.checkerframework.checker.nullness.qual.NonNull;
+
+/** Contains classes and methods to help with converting between {@link Row} and CSV strings. */
+class CsvRowConversions {
+
+ /** Converts between {@link Row} and CSV string using a {@link CSVFormat}. */
+ @AutoValue
+ abstract static class RowToCsv implements SerializableFunction<Row, String> {
+
+ static Builder builder() {
+ return new AutoValue_CsvRowConversions_RowToCsv.Builder();
+ }
+
+ /** The expected {@link Schema} of the {@link Row} input. */
+ abstract Schema getSchema();
+
+ /** The {@link CSVFormat} of the converted {@link Row} input. */
+ abstract CSVFormat getCSVFormat();
+
+ /** Converts a {@link Row} to a CSV string formatted using {@link #getCSVFormat()}. */
+ @Override
+ public String apply(Row input) {
+ Row safeInput = checkNotNull(input);
+ String[] header = getHeader();
+ Object[] values = new Object[header.length];
+ for (int i = 0; i < header.length; i++) {
+ values[i] = safeInput.getValue(header[i]);
+ }
+ return getCSVFormat().format(values);
+ }
+
+ @NonNull
+ String[] getHeader() {
+ return checkNotNull(getCSVFormat().getHeader());
+ }
+
+ @AutoValue.Builder
+ abstract static class Builder {
+
+ /** The expected {@link Schema} of the {@link Row} input. */
+ abstract Builder setSchema(Schema schema);
+
+ abstract Schema getSchema();
+
+ /** The {@link CSVFormat} of the converted {@link Row} input. */
+ abstract Builder setCSVFormat(CSVFormat format);
+
+ abstract CSVFormat getCSVFormat();
+
+ abstract RowToCsv autoBuild();
+
+ final RowToCsv build() {
+ checkArgument(getSchema().getFieldCount() > 0, "Schema has no fields");
+ setCSVFormat(
+ getCSVFormat()
+ .withSkipHeaderRecord()
Review Comment:
I wanted to keep CsvRowConversions and CsvIO independent. In the setting of CsvRowConversions.RowToCSV didn't want to assume anyone would remember to set withSkipHeaderRecord. Also, for a package-private class, I didn't want to throw an error either for a class that has nothing to do with the header other than lookups against the Schema and Row values. The withSkipHeaderRecord is an output driving parameter and one designed by CSVFormat in the setting of a single file.
--
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: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [beam] damondouglas commented on pull request #24630: [24469] Implement CsvIO.Write and supporting classes
Posted by "damondouglas (via GitHub)" <gi...@apache.org>.
damondouglas commented on PR #24630:
URL: https://github.com/apache/beam/pull/24630#issuecomment-1405673452
Run Java PreCommit
--
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: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [beam] lukecwik commented on a diff in pull request #24630: [24469] Implement CsvIO.Write and supporting classes
Posted by GitBox <gi...@apache.org>.
lukecwik commented on code in PR #24630:
URL: https://github.com/apache/beam/pull/24630#discussion_r1050162885
##########
sdks/java/io/csv/src/main/java/org/apache/beam/sdk/io/csv/CsvUtils.java:
##########
@@ -0,0 +1,270 @@
+/*
+ * 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.sdk.io.csv;
+
+import java.io.IOException;
+import java.nio.charset.StandardCharsets;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Set;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.Field;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.SimpleFunction;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableSet;
+import org.apache.commons.csv.CSVFormat;
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+/** Utilities to convert between CSV records, Beam rows, and user types. */
+public final class CsvUtils {
Review Comment:
```suggestion
final class CsvUtils {
```
##########
sdks/java/io/csv/src/main/java/org/apache/beam/sdk/io/csv/CsvPayloadSerializer.java:
##########
@@ -0,0 +1,96 @@
+/*
+ * 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.sdk.io.csv;
+
+import java.io.IOException;
+import java.nio.charset.StandardCharsets;
+import java.util.List;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.io.payloads.PayloadSerializer;
+import org.apache.beam.sdk.values.Row;
+import org.apache.commons.csv.CSVFormat;
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+/** Implementation of {@link PayloadSerializer} for a CSV format. */
+public class CsvPayloadSerializer implements PayloadSerializer {
Review Comment:
```suggestion
class CsvPayloadSerializer implements PayloadSerializer {
```
--
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: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [beam] github-actions[bot] commented on pull request #24630: [24469] Implement CsvIO.Write and supporting classes
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on PR #24630:
URL: https://github.com/apache/beam/pull/24630#issuecomment-1398298973
Assigning new set of reviewers because Pr has gone too long without review. If you would like to opt out of this review, comment `assign to next reviewer`:
R: @robertwb for label java.
R: @Abacn for label build.
R: @Abacn for label io.
Available commands:
- `stop reviewer notifications` - opt out of the automated review tooling
- `remind me after tests pass` - tag the comment author after tests pass
- `waiting on author` - shift the attention set back to the author (any comment or push by the author will return the attention set to the reviewers)
--
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: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [beam] damondouglas commented on pull request #24630: [24469] Implement CsvIO.Write and supporting classes
Posted by GitBox <gi...@apache.org>.
damondouglas commented on PR #24630:
URL: https://github.com/apache/beam/pull/24630#issuecomment-1386037805
Run Python_Runners PreCommit
--
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: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [beam] damondouglas commented on pull request #24630: [24469] Implement CsvIO.Write and supporting classes
Posted by GitBox <gi...@apache.org>.
damondouglas commented on PR #24630:
URL: https://github.com/apache/beam/pull/24630#issuecomment-1351882696
> Note that WriteFiles supports a lot of additional properties like dynamic destinations, no spilling, ... that users typically want. They are pretty easy to add since they are just boilerplate builder stuff. Your call if you want to add them now or later.
@lukecwik I'm glad you suggested this as it inspired me to remove fields shared between CsvIO.Write and FileIO.Write and instead just maintain an inner instance of FileIO.Write. For example, CsvIO.Write's withCompression calls the inner instance of FileIO.Write's withCompression instead of maintaining a redundant compression property.
--
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: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [beam] lukecwik commented on a diff in pull request #24630: [24469] Implement CsvIO.Write and supporting classes
Posted by GitBox <gi...@apache.org>.
lukecwik commented on code in PR #24630:
URL: https://github.com/apache/beam/pull/24630#discussion_r1046193793
##########
sdks/java/io/csv/src/main/java/org/apache/beam/sdk/io/csv/CsvIO.java:
##########
@@ -0,0 +1,384 @@
+/*
+ * 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.sdk.io.csv;
+
+import static java.nio.charset.StandardCharsets.UTF_8;
+import static org.apache.beam.sdk.values.TypeDescriptors.rows;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+
+import com.google.auto.value.AutoValue;
+import java.io.BufferedWriter;
+import java.io.IOException;
+import java.io.OutputStreamWriter;
+import java.io.PrintWriter;
+import java.nio.channels.Channels;
+import java.nio.channels.WritableByteChannel;
+import java.util.List;
+import java.util.Objects;
+import java.util.Optional;
+import org.apache.beam.sdk.io.Compression;
+import org.apache.beam.sdk.io.FileIO;
+import org.apache.beam.sdk.io.FileSystems;
+import org.apache.beam.sdk.io.fs.ResourceId;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.transforms.MapElements;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.SimpleFunction;
+import org.apache.beam.sdk.transforms.display.DisplayData;
+import org.apache.beam.sdk.transforms.display.HasDisplayData;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PDone;
+import org.apache.beam.sdk.values.Row;
+import org.apache.commons.csv.CSVFormat;
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+/** Transforms for reading and writing CSV files. */
+@SuppressWarnings({
+ "unused",
+ "nullness" // TODO(https://github.com/apache/beam/issues/20497)
+})
+public class CsvIO {
+
+ static final String DEFAULT_FILENAME_SUFFIX = ".csv";
+
+ /** Instantiates a {@link Write} for writing user types in CSV format. */
+ public static <T> Write<T> write() {
+ return new AutoValue_CsvIO_Write.Builder<T>().build();
+ }
+
+ /** Instantiates a {@link Write} for {@link Row}s in CSV format. */
+ public static Write<Row> writeRows() {
+ return new AutoValue_CsvIO_Write.Builder<Row>().build();
+ }
+
+ /** Implementation of {@link FileIO.Sink}. */
+ @AutoValue
+ public abstract static class Sink<T> implements FileIO.Sink<T> {
+
+ public static <T> Builder<T> builder() {
+ return new AutoValue_CsvIO_Sink.Builder<>();
+ }
+
+ /**
+ * Not to be confused with the CSV header, it is content written to the top of every sharded
+ * file prior to the header. In the example below, all the text proceeding the header
+ * 'column1,column2,column3' is the preamble.
+ *
+ * <p>Fake company, Inc. Lab experiment: abcdefg123456 Experiment date: 2022-12-05 Operator:
+ * John Doe
+ *
+ * <p>column1,column2,column3 1,2,3 4,5,6
+ */
+ public Sink<T> withPreamble(String preamble) {
+ return toBuilder().setPreamble(preamble).build();
+ }
+
+ private transient @Nullable PrintWriter writer;
+
+ /**
+ * Opens a {@link WritableByteChannel} for writing CSV files. Writes the {@link #getPreamble()}
+ * if available followed by the {@link #getHeader()}.
+ */
+ @Override
+ public void open(WritableByteChannel channel) throws IOException {
+ writer =
+ new PrintWriter(
+ new BufferedWriter(new OutputStreamWriter(Channels.newOutputStream(channel), UTF_8)));
+ if (getPreamble() != null) {
+ writer.println(getPreamble());
+ }
+ writer.println(getHeader());
+ }
+
+ /** Serializes and writes the {@param element} to a file. */
+ @Override
+ public void write(T element) throws IOException {
+ if (writer == null) {
+ throw new IllegalStateException(
+ String.format("%s writer is null", PrintWriter.class.getName()));
+ }
+ String line = getFormatFunction().apply(element);
+ writer.println(line);
+ }
+
+ @Override
+ public void flush() throws IOException {
+ if (writer == null) {
+ throw new IllegalStateException(
+ String.format("%s writer is null", PrintWriter.class.getName()));
+ }
+ writer.flush();
+ }
+
+ /**
+ * Not to be confused with the CSV header, it is content written to the top of every sharded
+ * file prior to the header. In the example below, all the text proceeding the header
+ * 'column1,column2,column3' is the preamble.
+ *
+ * <p>Fake company, Inc. Lab experiment: abcdefg123456 Experiment date: 2022-12-05 Operator:
+ * John Doe
+ *
+ * <p>column1,column2,column3 1,2,3 4,5,6
+ */
Review Comment:
Ditto for `<pre>{@code ...`
##########
sdks/java/io/csv/src/main/java/org/apache/beam/sdk/io/csv/CsvIO.java:
##########
@@ -0,0 +1,384 @@
+/*
+ * 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.sdk.io.csv;
+
+import static java.nio.charset.StandardCharsets.UTF_8;
+import static org.apache.beam.sdk.values.TypeDescriptors.rows;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+
+import com.google.auto.value.AutoValue;
+import java.io.BufferedWriter;
+import java.io.IOException;
+import java.io.OutputStreamWriter;
+import java.io.PrintWriter;
+import java.nio.channels.Channels;
+import java.nio.channels.WritableByteChannel;
+import java.util.List;
+import java.util.Objects;
+import java.util.Optional;
+import org.apache.beam.sdk.io.Compression;
+import org.apache.beam.sdk.io.FileIO;
+import org.apache.beam.sdk.io.FileSystems;
+import org.apache.beam.sdk.io.fs.ResourceId;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.transforms.MapElements;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.SimpleFunction;
+import org.apache.beam.sdk.transforms.display.DisplayData;
+import org.apache.beam.sdk.transforms.display.HasDisplayData;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PDone;
+import org.apache.beam.sdk.values.Row;
+import org.apache.commons.csv.CSVFormat;
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+/** Transforms for reading and writing CSV files. */
+@SuppressWarnings({
+ "unused",
+ "nullness" // TODO(https://github.com/apache/beam/issues/20497)
Review Comment:
Could we resolve the nullness concerns in new code?
##########
sdks/java/io/csv/src/main/java/org/apache/beam/sdk/io/csv/CsvIO.java:
##########
@@ -0,0 +1,384 @@
+/*
+ * 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.sdk.io.csv;
+
+import static java.nio.charset.StandardCharsets.UTF_8;
+import static org.apache.beam.sdk.values.TypeDescriptors.rows;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+
+import com.google.auto.value.AutoValue;
+import java.io.BufferedWriter;
+import java.io.IOException;
+import java.io.OutputStreamWriter;
+import java.io.PrintWriter;
+import java.nio.channels.Channels;
+import java.nio.channels.WritableByteChannel;
+import java.util.List;
+import java.util.Objects;
+import java.util.Optional;
+import org.apache.beam.sdk.io.Compression;
+import org.apache.beam.sdk.io.FileIO;
+import org.apache.beam.sdk.io.FileSystems;
+import org.apache.beam.sdk.io.fs.ResourceId;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.transforms.MapElements;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.SimpleFunction;
+import org.apache.beam.sdk.transforms.display.DisplayData;
+import org.apache.beam.sdk.transforms.display.HasDisplayData;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PDone;
+import org.apache.beam.sdk.values.Row;
+import org.apache.commons.csv.CSVFormat;
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+/** Transforms for reading and writing CSV files. */
+@SuppressWarnings({
+ "unused",
+ "nullness" // TODO(https://github.com/apache/beam/issues/20497)
+})
+public class CsvIO {
+
+ static final String DEFAULT_FILENAME_SUFFIX = ".csv";
+
+ /** Instantiates a {@link Write} for writing user types in CSV format. */
+ public static <T> Write<T> write() {
+ return new AutoValue_CsvIO_Write.Builder<T>().build();
+ }
+
+ /** Instantiates a {@link Write} for {@link Row}s in CSV format. */
+ public static Write<Row> writeRows() {
+ return new AutoValue_CsvIO_Write.Builder<Row>().build();
+ }
+
+ /** Implementation of {@link FileIO.Sink}. */
+ @AutoValue
+ public abstract static class Sink<T> implements FileIO.Sink<T> {
+
+ public static <T> Builder<T> builder() {
+ return new AutoValue_CsvIO_Sink.Builder<>();
+ }
+
+ /**
+ * Not to be confused with the CSV header, it is content written to the top of every sharded
+ * file prior to the header. In the example below, all the text proceeding the header
+ * 'column1,column2,column3' is the preamble.
+ *
+ * <p>Fake company, Inc. Lab experiment: abcdefg123456 Experiment date: 2022-12-05 Operator:
+ * John Doe
+ *
+ * <p>column1,column2,column3 1,2,3 4,5,6
+ */
+ public Sink<T> withPreamble(String preamble) {
+ return toBuilder().setPreamble(preamble).build();
+ }
+
+ private transient @Nullable PrintWriter writer;
+
+ /**
+ * Opens a {@link WritableByteChannel} for writing CSV files. Writes the {@link #getPreamble()}
+ * if available followed by the {@link #getHeader()}.
+ */
+ @Override
+ public void open(WritableByteChannel channel) throws IOException {
+ writer =
+ new PrintWriter(
+ new BufferedWriter(new OutputStreamWriter(Channels.newOutputStream(channel), UTF_8)));
+ if (getPreamble() != null) {
+ writer.println(getPreamble());
+ }
+ writer.println(getHeader());
+ }
+
+ /** Serializes and writes the {@param element} to a file. */
+ @Override
+ public void write(T element) throws IOException {
+ if (writer == null) {
+ throw new IllegalStateException(
+ String.format("%s writer is null", PrintWriter.class.getName()));
+ }
+ String line = getFormatFunction().apply(element);
+ writer.println(line);
+ }
+
+ @Override
+ public void flush() throws IOException {
+ if (writer == null) {
+ throw new IllegalStateException(
+ String.format("%s writer is null", PrintWriter.class.getName()));
+ }
+ writer.flush();
+ }
+
+ /**
+ * Not to be confused with the CSV header, it is content written to the top of every sharded
+ * file prior to the header. In the example below, all the text proceeding the header
+ * 'column1,column2,column3' is the preamble.
+ *
+ * <p>Fake company, Inc. Lab experiment: abcdefg123456 Experiment date: 2022-12-05 Operator:
+ * John Doe
+ *
+ * <p>column1,column2,column3 1,2,3 4,5,6
+ */
+ abstract @Nullable String getPreamble();
+
+ /**
+ * The column names of the CSV file written at the top line of each shard after the preamble, if
+ * available.
+ */
+ abstract String getHeader();
+
+ /** A {@link SimpleFunction} for converting a {@param T} to a CSV formatted string. */
+ abstract SimpleFunction<T, String> getFormatFunction();
+
+ abstract Builder<T> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<T> {
+
+ /**
+ * Not to be confused with the CSV header, it is content written to the top of every sharded
+ * file prior to the header. In the example below, all the text proceeding the header
+ * 'column1,column2,column3' is the preamble.
+ *
+ * <p>Fake company, Inc. Lab experiment: abcdefg123456 Experiment date: 2022-12-05 Operator:
+ * John Doe
+ *
+ * <p>column1,column2,column3 1,2,3 4,5,6
+ */
+ abstract Builder<T> setPreamble(String value);
+
+ /**
+ * The column names of the CSV file written at the top line of each shard after the preamble,
+ * if available.
+ */
+ abstract Builder<T> setHeader(String value);
+
+ /** A {@link SimpleFunction} for converting a {@param T} to a CSV formatted string. */
+ abstract Builder<T> setFormatFunction(SimpleFunction<T, String> value);
+
+ abstract Sink<T> build();
+ }
+ }
+
+ @AutoValue
+ public abstract static class Write<T> extends PTransform<PCollection<T>, PDone>
+ implements HasDisplayData {
+
+ /** Specifies a common prefix for all generated files. */
+ public Write<T> to(String filenamePrefix) {
+ return toBuilder().setFilenamePrefix(filenamePrefix).build();
+ }
+
+ /** The {@link CSVFormat} of the destination CSV file data. */
+ public Write<T> withCSVFormat(CSVFormat format) {
+ return toBuilder().setCSVFormat(format).build();
+ }
+
+ /**
+ * Not to be confused with the CSV header, it is content written to the top of every sharded
+ * file prior to the header. In the example below, all the text proceeding the header
+ * 'column1,column2,column3' is the preamble.
+ *
+ * <p>Fake company, Inc. Lab experiment: abcdefg123456 Experiment date: 2022-12-05 Operator:
Review Comment:
`<pre>{@code....`
##########
sdks/java/io/csv/src/main/java/org/apache/beam/sdk/io/csv/CsvIO.java:
##########
@@ -0,0 +1,384 @@
+/*
+ * 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.sdk.io.csv;
+
+import static java.nio.charset.StandardCharsets.UTF_8;
+import static org.apache.beam.sdk.values.TypeDescriptors.rows;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+
+import com.google.auto.value.AutoValue;
+import java.io.BufferedWriter;
+import java.io.IOException;
+import java.io.OutputStreamWriter;
+import java.io.PrintWriter;
+import java.nio.channels.Channels;
+import java.nio.channels.WritableByteChannel;
+import java.util.List;
+import java.util.Objects;
+import java.util.Optional;
+import org.apache.beam.sdk.io.Compression;
+import org.apache.beam.sdk.io.FileIO;
+import org.apache.beam.sdk.io.FileSystems;
+import org.apache.beam.sdk.io.fs.ResourceId;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.transforms.MapElements;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.SimpleFunction;
+import org.apache.beam.sdk.transforms.display.DisplayData;
+import org.apache.beam.sdk.transforms.display.HasDisplayData;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PDone;
+import org.apache.beam.sdk.values.Row;
+import org.apache.commons.csv.CSVFormat;
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+/** Transforms for reading and writing CSV files. */
+@SuppressWarnings({
+ "unused",
+ "nullness" // TODO(https://github.com/apache/beam/issues/20497)
+})
+public class CsvIO {
+
+ static final String DEFAULT_FILENAME_SUFFIX = ".csv";
+
+ /** Instantiates a {@link Write} for writing user types in CSV format. */
+ public static <T> Write<T> write() {
+ return new AutoValue_CsvIO_Write.Builder<T>().build();
+ }
+
+ /** Instantiates a {@link Write} for {@link Row}s in CSV format. */
+ public static Write<Row> writeRows() {
+ return new AutoValue_CsvIO_Write.Builder<Row>().build();
+ }
+
+ /** Implementation of {@link FileIO.Sink}. */
+ @AutoValue
+ public abstract static class Sink<T> implements FileIO.Sink<T> {
+
+ public static <T> Builder<T> builder() {
+ return new AutoValue_CsvIO_Sink.Builder<>();
+ }
+
+ /**
+ * Not to be confused with the CSV header, it is content written to the top of every sharded
+ * file prior to the header. In the example below, all the text proceeding the header
+ * 'column1,column2,column3' is the preamble.
+ *
+ * <p>Fake company, Inc. Lab experiment: abcdefg123456 Experiment date: 2022-12-05 Operator:
+ * John Doe
+ *
+ * <p>column1,column2,column3 1,2,3 4,5,6
+ */
+ public Sink<T> withPreamble(String preamble) {
+ return toBuilder().setPreamble(preamble).build();
+ }
+
+ private transient @Nullable PrintWriter writer;
+
+ /**
+ * Opens a {@link WritableByteChannel} for writing CSV files. Writes the {@link #getPreamble()}
+ * if available followed by the {@link #getHeader()}.
+ */
+ @Override
+ public void open(WritableByteChannel channel) throws IOException {
+ writer =
+ new PrintWriter(
+ new BufferedWriter(new OutputStreamWriter(Channels.newOutputStream(channel), UTF_8)));
+ if (getPreamble() != null) {
+ writer.println(getPreamble());
+ }
+ writer.println(getHeader());
+ }
+
+ /** Serializes and writes the {@param element} to a file. */
+ @Override
+ public void write(T element) throws IOException {
+ if (writer == null) {
+ throw new IllegalStateException(
+ String.format("%s writer is null", PrintWriter.class.getName()));
+ }
+ String line = getFormatFunction().apply(element);
+ writer.println(line);
+ }
+
+ @Override
+ public void flush() throws IOException {
+ if (writer == null) {
+ throw new IllegalStateException(
+ String.format("%s writer is null", PrintWriter.class.getName()));
+ }
+ writer.flush();
+ }
+
+ /**
+ * Not to be confused with the CSV header, it is content written to the top of every sharded
+ * file prior to the header. In the example below, all the text proceeding the header
+ * 'column1,column2,column3' is the preamble.
+ *
+ * <p>Fake company, Inc. Lab experiment: abcdefg123456 Experiment date: 2022-12-05 Operator:
+ * John Doe
+ *
+ * <p>column1,column2,column3 1,2,3 4,5,6
+ */
+ abstract @Nullable String getPreamble();
+
+ /**
+ * The column names of the CSV file written at the top line of each shard after the preamble, if
+ * available.
+ */
+ abstract String getHeader();
+
+ /** A {@link SimpleFunction} for converting a {@param T} to a CSV formatted string. */
+ abstract SimpleFunction<T, String> getFormatFunction();
+
+ abstract Builder<T> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<T> {
+
+ /**
+ * Not to be confused with the CSV header, it is content written to the top of every sharded
+ * file prior to the header. In the example below, all the text proceeding the header
+ * 'column1,column2,column3' is the preamble.
+ *
+ * <p>Fake company, Inc. Lab experiment: abcdefg123456 Experiment date: 2022-12-05 Operator:
+ * John Doe
+ *
+ * <p>column1,column2,column3 1,2,3 4,5,6
+ */
+ abstract Builder<T> setPreamble(String value);
+
+ /**
+ * The column names of the CSV file written at the top line of each shard after the preamble,
+ * if available.
+ */
+ abstract Builder<T> setHeader(String value);
+
+ /** A {@link SimpleFunction} for converting a {@param T} to a CSV formatted string. */
+ abstract Builder<T> setFormatFunction(SimpleFunction<T, String> value);
+
+ abstract Sink<T> build();
+ }
+ }
+
+ @AutoValue
+ public abstract static class Write<T> extends PTransform<PCollection<T>, PDone>
+ implements HasDisplayData {
+
+ /** Specifies a common prefix for all generated files. */
+ public Write<T> to(String filenamePrefix) {
+ return toBuilder().setFilenamePrefix(filenamePrefix).build();
+ }
+
+ /** The {@link CSVFormat} of the destination CSV file data. */
+ public Write<T> withCSVFormat(CSVFormat format) {
+ return toBuilder().setCSVFormat(format).build();
+ }
+
+ /**
+ * Not to be confused with the CSV header, it is content written to the top of every sharded
+ * file prior to the header. In the example below, all the text proceeding the header
+ * 'column1,column2,column3' is the preamble.
+ *
+ * <p>Fake company, Inc. Lab experiment: abcdefg123456 Experiment date: 2022-12-05 Operator:
+ * John Doe
+ *
+ * <p>column1,column2,column3 1,2,3 4,5,6
+ */
+ public Write<T> withPreamble(String preamble) {
+ return toBuilder().setPreamble(preamble).build();
+ }
+
+ /** Specifies the {@link Compression} of all generated shard files. */
+ public Write<T> withCompression(Compression compression) {
+ return toBuilder().setCompression(compression).build();
+ }
+
+ /**
+ * Specifies to use a given fixed number of shards per window. See {@link
+ * FileIO.Write#withNumShards(int)} for details.
+ */
+ public Write<T> withNumShards(Integer numShards) {
+ return toBuilder().setNumShards(numShards).build();
+ }
+
+ /**
+ * Specifies a directory into which all temporary files will be placed. See {@link
+ * FileIO.Write#withTempDirectory(String)}.
+ */
+ public Write<T> withTempDirectory(ResourceId value) {
+ return toBuilder().setTempDirectory(value).build();
+ }
+
+ abstract @Nullable String getFilenamePrefix();
+
+ abstract CSVFormat getCSVFormat();
+
+ abstract @Nullable String getPreamble();
+
+ abstract @Nullable Compression getCompression();
+
+ abstract @Nullable Integer getNumShards();
+
+ abstract String getFilenameSuffix();
+
+ abstract @Nullable ResourceId getTempDirectory();
+
+ abstract @Nullable List<String> getSchemaFields();
+
+ abstract Builder<T> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<T> {
+
+ abstract Builder<T> setFilenamePrefix(String value);
+
+ abstract Builder<T> setCSVFormat(CSVFormat value);
+
+ abstract Optional<CSVFormat> getCSVFormat();
+
+ abstract Builder<T> setPreamble(String value);
+
+ abstract Builder<T> setCompression(Compression value);
+
+ abstract Builder<T> setNumShards(Integer value);
+
+ abstract Builder<T> setFilenameSuffix(String value);
+
+ abstract Optional<String> getFilenameSuffix();
+
+ abstract Builder<T> setTempDirectory(ResourceId value);
+
+ abstract Builder<T> setSchemaFields(List<String> value);
+
+ abstract Write<T> autoBuild();
+
+ final Write<T> build() {
+ if (!getCSVFormat().isPresent()) {
+ setCSVFormat(CSVFormat.DEFAULT);
+ }
+
+ if (!getFilenameSuffix().isPresent()) {
+ setFilenameSuffix(DEFAULT_FILENAME_SUFFIX);
+ }
+
+ return autoBuild();
+ }
+ }
+
+ @Override
+ public void populateDisplayData(DisplayData.Builder builder) {
+ if (getFilenamePrefix() != null) {
+ builder.add(DisplayData.item("filenamePrefix", getFilenamePrefix()));
+ }
+ if (getCSVFormat() != null) {
+ builder.add(DisplayData.item("csvFormat", getCSVFormat().toString()));
+ }
+ if (getPreamble() != null) {
+ builder.add(DisplayData.item("preamble", getPreamble()));
+ }
+ if (getCompression() != null) {
+ builder.add(DisplayData.item("compression", getCompression().name()));
+ }
+ if (getNumShards() != null) {
+ builder.add(DisplayData.item("numShards", getNumShards()));
+ }
+
+ builder.add(DisplayData.item("filenameSuffix", getFilenameSuffix()));
+
+ if (getTempDirectory() != null) {
+ builder.add(DisplayData.item("tempDirectory", getTempDirectory().getFilename()));
+ }
+ }
+
+ @Override
+ public PDone expand(PCollection<T> input) {
+ if (!input.hasSchema()) {
+ throw new IllegalArgumentException(
+ String.format(
+ "%s requires an input Schema. Note that only Row or user classes are supported. Consider using TextIO or FileIO directly when writing primitive types",
+ Write.class.getName()));
+ }
+
+ Schema schema = input.getSchema();
+ SerializableFunction<T, Row> toRowFn = input.getToRowFunction();
+ PCollection<Row> rows =
+ input.apply(MapElements.into(rows()).via(toRowFn)).setRowSchema(schema);
+ rows.apply("writeRowsToCsv", buildFileIOWrite().via(buildSink(schema)));
Review Comment:
```suggestion
rows.apply("Write Rows To Csv", buildFileIOWrite().via(buildSink(schema)));
```
##########
sdks/java/io/csv/src/main/java/org/apache/beam/sdk/io/csv/CsvPayloadSerializer.java:
##########
@@ -0,0 +1,97 @@
+/*
+ * 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.sdk.io.csv;
+
+import java.io.IOException;
+import java.nio.charset.StandardCharsets;
+import java.util.List;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.io.payloads.PayloadSerializer;
+import org.apache.beam.sdk.values.Row;
+import org.apache.commons.csv.CSVFormat;
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+/** Implementation of {@link PayloadSerializer} for a CSV format. */
+public class CsvPayloadSerializer implements PayloadSerializer {
+
+ private final CSVFormat csvFormat;
+ private final List<String> schemaFields;
+
+ /**
+ * Instantiates a {@link PayloadSerializer} for a CSV format. {@param csvFormat} defaults to
+ * {@link CSVFormat#DEFAULT} and {@param schemaFields} defaults to {@link Schema#sorted()} {@link
+ * Schema#getFieldNames()}.
+ */
+ CsvPayloadSerializer(
+ Schema schema, @Nullable CSVFormat csvFormat, @Nullable List<String> schemaFields) {
+ CsvUtils.validateSchema(schema);
+ CsvUtils.validateHeaderAgainstSchema(schemaFields, schema);
+ if (csvFormat == null) {
+ csvFormat = CSVFormat.DEFAULT;
+ }
+ this.csvFormat = csvFormat;
+
+ if (schemaFields == null) {
+ schemaFields = schema.sorted().getFieldNames();
+ }
+ this.schemaFields = schemaFields;
+ }
+
+ CSVFormat getCsvFormat() {
+ return csvFormat;
+ }
+
+ List<String> getSchemaFields() {
+ return schemaFields;
+ }
+
+ /**
+ * Serializes a {@link Row} to a CSV byte[] record using a {@link CSVFormat}. The schemaFields
+ * provided or generated by the constructor drive the subset and order of serialized {@link Row}
+ * fields.
+ */
+ @Override
+ public byte[] serialize(Row row) {
+ StringBuilder builder = new StringBuilder();
+ try {
+ boolean newRecord = true;
+ for (String name : schemaFields) {
+ Object value = row.getValue(name);
+ if (value == null) {
Review Comment:
I would have expected that the CSVFormat converted nulls based upon https://commons.apache.org/proper/commons-csv/apidocs/org/apache/commons/csv/CSVFormat.Builder.html#setNullString-java.lang.String-
##########
sdks/java/io/csv/src/main/java/org/apache/beam/sdk/io/csv/CsvIO.java:
##########
@@ -0,0 +1,384 @@
+/*
+ * 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.sdk.io.csv;
+
+import static java.nio.charset.StandardCharsets.UTF_8;
+import static org.apache.beam.sdk.values.TypeDescriptors.rows;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+
+import com.google.auto.value.AutoValue;
+import java.io.BufferedWriter;
+import java.io.IOException;
+import java.io.OutputStreamWriter;
+import java.io.PrintWriter;
+import java.nio.channels.Channels;
+import java.nio.channels.WritableByteChannel;
+import java.util.List;
+import java.util.Objects;
+import java.util.Optional;
+import org.apache.beam.sdk.io.Compression;
+import org.apache.beam.sdk.io.FileIO;
+import org.apache.beam.sdk.io.FileSystems;
+import org.apache.beam.sdk.io.fs.ResourceId;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.transforms.MapElements;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.SimpleFunction;
+import org.apache.beam.sdk.transforms.display.DisplayData;
+import org.apache.beam.sdk.transforms.display.HasDisplayData;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PDone;
+import org.apache.beam.sdk.values.Row;
+import org.apache.commons.csv.CSVFormat;
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+/** Transforms for reading and writing CSV files. */
+@SuppressWarnings({
+ "unused",
+ "nullness" // TODO(https://github.com/apache/beam/issues/20497)
+})
+public class CsvIO {
+
+ static final String DEFAULT_FILENAME_SUFFIX = ".csv";
+
+ /** Instantiates a {@link Write} for writing user types in CSV format. */
+ public static <T> Write<T> write() {
+ return new AutoValue_CsvIO_Write.Builder<T>().build();
+ }
+
+ /** Instantiates a {@link Write} for {@link Row}s in CSV format. */
+ public static Write<Row> writeRows() {
+ return new AutoValue_CsvIO_Write.Builder<Row>().build();
+ }
+
+ /** Implementation of {@link FileIO.Sink}. */
+ @AutoValue
+ public abstract static class Sink<T> implements FileIO.Sink<T> {
+
+ public static <T> Builder<T> builder() {
+ return new AutoValue_CsvIO_Sink.Builder<>();
+ }
+
+ /**
+ * Not to be confused with the CSV header, it is content written to the top of every sharded
+ * file prior to the header. In the example below, all the text proceeding the header
+ * 'column1,column2,column3' is the preamble.
+ *
+ * <p>Fake company, Inc. Lab experiment: abcdefg123456 Experiment date: 2022-12-05 Operator:
+ * John Doe
+ *
+ * <p>column1,column2,column3 1,2,3 4,5,6
+ */
+ public Sink<T> withPreamble(String preamble) {
+ return toBuilder().setPreamble(preamble).build();
+ }
+
+ private transient @Nullable PrintWriter writer;
+
+ /**
+ * Opens a {@link WritableByteChannel} for writing CSV files. Writes the {@link #getPreamble()}
+ * if available followed by the {@link #getHeader()}.
+ */
+ @Override
+ public void open(WritableByteChannel channel) throws IOException {
+ writer =
+ new PrintWriter(
+ new BufferedWriter(new OutputStreamWriter(Channels.newOutputStream(channel), UTF_8)));
+ if (getPreamble() != null) {
+ writer.println(getPreamble());
+ }
+ writer.println(getHeader());
+ }
+
+ /** Serializes and writes the {@param element} to a file. */
+ @Override
+ public void write(T element) throws IOException {
+ if (writer == null) {
+ throw new IllegalStateException(
+ String.format("%s writer is null", PrintWriter.class.getName()));
+ }
+ String line = getFormatFunction().apply(element);
+ writer.println(line);
+ }
+
+ @Override
+ public void flush() throws IOException {
+ if (writer == null) {
+ throw new IllegalStateException(
+ String.format("%s writer is null", PrintWriter.class.getName()));
+ }
+ writer.flush();
+ }
+
+ /**
+ * Not to be confused with the CSV header, it is content written to the top of every sharded
+ * file prior to the header. In the example below, all the text proceeding the header
+ * 'column1,column2,column3' is the preamble.
+ *
+ * <p>Fake company, Inc. Lab experiment: abcdefg123456 Experiment date: 2022-12-05 Operator:
+ * John Doe
+ *
+ * <p>column1,column2,column3 1,2,3 4,5,6
+ */
+ abstract @Nullable String getPreamble();
+
+ /**
+ * The column names of the CSV file written at the top line of each shard after the preamble, if
+ * available.
+ */
+ abstract String getHeader();
+
+ /** A {@link SimpleFunction} for converting a {@param T} to a CSV formatted string. */
+ abstract SimpleFunction<T, String> getFormatFunction();
+
+ abstract Builder<T> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<T> {
+
+ /**
+ * Not to be confused with the CSV header, it is content written to the top of every sharded
+ * file prior to the header. In the example below, all the text proceeding the header
+ * 'column1,column2,column3' is the preamble.
+ *
+ * <p>Fake company, Inc. Lab experiment: abcdefg123456 Experiment date: 2022-12-05 Operator:
+ * John Doe
+ *
+ * <p>column1,column2,column3 1,2,3 4,5,6
+ */
+ abstract Builder<T> setPreamble(String value);
+
+ /**
+ * The column names of the CSV file written at the top line of each shard after the preamble,
+ * if available.
+ */
+ abstract Builder<T> setHeader(String value);
+
+ /** A {@link SimpleFunction} for converting a {@param T} to a CSV formatted string. */
+ abstract Builder<T> setFormatFunction(SimpleFunction<T, String> value);
+
+ abstract Sink<T> build();
+ }
+ }
+
+ @AutoValue
+ public abstract static class Write<T> extends PTransform<PCollection<T>, PDone>
+ implements HasDisplayData {
+
+ /** Specifies a common prefix for all generated files. */
+ public Write<T> to(String filenamePrefix) {
+ return toBuilder().setFilenamePrefix(filenamePrefix).build();
+ }
+
+ /** The {@link CSVFormat} of the destination CSV file data. */
+ public Write<T> withCSVFormat(CSVFormat format) {
+ return toBuilder().setCSVFormat(format).build();
+ }
+
+ /**
+ * Not to be confused with the CSV header, it is content written to the top of every sharded
+ * file prior to the header. In the example below, all the text proceeding the header
+ * 'column1,column2,column3' is the preamble.
+ *
+ * <p>Fake company, Inc. Lab experiment: abcdefg123456 Experiment date: 2022-12-05 Operator:
+ * John Doe
+ *
+ * <p>column1,column2,column3 1,2,3 4,5,6
+ */
+ public Write<T> withPreamble(String preamble) {
+ return toBuilder().setPreamble(preamble).build();
+ }
+
+ /** Specifies the {@link Compression} of all generated shard files. */
+ public Write<T> withCompression(Compression compression) {
+ return toBuilder().setCompression(compression).build();
+ }
+
+ /**
+ * Specifies to use a given fixed number of shards per window. See {@link
+ * FileIO.Write#withNumShards(int)} for details.
+ */
+ public Write<T> withNumShards(Integer numShards) {
+ return toBuilder().setNumShards(numShards).build();
+ }
+
+ /**
+ * Specifies a directory into which all temporary files will be placed. See {@link
+ * FileIO.Write#withTempDirectory(String)}.
+ */
+ public Write<T> withTempDirectory(ResourceId value) {
+ return toBuilder().setTempDirectory(value).build();
+ }
+
+ abstract @Nullable String getFilenamePrefix();
+
+ abstract CSVFormat getCSVFormat();
+
+ abstract @Nullable String getPreamble();
+
+ abstract @Nullable Compression getCompression();
+
+ abstract @Nullable Integer getNumShards();
+
+ abstract String getFilenameSuffix();
+
+ abstract @Nullable ResourceId getTempDirectory();
+
+ abstract @Nullable List<String> getSchemaFields();
+
+ abstract Builder<T> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<T> {
+
+ abstract Builder<T> setFilenamePrefix(String value);
+
+ abstract Builder<T> setCSVFormat(CSVFormat value);
+
+ abstract Optional<CSVFormat> getCSVFormat();
+
+ abstract Builder<T> setPreamble(String value);
+
+ abstract Builder<T> setCompression(Compression value);
+
+ abstract Builder<T> setNumShards(Integer value);
+
+ abstract Builder<T> setFilenameSuffix(String value);
+
+ abstract Optional<String> getFilenameSuffix();
+
+ abstract Builder<T> setTempDirectory(ResourceId value);
+
+ abstract Builder<T> setSchemaFields(List<String> value);
+
+ abstract Write<T> autoBuild();
+
+ final Write<T> build() {
+ if (!getCSVFormat().isPresent()) {
+ setCSVFormat(CSVFormat.DEFAULT);
+ }
+
+ if (!getFilenameSuffix().isPresent()) {
+ setFilenameSuffix(DEFAULT_FILENAME_SUFFIX);
Review Comment:
Ditto for filename suffix to always start by setting our default and then allowing users to override it if they want.
##########
sdks/java/io/csv/src/main/java/org/apache/beam/sdk/io/csv/CsvPayloadSerializer.java:
##########
@@ -0,0 +1,97 @@
+/*
+ * 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.sdk.io.csv;
+
+import java.io.IOException;
+import java.nio.charset.StandardCharsets;
+import java.util.List;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.io.payloads.PayloadSerializer;
+import org.apache.beam.sdk.values.Row;
+import org.apache.commons.csv.CSVFormat;
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+/** Implementation of {@link PayloadSerializer} for a CSV format. */
+public class CsvPayloadSerializer implements PayloadSerializer {
+
+ private final CSVFormat csvFormat;
+ private final List<String> schemaFields;
+
+ /**
+ * Instantiates a {@link PayloadSerializer} for a CSV format. {@param csvFormat} defaults to
+ * {@link CSVFormat#DEFAULT} and {@param schemaFields} defaults to {@link Schema#sorted()} {@link
+ * Schema#getFieldNames()}.
+ */
+ CsvPayloadSerializer(
+ Schema schema, @Nullable CSVFormat csvFormat, @Nullable List<String> schemaFields) {
+ CsvUtils.validateSchema(schema);
+ CsvUtils.validateHeaderAgainstSchema(schemaFields, schema);
+ if (csvFormat == null) {
+ csvFormat = CSVFormat.DEFAULT;
+ }
+ this.csvFormat = csvFormat;
+
+ if (schemaFields == null) {
+ schemaFields = schema.sorted().getFieldNames();
+ }
+ this.schemaFields = schemaFields;
+ }
+
+ CSVFormat getCsvFormat() {
+ return csvFormat;
+ }
+
+ List<String> getSchemaFields() {
+ return schemaFields;
+ }
+
+ /**
+ * Serializes a {@link Row} to a CSV byte[] record using a {@link CSVFormat}. The schemaFields
+ * provided or generated by the constructor drive the subset and order of serialized {@link Row}
+ * fields.
+ */
+ @Override
+ public byte[] serialize(Row row) {
+ StringBuilder builder = new StringBuilder();
+ try {
+ boolean newRecord = true;
+ for (String name : schemaFields) {
+ Object value = row.getValue(name);
+ if (value == null) {
+ value = "";
+ }
+ csvFormat.print(value, builder, newRecord);
+ newRecord = false;
+ }
+ } catch (IOException e) {
+ throw new IllegalStateException(e);
+ }
+ return builder.toString().getBytes(StandardCharsets.UTF_8);
+ }
+
+ /**
+ * Converts {@param bytes} to a {@link Row} based on the {@link CSVFormat} and schemaFields
+ * parameter in the constructor. Not yet implemented. See
+ * https://github.com/apache/beam/issues/24552.
Review Comment:
use `<a href=` tag
##########
sdks/java/io/csv/src/main/java/org/apache/beam/sdk/io/csv/CsvUtils.java:
##########
@@ -0,0 +1,250 @@
+/*
+ * 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.sdk.io.csv;
+
+import java.io.IOException;
+import java.nio.charset.StandardCharsets;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Set;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.Field;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.transforms.SimpleFunction;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableSet;
+import org.apache.commons.csv.CSVFormat;
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+/** Utilities to convert between CSV records, Beam rows, and user types. */
+public final class CsvUtils {
+
+ public static final Set<FieldType> VALID_FIELD_TYPE_SET =
Review Comment:
We'll want to describe in CsvIO header the field types that are supported when converting to/from schema.
##########
sdks/java/io/csv/src/main/java/org/apache/beam/sdk/io/csv/CsvIO.java:
##########
@@ -0,0 +1,384 @@
+/*
+ * 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.sdk.io.csv;
+
+import static java.nio.charset.StandardCharsets.UTF_8;
+import static org.apache.beam.sdk.values.TypeDescriptors.rows;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+
+import com.google.auto.value.AutoValue;
+import java.io.BufferedWriter;
+import java.io.IOException;
+import java.io.OutputStreamWriter;
+import java.io.PrintWriter;
+import java.nio.channels.Channels;
+import java.nio.channels.WritableByteChannel;
+import java.util.List;
+import java.util.Objects;
+import java.util.Optional;
+import org.apache.beam.sdk.io.Compression;
+import org.apache.beam.sdk.io.FileIO;
+import org.apache.beam.sdk.io.FileSystems;
+import org.apache.beam.sdk.io.fs.ResourceId;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.transforms.MapElements;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.SimpleFunction;
+import org.apache.beam.sdk.transforms.display.DisplayData;
+import org.apache.beam.sdk.transforms.display.HasDisplayData;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PDone;
+import org.apache.beam.sdk.values.Row;
+import org.apache.commons.csv.CSVFormat;
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+/** Transforms for reading and writing CSV files. */
+@SuppressWarnings({
+ "unused",
+ "nullness" // TODO(https://github.com/apache/beam/issues/20497)
+})
+public class CsvIO {
+
+ static final String DEFAULT_FILENAME_SUFFIX = ".csv";
+
+ /** Instantiates a {@link Write} for writing user types in CSV format. */
+ public static <T> Write<T> write() {
+ return new AutoValue_CsvIO_Write.Builder<T>().build();
+ }
+
+ /** Instantiates a {@link Write} for {@link Row}s in CSV format. */
+ public static Write<Row> writeRows() {
+ return new AutoValue_CsvIO_Write.Builder<Row>().build();
+ }
+
+ /** Implementation of {@link FileIO.Sink}. */
+ @AutoValue
+ public abstract static class Sink<T> implements FileIO.Sink<T> {
+
+ public static <T> Builder<T> builder() {
+ return new AutoValue_CsvIO_Sink.Builder<>();
+ }
+
+ /**
+ * Not to be confused with the CSV header, it is content written to the top of every sharded
+ * file prior to the header. In the example below, all the text proceeding the header
+ * 'column1,column2,column3' is the preamble.
+ *
+ * <p>Fake company, Inc. Lab experiment: abcdefg123456 Experiment date: 2022-12-05 Operator:
+ * John Doe
+ *
+ * <p>column1,column2,column3 1,2,3 4,5,6
Review Comment:
consider using `<pre>{@code ...`
See https://github.com/robertwb/incubator-beam/blob/9a36490de8129359106baf37e1d0b071e19e303a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableIO.java#L93 for an example when formatting this.
##########
sdks/java/io/csv/src/main/java/org/apache/beam/sdk/io/csv/CsvIO.java:
##########
@@ -0,0 +1,384 @@
+/*
+ * 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.sdk.io.csv;
+
+import static java.nio.charset.StandardCharsets.UTF_8;
+import static org.apache.beam.sdk.values.TypeDescriptors.rows;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+
+import com.google.auto.value.AutoValue;
+import java.io.BufferedWriter;
+import java.io.IOException;
+import java.io.OutputStreamWriter;
+import java.io.PrintWriter;
+import java.nio.channels.Channels;
+import java.nio.channels.WritableByteChannel;
+import java.util.List;
+import java.util.Objects;
+import java.util.Optional;
+import org.apache.beam.sdk.io.Compression;
+import org.apache.beam.sdk.io.FileIO;
+import org.apache.beam.sdk.io.FileSystems;
+import org.apache.beam.sdk.io.fs.ResourceId;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.transforms.MapElements;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.SimpleFunction;
+import org.apache.beam.sdk.transforms.display.DisplayData;
+import org.apache.beam.sdk.transforms.display.HasDisplayData;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PDone;
+import org.apache.beam.sdk.values.Row;
+import org.apache.commons.csv.CSVFormat;
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+/** Transforms for reading and writing CSV files. */
+@SuppressWarnings({
+ "unused",
+ "nullness" // TODO(https://github.com/apache/beam/issues/20497)
+})
+public class CsvIO {
+
+ static final String DEFAULT_FILENAME_SUFFIX = ".csv";
+
+ /** Instantiates a {@link Write} for writing user types in CSV format. */
+ public static <T> Write<T> write() {
+ return new AutoValue_CsvIO_Write.Builder<T>().build();
+ }
+
+ /** Instantiates a {@link Write} for {@link Row}s in CSV format. */
+ public static Write<Row> writeRows() {
+ return new AutoValue_CsvIO_Write.Builder<Row>().build();
+ }
+
+ /** Implementation of {@link FileIO.Sink}. */
+ @AutoValue
+ public abstract static class Sink<T> implements FileIO.Sink<T> {
+
+ public static <T> Builder<T> builder() {
+ return new AutoValue_CsvIO_Sink.Builder<>();
+ }
+
+ /**
+ * Not to be confused with the CSV header, it is content written to the top of every sharded
+ * file prior to the header. In the example below, all the text proceeding the header
+ * 'column1,column2,column3' is the preamble.
+ *
+ * <p>Fake company, Inc. Lab experiment: abcdefg123456 Experiment date: 2022-12-05 Operator:
+ * John Doe
+ *
+ * <p>column1,column2,column3 1,2,3 4,5,6
+ */
+ public Sink<T> withPreamble(String preamble) {
+ return toBuilder().setPreamble(preamble).build();
+ }
+
+ private transient @Nullable PrintWriter writer;
+
+ /**
+ * Opens a {@link WritableByteChannel} for writing CSV files. Writes the {@link #getPreamble()}
+ * if available followed by the {@link #getHeader()}.
+ */
+ @Override
+ public void open(WritableByteChannel channel) throws IOException {
+ writer =
+ new PrintWriter(
+ new BufferedWriter(new OutputStreamWriter(Channels.newOutputStream(channel), UTF_8)));
+ if (getPreamble() != null) {
+ writer.println(getPreamble());
+ }
+ writer.println(getHeader());
+ }
+
+ /** Serializes and writes the {@param element} to a file. */
+ @Override
+ public void write(T element) throws IOException {
+ if (writer == null) {
+ throw new IllegalStateException(
+ String.format("%s writer is null", PrintWriter.class.getName()));
+ }
+ String line = getFormatFunction().apply(element);
+ writer.println(line);
+ }
+
+ @Override
+ public void flush() throws IOException {
+ if (writer == null) {
+ throw new IllegalStateException(
+ String.format("%s writer is null", PrintWriter.class.getName()));
+ }
+ writer.flush();
+ }
+
+ /**
+ * Not to be confused with the CSV header, it is content written to the top of every sharded
+ * file prior to the header. In the example below, all the text proceeding the header
+ * 'column1,column2,column3' is the preamble.
+ *
+ * <p>Fake company, Inc. Lab experiment: abcdefg123456 Experiment date: 2022-12-05 Operator:
+ * John Doe
+ *
+ * <p>column1,column2,column3 1,2,3 4,5,6
+ */
+ abstract @Nullable String getPreamble();
+
+ /**
+ * The column names of the CSV file written at the top line of each shard after the preamble, if
+ * available.
+ */
+ abstract String getHeader();
+
+ /** A {@link SimpleFunction} for converting a {@param T} to a CSV formatted string. */
+ abstract SimpleFunction<T, String> getFormatFunction();
+
+ abstract Builder<T> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<T> {
+
+ /**
+ * Not to be confused with the CSV header, it is content written to the top of every sharded
+ * file prior to the header. In the example below, all the text proceeding the header
+ * 'column1,column2,column3' is the preamble.
+ *
+ * <p>Fake company, Inc. Lab experiment: abcdefg123456 Experiment date: 2022-12-05 Operator:
+ * John Doe
+ *
+ * <p>column1,column2,column3 1,2,3 4,5,6
Review Comment:
Ditto for `<pre>{@code ...`
##########
sdks/java/io/csv/src/main/java/org/apache/beam/sdk/io/csv/CsvIO.java:
##########
@@ -0,0 +1,384 @@
+/*
+ * 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.sdk.io.csv;
+
+import static java.nio.charset.StandardCharsets.UTF_8;
+import static org.apache.beam.sdk.values.TypeDescriptors.rows;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+
+import com.google.auto.value.AutoValue;
+import java.io.BufferedWriter;
+import java.io.IOException;
+import java.io.OutputStreamWriter;
+import java.io.PrintWriter;
+import java.nio.channels.Channels;
+import java.nio.channels.WritableByteChannel;
+import java.util.List;
+import java.util.Objects;
+import java.util.Optional;
+import org.apache.beam.sdk.io.Compression;
+import org.apache.beam.sdk.io.FileIO;
+import org.apache.beam.sdk.io.FileSystems;
+import org.apache.beam.sdk.io.fs.ResourceId;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.transforms.MapElements;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.SimpleFunction;
+import org.apache.beam.sdk.transforms.display.DisplayData;
+import org.apache.beam.sdk.transforms.display.HasDisplayData;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PDone;
+import org.apache.beam.sdk.values.Row;
+import org.apache.commons.csv.CSVFormat;
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+/** Transforms for reading and writing CSV files. */
+@SuppressWarnings({
+ "unused",
+ "nullness" // TODO(https://github.com/apache/beam/issues/20497)
+})
+public class CsvIO {
+
+ static final String DEFAULT_FILENAME_SUFFIX = ".csv";
+
+ /** Instantiates a {@link Write} for writing user types in CSV format. */
+ public static <T> Write<T> write() {
+ return new AutoValue_CsvIO_Write.Builder<T>().build();
+ }
+
+ /** Instantiates a {@link Write} for {@link Row}s in CSV format. */
+ public static Write<Row> writeRows() {
+ return new AutoValue_CsvIO_Write.Builder<Row>().build();
+ }
+
+ /** Implementation of {@link FileIO.Sink}. */
+ @AutoValue
+ public abstract static class Sink<T> implements FileIO.Sink<T> {
+
+ public static <T> Builder<T> builder() {
+ return new AutoValue_CsvIO_Sink.Builder<>();
+ }
+
+ /**
+ * Not to be confused with the CSV header, it is content written to the top of every sharded
+ * file prior to the header. In the example below, all the text proceeding the header
+ * 'column1,column2,column3' is the preamble.
+ *
+ * <p>Fake company, Inc. Lab experiment: abcdefg123456 Experiment date: 2022-12-05 Operator:
+ * John Doe
+ *
+ * <p>column1,column2,column3 1,2,3 4,5,6
+ */
+ public Sink<T> withPreamble(String preamble) {
+ return toBuilder().setPreamble(preamble).build();
+ }
+
+ private transient @Nullable PrintWriter writer;
+
+ /**
+ * Opens a {@link WritableByteChannel} for writing CSV files. Writes the {@link #getPreamble()}
+ * if available followed by the {@link #getHeader()}.
+ */
+ @Override
+ public void open(WritableByteChannel channel) throws IOException {
+ writer =
+ new PrintWriter(
+ new BufferedWriter(new OutputStreamWriter(Channels.newOutputStream(channel), UTF_8)));
+ if (getPreamble() != null) {
+ writer.println(getPreamble());
+ }
+ writer.println(getHeader());
+ }
+
+ /** Serializes and writes the {@param element} to a file. */
+ @Override
+ public void write(T element) throws IOException {
+ if (writer == null) {
+ throw new IllegalStateException(
+ String.format("%s writer is null", PrintWriter.class.getName()));
+ }
+ String line = getFormatFunction().apply(element);
+ writer.println(line);
+ }
+
+ @Override
+ public void flush() throws IOException {
+ if (writer == null) {
+ throw new IllegalStateException(
+ String.format("%s writer is null", PrintWriter.class.getName()));
+ }
+ writer.flush();
+ }
+
+ /**
+ * Not to be confused with the CSV header, it is content written to the top of every sharded
+ * file prior to the header. In the example below, all the text proceeding the header
+ * 'column1,column2,column3' is the preamble.
+ *
+ * <p>Fake company, Inc. Lab experiment: abcdefg123456 Experiment date: 2022-12-05 Operator:
+ * John Doe
+ *
+ * <p>column1,column2,column3 1,2,3 4,5,6
+ */
+ abstract @Nullable String getPreamble();
+
+ /**
+ * The column names of the CSV file written at the top line of each shard after the preamble, if
+ * available.
+ */
+ abstract String getHeader();
+
+ /** A {@link SimpleFunction} for converting a {@param T} to a CSV formatted string. */
+ abstract SimpleFunction<T, String> getFormatFunction();
+
+ abstract Builder<T> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<T> {
+
+ /**
+ * Not to be confused with the CSV header, it is content written to the top of every sharded
+ * file prior to the header. In the example below, all the text proceeding the header
+ * 'column1,column2,column3' is the preamble.
+ *
+ * <p>Fake company, Inc. Lab experiment: abcdefg123456 Experiment date: 2022-12-05 Operator:
+ * John Doe
+ *
+ * <p>column1,column2,column3 1,2,3 4,5,6
+ */
+ abstract Builder<T> setPreamble(String value);
+
+ /**
+ * The column names of the CSV file written at the top line of each shard after the preamble,
+ * if available.
+ */
+ abstract Builder<T> setHeader(String value);
+
+ /** A {@link SimpleFunction} for converting a {@param T} to a CSV formatted string. */
+ abstract Builder<T> setFormatFunction(SimpleFunction<T, String> value);
+
+ abstract Sink<T> build();
+ }
+ }
+
+ @AutoValue
+ public abstract static class Write<T> extends PTransform<PCollection<T>, PDone>
+ implements HasDisplayData {
+
+ /** Specifies a common prefix for all generated files. */
+ public Write<T> to(String filenamePrefix) {
+ return toBuilder().setFilenamePrefix(filenamePrefix).build();
+ }
+
+ /** The {@link CSVFormat} of the destination CSV file data. */
+ public Write<T> withCSVFormat(CSVFormat format) {
+ return toBuilder().setCSVFormat(format).build();
+ }
+
+ /**
+ * Not to be confused with the CSV header, it is content written to the top of every sharded
+ * file prior to the header. In the example below, all the text proceeding the header
+ * 'column1,column2,column3' is the preamble.
+ *
+ * <p>Fake company, Inc. Lab experiment: abcdefg123456 Experiment date: 2022-12-05 Operator:
+ * John Doe
+ *
+ * <p>column1,column2,column3 1,2,3 4,5,6
+ */
+ public Write<T> withPreamble(String preamble) {
+ return toBuilder().setPreamble(preamble).build();
+ }
+
+ /** Specifies the {@link Compression} of all generated shard files. */
+ public Write<T> withCompression(Compression compression) {
+ return toBuilder().setCompression(compression).build();
+ }
+
+ /**
+ * Specifies to use a given fixed number of shards per window. See {@link
+ * FileIO.Write#withNumShards(int)} for details.
+ */
+ public Write<T> withNumShards(Integer numShards) {
+ return toBuilder().setNumShards(numShards).build();
+ }
+
+ /**
+ * Specifies a directory into which all temporary files will be placed. See {@link
+ * FileIO.Write#withTempDirectory(String)}.
+ */
+ public Write<T> withTempDirectory(ResourceId value) {
+ return toBuilder().setTempDirectory(value).build();
+ }
+
+ abstract @Nullable String getFilenamePrefix();
+
+ abstract CSVFormat getCSVFormat();
+
+ abstract @Nullable String getPreamble();
+
+ abstract @Nullable Compression getCompression();
+
+ abstract @Nullable Integer getNumShards();
+
+ abstract String getFilenameSuffix();
+
+ abstract @Nullable ResourceId getTempDirectory();
+
+ abstract @Nullable List<String> getSchemaFields();
+
+ abstract Builder<T> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<T> {
+
+ abstract Builder<T> setFilenamePrefix(String value);
+
+ abstract Builder<T> setCSVFormat(CSVFormat value);
+
+ abstract Optional<CSVFormat> getCSVFormat();
+
+ abstract Builder<T> setPreamble(String value);
+
+ abstract Builder<T> setCompression(Compression value);
+
+ abstract Builder<T> setNumShards(Integer value);
+
+ abstract Builder<T> setFilenameSuffix(String value);
+
+ abstract Optional<String> getFilenameSuffix();
+
+ abstract Builder<T> setTempDirectory(ResourceId value);
+
+ abstract Builder<T> setSchemaFields(List<String> value);
+
+ abstract Write<T> autoBuild();
+
+ final Write<T> build() {
+ if (!getCSVFormat().isPresent()) {
+ setCSVFormat(CSVFormat.DEFAULT);
+ }
+
+ if (!getFilenameSuffix().isPresent()) {
+ setFilenameSuffix(DEFAULT_FILENAME_SUFFIX);
+ }
+
+ return autoBuild();
+ }
+ }
+
+ @Override
+ public void populateDisplayData(DisplayData.Builder builder) {
+ if (getFilenamePrefix() != null) {
+ builder.add(DisplayData.item("filenamePrefix", getFilenamePrefix()));
+ }
+ if (getCSVFormat() != null) {
+ builder.add(DisplayData.item("csvFormat", getCSVFormat().toString()));
+ }
+ if (getPreamble() != null) {
+ builder.add(DisplayData.item("preamble", getPreamble()));
+ }
+ if (getCompression() != null) {
+ builder.add(DisplayData.item("compression", getCompression().name()));
+ }
+ if (getNumShards() != null) {
+ builder.add(DisplayData.item("numShards", getNumShards()));
+ }
+
+ builder.add(DisplayData.item("filenameSuffix", getFilenameSuffix()));
+
+ if (getTempDirectory() != null) {
+ builder.add(DisplayData.item("tempDirectory", getTempDirectory().getFilename()));
+ }
+ }
+
+ @Override
+ public PDone expand(PCollection<T> input) {
+ if (!input.hasSchema()) {
+ throw new IllegalArgumentException(
+ String.format(
+ "%s requires an input Schema. Note that only Row or user classes are supported. Consider using TextIO or FileIO directly when writing primitive types",
+ Write.class.getName()));
+ }
+
+ Schema schema = input.getSchema();
+ SerializableFunction<T, Row> toRowFn = input.getToRowFunction();
+ PCollection<Row> rows =
+ input.apply(MapElements.into(rows()).via(toRowFn)).setRowSchema(schema);
Review Comment:
It is always good to give names for each PTransform that is applied so that it is easier to understand what is occurring.
```suggestion
input.apply("To Rows", MapElements.into(rows()).via(toRowFn)).setRowSchema(schema);
```
##########
sdks/java/io/csv/src/main/java/org/apache/beam/sdk/io/csv/CsvIO.java:
##########
@@ -0,0 +1,384 @@
+/*
+ * 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.sdk.io.csv;
+
+import static java.nio.charset.StandardCharsets.UTF_8;
+import static org.apache.beam.sdk.values.TypeDescriptors.rows;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+
+import com.google.auto.value.AutoValue;
+import java.io.BufferedWriter;
+import java.io.IOException;
+import java.io.OutputStreamWriter;
+import java.io.PrintWriter;
+import java.nio.channels.Channels;
+import java.nio.channels.WritableByteChannel;
+import java.util.List;
+import java.util.Objects;
+import java.util.Optional;
+import org.apache.beam.sdk.io.Compression;
+import org.apache.beam.sdk.io.FileIO;
+import org.apache.beam.sdk.io.FileSystems;
+import org.apache.beam.sdk.io.fs.ResourceId;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.transforms.MapElements;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.SimpleFunction;
+import org.apache.beam.sdk.transforms.display.DisplayData;
+import org.apache.beam.sdk.transforms.display.HasDisplayData;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PDone;
+import org.apache.beam.sdk.values.Row;
+import org.apache.commons.csv.CSVFormat;
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+/** Transforms for reading and writing CSV files. */
+@SuppressWarnings({
+ "unused",
+ "nullness" // TODO(https://github.com/apache/beam/issues/20497)
+})
+public class CsvIO {
+
+ static final String DEFAULT_FILENAME_SUFFIX = ".csv";
+
+ /** Instantiates a {@link Write} for writing user types in CSV format. */
+ public static <T> Write<T> write() {
+ return new AutoValue_CsvIO_Write.Builder<T>().build();
+ }
+
+ /** Instantiates a {@link Write} for {@link Row}s in CSV format. */
+ public static Write<Row> writeRows() {
+ return new AutoValue_CsvIO_Write.Builder<Row>().build();
+ }
+
+ /** Implementation of {@link FileIO.Sink}. */
+ @AutoValue
+ public abstract static class Sink<T> implements FileIO.Sink<T> {
+
+ public static <T> Builder<T> builder() {
+ return new AutoValue_CsvIO_Sink.Builder<>();
+ }
+
+ /**
+ * Not to be confused with the CSV header, it is content written to the top of every sharded
+ * file prior to the header. In the example below, all the text proceeding the header
+ * 'column1,column2,column3' is the preamble.
+ *
+ * <p>Fake company, Inc. Lab experiment: abcdefg123456 Experiment date: 2022-12-05 Operator:
+ * John Doe
+ *
+ * <p>column1,column2,column3 1,2,3 4,5,6
+ */
+ public Sink<T> withPreamble(String preamble) {
+ return toBuilder().setPreamble(preamble).build();
+ }
+
+ private transient @Nullable PrintWriter writer;
+
+ /**
+ * Opens a {@link WritableByteChannel} for writing CSV files. Writes the {@link #getPreamble()}
+ * if available followed by the {@link #getHeader()}.
+ */
+ @Override
+ public void open(WritableByteChannel channel) throws IOException {
+ writer =
+ new PrintWriter(
+ new BufferedWriter(new OutputStreamWriter(Channels.newOutputStream(channel), UTF_8)));
+ if (getPreamble() != null) {
+ writer.println(getPreamble());
+ }
+ writer.println(getHeader());
+ }
+
+ /** Serializes and writes the {@param element} to a file. */
+ @Override
+ public void write(T element) throws IOException {
+ if (writer == null) {
+ throw new IllegalStateException(
+ String.format("%s writer is null", PrintWriter.class.getName()));
+ }
+ String line = getFormatFunction().apply(element);
+ writer.println(line);
+ }
+
+ @Override
+ public void flush() throws IOException {
+ if (writer == null) {
+ throw new IllegalStateException(
+ String.format("%s writer is null", PrintWriter.class.getName()));
+ }
+ writer.flush();
+ }
+
+ /**
+ * Not to be confused with the CSV header, it is content written to the top of every sharded
+ * file prior to the header. In the example below, all the text proceeding the header
+ * 'column1,column2,column3' is the preamble.
+ *
+ * <p>Fake company, Inc. Lab experiment: abcdefg123456 Experiment date: 2022-12-05 Operator:
+ * John Doe
+ *
+ * <p>column1,column2,column3 1,2,3 4,5,6
+ */
+ abstract @Nullable String getPreamble();
+
+ /**
+ * The column names of the CSV file written at the top line of each shard after the preamble, if
+ * available.
+ */
+ abstract String getHeader();
+
+ /** A {@link SimpleFunction} for converting a {@param T} to a CSV formatted string. */
+ abstract SimpleFunction<T, String> getFormatFunction();
Review Comment:
Use SerializableFunction or ProcessFunction depending on whether you want to allow the user to throw an exception or not.
##########
sdks/java/io/csv/src/main/java/org/apache/beam/sdk/io/csv/CsvIO.java:
##########
@@ -0,0 +1,384 @@
+/*
+ * 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.sdk.io.csv;
+
+import static java.nio.charset.StandardCharsets.UTF_8;
+import static org.apache.beam.sdk.values.TypeDescriptors.rows;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+
+import com.google.auto.value.AutoValue;
+import java.io.BufferedWriter;
+import java.io.IOException;
+import java.io.OutputStreamWriter;
+import java.io.PrintWriter;
+import java.nio.channels.Channels;
+import java.nio.channels.WritableByteChannel;
+import java.util.List;
+import java.util.Objects;
+import java.util.Optional;
+import org.apache.beam.sdk.io.Compression;
+import org.apache.beam.sdk.io.FileIO;
+import org.apache.beam.sdk.io.FileSystems;
+import org.apache.beam.sdk.io.fs.ResourceId;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.transforms.MapElements;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.SimpleFunction;
+import org.apache.beam.sdk.transforms.display.DisplayData;
+import org.apache.beam.sdk.transforms.display.HasDisplayData;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PDone;
+import org.apache.beam.sdk.values.Row;
+import org.apache.commons.csv.CSVFormat;
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+/** Transforms for reading and writing CSV files. */
+@SuppressWarnings({
+ "unused",
+ "nullness" // TODO(https://github.com/apache/beam/issues/20497)
+})
+public class CsvIO {
+
+ static final String DEFAULT_FILENAME_SUFFIX = ".csv";
+
+ /** Instantiates a {@link Write} for writing user types in CSV format. */
+ public static <T> Write<T> write() {
+ return new AutoValue_CsvIO_Write.Builder<T>().build();
+ }
+
+ /** Instantiates a {@link Write} for {@link Row}s in CSV format. */
+ public static Write<Row> writeRows() {
+ return new AutoValue_CsvIO_Write.Builder<Row>().build();
+ }
+
+ /** Implementation of {@link FileIO.Sink}. */
+ @AutoValue
+ public abstract static class Sink<T> implements FileIO.Sink<T> {
+
+ public static <T> Builder<T> builder() {
+ return new AutoValue_CsvIO_Sink.Builder<>();
+ }
+
+ /**
+ * Not to be confused with the CSV header, it is content written to the top of every sharded
+ * file prior to the header. In the example below, all the text proceeding the header
+ * 'column1,column2,column3' is the preamble.
+ *
+ * <p>Fake company, Inc. Lab experiment: abcdefg123456 Experiment date: 2022-12-05 Operator:
+ * John Doe
+ *
+ * <p>column1,column2,column3 1,2,3 4,5,6
+ */
+ public Sink<T> withPreamble(String preamble) {
+ return toBuilder().setPreamble(preamble).build();
+ }
+
+ private transient @Nullable PrintWriter writer;
+
+ /**
+ * Opens a {@link WritableByteChannel} for writing CSV files. Writes the {@link #getPreamble()}
+ * if available followed by the {@link #getHeader()}.
+ */
+ @Override
+ public void open(WritableByteChannel channel) throws IOException {
+ writer =
+ new PrintWriter(
+ new BufferedWriter(new OutputStreamWriter(Channels.newOutputStream(channel), UTF_8)));
+ if (getPreamble() != null) {
+ writer.println(getPreamble());
+ }
+ writer.println(getHeader());
+ }
+
+ /** Serializes and writes the {@param element} to a file. */
+ @Override
+ public void write(T element) throws IOException {
+ if (writer == null) {
+ throw new IllegalStateException(
+ String.format("%s writer is null", PrintWriter.class.getName()));
+ }
+ String line = getFormatFunction().apply(element);
+ writer.println(line);
+ }
+
+ @Override
+ public void flush() throws IOException {
+ if (writer == null) {
+ throw new IllegalStateException(
+ String.format("%s writer is null", PrintWriter.class.getName()));
+ }
+ writer.flush();
+ }
+
+ /**
+ * Not to be confused with the CSV header, it is content written to the top of every sharded
+ * file prior to the header. In the example below, all the text proceeding the header
+ * 'column1,column2,column3' is the preamble.
+ *
+ * <p>Fake company, Inc. Lab experiment: abcdefg123456 Experiment date: 2022-12-05 Operator:
+ * John Doe
+ *
+ * <p>column1,column2,column3 1,2,3 4,5,6
+ */
+ abstract @Nullable String getPreamble();
+
+ /**
+ * The column names of the CSV file written at the top line of each shard after the preamble, if
+ * available.
+ */
+ abstract String getHeader();
+
+ /** A {@link SimpleFunction} for converting a {@param T} to a CSV formatted string. */
+ abstract SimpleFunction<T, String> getFormatFunction();
+
+ abstract Builder<T> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<T> {
+
+ /**
+ * Not to be confused with the CSV header, it is content written to the top of every sharded
+ * file prior to the header. In the example below, all the text proceeding the header
+ * 'column1,column2,column3' is the preamble.
+ *
+ * <p>Fake company, Inc. Lab experiment: abcdefg123456 Experiment date: 2022-12-05 Operator:
+ * John Doe
+ *
+ * <p>column1,column2,column3 1,2,3 4,5,6
+ */
+ abstract Builder<T> setPreamble(String value);
+
+ /**
+ * The column names of the CSV file written at the top line of each shard after the preamble,
+ * if available.
+ */
+ abstract Builder<T> setHeader(String value);
+
+ /** A {@link SimpleFunction} for converting a {@param T} to a CSV formatted string. */
+ abstract Builder<T> setFormatFunction(SimpleFunction<T, String> value);
+
+ abstract Sink<T> build();
+ }
+ }
+
+ @AutoValue
+ public abstract static class Write<T> extends PTransform<PCollection<T>, PDone>
+ implements HasDisplayData {
+
+ /** Specifies a common prefix for all generated files. */
+ public Write<T> to(String filenamePrefix) {
+ return toBuilder().setFilenamePrefix(filenamePrefix).build();
+ }
+
+ /** The {@link CSVFormat} of the destination CSV file data. */
+ public Write<T> withCSVFormat(CSVFormat format) {
+ return toBuilder().setCSVFormat(format).build();
+ }
+
+ /**
+ * Not to be confused with the CSV header, it is content written to the top of every sharded
+ * file prior to the header. In the example below, all the text proceeding the header
+ * 'column1,column2,column3' is the preamble.
+ *
+ * <p>Fake company, Inc. Lab experiment: abcdefg123456 Experiment date: 2022-12-05 Operator:
+ * John Doe
+ *
+ * <p>column1,column2,column3 1,2,3 4,5,6
+ */
+ public Write<T> withPreamble(String preamble) {
+ return toBuilder().setPreamble(preamble).build();
+ }
+
+ /** Specifies the {@link Compression} of all generated shard files. */
+ public Write<T> withCompression(Compression compression) {
+ return toBuilder().setCompression(compression).build();
+ }
+
+ /**
+ * Specifies to use a given fixed number of shards per window. See {@link
+ * FileIO.Write#withNumShards(int)} for details.
+ */
+ public Write<T> withNumShards(Integer numShards) {
+ return toBuilder().setNumShards(numShards).build();
+ }
+
+ /**
+ * Specifies a directory into which all temporary files will be placed. See {@link
+ * FileIO.Write#withTempDirectory(String)}.
+ */
+ public Write<T> withTempDirectory(ResourceId value) {
+ return toBuilder().setTempDirectory(value).build();
+ }
+
+ abstract @Nullable String getFilenamePrefix();
+
+ abstract CSVFormat getCSVFormat();
+
+ abstract @Nullable String getPreamble();
+
+ abstract @Nullable Compression getCompression();
+
+ abstract @Nullable Integer getNumShards();
+
+ abstract String getFilenameSuffix();
+
+ abstract @Nullable ResourceId getTempDirectory();
+
+ abstract @Nullable List<String> getSchemaFields();
+
+ abstract Builder<T> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<T> {
+
+ abstract Builder<T> setFilenamePrefix(String value);
+
+ abstract Builder<T> setCSVFormat(CSVFormat value);
+
+ abstract Optional<CSVFormat> getCSVFormat();
+
+ abstract Builder<T> setPreamble(String value);
+
+ abstract Builder<T> setCompression(Compression value);
+
+ abstract Builder<T> setNumShards(Integer value);
+
+ abstract Builder<T> setFilenameSuffix(String value);
+
+ abstract Optional<String> getFilenameSuffix();
+
+ abstract Builder<T> setTempDirectory(ResourceId value);
+
+ abstract Builder<T> setSchemaFields(List<String> value);
+
+ abstract Write<T> autoBuild();
+
+ final Write<T> build() {
+ if (!getCSVFormat().isPresent()) {
+ setCSVFormat(CSVFormat.DEFAULT);
Review Comment:
Do you differentiate between the user setting CSVFormat.DEFAULT and us setting it?
If not, consider dropping Optional and ensuring that we set this everywhere when constructing the builder.
##########
sdks/java/io/csv/src/main/java/org/apache/beam/sdk/io/csv/CsvPayloadSerializer.java:
##########
@@ -0,0 +1,97 @@
+/*
+ * 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.sdk.io.csv;
+
+import java.io.IOException;
+import java.nio.charset.StandardCharsets;
+import java.util.List;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.io.payloads.PayloadSerializer;
+import org.apache.beam.sdk.values.Row;
+import org.apache.commons.csv.CSVFormat;
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+/** Implementation of {@link PayloadSerializer} for a CSV format. */
+public class CsvPayloadSerializer implements PayloadSerializer {
+
+ private final CSVFormat csvFormat;
+ private final List<String> schemaFields;
+
+ /**
+ * Instantiates a {@link PayloadSerializer} for a CSV format. {@param csvFormat} defaults to
+ * {@link CSVFormat#DEFAULT} and {@param schemaFields} defaults to {@link Schema#sorted()} {@link
+ * Schema#getFieldNames()}.
+ */
+ CsvPayloadSerializer(
Review Comment:
It is easier to reason about this if you create constructors with only non-nullable fields and have the default operation being done OR have the caller own Schema->fields conversion and passing in CSVFormat#DEFAULT
Note, will CSVFormat ever actually be null?
##########
sdks/java/io/csv/src/main/java/org/apache/beam/sdk/io/csv/CsvIO.java:
##########
@@ -0,0 +1,384 @@
+/*
+ * 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.sdk.io.csv;
+
+import static java.nio.charset.StandardCharsets.UTF_8;
+import static org.apache.beam.sdk.values.TypeDescriptors.rows;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+
+import com.google.auto.value.AutoValue;
+import java.io.BufferedWriter;
+import java.io.IOException;
+import java.io.OutputStreamWriter;
+import java.io.PrintWriter;
+import java.nio.channels.Channels;
+import java.nio.channels.WritableByteChannel;
+import java.util.List;
+import java.util.Objects;
+import java.util.Optional;
+import org.apache.beam.sdk.io.Compression;
+import org.apache.beam.sdk.io.FileIO;
+import org.apache.beam.sdk.io.FileSystems;
+import org.apache.beam.sdk.io.fs.ResourceId;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.transforms.MapElements;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.SimpleFunction;
+import org.apache.beam.sdk.transforms.display.DisplayData;
+import org.apache.beam.sdk.transforms.display.HasDisplayData;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PDone;
+import org.apache.beam.sdk.values.Row;
+import org.apache.commons.csv.CSVFormat;
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+/** Transforms for reading and writing CSV files. */
+@SuppressWarnings({
+ "unused",
+ "nullness" // TODO(https://github.com/apache/beam/issues/20497)
+})
+public class CsvIO {
+
+ static final String DEFAULT_FILENAME_SUFFIX = ".csv";
+
+ /** Instantiates a {@link Write} for writing user types in CSV format. */
+ public static <T> Write<T> write() {
+ return new AutoValue_CsvIO_Write.Builder<T>().build();
+ }
+
+ /** Instantiates a {@link Write} for {@link Row}s in CSV format. */
+ public static Write<Row> writeRows() {
+ return new AutoValue_CsvIO_Write.Builder<Row>().build();
+ }
+
+ /** Implementation of {@link FileIO.Sink}. */
+ @AutoValue
+ public abstract static class Sink<T> implements FileIO.Sink<T> {
+
+ public static <T> Builder<T> builder() {
+ return new AutoValue_CsvIO_Sink.Builder<>();
+ }
+
+ /**
+ * Not to be confused with the CSV header, it is content written to the top of every sharded
+ * file prior to the header. In the example below, all the text proceeding the header
+ * 'column1,column2,column3' is the preamble.
+ *
+ * <p>Fake company, Inc. Lab experiment: abcdefg123456 Experiment date: 2022-12-05 Operator:
+ * John Doe
+ *
+ * <p>column1,column2,column3 1,2,3 4,5,6
+ */
+ public Sink<T> withPreamble(String preamble) {
+ return toBuilder().setPreamble(preamble).build();
+ }
+
+ private transient @Nullable PrintWriter writer;
+
+ /**
+ * Opens a {@link WritableByteChannel} for writing CSV files. Writes the {@link #getPreamble()}
+ * if available followed by the {@link #getHeader()}.
+ */
+ @Override
+ public void open(WritableByteChannel channel) throws IOException {
+ writer =
+ new PrintWriter(
+ new BufferedWriter(new OutputStreamWriter(Channels.newOutputStream(channel), UTF_8)));
+ if (getPreamble() != null) {
+ writer.println(getPreamble());
+ }
+ writer.println(getHeader());
+ }
+
+ /** Serializes and writes the {@param element} to a file. */
+ @Override
+ public void write(T element) throws IOException {
+ if (writer == null) {
+ throw new IllegalStateException(
+ String.format("%s writer is null", PrintWriter.class.getName()));
+ }
+ String line = getFormatFunction().apply(element);
+ writer.println(line);
+ }
+
+ @Override
+ public void flush() throws IOException {
+ if (writer == null) {
+ throw new IllegalStateException(
+ String.format("%s writer is null", PrintWriter.class.getName()));
+ }
+ writer.flush();
+ }
+
+ /**
+ * Not to be confused with the CSV header, it is content written to the top of every sharded
+ * file prior to the header. In the example below, all the text proceeding the header
+ * 'column1,column2,column3' is the preamble.
+ *
+ * <p>Fake company, Inc. Lab experiment: abcdefg123456 Experiment date: 2022-12-05 Operator:
+ * John Doe
+ *
+ * <p>column1,column2,column3 1,2,3 4,5,6
+ */
+ abstract @Nullable String getPreamble();
+
+ /**
+ * The column names of the CSV file written at the top line of each shard after the preamble, if
+ * available.
+ */
+ abstract String getHeader();
+
+ /** A {@link SimpleFunction} for converting a {@param T} to a CSV formatted string. */
+ abstract SimpleFunction<T, String> getFormatFunction();
+
+ abstract Builder<T> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<T> {
+
+ /**
+ * Not to be confused with the CSV header, it is content written to the top of every sharded
+ * file prior to the header. In the example below, all the text proceeding the header
+ * 'column1,column2,column3' is the preamble.
+ *
+ * <p>Fake company, Inc. Lab experiment: abcdefg123456 Experiment date: 2022-12-05 Operator:
+ * John Doe
+ *
+ * <p>column1,column2,column3 1,2,3 4,5,6
+ */
+ abstract Builder<T> setPreamble(String value);
+
+ /**
+ * The column names of the CSV file written at the top line of each shard after the preamble,
+ * if available.
+ */
+ abstract Builder<T> setHeader(String value);
+
+ /** A {@link SimpleFunction} for converting a {@param T} to a CSV formatted string. */
+ abstract Builder<T> setFormatFunction(SimpleFunction<T, String> value);
+
+ abstract Sink<T> build();
+ }
+ }
+
+ @AutoValue
+ public abstract static class Write<T> extends PTransform<PCollection<T>, PDone>
+ implements HasDisplayData {
+
+ /** Specifies a common prefix for all generated files. */
+ public Write<T> to(String filenamePrefix) {
+ return toBuilder().setFilenamePrefix(filenamePrefix).build();
+ }
+
+ /** The {@link CSVFormat} of the destination CSV file data. */
+ public Write<T> withCSVFormat(CSVFormat format) {
+ return toBuilder().setCSVFormat(format).build();
+ }
+
+ /**
+ * Not to be confused with the CSV header, it is content written to the top of every sharded
+ * file prior to the header. In the example below, all the text proceeding the header
+ * 'column1,column2,column3' is the preamble.
+ *
+ * <p>Fake company, Inc. Lab experiment: abcdefg123456 Experiment date: 2022-12-05 Operator:
+ * John Doe
+ *
+ * <p>column1,column2,column3 1,2,3 4,5,6
+ */
+ public Write<T> withPreamble(String preamble) {
+ return toBuilder().setPreamble(preamble).build();
+ }
+
+ /** Specifies the {@link Compression} of all generated shard files. */
+ public Write<T> withCompression(Compression compression) {
+ return toBuilder().setCompression(compression).build();
+ }
+
+ /**
+ * Specifies to use a given fixed number of shards per window. See {@link
+ * FileIO.Write#withNumShards(int)} for details.
+ */
+ public Write<T> withNumShards(Integer numShards) {
+ return toBuilder().setNumShards(numShards).build();
+ }
+
+ /**
+ * Specifies a directory into which all temporary files will be placed. See {@link
+ * FileIO.Write#withTempDirectory(String)}.
+ */
+ public Write<T> withTempDirectory(ResourceId value) {
+ return toBuilder().setTempDirectory(value).build();
+ }
+
+ abstract @Nullable String getFilenamePrefix();
+
+ abstract CSVFormat getCSVFormat();
+
+ abstract @Nullable String getPreamble();
+
+ abstract @Nullable Compression getCompression();
+
+ abstract @Nullable Integer getNumShards();
+
+ abstract String getFilenameSuffix();
+
+ abstract @Nullable ResourceId getTempDirectory();
+
+ abstract @Nullable List<String> getSchemaFields();
+
+ abstract Builder<T> toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder<T> {
+
+ abstract Builder<T> setFilenamePrefix(String value);
+
+ abstract Builder<T> setCSVFormat(CSVFormat value);
+
+ abstract Optional<CSVFormat> getCSVFormat();
+
+ abstract Builder<T> setPreamble(String value);
+
+ abstract Builder<T> setCompression(Compression value);
+
+ abstract Builder<T> setNumShards(Integer value);
+
+ abstract Builder<T> setFilenameSuffix(String value);
+
+ abstract Optional<String> getFilenameSuffix();
+
+ abstract Builder<T> setTempDirectory(ResourceId value);
+
+ abstract Builder<T> setSchemaFields(List<String> value);
+
+ abstract Write<T> autoBuild();
+
+ final Write<T> build() {
+ if (!getCSVFormat().isPresent()) {
+ setCSVFormat(CSVFormat.DEFAULT);
+ }
+
+ if (!getFilenameSuffix().isPresent()) {
+ setFilenameSuffix(DEFAULT_FILENAME_SUFFIX);
+ }
+
+ return autoBuild();
+ }
+ }
+
+ @Override
+ public void populateDisplayData(DisplayData.Builder builder) {
+ if (getFilenamePrefix() != null) {
+ builder.add(DisplayData.item("filenamePrefix", getFilenamePrefix()));
+ }
+ if (getCSVFormat() != null) {
+ builder.add(DisplayData.item("csvFormat", getCSVFormat().toString()));
+ }
+ if (getPreamble() != null) {
+ builder.add(DisplayData.item("preamble", getPreamble()));
+ }
+ if (getCompression() != null) {
+ builder.add(DisplayData.item("compression", getCompression().name()));
+ }
+ if (getNumShards() != null) {
+ builder.add(DisplayData.item("numShards", getNumShards()));
+ }
+
+ builder.add(DisplayData.item("filenameSuffix", getFilenameSuffix()));
+
+ if (getTempDirectory() != null) {
+ builder.add(DisplayData.item("tempDirectory", getTempDirectory().getFilename()));
+ }
+ }
+
+ @Override
+ public PDone expand(PCollection<T> input) {
+ if (!input.hasSchema()) {
+ throw new IllegalArgumentException(
+ String.format(
+ "%s requires an input Schema. Note that only Row or user classes are supported. Consider using TextIO or FileIO directly when writing primitive types",
+ Write.class.getName()));
+ }
+
+ Schema schema = input.getSchema();
+ SerializableFunction<T, Row> toRowFn = input.getToRowFunction();
+ PCollection<Row> rows =
+ input.apply(MapElements.into(rows()).via(toRowFn)).setRowSchema(schema);
+ rows.apply("writeRowsToCsv", buildFileIOWrite().via(buildSink(schema)));
+ return PDone.in(input.getPipeline());
+ }
+
+ /**
+ * Builds a header using {@link CSVFormat} based on either a {@link Schema#sorted()} {@link
+ * Schema#getFieldNames()} if {@link #getSchemaFields()} is null or {@link #getSchemaFields()}.
+ */
+ String buildHeader(Schema schema) {
+ if (getSchemaFields() != null) {
+ return CsvUtils.buildHeaderFrom(getSchemaFields(), getCSVFormat());
+ }
+ return CsvUtils.buildHeaderFrom(schema.sorted(), getCSVFormat());
+ }
+
+ /** Builds a {@link Sink} for writing {@link Row} serialized using {@link CSVFormat}. */
+ Sink<Row> buildSink(Schema schema) {
+ List<String> schemaFields = null;
+ String header = null;
+ if (getSchemaFields() != null) {
+ schemaFields = getSchemaFields();
+ }
+ return Sink.<Row>builder()
+ .setPreamble(getPreamble())
+ .setHeader(buildHeader(schema))
+ .setFormatFunction(
+ CsvUtils.getRowToCsvStringFunction(schema, getCSVFormat(), schemaFields))
+ .build();
+ }
+
+ /** Builds a {@link FileIO.Write} with a {@link Sink}. */
+ FileIO.Write<Void, Row> buildFileIOWrite() {
+ checkArgument(getFilenamePrefix() != null, "to() is required");
+
+ ResourceId prefix =
+ FileSystems.matchNewResource(getFilenamePrefix(), false /* isDirectory */);
+
+ FileIO.Write<Void, Row> write =
+ FileIO.<Row>write()
+ .to(prefix.getCurrentDirectory().toString())
+ .withPrefix(Objects.requireNonNull(prefix.getFilename()));
+
+ if (getCompression() != null) {
+ write = write.withCompression(getCompression());
+ }
+
+ if (getNumShards() != null) {
+ write = write.withNumShards(getNumShards());
+ }
+
+ if (getFilenameSuffix() != null) {
+ write = write.withSuffix(getFilenameSuffix());
+ }
+
+ if (getTempDirectory() != null) {
+ write = write.withTempDirectory(Objects.requireNonNull(getTempDirectory().getFilename()));
+ }
+
+ return write;
+ }
Review Comment:
nit: Consider inlining these into expand since these are small methods and we don't expect them to be used outside of expand()
##########
sdks/java/io/csv/src/main/java/org/apache/beam/sdk/io/csv/CsvUtils.java:
##########
@@ -0,0 +1,250 @@
+/*
+ * 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.sdk.io.csv;
+
+import java.io.IOException;
+import java.nio.charset.StandardCharsets;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Set;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.Field;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.transforms.SimpleFunction;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableSet;
+import org.apache.commons.csv.CSVFormat;
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+/** Utilities to convert between CSV records, Beam rows, and user types. */
+public final class CsvUtils {
Review Comment:
Do we want this to be public?
##########
sdks/java/io/csv/src/main/java/org/apache/beam/sdk/io/csv/CsvUtils.java:
##########
@@ -0,0 +1,250 @@
+/*
+ * 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.sdk.io.csv;
+
+import java.io.IOException;
+import java.nio.charset.StandardCharsets;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Set;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.Field;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.transforms.SimpleFunction;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableSet;
+import org.apache.commons.csv.CSVFormat;
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+/** Utilities to convert between CSV records, Beam rows, and user types. */
+public final class CsvUtils {
+
+ public static final Set<FieldType> VALID_FIELD_TYPE_SET =
+ ImmutableSet.of(
+ FieldType.BYTE,
+ FieldType.BOOLEAN,
+ FieldType.DATETIME,
+ FieldType.DECIMAL,
+ FieldType.DOUBLE,
+ FieldType.INT16,
+ FieldType.INT32,
+ FieldType.INT64,
+ FieldType.FLOAT,
+ FieldType.STRING);
+
+ /**
+ * Planned for: TODO(https://github.com/apache/beam/issues/24552) Returns a {@link SimpleFunction}
+ * that converts a CSV byte[] record to a Beam {@link Row}. Providing {@param schemaFields}
+ * determines the subset and order of {@param beamSchema} fields expected from the CSV byte[]
+ * record. Otherwise, the expected order derives from the {@link Schema#sorted()} field order.
+ */
+ public static SimpleFunction<byte[], Row> getCsvBytesToRowFunction(
+ Schema beamSchema, @Nullable CSVFormat format, @Nullable List<String> schemaFields) {
+ CsvPayloadSerializer payloadSerializer =
+ new CsvPayloadSerializer(beamSchema, format, schemaFields);
+ return new CsvBytesToRowFn(payloadSerializer);
+ }
+
+ /**
+ * Planned for: TODO(https://github.com/apache/beam/issues/24552) Returns a {@link SimpleFunction}
+ * that converts a CSV String record to a Beam {@link Row}. Providing {@param schemaFields}
+ * determines the subset and order of {@param beamSchema} fields expected from the CSV String
+ * record. Otherwise, the expected order derives from the {@link Schema#sorted()} field order.
+ */
+ public static SimpleFunction<String, Row> getCsvStringToRowFunction(
+ Schema beamSchema, @Nullable CSVFormat format, @Nullable List<String> schemaFields) {
+ CsvPayloadSerializer payloadSerializer =
+ new CsvPayloadSerializer(beamSchema, format, schemaFields);
+ return new CsvStringToRowFn(payloadSerializer);
+ }
+
+ /**
+ * Returns a {@link SimpleFunction} that converts a {@link Row} to a CSV byte[] record. Providing
+ * {@param schemaFields} determines the subset and order of {@param beamSchema} fields of the
+ * resulting CSV byte[] record. Otherwise, the order derives from the {@link Schema#sorted()}
+ * field order.
+ */
+ public static SimpleFunction<Row, byte[]> getRowToCsvBytesFunction(
+ Schema beamSchema, @Nullable CSVFormat format, @Nullable List<String> schemaFields) {
+ CsvPayloadSerializer payloadSerializer =
+ new CsvPayloadSerializer(beamSchema, format, schemaFields);
+ return new RowToCsvBytesFn(payloadSerializer);
+ }
+
+ /**
+ * Returns a {@link SimpleFunction} that converts a {@link Row} to a CSV String record. Providing
+ * {@param schemaFields} determines the subset and order of {@param beamSchema} fields of the
+ * resulting CSV String record. Otherwise, the order derives from the {@link Schema#sorted()}
+ * field order.
+ */
+ public static SimpleFunction<Row, String> getRowToCsvStringFunction(
+ Schema beamSchema, @Nullable CSVFormat format, @Nullable List<String> schemaFields) {
+
+ CsvPayloadSerializer payloadSerializer =
+ new CsvPayloadSerializer(beamSchema, format, schemaFields);
+ return new RowToCsvStringFn(payloadSerializer);
+ }
+
+ /** Formats {@param columns} into a header String based on {@link CSVFormat}. */
+ static String buildHeaderFrom(List<String> columns, CSVFormat csvFormat) {
+ StringBuilder builder = new StringBuilder();
+ try {
+ boolean newRecord = true;
+ for (String name : columns) {
+ csvFormat.print(name, builder, newRecord);
+ newRecord = false;
+ }
+ } catch (IOException e) {
+ throw new IllegalStateException(e);
+ }
+ return builder.toString();
+ }
+
+ /**
+ * Formats {@link Schema#sorted()} list of fields into a header String based on {@link CSVFormat}.
Review Comment:
```suggestion
* Formats {@link Schema#sorted()} list of fields into a header based upon {@link CSVFormat}.
```
##########
sdks/java/io/csv/src/main/java/org/apache/beam/sdk/io/csv/package-info.java:
##########
@@ -0,0 +1,20 @@
+/*
+ * 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.
+ */
+
+/** Transforms for reading and writing Csv files. */
Review Comment:
```suggestion
/** Transforms for reading and writing CSV files. */
```
##########
sdks/java/io/csv/src/main/java/org/apache/beam/sdk/io/csv/CsvIO.java:
##########
@@ -0,0 +1,384 @@
+/*
+ * 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.sdk.io.csv;
+
+import static java.nio.charset.StandardCharsets.UTF_8;
+import static org.apache.beam.sdk.values.TypeDescriptors.rows;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+
+import com.google.auto.value.AutoValue;
+import java.io.BufferedWriter;
+import java.io.IOException;
+import java.io.OutputStreamWriter;
+import java.io.PrintWriter;
+import java.nio.channels.Channels;
+import java.nio.channels.WritableByteChannel;
+import java.util.List;
+import java.util.Objects;
+import java.util.Optional;
+import org.apache.beam.sdk.io.Compression;
+import org.apache.beam.sdk.io.FileIO;
+import org.apache.beam.sdk.io.FileSystems;
+import org.apache.beam.sdk.io.fs.ResourceId;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.transforms.MapElements;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.SimpleFunction;
+import org.apache.beam.sdk.transforms.display.DisplayData;
+import org.apache.beam.sdk.transforms.display.HasDisplayData;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PDone;
+import org.apache.beam.sdk.values.Row;
+import org.apache.commons.csv.CSVFormat;
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+/** Transforms for reading and writing CSV files. */
Review Comment:
Please add:
* sample snippets
* call out supported/unsupported schema conversions
* any other notable limitations
##########
sdks/java/io/csv/src/test/java/org/apache/beam/sdk/io/csv/CsvFormatLogicalTypeTest.java:
##########
@@ -0,0 +1,47 @@
+/*
+ * 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.sdk.io.csv;
+
+import static org.junit.Assert.assertEquals;
+
+import org.apache.commons.csv.CSVFormat;
+import org.apache.commons.csv.CSVFormat.Predefined;
+import org.junit.Test;
+
+/** Tests for {@link CsvFormatLogicalType}. */
+public class CsvFormatLogicalTypeTest {
Review Comment:
```suggestion
@RunWith(JUnit4.class)
public class CsvFormatLogicalTypeTest {
```
--
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: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [beam] lukecwik commented on a diff in pull request #24630: [24469] Implement CsvIO.Write and supporting classes
Posted by "lukecwik (via GitHub)" <gi...@apache.org>.
lukecwik commented on code in PR #24630:
URL: https://github.com/apache/beam/pull/24630#discussion_r1089546389
##########
sdks/java/io/csv/src/main/java/org/apache/beam/sdk/io/csv/CsvIO.java:
##########
@@ -280,55 +281,53 @@ public Write<T> withCompression(Compression compression) {
return toBuilder().setTextIOWrite(getTextIOWrite().withCompression(compression)).build();
}
+ /** Specifies all data written without spilling, simplifying the pipeline. */
Review Comment:
Note that the parent javadoc states:
```
Whether to skip the spilling of data caused by having maxNumWritersPerBundle.
```
Maybe we could fix it and this to provide greater detail as to what this parameter actually does.
##########
sdks/java/io/csv/src/main/java/org/apache/beam/sdk/io/csv/CsvIO.java:
##########
@@ -415,28 +413,42 @@ public WriteFilesResult<String> expand(PCollection<T> input) {
PCollection<String> csv = rows.apply("To CSV", MapElements.into(strings()).via(toCsvFn));
- return csv.apply("Write CSV", getTextIOWrite().withHeader(header).withOutputFilenames());
+ return csv.apply("Write CSV", write.withOutputFilenames());
}
- CSVFormat applyRequiredCSVFormatSettings(Schema schema) {
- CSVFormat csvFormat = getCSVFormat().withSkipHeaderRecord();
+ private static CSVFormat buildHeaderFromSchemaIfNeeded(CSVFormat csvFormat, Schema schema) {
if (csvFormat.getHeader() == null) {
csvFormat = csvFormat.withHeader(schema.sorted().getFieldNames().toArray(new String[0]));
}
+
return csvFormat;
}
- private static String formatHeader(CSVFormat csvFormat) {
+ private static TextIO.Write writeWithCSVFormatHeaderAndComments(
+ CSVFormat csvFormat, TextIO.Write write) {
+
+ if (csvFormat.getSkipHeaderRecord()) {
+ return write;
+ }
+
String[] header = requireNonNull(csvFormat.getHeader());
List<String> result = new ArrayList<>();
if (csvFormat.getHeaderComments() != null) {
for (String comment : csvFormat.getHeaderComments()) {
result.add(csvFormat.getCommentMarker() + " " + comment);
}
}
+
CSVFormat withoutHeaderComments = csvFormat.withHeaderComments();
Review Comment:
nit: in follow-up did you mean to name this `withHeaderComments`?
--
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: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [beam] github-actions[bot] commented on pull request #24630: [24469] Implement CsvIO.Write and supporting classes
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on PR #24630:
URL: https://github.com/apache/beam/pull/24630#issuecomment-1369700879
Assigning new set of reviewers because Pr has gone too long without review. If you would like to opt out of this review, comment `assign to next reviewer`:
R: @kennknowles for label java.
R: @damccorm for label build.
R: @johnjcasey for label io.
Available commands:
- `stop reviewer notifications` - opt out of the automated review tooling
- `remind me after tests pass` - tag the comment author after tests pass
- `waiting on author` - shift the attention set back to the author (any comment or push by the author will return the attention set to the reviewers)
--
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: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org