| 707 | |
| 708 | |
| 709 | def test_frame_select_complex2(tmp_path): |
| 710 | pp = tmp_path / "params.hdf" |
| 711 | hh = tmp_path / "hist.hdf" |
| 712 | |
| 713 | # use non-trivial selection criteria |
| 714 | params = DataFrame({"A": [1, 1, 2, 2, 3]}) |
| 715 | params.to_hdf(pp, key="df", mode="w", format="table", data_columns=["A"]) |
| 716 | |
| 717 | selection = read_hdf(pp, "df", where="A=[2,3]") |
| 718 | hist = DataFrame( |
| 719 | np.random.default_rng(2).standard_normal((25, 1)), |
| 720 | columns=["data"], |
| 721 | index=MultiIndex.from_tuples( |
| 722 | [(i, j) for i in range(5) for j in range(5)], names=["l1", "l2"] |
| 723 | ), |
| 724 | ) |
| 725 | |
| 726 | hist.to_hdf(hh, key="df", mode="w", format="table") |
| 727 | |
| 728 | expected = read_hdf(hh, "df", where="l1=[2, 3, 4]") |
| 729 | |
| 730 | # scope with list like |
| 731 | l0 = selection.index.tolist() # noqa: F841 |
| 732 | with HDFStore(hh) as store: |
| 733 | result = store.select("df", where="l1=l0") |
| 734 | tm.assert_frame_equal(result, expected) |
| 735 | |
| 736 | result = read_hdf(hh, "df", where="l1=l0") |
| 737 | tm.assert_frame_equal(result, expected) |
| 738 | |
| 739 | # index |
| 740 | index = selection.index # noqa: F841 |
| 741 | result = read_hdf(hh, "df", where="l1=index") |
| 742 | tm.assert_frame_equal(result, expected) |
| 743 | |
| 744 | result = read_hdf(hh, "df", where="l1=selection.index") |
| 745 | tm.assert_frame_equal(result, expected) |
| 746 | |
| 747 | result = read_hdf(hh, "df", where="l1=selection.index.tolist()") |
| 748 | tm.assert_frame_equal(result, expected) |
| 749 | |
| 750 | result = read_hdf(hh, "df", where="l1=list(selection.index)") |
| 751 | tm.assert_frame_equal(result, expected) |
| 752 | |
| 753 | # scope with index |
| 754 | with HDFStore(hh) as store: |
| 755 | result = store.select("df", where="l1=index") |
| 756 | tm.assert_frame_equal(result, expected) |
| 757 | |
| 758 | result = store.select("df", where="l1=selection.index") |
| 759 | tm.assert_frame_equal(result, expected) |
| 760 | |
| 761 | result = store.select("df", where="l1=selection.index.tolist()") |
| 762 | tm.assert_frame_equal(result, expected) |
| 763 | |
| 764 | result = store.select("df", where="l1=list(selection.index)") |
| 765 | tm.assert_frame_equal(result, expected) |
| 766 | |