[WIP] Apply obstore as storage backend

Dockerfile.dev

@@ -40,6 +40,7 @@ RUN SETUPTOOLS_SCM_PRETEND_VERSION_FOR_FLYTEKIT=$PSEUDO_VERSION \
             -e /flytekit \
             -e /flytekit/plugins/flytekit-deck-standard \
             -e /flytekit/plugins/flytekit-flyteinteractive \
+            obstore==0.3.0b9 \
             markdown \
             pandas \
             pillow \

flytekit/core/obstore_filesystem.py

@@ -0,0 +1,36 @@
+"""
+Classes that overrides the AsyncFsspecStore that specify the filesystem specific parameters
+"""
+
+from obstore.fsspec import AsyncFsspecStore
+
+DEFAULT_BLOCK_SIZE = 5 * 2**20
+
+
+class ObstoreS3FileSystem(AsyncFsspecStore):
+    """
+    Add following property used in S3FileSystem
+    """
+
+    root_marker = ""
+    blocksize = DEFAULT_BLOCK_SIZE
+    protocol = ("s3", "s3a")
+    _extra_tokenize_attributes = ("default_block_size",)
+
+
+class ObstoreGCSFileSystem(AsyncFsspecStore):
+    """
+    Add following property used in GCSFileSystem
+    """
+
+    scopes = {"read_only", "read_write", "full_control"}
+    blocksize = DEFAULT_BLOCK_SIZE
+    protocol = "gcs", "gs"
+
+
+class ObstoreAzureBlobFileSystem(AsyncFsspecStore):
+    """
+    Add following property used in AzureBlobFileSystem
+    """
+
+    protocol = "abfs"

flytekit/types/structured/basic_dfs.py

@@ -9,7 +9,7 @@
 
 from flytekit import FlyteContext, lazy_module, logger
 from flytekit.configuration import DataConfig
-from flytekit.core.data_persistence import get_fsspec_storage_options
+from flytekit.core.data_persistence import get_fsspec_storage_options, split_path
 from flytekit.models import literals
 from flytekit.models.literals import StructuredDatasetMetadata
 from flytekit.models.types import StructuredDatasetType
@@ -37,7 +37,13 @@ def get_pandas_storage_options(
     from pandas.io.common import is_fsspec_url
 
     if is_fsspec_url(uri):
-        return get_fsspec_storage_options(protocol=get_protocol(uri), data_config=data_config, anonymous=anonymous)
+        bucket, _ = split_path(uri)
+        return get_fsspec_storage_options(
+            protocol=get_protocol(uri),
+            data_config=data_config,
+            anonymous=anonymous,
+            bucket=bucket,
+        )
 
     # Pandas does not allow storage_options for non-fsspec paths e.g. local.
     return None

plugins/flytekit-polars/flytekitplugins/polars/sd_transformers.py

@@ -2,7 +2,7 @@
 import typing
 
 from flytekit import FlyteContext, lazy_module
-from flytekit.core.data_persistence import get_fsspec_storage_options
+from flytekit.core.data_persistence import get_fsspec_storage_options, split_path
 from flytekit.models import literals
 from flytekit.models.literals import StructuredDatasetMetadata
 from flytekit.models.types import StructuredDatasetType
@@ -91,10 +91,11 @@ def decode(
         current_task_metadata: StructuredDatasetMetadata,
     ) -> pl.DataFrame:
         uri = flyte_value.uri
-
+        bucket, _ = split_path(uri)
         kwargs = get_fsspec_storage_options(
             protocol=fsspec_utils.get_protocol(uri),
             data_config=ctx.file_access.data_config,
+            bucket=bucket,
         )
         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]
@@ -153,10 +154,11 @@ def decode(
         current_task_metadata: StructuredDatasetMetadata,
     ) -> pl.LazyFrame:
         uri = flyte_value.uri
-
+        bucket, _ = split_path(uri)
         kwargs = get_fsspec_storage_options(
             protocol=fsspec_utils.get_protocol(uri),
             data_config=ctx.file_access.data_config,
+            bucket=bucket,
         )
         # use read_parquet instead of scan_parquet for now because scan_parquet currently doesn't work with fsspec:
         # https://github.com/pola-rs/polars/issues/16737

pyproject.toml

@@ -39,6 +39,7 @@ dependencies = [
     "marshmallow-jsonschema>=0.12.0",
     "mashumaro>=3.15",
     "msgpack>=1.1.0",
+    "obstore==0.3.0b10",
     "protobuf!=4.25.0",
     "pygments",
     "python-json-logger>=2.0.0",