flyte 0.0.1b0__py3-none-any.whl

flyte/io/structured_dataset/__init__.py

@@ -0,0 +1,129 @@
+"""
+Flytekit StructuredDataset
+==========================================================
+.. currentmodule:: flytekit.types.structured
+
+.. autosummary::
+   :template: custom.rst
+   :toctree: generated/
+
+    StructuredDataset
+    StructuredDatasetDecoder
+    StructuredDatasetEncoder
+"""
+
+import functools
+
+from flyte._logging import logger
+from flyte._utils.lazy_module import is_imported
+
+from .structured_dataset import (
+    DuplicateHandlerError,
+    StructuredDataset,
+    StructuredDatasetDecoder,
+    StructuredDatasetEncoder,
+    StructuredDatasetTransformerEngine,
+)
+
+
+@functools.lru_cache(maxsize=None)
+def register_csv_handlers():
+    from .basic_dfs import CSVToPandasDecodingHandler, PandasToCSVEncodingHandler
+
+    StructuredDatasetTransformerEngine.register(PandasToCSVEncodingHandler(), default_format_for_type=True)
+    StructuredDatasetTransformerEngine.register(CSVToPandasDecodingHandler(), default_format_for_type=True)
+
+
+@functools.lru_cache(maxsize=None)
+def register_pandas_handlers():
+    import pandas as pd
+
+    from flyte.types._renderer import TopFrameRenderer
+
+    from .basic_dfs import PandasToParquetEncodingHandler, ParquetToPandasDecodingHandler
+
+    StructuredDatasetTransformerEngine.register(PandasToParquetEncodingHandler(), default_format_for_type=True)
+    StructuredDatasetTransformerEngine.register(ParquetToPandasDecodingHandler(), default_format_for_type=True)
+    StructuredDatasetTransformerEngine.register_renderer(pd.DataFrame, TopFrameRenderer())
+
+
+@functools.lru_cache(maxsize=None)
+def register_arrow_handlers():
+    import pyarrow as pa
+
+    from flyte.types._renderer import ArrowRenderer
+
+    from .basic_dfs import ArrowToParquetEncodingHandler, ParquetToArrowDecodingHandler
+
+    StructuredDatasetTransformerEngine.register(ArrowToParquetEncodingHandler(), default_format_for_type=True)
+    StructuredDatasetTransformerEngine.register(ParquetToArrowDecodingHandler(), default_format_for_type=True)
+    StructuredDatasetTransformerEngine.register_renderer(pa.Table, ArrowRenderer())
+
+
+@functools.lru_cache(maxsize=None)
+def register_bigquery_handlers():
+    try:
+        from .bigquery import (
+            ArrowToBQEncodingHandlers,
+            BQToArrowDecodingHandler,
+            BQToPandasDecodingHandler,
+            PandasToBQEncodingHandlers,
+        )
+
+        StructuredDatasetTransformerEngine.register(PandasToBQEncodingHandlers())
+        StructuredDatasetTransformerEngine.register(BQToPandasDecodingHandler())
+        StructuredDatasetTransformerEngine.register(ArrowToBQEncodingHandlers())
+        StructuredDatasetTransformerEngine.register(BQToArrowDecodingHandler())
+    except ImportError:
+        logger.info(
+            "We won't register bigquery handler for structured dataset because "
+            "we can't find the packages google-cloud-bigquery-storage and google-cloud-bigquery"
+        )
+
+
+@functools.lru_cache(maxsize=None)
+def register_snowflake_handlers():
+    try:
+        from .snowflake import PandasToSnowflakeEncodingHandlers, SnowflakeToPandasDecodingHandler
+
+        StructuredDatasetTransformerEngine.register(SnowflakeToPandasDecodingHandler())
+        StructuredDatasetTransformerEngine.register(PandasToSnowflakeEncodingHandlers())
+
+    except ImportError:
+        logger.info(
+            "We won't register snowflake handler for structured dataset because "
+            "we can't find package snowflake-connector-python"
+        )
+
+
+def lazy_import_structured_dataset_handler():
+    if is_imported("pandas"):
+        try:
+            register_pandas_handlers()
+            register_csv_handlers()
+        except DuplicateHandlerError:
+            logger.debug("Transformer for pandas is already registered.")
+    if is_imported("pyarrow"):
+        try:
+            register_arrow_handlers()
+        except DuplicateHandlerError:
+            logger.debug("Transformer for arrow is already registered.")
+    if is_imported("google.cloud.bigquery"):
+        try:
+            register_bigquery_handlers()
+        except DuplicateHandlerError:
+            logger.debug("Transformer for bigquery is already registered.")
+    if is_imported("snowflake.connector"):
+        try:
+            register_snowflake_handlers()
+        except DuplicateHandlerError:
+            logger.debug("Transformer for snowflake is already registered.")
+
+
+__all__ = [
+    "StructuredDataset",
+    "StructuredDatasetDecoder",
+    "StructuredDatasetEncoder",
+    "StructuredDatasetTransformerEngine",
+    "lazy_import_structured_dataset_handler",
+]

flyte/io/structured_dataset/basic_dfs.py

@@ -0,0 +1,219 @@
+import os
+import typing
+from pathlib import Path
+from typing import TypeVar
+
+from flyteidl.core import literals_pb2, types_pb2
+from fsspec.core import split_protocol, strip_protocol
+
+import flyte.storage as storage
+from flyte._logging import logger
+from flyte._utils import lazy_module
+from flyte.io.structured_dataset.structured_dataset import (
+    CSV,
+    PARQUET,
+    StructuredDataset,
+    StructuredDatasetDecoder,
+    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")
+
+
+# pr: add back after storage
+def get_pandas_storage_options(uri: str, data_config=None, anonymous: bool = False) -> typing.Optional[typing.Dict]:
+    from pandas.io.common import is_fsspec_url  # type: ignore
+
+    if is_fsspec_url(uri):
+        if uri.startswith("s3"):
+            # pr: after storage, replace with real call to get_fsspec_storage_options
+            return {
+                "cache_regions": True,
+                "client_kwargs": {"endpoint_url": "http://localhost:30002"},
+                "key": "minio",
+                "secret": "miniostorage",
+            }
+        return {}
+
+    # Pandas does not allow storage_options for non-fsspec paths e.g. local.
+    return None
+
+
+class PandasToCSVEncodingHandler(StructuredDatasetEncoder):
+    def __init__(self):
+        super().__init__(pd.DataFrame, None, CSV)
+
+    async def encode(
+        self,
+        structured_dataset: StructuredDataset,
+        structured_dataset_type: types_pb2.StructuredDatasetType,
+    ) -> literals_pb2.StructuredDataset:
+        if not structured_dataset.uri:
+            from flyte._context import internal_ctx
+
+            ctx = internal_ctx()
+            uri = ctx.raw_data.get_random_remote_path()
+        else:
+            uri = typing.cast(str, structured_dataset.uri)
+
+        if not storage.is_remote(uri):
+            Path(uri).mkdir(parents=True, exist_ok=True)
+        path = os.path.join(uri, ".csv")
+        df = typing.cast(pd.DataFrame, structured_dataset.dataframe)
+        df.to_csv(
+            path,
+            index=False,
+            storage_options=get_pandas_storage_options(uri=path, data_config=None),
+        )
+        structured_dataset_type.format = CSV
+        return literals_pb2.StructuredDataset(
+            uri=uri, metadata=literals_pb2.StructuredDatasetMetadata(structured_dataset_type)
+        )
+
+
+class CSVToPandasDecodingHandler(StructuredDatasetDecoder):
+    def __init__(self):
+        super().__init__(pd.DataFrame, None, CSV)
+
+    async def decode(
+        self,
+        proto_value: literals_pb2.StructuredDataset,
+        current_task_metadata: literals_pb2.StructuredDatasetMetadata,
+    ) -> "pd.DataFrame":
+        from botocore.exceptions import NoCredentialsError
+
+        uri = proto_value.uri
+        columns = None
+        kwargs = get_pandas_storage_options(uri=uri, data_config=None)
+        path = os.path.join(uri, ".csv")
+        if current_task_metadata.structured_dataset_type and current_task_metadata.structured_dataset_type.columns:
+            columns = [c.name for c in current_task_metadata.structured_dataset_type.columns]
+        try:
+            return pd.read_csv(path, usecols=columns, storage_options=kwargs)
+        except NoCredentialsError:
+            logger.debug("S3 source detected, attempting anonymous S3 access")
+            kwargs = get_pandas_storage_options(uri=uri, data_config=None, anonymous=True)
+            return pd.read_csv(path, usecols=columns, storage_options=kwargs)
+
+
+class PandasToParquetEncodingHandler(StructuredDatasetEncoder):
+    def __init__(self):
+        super().__init__(pd.DataFrame, None, PARQUET)
+
+    async def encode(
+        self,
+        structured_dataset: StructuredDataset,
+        structured_dataset_type: types_pb2.StructuredDatasetType,
+    ) -> literals_pb2.StructuredDataset:
+        if not structured_dataset.uri:
+            from flyte._context import internal_ctx
+
+            ctx = internal_ctx()
+            uri = str(ctx.raw_data.get_random_remote_path())
+        else:
+            uri = typing.cast(str, structured_dataset.uri)
+
+        if not storage.is_remote(uri):
+            Path(uri).mkdir(parents=True, exist_ok=True)
+        path = os.path.join(uri, f"{0:05}")
+        df = typing.cast(pd.DataFrame, structured_dataset.dataframe)
+        df.to_parquet(
+            path,
+            coerce_timestamps="us",
+            allow_truncated_timestamps=False,
+            storage_options=get_pandas_storage_options(uri=path, data_config=None),
+        )
+        structured_dataset_type.format = PARQUET
+        return literals_pb2.StructuredDataset(
+            uri=uri, metadata=literals_pb2.StructuredDatasetMetadata(structured_dataset_type=structured_dataset_type)
+        )
+
+
+class ParquetToPandasDecodingHandler(StructuredDatasetDecoder):
+    def __init__(self):
+        super().__init__(pd.DataFrame, None, PARQUET)
+
+    async def decode(
+        self,
+        flyte_value: literals_pb2.StructuredDataset,
+        current_task_metadata: literals_pb2.StructuredDatasetMetadata,
+    ) -> "pd.DataFrame":
+        from botocore.exceptions import NoCredentialsError
+
+        uri = flyte_value.uri
+        columns = None
+        kwargs = get_pandas_storage_options(uri=uri, data_config=None)
+        if current_task_metadata.structured_dataset_type and current_task_metadata.structured_dataset_type.columns:
+            columns = [c.name for c in current_task_metadata.structured_dataset_type.columns]
+        try:
+            return pd.read_parquet(uri, columns=columns, storage_options=kwargs)
+        except NoCredentialsError:
+            logger.debug("S3 source detected, attempting anonymous S3 access")
+            kwargs = get_pandas_storage_options(uri=uri, data_config=None, anonymous=True)
+            return pd.read_parquet(uri, columns=columns, storage_options=kwargs)
+
+
+class ArrowToParquetEncodingHandler(StructuredDatasetEncoder):
+    def __init__(self):
+        super().__init__(pa.Table, None, PARQUET)
+
+    async def encode(
+        self,
+        structured_dataset: StructuredDataset,
+        structured_dataset_type: types_pb2.StructuredDatasetType,
+    ) -> literals_pb2.StructuredDataset:
+        import pyarrow.parquet as pq
+
+        if not structured_dataset.uri:
+            from flyte._context import internal_ctx
+
+            ctx = internal_ctx()
+            uri = ctx.raw_data.get_random_remote_path()
+        else:
+            uri = typing.cast(str, structured_dataset.uri)
+
+        if not storage.is_remote(uri):
+            Path(uri).mkdir(parents=True, exist_ok=True)
+        path = os.path.join(uri, f"{0:05}")
+        filesystem = storage.get_underlying_filesystem(path=path)
+        pq.write_table(structured_dataset.dataframe, strip_protocol(path), filesystem=filesystem)
+        return literals_pb2.StructuredDataset(
+            uri=uri, metadata=literals_pb2.StructuredDatasetMetadata(structured_dataset_type)
+        )
+
+
+class ParquetToArrowDecodingHandler(StructuredDatasetDecoder):
+    def __init__(self):
+        super().__init__(pa.Table, None, PARQUET)
+
+    async def decode(
+        self,
+        proto_value: literals_pb2.StructuredDataset,
+        current_task_metadata: literals_pb2.StructuredDatasetMetadata,
+    ) -> "pa.Table":
+        import pyarrow.parquet as pq
+        from botocore.exceptions import NoCredentialsError
+
+        uri = proto_value.uri
+        if not storage.is_remote(uri):
+            Path(uri).parent.mkdir(parents=True, exist_ok=True)
+        _, path = split_protocol(uri)
+
+        columns = None
+        if current_task_metadata.structured_dataset_type and current_task_metadata.structured_dataset_type.columns:
+            columns = [c.name for c in current_task_metadata.structured_dataset_type.columns]
+        try:
+            return pq.read_table(path, columns=columns)
+        except NoCredentialsError as e:
+            logger.debug("S3 source detected, attempting anonymous S3 access")
+            fs = storage.get_underlying_filesystem(path=uri, anonymous=True)
+            if fs is not None:
+                return pq.read_table(path, filesystem=fs, columns=columns)
+            raise e