From cd5b86f456a666bc8a91ac91bb9887c860897b9f Mon Sep 17 00:00:00 2001 From: Lisa Date: Mon, 22 Nov 2021 18:31:29 +0800 Subject: [PATCH 01/15] fix some error of type in flytekit/types Signed-off-by: Lisa --- flytekit/types/directory/__init__.py | 3 ++- flytekit/types/file/__init__.py | 30 ++++++++++++++++++--------- flytekit/types/file/file.py | 7 ++++--- flytekit/types/pickle/pickle.py | 4 ++-- flytekit/types/schema/types.py | 25 +++++++++++++--------- flytekit/types/schema/types_pandas.py | 8 +++---- 6 files changed, 47 insertions(+), 30 deletions(-) diff --git a/flytekit/types/directory/__init__.py b/flytekit/types/directory/__init__.py index 4edb5f9205..97d8ab57ce 100644 --- a/flytekit/types/directory/__init__.py +++ b/flytekit/types/directory/__init__.py @@ -18,7 +18,8 @@ # The following section provides some predefined aliases for commonly used FlyteDirectory formats. -TensorboardLogs = FlyteDirectory[typing.TypeVar("tensorboard")] +tensorboard = typing.TypeVar("tensorboard") +TensorboardLogs = FlyteDirectory[tensorboard] """ This type can be used to denote that the output is a folder that contains logs that can be loaded in tensorboard. this is usually the SummaryWriter output in pytorch or Keras callbacks which record the history readable by diff --git a/flytekit/types/file/__init__.py b/flytekit/types/file/__init__.py index 2b65efbcd6..81796fc49e 100644 --- a/flytekit/types/file/__init__.py +++ b/flytekit/types/file/__init__.py @@ -28,59 +28,69 @@ # This makes their usage extremely simple for the users. Please keep the list sorted. -HDF5EncodedFile = FlyteFile[typing.TypeVar("hdf5")] +hdf5 = typing.TypeVar("hdf5") +HDF5EncodedFile = FlyteFile[hdf5] """ This can be used to denote that the returned file is of type hdf5 and can be received by other tasks that accept an hdf5 format. This is usually useful for serializing Tensorflow models """ -HTMLPage = FlyteFile[typing.TypeVar("html")] +html = typing.TypeVar("html") +HTMLPage = FlyteFile[html] """ Can be used to receive or return an PNGImage. The underlying type is a FlyteFile, type. This is just a decoration and useful for attaching content type information with the file and automatically documenting code. """ -JoblibSerializedFile = FlyteFile[typing.TypeVar("joblib")] +joblib = typing.TypeVar("joblib") +JoblibSerializedFile = FlyteFile[joblib] """ This File represents a file that was serialized using `joblib.dump` method can be loaded back using `joblib.load` """ -JPEGImageFile = FlyteFile[typing.TypeVar("jpeg")] +jpeg = typing.TypeVar("jpeg") +JPEGImageFile = FlyteFile[jpeg] """ Can be used to receive or return an JPEGImage. The underlying type is a FlyteFile, type. This is just a decoration and useful for attaching content type information with the file and automatically documenting code. """ -PDFFile = FlyteFile[typing.TypeVar("pdf")] +pdf = typing.TypeVar("pdf") +PDFFile = FlyteFile[pdf] """ Can be used to receive or return an PDFFile. The underlying type is a FlyteFile, type. This is just a decoration and useful for attaching content type information with the file and automatically documenting code. """ -PNGImageFile = FlyteFile[typing.TypeVar("png")] +png = typing.TypeVar("png") +PNGImageFile = FlyteFile[png] """ Can be used to receive or return an PNGImage. The underlying type is a FlyteFile, type. This is just a decoration and useful for attaching content type information with the file and automatically documenting code. """ -PythonPickledFile = FlyteFile[typing.TypeVar("python-pickle")] +python_pickle = typing.TypeVar("python_pickle") +PythonPickledFile = FlyteFile[python_pickle] """ This type can be used when a serialized python pickled object is returned and shared between tasks. This only adds metadata to the file in Flyte, but does not really carry any object information """ -PythonNotebook = FlyteFile[typing.TypeVar("ipynb")] +ipynb = typing.TypeVar("ipynb") +PythonNotebook = FlyteFile[ipynb] """ This type is used to identify a python notebook file """ -SVGImageFile = FlyteFile[typing.TypeVar("svg")] +svg = typing.TypeVar("svg") +SVGImageFile = FlyteFile[svg] """ Can be used to receive or return an SVGImage. The underlying type is a FlyteFile, type. This is just a decoration and useful for attaching content type information with the file and automatically documenting code. """ -CSVFile = FlyteFile[typing.TypeVar("csv")] +csv = typing.TypeVar("csv") +CSVFile = FlyteFile[csv] """ Can be used to receive or return a CSVFile. The underlying type is a FlyteFile, type. This is just a decoration and useful for attaching content type information with the file and automatically documenting code. diff --git a/flytekit/types/file/file.py b/flytekit/types/file/file.py index 1fbcee049a..9b02692552 100644 --- a/flytekit/types/file/file.py +++ b/flytekit/types/file/file.py @@ -3,6 +3,7 @@ import os import pathlib import typing +from pydoc import locate from flytekit.core.context_manager import FlyteContext from flytekit.core.type_engine import TypeEngine, TypeTransformer @@ -220,7 +221,7 @@ def __init__(self): super().__init__(name="FlyteFilePath", t=FlyteFile) @staticmethod - def get_format(t: typing.Union[typing.Type[FlyteFile]]) -> str: + def get_format(t: typing.Union[typing.Type[FlyteFile], os.PathLike]) -> str: if t is os.PathLike: return "" return t.extension() @@ -342,14 +343,14 @@ def _downloader(): return ctx.file_access.get_data(uri, local_path, is_multipart=False) expected_format = FlyteFilePathTransformer.get_format(expected_python_type) - ff = FlyteFile[expected_format](local_path, _downloader) + ff = FlyteFile[locate(expected_format)](local_path, _downloader) ff._remote_source = uri return ff def guess_python_type(self, literal_type: LiteralType) -> typing.Type[FlyteFile[typing.Any]]: if literal_type.blob is not None and literal_type.blob.dimensionality == BlobType.BlobDimensionality.SINGLE: - return FlyteFile[typing.TypeVar(literal_type.blob.format)] + return FlyteFile[literal_type.blob.format] raise ValueError(f"Transformer {self} cannot reverse {literal_type}") diff --git a/flytekit/types/pickle/pickle.py b/flytekit/types/pickle/pickle.py index 8251d111bb..09a48395b9 100644 --- a/flytekit/types/pickle/pickle.py +++ b/flytekit/types/pickle/pickle.py @@ -20,8 +20,8 @@ class FlytePickle(typing.Generic[T]): """ @classmethod - def python_type(cls) -> None: - return None + def python_type(cls) -> typing.Type: + return type(None) def __class_getitem__(cls, python_type: typing.Type) -> typing.Type[T]: if python_type is None: diff --git a/flytekit/types/schema/types.py b/flytekit/types/schema/types.py index 89b6da3a90..a149329989 100644 --- a/flytekit/types/schema/types.py +++ b/flytekit/types/schema/types.py @@ -7,16 +7,20 @@ from dataclasses import dataclass from enum import Enum from typing import Type +from pathlib import Path import numpy as _np from flytekit.core.context_manager import FlyteContext, FlyteContextManager -from flytekit.core.type_engine import T, TypeEngine, TypeTransformer +from flytekit.core.type_engine import TypeEngine, TypeTransformer from flytekit.models.literals import Literal, Scalar, Schema from flytekit.models.types import LiteralType, SchemaType from flytekit.plugins import pandas +T = typing.TypeVar("T") + + class SchemaFormat(Enum): """ Represents the the schema storage format (at rest). @@ -35,7 +39,7 @@ class SchemaOpenMode(Enum): WRITE = "w" -def generate_ordered_files(directory: os.PathLike, n: int) -> str: +def generate_ordered_files(directory: os.PathLike, n: int) -> typing.Generator[str, None, None]: for i in range(n): yield os.path.join(directory, f"{i:05}") @@ -71,12 +75,12 @@ def all(self, **kwargs) -> T: class SchemaWriter(typing.Generic[T]): - def __init__(self, to_path: str, cols: typing.Dict[str, type], fmt: SchemaFormat): + def __init__(self, to_path: str, cols: typing.Optional[typing.Dict[str, type]], fmt: SchemaFormat): self._to_path = to_path self._fmt = fmt self._columns = cols # TODO This should be change to send a stop instead of hardcoded to 1024 - self._file_name_gen = generate_ordered_files(self._to_path, 1024) + self._file_name_gen = generate_ordered_files(Path(self._to_path), 1024) @property def to_path(self) -> str: @@ -105,7 +109,7 @@ def iter(self, **kwargs) -> typing.Generator[T, None, None]: with os.scandir(self._from_path) as it: for entry in it: if not entry.name.startswith(".") and entry.is_file(): - yield self._read(entry.path, **kwargs) + yield self._read(Path(entry.path), **kwargs) def all(self, **kwargs) -> T: files = [] @@ -114,7 +118,7 @@ def all(self, **kwargs) -> T: if not entry.name.startswith(".") and entry.is_file(): files.append(entry.path) - return self._read(*files, **kwargs) + return self._read(Path(*files), **kwargs) class LocalIOSchemaWriter(SchemaWriter[T]): @@ -219,10 +223,10 @@ def format(cls) -> SchemaFormat: def __init__( self, + downloader: typing.Callable[[str, os.PathLike], None], local_path: os.PathLike = None, remote_path: str = None, supported_mode: SchemaOpenMode = SchemaOpenMode.WRITE, - downloader: typing.Callable[[str, os.PathLike], None] = None, ): if supported_mode == SchemaOpenMode.READ and remote_path is None: @@ -282,8 +286,8 @@ def open( self._downloader(self.remote_path, self.local_path) self._downloaded = True if mode == SchemaOpenMode.WRITE: - return h.writer(self.local_path, self.columns(), self.format()) - return h.reader(self.local_path, self.columns(), self.format()) + return h.writer(typing.cast(str, self.local_path), self.columns(), self.format()) + return h.reader(typing.cast(str, self.local_path), self.columns(), self.format()) # Remote IO is handled. So we will just pass the remote reference to the object if mode == SchemaOpenMode.WRITE: @@ -298,6 +302,7 @@ def as_readonly(self) -> FlyteSchema: # Dummy path is ok, as we will assume data is already downloaded and will not download again remote_path=self.remote_path if self.remote_path else "", supported_mode=SchemaOpenMode.READ, + downloader=self._downloader, ) s._downloaded = True return s @@ -385,7 +390,7 @@ def downloader(x, y): def guess_python_type(self, literal_type: LiteralType) -> Type[T]: if not literal_type.schema: raise ValueError(f"Cannot reverse {literal_type}") - columns: dict[Type] = {} + columns: typing.Dict[str, Type] = {} for literal_column in literal_type.schema.columns: if literal_column.type == SchemaType.SchemaColumn.SchemaColumnType.INTEGER: columns[literal_column.name] = int diff --git a/flytekit/types/schema/types_pandas.py b/flytekit/types/schema/types_pandas.py index 41a5423c08..01aa8bf44a 100644 --- a/flytekit/types/schema/types_pandas.py +++ b/flytekit/types/schema/types_pandas.py @@ -15,8 +15,8 @@ class ParquetIO(object): PARQUET_ENGINE = "pyarrow" - def _read(self, chunk: os.PathLike, columns: typing.List[str], **kwargs) -> pandas.DataFrame: - return pandas.read_parquet(chunk, columns=columns, engine=self.PARQUET_ENGINE, **kwargs) + def _read(self, chunk: os.PathLike, _columns: typing.List[str], **kwargs) -> pandas.DataFrame: + return pandas.read_parquet(chunk, columns=_columns, engine=self.PARQUET_ENGINE, **kwargs) def read(self, *files: os.PathLike, columns: typing.List[str] = None, **kwargs) -> pandas.DataFrame: frames = [self._read(chunk=f, columns=columns, **kwargs) for f in files if os.path.getsize(f) > 0] @@ -59,13 +59,13 @@ def write( class FastParquetIO(ParquetIO): PARQUET_ENGINE = "fastparquet" - def _read(self, chunk: os.PathLike, columns: typing.List[str], **kwargs) -> pandas.DataFrame: + def _read(self, chunk: os.PathLike, _columns: typing.List[str], **kwargs) -> pandas.DataFrame: from fastparquet import ParquetFile as _ParquetFile from fastparquet import thrift_structures as _ts # TODO Follow up to figure out if this is not needed anymore # https://github.com/dask/fastparquet/issues/414#issuecomment-478983811 - df = pandas.read_parquet(chunk, columns=columns, engine=self.PARQUET_ENGINE, index=False) + df = pandas.read_parquet(chunk, columns=_columns, engine=self.PARQUET_ENGINE, index=False) df_column_types = df.dtypes pf = _ParquetFile(chunk) schema_column_dtypes = {l.name: l.type for l in list(pf.schema.schema_elements)} From 941724da37be9ea949dfb98ed939b096e7453291 Mon Sep 17 00:00:00 2001 From: Lisa Date: Sat, 27 Nov 2021 19:31:48 +0800 Subject: [PATCH 02/15] fix base_task fail Signed-off-by: Lisa --- flytekit/types/schema/types.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/flytekit/types/schema/types.py b/flytekit/types/schema/types.py index a149329989..391367840d 100644 --- a/flytekit/types/schema/types.py +++ b/flytekit/types/schema/types.py @@ -6,8 +6,8 @@ from abc import abstractmethod from dataclasses import dataclass from enum import Enum -from typing import Type from pathlib import Path +from typing import Type import numpy as _np @@ -17,7 +17,6 @@ from flytekit.models.types import LiteralType, SchemaType from flytekit.plugins import pandas - T = typing.TypeVar("T") From a3ebbe743e9286660fb2e314357f542816a355a9 Mon Sep 17 00:00:00 2001 From: Lisa Date: Sat, 27 Nov 2021 19:32:09 +0800 Subject: [PATCH 03/15] fix base_task fail Signed-off-by: Lisa --- flytekit/types/schema/types_pandas.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/flytekit/types/schema/types_pandas.py b/flytekit/types/schema/types_pandas.py index 01aa8bf44a..7ffd37f34b 100644 --- a/flytekit/types/schema/types_pandas.py +++ b/flytekit/types/schema/types_pandas.py @@ -19,7 +19,7 @@ def _read(self, chunk: os.PathLike, _columns: typing.List[str], **kwargs) -> pan return pandas.read_parquet(chunk, columns=_columns, engine=self.PARQUET_ENGINE, **kwargs) def read(self, *files: os.PathLike, columns: typing.List[str] = None, **kwargs) -> pandas.DataFrame: - frames = [self._read(chunk=f, columns=columns, **kwargs) for f in files if os.path.getsize(f) > 0] + frames = [self._read(chunk=f, _columns=columns, **kwargs) for f in files if os.path.getsize(f) > 0] if len(frames) == 1: return frames[0] elif len(frames) > 1: From d25f54ec8bb48f5a3d50d7dcc0aef0f581a4ae4c Mon Sep 17 00:00:00 2001 From: Lisa Date: Sat, 27 Nov 2021 20:42:31 +0800 Subject: [PATCH 04/15] fix pytest fail Signed-off-by: Lisa --- flytekit/types/schema/types.py | 1 + tests/flytekit/unit/core/test_local_cache.py | 2 +- tests/flytekit/unit/core/test_type_hints.py | 4 ++-- 3 files changed, 4 insertions(+), 3 deletions(-) diff --git a/flytekit/types/schema/types.py b/flytekit/types/schema/types.py index 391367840d..8ec3713b39 100644 --- a/flytekit/types/schema/types.py +++ b/flytekit/types/schema/types.py @@ -362,6 +362,7 @@ def to_literal( return Literal(scalar=Scalar(schema=Schema(remote_path, self._get_schema_type(python_type)))) schema = python_type( + downloader=None, local_path=ctx.file_access.get_random_local_directory(), remote_path=ctx.file_access.get_random_remote_directory(), ) diff --git a/tests/flytekit/unit/core/test_local_cache.py b/tests/flytekit/unit/core/test_local_cache.py index c43421bbf9..fd97e028e4 100644 --- a/tests/flytekit/unit/core/test_local_cache.py +++ b/tests/flytekit/unit/core/test_local_cache.py @@ -196,7 +196,7 @@ def t1() -> schema1: global n_cached_task_calls n_cached_task_calls += 1 - s = schema1() + s = schema1(downloader=None) s.open().write(pandas.DataFrame(data={"x": [1, 2], "y": ["3", "4"]})) return s diff --git a/tests/flytekit/unit/core/test_type_hints.py b/tests/flytekit/unit/core/test_type_hints.py index e9f54c829b..db00062921 100644 --- a/tests/flytekit/unit/core/test_type_hints.py +++ b/tests/flytekit/unit/core/test_type_hints.py @@ -844,7 +844,7 @@ def test_wf_typed_schema(): @task def t1() -> schema1: - s = schema1() + s = schema1(None) s.open().write(pandas.DataFrame(data={"x": [1, 2], "y": ["3", "4"]})) return s @@ -881,7 +881,7 @@ def test_wf_schema_to_df(): @task def t1() -> schema1: - s = schema1() + s = schema1(None) s.open().write(pandas.DataFrame(data={"x": [1, 2], "y": ["3", "4"]})) return s From ef711e1a9c8090efd303c5d82058f72f274ae168 Mon Sep 17 00:00:00 2001 From: aeioulisa Date: Mon, 29 Nov 2021 10:30:30 +0800 Subject: [PATCH 05/15] reset downloader Signed-off-by: Lisa --- flytekit/types/schema/types.py | 4 +--- tests/flytekit/unit/core/test_local_cache.py | 2 +- tests/flytekit/unit/core/test_type_hints.py | 4 ++-- 3 files changed, 4 insertions(+), 6 deletions(-) diff --git a/flytekit/types/schema/types.py b/flytekit/types/schema/types.py index 8ec3713b39..813d0cfe81 100644 --- a/flytekit/types/schema/types.py +++ b/flytekit/types/schema/types.py @@ -222,10 +222,10 @@ def format(cls) -> SchemaFormat: def __init__( self, - downloader: typing.Callable[[str, os.PathLike], None], local_path: os.PathLike = None, remote_path: str = None, supported_mode: SchemaOpenMode = SchemaOpenMode.WRITE, + downloader: typing.Callable[[str, os.PathLike], None] = None, ): if supported_mode == SchemaOpenMode.READ and remote_path is None: @@ -301,7 +301,6 @@ def as_readonly(self) -> FlyteSchema: # Dummy path is ok, as we will assume data is already downloaded and will not download again remote_path=self.remote_path if self.remote_path else "", supported_mode=SchemaOpenMode.READ, - downloader=self._downloader, ) s._downloaded = True return s @@ -362,7 +361,6 @@ def to_literal( return Literal(scalar=Scalar(schema=Schema(remote_path, self._get_schema_type(python_type)))) schema = python_type( - downloader=None, local_path=ctx.file_access.get_random_local_directory(), remote_path=ctx.file_access.get_random_remote_directory(), ) diff --git a/tests/flytekit/unit/core/test_local_cache.py b/tests/flytekit/unit/core/test_local_cache.py index fd97e028e4..c43421bbf9 100644 --- a/tests/flytekit/unit/core/test_local_cache.py +++ b/tests/flytekit/unit/core/test_local_cache.py @@ -196,7 +196,7 @@ def t1() -> schema1: global n_cached_task_calls n_cached_task_calls += 1 - s = schema1(downloader=None) + s = schema1() s.open().write(pandas.DataFrame(data={"x": [1, 2], "y": ["3", "4"]})) return s diff --git a/tests/flytekit/unit/core/test_type_hints.py b/tests/flytekit/unit/core/test_type_hints.py index db00062921..e9f54c829b 100644 --- a/tests/flytekit/unit/core/test_type_hints.py +++ b/tests/flytekit/unit/core/test_type_hints.py @@ -844,7 +844,7 @@ def test_wf_typed_schema(): @task def t1() -> schema1: - s = schema1(None) + s = schema1() s.open().write(pandas.DataFrame(data={"x": [1, 2], "y": ["3", "4"]})) return s @@ -881,7 +881,7 @@ def test_wf_schema_to_df(): @task def t1() -> schema1: - s = schema1(None) + s = schema1() s.open().write(pandas.DataFrame(data={"x": [1, 2], "y": ["3", "4"]})) return s From 1248ce0e1e92e34c15055ac836d214384b3fb901 Mon Sep 17 00:00:00 2001 From: Lisa Date: Tue, 30 Nov 2021 16:59:11 +0800 Subject: [PATCH 06/15] mypy error Signed-off-by: Lisa --- flytekit/types/directory/types.py | 2 +- flytekit/types/pickle/pickle.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/flytekit/types/directory/types.py b/flytekit/types/directory/types.py index a0112c1eea..37f958c8f3 100644 --- a/flytekit/types/directory/types.py +++ b/flytekit/types/directory/types.py @@ -300,7 +300,7 @@ def guess_python_type(self, literal_type: LiteralType) -> typing.Type[FlyteDirec literal_type.blob is not None and literal_type.blob.dimensionality == _core_types.BlobType.BlobDimensionality.MULTIPART ): - return FlyteDirectory[typing.TypeVar(literal_type.blob.format)] + return FlyteDirectory[literal_type.blob.format] raise ValueError(f"Transformer {self} cannot reverse {literal_type}") diff --git a/flytekit/types/pickle/pickle.py b/flytekit/types/pickle/pickle.py index 09a48395b9..1d175c343a 100644 --- a/flytekit/types/pickle/pickle.py +++ b/flytekit/types/pickle/pickle.py @@ -35,7 +35,7 @@ class _SpecificFormatClass(FlytePickle): def python_type(cls) -> typing.Type: return python_type - return _SpecificFormatClass + return _SpecificFormatClass.python_type() class FlytePickleTransformer(TypeTransformer[FlytePickle]): From 7f7edc3c13ccbdde55682eec308770b7f4f87996 Mon Sep 17 00:00:00 2001 From: Lisa Date: Tue, 30 Nov 2021 17:27:06 +0800 Subject: [PATCH 07/15] error of no extension attribute Signed-off-by: Lisa --- flytekit/types/file/file.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/flytekit/types/file/file.py b/flytekit/types/file/file.py index 9b02692552..149b38326e 100644 --- a/flytekit/types/file/file.py +++ b/flytekit/types/file/file.py @@ -224,7 +224,7 @@ def __init__(self): def get_format(t: typing.Union[typing.Type[FlyteFile], os.PathLike]) -> str: if t is os.PathLike: return "" - return t.extension() + return typing.cast(FlyteFile, t).extension() def _blob_type(self, format: str) -> BlobType: return BlobType(format=format, dimensionality=BlobType.BlobDimensionality.SINGLE) From cc808d0964ec21041e93087ab4c05fe8a40dfe1c Mon Sep 17 00:00:00 2001 From: Lisa Date: Tue, 30 Nov 2021 17:30:29 +0800 Subject: [PATCH 08/15] error of downloader is not callabe Signed-off-by: Lisa --- flytekit/types/schema/types.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/flytekit/types/schema/types.py b/flytekit/types/schema/types.py index 813d0cfe81..a41b63e7b0 100644 --- a/flytekit/types/schema/types.py +++ b/flytekit/types/schema/types.py @@ -280,6 +280,8 @@ def open( if not h.handles_remote_io: # The Schema Handler does not manage its own IO, and this it will expect the files are on local file-system if self._supported_mode == SchemaOpenMode.READ and not self._downloaded: + if self._downloader is None: + raise AssertionError("downloader cannot be None in read mode!") # Only for readable objects if they are not downloaded already, we should download them # Write objects should already have everything written to self._downloader(self.remote_path, self.local_path) From 3a3f33681dc786227970717d46c193ddc3ea59c0 Mon Sep 17 00:00:00 2001 From: Lisa Date: Tue, 30 Nov 2021 17:31:37 +0800 Subject: [PATCH 09/15] revert _columns Signed-off-by: Lisa --- flytekit/types/schema/types_pandas.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/flytekit/types/schema/types_pandas.py b/flytekit/types/schema/types_pandas.py index 7ffd37f34b..41a5423c08 100644 --- a/flytekit/types/schema/types_pandas.py +++ b/flytekit/types/schema/types_pandas.py @@ -15,11 +15,11 @@ class ParquetIO(object): PARQUET_ENGINE = "pyarrow" - def _read(self, chunk: os.PathLike, _columns: typing.List[str], **kwargs) -> pandas.DataFrame: - return pandas.read_parquet(chunk, columns=_columns, engine=self.PARQUET_ENGINE, **kwargs) + def _read(self, chunk: os.PathLike, columns: typing.List[str], **kwargs) -> pandas.DataFrame: + return pandas.read_parquet(chunk, columns=columns, engine=self.PARQUET_ENGINE, **kwargs) def read(self, *files: os.PathLike, columns: typing.List[str] = None, **kwargs) -> pandas.DataFrame: - frames = [self._read(chunk=f, _columns=columns, **kwargs) for f in files if os.path.getsize(f) > 0] + frames = [self._read(chunk=f, columns=columns, **kwargs) for f in files if os.path.getsize(f) > 0] if len(frames) == 1: return frames[0] elif len(frames) > 1: @@ -59,13 +59,13 @@ def write( class FastParquetIO(ParquetIO): PARQUET_ENGINE = "fastparquet" - def _read(self, chunk: os.PathLike, _columns: typing.List[str], **kwargs) -> pandas.DataFrame: + def _read(self, chunk: os.PathLike, columns: typing.List[str], **kwargs) -> pandas.DataFrame: from fastparquet import ParquetFile as _ParquetFile from fastparquet import thrift_structures as _ts # TODO Follow up to figure out if this is not needed anymore # https://github.com/dask/fastparquet/issues/414#issuecomment-478983811 - df = pandas.read_parquet(chunk, columns=_columns, engine=self.PARQUET_ENGINE, index=False) + df = pandas.read_parquet(chunk, columns=columns, engine=self.PARQUET_ENGINE, index=False) df_column_types = df.dtypes pf = _ParquetFile(chunk) schema_column_dtypes = {l.name: l.type for l in list(pf.schema.schema_elements)} From 8807afce109df032ef72430e3623784c493f4e7d Mon Sep 17 00:00:00 2001 From: Lisa Date: Tue, 30 Nov 2021 17:35:05 +0800 Subject: [PATCH 10/15] error of incompatible type Signed-off-by: Lisa --- flytekit/types/schema/types_pandas.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/flytekit/types/schema/types_pandas.py b/flytekit/types/schema/types_pandas.py index 41a5423c08..0edf024b08 100644 --- a/flytekit/types/schema/types_pandas.py +++ b/flytekit/types/schema/types_pandas.py @@ -15,7 +15,7 @@ class ParquetIO(object): PARQUET_ENGINE = "pyarrow" - def _read(self, chunk: os.PathLike, columns: typing.List[str], **kwargs) -> pandas.DataFrame: + def _read(self, chunk: os.PathLike, columns: typing.Optional[typing.List[str]], **kwargs) -> pandas.DataFrame: return pandas.read_parquet(chunk, columns=columns, engine=self.PARQUET_ENGINE, **kwargs) def read(self, *files: os.PathLike, columns: typing.List[str] = None, **kwargs) -> pandas.DataFrame: @@ -59,7 +59,7 @@ def write( class FastParquetIO(ParquetIO): PARQUET_ENGINE = "fastparquet" - def _read(self, chunk: os.PathLike, columns: typing.List[str], **kwargs) -> pandas.DataFrame: + def _read(self, chunk: os.PathLike, columns: typing.Optional[typing.List[str]], **kwargs) -> pandas.DataFrame: from fastparquet import ParquetFile as _ParquetFile from fastparquet import thrift_structures as _ts From ca3f8e2d058935f2543f0b3c849ad3d2a4bc488c Mon Sep 17 00:00:00 2001 From: Lisa Date: Tue, 30 Nov 2021 20:30:47 +0800 Subject: [PATCH 11/15] error of incompatible return value type Signed-off-by: Lisa --- flytekit/types/pickle/pickle.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/flytekit/types/pickle/pickle.py b/flytekit/types/pickle/pickle.py index 1d175c343a..9219d3a8b4 100644 --- a/flytekit/types/pickle/pickle.py +++ b/flytekit/types/pickle/pickle.py @@ -23,7 +23,7 @@ class FlytePickle(typing.Generic[T]): def python_type(cls) -> typing.Type: return type(None) - def __class_getitem__(cls, python_type: typing.Type) -> typing.Type[T]: + def __class_getitem__(cls, python_type: typing.Type) -> typing.Type: if python_type is None: return cls @@ -35,7 +35,7 @@ class _SpecificFormatClass(FlytePickle): def python_type(cls) -> typing.Type: return python_type - return _SpecificFormatClass.python_type() + return _SpecificFormatClass class FlytePickleTransformer(TypeTransformer[FlytePickle]): From 942d5fda2092f8bdc1e91802771126787c0831c5 Mon Sep 17 00:00:00 2001 From: Lisa Date: Wed, 1 Dec 2021 16:39:30 +0800 Subject: [PATCH 12/15] error of generic Signed-off-by: Lisa --- flytekit/types/directory/types.py | 6 +++--- flytekit/types/file/file.py | 6 +++--- flytekit/types/schema/types.py | 2 +- 3 files changed, 7 insertions(+), 7 deletions(-) diff --git a/flytekit/types/directory/types.py b/flytekit/types/directory/types.py index 37f958c8f3..81f4fb0fd0 100644 --- a/flytekit/types/directory/types.py +++ b/flytekit/types/directory/types.py @@ -133,7 +133,7 @@ def __fspath__(self): def extension(cls) -> str: return "" - def __class_getitem__(cls, item: typing.Type) -> typing.Type[FlyteDirectory]: + def __class_getitem__(cls, item: typing.Union[typing.Type, str]) -> typing.Type[FlyteDirectory]: if item is None: return cls item_string = str(item) @@ -290,7 +290,7 @@ def _downloader(): expected_format = self.get_format(expected_python_type) - fd = FlyteDirectory[expected_format](local_folder, _downloader) + fd = FlyteDirectory.__class_getitem__(expected_format)(local_folder, _downloader) fd._remote_source = uri return fd @@ -300,7 +300,7 @@ def guess_python_type(self, literal_type: LiteralType) -> typing.Type[FlyteDirec literal_type.blob is not None and literal_type.blob.dimensionality == _core_types.BlobType.BlobDimensionality.MULTIPART ): - return FlyteDirectory[literal_type.blob.format] + return FlyteDirectory.__class_getitem__(literal_type.blob.format) raise ValueError(f"Transformer {self} cannot reverse {literal_type}") diff --git a/flytekit/types/file/file.py b/flytekit/types/file/file.py index 149b38326e..2315182e09 100644 --- a/flytekit/types/file/file.py +++ b/flytekit/types/file/file.py @@ -138,7 +138,7 @@ def t2() -> flytekit_typing.FlyteFile["csv"]: def extension(cls) -> str: return "" - def __class_getitem__(cls, item: typing.Type) -> typing.Type[FlyteFile]: + def __class_getitem__(cls, item: typing.Union[str, typing.Type]) -> typing.Type[FlyteFile]: if item is None: return cls item_string = str(item) @@ -343,14 +343,14 @@ def _downloader(): return ctx.file_access.get_data(uri, local_path, is_multipart=False) expected_format = FlyteFilePathTransformer.get_format(expected_python_type) - ff = FlyteFile[locate(expected_format)](local_path, _downloader) + ff = FlyteFile.__class_getitem__(expected_format)(local_path, _downloader) ff._remote_source = uri return ff def guess_python_type(self, literal_type: LiteralType) -> typing.Type[FlyteFile[typing.Any]]: if literal_type.blob is not None and literal_type.blob.dimensionality == BlobType.BlobDimensionality.SINGLE: - return FlyteFile[literal_type.blob.format] + return FlyteFile.__class_getitem__(literal_type.blob.format) raise ValueError(f"Transformer {self} cannot reverse {literal_type}") diff --git a/flytekit/types/schema/types.py b/flytekit/types/schema/types.py index a41b63e7b0..b1b3ec83dd 100644 --- a/flytekit/types/schema/types.py +++ b/flytekit/types/schema/types.py @@ -406,7 +406,7 @@ def guess_python_type(self, literal_type: LiteralType) -> Type[T]: columns[literal_column.name] = bool else: raise ValueError(f"Unknown schema column type {literal_column}") - return FlyteSchema[columns] + return FlyteSchema.__class_getitem__(columns) TypeEngine.register(FlyteSchemaTransformer()) From dba5840eeb158576814035c15024455c693d69d9 Mon Sep 17 00:00:00 2001 From: Lisa Date: Wed, 1 Dec 2021 16:41:16 +0800 Subject: [PATCH 13/15] incompatible error Signed-off-by: Lisa --- flytekit/types/schema/types.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/flytekit/types/schema/types.py b/flytekit/types/schema/types.py index b1b3ec83dd..1b58148b33 100644 --- a/flytekit/types/schema/types.py +++ b/flytekit/types/schema/types.py @@ -387,7 +387,7 @@ def downloader(x, y): supported_mode=SchemaOpenMode.READ, ) - def guess_python_type(self, literal_type: LiteralType) -> Type[T]: + def guess_python_type(self, literal_type: LiteralType) -> Type[FlyteSchema]: if not literal_type.schema: raise ValueError(f"Cannot reverse {literal_type}") columns: typing.Dict[str, Type] = {} From 8963e08d451abf2bb402fa495fc4ca49c4f7bd0d Mon Sep 17 00:00:00 2001 From: Lisa Date: Wed, 1 Dec 2021 16:44:16 +0800 Subject: [PATCH 14/15] remove unuse import Signed-off-by: Lisa --- flytekit/types/file/file.py | 1 - 1 file changed, 1 deletion(-) diff --git a/flytekit/types/file/file.py b/flytekit/types/file/file.py index 2315182e09..eb64e5943d 100644 --- a/flytekit/types/file/file.py +++ b/flytekit/types/file/file.py @@ -3,7 +3,6 @@ import os import pathlib import typing -from pydoc import locate from flytekit.core.context_manager import FlyteContext from flytekit.core.type_engine import TypeEngine, TypeTransformer From 26bb8536674513bd33e8a76ac2eb5f5318b1d24b Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Fri, 3 Dec 2021 02:06:28 +0800 Subject: [PATCH 15/15] Fix tests Signed-off-by: Kevin Su --- flytekit/types/schema/types.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/flytekit/types/schema/types.py b/flytekit/types/schema/types.py index 1b58148b33..ff33d852b1 100644 --- a/flytekit/types/schema/types.py +++ b/flytekit/types/schema/types.py @@ -111,13 +111,13 @@ def iter(self, **kwargs) -> typing.Generator[T, None, None]: yield self._read(Path(entry.path), **kwargs) def all(self, **kwargs) -> T: - files = [] + files: typing.List[os.PathLike] = [] with os.scandir(self._from_path) as it: for entry in it: if not entry.name.startswith(".") and entry.is_file(): - files.append(entry.path) + files.append(Path(entry.path)) - return self._read(Path(*files), **kwargs) + return self._read(*files, **kwargs) class LocalIOSchemaWriter(SchemaWriter[T]):