[REVIEW] Refactor memory_usage to improve performance

[galipremsagar]

Refactored Frame.memory_usage to return a tuple of lists: (column_names, memory_usages).
Motivation for this change is to remove redundent steps i.e., dict(Frame.memory_usage)->unpack & dict (DataFrame.memory_usage)->unpack dict(in Series.init). Choosing to remove dict being returned by Frame.memory_usage will now result in a 5-10% faster execution of external API. Not a huge speedup but dask memory usage calculations that go through
sizeof_cudf_dataframe are now 200x faster.it quickly adds up when dask_cudf is used.

In [1]: import cudf

In [2]: df = cudf.DataFrame({'a':[1, 2, 3], 'b':['a', 'b', 'c'], 'd':[111, 123, 123]})

# THIS PR
In [3]: %timeit df.memory_usage()
198 µs ± 3.44 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)


# branch-22.06
In [3]: %timeit df.memory_usage()
219 µs ± 726 ns per loop (mean ± std. dev. of 7 runs, 1,000 loops each)


# branch-22.06
In [4]: %timeit dask_cudf.backends.sizeof_cudf_dataframe(df)
377 µs ± 5.67 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

# this PR
In [6]: %timeit dask_cudf.backends.sizeof_cudf_dataframe(df)
1.8 µs ± 14 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each)

[codecov]

Codecov Report

❗ No coverage uploaded for pull request base (branch-22.06@2d8d913). Click here to learn what that means.
The diff coverage is n/a.

@@               Coverage Diff               @@
##             branch-22.06   #10537   +/-   ##
===============================================
  Coverage                ?   86.33%           
===============================================
  Files                   ?      140           
  Lines                   ?    22300           
  Branches                ?        0           
===============================================
  Hits                    ?    19252           
  Misses                  ?     3048           
  Partials                ?        0           

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 2d8d913...87589b1. Read the comment docs.

[vyasr]

If this optimization is really noticeable, then I would suggest the following additional changes:

• Remove all inheritance logic from this method, since the amount of shared logic is going to be minimal. Just implement it separately in each of the 4 leaf classes: DataFrame, Series, Index, and MultiIndex.
• Don't even bother constructing the list of column names for anything except DataFrame. You don't need it.
• MultiIndex definition: sum([col.memory_usage for col in self._data.columns])
• GenericIndex definition: self._column.memory_usage
• Series implementation: self._column.memory_usage + (self._index.memory_usage if index else 0)
• DataFrame implementation: Basically what currently exists except inlining the relevant logic from the current Frame and IndexedFrame implementations.

[galipremsagar]

@vyasr Updated the implementations accordingly but kept the inheritances in place for the sake of docstrings. If you want me to remove that as well and duplicate the docstrings in the actual methods let me know, I'm fine with either of the approaches.

[bdice]

LGTM. One minor suggestion.

[vyasr]

Thanks for that catch @bdice.

@galipremsagar could you post new benchmarks so that we know what we're getting out of this?

[galipremsagar]

@galipremsagar could you post new benchmarks so that we know what we're getting out of this?

In [1]: import cudf

In [2]: df = cudf.DataFrame({'a':[1, 2, 3], 'b':['a', 'b', 'c'], 'd':[111, 123, 123]})

In [3]: df = df.set_index(['a', 'd'])

In [4]: x = df.index

# branch-22.06

In [5]: %timeit x.memory_usage()
12 µs ± 48.4 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)

In [6]: %timeit df.memory_usage()
232 µs ± 1.55 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)



# Now
In [7]: %timeit x.memory_usage()
10.3 µs ± 70.2 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)


In [6]: %timeit df.memory_usage()
217 µs ± 496 ns per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

[vyasr]

I'm proposing some alternative solutions offline since I don't think this is quite addressing the underlying issue.

[galipremsagar]

Discussed offline with @vyasr and decided to optimize sizeof_cudf_dataframe:

In [4]: df = cudf.DataFrame({'a':[1, 2, 3], 'b':['a', 'b', 'c'], 'd':[111, 123, 123]})

# branch-22.06
In [4]: %timeit dask_cudf.backends.sizeof_cudf_dataframe(df)
377 µs ± 5.67 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

# this PR
In [6]: %timeit dask_cudf.backends.sizeof_cudf_dataframe(df)
1.8 µs ± 14 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each)

[vyasr]

Awesome updating the registered Dask function should hopefully have a much bigger impact on Dask workloads than the relatively minor optimizations that we can make to memory_usage.

[galipremsagar]

rerun tests

[jakirkham]

Had a couple questions below. Likely just missing context

[jakirkham]

Why is this not implemented now?

[galipremsagar]

We have moved all their implementations to their respective Series/DataFrame/MultiIndex/Index .memory_usage methods directly. Kept these parent class methods for the sake of docstrings.

[jakirkham]

Ok so a user is unlikely to run into this error since they would be using one of these other public facing objects. Is that right?

[galipremsagar]

Yup, that's correct.

[galipremsagar]

@gpucibot merge