diff --git a/CHANGELOG.md b/CHANGELOG.md index e464228b9..a76113f1a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,8 @@ # Changelog ## [Unreleased] +- [PERF] Speed up balanced homogeneous NumPy `pivot_longer` reshapes by + slicing value groups positionally. - Issue #1656, PR #1668 @samukweku - [ENH] Avoid copying column data during `conditional_join` input validation. - Issue #1645, PR #1642 @samukweku - [ENH] Speed up `conditional_join` with an unsorted right join key and diff --git a/janitor/functions/pivot.py b/janitor/functions/pivot.py index 74ef094d6..d8c649c4d 100644 --- a/janitor/functions/pivot.py +++ b/janitor/functions/pivot.py @@ -1772,27 +1772,45 @@ def _stack_dot_value_multiple_labels( df = df.sort_index(axis=1) contents = {} indexer = None - for label in range(_value.size): - frame = df.loc[:, [label]] - any_extension_array = ( - frame.dtypes.map(is_extension_array_dtype).any(axis=None).item() - ) - if sort_by_appearance and any_extension_array: - frame = _build_content_extension_array(df=frame) - if indexer is None: - indexer = _build_indexer_reorder_contents( - length=len_df, - reps=reps, - ) - frame = frame.take(indexer) - elif any_extension_array: - frame = _build_content_extension_array(df=frame) - elif sort_by_appearance: - frame = frame._values.ravel(order="C") - else: - frame = frame._values.ravel(order="F") - label = _value[label] - contents[label] = frame + balanced = _index.size == max_count + # A balanced single-block NumPy frame is already laid out in contiguous + # value groups, so positional slices avoid repeated DataFrame selection + # and dtype inspection. Keep fragmented and extension-backed frames on + # the dtype-preserving path below. + manager = df._mgr + homogeneous_numpy = balanced and manager.is_single_block + homogeneous_numpy = homogeneous_numpy and not is_extension_array_dtype( + manager.blocks[0].dtype + ) + if homogeneous_numpy: + values = df._values + order = "C" if sort_by_appearance else "F" + for position, label in enumerate(_value): + start = position * reps + stop = start + reps + contents[label] = values[:, start:stop].ravel(order=order) + else: + for label in range(_value.size): + frame = df.loc[:, [label]] + any_extension_array = ( + frame.dtypes.map(is_extension_array_dtype).any(axis=None).item() + ) + if sort_by_appearance and any_extension_array: + frame = _build_content_extension_array(df=frame) + if indexer is None: + indexer = _build_indexer_reorder_contents( + length=len_df, + reps=reps, + ) + frame = frame.take(indexer) + elif any_extension_array: + frame = _build_content_extension_array(df=frame) + elif sort_by_appearance: + frame = frame._values.ravel(order="C") + else: + frame = frame._values.ravel(order="F") + label = _value[label] + contents[label] = frame nulls = _build_nulls(contents=contents, dropna=dropna) index = _build_index( index=index, diff --git a/tests/functions/test_pivot_longer.py b/tests/functions/test_pivot_longer.py index 9741242dc..bf0edfd00 100644 --- a/tests/functions/test_pivot_longer.py +++ b/tests/functions/test_pivot_longer.py @@ -1753,3 +1753,42 @@ def test_dropna_sort_by_appearance(): ) assert_frame_equal(actual, expected) + + +@pytest.mark.parametrize( + "dtype", + ["int64", "float64", "bool", "datetime64[ns]", "timedelta64[ns]"], +) +@pytest.mark.parametrize("sort_by_appearance", [False, True]) +def test_balanced_homogeneous_numpy_multiple_values(dtype, sort_by_appearance): + """Preserve order and dtype for balanced homogeneous NumPy values.""" + values = np.arange(12).reshape(3, 4) + if dtype == "bool": + values = values % 2 == 0 + elif dtype == "datetime64[ns]": + values = values.astype("timedelta64[D]") + np.datetime64("2020-01-01") + else: + values = values.astype(dtype) + df = pd.DataFrame(values, columns=["x_1", "x_2", "y_1", "y_2"]) + + actual = df.pivot_longer( + names_to=(".value", "position"), + names_sep="_", + sort_by_appearance=sort_by_appearance, + ) + + order = "C" if sort_by_appearance else "F" + positions = np.array(["1", "2"]) + if sort_by_appearance: + positions = np.tile(positions, len(df)) + else: + positions = positions.repeat(len(df)) + expected = pd.DataFrame( + { + "position": positions, + "x": df.loc[:, ["x_1", "x_2"]]._values.ravel(order=order), + "y": df.loc[:, ["y_1", "y_2"]]._values.ravel(order=order), + } + ) + + assert_frame_equal(actual, expected)