|
1 | 1 | from utils_pytest import * |
2 | 2 |
|
3 | 3 |
|
| 4 | +# show that re-setting the same partition spec |
| 5 | +# doesn't add a new spec, instead re-uses |
| 6 | +def test_re_set_partition_fields( |
| 7 | + extension, s3, with_default_location, pg_conn, superuser_conn |
| 8 | +): |
| 9 | + |
| 10 | + run_command(f"""CREATE SCHEMA test_re_set_partition_fields;""", pg_conn) |
| 11 | + |
| 12 | + run_command( |
| 13 | + "CREATE TABLE test_re_set_partition_fields.tbl(key int, value text) USING iceberg", |
| 14 | + pg_conn, |
| 15 | + ) |
| 16 | + pg_conn.commit() |
| 17 | + |
| 18 | + # when there are no partitions, the default spec-id is used, and there are no corresponding entries |
| 19 | + # in partition_specs |
| 20 | + default_spec_id = run_query( |
| 21 | + "SELECT default_spec_id FROM lake_iceberg.tables_internal WHERE table_name = 'test_re_set_partition_fields.tbl'::regclass", |
| 22 | + superuser_conn, |
| 23 | + ) |
| 24 | + assert default_spec_id[0][0] == 0 |
| 25 | + assert ( |
| 26 | + default_table_partition_spec_id(pg_conn, "test_re_set_partition_fields", "tbl") |
| 27 | + == 0 |
| 28 | + ) |
| 29 | + assert ( |
| 30 | + len(table_partition_specs(pg_conn, "test_re_set_partition_fields", "tbl")) == 1 |
| 31 | + ) |
| 32 | + |
| 33 | + # now, add a spec, and make sure it is pushed |
| 34 | + run_command( |
| 35 | + f"ALTER TABLE test_re_set_partition_fields.tbl OPTIONS (ADD partition_by 'key')", |
| 36 | + pg_conn, |
| 37 | + ) |
| 38 | + pg_conn.commit() |
| 39 | + |
| 40 | + default_spec_id = run_query( |
| 41 | + "SELECT default_spec_id FROM lake_iceberg.tables_internal WHERE table_name = 'test_re_set_partition_fields.tbl'::regclass", |
| 42 | + superuser_conn, |
| 43 | + ) |
| 44 | + assert default_spec_id[0][0] == 1 |
| 45 | + assert ( |
| 46 | + default_table_partition_spec_id(pg_conn, "test_re_set_partition_fields", "tbl") |
| 47 | + == 1 |
| 48 | + ) |
| 49 | + assert ( |
| 50 | + len(table_partition_specs(pg_conn, "test_re_set_partition_fields", "tbl")) == 2 |
| 51 | + ) |
| 52 | + |
| 53 | + # now, drop the spec, and make sure it is pushed |
| 54 | + run_command( |
| 55 | + f"ALTER TABLE test_re_set_partition_fields.tbl OPTIONS (DROP partition_by)", |
| 56 | + pg_conn, |
| 57 | + ) |
| 58 | + pg_conn.commit() |
| 59 | + |
| 60 | + default_spec_id = run_query( |
| 61 | + "SELECT default_spec_id FROM lake_iceberg.tables_internal WHERE table_name = 'test_re_set_partition_fields.tbl'::regclass", |
| 62 | + superuser_conn, |
| 63 | + ) |
| 64 | + assert default_spec_id[0][0] == 0 |
| 65 | + assert ( |
| 66 | + default_table_partition_spec_id(pg_conn, "test_re_set_partition_fields", "tbl") |
| 67 | + == 0 |
| 68 | + ) |
| 69 | + assert ( |
| 70 | + len(table_partition_specs(pg_conn, "test_re_set_partition_fields", "tbl")) == 2 |
| 71 | + ) |
| 72 | + |
| 73 | + # now, add another partition spec |
| 74 | + run_command( |
| 75 | + f"ALTER TABLE test_re_set_partition_fields.tbl OPTIONS (ADD partition_by 'bucket(10, key)')", |
| 76 | + pg_conn, |
| 77 | + ) |
| 78 | + pg_conn.commit() |
| 79 | + |
| 80 | + default_spec_id = run_query( |
| 81 | + "SELECT default_spec_id FROM lake_iceberg.tables_internal WHERE table_name = 'test_re_set_partition_fields.tbl'::regclass", |
| 82 | + superuser_conn, |
| 83 | + ) |
| 84 | + assert default_spec_id[0][0] == 2 |
| 85 | + assert ( |
| 86 | + default_table_partition_spec_id(pg_conn, "test_re_set_partition_fields", "tbl") |
| 87 | + == 2 |
| 88 | + ) |
| 89 | + assert ( |
| 90 | + len(table_partition_specs(pg_conn, "test_re_set_partition_fields", "tbl")) == 3 |
| 91 | + ) |
| 92 | + |
| 93 | + # setting back to existing one should not push a new id |
| 94 | + # now, add a spec, and make sure it is pushed |
| 95 | + run_command( |
| 96 | + f"ALTER TABLE test_re_set_partition_fields.tbl OPTIONS (SET partition_by 'key')", |
| 97 | + pg_conn, |
| 98 | + ) |
| 99 | + pg_conn.commit() |
| 100 | + |
| 101 | + default_spec_id = run_query( |
| 102 | + "SELECT default_spec_id FROM lake_iceberg.tables_internal WHERE table_name = 'test_re_set_partition_fields.tbl'::regclass", |
| 103 | + superuser_conn, |
| 104 | + ) |
| 105 | + assert default_spec_id[0][0] == 1 |
| 106 | + assert ( |
| 107 | + default_table_partition_spec_id(pg_conn, "test_re_set_partition_fields", "tbl") |
| 108 | + == 1 |
| 109 | + ) |
| 110 | + assert ( |
| 111 | + len(table_partition_specs(pg_conn, "test_re_set_partition_fields", "tbl")) == 3 |
| 112 | + ) |
| 113 | + |
| 114 | + # a spec with two fields swapped are considered separate specs |
| 115 | + run_command( |
| 116 | + f"ALTER TABLE test_re_set_partition_fields.tbl OPTIONS (SET partition_by 'truncate(100, key), bucket(50, value)')", |
| 117 | + pg_conn, |
| 118 | + ) |
| 119 | + run_command( |
| 120 | + f"ALTER TABLE test_re_set_partition_fields.tbl OPTIONS (SET partition_by 'bucket(50, value), truncate(100, key)')", |
| 121 | + pg_conn, |
| 122 | + ) |
| 123 | + pg_conn.commit() |
| 124 | + |
| 125 | + default_spec_id = run_query( |
| 126 | + "SELECT default_spec_id FROM lake_iceberg.tables_internal WHERE table_name = 'test_re_set_partition_fields.tbl'::regclass", |
| 127 | + superuser_conn, |
| 128 | + ) |
| 129 | + assert default_spec_id[0][0] == 4 |
| 130 | + assert ( |
| 131 | + default_table_partition_spec_id(pg_conn, "test_re_set_partition_fields", "tbl") |
| 132 | + == 4 |
| 133 | + ) |
| 134 | + assert ( |
| 135 | + len(table_partition_specs(pg_conn, "test_re_set_partition_fields", "tbl")) == 5 |
| 136 | + ) |
| 137 | + |
| 138 | + run_command( |
| 139 | + f"ALTER TABLE test_re_set_partition_fields.tbl OPTIONS (SET partition_by 'truncate(100, key), bucket(50, value)')", |
| 140 | + pg_conn, |
| 141 | + ) |
| 142 | + pg_conn.commit() |
| 143 | + |
| 144 | + default_spec_id = run_query( |
| 145 | + "SELECT default_spec_id FROM lake_iceberg.tables_internal WHERE table_name = 'test_re_set_partition_fields.tbl'::regclass", |
| 146 | + superuser_conn, |
| 147 | + ) |
| 148 | + assert default_spec_id[0][0] == 3 |
| 149 | + assert ( |
| 150 | + default_table_partition_spec_id(pg_conn, "test_re_set_partition_fields", "tbl") |
| 151 | + == 3 |
| 152 | + ) |
| 153 | + assert ( |
| 154 | + len(table_partition_specs(pg_conn, "test_re_set_partition_fields", "tbl")) == 5 |
| 155 | + ) |
| 156 | + |
| 157 | + # finally, drop the spec and finish the test |
| 158 | + run_command( |
| 159 | + f"ALTER TABLE test_re_set_partition_fields.tbl OPTIONS (DROP partition_by)", |
| 160 | + pg_conn, |
| 161 | + ) |
| 162 | + pg_conn.commit() |
| 163 | + |
| 164 | + default_spec_id = run_query( |
| 165 | + "SELECT default_spec_id FROM lake_iceberg.tables_internal WHERE table_name = 'test_re_set_partition_fields.tbl'::regclass", |
| 166 | + superuser_conn, |
| 167 | + ) |
| 168 | + assert default_spec_id[0][0] == 0 |
| 169 | + assert ( |
| 170 | + default_table_partition_spec_id(pg_conn, "test_re_set_partition_fields", "tbl") |
| 171 | + == 0 |
| 172 | + ) |
| 173 | + assert ( |
| 174 | + len(table_partition_specs(pg_conn, "test_re_set_partition_fields", "tbl")) == 5 |
| 175 | + ) |
| 176 | + |
| 177 | + run_command(f"""DROP SCHEMA test_re_set_partition_fields CASCADE;""", pg_conn) |
| 178 | + pg_conn.commit() |
| 179 | + |
| 180 | + |
4 | 181 | def test_ddl_with_identity_partitions(extension, s3, with_default_location, pg_conn): |
5 | 182 | run_command(f"""CREATE SCHEMA test_ddl_with_partitions;""", pg_conn) |
6 | 183 |
|
@@ -455,3 +632,26 @@ def test_partition_by_with_collation(extension, s3, with_default_location, pg_co |
455 | 632 |
|
456 | 633 | run_command("DROP COLLATION s_coll;", pg_conn) |
457 | 634 | pg_conn.commit() |
| 635 | + |
| 636 | + |
| 637 | +def table_metadata(pg_conn, table_namespace, table_name): |
| 638 | + metadata_location = run_query( |
| 639 | + f"SELECT metadata_location FROM iceberg_tables WHERE table_name = '{table_name}' and table_namespace = '{table_namespace}'", |
| 640 | + pg_conn, |
| 641 | + )[0][0] |
| 642 | + |
| 643 | + pg_query = f"SELECT * FROM lake_iceberg.metadata('{metadata_location}')" |
| 644 | + |
| 645 | + metadata = run_query(pg_query, pg_conn)[0][0] |
| 646 | + |
| 647 | + return metadata |
| 648 | + |
| 649 | + |
| 650 | +def table_partition_specs(pg_conn, table_namespace, table_name): |
| 651 | + metadata = table_metadata(pg_conn, table_namespace, table_name) |
| 652 | + return metadata["partition-specs"] |
| 653 | + |
| 654 | + |
| 655 | +def default_table_partition_spec_id(pg_conn, table_namespace, table_name): |
| 656 | + metadata = table_metadata(pg_conn, table_namespace, table_name) |
| 657 | + return metadata["default-spec-id"] |
0 commit comments