SNOW-900365: Add Python 3.11 merge gate (#1040)

snowflakedb · Sep 6, 2023 · 27985eb · 27985eb
1 parent 78b37d8
commit 27985eb
Show file tree

Hide file tree

Showing 6 changed files with 69 additions and 16 deletions.
diff --git a/.github/workflows/precommit.yml b/.github/workflows/precommit.yml
@@ -72,7 +72,7 @@ jobs:
             download_name: macos
           - image_name: windows-latest
             download_name: windows
-        python-version: ["3.8", "3.9", "3.10"]
+        python-version: ["3.8", "3.9", "3.10", "3.11"]
         cloud-provider: [aws, azure, gcp]
     steps:
       - name: Checkout Code
@@ -101,7 +101,7 @@ jobs:
         run: python -m pip install -U setuptools pip wheel
       - name: Install tox
         run: python -m pip install tox
-      - if: ${{ contains('macos', matrix.os.download_name) }}
+      - if: ${{ matrix.python-version != '3.11' && contains('macos', matrix.os.download_name) }}
         name: Run doctests
         run: python -m tox -e "py${PYTHON_VERSION}-doctest-notudf-ci"
         env:
@@ -112,14 +112,25 @@ jobs:
           # Specify SNOWFLAKE_IS_PYTHON_RUNTIME_TEST: 1 when adding >= python3.11 with no server-side support
           # For example, see https://github.com/snowflakedb/snowpark-python/pull/681
         shell: bash
-      - name: Run tests (excluding doctests)
+      - if: ${{ matrix.python-version != '3.11' }}
+        name: Run tests (excluding doctests)
         run: python -m tox -e "py${PYTHON_VERSION/\./}-notdoctest-ci"
         env:
           PYTHON_VERSION: ${{ matrix.python-version }}
           cloud_provider: ${{ matrix.cloud-provider }}
           PYTEST_ADDOPTS: --color=yes --tb=short
           TOX_PARALLEL_NO_SPINNER: 1
         shell: bash
+      - if: ${{ matrix.python-version == '3.11' }}
+        name: Run tests (excluding udf and doctest tests)
+        run: python -m tox -e "py${PYTHON_VERSION/\./}-notudfdoctest"
+        env:
+          PYTHON_VERSION: ${{ matrix.python-version }}
+          cloud_provider: ${{ matrix.cloud-provider }}
+          PYTEST_ADDOPTS: --color=yes --tb=short
+          TOX_PARALLEL_NO_SPINNER: 1
+          SNOWFLAKE_IS_PYTHON_RUNTIME_TEST: 1
+        shell: bash
       - name: Combine coverages
         run: python -m tox -e coverage --skip-missing-interpreters false
         shell: bash
@@ -182,7 +193,7 @@ jobs:
         os:
           - image_name: macos-latest
             download_name: macos  # it includes doctest
-        python-version: ["3.8", "3.9", "3.10"]
+        python-version: ["3.8", "3.9", "3.10", "3.11"]
         cloud-provider: [aws]
     steps:
       - name: Checkout Code
@@ -211,7 +222,7 @@ jobs:
         run: python -m pip install -U setuptools pip wheel
       - name: Install tox
         run: python -m pip install tox
-      - if: ${{ contains('macos', matrix.os.download_name) }}
+      - if: ${{ matrix.python-version != '3.11' && contains('macos', matrix.os.download_name) }}
         name: Run doctests
         run: python -m tox -e "py${PYTHON_VERSION}-doctest-notudf-ci"
         env:
@@ -220,14 +231,25 @@ jobs:
           PYTEST_ADDOPTS: --color=yes --tb=short --disable_sql_simplifier
           TOX_PARALLEL_NO_SPINNER: 1
         shell: bash
-      - name: Run tests (excluding doctests)
+      - if: ${{ matrix.python-version != '3.11' }}
+        name: Run tests (excluding doctests)
         run: python -m tox -e "py${PYTHON_VERSION/\./}-notdoctest-ci"
         env:
           PYTHON_VERSION: ${{ matrix.python-version }}
           cloud_provider: ${{ matrix.cloud-provider }}
           PYTEST_ADDOPTS: --color=yes --tb=short --disable_sql_simplifier
           TOX_PARALLEL_NO_SPINNER: 1
         shell: bash
+      - if: ${{ matrix.python-version == '3.11' }}
+        name: Run tests (excluding udf and doctest tests)
+        run: python -m tox -e "py${PYTHON_VERSION/\./}-notudfdoctest"
+        env:
+          PYTHON_VERSION: ${{ matrix.python-version }}
+          cloud_provider: ${{ matrix.cloud-provider }}
+          PYTEST_ADDOPTS: --color=yes --tb=short --disable_sql_simplifier
+          TOX_PARALLEL_NO_SPINNER: 1
+          SNOWFLAKE_IS_PYTHON_RUNTIME_TEST: 1
+        shell: bash
       - name: Combine coverages
         run: python -m tox -e coverage --skip-missing-interpreters false
         shell: bash

diff --git a/src/snowflake/snowpark/_internal/code_generation.py b/src/snowflake/snowpark/_internal/code_generation.py
@@ -180,12 +180,12 @@ def get_class_references(
 
 def extract_func_global_refs(code: CodeType) -> Set[str]:
     # inspired by cloudpickle to recursively extract all the global references used by the target func's code object
-    co_names = code.co_names
+    # check: https://github.com/cloudpipe/cloudpickle/commit/6a0e12d058d1bd3ab26ec000ac2249b4ee7e9c9f
     out_names = set()
     for instr in dis.get_instructions(code):
         op = instr.opcode
         if op in GLOBAL_OPS:
-            out_names.add(co_names[instr.arg])
+            out_names.add(instr.argval)
 
     if code.co_consts:
         for const in code.co_consts:

diff --git a/src/snowflake/snowpark/types.py b/src/snowflake/snowpark/types.py
@@ -485,11 +485,25 @@ class PandasSeries(pandas.Series, Generic[_T]):
 
     _TT = TypeVarTuple("_TT")
 
-    class PandasDataFrame(pandas.DataFrame, Generic[_TT]):
-        """
-        The type hint for annotating Pandas DataFrame data when registering UDFs.
-        The input should be a list of data types for all columns in order.
-        It cannot be used to annotate the return value of a Pandas UDF.
-        """
+    if sys.version_info >= (3, 11):
+        from typing import Unpack
 
-        pass
+        class PandasDataFrame(pandas.DataFrame, Generic[Unpack[_TT]]):
+            """
+            The type hint for annotating Pandas DataFrame data when registering UDFs.
+            The input should be a list of data types for all columns in order.
+            It cannot be used to annotate the return value of a Pandas UDF.
+            """
+
+            pass
+
+    else:
+
+        class PandasDataFrame(pandas.DataFrame, Generic[_TT]):
+            """
+            The type hint for annotating Pandas DataFrame data when registering UDFs.
+            The input should be a list of data types for all columns in order.
+            It cannot be used to annotate the return value of a Pandas UDF.
+            """
+
+            pass
diff --git a/tests/integ/test_dataframe.py b/tests/integ/test_dataframe.py
@@ -593,6 +593,7 @@ def process(self, n: int) -> Iterable[Tuple[int, int]]:
     )
 
 
+@pytest.mark.udf
 def test_select_with_table_function_column_overlap(session):
     df = session.create_dataframe([[1, 2, 3], [4, 5, 6]], schema=["A", "B", "C"])
 
@@ -2369,6 +2370,7 @@ def test_save_as_table_nullable_test(session, save_mode, table_type):
         Utils.drop_table(session, table_name)
 
 
+@pytest.mark.udf
 @pytest.mark.parametrize("table_type", ["", "temp", "temporary", "transient"])
 @pytest.mark.parametrize(
     "save_mode", ["append", "overwrite", "ignore", "errorifexists"]

diff --git a/tests/integ/test_packaging.py b/tests/integ/test_packaging.py
@@ -175,6 +175,7 @@ def test_patch_on_get_available_versions_for_packages(session):
     assert "catboost" not in returned
 
 
+@pytest.mark.udf
 @pytest.mark.skipif(
     (not is_pandas_and_numpy_available) or IS_IN_STORED_PROC,
     reason="numpy and pandas are required",
@@ -254,6 +255,7 @@ def is_yaml_available() -> bool:
     session.clear_packages()
 
 
+@pytest.mark.udf
 def test_add_packages_with_underscore(session):
     packages = ["spacy-model-en_core_web_sm", "typing_extensions"]
     count = (
@@ -315,6 +317,7 @@ def test_add_packages_negative(session, caplog):
         session.remove_package("python-dateutil")
 
 
+@pytest.mark.udf
 @pytest.mark.skipif(
     (not is_pandas_and_numpy_available) or IS_IN_STORED_PROC,
     reason="numpy and pandas are required",
@@ -409,6 +412,7 @@ def test_add_unsupported_requirements_twice_should_not_fail_for_same_requirement
         assert "pyyaml" in package_set
 
 
+@pytest.mark.udf
 @pytest.mark.skipif(
     IS_IN_STORED_PROC,
     reason="Subprocess calls are not allowed within stored procedures.",
@@ -442,6 +446,7 @@ def arch_function() -> list:
         )
 
 
+@pytest.mark.udf
 @pytest.mark.skipif(
     IS_IN_STORED_PROC,
     reason="Subprocess calls are not allowed within stored procedures.",
@@ -470,6 +475,7 @@ def import_scikit_fuzzy() -> str:
     Utils.check_answer(session.sql(f"select {udf_name}()"), [Row("0.4.2")])
 
 
+@pytest.mark.udf
 @pytest.mark.skipif(
     IS_IN_STORED_PROC,
     reason="Subprocess calls are not allowed within stored procedures.",
@@ -499,6 +505,7 @@ def run_scikit_fuzzy(_: Session) -> str:
     assert run_scikit_fuzzy(session) == "0.4.2"
 
 
+@pytest.mark.udf
 @pytest.mark.skipif(
     IS_IN_STORED_PROC,
     reason="Subprocess calls are not allowed within stored procedures.",
@@ -568,6 +575,7 @@ def requirements_file_with_local_path():
             os.remove(path)
 
 
+@pytest.mark.udf
 @pytest.mark.skipif(
     IS_IN_STORED_PROC,
     reason="matplotlib required",
@@ -656,6 +664,7 @@ def test_add_requirements_with_ranged_requirements_in_yaml(session, ranged_yaml_
         session.add_requirements(ranged_yaml_file)
 
 
+@pytest.mark.udf
 @pytest.mark.skipif(
     IS_IN_STORED_PROC,
     reason="Subprocess calls are not allowed within stored procedures.",
@@ -684,6 +693,7 @@ def check_if_package_works() -> str:
         )
 
 
+@pytest.mark.udf
 @pytest.mark.skipif(
     IS_IN_STORED_PROC,
     reason="Subprocess calls are not allowed within stored procedures.",
@@ -709,6 +719,7 @@ def check_if_package_works(session_):
         assert check_if_package_works() == "0.4.2"
 
 
+@pytest.mark.udf
 @pytest.mark.skipif(not is_dateutil_available, reason="dateutil is required")
 def test_add_import_package(session):
     def plus_one_month(x):
@@ -726,6 +737,7 @@ def plus_one_month(x):
     )
 
 
+@pytest.mark.udf
 @pytest.mark.skipif(
     IS_IN_STORED_PROC,
     reason="numpy and pandas are required",
@@ -759,6 +771,7 @@ def get_numpy_pandas_version() -> str:
     Utils.check_answer(session.sql(f"select {udf_name}()"), [Row("1.3.0")])
 
 
+@pytest.mark.udf
 @pytest.mark.skipif(
     IS_IN_STORED_PROC,
     reason="Subprocess calls are not allowed within stored procedures.",
@@ -816,6 +829,7 @@ def test_add_requirements_unsupported_with_cache_path_negative(
             session.add_requirements(test_files.test_unsupported_requirements_file)
 
 
+@pytest.mark.udf
 @pytest.mark.skipif(
     IS_IN_STORED_PROC,
     reason="Subprocess calls are not allowed within stored procedures.",
@@ -857,6 +871,7 @@ def get_numpy_pandas_version() -> str:
     Utils.check_answer(session.sql(f"select {udf_name}()"), [Row("0.4.2")])
 
 
+@pytest.mark.udf
 @pytest.mark.skipif(
     IS_IN_STORED_PROC,
     reason="Subprocess calls are not allowed within stored procedures.",

diff --git a/tox.ini b/tox.ini
@@ -94,7 +94,7 @@ commands = coverage combine
            coverage report -m
            coverage xml -o {env:COV_REPORT_DIR:{toxworkdir}}/coverage.xml
            coverage html -d {env:COV_REPORT_DIR:{toxworkdir}}/htmlcov --show-contexts
-depends = py38, py39, py310
+depends = py38, py39, py310, py311
 
 [testenv:docs]
 basepython = python3.8