static void BM_PostgresqlDecimalWrite()

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));
}