Refactor jni writer data sink

cpp/include/cudf/io/csv.hpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2020-2022, NVIDIA CORPORATION.
+ * Copyright (c) 2020-2023, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -1332,7 +1332,7 @@ class csv_writer_options {
   size_type _rows_per_chunk = std::numeric_limits<size_type>::max();
   // character to use for separating lines (default "\n")
   std::string _line_terminator = "\n";
-  // character to use for separating lines (default "\n")
+  // character to use for separating column values (default ",")
   char _inter_column_delimiter = ',';
   // string to use for values != 0 in INT8 types (default 'true')
   std::string _true_value = std::string{"true"};
@@ -1422,9 +1422,9 @@ class csv_writer_options {
   [[nodiscard]] std::string get_line_terminator() const { return _line_terminator; }
 
   /**
-   * @brief Returns character used for separating lines.
+   * @brief Returns character used for separating column values.
    *
-   * @return Character used for separating lines
+   * @return Character used for separating column values.
    */
   [[nodiscard]] char get_inter_column_delimiter() const { return _inter_column_delimiter; }
 
@@ -1479,9 +1479,9 @@ class csv_writer_options {
   void set_line_terminator(std::string term) { _line_terminator = term; }
 
   /**
-   * @brief Sets character used for separating lines.
+   * @brief Sets character used for separating column values.
    *
-   * @param delim Character to indicate delimiting
+   * @param delim Character to delimit column values
    */
   void set_inter_column_delimiter(char delim) { _inter_column_delimiter = delim; }
 
@@ -1498,6 +1498,13 @@ class csv_writer_options {
    * @param val String to represent values == 0 in INT8 types
    */
   void set_false_value(std::string val) { _false_value = val; }
+
+  /**
+   * @brief (Re)sets the table being written.
+   *
+   * @param table Table to be written
+   */
+  void set_table(table_view const& table) { _table = table; }
 };
 
 /**
@@ -1586,9 +1593,9 @@ class csv_writer_options_builder {
   }
 
   /**
-   * @brief Sets character used for separating lines.
+   * @brief Sets character used for separating column values.
    *
-   * @param delim Character to indicate delimiting
+   * @param delim Character to delimit column values
    * @return this for chaining
    */
   csv_writer_options_builder& inter_column_delimiter(char delim)

java/src/main/java/ai/rapids/cudf/CSVWriterOptions.java

@@ -0,0 +1,134 @@
+/*
+ *
+ *  Copyright (c) 2023, NVIDIA CORPORATION.
+ *
+ *  Licensed 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 ai.rapids.cudf;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+
+public class CSVWriterOptions {
+
+  private String[] columnNames;
+  private Boolean includeHeader = false;
+  private String rowDelimiter = "\n";
+  private byte fieldDelimiter = ',';
+  private String nullValue = "\\N";
+  private String falseValue = "false";
+  private String trueValue = "true";
+
+  private CSVWriterOptions(Builder builder) {
+    this.columnNames = builder.columnNames.toArray(new String[builder.columnNames.size()]);
+    this.nullValue = builder.nullValue;
+    this.includeHeader = builder.includeHeader;
+    this.fieldDelimiter = builder.fieldDelimiter;
+    this.rowDelimiter = builder.rowDelimiter;
+    this.falseValue = builder.falseValue;
+    this.trueValue = builder.trueValue;
+  }
+
+  public String[] getColumnNames() {
+    return columnNames;
+  }
+
+  public Boolean getIncludeHeader() {
+    return includeHeader;
+  }
+
+  public String getRowDelimiter() {
+    return rowDelimiter;
+  }
+
+  public byte getFieldDelimiter() {
+    return fieldDelimiter;
+  }
+
+  public String getNullValue() {
+    return nullValue;
+  }
+
+  public String getTrueValue() {
+    return trueValue;
+  }
+
+  public String getFalseValue() {
+    return falseValue;
+  }
+
+  public static Builder builder() {
+    return new Builder();
+  }
+
+  public static class Builder {
+
+    private List<String> columnNames = Collections.emptyList();
+    private Boolean includeHeader = false;
+    private String rowDelimiter = "\n";
+    private byte fieldDelimiter = ',';
+    private String nullValue = "\\N";
+    private String falseValue = "false";
+    private String trueValue = "true";
+
+    public CSVWriterOptions build() {
+      return new CSVWriterOptions(this);
+    }
+
+    public Builder withColumnNames(List<String> columnNames) {
+      this.columnNames = columnNames;
+      return this;
+    }
+
+    public Builder withColumnNames(String... columnNames) {
+      List<String> columnNamesList = new ArrayList<>();
+      for (String columnName : columnNames) {
+        columnNamesList.add(columnName);
+      }
+      return withColumnNames(columnNamesList);
+    }
+
+    public Builder withIncludeHeader(Boolean includeHeader) {
+      this.includeHeader = includeHeader;
+      return this;
+    }
+
+    public Builder withRowDelimiter(String rowDelimiter) {
+      this.rowDelimiter = rowDelimiter;
+      return this;
+    }
+
+    public Builder withFieldDelimiter(byte fieldDelimiter) {
+      this.fieldDelimiter = fieldDelimiter;
+      return this;
+    }
+
+    public Builder withNullValue(String nullValue) {
+      this.nullValue = nullValue;
+      return this;
+    }
+
+    public Builder withTrueValue(String trueValue) {
+      this.trueValue = trueValue;
+      return this;
+    }
+
+    public Builder withFalseValue(String falseValue) {
+      this.falseValue = falseValue;
+      return this;
+    }
+  }
+}

java/src/main/java/ai/rapids/cudf/Table.java

@@ -1,6 +1,6 @@
 /*
  *
- *  Copyright (c) 2019-2022, NVIDIA CORPORATION.
+ *  Copyright (c) 2019-2023, NVIDIA CORPORATION.
  *
  *  Licensed under the Apache License, Version 2.0 (the "License");
  *  you may not use this file except in compliance with the License.
@@ -857,6 +857,82 @@ public static Table readCSV(Schema schema, CSVOptions opts, HostMemoryBuffer buf
         opts.getFalseValues()));
   }
 
+  private static native void writeCSVToFile(long table,
+                                            String[] columnNames,
+                                            boolean includeHeader,
+                                            String rowDelimiter,
+                                            byte fieldDelimiter,
+                                            String nullValue,
+                                            String trueValue,
+                                            String falseValue,
+                                            String outputPath) throws CudfException;
+
+  public void writeCSVToFile(CSVWriterOptions options, String outputPath) {
+    writeCSVToFile(nativeHandle,
+                   options.getColumnNames(),
+                   options.getIncludeHeader(),
+                   options.getRowDelimiter(),
+                   options.getFieldDelimiter(),
+                   options.getNullValue(),
+                   options.getTrueValue(),
+                   options.getFalseValue(),
+                   outputPath);
+  }
+
+  private static native long startWriteCSVToBuffer(String[] columnNames,
+                                                   boolean includeHeader,
+                                                   String rowDelimiter,
+                                                   byte fieldDelimiter,
+                                                   String nullValue,
+                                                   String trueValue,
+                                                   String falseValue,
+                                                   HostBufferConsumer buffer) throws CudfException;
+
+  private static native void writeCSVChunkToBuffer(long writerHandle, long tableHandle);
+
+  private static native void endWriteCSVToBuffer(long writerHandle);
+
+  private static class CSVTableWriter implements TableWriter {
+    private long writerHandle;
+    private HostBufferConsumer consumer;
+
+    private CSVTableWriter(CSVWriterOptions options, HostBufferConsumer consumer) {
+      this.writerHandle = startWriteCSVToBuffer(options.getColumnNames(),
+                                                options.getIncludeHeader(),
+                                                options.getRowDelimiter(),
+                                                options.getFieldDelimiter(),
+                                                options.getNullValue(),
+                                                options.getTrueValue(),
+                                                options.getFalseValue(),
+                                                consumer);
+      this.consumer = consumer;
+    }
+
+    @Override
+    public void write(Table table) {
+      if (writerHandle == 0) {
+        throw new IllegalStateException("Writer was already closed");
+      }
+      writeCSVChunkToBuffer(writerHandle, table.nativeHandle);
+    }
+
+    @Override
+    public void close() throws CudfException {
+      if (writerHandle != 0) {
+        endWriteCSVToBuffer(writerHandle);
+        writerHandle = 0;
+      }
+      if (consumer != null) {
+        consumer.done();
+        consumer = null;
+      }
+    }
+  }
+
+  public static TableWriter getCSVBufferWriter(CSVWriterOptions options, HostBufferConsumer bufferConsumer) {
+    return new CSVTableWriter(options, bufferConsumer);
+  }
+
   /**
    * Read a JSON file using the default JSONOptions.
    * @param schema the schema of the file.  You may use Schema.INFERRED to infer the schema.