Skip to content

stemflow.lazyloading.open_db_connection

open_both_Xy_db_connection(X_train, y_train, duckdb_config)

Open a DuckDB connection. With one source (X_train), behaves like before and yields (X_obj, con). With two sources (X_train, Y_train), yields (X_rel_or_df, Y_rel_or_df, con) sharing the SAME connection.

Source code in stemflow/lazyloading/open_db_connection.py 
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
@contextmanager
def open_both_Xy_db_connection(X_train, y_train, duckdb_config):
    """
    Open a DuckDB connection. With one source (X_train), behaves like before and yields (X_obj, con).
    With two sources (X_train, Y_train), yields (X_rel_or_df, Y_rel_or_df, con) sharing the SAME connection.
    """
    con = None
    try:
        # one shared connection for both cases
        con = duckdb.connect(config=duckdb_config)
        # dual-source mode: expose BOTH in the SAME connection
        X_obj = _as_relation(con, X_train, "X_df", "xdb")
        Y_obj = _as_relation(con, y_train, "y_df", "ydb")
        yield X_obj, Y_obj, con

    finally:
        if con is not None:
            con.close()

open_db_connection(X_train, duckdb_config)

open the connection to databse using duckdb

Parameters:

  • X_train –

    Training variables. Can be either a pd.DataFrame object or a string that indicate the path to the database (.duckdb or .parquet).

  • self –

    The AdaSTEM/STEM model instance

return

The name of the table in the open connection to the dataframe

Source code in stemflow/lazyloading/open_db_connection.py 
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
@contextmanager
def open_db_connection(X_train, duckdb_config):
    """open the connection to databse using duckdb

    Args:
        X_train: Training variables. Can be either a pd.DataFrame object or a string that indicate the path to the database (.duckdb or .parquet).
        self: The AdaSTEM/STEM model instance

    return:
        The name of the table in the open connection to the dataframe
    """
    con = None
    try:
        if isinstance(X_train, pd.DataFrame):
            con = duckdb.connect(config=duckdb_config)
            the_only_table_name = X_train
        elif isinstance(X_train, str):
            if X_train.endswith('.duckdb'):
                con = duckdb.connect(X_train, read_only=True, config=duckdb_config)
                the_only_table_name = con.sql("""SHOW TABLES""").fetchall()[0][0]
                quoted_table = '"' + the_only_table_name.replace('"', '""') + '"'
                the_only_table_name = con.sql(f"select * from {quoted_table};")
            elif X_train.endswith('.parquet'):
                con = duckdb.connect(config=duckdb_config)
                the_only_table_name = con.read_parquet(X_train, hive_partitioning=False)

        yield the_only_table_name, con

    finally:
        if con is not None:
            con.close()