in c/driver/postgresql/postgresql_benchmark.cc [40:161]
static void BM_PostgresqlExecute(benchmark::State& state) {
const char* uri = std::getenv("ADBC_POSTGRESQL_TEST_URI");
if (!uri || !strcmp(uri, "")) {
state.SkipWithError("ADBC_POSTGRESQL_TEST_URI not set!");
return;
}
adbc_validation::Handle<struct AdbcDatabase> database;
struct AdbcError error;
ADBC_BENCHMARK_RETURN_NOT_OK(AdbcDatabaseNew(&database.value, &error));
ADBC_BENCHMARK_RETURN_NOT_OK(
AdbcDatabaseSetOption(&database.value, "uri", uri, &error));
ADBC_BENCHMARK_RETURN_NOT_OK(AdbcDatabaseInit(&database.value, &error));
adbc_validation::Handle<struct AdbcConnection> connection;
ADBC_BENCHMARK_RETURN_NOT_OK(AdbcConnectionNew(&connection.value, &error));
ADBC_BENCHMARK_RETURN_NOT_OK(
AdbcConnectionInit(&connection.value, &database.value, &error));
adbc_validation::Handle<struct AdbcStatement> statement;
ADBC_BENCHMARK_RETURN_NOT_OK(
AdbcStatementNew(&connection.value, &statement.value, &error));
const char* drop_query = "DROP TABLE IF EXISTS adbc_postgresql_ingest_benchmark";
ADBC_BENCHMARK_RETURN_NOT_OK(
AdbcStatementSetSqlQuery(&statement.value, drop_query, &error));
ADBC_BENCHMARK_RETURN_NOT_OK(
AdbcStatementExecuteQuery(&statement.value, nullptr, nullptr, &error));
adbc_validation::Handle<struct ArrowSchema> schema;
adbc_validation::Handle<struct ArrowArray> array;
struct ArrowError na_error;
ADBC_BENCHMARK_RETURN_NOT_OK(
adbc_validation::MakeSchema(&schema.value, {
{"bools", NANOARROW_TYPE_BOOL},
{"int16s", NANOARROW_TYPE_INT16},
{"int32s", NANOARROW_TYPE_INT32},
{"int64s", NANOARROW_TYPE_INT64},
{"floats", NANOARROW_TYPE_FLOAT},
{"doubles", NANOARROW_TYPE_DOUBLE},
}));
if (ArrowArrayInitFromSchema(&array.value, &schema.value, &na_error) != NANOARROW_OK) {
state.SkipWithError("Call to ArrowArrayInitFromSchema failed!");
error.release(&error);
return;
}
if (ArrowArrayStartAppending(&array.value) != NANOARROW_OK) {
state.SkipWithError("Call to ArrowArrayStartAppending failed!");
error.release(&error);
return;
}
const size_t n_zeros = 1000;
const size_t n_ones = 1000;
for (size_t i = 0; i < n_zeros; i++) {
// assumes fixed size primitive layouts for now
ArrowBufferAppendInt8(ArrowArrayBuffer(array.value.children[0], 1), 0);
ArrowBufferAppendInt16(ArrowArrayBuffer(array.value.children[1], 1), 0);
ArrowBufferAppendInt32(ArrowArrayBuffer(array.value.children[2], 1), 0);
ArrowBufferAppendInt64(ArrowArrayBuffer(array.value.children[3], 1), 0);
ArrowBufferAppendFloat(ArrowArrayBuffer(array.value.children[4], 1), 0.0);
ArrowBufferAppendDouble(ArrowArrayBuffer(array.value.children[5], 1), 0.0);
}
for (size_t i = 0; i < n_ones; i++) {
// assumes fixed size primitive layouts for now
ArrowBufferAppendInt8(ArrowArrayBuffer(array.value.children[0], 1), 1);
ArrowBufferAppendInt16(ArrowArrayBuffer(array.value.children[1], 1), 1);
ArrowBufferAppendInt32(ArrowArrayBuffer(array.value.children[2], 1), 1);
ArrowBufferAppendInt64(ArrowArrayBuffer(array.value.children[3], 1), 1);
ArrowBufferAppendFloat(ArrowArrayBuffer(array.value.children[4], 1), 1.0);
ArrowBufferAppendDouble(ArrowArrayBuffer(array.value.children[5], 1), 1.0);
}
for (int64_t i = 0; i < array.value.n_children; i++) {
array.value.children[i]->length = n_zeros + n_ones;
}
array.value.length = n_zeros + n_ones;
if (ArrowArrayFinishBuildingDefault(&array.value, &na_error) != NANOARROW_OK) {
state.SkipWithError("Call to ArrowArrayFinishBuildingDefault failed");
error.release(&error);
return;
}
const char* create_query =
"CREATE TABLE adbc_postgresql_ingest_benchmark (bools BOOLEAN, int16s SMALLINT, "
"int32s INTEGER, int64s BIGINT, floats REAL, doubles DOUBLE PRECISION)";
ADBC_BENCHMARK_RETURN_NOT_OK(
AdbcStatementSetSqlQuery(&statement.value, create_query, &error));
ADBC_BENCHMARK_RETURN_NOT_OK(
AdbcStatementExecuteQuery(&statement.value, nullptr, nullptr, &error));
adbc_validation::Handle<struct AdbcStatement> insert_stmt;
ADBC_BENCHMARK_RETURN_NOT_OK(
AdbcStatementNew(&connection.value, &insert_stmt.value, &error));
ADBC_BENCHMARK_RETURN_NOT_OK(
AdbcStatementSetOption(&insert_stmt.value, ADBC_INGEST_OPTION_TARGET_TABLE,
"adbc_postgresql_ingest_benchmark", &error));
ADBC_BENCHMARK_RETURN_NOT_OK(
AdbcStatementSetOption(&insert_stmt.value, ADBC_INGEST_OPTION_MODE,
ADBC_INGEST_OPTION_MODE_APPEND, &error));
for (auto _ : state) {
AdbcStatementBind(&insert_stmt.value, &array.value, &schema.value, &error);
AdbcStatementExecuteQuery(&insert_stmt.value, nullptr, nullptr, &error);
}
ADBC_BENCHMARK_RETURN_NOT_OK(
AdbcStatementSetSqlQuery(&statement.value, drop_query, &error));
ADBC_BENCHMARK_RETURN_NOT_OK(
AdbcStatementExecuteQuery(&statement.value, nullptr, nullptr, &error));
}