Added AzureBlobFileSystem support for StructuredDatasets

Signed-off-by: Nick Müller <[email protected]>

Author: Nick Müller

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)

flytekit/types/structured/structured_dataset.py

@@ -39,6 +39,7 @@
 # Protocols
 BIGQUERY = "bq"
 S3 = "s3"
+ABFS = "abfs"
 GCS = "gs"
 LOCAL = "/"
 

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)
 

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
     )

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)