squirrels 0.2.0rc1__py3-none-any.whl → 0.2.2__py3-none-any.whl

squirrels/arguments/run_time_args.py

@@ -7,6 +7,7 @@ from .init_time_args import ConnectionsArgs, ParametersArgs
 from ..user_base import User
 from ..parameters import Parameter
 from .._connection_set import ConnectionSetIO
+from .. import _utils as u
 
 
 @dataclass
@@ -20,7 +21,7 @@ class AuthArgs(ConnectionsArgs):
 class ContextArgs(ParametersArgs):
     user: User
     prms: dict[str, Parameter]
-    args: dict[str, Any]
+    traits: dict[str, Any]
 
 
 @dataclass
@@ -80,10 +81,20 @@ class ModelArgs(ModelDepsArgs):
         if dataframes is None:
             dataframes = {x: self.ref(x) for x in self.dependencies}
 
-        conn = sqlite3.connect(":memory:")
+        use_duckdb = u.use_duckdb()
+        if use_duckdb:
+            import duckdb
+            conn = duckdb.connect()
+        else:
+            conn = sqlite3.connect(":memory:")
+        
         try:
             for name, df in dataframes.items():
-                df.to_sql(name, conn, index=False)
-            return pd.read_sql(query, conn)
+                if use_duckdb:
+                    conn.execute(f"CREATE TABLE {name} AS FROM df")
+                else:
+                    df.to_sql(name, conn, index=False)
+            
+            return conn.execute(query).df() if use_duckdb else pd.read_sql(query, conn)
         finally:
             conn.close()