Add support pandas 2 in flytekit

dev-requirements.in

@@ -35,11 +35,15 @@ 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
+pandas
+pyarrow
+numpy
+scikit-learn
 types-requests
 prometheus-client

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)

setup.py

@@ -53,11 +53,8 @@
         # 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",
         "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/core/test_type_hints.py

@@ -858,7 +858,7 @@ def t1(a: int) -> pandas.DataFrame:
 
     @task
     def t2(df: pandas.DataFrame) -> pandas.DataFrame:
-        return df.append(pandas.DataFrame(data={"col1": [5, 10], "col2": [5, 10]}))
+        return pd.concat([df, pandas.DataFrame(data={"col1": [5, 10], "col2": [5, 10]})])
 
     @workflow
     def my_wf(a: int) -> pandas.DataFrame:

tests/flytekit/unit/lazy_module/test_lazy_module.py

@@ -0,0 +1,12 @@
+import pytest
+
+from flytekit.lazy_import.lazy_module import LazyModule, lazy_module
+
+
+def test_lazy_module():
+    mod = lazy_module("pandas")
+    assert mod.__name__ == "pandas"
+    mod = lazy_module("fake_module")
+    assert isinstance(mod, LazyModule)
+    with pytest.raises(ImportError, match="Module fake_module is not yet installed."):
+        print(mod.attr)