Add upsert in `_set_trial_attr_without_commit` for PostgreSQL by jaikumarm · Pull Request #6282 · optuna/optuna

jaikumarm · 2025-09-18T18:42:50Z

Motivation

This PR adds upsert in _set_trial_attr_without_commit for postgresql. It is an expansion of #5741 and #5703. These two PR's added upsert for sqlite and mysql for trial attrs, but postgresql was marked as TODO, this adds support for postgresql.

Description of the changes

From my research it seems sqlalchemy already has the upsert function on_conflict_do_update similar to sqlite and mysql. I just had to add import dialects and a conditional call for postgres.

I checked sqlalchemy docs it looks like 1.4.2 supports postgresql on_conflict_do_update, so don't think we need to change any dependencies.
https://docs.sqlalchemy.org/en/14/dialects/postgresql.html#sqlalchemy.dialects.postgresql.Insert.on_conflict_do_update

Verification of SQL changes

I followed the same verification example provided in #5741 and #5703 for verification

Code

import optuna
import logging

optuna.logging.set_verbosity(optuna.logging.WARNING)

postgresql_url = os.environ["OPTUNA_PG_TEST_STORAGE"]

study = optuna.create_study(storage=postgresql_url)
storage = study._storage


def objective(trial: optuna.Trial) -> float:
    storage.set_trial_system_attr(trial._trial_id, "system_attr_1", "value_1")

    logging.basicConfig()
    logging.getLogger("sqlalchemy.engine").setLevel(logging.INFO)
    storage.set_trial_system_attr(trial._trial_id, "system_attr_1", "new_value_1")
    storage.set_trial_system_attr(trial._trial_id, "system_attr_2", "value_2")
    logging.getLogger("sqlalchemy.engine").setLevel(logging.WARNING)

    return trial.suggest_float("x", 0.0, 1.0) ** 2


study.optimize(objective, n_trials=1)

Output

Before

INFO:sqlalchemy.engine.Engine:BEGIN (implicit)
INFO:sqlalchemy.engine.Engine:SELECT trials.trial_id AS trials_trial_id, trials.number AS trials_number, trials.study_id AS trials_study_id, trials.state AS trials_state, trials.datetime_start AS trials_datetime_start, trials.datetime_complete AS trials_datetime_complete 
FROM trials 
WHERE trials.trial_id = %(trial_id_1)s
INFO:sqlalchemy.engine.Engine:[cached since 0.02755s ago] {'trial_id_1': 3}
INFO:sqlalchemy.engine.Engine:SELECT trial_system_attributes.trial_system_attribute_id AS trial_system_attributes_trial_system_attribute_id, trial_system_attributes.trial_id AS trial_system_attributes_trial_id, trial_system_attributes.key AS trial_system_attributes_key, trial_system_attributes.value_json AS trial_system_attributes_value_json 
FROM trial_system_attributes 
WHERE trial_system_attributes.trial_id = %(trial_id_1)s AND trial_system_attributes.key = %(key_1)s
INFO:sqlalchemy.engine.Engine:[cached since 0.006932s ago] {'trial_id_1': 3, 'key_1': 'system_attr_1'}
INFO:sqlalchemy.engine.Engine:UPDATE trial_system_attributes SET value_json=%(value_json)s WHERE trial_system_attributes.trial_system_attribute_id = %(trial_system_attributes_trial_system_attribute_id)s
INFO:sqlalchemy.engine.Engine:[generated in 0.00016s] {'value_json': '"new_value_1"', 'trial_system_attributes_trial_system_attribute_id': 7}
INFO:sqlalchemy.engine.Engine:COMMIT
INFO:sqlalchemy.engine.Engine:BEGIN (implicit)
INFO:sqlalchemy.engine.Engine:SELECT trials.trial_id AS trials_trial_id, trials.number AS trials_number, trials.study_id AS trials_study_id, trials.state AS trials_state, trials.datetime_start AS trials_datetime_start, trials.datetime_complete AS trials_datetime_complete 
FROM trials 
WHERE trials.trial_id = %(trial_id_1)s
INFO:sqlalchemy.engine.Engine:[cached since 0.03367s ago] {'trial_id_1': 3}
INFO:sqlalchemy.engine.Engine:SELECT trial_system_attributes.trial_system_attribute_id AS trial_system_attributes_trial_system_attribute_id, trial_system_attributes.trial_id AS trial_system_attributes_trial_id, trial_system_attributes.key AS trial_system_attributes_key, trial_system_attributes.value_json AS trial_system_attributes_value_json 
FROM trial_system_attributes 
WHERE trial_system_attributes.trial_id = %(trial_id_1)s AND trial_system_attributes.key = %(key_1)s
INFO:sqlalchemy.engine.Engine:[cached since 0.01297s ago] {'trial_id_1': 3, 'key_1': 'system_attr_2'}
INFO:sqlalchemy.engine.Engine:INSERT INTO trial_system_attributes (trial_id, key, value_json) VALUES (%(trial_id)s, %(key)s, %(value_json)s) RETURNING trial_system_attributes.trial_system_attribute_id
INFO:sqlalchemy.engine.Engine:[cached since 0.01201s ago] {'trial_id': 3, 'key': 'system_attr_2', 'value_json': '"value_2"'}
INFO:sqlalchemy.engine.Engine:COMMIT

After

INFO:sqlalchemy.engine.Engine:BEGIN (implicit)
INFO:sqlalchemy.engine.Engine:SELECT trials.trial_id AS trials_trial_id, trials.number AS trials_number, trials.study_id AS trials_study_id, trials.state AS trials_state, trials.datetime_start AS trials_datetime_start, trials.datetime_complete AS trials_datetime_complete 
FROM trials 
WHERE trials.trial_id = %(trial_id_1)s
INFO:sqlalchemy.engine.Engine:[cached since 0.1034s ago] {'trial_id_1': 1}
INFO:sqlalchemy.engine.Engine:INSERT INTO trial_system_attributes (trial_id, key, value_json) VALUES (%(trial_id)s, %(key)s, %(value_json)s) ON CONFLICT (trial_id, key) DO UPDATE SET value_json = excluded.value_json RETURNING trial_system_attributes.trial_system_attribute_id
INFO:sqlalchemy.engine.Engine:[no key 0.00019s] {'trial_id': 1, 'key': 'system_attr_1', 'value_json': '"new_value_1"'}
INFO:sqlalchemy.engine.Engine:COMMIT
INFO:sqlalchemy.engine.Engine:BEGIN (implicit)
INFO:sqlalchemy.engine.Engine:SELECT trials.trial_id AS trials_trial_id, trials.number AS trials_number, trials.study_id AS trials_study_id, trials.state AS trials_state, trials.datetime_start AS trials_datetime_start, trials.datetime_complete AS trials_datetime_complete 
FROM trials 
WHERE trials.trial_id = %(trial_id_1)s
INFO:sqlalchemy.engine.Engine:[cached since 0.1089s ago] {'trial_id_1': 1}
INFO:sqlalchemy.engine.Engine:INSERT INTO trial_system_attributes (trial_id, key, value_json) VALUES (%(trial_id)s, %(key)s, %(value_json)s) ON CONFLICT (trial_id, key) DO UPDATE SET value_json = excluded.value_json RETURNING trial_system_attributes.trial_system_attribute_id
INFO:sqlalchemy.engine.Engine:[no key 0.00019s] {'trial_id': 1, 'key': 'system_attr_2', 'value_json': '"value_2"'}
INFO:sqlalchemy.engine.Engine:COMMIT

Performance Check

Code

import os
import time

import optuna


postgresql_url = os.environ["OPTUNA_PG_TEST_STORAGE"]
sqlite_url = "sqlite:///test-pg-upsert.db"
study_name = "test-pg-upsert"

sqlite_study = optuna.create_study(storage=sqlite_url, study_name=study_name)
postgresql_study = optuna.create_study(storage=postgresql_url, study_name=study_name)

sqlite_trial = sqlite_study.ask()
postgresql_trial = postgresql_study.ask()

start = time.time()
for i in range(1000):
    sqlite_trial.set_user_attr("foo", i)
print(f"{sqlite_trial.user_attrs=}")
print(f"SQLite: {(time.time() - start) * 1000:.2f}[ms]")

start = time.time()
for i in range(1000):
    postgresql_trial.set_user_attr("foo", i)
print(f"PostgreSQL: {(time.time() - start) * 1000:.2f}[ms]")
print(f"{postgresql_trial.user_attrs=}")

Output

Release 4.5.0

[I 2025-09-18 10:36:20,446] A new study created in RDB with name: test-pg-upsert
[I 2025-09-18 10:36:20,842] A new study created in RDB with name: test-pg-upsert
sqlite_trial.user_attrs={'foo': 999}
SQLite: 3091.52[ms]
PostgreSQL: 6331.63[ms]
postgresql_trial.user_attrs={'foo': 999}

Master 4.6.0.dev0 with fix

[I 2025-09-18 10:41:08,638] A new study created in RDB with name: test-pg-upsert
[I 2025-09-18 10:41:08,882] A new study created in RDB with name: test-pg-upsert
sqlite_trial.user_attrs={'foo': 999}
SQLite: 2858.08[ms]
PostgreSQL: 3689.49[ms]
postgresql_trial.user_attrs={'foo': 999}

this is my first PR for optuna, so not sure If it has everything thats needed, if not please let me know.

c-bata · 2025-09-19T03:43:33Z

@kAIto47802 Could you review this PR?

github-actions · 2025-09-28T23:05:39Z

This pull request has not seen any recent activity.

y0z · 2025-10-02T05:17:09Z

@kAIto47802, @c-bata
Let me reassign this PR, as you don't have time to review it.

@nabenabe0928 @not522, could you review this PR?

not522

Sorry for the late response. It's almost LGTM. Could you update the TODO comment?

optuna/storages/_rdb/storage.py

Co-authored-by: Naoto Mizuno <naotomizuno@preferred.jp>

not522

Thank you! LGTM!

gen740

LGTM!

TBH, in my environment (MacOS 26 Tahoe, MacBook Pro M3 Max), there is not much difference in performance between master and this PR.

This PR

Benchmark 1: python3 pr6282.py
  Time (mean ± σ):     727.3 ms ±  11.3 ms    [User: 590.3 ms, System: 58.1 ms]
  Range (min … max):   715.5 ms … 745.9 ms    10 runs

Master

Benchmark 1: python3 pr6282.py
  Time (mean ± σ):     745.4 ms ±  18.8 ms    [User: 581.6 ms, System: 63.4 ms]
  Range (min … max):   710.2 ms … 782.1 ms    10 runs

not522 · 2025-10-03T05:50:34Z

Note: I confirmed that this PR ensures faster performance if value updates occur. (8.02s -> 7.87s)

import optuna

postgresql_url = "postgresql+psycopg2://user:test@127.0.0.1/optunatest"
study = optuna.create_study(storage=postgresql_url, study_name="test")
storage = study._storage

def objective(trial: optuna.Trial) -> float:
    storage.set_trial_system_attr(trial._trial_id, "system_attr_1", "value_1")
    storage.set_trial_system_attr(trial._trial_id, "system_attr_1", "new_value_1")
    storage.set_trial_system_attr(trial._trial_id, "system_attr_2", "value_2")
    return trial.suggest_float("x", 0.0, 1.0) ** 2

study.optimize(objective, n_trials=1000)

jaikumarm added 2 commits September 18, 2025 11:10

adding upsert in _set_trial_attr_without_commit for postgresql

aff7195

moving import order to fix isort check failure

459b255

nabenabe0928 changed the title ~~adding upsert in _set_trial_attr_without_commit for postgresql~~ Add upsert in _set_trial_attr_without_commit for PostgreSQL Sep 19, 2025

c-bata assigned c-bata and kAIto47802 Sep 19, 2025

c-bata added the enhancement Change that does not break compatibility and not affect public interfaces, but improves performance. label Sep 19, 2025

github-actions bot added the stale Exempt from stale bot labeling. label Sep 28, 2025

c-bata removed the stale Exempt from stale bot labeling. label Sep 29, 2025

y0z assigned not522 and nabenabe0928 and unassigned c-bata and kAIto47802 Oct 2, 2025

not522 reviewed Oct 2, 2025

View reviewed changes

optuna/storages/_rdb/storage.py Outdated Show resolved Hide resolved

update TODO to remove PostgreSQL

f0a255e

Co-authored-by: Naoto Mizuno <naotomizuno@preferred.jp>

not522 approved these changes Oct 3, 2025

View reviewed changes

not522 removed their assignment Oct 3, 2025

gen740 assigned gen740 and unassigned nabenabe0928 Oct 3, 2025

gen740 approved these changes Oct 3, 2025

View reviewed changes

gen740 added this to the v4.6.0 milestone Oct 3, 2025

gen740 merged commit b7e92b4 into optuna:master Oct 3, 2025
14 checks passed

gen740 removed their assignment Oct 3, 2025

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Add upsert in `_set_trial_attr_without_commit` for PostgreSQL#6282

Add upsert in `_set_trial_attr_without_commit` for PostgreSQL#6282
gen740 merged 3 commits intooptuna:masterfrom
jaikumarm:fix/postgresql-upsert-in-set_trial_system_attr

jaikumarm commented Sep 18, 2025

Uh oh!

c-bata commented Sep 19, 2025

Uh oh!

github-actions bot commented Sep 28, 2025

Uh oh!

y0z commented Oct 2, 2025 •

edited

Loading

Uh oh!

not522 left a comment

Uh oh!

Uh oh!

not522 left a comment

Uh oh!

gen740 left a comment

Uh oh!

Uh oh!

not522 commented Oct 3, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

7 participants

Uh oh!

Conversation

jaikumarm commented Sep 18, 2025

Motivation

Description of the changes

Verification of SQL changes

Code

Output

Before

After

Performance Check

Code

Output

Release 4.5.0

Master 4.6.0.dev0 with fix

Uh oh!

c-bata commented Sep 19, 2025

Uh oh!

github-actions bot commented Sep 28, 2025

Uh oh!

y0z commented Oct 2, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

not522 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

not522 left a comment

Choose a reason for hiding this comment

Uh oh!

gen740 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

not522 commented Oct 3, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

7 participants

y0z commented Oct 2, 2025 •

edited

Loading