test_combine_concat.py 2.6 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162
  1. import numpy as np
  2. import pytest
  3. import pandas as pd
  4. import pandas._testing as tm
  5. from pandas.core.arrays.sparse import SparseArray
  6. class TestSparseArrayConcat:
  7. @pytest.mark.parametrize("kind", ["integer", "block"])
  8. def test_basic(self, kind):
  9. a = SparseArray([1, 0, 0, 2], kind=kind)
  10. b = SparseArray([1, 0, 2, 2], kind=kind)
  11. result = SparseArray._concat_same_type([a, b])
  12. # Can't make any assertions about the sparse index itself
  13. # since we aren't don't merge sparse blocs across arrays
  14. # in to_concat
  15. expected = np.array([1, 2, 1, 2, 2], dtype="int64")
  16. tm.assert_numpy_array_equal(result.sp_values, expected)
  17. assert result.kind == kind
  18. @pytest.mark.parametrize("kind", ["integer", "block"])
  19. def test_uses_first_kind(self, kind):
  20. other = "integer" if kind == "block" else "block"
  21. a = SparseArray([1, 0, 0, 2], kind=kind)
  22. b = SparseArray([1, 0, 2, 2], kind=other)
  23. result = SparseArray._concat_same_type([a, b])
  24. expected = np.array([1, 2, 1, 2, 2], dtype="int64")
  25. tm.assert_numpy_array_equal(result.sp_values, expected)
  26. assert result.kind == kind
  27. @pytest.mark.parametrize(
  28. "other, expected_dtype",
  29. [
  30. # compatible dtype -> preserve sparse
  31. (pd.Series([3, 4, 5], dtype="int64"), pd.SparseDtype("int64", 0)),
  32. # (pd.Series([3, 4, 5], dtype="Int64"), pd.SparseDtype("int64", 0)),
  33. # incompatible dtype -> Sparse[common dtype]
  34. (pd.Series([1.5, 2.5, 3.5], dtype="float64"), pd.SparseDtype("float64", 0)),
  35. # incompatible dtype -> Sparse[object] dtype
  36. (pd.Series(["a", "b", "c"], dtype=object), pd.SparseDtype(object, 0)),
  37. # categorical with compatible categories -> dtype of the categories
  38. (pd.Series([3, 4, 5], dtype="category"), np.dtype("int64")),
  39. (pd.Series([1.5, 2.5, 3.5], dtype="category"), np.dtype("float64")),
  40. # categorical with incompatible categories -> object dtype
  41. (pd.Series(["a", "b", "c"], dtype="category"), np.dtype(object)),
  42. ],
  43. )
  44. def test_concat_with_non_sparse(other, expected_dtype):
  45. # https://github.com/pandas-dev/pandas/issues/34336
  46. s_sparse = pd.Series([1, 0, 2], dtype=pd.SparseDtype("int64", 0))
  47. result = pd.concat([s_sparse, other], ignore_index=True)
  48. expected = pd.Series(list(s_sparse) + list(other)).astype(expected_dtype)
  49. tm.assert_series_equal(result, expected)
  50. result = pd.concat([other, s_sparse], ignore_index=True)
  51. expected = pd.Series(list(other) + list(s_sparse)).astype(expected_dtype)
  52. tm.assert_series_equal(result, expected)