Repository: ibis-project/ibis

Issue:
feat: Make creating UDF's optional
### Is your feature request related to a problem?

When opening a Snowflake connection, there is some code that creates UDFs in snowflake. For reading tables/executing plain SQL on snowflake they do not seem to be required, so I would like to turn that behaviour off and save some seconds on establishing the initial connection.

### Describe the solution you'd like

I see this code here https://github.com/ibis-project/ibis/blob/master/ibis/backends/snowflake/__init__.py#L230. It would be great to have a connection parameter that allows me to skip this part (e.g. `udf_support=False`). 

### What version of ibis are you running?

6.0.0

### What backend(s) are you using, if any?

Snowflake, BigQuery, ...

### Code of Conduct

- [X] I agree to follow this project's Code of Conduct


Discussion:
@tdoehmen Welcome! 👋🏻 

Thanks for opening an issue about this. We can definitely make automatic creation of UDFs in the snowflake backend optional in the next release (6.1.0).
Thanks that sounds great! I didn't look closely into other connectors, but is that something specific to Snowflake or would it make sense to make this behaviour optional across other connectors as well (e.g. BigQuery)? 
I think this is probably specific to snowflake due to its lack of various map functions.

As soon as snowflake adds:

- `object_merge` (combine two objects into one)
- `object_values` (extract the values of an object into an array)
- `object_from_arrays` (create an object from a keys array and a values array)
- `array_zip` (zip N arrays into an array of objects, each with N fields)

we can remove the code to add these UDFs.
@tdoehmen Just out of curiosity what is the time difference you're observing when you turn off UDF creation?
When I time in two separate IPython sessions these I'm only seeing a few milliseconds' difference:

#### **Without** UDFs

```
❯ ipython

In [1]: from ibis.interactive import *

In [2]: import os

In [3]: %time con = ibis.connect(os.environ["SNOWFLAKE_URL"], enable_object_udfs=False)
CPU times: user 857 ms, sys: 2.15 s, total: 3.01 s
Wall time: 637 ms
```

#### **With** UDFs

```
❯ ipython

In [1]: from ibis.interactive import *

In [2]: import os

In [3]: %time con = ibis.connect(os.environ["SNOWFLAKE_URL"], enable_object_udfs=True)
CPU times: user 863 ms, sys: 2.15 s, total: 3.02 s
Wall time: 640 ms
```
According to my logs, it takes about 0.7s to create the UDFs. It's relatively independent of whether the UDFs already exist or not. However, I haven't compared to manually turning off UDF creation in the code yet. Is there a dev branch from which I could install the version with the off-switch? Then I could try it out.

```
2023-07-20 12:02:03,893 INFO execute(806) query: [CREATE DATABASE IF NOT EXISTS ibis_udfs]
2023-07-20 12:02:04,039 INFO execute(819) query execution done
2023-07-20 12:02:04,040 INFO _init_result_and_meta(961) Number of results in first chunk: 1
2023-07-20 12:02:04,040 INFO execute(806) query: [CREATE OR REPLACE FUNCTION ibis_udfs.public.object_merge("obj1" OBJECT, "obj2" O...]
2023-07-20 12:02:04,148 INFO execute(819) query execution done
2023-07-20 12:02:04,149 INFO _init_result_and_meta(961) Number of results in first chunk: 1
2023-07-20 12:02:04,149 INFO execute(806) query: [CREATE OR REPLACE FUNCTION ibis_udfs.public.object_values("obj" OBJECT) RETURNS ...]
2023-07-20 12:02:04,249 INFO execute(819) query execution done
2023-07-20 12:02:04,250 INFO _init_result_and_meta(961) Number of results in first chunk: 1
2023-07-20 12:02:04,250 INFO execute(806) query: [CREATE OR REPLACE FUNCTION ibis_udfs.public.object_from_arrays("ks" ARRAY, "vs" ...]
2023-07-20 12:02:04,384 INFO execute(819) query execution done
2023-07-20 12:02:04,385 INFO _init_result_and_meta(961) Number of results in first chunk: 1
2023-07-20 12:02:04,385 INFO execute(806) query: [CREATE OR REPLACE FUNCTION ibis_udfs.public.array_zip("arrays" ARRAY) RETURNS AR...]
2023-07-20 12:02:04,527 INFO execute(819) query execution done
2023-07-20 12:02:04,528 INFO _init_result_and_meta(961) Number of results in first chunk: 1
2023-07-20 12:02:04,528 INFO execute(806) query: [USE SCHEMA "SNOWFLAKE_SAMPLE_DATA"."TPCH_SF1"]
2023-07-20 12:02:04,604 INFO execute(819) query execution done
2023-07-20 12:02:04,605 INFO _init_result_and_meta(961) Number of results in first chunk: 1
```
@cpcloud random gentle ping on this 
@tdoehmen Can you give #6722 a try and see if that resolves the issue for you?
The PR is in draft mode, since I'm not super inclined to add this flag unless we can somehow demonstrate the latency hit is due to creating those UDFs.
Hey, I checked out your branch. That's doing exactly what I was hoping for. Below are the runtimes for creating a connection to snowflake and reading a snowflake table, averaged across 10 runs:

ibis 6.0, udf disabled: 4.67
ibis 6.0: 5.44
ibis 6.1: 5.51

I think it is quite evident that there is an overhead associated with the UDF creation. I know it's relatively small (~0.7s), but for smaller queries it actually starts to matter. Apart from that, when I'm using ibis on client's warehouses, I simply wouldn't want it to make changes to the clients database on it's own. 
@tdoehmen Thanks for getting back! Glad to see that option works for you. I will go ahead and get that PR reviewed.

Resolving patch:
diff --git a/ibis/backends/snowflake/__init__.py b/ibis/backends/snowflake/__init__.py
--- a/ibis/backends/snowflake/__init__.py
+++ b/ibis/backends/snowflake/__init__.py
@@ -163,6 +163,7 @@ def do_connect(
         password: str | None = None,
         authenticator: str | None = None,
         connect_args: Mapping[str, Any] | None = None,
+        create_object_udfs: bool = True,
         **kwargs: Any,
     ):
         """Connect to Snowflake.
@@ -189,9 +190,12 @@ def do_connect(
             database connection is made. This means that
             `ibis.snowflake.connect(...)` can succeed, while subsequent API
             calls fail if the authentication fails for any reason.
+        create_object_udfs
+            Enable object UDF extensions defined by ibis on the first
+            connection to the database.
         connect_args
             Additional arguments passed to the SQLAlchemy engine creation call.
-        kwargs:
+        kwargs
             Additional arguments passed to the SQLAlchemy URL constructor.
             See https://docs.snowflake.com/en/developer-guide/python-connector/sqlalchemy#additional-connection-parameters
             for more details
@@ -241,23 +245,24 @@ def do_connect(
         @sa.event.listens_for(engine, "connect")
         def connect(dbapi_connection, connection_record):
             """Register UDFs on a `"connect"` event."""
-            dialect = engine.dialect
-            quote = dialect.preparer(dialect).quote_identifier
-            with dbapi_connection.cursor() as cur:
-                database, schema = cur.execute(
-                    "SELECT CURRENT_DATABASE(), CURRENT_SCHEMA()"
-                ).fetchone()
-                try:
-                    cur.execute("CREATE DATABASE IF NOT EXISTS ibis_udfs")
-                    cur.execute(f"USE SCHEMA {quote(database)}.{quote(schema)}")
-                    for name, defn in _SNOWFLAKE_MAP_UDFS.items():
-                        cur.execute(self._make_udf(name, defn))
-                except Exception as e:  # noqa: BLE001
-                    warnings.warn(
-                        f"Unable to create map UDFs, some functionality will not work: {e}"
-                    )
+            if create_object_udfs:
+                with dbapi_connection.cursor() as cur:
+                    database, schema = cur.execute(
+                        "SELECT CURRENT_DATABASE(), CURRENT_SCHEMA()"
+                    ).fetchone()
+                    try:
+                        cur.execute("CREATE DATABASE IF NOT EXISTS ibis_udfs")
+                        # snowflake activates a database on creation, so reset
+                        # it back to the original database and schema
+                        cur.execute(f"USE SCHEMA {database}.{schema}")
+                        for name, defn in _SNOWFLAKE_MAP_UDFS.items():
+                            cur.execute(self._make_udf(name, defn))
+                    except Exception as e:  # noqa: BLE001
+                        warnings.warn(
+                            f"Unable to create map UDFs, some functionality will not work: {e}"
+                        )
 
-        res = super().do_connect(engine)
+        super().do_connect(engine)
 
         def normalize_name(name):
             if name is None:
@@ -270,7 +275,6 @@ def normalize_name(name):
                 return name
 
         self.con.dialect.normalize_name = normalize_name
-        return res
 
     def _get_udf_source(self, udf_node: ops.ScalarUDF):
         name = type(udf_node).__name__


Test patch: