maxframe 1.0.0rc4__cp39-cp39-win_amd64.whl → 1.1.1__cp39-cp39-win_amd64.whl

maxframe/dataframe/extensions/flatjson.py

@@ -0,0 +1,131 @@
+# Copyright 1999-2024 Alibaba Group Holding Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import List
+
+from ... import opcodes
+from ...core import OutputType
+from ...serialization.serializables import ListField
+from ...serialization.serializables.field_type import FieldTypes
+from ..core import DataFrame
+from ..operators import DataFrameOperator, DataFrameOperatorMixin
+from ..utils import make_dtypes, parse_index
+
+
+class SeriesFlatJSONOperator(DataFrameOperator, DataFrameOperatorMixin):
+    _op_type_ = opcodes.FLATJSON
+
+    query_paths = ListField("query_paths", field_type=FieldTypes.string, default=None)
+
+    def __call__(self, series, dtypes):
+        if self._output_types[0] == OutputType.series:
+            name, dtype = dtypes
+            return self.new_series(
+                [series],
+                shape=series.shape,
+                index_value=series.index_value,
+                name=name,
+                dtype=dtype,
+            )
+        return self.new_dataframe(
+            [series],
+            shape=(series.shape[0], len(dtypes)),
+            index_value=series.index_value,
+            columns_value=parse_index(dtypes.index, store_data=True),
+            dtypes=make_dtypes(dtypes),
+        )
+
+
+def series_flatjson(
+    series,
+    query_paths: List[str],
+    dtypes=None,
+    dtype=None,
+    name: str = None,
+) -> DataFrame:
+    """
+    Flat JSON object in the series to a dataframe according to JSON query.
+
+    Parameters
+    ----------
+    series : Series
+        The series of json strings.
+
+    query_paths: List[str] or str
+        The JSON query paths for each generated column. The path format should follow
+        [RFC9535](https://datatracker.ietf.org/doc/rfc9535/).
+
+    dtypes : Series, default None
+        Specify dtypes of returned DataFrame. Can't work with dtype.
+
+    dtype : numpy.dtype, default None
+        Specify dtype of returned Series. Can't work with dtypes.
+
+    name : str, default None
+        Specify name of the returned Series.
+
+    Returns
+    -------
+    DataFrame or Series
+        Result of DataFrame when dtypes specified, else Series.
+
+    Examples
+    --------
+    >>> import maxframe.dataframe as md
+    >>> import pandas as pd
+    >>> s = md.Series(
+    ...     [
+    ...         '{"age": 24, "gender": "male", "graduated": false}',
+    ...         '{"age": 25, "gender": "female", "graduated": true}',
+    ...     ]
+    ... )
+    >>> s.execute()
+    0    {"age": 24, "gender": "male", "graduated": false}
+    1    {"age": 25, "gender": "female", "graduated": true}
+    dtype: object
+
+    >>> df = s.mf.flatjson(
+    ...    ["$.age", "$.gender", "$.graduated"],
+    ...    dtypes=pd.Series(["int32", "object", "bool"], index=["age", "gender", "graduated"]),
+    ... )
+    >>> df.execute()
+        age  gender  graduated
+    0   24    male       True
+    1   25  female       True
+
+    >>> s2 = s.mf.flatjson("$.age", name="age", dtype="int32")
+    >>> s2.execute()
+    0    24
+    1    25
+    Name: age, dtype: int32
+    """
+    if isinstance(query_paths, str):
+        query_paths = [query_paths]
+    if dtypes is not None and dtype is not None:
+        raise ValueError("Both dtypes and dtype cannot be specified at the same time.")
+    if dtype is not None:
+        if len(query_paths) != 1:
+            raise ValueError("query_paths should have only one path if dtype is set")
+        output_type = OutputType.series
+    elif dtypes is not None:
+        if len(dtypes) != len(query_paths):
+            raise ValueError("query_paths and dtypes should have same length")
+        output_type = OutputType.dataframe
+    else:
+        raise ValueError("dtypes or dtype should be specified")
+
+    dtypes = (name, dtype) if dtype is not None else dtypes
+    return SeriesFlatJSONOperator(query_paths=query_paths, _output_types=[output_type])(
+        series, dtypes
+    )

maxframe/dataframe/extensions/flatmap.py

@@ -17,17 +17,17 @@ from typing import Callable
 import numpy as np
 import pandas as pd
 
-from maxframe import opcodes
-from maxframe.core import OutputType
-from maxframe.dataframe.core import DataFrame, IndexValue
-from maxframe.dataframe.operators import DataFrameOperator, DataFrameOperatorMixin
-from maxframe.dataframe.utils import make_dtypes, parse_index
-from maxframe.serialization.serializables import (
+from ... import opcodes
+from ...core import OutputType
+from ...serialization.serializables import (
     BoolField,
     DictField,
     FunctionField,
     TupleField,
 )
+from ..core import DataFrame
+from ..operators import DataFrameOperator, DataFrameOperatorMixin
+from ..utils import gen_unknown_index_value, make_dtypes, parse_index
 
 
 class DataFrameFlatMapOperator(DataFrameOperator, DataFrameOperatorMixin):
@@ -41,22 +41,12 @@ class DataFrameFlatMapOperator(DataFrameOperator, DataFrameOperatorMixin):
     def __init__(self, output_types=None, **kw):
         super().__init__(_output_types=output_types, **kw)
 
-    @staticmethod
-    def _gen_flattening_index_value(index_value, *args) -> IndexValue:
-        pd_index = index_value.to_pandas()
-        if not isinstance(pd_index, pd.MultiIndex):
-            # for func return multi rows, will copy indexes
-            return parse_index(pd.Index([], dtype=pd_index.dtype), *args)
-        # multi index will keep the same level and types
-        return parse_index(
-            pd.MultiIndex.from_arrays([c[:0] for c in pd_index.levels]), *args
-        )
-
     def _call_dataframe(self, df: DataFrame, dtypes: pd.Series):
         dtypes = make_dtypes(dtypes)
-        index_value = self._gen_flattening_index_value(
+        index_value = gen_unknown_index_value(
             df.index_value,
             (df.key, df.index_value.key, self.func),
+            normalize_range_index=True,
         )
         return self.new_dataframe(
             [df],
@@ -67,9 +57,10 @@ class DataFrameFlatMapOperator(DataFrameOperator, DataFrameOperatorMixin):
         )
 
     def _call_series_or_index(self, series, dtypes=None):
-        index_value = self._gen_flattening_index_value(
+        index_value = gen_unknown_index_value(
             series.index_value,
             (series.key, series.index_value.key, self.func),
+            normalize_range_index=True,
         )
 
         if self.output_types[0] == OutputType.series:
@@ -114,9 +105,6 @@ def df_flatmap(dataframe, func: Callable, dtypes=None, raw=False, args=(), **kwa
 
     Parameters
     ----------
-    dataframe : DataFrame
-        The DataFrame to which the function will be applied.
-
     func : Callable
         Function to apply to each row of the DataFrame. It should accept a Series (or an array if `raw=True`)
         representing a row and return a list or iterable of values.
@@ -143,8 +131,8 @@ def df_flatmap(dataframe, func: Callable, dtypes=None, raw=False, args=(), **kwa
 
     Notes
     -----
-    The `func` must return an iterable of values for each input row. The index of the resulting DataFrame will be
-     repeated based on the number of output rows generated by `func`.
+    The ``func`` must return an iterable of values for each input row. The index of the resulting DataFrame will be
+    repeated based on the number of output rows generated by `func`.
 
     Examples
     --------
@@ -226,12 +214,9 @@ def series_flatmap(
 
     Parameters
     ----------
-    series : Series
-        The series to which the function will be applied.
-
     func : Callable
         Function to apply to each element of the Series. It should accept a scalar value
-        (or an array if `raw=True`) and return a list or iterable of values.
+        (or an array if ``raw=True``) and return a list or iterable of values.
 
     dtypes : Series, default None
         Specify dtypes of returned DataFrame. Can't work with dtype.
@@ -243,10 +228,10 @@ def series_flatmap(
         Specify name of the returned Series.
 
     args : tuple
-        Positional arguments to pass to `func`.
+        Positional arguments to pass to ``func``.
 
     **kwargs
-        Additional keyword arguments to pass as keywords arguments to `func`.
+        Additional keyword arguments to pass as keywords arguments to ``func``.
 
     Returns
     -------
@@ -255,9 +240,11 @@ def series_flatmap(
 
     Notes
     -----
-    The `func` must return an iterable of values for each input element. If `dtypes` is specified,
-    `flatmap` will return a DataFrame, if `dtype` and `name` is specified, a Series will be returned. The index of
-     the resulting DataFrame/Series will be repeated based on the number of output rows generated by `func`.
+    The ``func`` must return an iterable of values for each input element. If ``dtypes`` is specified,
+    `flatmap` will return a DataFrame, if ``dtype`` and ``name`` is specified, a Series will be returned.
+
+    The index of the resulting DataFrame/Series will be repeated based on the number of output rows generated
+    by ``func``.
 
     Examples
     --------
@@ -275,11 +262,7 @@ def series_flatmap(
     >>> def generate_values_array(x):
     ...     return [x * 2, x * 3]
 
-    >>> def generate_values_in_generator(x):
-    ...     yield pd.Series([x * 2, x * 4])
-    ...     yield pd.Series([x * 3, x * 5])
-
-    Specify `dtype` with a function which returns list to return more than one elements as a Series:
+    Specify ``dtype`` with a function which returns list to return more elements as a Series:
 
     >>> df['A'].mf.flatmap(generate_values_array, dtype="int", name="C").execute()
         0    2
@@ -290,7 +273,12 @@ def series_flatmap(
         2    9
         Name: C, dtype: int64
 
-    Specify `dtypes` to return multi columns as a DataFrame:
+    Specify ``dtypes`` to return multi columns as a DataFrame:
+
+
+    >>> def generate_values_in_generator(x):
+    ...     yield pd.Series([x * 2, x * 4])
+    ...     yield pd.Series([x * 3, x * 5])
 
     >>> df['A'].mf.flatmap(generate_values_in_generator, dtypes={"A": "int", "B": "int"}).execute()
            A   B
@@ -302,7 +290,7 @@ def series_flatmap(
         2  9  15
     """
 
-    if dtypes and dtype:
+    if dtypes is not None and dtype is not None:
         raise ValueError("Both dtypes and dtype cannot be specified at the same time.")
 
     dtypes = (name, dtype) if dtype is not None else dtypes

maxframe/dataframe/extensions/reshuffle.py

@@ -38,7 +38,7 @@ class DataFrameReshuffle(DataFrameOperator, DataFrameOperatorMixin):
         else:
             idx_value = df.index_value
             if isinstance(idx_value.value, IndexValue.RangeIndex):
-                idx_value = parse_index(pd.Int64Index([0]))
+                idx_value = parse_index(pd.RangeIndex(1))
         params = df.params
         params["index_value"] = idx_value
         self._output_types = get_output_types(df)

maxframe/dataframe/extensions/tests/test_apply_chunk.py

@@ -0,0 +1,186 @@
+# Copyright 1999-2024 Alibaba Group Holding Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import numpy as np
+import pandas as pd
+import pytest
+
+from ....udf import MarkedFunction
+from ... import DataFrame
+from ...core import DATAFRAME_TYPE, SERIES_TYPE
+
+
+@pytest.fixture
+def df1():
+    return DataFrame({"a": [1, 2, 3], "b": [1, 2, 3], "c": [1, 2, 3]})
+
+
+@pytest.fixture
+def df2():
+    return DataFrame([[1, 2, 3], [1, 2, 3], [1, 2, 3]], columns=["a", "b", "c"])
+
+
+@pytest.fixture
+def df3():
+    return DataFrame(
+        [[1, 2, 3], [1, 2, 3], [1, 2, 3]],
+        columns=["a", "b", "c"],
+        index=pd.MultiIndex.from_arrays([[1, 2, 3], [1, 2, 3]], names=["A", "B"]),
+    )
+
+
+def test_apply_chunk_infer_dtypes_and_index(df1, df2, df3):
+    # dataframe -> dataframe filter
+    result = df3.mf.apply_chunk(
+        lambda data: data.query("A > 1"), batch_rows=2, output_type="dataframe"
+    )
+    assert isinstance(result, DATAFRAME_TYPE)
+    assert df3.index_value.key != result.index_value.key
+    assert df3.index_value.to_pandas().names == result.index_value.to_pandas().names
+
+    # dataframe -> dataframe keep same
+    result = df1.mf.apply_chunk(
+        lambda data: data, batch_rows=2, output_type="dataframe"
+    )
+    assert isinstance(result, DATAFRAME_TYPE)
+    assert result.index_value is df1.index_value
+
+    # dataframe -> dataframe ufunc with arguments
+    result = df1.mf.apply_chunk(
+        np.add, batch_rows=2, args=(2,), output_type="dataframe"
+    )
+    assert isinstance(result, DATAFRAME_TYPE)
+    assert result.index_value is df1.index_value
+    assert result.dtypes.equals(df1.dtypes)
+    assert result.shape == df1.shape
+
+    # dataframe -> series ufunc return series
+    result = df1.mf.apply_chunk(np.sum, batch_rows=2)
+    assert isinstance(result, SERIES_TYPE)
+    assert result.index_value is not df1.index_value
+
+    # series -> series
+    result = df3.a.mf.apply_chunk(lambda data: data, batch_rows=2, output_type="series")
+    assert isinstance(result, SERIES_TYPE)
+    assert df3.a.index_value is result.index_value
+
+    result = df3.a.mf.apply_chunk(
+        np.sum, batch_rows=2, output_type="series", dtype=np.int64, name="sum"
+    )
+    assert isinstance(result, SERIES_TYPE)
+    assert isinstance(result.index_value.to_pandas(), pd.RangeIndex)
+
+    # general functions
+    def process(data, param, k):
+        return data * param * k
+
+    result = df2.mf.apply_chunk(
+        process, batch_rows=3, output_type="dataframe", args=(4,), k=1
+    )
+    assert result.index_value is df2.index_value
+    assert result.dtypes.equals(df2.dtypes)
+
+    # mark functions
+    from ....udf import with_python_requirements, with_resources
+
+    @with_resources("empty.txt")
+    @with_python_requirements("numpy")
+    def process(data, k):
+        return data
+
+    result = df1.mf.apply_chunk(process, batch_rows=3, output_type="dataframe", k=1)
+    assert result.index_value is df1.index_value
+    assert result.dtypes.equals(df1.dtypes)
+    assert isinstance(result.op.func, MarkedFunction)
+    assert result.op.func is not process
+    assert result.op.func.resources is process.resources
+    assert result.op.func.pythonpacks is process.pythonpacks
+
+    def func_series_ret_series(data):
+        return pd.DataFrame([data, data])
+
+    result = df3.a.mf.apply_chunk(
+        func_series_ret_series, batch_rows=2, output_type="dataframe"
+    )
+    assert isinstance(result, DATAFRAME_TYPE)
+    assert result.op.func is func_series_ret_series
+
+
+def test_apply_test(df1):
+    def process(x, param):
+        return x * param
+
+    result = df1.a.mf.apply_chunk(
+        process, batch_rows=2, output_type="series", args=(5,)
+    )
+    assert isinstance(result, SERIES_TYPE)
+
+
+def test_apply_chunk(df1):
+    keys = [1, 2]
+
+    def f(x, keys):
+        if x["a"] in keys:
+            return [1, 0]
+        else:
+            return [0, 1]
+
+    result = df1[["a"]].mf.apply_chunk(
+        f,
+        output_type="dataframe",
+        dtypes=pd.Series(["int64", "int64"]),
+        batch_rows=5,
+        keys=keys,
+    )
+    assert result.shape == (np.nan, 2)
+    assert df1.index_value.key != result.index_value.key
+
+    # dataframe return series
+    result = df1.mf.apply_chunk(
+        lambda x: x.a,
+        output_type="series",
+        dtype="int64",
+        batch_rows=5,
+    )
+    assert result.shape == (np.nan,)
+    assert df1.index_value.key == result.index_value.key
+    assert df1.a.index_value.key == result.index_value.key
+
+    # return dataframe with given dtypes
+    result = df1.a.mf.apply_chunk(
+        lambda x: pd.concat([x, x], axis=1),
+        output_type="dataframe",
+        dtypes=pd.Series(["int64", "int64"]),
+        batch_rows=5,
+    )
+    assert result.shape == (np.nan, 2)
+    assert df1.a.index_value.key != result.index_value.key
+
+    # return series but as dataframe
+    result = df1.a.mf.apply_chunk(
+        lambda x: pd.concat([x, x], axis=0),
+        output_type="dataframe",
+        dtypes={"c": np.int_},
+        batch_rows=5,
+    )
+    assert result.shape == (np.nan, 1)
+
+
+def test_apply_chunk_exception(df1):
+    with pytest.raises(ValueError):
+        df1.mf.apply_chunk(lambda data: data, batch_rows=-1, output_type="dataframe")
+
+    with pytest.raises(TypeError):
+        df1.mf.apply_chunk(
+            lambda data: data, batch_rows=object(), output_type="dataframe"
+        )

maxframe/dataframe/extensions/tests/test_extensions.py

@@ -16,8 +16,9 @@ import pandas as pd
 import pytest
 
 from .... import dataframe as md
+from ....tests.utils import assert_mf_index_dtype
 from ... import DataFrame
-from ...core import IndexValue
+from ...core import DATAFRAME_TYPE, SERIES_TYPE, IndexValue
 from ..reshuffle import DataFrameReshuffle
 
 
@@ -32,7 +33,7 @@ def test_reshuffle():
 
     r = mdf.mf.reshuffle()
     assert isinstance(r.op, DataFrameReshuffle)
-    assert isinstance(r.index_value.value, IndexValue.Int64Index)
+    assert_mf_index_dtype(r.index_value.value, np.int64)
 
     r = mdf.mf.reshuffle(ignore_index=True)
     assert isinstance(r.op, DataFrameReshuffle)
@@ -97,3 +98,46 @@ def test_flatmap(df1, df2, df3):
     assert apply_s.shape == (np.nan, 2)
     assert df3.index_value.key != apply_s.index_value.key
     assert df3.key != apply_s.index_value.key
+
+
+def test_flatjson():
+    s1 = md.Series(["{{'a': 1, 'b': false}}"], index=[1])
+    df1 = s1.mf.flatjson(
+        ["$.a", "$.b"], dtypes=pd.Series(["int32", "bool"], index=["a", "b"])
+    )
+    assert df1.shape == (1, 2)
+    assert df1.index_value.key == s1.index_value.key
+    assert isinstance(df1, DATAFRAME_TYPE)
+    assert list(df1.dtypes) == [np.dtype("int32"), np.dtype("bool")]
+    assert list(df1.dtypes.index) == ["a", "b"]
+
+    df2 = s1.mf.flatjson(["$.a"], dtypes=pd.Series(["int32"], index=["a"]))
+    assert df2.shape == (1, 1)
+    assert df2.index_value.key == s1.index_value.key
+    assert isinstance(df2, DATAFRAME_TYPE)
+    assert list(df2.dtypes) == [np.dtype("int32")]
+    assert list(df2.dtypes.index) == ["a"]
+
+    s2 = s1.mf.flatjson("$.a", dtype="int32", name="a")
+    assert s2.shape == (1,)
+    assert s2.index_value.key == s1.index_value.key
+    assert isinstance(s2, SERIES_TYPE)
+    assert s2.dtype == np.dtype("int32")
+    assert s2.name == "a"
+
+    with pytest.raises(ValueError):
+        s1.mf.flatjson([], dtypes=pd.Series(["int32", "bool"], index=["a", "b"]))
+    with pytest.raises(ValueError):
+        s1.mf.flatjson(["$.a"], dtypes=pd.Series(["int32", "bool"], index=["a", "b"]))
+    with pytest.raises(ValueError):
+        s1.mf.flatjson(["$.a"], dtypes=pd.Series(["int32", "bool"], index=["a", "b"]))
+    with pytest.raises(ValueError):
+        s1.mf.flatjson(["$.a", "$.b"], dtypes=pd.Series(["bool"], index=["b"]))
+    with pytest.raises(ValueError):
+        s1.mf.flatjson(
+            ["$.a"],
+            dtype="int32",
+            dtypes=pd.Series(["int32"], index=["a"]),
+        )
+    with pytest.raises(ValueError):
+        s1.mf.flatjson(["$.a"])

maxframe/dataframe/groupby/__init__.py

@@ -55,6 +55,7 @@ def _install():
         setattr(cls, "kurtosis", lambda groupby, **kw: agg(groupby, "kurtosis", **kw))
         setattr(cls, "sem", lambda groupby, **kw: agg(groupby, "sem", **kw))
         setattr(cls, "nunique", lambda groupby, **kw: agg(groupby, "nunique", **kw))
+        setattr(cls, "median", lambda groupby, **kw: agg(groupby, "median", **kw))
 
         setattr(cls, "apply", groupby_apply)
         setattr(cls, "transform", groupby_transform)

maxframe/dataframe/groupby/aggregation.py

@@ -79,6 +79,7 @@ _agg_functions = {
     "kurt": lambda x, bias=False: x.kurt(bias=bias),
     "kurtosis": lambda x, bias=False: x.kurtosis(bias=bias),
     "nunique": lambda x: x.nunique(),
+    "median": lambda x: x.median(),
 }
 _series_col_name = "col_name"
 

maxframe/dataframe/groupby/apply.py

@@ -28,7 +28,13 @@ from ...serialization.serializables import (
 )
 from ...utils import get_func_token, quiet_stdio, tokenize
 from ..operators import DataFrameOperator, DataFrameOperatorMixin
-from ..utils import make_dtype, make_dtypes, parse_index, validate_output_types
+from ..utils import (
+    copy_func_scheduling_hints,
+    make_dtype,
+    make_dtypes,
+    parse_index,
+    validate_output_types,
+)
 
 
 class GroupByApplyLogicKeyGeneratorMixin(OperatorLogicKeyGeneratorMixin):
@@ -56,6 +62,8 @@ class GroupByApply(
 
     def __init__(self, output_types=None, **kw):
         super().__init__(_output_types=output_types, **kw)
+        if hasattr(self, "func"):
+            copy_func_scheduling_hints(self.func, self)
 
     def _update_key(self):
         values = [v for v in self._values_ if v is not self.func] + [

maxframe/dataframe/groupby/core.py

@@ -28,7 +28,7 @@ from ..utils import build_df, build_series, parse_index
 
 cudf = lazy_import("cudf")
 
-_GROUP_KEYS_NO_DEFAULT = pd_release_version >= (1, 5, 0)
+_GROUP_KEYS_NO_DEFAULT = pd_release_version[:2] == (1, 5)
 _default_group_keys = no_default if _GROUP_KEYS_NO_DEFAULT else True
 
 

maxframe/dataframe/groupby/fill.py

@@ -35,12 +35,15 @@ class GroupByFillOperator(DataFrameOperator, DataFrameOperatorMixin):
         func_name = getattr(self, "_func_name")
 
         if func_name == "fillna":
+            kw = {}
+            if self.axis is not None:
+                kw["axis"] = self.axis
             result_df = mock_groupby.fillna(
                 value=self.value,
                 method=self.method,
-                axis=self.axis,
                 limit=self.limit,
                 downcast=self.downcast,
+                **kw,
             )
         else:
             result_df = getattr(mock_groupby, func_name)(limit=self.limit)

maxframe/dataframe/groupby/getitem.py

@@ -88,5 +88,11 @@ def df_groupby_getitem(df_groupby, item):
     if df_groupby.selection:
         raise IndexError(f"Column(s) {df_groupby.selection!r} already selected")
 
+    if (
+        isinstance(item, tuple)
+        and item not in df_groupby.dtypes
+        and item not in df_groupby.index.names
+    ):
+        item = list(item)
     op = GroupByIndex(selection=item, output_types=output_types)
     return op(df_groupby)

maxframe/dataframe/groupby/tests/test_groupby.py

@@ -230,7 +230,7 @@ def test_groupby_transform():
     assert r.op._op_type_ == opcodes.TRANSFORM
     assert r.op.output_types[0] == OutputType.dataframe
 
-    r = mdf.groupby("b").transform(["cummax", "cumcount"], _call_agg=True)
+    r = mdf[list("abde")].groupby("b").transform(["cummax", "cumcount"], _call_agg=True)
     assert r.shape == (np.nan, 6)
     assert r.op._op_type_ == opcodes.TRANSFORM
     assert r.op.output_types[0] == OutputType.dataframe

maxframe/dataframe/groupby/transform.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import logging
+
 import numpy as np
 import pandas as pd
 
@@ -20,7 +22,9 @@ from ...core import OutputType
 from ...serialization.serializables import AnyField, BoolField, DictField, TupleField
 from ...utils import quiet_stdio
 from ..operators import DataFrameOperator, DataFrameOperatorMixin
-from ..utils import parse_index
+from ..utils import copy_func_scheduling_hints, parse_index
+
+logger = logging.getLogger(__name__)
 
 
 class GroupByTransform(DataFrameOperator, DataFrameOperatorMixin):
@@ -35,6 +39,8 @@ class GroupByTransform(DataFrameOperator, DataFrameOperatorMixin):
 
     def __init__(self, output_types=None, **kw):
         super().__init__(_output_types=output_types, **kw)
+        if hasattr(self, "func"):
+            copy_func_scheduling_hints(self.func, self)
 
     def _infer_df_func_returns(self, in_groupby, dtypes, index):
         index_value, output_types, new_dtypes = None, None, None
@@ -65,7 +71,7 @@ class GroupByTransform(DataFrameOperator, DataFrameOperatorMixin):
                 output_types = [OutputType.series]
                 new_dtypes = new_dtypes or (infer_df.name, infer_df.dtype)
         except:  # noqa: E722  # nosec
-            pass
+            logger.info("Exception raised while inferring df_func", exc_info=True)
 
         self.output_types = output_types if not self.output_types else self.output_types
         dtypes = new_dtypes if dtypes is None else dtypes