Cotransform dict DataFrames, improve transformer deco, add no effect …

…checkpoint (#30)

* setup fugue sql

* SQL: add basic extensions and tests

* update

* update

* clean up sql files

* fix syntax, add save load

* add test for load

* FugueSQLWorkflow

* update version

* Add pandas udf support, add SQL persist broadcast

* update

* update

* update

* update

* update

* make transformer schema more flexible

* Remove ABC from Transformers

* cotransform on dict like dataframes, add no effect checkpoint

* change deco lambda

* add more tests

README.md

@@ -8,3 +8,13 @@
 [![Doc](https://readthedocs.org/projects/fugue/badge)](https://triad.readthedocs.org)
 
 An abstraction layer for distributed computation
+
+## Tutorials
+
+### The simplest way to start:
+[![Binder](https://mybinder.org/badge_logo.svg)](https://mybinder.org/v2/gh/fugue-project/tutorials/master)
+
+### Run the tutorial using docker
+```
+docker run -p 8888:8888 fugueproject/tutorials:latest
+```

fugue/__init__.py

@@ -1,6 +1,9 @@
 # flake8: noqa
 
-__version__ = "0.2.8"
+__version__ = "0.2.9"
+
+from triad.collections import Schema
+from triad.collections.fs import FileSystem
 
 from fugue.collections.partition import PartitionCursor, PartitionSpec
 from fugue.dataframe.array_dataframe import ArrayDataFrame
@@ -15,6 +18,11 @@
 from fugue.extensions.creator import Creator, creator
 from fugue.extensions.outputter import Outputter, outputter
 from fugue.extensions.processor import Processor, processor
-from fugue.extensions.transformer import CoTransformer, Transformer, transformer
+from fugue.extensions.transformer import (
+    CoTransformer,
+    Transformer,
+    transformer,
+    cotransformer,
+)
 from fugue.workflow.workflow import FugueWorkflow, WorkflowDataFrame
 from fugue.workflow.workflow_context import FugueWorkflowContext

fugue/dataframe/dataframe.py

@@ -169,7 +169,11 @@ def show(
                 print("")
             if len(self.metadata) > 0:
                 print("Metadata:")
-                print(self.metadata.to_json(indent=True))
+                try:
+                    # try pretty print, but if not convertible to json, print original
+                    print(self.metadata.to_json(indent=True))
+                except Exception:  # pragma: no cover
+                    print(self.metadata)
                 print("")
 
     def head(self, n: int, columns: Optional[List[str]] = None) -> List[Any]:

fugue/execution/execution_engine.py

@@ -17,6 +17,7 @@
 from triad.exceptions import InvalidOperationError
 from triad.utils.assertion import assert_or_throw
 from triad.utils.convert import to_size
+from triad.utils.string import validate_triad_var_name
 
 _DEFAULT_JOIN_KEYS: List[str] = []
 
@@ -116,8 +117,10 @@ def serialize_by_partition(
         partition_spec: PartitionSpec,
         df_name: str,
         temp_path: Optional[str] = None,
-        to_file_threshold: int = -1,
+        to_file_threshold: Any = -1,
+        has_name: bool = False,
     ) -> DataFrame:
+        to_file_threshold = _get_file_threshold(to_file_threshold)
         on = list(filter(lambda k: k in df.schema, partition_spec.partition_by))
         presort = list(
             filter(lambda p: p[0] in df.schema, partition_spec.presort.items())
@@ -135,7 +138,8 @@ def serialize_by_partition(
         metadata = dict(
             serialized=True,
             serialized_cols={df_name: col_name},
-            schemas={df_name: df.schema},
+            schemas={df_name: str(df.schema)},
+            serialized_has_name=has_name,
         )
         return self.map(df, s.run, output_schema, partition_spec, metadata)
 
@@ -147,16 +151,15 @@ def zip(
         partition_spec: PartitionSpec = EMPTY_PARTITION_SPEC,
         temp_path: Optional[str] = None,
         to_file_threshold: Any = -1,
+        df1_name: Optional[str] = None,
+        df2_name: Optional[str] = None,
     ):
         on = list(partition_spec.partition_by)
         how = how.lower()
         assert_or_throw(
             "semi" not in how and "anti" not in how,
             InvalidOperationError("zip does not support semi or anti joins"),
         )
-        to_file_threshold = (
-            -1 if to_file_threshold == -1 else to_size(to_file_threshold)
-        )
         serialized_cols: Dict[str, Any] = {}
         schemas: Dict[str, Any] = {}
         if len(on) == 0:
@@ -171,14 +174,17 @@ def zip(
             )
         partition_spec = PartitionSpec(partition_spec, by=on)
 
-        def update_df(df: DataFrame) -> DataFrame:
+        def update_df(df: DataFrame, name: Optional[str]) -> DataFrame:
+            if name is None:
+                name = f"_{len(serialized_cols)}"
             if not df.metadata.get("serialized", False):
                 df = self.serialize_by_partition(
                     df,
                     partition_spec,
-                    f"_{len(serialized_cols)}",
+                    name,
                     temp_path,
                     to_file_threshold,
+                    has_name=name is not None,
                 )
             for k in df.metadata["serialized_cols"].keys():
                 assert_or_throw(
@@ -188,13 +194,58 @@ def update_df(df: DataFrame) -> DataFrame:
                 schemas[k] = df.metadata["schemas"][k]
             return df
 
-        df1 = update_df(df1)
-        df2 = update_df(df2)
+        df1 = update_df(df1, df1_name)
+        df2 = update_df(df2, df2_name)
         metadata = dict(
-            serialized=True, serialized_cols=serialized_cols, schemas=schemas
+            serialized=True,
+            serialized_cols=serialized_cols,
+            schemas=schemas,
+            serialized_has_name=df1_name is not None or df2_name is not None,
         )
         return self.join(df1, df2, how=how, on=on, metadata=metadata)
 
+    def zip_all(
+        self,
+        dfs: DataFrames,
+        how: str = "inner",
+        partition_spec: PartitionSpec = EMPTY_PARTITION_SPEC,
+        temp_path: Optional[str] = None,
+        to_file_threshold: Any = -1,
+    ) -> DataFrame:
+        assert_or_throw(len(dfs) > 0, "can't zip 0 dataframes")
+        pairs = list(dfs.items())
+        has_name = dfs.has_key
+        if len(dfs) == 1:
+            return self.serialize_by_partition(
+                pairs[0][1],
+                partition_spec,
+                pairs[0][0],
+                temp_path,
+                to_file_threshold,
+                has_name=has_name,
+            )
+        df = self.zip(
+            pairs[0][1],
+            pairs[1][1],
+            how=how,
+            partition_spec=partition_spec,
+            temp_path=temp_path,
+            to_file_threshold=to_file_threshold,
+            df1_name=pairs[0][0] if has_name else None,
+            df2_name=pairs[1][0] if has_name else None,
+        )
+        for i in range(2, len(dfs)):
+            df = self.zip(
+                df,
+                pairs[i][1],
+                how=how,
+                partition_spec=partition_spec,
+                temp_path=temp_path,
+                to_file_threshold=to_file_threshold,
+                df2_name=pairs[i][0] if has_name else None,
+            )
+        return df
+
     def comap(
         self,
         df: DataFrame,
@@ -243,9 +294,19 @@ def __deepcopy__(self, memo: Any) -> "ExecutionEngine":
         return self
 
 
+def _get_file_threshold(size: Any) -> int:
+    if size is None:
+        return -1
+    if isinstance(size, int):
+        return size
+    return to_size(size)
+
+
 def _df_name_to_serialize_col(name: str):
     assert_or_throw(name is not None, "Dataframe name can't be None")
-    return "__blob__" + name + "__"
+    name = "__blob__" + name + "__"
+    assert_or_throw(validate_triad_var_name(name), "Invalid name " + name)
+    return name
 
 
 class _PartitionSerializer(object):
@@ -271,19 +332,18 @@ def __init__(
     ):
         self.schemas = df.metadata["schemas"]
         self.df_idx = [
-            (df.schema.index_of_key(v), self.schemas[k])
+            (df.schema.index_of_key(v), k, self.schemas[k])
             for k, v in df.metadata["serialized_cols"].items()
         ]
+        self.named = df.metadata.get("serialized_has_name", False)
         self.func = func
         self._on_init = on_init
 
     def on_init(self, partition_no, df: DataFrame) -> None:
         if self._on_init is None:
             return
         # TODO: currently, get_output_schema only gets empty dataframes
-        empty_dfs = DataFrames(
-            {k: ArrayDataFrame([], v) for k, v in self.schemas.items()}
-        )
+        empty_dfs = _generate_comap_empty_dfs(self.schemas, self.named)
         self._on_init(partition_no, empty_dfs)
 
     def run(self, cursor: PartitionCursor, df: LocalDataFrame) -> LocalDataFrame:
@@ -295,11 +355,21 @@ def run(self, cursor: PartitionCursor, df: LocalDataFrame) -> LocalDataFrame:
         dfs = DataFrames(list(self._get_dfs(data[0])))
         return self.func(cursor, dfs)
 
-    def _get_dfs(self, row: Any) -> Iterable[DataFrame]:
-        for k, v in self.df_idx:
+    def _get_dfs(self, row: Any) -> Iterable[Any]:
+        for k, name, v in self.df_idx:
             if row[k] is None:
-                yield ArrayDataFrame([], v)
+                df: DataFrame = ArrayDataFrame([], v)
             else:
-                df = deserialize_df(row[k])
+                df = deserialize_df(row[k])  # type: ignore
                 assert df is not None
+            if self.named:
+                yield name, df
+            else:
                 yield df
+
+
+def _generate_comap_empty_dfs(schemas: Any, named: bool) -> DataFrames:
+    if named:
+        return DataFrames({k: ArrayDataFrame([], v) for k, v in schemas.items()})
+    else:
+        return DataFrames([ArrayDataFrame([], v) for v in schemas.values()])

fugue/extensions/__init__.py

@@ -1,5 +1,10 @@
 # flake8: noqa
-from fugue.extensions.transformer import Transformer, CoTransformer, transformer
+from fugue.extensions.transformer import (
+    Transformer,
+    CoTransformer,
+    transformer,
+    cotransformer,
+)
 from fugue.extensions.creator import Creator, creator
 from fugue.extensions.processor import Processor, processor
 from fugue.extensions.outputter import Outputter, outputter

fugue/extensions/builtins/outputters.py

@@ -10,7 +10,7 @@ def process(self, dfs: DataFrames) -> None:
             df.show(
                 self.params.get("rows", 10),
                 self.params.get("show_count", False),
-                title=self.params.get("title", ""),
+                title=self.params.get_or_none("title", str),
             )
 
 

fugue/extensions/builtins/processors.py

@@ -10,6 +10,7 @@
 )
 from fugue.exceptions import FugueWorkflowError
 from fugue.execution import SQLEngine
+from fugue.execution.execution_engine import _generate_comap_empty_dfs
 from fugue.extensions.processor import Processor
 from fugue.extensions.transformer import CoTransformer, Transformer, to_transformer
 from triad.collections import ParamDict
@@ -56,8 +57,8 @@ def cotransform(self, df: DataFrame, tf: CoTransformer) -> DataFrame:
         )
         tf._key_schema = df.schema - list(df.metadata["serialized_cols"].values())
         # TODO: currently, get_output_schema only gets empty dataframes
-        empty_dfs = DataFrames(
-            {k: ArrayDataFrame([], v) for k, v in df.metadata["schemas"].items()}
+        empty_dfs = _generate_comap_empty_dfs(
+            df.metadata["schemas"], df.metadata.get("serialized_has_name", False)
         )
         tf._output_schema = Schema(tf.get_output_schema(empty_dfs))
         tr = _CoTransformerRunner(df, tf, self._ignore_errors)
@@ -100,30 +101,13 @@ def process(self, dfs: DataFrames) -> DataFrame:
         # TODO: this should also search on workflow conf
         temp_path = self.params.get_or_none("temp_path", str)
         to_file_threshold = self.params.get_or_none("to_file_threshold", object)
-        if to_file_threshold is None:  # pragma: no cover
-            to_file_threshold = -1
-        if len(dfs) == 1:
-            return self.execution_engine.serialize_by_partition(
-                dfs[0], partition_spec, "_0", temp_path, to_file_threshold
-            )
-        df = self.execution_engine.zip(
-            dfs[0],
-            dfs[1],
+        return self.execution_engine.zip_all(
+            dfs,
             how=how,
             partition_spec=partition_spec,
             temp_path=temp_path,
             to_file_threshold=to_file_threshold,
         )
-        for i in range(2, len(dfs)):
-            df = self.execution_engine.zip(
-                df,
-                dfs[i],
-                how=how,
-                partition_spec=partition_spec,
-                temp_path=temp_path,
-                to_file_threshold=to_file_threshold,
-            )
-        return df
 
 
 class Rename(Processor):

fugue/extensions/transformer/convert.py

@@ -77,11 +77,13 @@ def __uuid__(self) -> str:
         return to_uuid(self._wrapper.__uuid__(), self._output_schema_arg)
 
     def _parse_schema(self, obj: Any, df: DataFrame) -> Schema:
+        if callable(obj):
+            return obj(df, **self.params)
         if isinstance(obj, str):
             return df.schema.transform(obj)
         if isinstance(obj, List):
             return df.schema.transform(*obj)
-        return df.schema.transform(obj)
+        raise NotImplementedError  # pragma: no cover
 
     @staticmethod
     def from_func(func: Callable, schema: Any) -> "_FuncAsTransformer":
@@ -98,14 +100,30 @@ def from_func(func: Callable, schema: Any) -> "_FuncAsTransformer":
 
 class _FuncAsCoTransformer(CoTransformer):
     def get_output_schema(self, dfs: DataFrames) -> Any:
-        return self._parse_schema(self._output_schema_arg)  # type: ignore
+        return self._parse_schema(self._output_schema_arg, dfs)  # type: ignore
 
     @no_type_check
     def transform(self, dfs: DataFrames) -> LocalDataFrame:
-        args: List[Any] = [dfs] if self._dfs_input else list(dfs.values())
-        return self._wrapper.run(  # type: ignore
-            args, self.params, ignore_unknown=False, output_schema=self.output_schema
-        )
+        if self._dfs_input:  # function has DataFrames input
+            return self._wrapper.run(  # type: ignore
+                [dfs],
+                self.params,
+                ignore_unknown=False,
+                output_schema=self.output_schema,
+            )
+        if not dfs.has_key:  # input does not have key
+            return self._wrapper.run(  # type: ignore
+                list(dfs.values()),
+                self.params,
+                ignore_unknown=False,
+                output_schema=self.output_schema,
+            )
+        else:  # input DataFrames has key
+            p = dict(dfs)
+            p.update(self.params)
+            return self._wrapper.run(  # type: ignore
+                [], p, ignore_unknown=False, output_schema=self.output_schema
+            )
 
     def __call__(self, *args: Any, **kwargs: Any) -> Any:
         return self._wrapper(*args, **kwargs)  # type: ignore
@@ -116,15 +134,15 @@ def __uuid__(self) -> str:
             self._wrapper.__uuid__(), self._output_schema_arg, self._dfs_input
         )
 
-    def _parse_schema(self, obj: Any) -> Schema:
-        if obj is None:
-            return Schema()
+    def _parse_schema(self, obj: Any, dfs: DataFrames) -> Schema:
+        if callable(obj):
+            return obj(dfs, **self.params)
         if isinstance(obj, str):
             return Schema(obj)
         if isinstance(obj, List):
             s = Schema()
             for x in obj:
-                s += self._parse_schema(x)
+                s += self._parse_schema(x, dfs)
             return s
         return Schema(obj)