Add support for dumping to file

Configurable using 'dumpdir', 'dumpprefix' and 'dumponserver'.

Author: joerivanruth

README.md

@@ -141,6 +141,21 @@ JDBC data source. The following configuration options are supported:
 * **immediatecommit**: commit immediately after each batch instead of only at
   the end.
 
+* **dumpdir**: if this is set, the data is written to this directory instead of
+  being uploaded to MonetDB. Every partition will create a subdirectory containing
+  one sql file `copy.sql` with the COPY INTO statement, plus one binary file
+  per column.
+
+  Note: A **url** must still be provided. It will be used to determine the
+  exact data types.
+  
+* **dumpprefix**: used with **dumpdir**. By default, the `copy.sql` generated by
+  **dumpdir** will refer to the binary files with an absolute path based on **dumpdir**.
+  **dumpprefix** can be used to override that.
+
+* **dumponserver**: use `ON SERVER` in `copy.sql` rather than `ON CLIENT`.
+  Default 'false'.
+
 ## Building
 
 To build, simply run `make`. This will create `monetdb-spark-X.Y.Z-fat.jar`

monetdb-spark/src/main/java/org/monetdb/spark/bincopy/BinCopyDataWriterFactory.java

@@ -22,6 +22,7 @@
 import org.monetdb.spark.workerside.MonetDataWriter;
 import org.monetdb.spark.workerside.Step;
 
+import java.io.IOException;
 import java.io.Serial;
 import java.io.Serializable;
 import java.sql.SQLException;
@@ -49,10 +50,10 @@ public DataWriter<InternalRow> createWriter(int partitionId, long taskId) {
 			String identifier = "part" + partitionId + "-task" + taskId;
 			Collector collector = new Collector();
 			sqlstmt.identifier(identifier);
-			BinCopyUploader uploader = new BinCopyUploader(parms.getDestination(), collector, sqlstmt);
+			Uploader uploader = parms.getDumpdir() == null ? new BinCopyUploader(parms.getDestination(), collector, sqlstmt) : new BinCopyFileDump(parms.getDumpdir(), parms.getDumpPrefix(), collector, sqlstmt, partitionId, taskId);
 			collector.registerWithConverters(steps);
 			return new MonetDataWriter(collector, steps, uploader, parms.isImmediateCommit(), identifier, parms.getBatchSize());
-		} catch (SQLException e) {
+		} catch (SQLException | IOException e) {
 			throw new RuntimeException(e);
 		}
 	}

monetdb-spark/src/main/java/org/monetdb/spark/bincopy/BinCopyFileDump.java

@@ -0,0 +1,173 @@
+/*
+ * SPDX-License-Identifier: MPL-2.0
+ *
+ * This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0.  If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/.
+ *
+ * Copyright MonetDB Solutions B.V.
+ */
+
+package org.monetdb.spark.bincopy;
+
+import org.jetbrains.annotations.NotNull;
+import org.monetdb.spark.workerside.Collector;
+
+import java.io.File;
+import java.io.IOException;
+import java.io.OutputStream;
+import java.nio.charset.StandardCharsets;
+import java.nio.file.FileAlreadyExistsException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+import java.sql.SQLException;
+import java.util.ArrayList;
+
+public class BinCopyFileDump implements Uploader {
+	private final Path destDir;
+	private final String copyFileName;
+	private final String dataFilePrefix;
+	private final String dataFileSuffix;
+	private final Collector collector;
+	private final BinCopySql sqlstmt;
+	private final ArrayList<OutputStream> outs;
+	private boolean closed;
+	private Runnable onStartUpload;
+	private Runnable onEndUpload;
+
+	public BinCopyFileDump(String dumpdir, String dumpPrefix, Collector collector, BinCopySql sqlstmt, int partitionId, long taskId) throws IOException {
+		this.collector = collector;
+		this.sqlstmt = sqlstmt;
+		String part = "part" + partitionId;
+		String suffix = "tmp" + taskId;
+
+		Path topDir = Paths.get(dumpdir);
+		destDir = topDir.resolve(part);
+		copyFileName = "copy.sql";
+		dataFilePrefix = "col";
+		dataFileSuffix = ".bin";
+
+		boolean haveDumpPrefix = dumpPrefix != null && !dumpPrefix.isEmpty();
+		String prefix = (haveDumpPrefix ? dumpPrefix : ".") + File.separator + part + File.separator + dataFilePrefix;
+		sqlstmt.nameMapper(i -> getDataFileName(i, prefix, dataFileSuffix));
+
+		try {
+			Files.createDirectory(topDir);
+		} catch (FileAlreadyExistsException ignored) {
+		}
+		Files.createDirectory(destDir);
+		String sql = sqlstmt.toString();
+		Files.writeString(getCopyFilePath(), sql + ";", StandardCharsets.UTF_8);
+
+		outs = new ArrayList<>(sqlstmt.getColumnCount());
+		openOutFiles();
+	}
+
+	private @NotNull Path getCopyFilePath() {
+		return destDir.resolve(copyFileName);
+	}
+
+	private @NotNull Path getDataFilePath(int i) {
+		String name = getDataFileName(i, dataFilePrefix, dataFileSuffix);
+		return destDir.resolve(name);
+	}
+
+	private @NotNull String getDataFileName(int i, String prefix, String suffix) {
+		StringBuilder builder = new StringBuilder(prefix).append(i).append(".");
+		for (char c : sqlstmt.getColumnName(i).toCharArray()) {
+			if (!Character.isAlphabetic(c) && !Character.isDigit(c) && c != '_')
+				c = '~';
+			builder.append(c);
+		}
+		builder.append(suffix);
+		return builder.toString();
+	}
+
+	private void openOutFiles() throws IOException {
+		int n = sqlstmt.getColumnCount();
+		int i = 0;
+		try {
+			while (i < n) {
+				Path p = getDataFilePath(i);
+				OutputStream s = Files.newOutputStream(p);
+				outs.add(s);
+				i++;
+			}
+		} catch (IOException e) {
+			closeOutputFiles(true);
+		}
+	}
+
+	private void closeOutputFiles(boolean ignoreErrors) throws IOException {
+		if (closed)
+			return;
+		IOException toThrow = null;
+		for (int i = 0; i < outs.size(); i++) {
+			OutputStream s = outs.get(i);
+			if (s != null) {
+				try {
+					s.close();
+				} catch (IOException e) {
+					if (toThrow == null)
+						toThrow = e;
+				}
+				outs.set(i, null);
+			}
+		}
+		closed = true;
+		if (toThrow != null && !ignoreErrors)
+			throw toThrow;
+	}
+
+	@Override
+	public void uploadBatch() throws SQLException, IOException {
+		if (onStartUpload != null)
+			onStartUpload.run();
+		try {
+			for (int i = 0; i < outs.size(); i++) {
+				collector.writeTo(i, outs.get(i));
+			}
+		} finally {
+			if (onEndUpload != null)
+				onEndUpload.run();
+		}
+	}
+
+	@Override
+	public void commit() throws SQLException, IOException {
+		for (OutputStream s : outs) {
+			s.flush();
+		}
+		closeOutputFiles(false);
+	}
+
+	@Override
+	public void close() throws IOException {
+		if (closed) // was committed
+			return;
+		closeOutputFiles(true);
+		silentDeleteFile(getCopyFilePath());
+		for (int i = 0; i < sqlstmt.getColumnCount(); i++)
+			silentDeleteFile(getDataFilePath(i));
+		Files.delete(destDir);
+	}
+
+	private void silentDeleteFile(@NotNull Path p) {
+		try {
+			Files.delete(p);
+		} catch (IOException ignored) {
+		}
+	}
+
+	@Override
+	public void setOnStartUpload(Runnable callback) {
+		onStartUpload = callback;
+	}
+
+	@Override
+	public void setOnEndUpload(Runnable callback) {
+		onEndUpload = callback;
+
+	}
+}

monetdb-spark/src/main/java/org/monetdb/spark/bincopy/BinCopySql.java

@@ -13,14 +13,14 @@
 import java.io.PrintWriter;
 import java.io.Serializable;
 import java.io.StringWriter;
+import java.util.function.IntFunction;
 
 public class BinCopySql implements Serializable {
-	String quotedTableName;
-	String[] columnNames;
-	String identifier;
-	String prefix = "";
-	String suffix = "";
-	boolean onClient = true;
+	private final String quotedTableName;
+	private final String[] columnNames;
+	private String identifier;
+	private boolean onClient = true;
+	private transient IntFunction<String> nameMapper;
 
 	public BinCopySql(String quotedTableName, String[] columnNames) {
 		this.quotedTableName = quotedTableName;
@@ -32,13 +32,8 @@ public BinCopySql identifier(String identifier) {
 		return this;
 	}
 
-	public BinCopySql suffix(String suffix) {
-		this.suffix = suffix;
-		return this;
-	}
-
-	public BinCopySql prefix(String prefix) {
-		this.prefix = prefix;
+	public BinCopySql nameMapper(IntFunction<String> f) {
+		nameMapper = f;
 		return this;
 	}
 
@@ -51,6 +46,14 @@ public BinCopySql onServer(boolean onServer) {
 		return onClient(!onServer);
 	}
 
+	public int getColumnCount() {
+		return columnNames.length;
+	}
+
+	public String getColumnName(int i) {
+		return columnNames[i];
+	}
+
 	@Override
 	public String toString() {
 		String sep;
@@ -67,22 +70,23 @@ public String toString() {
 		}
 		pw.println(")");
 
-		pw.print("FROM ");
-		final String qprefix, qsuffix;
-		if (prefix.contains("\\") || prefix.contains("'") || suffix.contains("\\") || suffix.contains("'")) {
-			qprefix = "R'" + prefix.replace("'", "''");
-			qsuffix = suffix.replace("'", "''") + "'";
-		} else {
-			qprefix = "'" + prefix;
-			qsuffix = suffix + "'";
-		}
-		sep = "";
+		pw.println("FROM");
+		sep = "\t";
 		for (int i = 0; i < columnNames.length; i++) {
-			pw.printf("%s%s%d%s", sep, qprefix, i, qsuffix);
-			sep = ", ";
+			String name = nameMapper != null ? nameMapper.apply(i) : Integer.toString(i);
+			String escaped = escape(name);
+			pw.printf("%s%s", sep, escaped);
+			sep = ",\n\t";
 		} pw.println();
 		pw.println(onClient ? "ON CLIENT" : "ON SERVER");
 
 		return sw.toString();
 	}
+
+	private String escape(String s) {
+		if (s.contains("'") || s.contains("\\"))
+			return "R'" + s.replace("'", "''") + "'";
+		else
+			return "'" + s + "'";
+	}
 }

monetdb-spark/src/main/java/org/monetdb/spark/bincopy/BinCopyUploader.java

@@ -17,7 +17,7 @@
 import java.sql.PreparedStatement;
 import java.sql.SQLException;
 
-public class BinCopyUploader {
+public class BinCopyUploader implements Uploader {
 	private final Collector collector;
 	private final MonetConnection conn;
 	private final PreparedStatement stmt;
@@ -33,27 +33,32 @@ public BinCopyUploader(Destination dest, Collector collector, BinCopySql sqlstmt
 		stmt = conn.prepareStatement(sql);
 	}
 
+	@Override
 	public void uploadBatch() throws SQLException {
 		if (collector.getRowCount() == 0)
 			return;
 		stmt.execute();
 		collector.clear();
 	}
 
+	@Override
 	public void commit() throws SQLException {
 		uploadBatch();
 		conn.commit();
 	}
 
+	@Override
 	public void close() throws SQLException {
 		stmt.close();
 		conn.close();
 	}
 
+	@Override
 	public void setOnStartUpload(Runnable callback) {
 		collector.setOnStartUpload(callback);
 	}
 
+	@Override
 	public void setOnEndUpload(Runnable callback) {
 		collector.setOnEndUpload(callback);
 	}

monetdb-spark/src/main/java/org/monetdb/spark/bincopy/Uploader.java

@@ -0,0 +1,26 @@
+/*
+ * SPDX-License-Identifier: MPL-2.0
+ *
+ * This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0.  If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/.
+ *
+ * Copyright MonetDB Solutions B.V.
+ */
+
+package org.monetdb.spark.bincopy;
+
+import java.io.IOException;
+import java.sql.SQLException;
+
+public interface Uploader {
+	void uploadBatch() throws SQLException, IOException;
+
+	void commit() throws SQLException, IOException;
+
+	void close() throws SQLException, IOException;
+
+	void setOnStartUpload(Runnable callback);
+
+	void setOnEndUpload(Runnable callback);
+}