| 1915 | |
| 1916 | |
| 1917 | def test_table_unify_dictionaries(): |
| 1918 | batch1 = pa.record_batch([ |
| 1919 | pa.array(["foo", "bar", None, "foo"]).dictionary_encode(), |
| 1920 | pa.array([123, 456, 456, 789]).dictionary_encode(), |
| 1921 | pa.array([True, False, None, None])], names=['a', 'b', 'c']) |
| 1922 | batch2 = pa.record_batch([ |
| 1923 | pa.array(["quux", "foo", None, "quux"]).dictionary_encode(), |
| 1924 | pa.array([456, 789, 789, None]).dictionary_encode(), |
| 1925 | pa.array([False, None, None, True])], names=['a', 'b', 'c']) |
| 1926 | |
| 1927 | table = pa.Table.from_batches([batch1, batch2]) |
| 1928 | table = table.replace_schema_metadata({b"key1": b"value1"}) |
| 1929 | assert table.column(0).chunk(0).dictionary.equals( |
| 1930 | pa.array(["foo", "bar"])) |
| 1931 | assert table.column(0).chunk(1).dictionary.equals( |
| 1932 | pa.array(["quux", "foo"])) |
| 1933 | assert table.column(1).chunk(0).dictionary.equals( |
| 1934 | pa.array([123, 456, 789])) |
| 1935 | assert table.column(1).chunk(1).dictionary.equals( |
| 1936 | pa.array([456, 789])) |
| 1937 | |
| 1938 | table = table.unify_dictionaries(pa.default_memory_pool()) |
| 1939 | expected_dict_0 = pa.array(["foo", "bar", "quux"]) |
| 1940 | expected_dict_1 = pa.array([123, 456, 789]) |
| 1941 | assert table.column(0).chunk(0).dictionary.equals(expected_dict_0) |
| 1942 | assert table.column(0).chunk(1).dictionary.equals(expected_dict_0) |
| 1943 | assert table.column(1).chunk(0).dictionary.equals(expected_dict_1) |
| 1944 | assert table.column(1).chunk(1).dictionary.equals(expected_dict_1) |
| 1945 | |
| 1946 | assert table.to_pydict() == { |
| 1947 | 'a': ["foo", "bar", None, "foo", "quux", "foo", None, "quux"], |
| 1948 | 'b': [123, 456, 456, 789, 456, 789, 789, None], |
| 1949 | 'c': [True, False, None, None, False, None, None, True], |
| 1950 | } |
| 1951 | assert table.schema.metadata == {b"key1": b"value1"} |
| 1952 | |
| 1953 | |
| 1954 | def test_table_maps_as_pydicts(): |