test_util.py 2.8 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879
  1. import numpy as np
  2. import pytest
  3. from pandas import (
  4. Index,
  5. date_range,
  6. )
  7. import pandas._testing as tm
  8. from pandas.core.reshape.util import cartesian_product
  9. class TestCartesianProduct:
  10. def test_simple(self):
  11. x, y = list("ABC"), [1, 22]
  12. result1, result2 = cartesian_product([x, y])
  13. expected1 = np.array(["A", "A", "B", "B", "C", "C"])
  14. expected2 = np.array([1, 22, 1, 22, 1, 22])
  15. tm.assert_numpy_array_equal(result1, expected1)
  16. tm.assert_numpy_array_equal(result2, expected2)
  17. def test_datetimeindex(self):
  18. # regression test for GitHub issue #6439
  19. # make sure that the ordering on datetimeindex is consistent
  20. x = date_range("2000-01-01", periods=2)
  21. result1, result2 = (Index(y).day for y in cartesian_product([x, x]))
  22. expected1 = Index([1, 1, 2, 2], dtype=np.int32)
  23. expected2 = Index([1, 2, 1, 2], dtype=np.int32)
  24. tm.assert_index_equal(result1, expected1)
  25. tm.assert_index_equal(result2, expected2)
  26. def test_tzaware_retained(self):
  27. x = date_range("2000-01-01", periods=2, tz="US/Pacific")
  28. y = np.array([3, 4])
  29. result1, result2 = cartesian_product([x, y])
  30. expected = x.repeat(2)
  31. tm.assert_index_equal(result1, expected)
  32. def test_tzaware_retained_categorical(self):
  33. x = date_range("2000-01-01", periods=2, tz="US/Pacific").astype("category")
  34. y = np.array([3, 4])
  35. result1, result2 = cartesian_product([x, y])
  36. expected = x.repeat(2)
  37. tm.assert_index_equal(result1, expected)
  38. @pytest.mark.parametrize("x, y", [[[], []], [[0, 1], []], [[], ["a", "b", "c"]]])
  39. def test_empty(self, x, y):
  40. # product of empty factors
  41. expected1 = np.array([], dtype=np.asarray(x).dtype)
  42. expected2 = np.array([], dtype=np.asarray(y).dtype)
  43. result1, result2 = cartesian_product([x, y])
  44. tm.assert_numpy_array_equal(result1, expected1)
  45. tm.assert_numpy_array_equal(result2, expected2)
  46. def test_empty_input(self):
  47. # empty product (empty input):
  48. result = cartesian_product([])
  49. expected = []
  50. assert result == expected
  51. @pytest.mark.parametrize(
  52. "X", [1, [1], [1, 2], [[1], 2], "a", ["a"], ["a", "b"], [["a"], "b"]]
  53. )
  54. def test_invalid_input(self, X):
  55. msg = "Input must be a list-like of list-likes"
  56. with pytest.raises(TypeError, match=msg):
  57. cartesian_product(X=X)
  58. def test_exceed_product_space(self):
  59. # GH31355: raise useful error when produce space is too large
  60. msg = "Product space too large to allocate arrays!"
  61. with pytest.raises(ValueError, match=msg):
  62. dims = [np.arange(0, 22, dtype=np.int16) for i in range(12)] + [
  63. (np.arange(15128, dtype=np.int16)),
  64. ]
  65. cartesian_product(X=dims)