maxframe 1.3.0__cp311-cp311-win32.whl → 2.0.0__cp311-cp311-win32.whl

maxframe/dataframe/misc/rechunk.py

@@ -0,0 +1,59 @@
+# Copyright 1999-2025 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 ... import opcodes
+from ...core import OutputType
+from ...serialization.serializables import AnyField
+from ...tensor.rechunk.rechunk import chunk_size_type
+from ...typing_ import TileableType
+from ..core import DATAFRAME_TYPE
+from ..operators import DataFrameOperator, DataFrameOperatorMixin
+
+
+class DataFrameRechunk(DataFrameOperator, DataFrameOperatorMixin):
+    _op_type_ = opcodes.RECHUNK
+
+    chunk_size = AnyField("chunk_size")
+
+    def __call__(self, x):
+        if isinstance(x, DATAFRAME_TYPE):
+            return self.new_dataframe(
+                [x],
+                shape=x.shape,
+                dtypes=x.dtypes,
+                columns_value=x.columns_value,
+                index_value=x.index_value,
+            )
+        else:
+            self.output_types = x.op.output_types
+            f = (
+                self.new_series
+                if self.output_types[0] == OutputType.series
+                else self.new_index
+            )
+            return f(
+                [x],
+                shape=x.shape,
+                dtype=x.dtype,
+                index_value=x.index_value,
+                name=x.name,
+            )
+
+
+def rechunk(a: TileableType, chunk_size: chunk_size_type, reassign_worker=False):
+    op = DataFrameRechunk(
+        chunk_size=chunk_size,
+        reassign_worker=reassign_worker,
+    )
+    return op(a)

maxframe/dataframe/misc/shift.py

@@ -12,10 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import List
+
 import pandas as pd
 
 from ... import opcodes
-from ...core import OutputType
+from ...core import EntityData, OutputType
 from ...serialization.serializables import AnyField, Int8Field, Int64Field, KeyField
 from ...utils import no_default, pd_release_version
 from ..operators import DataFrameOperator, DataFrameOperatorMixin
@@ -39,9 +41,10 @@ class DataFrameShift(DataFrameOperator, DataFrameOperatorMixin):
     def input(self):
         return self._input
 
-    def _set_inputs(self, inputs):
-        super()._set_inputs(inputs)
-        self._input = self._inputs[0]
+    @classmethod
+    def _set_inputs(cls, op: "DataFrameShift", inputs: List[EntityData]):
+        super()._set_inputs(op, inputs)
+        op._input = op._inputs[0]
 
     def _call_dataframe(self, df):
         test_df = build_df(df)

maxframe/dataframe/misc/stack.py

@@ -18,6 +18,7 @@ import numpy as np
 import pandas as pd
 
 from ... import opcodes
+from ...core import EntityData
 from ...serialization.serializables import AnyField, BoolField, KeyField
 from ..operators import DataFrameOperator, DataFrameOperatorMixin
 from ..utils import build_df, parse_index
@@ -30,9 +31,10 @@ class DataFrameStack(DataFrameOperator, DataFrameOperatorMixin):
     level = AnyField("level", default=None)
     dropna = BoolField("dropna", default=None)
 
-    def _set_inputs(self, inputs):
-        super()._set_inputs(inputs)
-        self._input_df = self._inputs[0]
+    @classmethod
+    def _set_inputs(cls, op: "DataFrameStack", inputs: List[EntityData]):
+        super()._set_inputs(op, inputs)
+        op._input_df = op._inputs[0]
 
     @classmethod
     def _calc_size(cls, size: int, level: Union[List, int], dtypes: pd.Series):

maxframe/dataframe/misc/tests/test_misc.py

@@ -31,6 +31,166 @@ from ...datasource.series import from_pandas as from_pandas_series
 from .. import astype, cut
 
 
+def test_dataframe_apply():
+    cols = [chr(ord("A") + i) for i in range(10)]
+    df_raw = pd.DataFrame(dict((c, [i**2 for i in range(20)]) for c in cols))
+
+    df = from_pandas_df(df_raw, chunk_size=5)
+
+    def df_func_with_err(v):
+        assert len(v) > 2
+        return v.sort_values()
+
+    def df_series_func_with_err(v):
+        assert len(v) > 2
+        return 0
+
+    with pytest.raises(TypeError):
+        df.apply(df_func_with_err)
+
+    r = df.apply(df_func_with_err, output_type="dataframe", dtypes=df_raw.dtypes)
+    assert r.shape == (np.nan, df.shape[-1])
+    assert r.op._op_type_ == opcodes.APPLY
+    assert r.op.output_types[0] == OutputType.dataframe
+    assert r.op.elementwise is False
+
+    r = df.apply(
+        df_series_func_with_err, output_type="series", dtype=object, name="output"
+    )
+    assert r.dtype == np.dtype("O")
+    assert r.shape == (df.shape[-1],)
+    assert r.op._op_type_ == opcodes.APPLY
+    assert r.op.output_types[0] == OutputType.series
+    assert r.op.elementwise is False
+
+    r = df.apply("ffill")
+    assert r.op._op_type_ == opcodes.FILL_NA
+
+    r = df.apply(np.sqrt)
+    assert all(v == np.dtype("float64") for v in r.dtypes) is True
+    assert r.shape == df.shape
+    assert r.op._op_type_ == opcodes.APPLY
+    assert r.op.output_types[0] == OutputType.dataframe
+    assert r.op.elementwise is True
+
+    r = df.apply(lambda x: pd.Series([1, 2]))
+    assert all(v == np.dtype("int64") for v in r.dtypes) is True
+    assert r.shape == (np.nan, df.shape[1])
+    assert r.op.output_types[0] == OutputType.dataframe
+    assert r.op.elementwise is False
+
+    r = df.apply(np.sum, axis="index")
+    assert np.dtype("int64") == r.dtype
+    assert r.shape == (df.shape[1],)
+    assert r.op.output_types[0] == OutputType.series
+    assert r.op.elementwise is False
+
+    r = df.apply(np.sum, axis="columns")
+    assert np.dtype("int64") == r.dtype
+    assert r.shape == (df.shape[0],)
+    assert r.op.output_types[0] == OutputType.series
+    assert r.op.elementwise is False
+
+    r = df.apply(lambda x: pd.Series([1, 2], index=["foo", "bar"]), axis=1)
+    assert all(v == np.dtype("int64") for v in r.dtypes) is True
+    assert r.shape == (df.shape[0], 2)
+    assert r.op.output_types[0] == OutputType.dataframe
+    assert r.op.elementwise is False
+
+    r = df.apply(lambda x: [1, 2], axis=1, result_type="expand")
+    assert all(v == np.dtype("int64") for v in r.dtypes) is True
+    assert r.shape == (df.shape[0], 2)
+    assert r.op.output_types[0] == OutputType.dataframe
+    assert r.op.elementwise is False
+
+    r = df.apply(lambda x: list(range(10)), axis=1, result_type="reduce")
+    assert np.dtype("object") == r.dtype
+    assert r.shape == (df.shape[0],)
+    assert r.op.output_types[0] == OutputType.series
+    assert r.op.elementwise is False
+
+    r = df.apply(lambda x: list(range(10)), axis=1, result_type="broadcast")
+    assert all(v == np.dtype("int64") for v in r.dtypes) is True
+    assert r.shape == (df.shape[0], 10)
+    assert r.op.output_types[0] == OutputType.dataframe
+    assert r.op.elementwise is False
+
+    raw = pd.DataFrame({"a": [np.array([1, 2, 3]), np.array([4, 5, 6])]})
+    df = from_pandas_df(raw)
+    df2 = df.apply(
+        lambda x: x["a"].astype(pd.Series),
+        axis=1,
+        output_type="dataframe",
+        dtypes=pd.Series([np.dtype(float)] * 3),
+    )
+    assert df2.ndim == 2
+
+
+def test_series_apply():
+    idxes = [chr(ord("A") + i) for i in range(20)]
+    s_raw = pd.Series([i**2 for i in range(20)], index=idxes)
+
+    series = from_pandas_series(s_raw, chunk_size=5)
+
+    r = series.apply("add", args=(1,))
+    assert r.op._op_type_ == opcodes.ADD
+
+    r = series.apply(np.sqrt)
+    assert np.dtype("float64") == r.dtype
+    assert r.shape == series.shape
+    assert r.index_value is series.index_value
+    assert r.op._op_type_ == opcodes.APPLY
+    assert r.op.output_types[0] == OutputType.series
+
+    r = series.apply("sqrt")
+    assert np.dtype("float64") == r.dtype
+    assert r.shape == series.shape
+    assert r.op._op_type_ == opcodes.APPLY
+    assert r.op.output_types[0] == OutputType.series
+
+    r = series.apply(lambda x: [x, x + 1], convert_dtype=False)
+    assert np.dtype("object") == r.dtype
+    assert r.shape == series.shape
+    assert r.op._op_type_ == opcodes.APPLY
+    assert r.op.output_types[0] == OutputType.series
+
+    s_raw2 = pd.Series([np.array([1, 2, 3]), np.array([4, 5, 6])])
+    series = from_pandas_series(s_raw2)
+
+    r = series.apply(np.sum)
+    assert r.dtype == np.dtype(object)
+
+    r = series.apply(lambda x: pd.Series([1]), output_type="dataframe")
+    expected = s_raw2.apply(lambda x: pd.Series([1]))
+    pd.testing.assert_series_equal(r.dtypes, expected.dtypes)
+
+    dtypes = pd.Series([np.dtype(float)] * 3)
+    r = series.apply(pd.Series, output_type="dataframe", dtypes=dtypes)
+    assert r.ndim == 2
+    pd.testing.assert_series_equal(r.dtypes, dtypes)
+    assert r.shape == (2, 3)
+
+    def apply_with_error(_):
+        raise ValueError
+
+    r = series.apply(apply_with_error, output_type="dataframe", dtypes=dtypes)
+    assert r.ndim == 2
+
+    r = series.apply(
+        pd.Series, output_type="dataframe", dtypes=dtypes, index=pd.RangeIndex(2)
+    )
+    assert r.ndim == 2
+    pd.testing.assert_series_equal(r.dtypes, dtypes)
+    assert r.shape == (2, 3)
+
+    with pytest.raises(AttributeError, match="abc"):
+        series.apply("abc")
+
+    with pytest.raises(TypeError):
+        # dtypes not provided
+        series.apply(lambda x: x.tolist(), output_type="dataframe")
+
+
 def test_transform():
     cols = [chr(ord("A") + i) for i in range(10)]
     df_raw = pd.DataFrame(dict((c, [i**2 for i in range(20)]) for c in cols))
@@ -339,6 +499,11 @@ def test_get_dummies():
     r = get_dummies(df)
     assert isinstance(r, DATAFRAME_TYPE)
 
+    raw = pd.Series(["a", "a", "b", "c"])
+    ms = from_pandas_series(raw, chunk_size=2)
+    r = get_dummies(ms)
+    assert isinstance(r, DATAFRAME_TYPE)
+
 
 def test_to_numeric():
     raw = pd.DataFrame({"a": [1.0, 2, 3, -3]})
@@ -381,7 +546,7 @@ def test_apply():
 
     keys = [1, 2]
 
-    @with_running_options(engine="spe")
+    @with_running_options(engine="spe", memory="40GB")
     def f(x, keys):
         if x["a"] in keys:
             return [1, 0]
@@ -398,6 +563,7 @@ def test_apply():
     )
     assert apply_df.shape == (3, 2)
     assert apply_df.op.expect_engine == "SPE"
+    assert apply_df.op.expect_resources == {"cpu": 1, "memory": "40GB", "gpu": 0}
 
 
 def test_pivot_table():

maxframe/dataframe/misc/transform.py

@@ -12,23 +12,24 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any, Union
+from typing import Any, MutableMapping, Union
 
 import numpy as np
-import pandas as pd
 from pandas import DataFrame, Series
 
 from ... import opcodes
 from ...core import OutputType
 from ...serialization.serializables import AnyField, BoolField, DictField, TupleField
-from ...utils import pd_release_version, quiet_stdio
+from ...udf import BuiltinFunction, MarkedFunction
+from ...utils import copy_if_possible, pd_release_version
 from ..core import DATAFRAME_TYPE
 from ..operators import DataFrameOperator, DataFrameOperatorMixin
 from ..utils import (
+    InferredDataFrameMeta,
     build_df,
     build_series,
     copy_func_scheduling_hints,
-    make_dtypes,
+    infer_dataframe_return_value,
     pack_func_args,
     parse_index,
     validate_axis,
@@ -53,63 +54,55 @@ class TransformOperator(DataFrameOperator, DataFrameOperatorMixin):
         if hasattr(self, "func"):
             copy_func_scheduling_hints(self.func, self)
 
-    def _infer_df_func_returns(self, df, dtypes):
-        packed_funcs = self.func
-        test_df = _build_stub_pandas_obj(df, self.output_types[0])
-        if self.output_types[0] == OutputType.dataframe:
-            try:
-                with np.errstate(all="ignore"), quiet_stdio():
-                    if self.call_agg:
-                        infer_df = test_df.agg(packed_funcs, axis=self.axis)
-                    else:
-                        infer_df = test_df.transform(packed_funcs, axis=self.axis)
-            except:  # noqa: E722
-                infer_df = None
-        else:
-            try:
-                with np.errstate(all="ignore"), quiet_stdio():
-                    if self.call_agg:
-                        infer_df = test_df.agg(packed_funcs)
-                    else:
-                        if not _with_convert_dtype:
-                            infer_df = test_df.transform(packed_funcs)
-                        else:  # pragma: no cover
-                            infer_df = test_df.transform(
-                                packed_funcs, convert_dtype=self.convert_dtype
-                            )
-            except:  # noqa: E722
-                infer_df = None
-
-        if infer_df is None and dtypes is None:
-            raise TypeError(
-                "Failed to infer dtype, please specify dtypes as arguments."
-            )
-
-        if infer_df is None:
-            is_df = self.output_types[0] == OutputType.dataframe
-        else:
-            is_df = isinstance(infer_df, pd.DataFrame)
-
-        if is_df:
-            new_dtypes = make_dtypes(dtypes) if dtypes is not None else infer_df.dtypes
-            self.output_types = [OutputType.dataframe]
-        else:
-            new_dtypes = (
-                dtypes if dtypes is not None else (infer_df.name, infer_df.dtype)
-            )
-            self.output_types = [OutputType.series]
-
-        return new_dtypes
+    def has_custom_code(self) -> bool:
+        return not isinstance(self.func, BuiltinFunction)
 
-    def __call__(self, df, dtypes=None, index=None, skip_infer=None):
+    def _infer_df_func_returns(
+        self, df, dtypes=None, dtype=None, name=None, index=None
+    ) -> InferredDataFrameMeta:
+        def infer_func(df_obj):
+            if self.call_agg:
+                return df_obj.agg(self.func, self.axis)
+            else:
+                return df_obj.transform(self.func, self.axis)
+
+        res = infer_dataframe_return_value(
+            df,
+            infer_func,
+            self.output_types[0] if self.output_types else None,
+            dtypes=dtypes,
+            dtype=dtype,
+            name=name,
+            index=index,
+            inherit_index=True,
+        )
+        res.check_absence("dtypes", "dtype")
+        return res
+
+    def __call__(
+        self, df, dtypes=None, dtype=None, name=None, index=None, skip_infer=None
+    ):
         axis = getattr(self, "axis", None) or 0
         self.axis = validate_axis(axis, df)
         if not skip_infer:
-            dtypes = self._infer_df_func_returns(df, dtypes)
+            inferred_meta = self._infer_df_func_returns(
+                df, dtypes=dtypes, dtype=dtype, name=name, index=index
+            )
+        else:
+            index_value = parse_index(index) if index else df.index_value
+            inferred_meta = InferredDataFrameMeta(
+                self.output_types[0],
+                dtypes=dtypes,
+                dtype=dtype,
+                name=name,
+                index_value=index_value,
+            )
 
+        self._output_types = [inferred_meta.output_type]
         if self.output_types[0] == OutputType.dataframe:
             new_shape = list(df.shape)
-            new_index_value = df.index_value
+            new_index_value = inferred_meta.index_value
+            dtypes = inferred_meta.dtypes
             if len(new_shape) == 1:
                 new_shape.append(len(dtypes) if dtypes is not None else np.nan)
             else:
@@ -118,6 +111,7 @@ class TransformOperator(DataFrameOperator, DataFrameOperatorMixin):
             if self.call_agg:
                 new_shape[self.axis] = np.nan
                 new_index_value = parse_index(None, (df.key, df.index_value.key))
+
             if dtypes is None:
                 columns_value = None
             else:
@@ -130,11 +124,6 @@ class TransformOperator(DataFrameOperator, DataFrameOperatorMixin):
                 columns_value=columns_value,
             )
         else:
-            if dtypes is not None:
-                name, dtype = dtypes
-            else:
-                name, dtype = None, None
-
             if isinstance(df, DATAFRAME_TYPE):
                 new_shape = (df.shape[1 - axis],)
                 new_index_value = [df.columns_value, df.index_value][axis]
@@ -145,15 +134,25 @@ class TransformOperator(DataFrameOperator, DataFrameOperatorMixin):
             return self.new_series(
                 [df],
                 shape=new_shape,
-                name=name,
-                dtype=dtype,
+                name=inferred_meta.name,
+                dtype=inferred_meta.dtype,
                 index_value=new_index_value,
             )
 
+    @classmethod
+    def estimate_size(
+        cls, ctx: MutableMapping[str, Union[int, float]], op: "TransformOperator"
+    ) -> None:
+        if isinstance(op.func, MarkedFunction):
+            ctx[op.outputs[0].key] = float("inf")
+        super().estimate_size(ctx, op)
+
 
 def get_packed_funcs(df, output_type, func, *args, **kwds) -> Any:
     stub_df = _build_stub_pandas_obj(df, output_type)
-    return pack_func_args(stub_df, func, *args, **kwds)
+    n_args = copy_if_possible(args)
+    n_kwds = copy_if_possible(kwds)
+    return pack_func_args(stub_df, func, *n_args, **n_kwds)
 
 
 def _build_stub_pandas_obj(df, output_type) -> Union[DataFrame, Series]:
@@ -241,7 +240,7 @@ def df_transform(df, func, axis=0, *args, dtypes=None, skip_infer=False, **kwarg
         axis=axis,
         args=args,
         kwds=kwargs,
-        output_types=[OutputType.dataframe],
+        output_types=[OutputType.dataframe] if not call_agg else None,
         call_agg=call_agg,
     )
     return op(df, dtypes=dtypes, skip_infer=skip_infer)
@@ -337,5 +336,4 @@ def series_transform(
         output_types=[OutputType.series],
         call_agg=call_agg,
     )
-    dtypes = (series.name, dtype) if dtype is not None else None
-    return op(series, dtypes=dtypes, skip_infer=skip_infer)
+    return op(series, dtype=dtype, name=series.name, skip_infer=skip_infer)

maxframe/dataframe/misc/value_counts.py

@@ -12,11 +12,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import List
+
 import numpy as np
 import pandas as pd
 
 from ... import opcodes
-from ...core import OutputType
+from ...core import EntityData, OutputType
 from ...serialization.serializables import BoolField, Int64Field, KeyField, StringField
 from ...utils import pd_release_version
 from ..operators import DataFrameOperator, DataFrameOperatorMixin
@@ -42,9 +44,10 @@ class DataFrameValueCounts(DataFrameOperator, DataFrameOperatorMixin):
         super().__init__(**kw)
         self.output_types = [OutputType.series]
 
-    def _set_inputs(self, inputs):
-        super()._set_inputs(inputs)
-        self.input = self._inputs[0]
+    @classmethod
+    def _set_inputs(cls, op: "DataFrameValueCounts", inputs: List[EntityData]):
+        super()._set_inputs(op, inputs)
+        op.input = op._inputs[0]
 
     def __call__(self, inp):
         test_series = build_series(inp).value_counts(normalize=self.normalize)

maxframe/dataframe/missing/dropna.py

@@ -37,6 +37,8 @@ class DataFrameDropNA(DataFrameOperator, DataFrameOperatorMixin):
     drop_directly = BoolField("drop_directly", default=None)
     # size of subset, used when how == 'any'
     subset_size = Int32Field("subset_size", default=None)
+    # if True, drop index
+    ignore_index = BoolField("ignore_index", default=False)
 
     def __init__(self, sparse=None, output_types=None, **kw):
         super().__init__(_output_types=output_types, sparse=sparse, **kw)
@@ -52,7 +54,13 @@ class DataFrameDropNA(DataFrameOperator, DataFrameOperatorMixin):
 
 
 def df_dropna(
-    df, axis=0, how=no_default, thresh=no_default, subset=None, inplace=False
+    df,
+    axis=0,
+    how=no_default,
+    thresh=no_default,
+    subset=None,
+    inplace=False,
+    ignore_index=False,
 ):
     """
     Remove missing values.
@@ -69,11 +77,6 @@ def df_dropna(
         * 0, or 'index' : Drop rows which contain missing values.
         * 1, or 'columns' : Drop columns which contain missing value.
 
-        .. versionchanged:: 1.0.0
-
-           Pass tuple or list to drop on multiple axes.
-           Only a single axis is allowed.
-
     how : {'any', 'all'}, default 'any'
         Determine if row or column is removed from DataFrame, when we have
         at least one NA or all NA.
@@ -88,6 +91,8 @@ def df_dropna(
         these would be a list of columns to include.
     inplace : bool, default False
         If True, do operation inplace and return None.
+    ignore_index : bool, default False
+        If True, the resulting axis will be labeled 0, 1, …, n - 1.
 
     Returns
     -------
@@ -168,6 +173,7 @@ def df_dropna(
         how=how,
         thresh=thresh,
         subset=subset,
+        ignore_index=ignore_index,
         output_types=[OutputType.dataframe],
     )
     out_df = op(df)
@@ -177,7 +183,7 @@ def df_dropna(
         return out_df
 
 
-def series_dropna(series, axis=0, inplace=False, how=None):
+def series_dropna(series, axis=0, inplace=False, how=None, ignore_index=False):
     """
     Return a new Series with missing values removed.
 
@@ -192,6 +198,8 @@ def series_dropna(series, axis=0, inplace=False, how=None):
         If True, do operation inplace and return None.
     how : str, optional
         Not in use. Kept for compatibility.
+    ignore_index : bool, default False
+        If True, the resulting axis will be labeled 0, 1, …, n - 1.
 
     Returns
     -------
@@ -253,6 +261,7 @@ def series_dropna(series, axis=0, inplace=False, how=None):
     op = DataFrameDropNA(
         axis=axis,
         how=how,
+        ignore_index=ignore_index,
         output_types=[OutputType.series],
     )
     out_series = op(series)

maxframe/dataframe/missing/fillna.py

@@ -12,18 +12,21 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import List
+
 import pandas as pd
 
 from ... import opcodes
-from ...core import ENTITY_TYPE, Entity, get_output_types
+from ...core import ENTITY_TYPE, Entity, EntityData, get_output_types
 from ...serialization.serializables import AnyField, Int64Field, StringField
 from ..core import DATAFRAME_TYPE, SERIES_TYPE
 from ..operators import DataFrameOperator, DataFrameOperatorMixin
 from ..utils import validate_axis
 
 
-class FillNA(DataFrameOperator, DataFrameOperatorMixin):
+class DataFrameFillNA(DataFrameOperator, DataFrameOperatorMixin):
     _op_type_ = opcodes.FILL_NA
+    _legacy_name = "FillNA"
 
     value = AnyField(
         "value", on_serialize=lambda x: x.data if isinstance(x, Entity) else x
@@ -33,13 +36,14 @@ class FillNA(DataFrameOperator, DataFrameOperatorMixin):
     limit = Int64Field("limit", default=None)
     downcast = AnyField("downcast", default=None)
 
-    def __init__(self, output_limit=1, **kw):
-        super().__init__(output_limit=output_limit, **kw)
+    def __init__(self, output_limit=1, output_types=None, **kw):
+        super().__init__(output_limit=output_limit, _output_types=output_types, **kw)
 
-    def _set_inputs(self, inputs):
-        super()._set_inputs(inputs)
-        if self.method is None and len(inputs) > 1:
-            self.value = self._inputs[1]
+    @classmethod
+    def _set_inputs(cls, op: "DataFrameFillNA", inputs: List[EntityData]):
+        super()._set_inputs(op, inputs)
+        if op.method is None and len(inputs) > 1:
+            op.value = op._inputs[1]
 
     def __call__(self, a, value_df=None):
         method = getattr(self, "method", None)
@@ -81,6 +85,10 @@ class FillNA(DataFrameOperator, DataFrameOperatorMixin):
             )
 
 
+# keep for import compatibility
+FillNA = DataFrameFillNA
+
+
 def fillna(
     df, value=None, method=None, axis=None, inplace=False, limit=None, downcast=None
 ):
@@ -198,7 +206,7 @@ def fillna(
     else:
         value_df = None
 
-    op = FillNA(
+    op = DataFrameFillNA(
         value=value,
         method=method,
         axis=axis,
@@ -267,7 +275,7 @@ def index_fillna(index, value=None, downcast=None):
     if isinstance(value, (list, pd.Series, SERIES_TYPE)):
         raise ValueError("'value' must be a scalar, passed: %s" % type(value))
 
-    op = FillNA(
+    op = DataFrameFillNA(
         value=value,
         downcast=downcast,
         output_types=get_output_types(index),