import numpy as np import pytest from pandas import ( Categorical, DataFrame, Series, _testing as tm, concat, read_hdf, ) from pandas.tests.io.pytables.common import ( _maybe_remove, ensure_clean_store, ) pytestmark = pytest.mark.single_cpu def test_categorical(setup_path): with ensure_clean_store(setup_path) as store: # Basic _maybe_remove(store, "s") s = Series( Categorical( ["a", "b", "b", "a", "a", "c"], categories=["a", "b", "c", "d"], ordered=False, ) ) store.append("s", s, format="table") result = store.select("s") tm.assert_series_equal(s, result) _maybe_remove(store, "s_ordered") s = Series( Categorical( ["a", "b", "b", "a", "a", "c"], categories=["a", "b", "c", "d"], ordered=True, ) ) store.append("s_ordered", s, format="table") result = store.select("s_ordered") tm.assert_series_equal(s, result) _maybe_remove(store, "df") df = DataFrame({"s": s, "vals": [1, 2, 3, 4, 5, 6]}) store.append("df", df, format="table") result = store.select("df") tm.assert_frame_equal(result, df) # Dtypes _maybe_remove(store, "si") s = Series([1, 1, 2, 2, 3, 4, 5]).astype("category") store.append("si", s) result = store.select("si") tm.assert_series_equal(result, s) _maybe_remove(store, "si2") s = Series([1, 1, np.nan, 2, 3, 4, 5]).astype("category") store.append("si2", s) result = store.select("si2") tm.assert_series_equal(result, s) # Multiple _maybe_remove(store, "df2") df2 = df.copy() df2["s2"] = Series(list("abcdefg")).astype("category") store.append("df2", df2) result = store.select("df2") tm.assert_frame_equal(result, df2) # Make sure the metadata is OK info = store.info() assert "/df2 " in info # df2._mgr.blocks[0] and df2._mgr.blocks[2] are Categorical assert "/df2/meta/values_block_0/meta" in info assert "/df2/meta/values_block_2/meta" in info # unordered _maybe_remove(store, "s2") s = Series( Categorical( ["a", "b", "b", "a", "a", "c"], categories=["a", "b", "c", "d"], ordered=False, ) ) store.append("s2", s, format="table") result = store.select("s2") tm.assert_series_equal(result, s) # Query _maybe_remove(store, "df3") store.append("df3", df, data_columns=["s"]) expected = df[df.s.isin(["b", "c"])] result = store.select("df3", where=['s in ["b","c"]']) tm.assert_frame_equal(result, expected) expected = df[df.s.isin(["b", "c"])] result = store.select("df3", where=['s = ["b","c"]']) tm.assert_frame_equal(result, expected) expected = df[df.s.isin(["d"])] result = store.select("df3", where=['s in ["d"]']) tm.assert_frame_equal(result, expected) expected = df[df.s.isin(["f"])] result = store.select("df3", where=['s in ["f"]']) tm.assert_frame_equal(result, expected) # Appending with same categories is ok store.append("df3", df) df = concat([df, df]) expected = df[df.s.isin(["b", "c"])] result = store.select("df3", where=['s in ["b","c"]']) tm.assert_frame_equal(result, expected) # Appending must have the same categories df3 = df.copy() df3["s"] = df3["s"].cat.remove_unused_categories() msg = "cannot append a categorical with different categories to the existing" with pytest.raises(ValueError, match=msg): store.append("df3", df3) # Remove, and make sure meta data is removed (its a recursive # removal so should be). result = store.select("df3/meta/s/meta") assert result is not None store.remove("df3") with pytest.raises( KeyError, match="'No object named df3/meta/s/meta in the file'" ): store.select("df3/meta/s/meta") def test_categorical_conversion(tmp_path, setup_path): # GH13322 # Check that read_hdf with categorical columns doesn't return rows if # where criteria isn't met. obsids = ["ESP_012345_6789", "ESP_987654_3210"] imgids = ["APF00006np", "APF0001imm"] data = [4.3, 9.8] # Test without categories df = DataFrame({"obsids": obsids, "imgids": imgids, "data": data}) # We are expecting an empty DataFrame matching types of df expected = df.iloc[[], :] path = tmp_path / setup_path df.to_hdf(path, key="df", format="table", data_columns=True) result = read_hdf(path, "df", where="obsids=B") tm.assert_frame_equal(result, expected) # Test with categories df.obsids = df.obsids.astype("category") df.imgids = df.imgids.astype("category") # We are expecting an empty DataFrame matching types of df expected = df.iloc[[], :] path = tmp_path / setup_path df.to_hdf(path, key="df", format="table", data_columns=True) result = read_hdf(path, "df", where="obsids=B") tm.assert_frame_equal(result, expected) def test_categorical_nan_only_columns(tmp_path, setup_path): # GH18413 # Check that read_hdf with categorical columns with NaN-only values can # be read back. df = DataFrame( { "a": ["a", "b", "c", np.nan], "b": [np.nan, np.nan, np.nan, np.nan], "c": [1, 2, 3, 4], "d": Series([None] * 4, dtype=object), } ) df["a"] = df.a.astype("category") df["b"] = df.b.astype("category") df["d"] = df.b.astype("category") expected = df path = tmp_path / setup_path df.to_hdf(path, key="df", format="table", data_columns=True) result = read_hdf(path, "df") tm.assert_frame_equal(result, expected) @pytest.mark.parametrize( "where, df, expected", [ ('col=="q"', DataFrame({"col": ["a", "b", "s"]}), DataFrame({"col": []})), ('col=="a"', DataFrame({"col": ["a", "b", "s"]}), DataFrame({"col": ["a"]})), ], ) def test_convert_value( tmp_path, setup_path, where: str, df: DataFrame, expected: DataFrame ): # GH39420 # Check that read_hdf with categorical columns can filter by where condition. df.col = df.col.astype("category") max_widths = {"col": 1} categorical_values = sorted(df.col.unique()) expected.col = expected.col.astype("category") expected.col = expected.col.cat.set_categories(categorical_values) path = tmp_path / setup_path df.to_hdf(path, key="df", format="table", min_itemsize=max_widths) result = read_hdf(path, where=where) tm.assert_frame_equal(result, expected)