From b2fef6796fa1956924fae8c71455f925abb6ad88 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 6 Apr 2026 18:48:40 +0000 Subject: [PATCH 1/2] [pre-commit.ci] pre-commit suggestions MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit updates: - [github.com/psf/black-pre-commit-mirror: 25.12.0 → 26.3.1](https://github.com/psf/black-pre-commit-mirror/compare/25.12.0...26.3.1) - [github.com/psf/black-pre-commit-mirror: 25.12.0 → 26.3.1](https://github.com/psf/black-pre-commit-mirror/compare/25.12.0...26.3.1) - [github.com/pycqa/isort: 7.0.0 → 8.0.1](https://github.com/pycqa/isort/compare/7.0.0...8.0.1) --- .pre-commit-config.yaml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 5f74eda7..f36090ae 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -12,17 +12,17 @@ repos: - id: check-toml - repo: https://github.com/psf/black-pre-commit-mirror - rev: 25.12.0 + rev: 26.3.1 hooks: - id: black - repo: https://github.com/psf/black-pre-commit-mirror - rev: 25.12.0 + rev: 26.3.1 hooks: - id: black-jupyter - repo: https://github.com/pycqa/isort - rev: 7.0.0 + rev: 8.0.1 hooks: - id: isort args: ["--profile", "black"] From 2a4212a511c4078b5475ee147a679fb22266f609 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 6 Apr 2026 18:52:10 +0000 Subject: [PATCH 2/2] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- relbench/base/task_autocomplete.py | 6 ++--- relbench/tasks/amazon.py | 42 ++++++++++-------------------- relbench/tasks/arxiv.py | 24 ++++++----------- relbench/tasks/avito.py | 24 ++++++----------- relbench/tasks/event.py | 18 +++++-------- relbench/tasks/f1.py | 24 ++++++----------- relbench/tasks/hm.py | 18 +++++-------- relbench/tasks/mimic.py | 6 ++--- relbench/tasks/ratebeer.py | 42 ++++++++++-------------------- relbench/tasks/stack.py | 30 +++++++-------------- relbench/tasks/trial.py | 30 +++++++-------------- tutorials/custom_task.ipynb | 6 ++--- 12 files changed, 90 insertions(+), 180 deletions(-) diff --git a/relbench/base/task_autocomplete.py b/relbench/base/task_autocomplete.py index 6c5c2238..59497921 100644 --- a/relbench/base/task_autocomplete.py +++ b/relbench/base/task_autocomplete.py @@ -173,8 +173,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab min_timestamp = timestamp_df["timestamp"].min() max_timestamp = timestamp_df["timestamp"].max() - df = duckdb.sql( - f""" + df = duckdb.sql(f""" SELECT entity_table.{self.time_col}, entity_table.{entity_col}, @@ -188,8 +187,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab WHERE entity_table.{self.time_col} > '{min_timestamp}' AND entity_table.{self.time_col} <= '{max_timestamp}' - """ - ).df() + """).df() if self.task_type == TaskType.MULTICLASS_CLASSIFICATION: df[self.target_col] = self.transform_target(df[self.target_col]) diff --git a/relbench/tasks/amazon.py b/relbench/tasks/amazon.py index 12dbcd52..79b3f0dd 100644 --- a/relbench/tasks/amazon.py +++ b/relbench/tasks/amazon.py @@ -34,8 +34,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab review = db.table_dict["review"].df timestamp_df = pd.DataFrame({"timestamp": timestamps}) - df = duckdb.sql( - f""" + df = duckdb.sql(f""" SELECT timestamp, customer_id, @@ -61,8 +60,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab review_time > timestamp - INTERVAL '{self.timedelta}' AND review_time <= timestamp ) - """ - ).df() + """).df() return Table( df=df, @@ -90,8 +88,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab review = db.table_dict["review"].df timestamp_df = pd.DataFrame({"timestamp": timestamps}) - df = duckdb.sql( - f""" + df = duckdb.sql(f""" SELECT timestamp, customer_id, @@ -120,8 +117,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab review_time > timestamp - INTERVAL '{self.timedelta}' AND review_time <= timestamp ) - """ - ).df() + """).df() return Table( df=df, @@ -149,8 +145,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab review = db.table_dict["review"].df timestamp_df = pd.DataFrame({"timestamp": timestamps}) - df = duckdb.sql( - f""" + df = duckdb.sql(f""" SELECT timestamp, product_id, @@ -176,8 +171,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab review_time > timestamp - INTERVAL '{self.timedelta}' AND review_time <= timestamp ) - """ - ).df() + """).df() return Table( df=df, @@ -205,8 +199,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab review = db.table_dict["review"].df timestamp_df = pd.DataFrame({"timestamp": timestamps}) - df = duckdb.sql( - f""" + df = duckdb.sql(f""" SELECT timestamp, product.product_id, @@ -222,8 +215,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab GROUP BY timestamp, product.product_id - """ - ).df() + """).df() return Table( df=df, @@ -252,8 +244,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab review = db.table_dict["review"].df timestamp_df = pd.DataFrame({"timestamp": timestamps}) - df = duckdb.sql( - f""" + df = duckdb.sql(f""" SELECT t.timestamp, review.customer_id, @@ -270,8 +261,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab GROUP BY t.timestamp, review.customer_id - """ - ).df() + """).df() return Table( df=df, @@ -303,8 +293,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab review = db.table_dict["review"].df timestamp_df = pd.DataFrame({"timestamp": timestamps}) - df = duckdb.sql( - f""" + df = duckdb.sql(f""" SELECT t.timestamp, review.customer_id, @@ -323,8 +312,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab GROUP BY t.timestamp, review.customer_id - """ - ).df() + """).df() return Table( df=df, @@ -360,8 +348,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab 300 # minimum length of review to be considered as detailed review ) - df = duckdb.sql( - f""" + df = duckdb.sql(f""" SELECT t.timestamp, review.customer_id, @@ -380,8 +367,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab GROUP BY t.timestamp, review.customer_id - """ - ).df() + """).df() return Table( df=df, diff --git a/relbench/tasks/arxiv.py b/relbench/tasks/arxiv.py index 6103d7d8..ea5448cc 100644 --- a/relbench/tasks/arxiv.py +++ b/relbench/tasks/arxiv.py @@ -33,8 +33,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab papers = db.table_dict["papers"].df citations = db.table_dict["citations"].df - df = duckdb.sql( - f""" + df = duckdb.sql(f""" WITH paper_citations AS ( SELECT t.timestamp AS date, @@ -54,8 +53,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab Paper_ID, CASE WHEN citation_count > 0 THEN 1 ELSE 0 END AS cited FROM paper_citations; - """ - ).df() + """).df() return Table( df=df, @@ -88,8 +86,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab paperCategories = db.table_dict["paperCategories"].df categories = db.table_dict["categories"].df - df = duckdb.sql( - f""" + df = duckdb.sql(f""" WITH author_pubs AS ( SELECT t.timestamp AS date, @@ -127,8 +124,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab LEFT JOIN categories c ON r.Primary_Category_ID = c.Category_ID WHERE r.rn = 1 ; - """ - ).df() + """).df() return Table( df=df, @@ -154,8 +150,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab timestamp_df = pd.DataFrame({"timestamp": timestamps}) paperAuthors = db.table_dict["paperAuthors"].df - df = duckdb.sql( - f""" + df = duckdb.sql(f""" WITH author_pubs AS ( SELECT t.timestamp AS date, @@ -168,8 +163,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab GROUP BY t.timestamp, pa.Author_ID ) SELECT date, Author_ID, publication_count FROM author_pubs; - """ - ).df() + """).df() return Table( df=df, @@ -199,8 +193,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab papers = db.table_dict["papers"].df citations = db.table_dict["citations"].df - df = duckdb.sql( - f""" + df = duckdb.sql(f""" WITH paper_co_citations AS ( SELECT t.timestamp AS date, @@ -226,8 +219,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab FROM paper_co_citations GROUP BY date, Paper_ID ; - """ - ).df() + """).df() return Table( df=df, diff --git a/relbench/tasks/avito.py b/relbench/tasks/avito.py index 7f03a622..1df2ea4f 100644 --- a/relbench/tasks/avito.py +++ b/relbench/tasks/avito.py @@ -32,8 +32,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab ads_info = db.table_dict["AdsInfo"].df search_stream = db.table_dict["SearchStream"].df timestamp_df = pd.DataFrame({"timestamp": timestamps}) - df = duckdb.sql( - f""" + df = duckdb.sql(f""" SELECT search_ads.AdID, t.timestamp, @@ -55,8 +54,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab search_ads.AdID HAVING SUM(search_ads.isClick) > 0 - """ - ).df() + """).df() return Table( df=df, @@ -81,8 +79,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab user_info = db.table_dict["UserInfo"].df visits_stream = db.table_dict["VisitStream"].df timestamp_df = pd.DataFrame({"timestamp": timestamps}) - df = duckdb.sql( - f""" + df = duckdb.sql(f""" SELECT visit_ads.UserID, t.timestamp, @@ -103,8 +100,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab GROUP BY t.timestamp, visit_ads.UserID - """ - ).df() + """).df() return Table( df=df, @@ -131,8 +127,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab search_info = db.table_dict["SearchInfo"].df search_stream = db.table_dict["SearchStream"].df timestamp_df = pd.DataFrame({"timestamp": timestamps}) - df = duckdb.sql( - f""" + df = duckdb.sql(f""" SELECT search_ads.UserID, t.timestamp, @@ -160,8 +155,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab GROUP BY t.timestamp, search_ads.UserID - """ - ).df() + """).df() return Table( df=df, @@ -191,8 +185,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab visits_stream = db.table_dict["VisitStream"].df timestamp_df = pd.DataFrame({"timestamp": timestamps}) - df = duckdb.sql( - f""" + df = duckdb.sql(f""" SELECT visit_ads.UserID, t.timestamp, @@ -213,8 +206,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab GROUP BY t.timestamp, visit_ads.UserID - """ - ).df() + """).df() return Table( df=df, fkey_col_to_pkey_table={ diff --git a/relbench/tasks/event.py b/relbench/tasks/event.py index 270a7ab1..8463c6cd 100644 --- a/relbench/tasks/event.py +++ b/relbench/tasks/event.py @@ -24,8 +24,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab event_interest = db.table_dict["event_interest"].df timestamp_df = pd.DataFrame({"timestamp": timestamps}) - df = duckdb.sql( - f"""SELECT + df = duckdb.sql(f"""SELECT t.timestamp, event_attendees.user_id AS user, SUM(CASE WHEN event_attendees.status IN ('yes', 'maybe') THEN 1 ELSE 0 END) AS target @@ -39,8 +38,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab GROUP BY t.timestamp, event_attendees.user_id - """ - ).df() + """).df() df = df.dropna(subset=["user"]) df["user"] = df["user"].astype(int) df = df.reset_index() @@ -86,8 +84,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab ) timestamp_df = pd.concat([new_row, timestamp_df], ignore_index=True) - df = duckdb.sql( - f""" + df = duckdb.sql(f""" WITH tb AS( SELECT t.timestamp AS timestamp, @@ -113,8 +110,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab tb WHERE prev_target = 1; - """ - ).df() + """).df() if eval_timestamp_len == 1: df = df[df.timestamp == df.timestamp.max()] @@ -152,8 +148,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab event_interest = db.table_dict["event_interest"].df timestamp_df = pd.DataFrame({"timestamp": timestamps}) - df = duckdb.sql( - f"""SELECT + df = duckdb.sql(f"""SELECT t.timestamp AS timestamp, event_attendees.user_id AS user, CASE @@ -170,8 +165,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab GROUP BY t.timestamp, event_attendees.user_id - """ - ).df() + """).df() df = df.dropna(subset=["user"]) df["user"] = df["user"].astype(int) diff --git a/relbench/tasks/f1.py b/relbench/tasks/f1.py index 85e5a4cd..2a1994f6 100644 --- a/relbench/tasks/f1.py +++ b/relbench/tasks/f1.py @@ -34,8 +34,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab results = db.table_dict["results"].df - df = duckdb.sql( - f""" + df = duckdb.sql(f""" SELECT t.timestamp as date, re.driverId as driverId, @@ -56,8 +55,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab GROUP BY t.timestamp, re.driverId ; - """ - ).df() + """).df() return Table( df=df, @@ -84,8 +82,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab results = db.table_dict["results"].df - df = duckdb.sql( - f""" + df = duckdb.sql(f""" SELECT t.timestamp as date, re.driverId as driverId, @@ -106,8 +103,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab GROUP BY t.timestamp, re.driverId ; - """ - ).df() + """).df() return Table( df=df, @@ -135,8 +131,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab qualifying = db.table_dict["qualifying"].df - df = duckdb.sql( - f""" + df = duckdb.sql(f""" SELECT t.timestamp as date, qu.driverId as driverId, @@ -160,8 +155,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab GROUP BY t.timestamp, qu.driverId ; - """ - ).df() + """).df() df["qualifying"] = df["qualifying"].astype("int64") @@ -195,8 +189,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab results = db.table_dict["results"].df races = db.table_dict["races"].df - df = duckdb.sql( - f""" + df = duckdb.sql(f""" SELECT t.timestamp as date, re.driverId as driverId, @@ -214,8 +207,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab re.raceId = race.raceId GROUP BY t.timestamp, re.driverId ; - """ - ).df() + """).df() return Table( df=df, diff --git a/relbench/tasks/hm.py b/relbench/tasks/hm.py index d40158eb..05359551 100644 --- a/relbench/tasks/hm.py +++ b/relbench/tasks/hm.py @@ -35,8 +35,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab transactions = db.table_dict["transactions"].df timestamp_df = pd.DataFrame({"timestamp": timestamps}) - df = duckdb.sql( - f""" + df = duckdb.sql(f""" SELECT t.timestamp, transactions.customer_id, @@ -51,8 +50,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab GROUP BY t.timestamp, transactions.customer_id - """ - ).df() + """).df() return Table( df=df, @@ -81,8 +79,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab transactions = db.table_dict["transactions"].df timestamp_df = pd.DataFrame({"timestamp": timestamps}) - df = duckdb.sql( - f""" + df = duckdb.sql(f""" SELECT timestamp, customer_id, @@ -108,8 +105,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab t_dat > timestamp - INTERVAL '{self.timedelta}' AND t_dat <= timestamp ) - """ - ).df() + """).df() return Table( df=df, @@ -136,8 +132,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab timestamp_df = pd.DataFrame({"timestamp": timestamps}) article = db.table_dict["article"].df - df = duckdb.sql( - f""" + df = duckdb.sql(f""" SELECT timestamp, article_id, @@ -155,8 +150,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab t_dat > timestamp AND t_dat <= timestamp + INTERVAL '{self.timedelta}' ) - """ - ).df() + """).df() return Table( df=df, diff --git a/relbench/tasks/mimic.py b/relbench/tasks/mimic.py index 446fc0b5..86b57dc1 100644 --- a/relbench/tasks/mimic.py +++ b/relbench/tasks/mimic.py @@ -39,8 +39,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab icu = db.table_dict["icustays"].df patients = db.table_dict["patients"].df - df = duckdb.sql( - f""" + df = duckdb.sql(f""" SELECT t.timestamp as intime, p.subject_id, @@ -62,8 +61,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab AND i.intime IS NOT NULL AND i.los_icu IS NOT NULL - """ - ).df() + """).df() # Ensure time column is datetime64[ns] for to_unix_time if self.time_col in df.columns: diff --git a/relbench/tasks/ratebeer.py b/relbench/tasks/ratebeer.py index 6e386340..c1a5815f 100644 --- a/relbench/tasks/ratebeer.py +++ b/relbench/tasks/ratebeer.py @@ -60,8 +60,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab beer_ratings = db.table_dict["beer_ratings"].df timestamp_df = pd.DataFrame({"timestamp": timestamps}) - df = duckdb.sql( - f""" + df = duckdb.sql(f""" SELECT t.timestamp, b.beer_id AS beer_id, @@ -90,8 +89,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab GROUP BY t.timestamp, b.beer_id - """ - ).df() + """).df() return Table( df=df, @@ -117,8 +115,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab beer_ratings = db.table_dict["beer_ratings"].df timestamp_df = pd.DataFrame({"timestamp": timestamps}) - df = duckdb.sql( - f""" + df = duckdb.sql(f""" SELECT t.timestamp, u.user_id AS user_id, @@ -147,8 +144,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab GROUP BY t.timestamp, u.user_id - """ - ).df() + """).df() return Table( df=df, @@ -174,8 +170,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab beers = db.table_dict["beers"].df timestamp_df = pd.DataFrame({"timestamp": timestamps}) - df = duckdb.sql( - f""" + df = duckdb.sql(f""" SELECT t.timestamp, brew.brewer_id AS brewer_id, @@ -204,8 +199,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab GROUP BY t.timestamp, brew.brewer_id - """ - ).df() + """).df() return Table( df=df, @@ -234,8 +228,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab beer_ratings = db.table_dict["beer_ratings"].df timestamp_df = pd.DataFrame({"timestamp": timestamps}) - df = duckdb.sql( - f""" + df = duckdb.sql(f""" SELECT t.timestamp, u.user_id, @@ -256,8 +249,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab AND br2.created_at > t.timestamp - INTERVAL '{self.timedelta}' AND br2.created_at <= t.timestamp ) - """ - ).df() + """).df() return Table( df=df, @@ -296,8 +288,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab beer_ratings = db.table_dict["beer_ratings"].df timestamp_df = pd.DataFrame({"timestamp": timestamps}) - df = duckdb.sql( - f""" + df = duckdb.sql(f""" SELECT t.timestamp, f.user_id, @@ -321,8 +312,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab GROUP BY t.timestamp, f.user_id - """ - ).df() + """).df() return Table( df=df, @@ -361,8 +351,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab beer_ratings = db.table_dict["beer_ratings"].df timestamp_df = pd.DataFrame({"timestamp": timestamps}) - df = duckdb.sql( - f""" + df = duckdb.sql(f""" SELECT t.timestamp, pr.user_id, @@ -393,8 +382,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab GROUP BY t.timestamp, pr.user_id - """ - ).df() + """).df() return Table( df=df, @@ -432,8 +420,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab beer_ratings = db.table_dict["beer_ratings"].df timestamp_df = pd.DataFrame({"timestamp": timestamps}) - df = duckdb.sql( - f""" + df = duckdb.sql(f""" SELECT t.timestamp, br.user_id, @@ -457,8 +444,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab GROUP BY t.timestamp, br.user_id - """ - ).df() + """).df() return Table( df=df, diff --git a/relbench/tasks/stack.py b/relbench/tasks/stack.py index c122bf11..47fbe0e3 100644 --- a/relbench/tasks/stack.py +++ b/relbench/tasks/stack.py @@ -36,8 +36,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab posts = db.table_dict["posts"].df users = db.table_dict["users"].df - df = duckdb.sql( - f""" + df = duckdb.sql(f""" WITH ALL_ENGAGEMENT AS ( SELECT @@ -93,8 +92,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab u.timestamp, u.id ; - """ - ).df() + """).df() return Table( df=df, @@ -121,8 +119,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab votes = db.table_dict["votes"].df posts = db.table_dict["posts"].df - df = duckdb.sql( - f""" + df = duckdb.sql(f""" SELECT t.timestamp, p.id AS PostId, @@ -148,8 +145,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab p.id ; - """ - ).df() + """).df() return Table( df=df, @@ -175,8 +171,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab users = db.table_dict["users"].df badges = db.table_dict["badges"].df - df = duckdb.sql( - f""" + df = duckdb.sql(f""" SELECT t.timestamp, u.Id as UserId, @@ -197,8 +192,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab GROUP BY t.timestamp, u.Id - """ - ).df() + """).df() # remove any IderId rows that are NaN df = df.dropna(subset=["UserId"]) @@ -239,8 +233,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab posts = db.table_dict["posts"].df comments = db.table_dict["comments"].df - df = duckdb.sql( - f""" + df = duckdb.sql(f""" SELECT t.timestamp, c.UserId as UserId, @@ -264,8 +257,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab GROUP BY t.timestamp, c.UserId - """ - ).df() + """).df() return Table( df=df, @@ -299,8 +291,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab posts = db.table_dict["posts"].df postLinks = db.table_dict["postLinks"].df - df = duckdb.sql( - f""" + df = duckdb.sql(f""" SELECT t.timestamp, pl.PostId as PostId, @@ -328,8 +319,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab GROUP BY t.timestamp, pl.PostId; - """ - ).df() + """).df() return Table( df=df, diff --git a/relbench/tasks/trial.py b/relbench/tasks/trial.py index 31116a8d..da92efa6 100644 --- a/relbench/tasks/trial.py +++ b/relbench/tasks/trial.py @@ -33,8 +33,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab outcomes = db.table_dict["outcomes"].df outcome_analyses = db.table_dict["outcome_analyses"].df - df = duckdb.sql( - f""" + df = duckdb.sql(f""" WITH TRIAL_INFO AS ( SELECT oa.nct_id, @@ -66,8 +65,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab and tr.date <= t.timestamp + INTERVAL '{self.timedelta}' WHERE tr.nct_id is not null GROUP BY t.timestamp, tr.nct_id; - """ - ).df() + """).df() return Table( df=df, @@ -94,8 +92,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab reported_event_totals = db.table_dict["reported_event_totals"].df studies = db.table_dict["studies"].df - df = duckdb.sql( - f""" + df = duckdb.sql(f""" WITH TRIAL_INFO AS ( SELECT r.nct_id, @@ -121,8 +118,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab and tr.date <= t.timestamp + INTERVAL '{self.timedelta}' WHERE tr.nct_id is not null and tr.subjects_affected is not null GROUP BY t.timestamp, tr.nct_id; - """ - ).df() + """).df() return Table( df=df, @@ -151,8 +147,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab studies = db.table_dict["studies"].df outcomes = db.table_dict["outcomes"].df - df = duckdb.sql( - f""" + df = duckdb.sql(f""" WITH TRIAL_INFO AS ( SELECT oa.nct_id, @@ -179,8 +174,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab and tr.date <= t.timestamp + INTERVAL '{self.timedelta}' WHERE fs.facility_id is not null GROUP BY t.timestamp, fs.facility_id; - """ - ).df() + """).df() return Table( df=df, @@ -208,8 +202,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab sponsors_studies = db.table_dict["sponsors_studies"].df condition_study = db.table_dict["conditions_studies"].df - df = duckdb.sql( - f""" + df = duckdb.sql(f""" SELECT t.timestamp, cs.condition_id, @@ -220,8 +213,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab ON cs.date > t.timestamp and cs.date <= t.timestamp + INTERVAL '{self.timedelta}' GROUP BY t.timestamp, cs.condition_id; - """ - ).df() + """).df() return Table( df=df, @@ -252,8 +244,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab sponsors_studies = db.table_dict["sponsors_studies"].df facility_study = db.table_dict["facilities_studies"].df - df = duckdb.sql( - f""" + df = duckdb.sql(f""" SELECT t.timestamp, fs.facility_id, @@ -264,8 +255,7 @@ def make_table(self, db: Database, timestamps: "pd.Series[pd.Timestamp]") -> Tab ON fs.date > t.timestamp and fs.date <= t.timestamp + INTERVAL '{self.timedelta}' GROUP BY t.timestamp, fs.facility_id; - """ - ).df() + """).df() return Table( df=df, diff --git a/tutorials/custom_task.ipynb b/tutorials/custom_task.ipynb index 9e9db276..c38fdcd0 100644 --- a/tutorials/custom_task.ipynb +++ b/tutorials/custom_task.ipynb @@ -189,8 +189,7 @@ " ################################################################################\n", " # This SQL query computes the labels for date, driverID pairs\n", " ################################################################################\n", - " df = duckdb.sql(\n", - " f\"\"\"\n", + " df = duckdb.sql(f\"\"\"\n", " SELECT\n", " t.timestamp as date,\n", " dri.driverId as driverId,\n", @@ -218,8 +217,7 @@ " GROUP BY t.timestamp, dri.driverId\n", "\n", " ;\n", - " \"\"\"\n", - " ).df()\n", + " \"\"\").df()\n", "\n", " ################################################################################\n", " # The task table is expressed via a Table object, same as used in Database\n",