Add benchmark for bulk ingestion

benchmark/graph.rb

@@ -26,14 +26,20 @@
   "integer",
   "string",
 ]
+types = [
+  "insert",
+  "select",
+]
 benchmarks.each do |benchmark|
   base_dir = File.join(__dir__, benchmark)
-  data = CSV.read(File.join(base_dir, "result.csv"),
-                  headers: true,
-                  converters: :all)
-  plotter = Charty.bar_plot(data: data,
-                            x: "N records",
-                            y: "Elapsed time (sec)",
-                            color: "Approach")
-  plotter.save(File.join(base_dir, "result.svg"))
+  types.each do |type|
+    data = CSV.read(File.join(base_dir, "#{type}.csv"),
+                    headers: true,
+                    converters: :all)
+    plotter = Charty.bar_plot(data: data,
+                              x: "N records",
+                              y: "Elapsed time (sec)",
+                              color: "Approach")
+    plotter.save(File.join(base_dir, "#{type}.svg"))
+  end
 end

benchmark/insert-copy.cc

@@ -0,0 +1,152 @@
+/*
+ * 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.
+ */
+
+#include <chrono>
+#include <cstdlib>
+#include <iostream>
+#include <vector>
+
+#include <libpq-fe.h>
+
+class ConnectionFinisher {
+   public:
+	ConnectionFinisher(PGconn* connection) : connection_(connection) {}
+	~ConnectionFinisher() { PQfinish(connection_); }
+
+   private:
+	PGconn* connection_;
+};
+
+class ResultClearner {
+   public:
+	ResultClearner(PGresult* result) : result_(result) {}
+	~ResultClearner() { PQclear(result_); }
+
+   private:
+	PGresult* result_;
+};
+
+int
+main(int argc, char** argv)
+{
+	std::string connectionString;
+	if (!std::getenv("PGDATABASE"))
+	{
+		connectionString += "dbname=afs_benchmark";
+	}
+	auto connection = PQconnectdb(connectionString.c_str());
+	ConnectionFinisher connectionFinisher(connection);
+	if (PQstatus(connection) != CONNECTION_OK)
+	{
+		std::cerr << "failed to connect: " << PQerrorMessage(connection) << std::endl;
+		return EXIT_FAILURE;
+	}
+
+	{
+		auto result = PQexec(connection, "DROP TABLE IF EXISTS data_insert");
+		ResultClearner resultClearner(result);
+		if (PQresultStatus(result) != PGRES_COMMAND_OK)
+		{
+			std::cerr << "failed to drop: " << PQerrorMessage(connection) << std::endl;
+			return EXIT_FAILURE;
+		}
+	}
+
+	{
+		auto result = PQexec(connection, "CREATE TABLE data_insert (LIKE data)");
+		ResultClearner resultClearner(result);
+		if (PQresultStatus(result) != PGRES_COMMAND_OK)
+		{
+			std::cerr << "failed to create: " << PQerrorMessage(connection) << std::endl;
+			return EXIT_FAILURE;
+		}
+	}
+
+	std::vector<std::string> buffers;
+	{
+		auto result = PQexec(connection, "COPY data TO STDOUT (FORMAT binary)");
+		ResultClearner resultClearner(result);
+		if (PQresultStatus(result) != PGRES_COPY_OUT)
+		{
+			std::cerr << "failed to copy to: " << PQerrorMessage(connection) << std::endl;
+			return EXIT_FAILURE;
+		}
+		while (true)
+		{
+			char* data;
+			auto size = PQgetCopyData(connection, &data, 0);
+			if (size == -1)
+			{
+				break;
+			}
+			if (size == -2)
+			{
+				std::cerr << "failed to read copy data: " << PQerrorMessage(connection)
+						  << std::endl;
+				return EXIT_FAILURE;
+			}
+			buffers.emplace_back(data, size);
+			free(data);
+		}
+	}
+	auto before = std::chrono::steady_clock::now();
+	{
+		auto result = PQexec(connection, "COPY data_insert FROM STDOUT (FORMAT binary)");
+		ResultClearner resultClearner(result);
+		if (PQresultStatus(result) != PGRES_COPY_IN)
+		{
+			std::cerr << "failed to copy from: " << PQerrorMessage(connection)
+					  << std::endl;
+			return EXIT_FAILURE;
+		}
+		for (const auto& buffer : buffers)
+		{
+			auto copyDataResult = PQputCopyData(connection, buffer.data(), buffer.size());
+			if (copyDataResult == -1)
+			{
+				std::cerr << "failed to put copy data: " << PQerrorMessage(connection)
+						  << std::endl;
+				return EXIT_FAILURE;
+			}
+		}
+		auto copyEndResult = PQputCopyEnd(connection, nullptr);
+		if (copyEndResult == -1)
+		{
+			std::cerr << "failed to end copy data: " << PQerrorMessage(connection)
+					  << std::endl;
+			return EXIT_FAILURE;
+		}
+	}
+	{
+		auto result = PQgetResult(connection);
+		ResultClearner resultClearner(result);
+		if (PQresultStatus(result) != PGRES_COMMAND_OK)
+		{
+			std::cerr << "failed to copy from: " << PQerrorMessage(connection)
+					  << std::endl;
+			return EXIT_FAILURE;
+		}
+	}
+	auto elapsedTime = std::chrono::duration_cast<std::chrono::milliseconds>(
+						   std::chrono::steady_clock::now() - before)
+	                       .count();
+	printf("%.3f\n", elapsedTime / 1000.0);
+
+	return EXIT_SUCCESS;
+}

benchmark/insert-flight-sql.rb

@@ -0,0 +1,54 @@
+#!/usr/bin/env ruby
+#
+# 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.
+
+require "time"
+
+require "arrow-flight-sql"
+
+call_options = ArrowFlight::CallOptions.new
+call_options.add_header("x-flight-sql-database",
+                        ENV["PGDATABASE"] || "afs_benchmark")
+client = ArrowFlight::Client.new("grpc://127.0.0.1:15432")
+client.authenticate_basic(ENV["PGUSER"] || ENV["USER"],
+                          ENV["PGPASSWORD"] || "",
+                          call_options)
+sql_client = ArrowFlightSQL::Client.new(client)
+
+sql_client.execute_update("DROP TABLE IF EXISTS data_insert", call_options)
+sql_client.execute_update("CREATE TABLE data_insert (LIKE data)", call_options)
+
+info = sql_client.execute("SELECT * FROM data", call_options)
+endpoint = info.endpoints.first
+reader = sql_client.do_get(endpoint.ticket, call_options)
+record_batches = []
+loop do
+  record_batch = reader.read_next
+  break if record_batch.nil?
+  record_batches << record_batch.data
+end
+
+before = Time.now
+sql_client.prepare("INSERT INTO data_insert VALUES ($1)",
+                   call_options) do |statement|
+  record_batches.each do |record_batch|
+    statement.record_batch = record_batch
+    statement.execute_update(call_options)
+  end
+end
+puts("%.3f" % (Time.now - before))

benchmark/insert.cc

@@ -0,0 +1,139 @@
+/*
+ * 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.
+ */
+
+#include <chrono>
+#include <cstdlib>
+#include <iostream>
+
+#include <libpq-fe.h>
+
+#include <catalog/pg_type_d.h>
+
+class ConnectionFinisher {
+   public:
+	ConnectionFinisher(PGconn* connection) : connection_(connection) {}
+	~ConnectionFinisher() { PQfinish(connection_); }
+
+   private:
+	PGconn* connection_;
+};
+
+class ResultClearner {
+   public:
+	ResultClearner(PGresult* result) : result_(result) {}
+	~ResultClearner() { PQclear(result_); }
+
+   private:
+	PGresult* result_;
+};
+
+int
+main(int argc, char** argv)
+{
+	std::string connectionString;
+	if (!std::getenv("PGDATABASE"))
+	{
+		connectionString += "dbname=afs_benchmark";
+	}
+	auto connection = PQconnectdb(connectionString.c_str());
+	ConnectionFinisher connectionFinisher(connection);
+	if (PQstatus(connection) != CONNECTION_OK)
+	{
+		std::cerr << "failed to connect: " << PQerrorMessage(connection) << std::endl;
+		return EXIT_FAILURE;
+	}
+
+	{
+		auto result = PQexec(connection, "DROP TABLE IF EXISTS data_insert");
+		ResultClearner resultClearner(result);
+		if (PQresultStatus(result) != PGRES_COMMAND_OK)
+		{
+			std::cerr << "failed to drop: " << PQerrorMessage(connection) << std::endl;
+			return EXIT_FAILURE;
+		}
+	}
+
+	{
+		auto result = PQexec(connection, "CREATE TABLE data_insert (LIKE data)");
+		ResultClearner resultClearner(result);
+		if (PQresultStatus(result) != PGRES_COMMAND_OK)
+		{
+			std::cerr << "failed to create: " << PQerrorMessage(connection) << std::endl;
+			return EXIT_FAILURE;
+		}
+	}
+
+	std::string insert = "INSERT INTO data_insert VALUES ";
+	{
+		auto result = PQexec(connection, "SELECT * FROM data");
+		ResultClearner resultClearner(result);
+		if (PQresultStatus(result) != PGRES_TUPLES_OK)
+		{
+			std::cerr << "failed to select: " << PQerrorMessage(connection) << std::endl;
+			return EXIT_FAILURE;
+		}
+		auto nTuples = PQntuples(result);
+		auto nFields = PQnfields(result);
+		for (int iTuple = 0; iTuple < nTuples; iTuple++)
+		{
+			if (iTuple > 0)
+			{
+				insert += ", ";
+			}
+			for (int iField = 0; iField < nFields; iField++)
+			{
+				if (PQgetisnull(result, iTuple, iField))
+				{
+					insert += "(null)";
+				}
+				else
+				{
+					insert += "(";
+					auto type = PQftype(result, iField);
+					if (type == TEXTOID)
+					{
+						insert += "'";
+					}
+					insert += PQgetvalue(result, iTuple, iField);
+					if (type == TEXTOID)
+					{
+						insert += "'";
+					}
+					insert += ")";
+				}
+			}
+		}
+	}
+	auto before = std::chrono::steady_clock::now();
+	{
+		auto result = PQexec(connection, insert.c_str());
+		ResultClearner resultClearner(result);
+		if (PQresultStatus(result) != PGRES_COMMAND_OK)
+		{
+			std::cerr << "failed to insert: " << PQerrorMessage(connection) << std::endl;
+			return EXIT_FAILURE;
+		}
+	}
+	auto elapsedTime = std::chrono::duration_cast<std::chrono::milliseconds>(
+						   std::chrono::steady_clock::now() - before)
+	                       .count();
+	printf("%.3f\n", elapsedTime / 1000.0);
+
+	return EXIT_SUCCESS;
+}