def test_load_csv_folder(self): e = self.engine native = NativeExecutionEngine() a = ArrayDataFrame([[6.1, 1.1]], "c:double,a:double") b = ArrayDataFrame([[2.1, 7.1], [4.1, 8.1]], "c:double,a:double") path = os.path.join(self.tmpdir, "a", "b") native.save_df(a, os.path.join(path, "a.csv"), format_hint="csv", header=True) native.save_df(b, os.path.join(path, "b.csv"), format_hint="csv", header=True) FileSystem().touch(os.path.join(path, "_SUCCESS")) c = e.load_df( path, format_hint="csv", header=True, infer_schema=True, columns=["a", "c"], ) df_eq(c, [[1.1, 6.1], [7.1, 2.1], [8.1, 4.1]], "a:double,c:double", throw=True)
def test_load_parquet_folder(self): e = self.engine native = NativeExecutionEngine() a = ArrayDataFrame([[6, 1]], "c:int,a:long") b = ArrayDataFrame([[2, 7], [4, 8]], "c:int,a:long") path = os.path.join(self.tmpdir, "a", "b") native.save_df(a, os.path.join(path, "a.parquet")) native.save_df(b, os.path.join(path, "b.parquet")) FileSystem().touch(os.path.join(path, "_SUCCESS")) c = e.load_df(path, format_hint="parquet", columns=["a", "c"]) df_eq(c, [[1, 6], [7, 2], [8, 4]], "a:long,c:int", throw=True)
def test_load_avro_folder(self): # TODO: switch to c:int,a:long when we can preserve schema to avro e = self.engine native = NativeExecutionEngine() a = ArrayDataFrame([[6, 1]], "c:long,a:long") b = ArrayDataFrame([[2, 7], [4, 8]], "c:long,a:long") path = os.path.join(self.tmpdir, "a", "b") native.save_df(a, os.path.join(path, "a.avro")) native.save_df(b, os.path.join(path, "b.avro")) FileSystem().touch(os.path.join(path, "_SUCCESS")) c = e.load_df(path, format_hint="avro", columns=["a", "c"]) df_eq(c, [[1, 6], [7, 2], [8, 4]], "a:long,c:long", throw=True)
class DaskExecutionEngine(ExecutionEngine): """The execution engine based on `Dask <https://docs.dask.org/>`_. Please read |ExecutionEngineTutorial| to understand this important Fugue concept :param conf: |ParamsLikeObject| defaults to None, read |FugueConfig| to learn Fugue specific options :Notice: You should setup Dask single machine or distributed environment in the :doc:`common <dask:setup>` way. Before initializing :class:`~.DaskExecutionEngine` """ def __init__(self, conf: Any = None): p = ParamDict(FUGUE_DASK_DEFAULT_CONF) p.update(ParamDict(conf)) super().__init__(p) self._fs = FileSystem() self._log = logging.getLogger() self._native = NativeExecutionEngine(conf=conf) def __repr__(self) -> str: return "DaskExecutionEngine" @property def log(self) -> logging.Logger: return self._log @property def fs(self) -> FileSystem: return self._fs @property def default_sql_engine(self) -> SQLEngine: return QPDDaskEngine(self) @property def pl_utils(self) -> DaskUtils: """Pandas-like dataframe utils""" return DaskUtils() def to_df(self, df: Any, schema: Any = None, metadata: Any = None) -> DaskDataFrame: """Convert a data structure to :class:`~fugue_dask.dataframe.DaskDataFrame` :param data: :class:`~fugue.dataframe.dataframe.DataFrame`, :class:`dask:dask.dataframe.DataFrame`, pandas DataFrame or list or iterable of arrays :param schema: |SchemaLikeObject|, defaults to None. :param metadata: |ParamsLikeObject|, defaults to None :return: engine compatible dataframe :Notice: * if the input is already :class:`~fugue_dask.dataframe.DaskDataFrame`, it should return itself * For list or iterable of arrays, ``schema`` must be specified * When ``schema`` is not None, a potential type cast may happen to ensure the dataframe's schema. * all other methods in the engine can take arbitrary dataframes and call this method to convert before doing anything """ default_partitions = self.conf.get_or_throw( FUGUE_DASK_CONF_DATAFRAME_DEFAULT_PARTITIONS, int) if isinstance(df, DataFrame): assert_or_throw( schema is None and metadata is None, ValueError( "schema and metadata must be None when df is a DataFrame"), ) if isinstance(df, DaskDataFrame): return df if isinstance(df, PandasDataFrame): return DaskDataFrame(df.native, df.schema, df.metadata, num_partitions=default_partitions) return DaskDataFrame( df.as_array(type_safe=True), df.schema, df.metadata, num_partitions=default_partitions, ) return DaskDataFrame(df, schema, metadata, num_partitions=default_partitions) def repartition(self, df: DataFrame, partition_spec: PartitionSpec) -> DaskDataFrame: df = self.to_df(df) if partition_spec.empty: return df if len(partition_spec.partition_by) > 0: return df p = partition_spec.get_num_partitions( **{ KEYWORD_ROWCOUNT: lambda: df.persist().count(), # type: ignore KEYWORD_CORECOUNT: lambda: 2, # TODO: remove this hard code }) if p > 0: return DaskDataFrame( df.native.repartition(npartitions=p), schema=df.schema, metadata=df.metadata, type_safe=False, ) return df def map( self, df: DataFrame, map_func: Callable[[PartitionCursor, LocalDataFrame], LocalDataFrame], output_schema: Any, partition_spec: PartitionSpec, metadata: Any = None, on_init: Optional[Callable[[int, DataFrame], Any]] = None, ) -> DataFrame: presort = partition_spec.presort presort_keys = list(presort.keys()) presort_asc = list(presort.values()) output_schema = Schema(output_schema) input_schema = df.schema on_init_once: Any = (None if on_init is None else RunOnce( on_init, lambda *args, **kwargs: to_uuid(id(on_init), id(args[0])))) def _map(pdf: Any) -> dd.DataFrame: if pdf.shape[0] == 0: return PandasDataFrame([], output_schema).as_pandas() if len(presort_keys) > 0: pdf = pdf.sort_values(presort_keys, ascending=presort_asc) input_df = PandasDataFrame(pdf.reset_index(drop=True), input_schema, pandas_df_wrapper=True) if on_init_once is not None: on_init_once(0, input_df) cursor = partition_spec.get_cursor(input_schema, 0) cursor.set(input_df.peek_array(), 0, 0) output_df = map_func(cursor, input_df) return output_df.as_pandas() df = self.to_df(df) if len(partition_spec.partition_by) == 0: pdf = self.repartition(df, partition_spec) result = pdf.native.map_partitions(_map, meta=output_schema.pandas_dtype) else: df = self.repartition( df, PartitionSpec(num=partition_spec.num_partitions)) result = self.pl_utils.safe_groupby_apply( df.native, partition_spec.partition_by, _map, meta=output_schema.pandas_dtype, ) return DaskDataFrame(result, output_schema, metadata) def broadcast(self, df: DataFrame) -> DataFrame: return self.to_df(df) def persist( self, df: DataFrame, lazy: bool = False, **kwargs: Any, ) -> DataFrame: return self.to_df(df).persist() def join( self, df1: DataFrame, df2: DataFrame, how: str, on: List[str] = _DEFAULT_JOIN_KEYS, metadata: Any = None, ) -> DataFrame: key_schema, output_schema = get_join_schemas(df1, df2, how=how, on=on) d = self.pl_utils.join( self.to_df(df1).native, self.to_df(df2).native, join_type=how, on=key_schema.names, ) return DaskDataFrame(d, output_schema, metadata) def union( self, df1: DataFrame, df2: DataFrame, distinct: bool = True, metadata: Any = None, ) -> DataFrame: assert_or_throw( df1.schema == df2.schema, lambda: ValueError(f"{df1.schema} != {df2.schema}"), ) d = self.pl_utils.union(self.to_df(df1).native, self.to_df(df2).native, unique=distinct) return DaskDataFrame(d, df1.schema, metadata) def subtract( self, df1: DataFrame, df2: DataFrame, distinct: bool = True, metadata: Any = None, ) -> DataFrame: assert_or_throw( distinct, NotImplementedError("EXCEPT ALL for DaskExecutionEngine")) assert_or_throw( df1.schema == df2.schema, lambda: ValueError(f"{df1.schema} != {df2.schema}"), ) d = self.pl_utils.except_df(self.to_df(df1).native, self.to_df(df2).native, unique=distinct) return DaskDataFrame(d, df1.schema, metadata) def intersect( self, df1: DataFrame, df2: DataFrame, distinct: bool = True, metadata: Any = None, ) -> DataFrame: assert_or_throw( distinct, NotImplementedError("INTERSECT ALL for DaskExecutionEngine")) assert_or_throw( df1.schema == df2.schema, lambda: ValueError(f"{df1.schema} != {df2.schema}"), ) d = self.pl_utils.intersect(self.to_df(df1).native, self.to_df(df2).native, unique=distinct) return DaskDataFrame(d, df1.schema, metadata) def distinct( self, df: DataFrame, metadata: Any = None, ) -> DataFrame: d = self.pl_utils.drop_duplicates(self.to_df(df).native) return DaskDataFrame(d, df.schema, metadata) def dropna( self, df: DataFrame, how: str = "any", thresh: int = None, subset: List[str] = None, metadata: Any = None, ) -> DataFrame: d = self.to_df(df).native.dropna(how=how, thresh=thresh, subset=subset) return DaskDataFrame(d, df.schema, metadata) def fillna( self, df: DataFrame, value: Any, subset: List[str] = None, metadata: Any = None, ) -> DataFrame: assert_or_throw( (not isinstance(value, list)) and (value is not None), ValueError("fillna value can not be a list or None"), ) if isinstance(value, dict): assert_or_throw( (None not in value.values()) and (any(value.values())), ValueError( "fillna dict can not contain None and needs at least one value" ), ) mapping = value else: # If subset is none, apply to all columns subset = subset or df.schema.names mapping = {col: value for col in subset} d = self.to_df(df).native.fillna(mapping) return DaskDataFrame(d, df.schema, metadata) def sample( self, df: DataFrame, n: Optional[int] = None, frac: Optional[float] = None, replace: bool = False, seed: Optional[int] = None, metadata: Any = None, ) -> DataFrame: assert_or_throw( (n is None and frac is not None) or (n is not None and frac is None), ValueError("one and only one of n and frac should be set"), ) # TODO: dask does not support sample by number of rows d = self.to_df(df).native.sample(n=n, frac=frac, replace=replace, random_state=seed) return DaskDataFrame(d, df.schema, metadata) def take( self, df: DataFrame, n: int, presort: str, na_position: str = "last", partition_spec: PartitionSpec = EMPTY_PARTITION_SPEC, metadata: Any = None, ) -> DataFrame: assert_or_throw( isinstance(n, int), ValueError("n needs to be an integer"), ) d = self.to_df(df).native meta = [(d[x].name, d[x].dtype) for x in d.columns] if presort: presort = parse_presort_exp(presort) # Use presort over partition_spec.presort if possible _presort: IndexedOrderedDict = presort or partition_spec.presort def _partition_take(partition, n, presort): if len(presort.keys()) > 0: partition = partition.sort_values( list(presort.keys()), ascending=list(presort.values()), na_position=na_position, ) return partition.head(n) if len(partition_spec.partition_by) == 0: if len(_presort.keys()) == 0: d = d.head(n) else: # Use the default partition d = (d.map_partitions( _partition_take, n, _presort, meta=meta).reset_index(drop=True).compute()) # compute() brings this to Pandas so we can use pandas d = d.sort_values( list(_presort.keys()), ascending=list(_presort.values()), na_position=na_position, ).head(n) else: d = (d.groupby(partition_spec.partition_by, dropna=False).apply( _partition_take, n=n, presort=_presort, meta=meta).reset_index(drop=True)) return DaskDataFrame(d, df.schema, metadata) def load_df( self, path: Union[str, List[str]], format_hint: Any = None, columns: Any = None, **kwargs: Any, ) -> DaskDataFrame: return self.to_df( load_df(path, format_hint=format_hint, columns=columns, fs=self.fs, **kwargs)) def save_df( self, df: DataFrame, path: str, format_hint: Any = None, mode: str = "overwrite", partition_spec: PartitionSpec = EMPTY_PARTITION_SPEC, force_single: bool = False, **kwargs: Any, ) -> None: if force_single: self._native.save_df( df, path=path, format_hint=format_hint, mode=mode, partition_spec=partition_spec, **kwargs, ) else: if not partition_spec.empty: self.log.warning( # pragma: no cover "partition_spec is not respected in %s.save_df", self) self.fs.makedirs(os.path.dirname(path), recreate=True) df = self.to_df(df) save_df(df, path, format_hint=format_hint, mode=mode, fs=self.fs, **kwargs)