Repository: ibis-project/ibis

Issue:
feat(snowflake): run ibis code inside of stored procedure
### Is your feature request related to a problem?

Writing stored procedures for Snowflake using Python [depends on the Snowpark API](https://docs.snowflake.com/en/sql-reference/stored-procedures-python), centering around a [session object](https://docs.snowflake.com/ko/developer-guide/snowpark/reference/python/api/snowflake.snowpark.Session.html) rather than a database connection.

### Describe the solution you'd like

To use Ibis for data wrangling within a Snowflake stored procedure, a new backend is required. It would presumably use `Session.sql` as the interface between SQLAlchemy-based expressions and Snowpark API (which has its own query builder).

### What version of ibis are you running?

5.0.0

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

Snowflake

### Code of Conduct

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


Discussion:
Hey @malthe 👋🏻!

Thanks for the issue.

Is your desire to run ibis code inside of a Python procedure, all tied together from an ibis session?

Can you express in pseudocode the kind workflow you want?

It sounds like it's something like this, but please correct me if I'm wrong:

```python
>>> # assume we can do something like this with the current snowflake backend
>>> con = ibis.snowflake.connect(...)
>>> stored_procedure = """
... import ibis
... # do ibis things? 🙃
... """
>>> # not sure if ☝🏻 is a string or some other kind of object
>>> func = con.register(stored_procedure) # not sure about this API at all
>>> t = con.tables.my_table
>>> res = func(t.some_column)  # ?
```

@malthe Would it make sense to chat on our gitter channel? https://gitter.im/ibis-dev/Lobby

It might be faster/easier to get to understanding there.
Using these Snowpark API-enabled stored procedures (created for example using the [sproc decorator](https://docs.snowflake.com/ko/developer-guide/snowpark/reference/python/api/snowflake.snowpark.functions.sproc.html)), we can run self-contained analytical models which pull data on their own accord using the session object.

The workflow could be that you create a stored procedure using just the regular Snowflake Python connector library – perhaps in a CI/CD pipeline – and use Ibis from within that stored procedure (i.e. Python function):

```python
@sproc(packages=["ibis-framework", "snowflake-snowpark-python"])
def my_proc(session: "snowflake.snowpark.Session"):
    # Using "connect" here feels weird, but nonetheless ...
    backend = ibis.snowpark.connect(session)
    ...
    # Use `backend` as you would any other backend object.
```

Ok, so here's the summary of the use case and some issues that need input from
someone from the Snowpark team.

To clarify, @malthe this isn't directed at you in any way. 

The issues I'm discussing below are all about collaborating with the Snowpark team!

### Use case
The goal is to be able to author ibis expressions **inside** of a stored procedure defined in a Snowpark session.

There are many challenges here that we need help from snowflake to tackle.

### What the Ibis team would need from Snowflake
1. Ibis supports Python 3.8-3.11; we would like it if the various snowflake libraries did too
2. Support nested data deserialization using arrow; getting result sets to the client has sub-optimal performance
3. Support more than one version of arrow
4. Collaboration to quickly resolve issues that come up during development and testing

### Happy Path
There are still some engineering questions on the ibis side assuming the above can be done

1. Should we reuse the pyspark compiler against a snowpark backend or generate
   SQL and hand that over to snowpark?
2. Should we keep the snowflake backend around, or just cut over to snowpark? I'd prefer **not** to have two Snowflake backends.

### Less Happy Path
I really would like to avoid rebuilding any of what snowpark has done, but if there's no way for us to collaborate effectively on making end users happy we'll need to think about how to solve this problem

1. Can we support stored procs without snowpark?
2. If so, we'll probably still need a way to construct a snowflake backend instance from a snowpark session, is that possible/sane?

Again, I would 100% like to avoid any of that.

I hope we can start collaborating on this! It's super exciting!!

We would also be happy to hear about any concrete asks for ibis from the Snowpark team!
@cpcloud thanks for letting us know. I want to understand what you are trying to do. The gist is that ibis is trying to add Snowflake as a backend and you are considering various options, correct?

a) Use SQL directly and use the snowflake python connector to send this SQL over to execute over on Snowflake.
b) Use dataframe semantics and therefore, use Snowpark to do the dataframe math for you.

Both of these approaches seem fine to us so if there are issues that are blocking you from making progress, please file them in the respective repo. We do actively monitor these issues and will prioritize them along with our internal roadmap. 
> @cpcloud thanks for letting us know. I want to understand what you are trying to do. The gist is that ibis is trying to add Snowflake as a backend and you are considering various options, correct?
> 

No, we already have a very functional high operation coverage Snowflake backend based on `snowflake-sqlalchemy`.

We've got a Streamlit app [in our documentation](https://ibis-project.org/backends/support_matrix/) that shows coverage, and our Snowflake backend is one of the highest coverage backends.

> a) Use SQL directly and use the snowflake python connector to send this SQL over to execute over on Snowflake. b) Use dataframe semantics and therefore, use Snowpark to do the dataframe math for you.
> 
> Both of these approaches seem fine to us so if there are issues that are blocking you from making progress, please file them in the respective repo. We do actively monitor these issues and will prioritize them along with our internal roadmap.

As said at the beginning of my long comment above:

> ### Use case
> The goal is to be able to author ibis expressions inside of a stored procedure
> defined in a Snowpark session.

It seems like we can probably do all of this without Snowpark at all, and just use the underlying connection to construct an ibis backend.
@sfc-gh-achandrasekaran IIUC, for Snowflake customers to be able to run Ibis code inside their UDFs and stored procedures, `ibis-framework` needs to be added to the `PACKAGES` configuration value, but it looks like `ibis-framework` isn't available upstream currently. I'm happy to help if the Snowflake team needs any assistance making it available.
@malthe Quick update on this.

After #6156, you'll have the ability to create an ibis backend from a snowpark session. We test the entire ibis backend test suite against a backend created this way, so you can be confident that there's no obvious differences between connecting with a URL/parameters and connecting from a snowpark session.

Unfortunately this is only half the story for using ibis inside of a snowpark stored procedure. There are issues that need to be addressed on the snowflake/snowpark side of things.

1. At the very least, `snowflake-sqlalchemy` needs to exist in the snowflake conda channel. Snowflake uses implicit namespace packages for their Python libraries, and this isn't compatible with their `snowcli` packaging mechanism. End users would still need to package up all the rest of ibis's required dependencies that don't exist on the snowflake conda channel.
2. Ideally, `ibis-framework` (and its dependencies) would simply exist on their conda channel. This would still require `snowflake-sqlalchemy` to exist, but would not place the burden of packaging and uploading third party dependencies using on end users.
@malthe A related question: does your desired workflow involve mixing snowpark session code with ibis code? 

One tricky bit with #6156 is that the bind parameter syntax is different between snowpark and snowflake-sqlalchemy (I don't know why exactly) and there isn't an obvious sane workaround to avoid mutating the snowpark connection to make it work with ibis.

This means that without any further changes to #6156 you can't use the existing snowpark session for anything that involves bind parameters after you call `ibis.snowflake.from_snowpark`.

I think I can find a suitable workaround, but I also don't want to spend a bunch of time on that if it's not a problem that you expect to have.
Just a quick update saying that I have opened up a so-called idea on the [Snowflake Ideas](https://community.snowflake.com/s/ideas) website. It's not possible to link directly to the suggestion to include snowflake-sqlalchemy, but at least now people can search it up and upvote it.
Update on this: Unfortunately, it was too difficult to get connections from snowflake to play nicely as pure DBAPI implementations for use with snowflake-sqlalchemy.

However, we're likely to move away from snowflake-alchemy soon, which may make it easier to implement this feature.
**UPDATE:** 9.0 will ship with a `ibis.snowflake.from_snowpark(session)` API. We run the entire Ibis test suite against a backend created this way, so it should be possible to write Ibis code inside of a Snowpark stored procedure.

I'll investigate adding a test to verify the explicit behavior of running Ibis code inside of a stored proc.

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
@@ -222,8 +222,8 @@ def do_connect(self, create_object_udfs: bool = True, **kwargs: Any):
             create_object_udfs=create_object_udfs,
         )
 
-    def _setup_session(self, *, con, session_parameters, create_object_udfs: bool):
-        self.con = con
+    def _setup_session(self, *, session_parameters, create_object_udfs: bool):
+        con = self.con
 
         # enable multiple SQL statements by default
         session_parameters.setdefault("MULTI_STATEMENT_COUNT", 0)
@@ -318,12 +318,15 @@ def from_snowpark(cls, session, *, create_object_udfs: bool = True) -> Backend:
         │ ansonca01 │    16 │
         └───────────┴───────┘
         """
+        import snowflake.connector
+
         backend = cls(_from_snowpark=True)
-        backend._setup_session(
-            con=session._conn._conn,
-            session_parameters={},
-            create_object_udfs=create_object_udfs,
-        )
+        backend.con = session._conn._conn
+        with contextlib.suppress(snowflake.connector.errors.ProgrammingError):
+            # stored procs on snowflake don't allow session mutation it seems
+            backend._setup_session(
+                session_parameters={}, create_object_udfs=create_object_udfs
+            )
         return backend
 
     def reconnect(self) -> None:


Test patch:
diff --git a/ibis/backends/snowflake/tests/test_udf.py b/ibis/backends/snowflake/tests/test_udf.py
--- a/ibis/backends/snowflake/tests/test_udf.py
+++ b/ibis/backends/snowflake/tests/test_udf.py
@@ -1,11 +1,17 @@
 from __future__ import annotations
 
+import os
+import shutil
+import tempfile
+
 import pandas.testing as tm
 import pytest
 from pytest import param
 
+import ibis
 import ibis.expr.datatypes as dt
 from ibis import udf
+from ibis.backends.tests.errors import SnowflakeProgrammingError
 
 
 @udf.scalar.builtin
@@ -89,7 +95,6 @@ def test_builtin_agg_udf(con):
 
 
 def test_xgboost_model(con):
-    import ibis
     from ibis import _
 
     @udf.scalar.pandas(
@@ -150,3 +155,135 @@ def cases(value, mapping):
     assert not df.empty
     assert "predicted_price" in df.columns
     assert len(df) == diamonds.count().execute()
+
+
+def touch_package(pkgpath):
+    # nix timestamps every file to the UNIX epoch for reproducibility
+    # so we modify the utime of the _copied_ code since snowflake has
+    # some annoying checks for zipping files that do not allow files
+    # older than 1980 ¯\_(ツ)_/¯
+    os.utime(pkgpath, None)
+    for root, dirs, files in os.walk(pkgpath):
+        for path in dirs + files:
+            os.utime(os.path.join(root, path), None)
+
+
+def add_packages(d, session):
+    import parsy
+    import pyarrow_hotfix
+    import rich
+    import sqlglot
+
+    import ibis
+
+    for module in (rich, parsy, sqlglot, pyarrow_hotfix):
+        pkgname = module.__name__
+        pkgpath = os.path.join(d, pkgname)
+        shutil.copytree(os.path.dirname(module.__file__), pkgpath)
+        touch_package(pkgpath)
+        session.add_import(pkgname, import_path=pkgname)
+
+    # no need to touch the package because we're using the local version
+    shutil.copytree(os.path.dirname(ibis.__file__), "ibis")
+    session.add_import("ibis", import_path="ibis")
+
+
+@pytest.fixture
+def snowpark_session():
+    if not os.environ.get("SNOWFLAKE_SNOWPARK"):
+        pytest.skip("SNOWFLAKE_SNOWPARK is not set")
+    else:
+        sp = pytest.importorskip("snowflake.snowpark")
+
+        if connection_name := os.environ.get("SNOWFLAKE_DEFAULT_CONNECTION_NAME"):
+            builder = sp.Session.builder.config("connection_name", connection_name)
+        else:
+            builder = sp.Session.builder.configs(
+                {
+                    "user": os.environ["SNOWFLAKE_USER"],
+                    "account": os.environ["SNOWFLAKE_ACCOUNT"],
+                    "password": os.environ["SNOWFLAKE_PASSWORD"],
+                    "warehouse": os.environ["SNOWFLAKE_WAREHOUSE"],
+                    "database": os.environ["SNOWFLAKE_DATABASE"],
+                    "schema": os.environ["SNOWFLAKE_SCHEMA"],
+                }
+            )
+
+        session = builder.create()
+        session.custom_package_usage_config["enabled"] = True
+
+        pwd = os.getcwd()
+
+        with tempfile.TemporaryDirectory() as d:
+            os.chdir(d)
+
+            try:
+                add_packages(d, session)
+                yield session
+            finally:
+                os.chdir(pwd)
+                session.clear_imports()
+
+
+@pytest.mark.parametrize(
+    "execute_as",
+    [
+        "owner",
+        param("caller", marks=[pytest.mark.xfail(raises=SnowflakeProgrammingError)]),
+    ],
+)
+def test_ibis_inside_snowpark(snowpark_session, execute_as):
+    import snowflake.snowpark as sp
+
+    def ibis_sproc(session):
+        import ibis.backends.snowflake
+
+        con = ibis.backends.snowflake.Backend.from_snowpark(session)
+
+        expr = (
+            con.tables.functional_alltypes.group_by("string_col")
+            .agg(n=lambda t: t.count())
+            .order_by("string_col")
+        )
+
+        return session.sql(ibis.to_sql(expr))
+
+    expected = (
+        snowpark_session.table('"functional_alltypes"')
+        .group_by('"string_col"')
+        .count()
+        .rename("COUNT", '"n"')
+        .order_by('"string_col"')
+        .to_pandas()
+    )
+
+    local_result = ibis_sproc(snowpark_session).to_pandas()
+
+    tm.assert_frame_equal(local_result, expected)
+
+    name = ibis_sproc.__name__
+
+    snowpark_session.sproc.register(
+        ibis_sproc,
+        name=name,
+        execute_as=execute_as,
+        imports=["parsy", "rich", "sqlglot", "pyarrow_hotfix", "ibis"],
+        # empty struct here tells Snowflake to infer the return type from the
+        # return value of the function, which is required to be a Snowpark
+        # table in that case
+        return_type=sp.types.StructType(),
+        packages=[
+            "snowflake-snowpark-python",
+            "toolz",
+            "atpublic",
+            "bidict",
+            "pyarrow",
+            "pandas",
+            "numpy",
+        ],
+        replace=True,
+    )
+
+    remote_result = snowpark_session.call(name).to_pandas()
+
+    tm.assert_frame_equal(remote_result, local_result)