in c/driver/postgresql/postgresql_benchmark.cc [163:294]
static void BM_PostgresqlDecimalWrite(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;
constexpr enum ArrowType type = NANOARROW_TYPE_DECIMAL128;
constexpr int32_t bitwidth = 128;
constexpr int32_t precision = 38;
constexpr int32_t scale = 8;
constexpr size_t ncols = 5;
ArrowSchemaInit(&schema.value);
if (ArrowSchemaSetTypeStruct(&schema.value, ncols) != NANOARROW_OK) {
state.SkipWithError("Call to ArrowSchemaSetTypeStruct failed!");
error.release(&error);
return;
}
for (size_t i = 0; i < ncols; i++) {
if (PrivateArrowSchemaSetTypeDecimal(schema.value.children[i], type, precision,
scale) != NANOARROW_OK) {
state.SkipWithError("Call to ArrowSchemaSetTypeDecimal failed!");
error.release(&error);
return;
}
std::string colname = "col" + std::to_string(i);
if (ArrowSchemaSetName(schema.value.children[i], colname.c_str()) != NANOARROW_OK) {
state.SkipWithError("Call to ArrowSchemaSetName failed!");
error.release(&error);
return;
}
}
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;
}
constexpr size_t nrows = 1000;
struct ArrowDecimal decimal;
ArrowDecimalInit(&decimal, bitwidth, precision, scale);
for (size_t i = 0; i < nrows; i++) {
for (size_t j = 0; j < ncols; j++) {
ArrowDecimalSetInt(&decimal, i + j);
if (ArrowArrayAppendDecimal(array.value.children[j], &decimal) != NANOARROW_OK) {
state.SkipWithError("Call to ArrowArrayAppendDecimal failed");
error.release(&error);
return;
}
}
}
for (int64_t i = 0; i < array.value.n_children; i++) {
array.value.children[i]->length = nrows;
}
array.value.length = nrows;
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 (col0 DECIMAL(38, 8), "
"col1 DECIMAL(38, 8), col2 DECIMAL(38, 8), col3 DECIMAL(38, 8), col4 DECIMAL(38, "
"8))";
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));
}