test_concat.py 2.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869
  1. import numpy as np
  2. import pytest
  3. import pandas as pd
  4. import pandas._testing as tm
  5. @pytest.mark.parametrize(
  6. "to_concat_dtypes, result_dtype",
  7. [
  8. (["Int64", "Int64"], "Int64"),
  9. (["UInt64", "UInt64"], "UInt64"),
  10. (["Int8", "Int8"], "Int8"),
  11. (["Int8", "Int16"], "Int16"),
  12. (["UInt8", "Int8"], "Int16"),
  13. (["Int32", "UInt32"], "Int64"),
  14. (["Int64", "UInt64"], "Float64"),
  15. (["Int64", "boolean"], "object"),
  16. (["UInt8", "boolean"], "object"),
  17. ],
  18. )
  19. def test_concat_series(to_concat_dtypes, result_dtype):
  20. # we expect the same dtypes as we would get with non-masked inputs,
  21. # just masked where available.
  22. result = pd.concat([pd.Series([0, 1, pd.NA], dtype=t) for t in to_concat_dtypes])
  23. expected = pd.concat([pd.Series([0, 1, pd.NA], dtype=object)] * 2).astype(
  24. result_dtype
  25. )
  26. tm.assert_series_equal(result, expected)
  27. # order doesn't matter for result
  28. result = pd.concat(
  29. [pd.Series([0, 1, pd.NA], dtype=t) for t in to_concat_dtypes[::-1]]
  30. )
  31. expected = pd.concat([pd.Series([0, 1, pd.NA], dtype=object)] * 2).astype(
  32. result_dtype
  33. )
  34. tm.assert_series_equal(result, expected)
  35. @pytest.mark.parametrize(
  36. "to_concat_dtypes, result_dtype",
  37. [
  38. (["Int64", "int64"], "Int64"),
  39. (["UInt64", "uint64"], "UInt64"),
  40. (["Int8", "int8"], "Int8"),
  41. (["Int8", "int16"], "Int16"),
  42. (["UInt8", "int8"], "Int16"),
  43. (["Int32", "uint32"], "Int64"),
  44. (["Int64", "uint64"], "Float64"),
  45. (["Int64", "bool"], "object"),
  46. (["UInt8", "bool"], "object"),
  47. ],
  48. )
  49. def test_concat_series_with_numpy(to_concat_dtypes, result_dtype):
  50. # we expect the same dtypes as we would get with non-masked inputs,
  51. # just masked where available.
  52. s1 = pd.Series([0, 1, pd.NA], dtype=to_concat_dtypes[0])
  53. s2 = pd.Series(np.array([0, 1], dtype=to_concat_dtypes[1]))
  54. result = pd.concat([s1, s2], ignore_index=True)
  55. expected = pd.Series([0, 1, pd.NA, 0, 1], dtype=object).astype(result_dtype)
  56. tm.assert_series_equal(result, expected)
  57. # order doesn't matter for result
  58. result = pd.concat([s2, s1], ignore_index=True)
  59. expected = pd.Series([0, 1, 0, 1, pd.NA], dtype=object).astype(result_dtype)
  60. tm.assert_series_equal(result, expected)