Added AzureBlobFileSystem support for StructuredDatasets

flyteorg · Jul 21, 2022 · 81b3370 · 81b3370
1 parent c6e7237
commit 81b3370
Show file tree

Hide file tree

Showing 5 changed files with 14 additions and 4 deletions.
diff --git a/flytekit/types/structured/basic_dfs.py b/flytekit/types/structured/basic_dfs.py
@@ -12,6 +12,7 @@
 from flytekit.models.literals import StructuredDatasetMetadata
 from flytekit.models.types import StructuredDatasetType
 from flytekit.types.structured.structured_dataset import (
+    ABFS,
     GCS,
     LOCAL,
     PARQUET,
@@ -106,7 +107,7 @@ def decode(
 
 
 # Don't override default protocol
-for protocol in [LOCAL, S3, GCS]:
+for protocol in [LOCAL, S3, GCS, ABFS]:
     StructuredDatasetTransformerEngine.register(PandasToParquetEncodingHandler(protocol), default_for_type=False)
     StructuredDatasetTransformerEngine.register(ParquetToPandasDecodingHandler(protocol), default_for_type=False)
     StructuredDatasetTransformerEngine.register(ArrowToParquetEncodingHandler(protocol), default_for_type=False)

diff --git a/flytekit/types/structured/structured_dataset.py b/flytekit/types/structured/structured_dataset.py
@@ -39,6 +39,7 @@
 # Protocols
 BIGQUERY = "bq"
 S3 = "s3"
+ABFS = "abfs"
 GCS = "gs"
 LOCAL = "/"
 

diff --git a/plugins/flytekit-data-fsspec/flytekitplugins/fsspec/__init__.py b/plugins/flytekit-data-fsspec/flytekitplugins/fsspec/__init__.py
@@ -26,7 +26,7 @@
 
 from flytekit import StructuredDatasetTransformerEngine, logger
 from flytekit.configuration import internal
-from flytekit.types.structured.structured_dataset import GCS, S3
+from flytekit.types.structured.structured_dataset import ABFS, GCS, S3
 
 from .arrow import ArrowToParquetEncodingHandler, ParquetToArrowDecodingHandler
 from .pandas import PandasToParquetEncodingHandler, ParquetToPandasDecodingHandler
@@ -41,6 +41,9 @@ def _register(protocol: str):
     StructuredDatasetTransformerEngine.register(ParquetToArrowDecodingHandler(protocol), True, True)
 
 
+if importlib.util.find_spec("adlfs"):
+    _register(ABFS)
+
 if importlib.util.find_spec("s3fs"):
     _register(S3)
 

diff --git a/plugins/flytekit-polars/flytekitplugins/polars/sd_transformers.py b/plugins/flytekit-polars/flytekitplugins/polars/sd_transformers.py
@@ -7,6 +7,7 @@
 from flytekit.models.literals import StructuredDatasetMetadata
 from flytekit.models.types import StructuredDatasetType
 from flytekit.types.structured.structured_dataset import (
+    ABFS,
     GCS,
     LOCAL,
     PARQUET,
@@ -62,7 +63,7 @@ def decode(
         return pl.read_parquet(path)
 
 
-for protocol in [LOCAL, S3, GCS]:
+for protocol in [LOCAL, S3, GCS, ABFS]:
     StructuredDatasetTransformerEngine.register(
         PolarsDataFrameToParquetEncodingHandler(protocol), default_for_type=False
     )

diff --git a/plugins/flytekit-spark/flytekitplugins/spark/sd_transformers.py b/plugins/flytekit-spark/flytekitplugins/spark/sd_transformers.py
@@ -7,7 +7,11 @@
 from flytekit.models.literals import StructuredDatasetMetadata
 from flytekit.models.types import StructuredDatasetType
 from flytekit.types.structured.structured_dataset import (
+    ABFS,
+    GCS,
+    LOCAL,
     PARQUET,
+    S3,
     StructuredDataset,
     StructuredDatasetDecoder,
     StructuredDatasetEncoder,
@@ -48,6 +52,6 @@ def decode(
         return user_ctx.spark_session.read.parquet(flyte_value.uri)
 
 
-for protocol in ["/", "s3"]:
+for protocol in [LOCAL, S3, GCS, ABFS]:
     StructuredDatasetTransformerEngine.register(SparkToParquetEncodingHandler(protocol), default_for_type=False)
     StructuredDatasetTransformerEngine.register(ParquetToSparkDecodingHandler(protocol), default_for_type=False)