Add support pandas 2 in flytekit

.github/workflows/pythonbuild.yml

@@ -21,6 +21,12 @@ jobs:
       matrix:
         os: [ubuntu-latest, windows-latest, macos-latest]
         python-version: ["3.8", "3.11"]
+        install-pandas: ["true", "false"]
+        exclude:
+          - os: macos-latest
+            install-pandas: "false"
+          - os: windows-latest
+            install-pandas: "false"
     steps:
       - uses: actions/checkout@v4
       - name: Set up Python ${{ matrix.python-version }}
@@ -34,6 +40,12 @@ jobs:
           path: ~/.cache/pip
           # Look to see if there is a cache hit for the corresponding requirements files
           key: ${{ format('{0}-pip-{1}', runner.os, hashFiles('dev-requirements.in', 'requirements.in')) }}
+      - name: Install pandas v1
+        if: ${{ (matrix.python-version=='3.8') && (matrix.install-pandas=='true') }}
+        run: pip install "pandas<2.0.0"
+      - name: Install pandas v2
+        if: ${{ matrix.python-version=='3.11' && (matrix.install-pandas=='true')  }}
+        run: pip install "pandas>=2.0.0"
       - name: Install dependencies
         run: |
           make setup && pip freeze
@@ -106,7 +118,9 @@ jobs:
           # Look to see if there is a cache hit for the corresponding requirements files
           key: ${{ format('{0}-pip-{1}', runner.os, hashFiles('dev-requirements.in', 'requirements.in')) }}
       - name: Install dependencies
-        run: make setup && pip freeze
+        run: |
+          pip install pandas
+          make setup && pip freeze
       - name: Install FlyteCTL
         uses: unionai-oss/flytectl-setup-action@master
       - name: Setup Flyte Sandbox
@@ -217,6 +231,7 @@ jobs:
           key: ${{ format('{0}-pip-{1}', runner.os, hashFiles('dev-requirements.txt', format('plugins/{0}/requirements.txt', matrix.plugin-names ))) }}
       - name: Install dependencies
         run: |
+          pip install pandas
           make setup
           cd plugins/${{ matrix.plugin-names }}
           pip install .

dev-requirements.in

@@ -35,11 +35,13 @@ torch<=2.0.0; python_version>='3.11' or platform_system!='Windows'
 # Once a solution is found, this should be updated to support Windows as well.
 python-magic; (platform_system=='Darwin' or platform_system=='Linux')
 
-pillow
-scikit-learn
 types-protobuf
 types-croniter
 types-mock
 autoflake
+
+pillow
+numpy
+scikit-learn
 types-requests
 prometheus-client

flytekit/extras/sqlite3/task.py

@@ -5,17 +5,20 @@
 import tempfile
 import typing
 from dataclasses import dataclass
-
-import pandas as pd
-
-from flytekit import FlyteContext, kwtypes
+from flytekit import FlyteContext, kwtypes, lazy_module
 from flytekit.configuration import DefaultImages, SerializationSettings
 from flytekit.core.base_sql_task import SQLTask
 from flytekit.core.python_customized_container_task import PythonCustomizedContainerTask
 from flytekit.core.shim_task import ShimTaskExecutor
 from flytekit.models import task as task_models
 
 
+if typing.TYPE_CHECKING:
+    import pandas as pd
+else:
+    pd = lazy_module("pandas")
+
+
 def unarchive_file(local_path: str, to_dir: str):
     """
     Unarchive given archive and returns the unarchived file name. It is expected that only one file is unarchived.

flytekit/lazy_import/lazy_module.py

@@ -1,9 +1,19 @@
 import importlib.util
 import sys
+import types
 
 LAZY_MODULES = []
 
 
+class LazyModule(types.ModuleType):
+    def __init__(self, module_name: str):
+        super().__init__(module_name)
+        self._module_name = module_name
+
+    def __getattribute__(self, attr):
+        raise ImportError(f"Module {object.__getattribute__(self, '_module_name')} is not yet installed.")
+
+
 def is_imported(module_name):
     """
     This function is used to check if a module has been imported by the regular import.
@@ -24,6 +34,10 @@ def lazy_module(fullname):
         return sys.modules[fullname]
     # https://docs.python.org/3/library/importlib.html#implementing-lazy-imports
     spec = importlib.util.find_spec(fullname)
+    if spec is None:
+        # Return a dummy module if the module is not found in the python environment,
+        # so that we can raise a proper error when the user tries to access an attribute in the module.
+        return LazyModule(fullname)
     loader = importlib.util.LazyLoader(spec.loader)
     spec.loader = loader
     module = importlib.util.module_from_spec(spec)

flytekit/types/structured/basic_dfs.py

@@ -3,14 +3,11 @@
 from pathlib import Path
 from typing import TypeVar
 
-import pandas as pd
-import pyarrow as pa
-import pyarrow.parquet as pq
 from botocore.exceptions import NoCredentialsError
 from fsspec.core import split_protocol, strip_protocol
 from fsspec.utils import get_protocol
 
-from flytekit import FlyteContext, logger
+from flytekit import FlyteContext, lazy_module, logger
 from flytekit.configuration import DataConfig
 from flytekit.core.data_persistence import get_fsspec_storage_options
 from flytekit.models import literals
@@ -24,6 +21,13 @@
     StructuredDatasetEncoder,
 )
 
+if typing.TYPE_CHECKING:
+    import pandas as pd
+    import pyarrow as pa
+else:
+    pd = lazy_module("pandas")
+    pa = lazy_module("pyarrow")
+
 T = TypeVar("T")
 
 
@@ -70,7 +74,7 @@
         ctx: FlyteContext,
         flyte_value: literals.StructuredDataset,
         current_task_metadata: StructuredDatasetMetadata,
-    ) -> pd.DataFrame:
+    ) -> "pd.DataFrame":
         uri = flyte_value.uri
         columns = None
         kwargs = get_pandas_storage_options(uri=uri, data_config=ctx.file_access.data_config)
@@ -121,7 +125,7 @@
         ctx: FlyteContext,
         flyte_value: literals.StructuredDataset,
         current_task_metadata: StructuredDatasetMetadata,
-    ) -> pd.DataFrame:
+    ) -> "pd.DataFrame":
         uri = flyte_value.uri
         columns = None
         kwargs = get_pandas_storage_options(uri=uri, data_config=ctx.file_access.data_config)
@@ -145,9 +149,9 @@
         structured_dataset: StructuredDataset,
         structured_dataset_type: StructuredDatasetType,
     ) -> literals.StructuredDataset:
-        uri = typing.cast(str, structured_dataset.uri) or ctx.file_access.join(
-            ctx.file_access.raw_output_prefix, ctx.file_access.get_random_string()
-        )
+        import pyarrow.parquet as pq
+
+        uri = typing.cast(str, structured_dataset.uri) or ctx.file_access.get_random_remote_directory()
         if not ctx.file_access.is_remote(uri):
             Path(uri).mkdir(parents=True, exist_ok=True)
         path = os.path.join(uri, f"{0:05}")
@@ -165,7 +169,9 @@
         ctx: FlyteContext,
         flyte_value: literals.StructuredDataset,
         current_task_metadata: StructuredDatasetMetadata,
-    ) -> pa.Table:
+    ) -> "pa.Table":
+        import pyarrow.parquet as pq
+
         uri = flyte_value.uri
         if not ctx.file_access.is_remote(uri):
             Path(uri).parent.mkdir(parents=True, exist_ok=True)

flytekit/types/structured/bigquery.py

@@ -1,12 +1,10 @@
 import re
 import typing
 
-import pandas as pd
-import pyarrow as pa
 from google.cloud import bigquery, bigquery_storage
 from google.cloud.bigquery_storage_v1 import types
 
-from flytekit import FlyteContext
+from flytekit import FlyteContext, lazy_module
 from flytekit.models import literals
 from flytekit.models.types import StructuredDatasetType
 from flytekit.types.structured.structured_dataset import (
@@ -16,6 +14,13 @@
     StructuredDatasetMetadata,
 )
 
+if typing.TYPE_CHECKING:
+    import pandas as pd
+    import pyarrow as pa
+else:
+    pd = lazy_module("pandas")
+    pa = lazy_module("pyarrow")
+
 BIGQUERY = "bq"
 
 

flytekit/types/structured/structured_dataset.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import collections
+import sys
 import types
 import typing
 from abc import ABC, abstractmethod
@@ -150,7 +151,7 @@ def extract_cols_and_format(
                 if ordered_dict_cols is not None:
                     raise ValueError(f"Column information was already found {ordered_dict_cols}, cannot use {aa}")
                 ordered_dict_cols = aa
-            elif isinstance(aa, pa.lib.Schema):
+            elif "pyarrow" in sys.modules and isinstance(aa, pa.lib.Schema):
                 if pa_schema is not None:
                     raise ValueError(f"Arrow schema was already found {pa_schema}, cannot use {aa}")
                 pa_schema = aa

plugins/flytekit-duckdb/setup.py

@@ -4,7 +4,7 @@
 
 microlib_name = f"flytekitplugins-{PLUGIN_NAME}"
 
-plugin_requires = ["flytekit>=1.3.0b2,<2.0.0", "duckdb"]
+plugin_requires = ["flytekit>=1.3.0b2,<2.0.0", "duckdb", "pandas"]
 
 __version__ = "0.0.0+develop"
 

plugins/flytekit-mlflow/setup.py

@@ -5,7 +5,7 @@
 microlib_name = f"flytekitplugins-{PLUGIN_NAME}"
 
 # TODO: support mlflow 2.0+
-plugin_requires = ["flytekit>=1.1.0,<2.0.0", "plotly", "mlflow<2.0.0"]
+plugin_requires = ["flytekit>=1.1.0,<2.0.0", "plotly", "mlflow<2.0.0", "pandas"]
 
 __version__ = "0.0.0+develop"
 

plugins/flytekit-polars/setup.py

@@ -4,10 +4,7 @@
 
 microlib_name = f"flytekitplugins-{PLUGIN_NAME}"
 
-plugin_requires = [
-    "flytekit>=1.3.0b2,<2.0.0",
-    "polars>=0.8.27,<0.17.0",
-]
+plugin_requires = ["flytekit>=1.3.0b2,<2.0.0", "polars>=0.8.27,<0.17.0", "pandas"]
 
 __version__ = "0.0.0+develop"
 

plugins/flytekit-spark/setup.py

@@ -4,7 +4,7 @@
 
 microlib_name = f"flytekitplugins-{PLUGIN_NAME}"
 
-plugin_requires = ["flytekit>=1.3.0b2,<2.0.0", "pyspark>=3.0.0", "aiohttp", "flyteidl>=1.10.0"]
+plugin_requires = ["flytekit>=1.3.0b2,<2.0.0", "pyspark>=3.0.0", "aiohttp", "flyteidl>=1.10.0", "pandas"]
 
 __version__ = "0.0.0+develop"
 

plugins/flytekit-sqlalchemy/setup.py

@@ -4,7 +4,7 @@
 
 microlib_name = f"flytekitplugins-{PLUGIN_NAME}"
 
-plugin_requires = ["flytekit>=1.3.0b2,<2.0.0", "sqlalchemy>=1.4.7"]
+plugin_requires = ["flytekit>=1.3.0b2,<2.0.0", "sqlalchemy>=1.4.7", "pandas"]
 
 __version__ = "0.0.0+develop"
 

setup.py

@@ -53,11 +53,9 @@
         # TODO: remove upper-bound after fixing change in contract
         "marshmallow-jsonschema>=0.12.0",
         "mashumaro>=3.9.1",
-        "numpy",
-        "pandas>=1.0.0,<2.0.0",
         # TODO: Remove upper-bound after protobuf community fixes it. https://github.com/flyteorg/flyte/issues/4359
         "protobuf<4.25.0",
-        "pyarrow>=4.0.0,<11.0.0",
+        "pyarrow",
         "python-json-logger>=2.0.0",
         "pytimeparse>=1.1.8,<2.0.0",
         "pyyaml!=6.0.0,!=5.4.0,!=5.4.1",  # pyyaml is broken with cython 3: https://github.com/yaml/pyyaml/issues/601

tests/flytekit/unit/cli/pyflyte/test_run.py

@@ -3,7 +3,6 @@
 import pathlib
 import sys
 from enum import Enum
-
 import mock
 import pytest
 from click.testing import CliRunner
@@ -16,6 +15,8 @@
 from flytekit.interaction.click_types import FileParamType
 from flytekit.remote import FlyteRemote
 
+pytest.importorskip("pandas")
+
 WORKFLOW_FILE = os.path.join(os.path.dirname(os.path.realpath(__file__)), "workflow.py")
 REMOTE_WORKFLOW_FILE = "https://raw.githubusercontent.com/flyteorg/flytesnacks/8337b64b33df046b2f6e4cba03c74b7bdc0c4fb1/cookbook/core/flyte_basics/basic_workflow.py"
 IMPERATIVE_WORKFLOW_FILE = os.path.join(os.path.dirname(os.path.realpath(__file__)), "imperative_wf.py")

tests/flytekit/unit/core/test_checkpoint.py

@@ -7,6 +7,8 @@
 from flytekit.core.checkpointer import SyncCheckpoint
 from flytekit.core.local_cache import LocalTaskCache
 
+pytest.importorskip("pandas")
+
 
 def test_sync_checkpoint_write(tmpdir):
     td_path = Path(tmpdir)

tests/flytekit/unit/core/test_complex_nesting.py

@@ -13,6 +13,8 @@
 from flytekit.types.directory import FlyteDirectory
 from flytekit.types.file import FlyteFile
 
+pytest.importorskip("pandas")
+
 
 @dataclass
 class MyProxyConfiguration(DataClassJsonMixin):

tests/flytekit/unit/core/test_data_persistence.py

@@ -1,3 +1,4 @@
+import pytest
 import io
 import os
 import pathlib
@@ -6,11 +7,11 @@
 import tempfile
 
 import mock
-import pandas as pd
 from azure.identity import ClientSecretCredential, DefaultAzureCredential
-
 from flytekit.core.data_persistence import FileAccessProvider
 
+pd = pytest.importorskip("pandas")
+
 
 def test_get_manual_random_remote_path():
     fp = FileAccessProvider("/tmp", "s3://my-bucket")

tests/flytekit/unit/core/test_dataclass.py

@@ -1,11 +1,14 @@
 from dataclasses import dataclass
 from typing import List
 
+import pytest
 from dataclasses_json import DataClassJsonMixin
 
 from flytekit.core.task import task
 from flytekit.core.workflow import workflow
 
+pytest.importorskip("pandas")
+
 
 def test_dataclass():
     @dataclass

tests/flytekit/unit/core/test_imperative.py

@@ -1,9 +1,6 @@
+import pytest
 import typing
 from collections import OrderedDict
-
-import pandas as pd
-import pytest
-
 import flytekit.configuration
 from flytekit.configuration import Image, ImageConfig
 from flytekit.core.base_task import kwtypes
@@ -15,7 +12,13 @@
 from flytekit.models import literals as literal_models
 from flytekit.tools.translator import get_serializable
 from flytekit.types.file import FlyteFile
-from flytekit.types.schema import FlyteSchema
+
+if typing.TYPE_CHECKING:
+    import pandas as pd
+else:
+    pd = pytest.importorskip("pandas")
+from flytekit.types.schema import FlyteSchema  # noqa: E402
+
 
 default_img = Image(name="default", fqn="test", tag="tag")
 serialization_settings = flytekit.configuration.SerializationSettings(