Concatenate dictionary of objects along axis=1

[er-eis]

Description

Note: This work is heavily based off amanlai's PR raised here, wasn't able to base my branch off amanlai's due to deleted branch.

Closes #15115.
Unlike pandas.concat, cudf.concat doesn't work with a dictionary of objects. The following code raises an error.

d = {
    'first': cudf.DataFrame({'A': [1, 2], 'B': [3, 4]}),
    'second': cudf.DataFrame({'A': [5, 6], 'B': [7, 8]}),
}

cudf.concat(d, axis=1)

This commit resolves this issue.

Checklist

• I am familiar with the Contributing Guidelines.
• New or existing tests cover these changes.
• The documentation is up to date with these changes.
  ^ need me to create an entry in the CHANGELOG.md?

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[er-eis]

to the reviewer(s):

please take a look here 0736e4e#diff-1d997d95893af6a665d5803d179f472c673a7d4e1a0a04a305bd2f1c4a66d957L237

any idea why we were casting to 'float64' previously? tried tracking down why, see my commit message

[bdice]

/okay to test

[bdice]

any idea why we were casting to 'float64' previously?

My best guess is that this a historical artifact and is no longer necessary. I think your change to remove that cast is the right move -- thanks!

I triggered CI and we can see if it passes tests. If it fails, we can investigate further.

[er-eis]

My best guess is that this a historical artifact and is no longer necessary.

it weirds me out that the test has been passing all this time, but i'm a noob around these parts so i trust you!

[wence-]

Thanks for resurrecting this! This code is really fiddly, so great job unpicking things. I have a few suggestions and questions for cleanup.

[wence-]

/ok to test

[wence-]

/ok to test

[wence-]

Sorry, I made a sequence of mess-ups, but hopefully we've got there.

[wence-]

/ok to test

[er-eis]

@wence- do we prefer set().union(*(map(type, obj._data.keys()) for obj in objs)) over {type(name) for o in objs for name in o._data.keys()} ?

{type(name) for o in objs for name in o._data.keys()} is:

• an order of magnitude faster
• imo, more readable/pythonic

you can test the execution time by setting a breakpoint on line 427 and running this:

before = time.perf_counter()
set().union(*(map(type, obj._data.keys()) for obj in objs))
after_changed = time.perf_counter() - before
b = time.perf_counter()
{type(name) for o in objs for name in o._data.keys()}
a_changed = time.perf_counter() - b

[wence-]

@wence- do we prefer set().union(*(map(type, obj._data.keys()) for obj in objs)) over {type(name) for o in objs for name in o._data.keys()} ?

{type(name) for o in objs for name in o._data.keys()} is:

• an order of magnitude faster
• imo, more readable/pythonic

you can test the execution time by setting a breakpoint on line 427 and running this:

before = time.perf_counter()
set().union(*(map(type, obj._data.keys()) for obj in objs))
after_changed = time.perf_counter() - before
b = time.perf_counter()
{type(name) for o in objs for name in o._data.keys()}
a_changed = time.perf_counter() - b

Ah sorry. Too much time writing Haskell. Please go back to your approach.

[er-eis]

sure, making a commit now and merging in latest changes

[wence-]

/ok to test

[bdice]

It's best if we can avoid creating instances of GPU objects (cudf.DataFrame) in the parametrize arguments. Those are executed at test collection time rather than at test runtime, and make the test suite slow to launch due to a large number of small host-device copies. Let's defer construction using a pattern like this:

@pytest.mark.parametrize(
    "d",
    [
        {"first": {"A": [1, 2], "B": [3, 4]}},
        # ...
    ],
)
def test_concat_dictionary(d, axis):
    # Convert dict-of-dicts to dict-of-DataFrames to avoid raw GPU objects in the parameters
    d = {k: cudf.DataFrame(v) for k, v in input.items()}

    result = cudf.concat(d, axis=axis)
    expected = cudf.from_pandas(
        pd.concat({k: df.to_pandas() for k, df in d.items()}, axis=axis)
    )
    assert_eq(expected, result)

[er-eis]

will do! let me know if you'd like me to clean up the other tests in this file in the same manner

[er-eis]

i passed in the reference to each class for the tests, let me know if this causes weirdness during test collection and i'll make a simple map.

[bdice]

Yes, that’s totally fine. No need to refactor the rest of the file in this PR. A separate PR would be welcome if you’re interested.

[bdice]

/ok to test

[er-eis]

@bdice @wence- i think we're good to merge?

[bdice]

/merge

[bdice]

Thanks @er-eis! I think you mentioned you had a follow-up PR planned? Please feel free to open an issue documenting any next steps that are needed, even if you don't have time to contribute a PR.

[er-eis]

@bdice on it!

@bdice and @wence- , thanks for the great reviews, this was fun!