def concat(objs: List):
"""
Concat the results of partitioned dask task executions. This function guess the
types of resulting list, then calls the corresponding native dask concat functions.
Parameters
----------
objs: List
List of the partitioned dask task execution results, which will be concat.
Returns
-------
obj:
The concat result
"""
if is_arraylike(objs[0]):
res = array_concat(objs,
axes=[0]) # TODO: Add concat with args support
elif any((is_dataframe_like(objs[0]), is_series_like(objs[0]),
is_index_like(objs[0]))):
res = df_concat(objs)
else:
res = objs
return res.compute() if is_dask_collection(res) else res
def assert_dask_dtypes(ddf, res, numeric_equal=True):
"""Check that the dask metadata matches the result.
If `numeric_equal`, integer and floating dtypes compare equal. This is
useful due to the implicit conversion of integer to floating upon
encountering missingness, which is hard to infer statically."""
eq_type_sets = [{"O", "S", "U", "a"}] # treat object and strings alike
if numeric_equal:
eq_type_sets.append({"i", "f", "u"})
def eq_dtypes(a, b):
return any(a.kind in eq_types and b.kind in eq_types
for eq_types in eq_type_sets) or (a == b)
if not is_dask_collection(res) and is_dataframe_like(res):
for col, a, b in pd.concat([ddf._meta.dtypes, res.dtypes],
axis=1).itertuples():
assert eq_dtypes(a, b)
elif not is_dask_collection(res) and (is_index_like(res)
or is_series_like(res)):
a = ddf._meta.dtype
b = res.dtype
assert eq_dtypes(a, b)
else:
if hasattr(ddf._meta, "dtype"):
a = ddf._meta.dtype
if not hasattr(res, "dtype"):
assert np.isscalar(res)
b = np.dtype(type(res))
else:
b = res.dtype
assert eq_dtypes(a, b)
else:
assert type(ddf._meta) == type(res)
def index(x):
if is_index_like(x):
return x
try:
return x.index.get_level_values(0)
except AttributeError:
return x.index
def has_known_categories(x):
"""Returns whether the categories in `x` are known.
Parameters
----------
x : Series or CategoricalIndex
"""
x = getattr(x, "_meta", x)
if is_series_like(x):
return UNKNOWN_CATEGORIES not in x.cat.categories
elif is_index_like(x) and hasattr(x, "categories"):
return UNKNOWN_CATEGORIES not in x.categories
raise TypeError("Expected Series or CategoricalIndex")
def check_meta(x, meta, funcname=None, numeric_equal=True):
"""Check that the dask metadata matches the result.
If metadata matches, ``x`` is passed through unchanged. A nice error is
raised if metadata doesn't match.
Parameters
----------
x : DataFrame, Series, or Index
meta : DataFrame, Series, or Index
The expected metadata that ``x`` should match
funcname : str, optional
The name of the function in which the metadata was specified. If
provided, the function name will be included in the error message to be
more helpful to users.
numeric_equal : bool, optionl
If True, integer and floating dtypes compare equal. This is useful due
to panda's implicit conversion of integer to floating upon encountering
missingness, which is hard to infer statically.
"""
eq_types = {"i", "f", "u"} if numeric_equal else set()
def equal_dtypes(a, b):
if is_categorical_dtype(a) != is_categorical_dtype(b):
return False
if isinstance(a, str) and a == "-" or isinstance(b, str) and b == "-":
return False
if is_categorical_dtype(a) and is_categorical_dtype(b):
if UNKNOWN_CATEGORIES in a.categories or UNKNOWN_CATEGORIES in b.categories:
return True
return a == b
return (a.kind in eq_types and b.kind in eq_types) or is_dtype_equal(
a, b)
if not (is_dataframe_like(meta) or is_series_like(meta)
or is_index_like(meta)) or is_dask_collection(meta):
raise TypeError("Expected partition to be DataFrame, Series, or "
"Index, got `%s`" % typename(type(meta)))
# Notice, we use .__class__ as opposed to type() in order to support
# object proxies see <https://github.com/dask/dask/pull/6981>
if x.__class__ != meta.__class__:
errmsg = "Expected partition of type `{}` but got `{}`".format(
typename(type(meta)),
typename(type(x)),
)
elif is_dataframe_like(meta):
dtypes = pd.concat([x.dtypes, meta.dtypes], axis=1, sort=True)
bad_dtypes = [(repr(col), a, b)
for col, a, b in dtypes.fillna("-").itertuples()
if not equal_dtypes(a, b)]
if bad_dtypes:
errmsg = "Partition type: `{}`\n{}".format(
typename(type(meta)),
asciitable(["Column", "Found", "Expected"], bad_dtypes),
)
else:
check_matching_columns(meta, x)
return x
else:
if equal_dtypes(x.dtype, meta.dtype):
return x
errmsg = "Partition type: `{}`\n{}".format(
typename(type(meta)),
asciitable(["", "dtype"], [("Found", x.dtype),
("Expected", meta.dtype)]),
)
raise ValueError("Metadata mismatch found%s.\n\n"
"%s" %
((" in `%s`" % funcname if funcname else ""), errmsg))