collate-data-diff 0.11.2__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (54) hide show
  1. collate_data_diff-0.11.2.dist-info/LICENSE +18 -0
  2. collate_data_diff-0.11.2.dist-info/METADATA +77 -0
  3. collate_data_diff-0.11.2.dist-info/RECORD +54 -0
  4. collate_data_diff-0.11.2.dist-info/WHEEL +4 -0
  5. collate_data_diff-0.11.2.dist-info/entry_points.txt +3 -0
  6. data_diff/__init__.py +180 -0
  7. data_diff/__main__.py +618 -0
  8. data_diff/abcs/__init__.py +0 -0
  9. data_diff/abcs/compiler.py +13 -0
  10. data_diff/abcs/database_types.py +308 -0
  11. data_diff/cloud/__init__.py +2 -0
  12. data_diff/cloud/data_source.py +318 -0
  13. data_diff/cloud/datafold_api.py +304 -0
  14. data_diff/config.py +127 -0
  15. data_diff/databases/__init__.py +17 -0
  16. data_diff/databases/_connect.py +306 -0
  17. data_diff/databases/base.py +1291 -0
  18. data_diff/databases/bigquery.py +315 -0
  19. data_diff/databases/clickhouse.py +203 -0
  20. data_diff/databases/databricks.py +248 -0
  21. data_diff/databases/duckdb.py +192 -0
  22. data_diff/databases/mssql.py +229 -0
  23. data_diff/databases/mysql.py +159 -0
  24. data_diff/databases/oracle.py +195 -0
  25. data_diff/databases/postgresql.py +258 -0
  26. data_diff/databases/presto.py +197 -0
  27. data_diff/databases/redshift.py +217 -0
  28. data_diff/databases/snowflake.py +207 -0
  29. data_diff/databases/trino.py +50 -0
  30. data_diff/databases/vertica.py +160 -0
  31. data_diff/dbt.py +604 -0
  32. data_diff/dbt_config_validators.py +65 -0
  33. data_diff/dbt_parser.py +523 -0
  34. data_diff/diff_tables.py +416 -0
  35. data_diff/errors.py +74 -0
  36. data_diff/format.py +359 -0
  37. data_diff/hashdiff_tables.py +264 -0
  38. data_diff/info_tree.py +62 -0
  39. data_diff/joindiff_tables.py +399 -0
  40. data_diff/lexicographic_space.py +240 -0
  41. data_diff/parse_time.py +74 -0
  42. data_diff/py.typed +0 -0
  43. data_diff/queries/__init__.py +0 -0
  44. data_diff/queries/api.py +200 -0
  45. data_diff/queries/ast_classes.py +798 -0
  46. data_diff/queries/base.py +24 -0
  47. data_diff/queries/extras.py +29 -0
  48. data_diff/query_utils.py +56 -0
  49. data_diff/schema.py +52 -0
  50. data_diff/table_segment.py +286 -0
  51. data_diff/thread_utils.py +98 -0
  52. data_diff/tracking.py +237 -0
  53. data_diff/utils.py +625 -0
  54. data_diff/version.py +1 -0
@@ -0,0 +1,399 @@
1
+ """Provides classes for performing a table diff using JOIN"""
2
+
3
+ from decimal import Decimal
4
+ from functools import partial
5
+ import logging
6
+ from typing import List, Optional
7
+ from itertools import chain
8
+
9
+ import attrs
10
+
11
+ from data_diff.databases import Database, MsSQL, MySQL, BigQuery, Presto, Oracle, Snowflake, DuckDB
12
+ from data_diff.abcs.database_types import NumericType, DbPath
13
+ from data_diff.databases.base import Compiler
14
+ from data_diff.queries.api import (
15
+ table,
16
+ sum_,
17
+ and_,
18
+ if_,
19
+ or_,
20
+ outerjoin,
21
+ leftjoin,
22
+ rightjoin,
23
+ this,
24
+ when,
25
+ )
26
+ from data_diff.queries.ast_classes import Concat, Count, Expr, Random, TablePath, Code, ITable
27
+ from data_diff.queries.extras import NormalizeAsString
28
+ from data_diff.info_tree import InfoTree
29
+ from data_diff.query_utils import append_to_table, drop_table
30
+ from data_diff.utils import safezip
31
+ from data_diff.table_segment import TableSegment
32
+ from data_diff.diff_tables import TableDiffer, DiffResult
33
+ from data_diff.thread_utils import ThreadedYielder
34
+
35
+
36
+ logger = logging.getLogger("joindiff_tables")
37
+
38
+ TABLE_WRITE_LIMIT = 1000
39
+
40
+
41
+ def merge_dicts(dicts):
42
+ i = iter(dicts)
43
+ try:
44
+ res = next(i)
45
+ except StopIteration:
46
+ return {}
47
+
48
+ for d in i:
49
+ res.update(d)
50
+ return res
51
+
52
+
53
+ def sample(table_expr):
54
+ return table_expr.order_by(Random()).limit(10)
55
+
56
+
57
+ def create_temp_table(c: Compiler, path: TablePath, expr: Expr) -> str:
58
+ db = c.database
59
+ c: Compiler = attrs.evolve(c, root=False) # we're compiling fragments, not full queries
60
+ if isinstance(db, BigQuery):
61
+ return f"create table {c.dialect.compile(c, path)} OPTIONS(expiration_timestamp=TIMESTAMP_ADD(CURRENT_TIMESTAMP(), INTERVAL 1 DAY)) as {c.dialect.compile(c, expr)}"
62
+ elif isinstance(db, Presto):
63
+ return f"create table {c.dialect.compile(c, path)} as {c.dialect.compile(c, expr)}"
64
+ elif isinstance(db, Oracle):
65
+ return f"create global temporary table {c.dialect.compile(c, path)} as {c.dialect.compile(c, expr)}"
66
+ else:
67
+ return f"create temporary table {c.dialect.compile(c, path)} as {c.dialect.compile(c, expr)}"
68
+
69
+
70
+ def bool_to_int(x):
71
+ return if_(x, 1, 0)
72
+
73
+
74
+ def _outerjoin(db: Database, a: ITable, b: ITable, keys1: List[str], keys2: List[str], select_fields: dict) -> ITable:
75
+ on = [a[k1] == b[k2] for k1, k2 in safezip(keys1, keys2)]
76
+
77
+ is_exclusive_a = and_(b[k] == None for k in keys2)
78
+ is_exclusive_b = and_(a[k] == None for k in keys1)
79
+
80
+ if isinstance(db, MsSQL):
81
+ # There is no "IS NULL" or "ISNULL()" as expressions, only as conditions.
82
+ is_exclusive_a = when(is_exclusive_a).then(1).else_(0)
83
+ is_exclusive_b = when(is_exclusive_b).then(1).else_(0)
84
+
85
+ if isinstance(db, Oracle):
86
+ is_exclusive_a = bool_to_int(is_exclusive_a)
87
+ is_exclusive_b = bool_to_int(is_exclusive_b)
88
+
89
+ if isinstance(db, MySQL):
90
+ # No outer join
91
+ l = leftjoin(a, b).on(*on).select(is_exclusive_a=is_exclusive_a, is_exclusive_b=False, **select_fields)
92
+ r = rightjoin(a, b).on(*on).select(is_exclusive_a=False, is_exclusive_b=is_exclusive_b, **select_fields)
93
+ return l.union(r)
94
+
95
+ return outerjoin(a, b).on(*on).select(is_exclusive_a=is_exclusive_a, is_exclusive_b=is_exclusive_b, **select_fields)
96
+
97
+
98
+ def _slice_tuple(t, *sizes):
99
+ i = 0
100
+ for size in sizes:
101
+ yield t[i : i + size]
102
+ i += size
103
+ assert i == len(t)
104
+
105
+
106
+ def json_friendly_value(v):
107
+ if isinstance(v, Decimal):
108
+ return float(v)
109
+ return v
110
+
111
+
112
+ @attrs.define(frozen=False)
113
+ class JoinDiffer(TableDiffer):
114
+ """Finds the diff between two SQL tables in the same database, using JOINs.
115
+
116
+ The algorithm uses an OUTER JOIN (or equivalent) with extra checks and statistics.
117
+ The two tables must reside in the same database, and their primary keys must be unique and not null.
118
+
119
+ All parameters are optional.
120
+
121
+ Parameters:
122
+ threaded (bool): Enable/disable threaded diffing. Needed to take advantage of database threads.
123
+ max_threadpool_size (int): Maximum size of each threadpool. ``None`` means auto.
124
+ Only relevant when `threaded` is ``True``.
125
+ There may be many pools, so number of actual threads can be a lot higher.
126
+ validate_unique_key (bool): Enable/disable validating that the key columns are unique. (default: True)
127
+ If there are no UNIQUE constraints in the schema, it is done in a single query,
128
+ and can't be threaded, so it's very slow on non-cloud dbs.
129
+ sample_exclusive_rows (bool): Enable/disable sampling of exclusive rows. (default: False)
130
+ Creates a temporary table.
131
+ materialize_to_table (DbPath, optional): Path of new table to write diff results to. Disabled if not provided.
132
+ materialize_all_rows (bool): Materialize every row, not just those that are different. (default: False)
133
+ table_write_limit (int): Maximum number of rows to write when materializing, per thread.
134
+ skip_null_keys (bool): Skips diffing any rows with null PKs (displays a warning if any are null) (default: False)
135
+ """
136
+
137
+ validate_unique_key: bool = True
138
+ sample_exclusive_rows: bool = False
139
+ materialize_to_table: Optional[DbPath] = None
140
+ materialize_all_rows: bool = False
141
+ table_write_limit: int = TABLE_WRITE_LIMIT
142
+ skip_null_keys: bool = False
143
+
144
+ stats: dict = attrs.field(factory=dict)
145
+
146
+ def _diff_tables_root(self, table1: TableSegment, table2: TableSegment, info_tree: InfoTree) -> DiffResult:
147
+ db = table1.database
148
+
149
+ if table1.database is not table2.database:
150
+ raise ValueError("Join-diff only works when both tables are in the same database")
151
+
152
+ table1, table2 = self._threaded_call("with_schema", [table1, table2])
153
+
154
+ bg_funcs = [partial(self._test_duplicate_keys, table1, table2)] if self.validate_unique_key else []
155
+ if self.materialize_to_table:
156
+ drop_table(db, self.materialize_to_table)
157
+
158
+ with self._run_in_background(*bg_funcs):
159
+ if isinstance(db, (Snowflake, BigQuery, DuckDB)):
160
+ # Don't segment the table; let the database handling parallelization
161
+ yield from self._diff_segments(None, table1, table2, info_tree, None)
162
+ else:
163
+ yield from self._bisect_and_diff_tables(table1, table2, info_tree)
164
+ logger.info(f"Diffing complete: {table1.table_path} <> {table2.table_path}")
165
+ if self.materialize_to_table:
166
+ logger.info("Materialized diff to table '%s'.", ".".join(self.materialize_to_table))
167
+
168
+ def _diff_segments(
169
+ self,
170
+ ti: ThreadedYielder,
171
+ table1: TableSegment,
172
+ table2: TableSegment,
173
+ info_tree: InfoTree,
174
+ max_rows: int,
175
+ level=0,
176
+ segment_index=None,
177
+ segment_count=None,
178
+ ):
179
+ assert table1.database is table2.database
180
+
181
+ if segment_index or table1.min_key or max_rows:
182
+ logger.info(
183
+ ". " * level + f"Diffing segment {segment_index}/{segment_count}, "
184
+ f"key-range: {table1.min_key}..{table2.max_key}, "
185
+ f"size <= {max_rows}"
186
+ )
187
+
188
+ db = table1.database
189
+ diff_rows, a_cols, b_cols, is_diff_cols, all_rows = self._create_outer_join(table1, table2)
190
+
191
+ with self._run_in_background(
192
+ partial(self._collect_stats, 1, table1, info_tree),
193
+ partial(self._collect_stats, 2, table2, info_tree),
194
+ partial(self._test_null_keys, table1, table2),
195
+ partial(self._sample_and_count_exclusive, db, diff_rows, a_cols, b_cols, table1, table2),
196
+ partial(self._count_diff_per_column, db, diff_rows, list(a_cols), is_diff_cols, table1, table2),
197
+ partial(
198
+ self._materialize_diff,
199
+ db,
200
+ all_rows if self.materialize_all_rows else diff_rows,
201
+ segment_index=segment_index,
202
+ )
203
+ if self.materialize_to_table
204
+ else None,
205
+ ):
206
+ assert len(a_cols) == len(b_cols)
207
+ logger.debug(f"Querying for different rows: {table1.table_path}")
208
+ diff = db.query(diff_rows, list, log_message=table1.table_path)
209
+ info_tree.info.set_diff(diff, schema=tuple(diff_rows.schema.items()))
210
+ for is_xa, is_xb, *x in diff:
211
+ if is_xa and is_xb:
212
+ # Can't both be exclusive, meaning a pk is NULL
213
+ # This can happen if the explicit null test didn't finish running yet
214
+ if self.skip_null_keys:
215
+ # warning is thrown in explicit null test
216
+ continue
217
+ else:
218
+ raise ValueError("NULL values in one or more primary keys")
219
+ # _is_diff, a_row, b_row = _slice_tuple(x, len(is_diff_cols), len(a_cols), len(b_cols))
220
+ _is_diff, ab_row = _slice_tuple(x, len(is_diff_cols), len(a_cols) + len(b_cols))
221
+ a_row, b_row = ab_row[::2], ab_row[1::2]
222
+ assert len(a_row) == len(b_row)
223
+ if not is_xb:
224
+ yield "-", tuple(a_row)
225
+ if not is_xa:
226
+ yield "+", tuple(b_row)
227
+
228
+ def _test_duplicate_keys(self, table1: TableSegment, table2: TableSegment):
229
+ logger.debug(f"Testing for duplicate keys: {table1.table_path} <> {table2.table_path}")
230
+
231
+ # Test duplicate keys
232
+ for ts in [table1, table2]:
233
+ unique = (
234
+ ts.database.query_table_unique_columns(ts.table_path) if ts.database.SUPPORTS_UNIQUE_CONSTAINT else []
235
+ )
236
+
237
+ t = ts.make_select()
238
+ key_columns = ts.key_columns
239
+
240
+ unvalidated = list(set(key_columns) - set(unique))
241
+ if unvalidated:
242
+ logger.info(f"Validating that the are no duplicate keys in columns: {unvalidated} for {ts.table_path}")
243
+ # Validate that there are no duplicate keys
244
+ self.stats["validated_unique_keys"] = self.stats.get("validated_unique_keys", []) + [unvalidated]
245
+ q = t.select(total=Count(), total_distinct=Count(Concat(this[unvalidated]), distinct=True))
246
+ total, total_distinct = ts.database.query(q, tuple, log_message=ts.table_path)
247
+ if total != total_distinct:
248
+ raise ValueError("Duplicate primary keys")
249
+
250
+ def _test_null_keys(self, table1, table2):
251
+ logger.debug(f"Testing for null keys: {table1.table_path} <> {table2.table_path}")
252
+
253
+ # Test null keys
254
+ for ts in [table1, table2]:
255
+ t = ts.make_select()
256
+ key_columns = ts.key_columns
257
+
258
+ q = t.select(*this[key_columns]).where(or_(this[k] == None for k in key_columns))
259
+ nulls = ts.database.query(q, list, log_message=ts.table_path)
260
+ if nulls:
261
+ if self.skip_null_keys:
262
+ logger.warning(
263
+ f"NULL values in one or more primary keys of {ts.table_path}. Skipping rows with NULL keys."
264
+ )
265
+ else:
266
+ raise ValueError(f"NULL values in one or more primary keys of {ts.table_path}")
267
+
268
+ def _collect_stats(self, i, table_seg: TableSegment, info_tree: InfoTree):
269
+ logger.debug(f"Collecting stats for table #{i}: {table_seg.table_path}")
270
+ db = table_seg.database
271
+
272
+ # Metrics
273
+ col_exprs = merge_dicts(
274
+ {
275
+ # f"min_{c}": min_(this[c]),
276
+ # f"max_{c}": max_(this[c]),
277
+ }
278
+ if c in table_seg.key_columns
279
+ else {
280
+ f"sum_{c}": sum_(this[c]),
281
+ # f"avg_{c}": avg(this[c]),
282
+ # f"min_{c}": min_(this[c]),
283
+ # f"max_{c}": max_(this[c]),
284
+ }
285
+ for c in table_seg.relevant_columns
286
+ if isinstance(table_seg._schema[c], NumericType)
287
+ )
288
+ col_exprs["count"] = Count()
289
+
290
+ res = db.query(table_seg.make_select().select(**col_exprs), tuple, log_message=table_seg.table_path)
291
+
292
+ for col_name, value in safezip(col_exprs, res):
293
+ if value is not None:
294
+ value = json_friendly_value(value)
295
+ stat_name = f"table{i}_{col_name}"
296
+
297
+ if col_name == "count":
298
+ info_tree.info.rowcounts[i] = value
299
+
300
+ if stat_name in self.stats:
301
+ self.stats[stat_name] += value
302
+ else:
303
+ self.stats[stat_name] = value
304
+
305
+ logger.debug("Done collecting stats for table #%s: %s", i, table_seg.table_path)
306
+
307
+ def _create_outer_join(self, table1, table2):
308
+ db = table1.database
309
+ if db is not table2.database:
310
+ raise ValueError("Joindiff only applies to tables within the same database")
311
+
312
+ keys1 = table1.key_columns
313
+ keys2 = table2.key_columns
314
+ if len(keys1) != len(keys2):
315
+ raise ValueError("The provided key columns are of a different count")
316
+
317
+ cols1 = table1.relevant_columns
318
+ cols2 = table2.relevant_columns
319
+ if len(cols1) != len(cols2):
320
+ raise ValueError("The provided columns are of a different count")
321
+
322
+ a = table1.make_select()
323
+ b = table2.make_select()
324
+
325
+ is_diff_cols = {f"is_diff_{c1}": bool_to_int(a[c1].is_distinct_from(b[c2])) for c1, c2 in safezip(cols1, cols2)}
326
+
327
+ a_cols = {f"{c}_a": NormalizeAsString(a[c]) for c in cols1}
328
+ b_cols = {f"{c}_b": NormalizeAsString(b[c]) for c in cols2}
329
+ # Order columns as col1_a, col1_b, col2_a, col2_b, etc.
330
+ cols = {k: v for k, v in chain(*zip(a_cols.items(), b_cols.items()))}
331
+
332
+ all_rows = _outerjoin(db, a, b, keys1, keys2, {**is_diff_cols, **cols})
333
+ diff_rows = all_rows.where(or_(this[c] == 1 for c in is_diff_cols))
334
+ return diff_rows, a_cols, b_cols, is_diff_cols, all_rows
335
+
336
+ def _count_diff_per_column(
337
+ self,
338
+ db,
339
+ diff_rows,
340
+ cols,
341
+ is_diff_cols,
342
+ table1: Optional[TableSegment] = None,
343
+ table2: Optional[TableSegment] = None,
344
+ ):
345
+ logger.debug(f"Counting differences per column: {table1.table_path} <> {table2.table_path}")
346
+ is_diff_cols_counts = db.query(
347
+ diff_rows.select(sum_(this[c]) for c in is_diff_cols),
348
+ tuple,
349
+ log_message=f"{table1.table_path} <> {table2.table_path}",
350
+ )
351
+ diff_counts = {}
352
+ for name, count in safezip(cols, is_diff_cols_counts):
353
+ diff_counts[name] = diff_counts.get(name, 0) + (count or 0)
354
+ self.stats["diff_counts"] = diff_counts
355
+
356
+ def _sample_and_count_exclusive(
357
+ self,
358
+ db,
359
+ diff_rows,
360
+ a_cols,
361
+ b_cols,
362
+ table1: Optional[TableSegment] = None,
363
+ table2: Optional[TableSegment] = None,
364
+ ):
365
+ if isinstance(db, (Oracle, MsSQL)):
366
+ exclusive_rows_query = diff_rows.where((this.is_exclusive_a == 1) | (this.is_exclusive_b == 1))
367
+ else:
368
+ exclusive_rows_query = diff_rows.where(this.is_exclusive_a | this.is_exclusive_b)
369
+
370
+ if not self.sample_exclusive_rows:
371
+ logger.debug(f"Counting exclusive rows: {table1.table_path} <> {table2.table_path}")
372
+ self.stats["exclusive_count"] = db.query(
373
+ exclusive_rows_query.count(), int, log_message=f"{table1.table_path} <> {table2.table_path}"
374
+ )
375
+ return
376
+
377
+ logger.info("Counting and sampling exclusive rows")
378
+
379
+ def exclusive_rows(expr):
380
+ c = Compiler(db)
381
+ name = c.new_unique_table_name("temp_table")
382
+ exclusive_rows = table(name, schema=expr.source_table.schema)
383
+ yield Code(create_temp_table(c, exclusive_rows, expr.limit(self.table_write_limit)))
384
+
385
+ count = yield exclusive_rows.count()
386
+ self.stats["exclusive_count"] = self.stats.get("exclusive_count", 0) + count[0][0]
387
+ sample_rows = yield sample(exclusive_rows.select(*this[list(a_cols)], *this[list(b_cols)]))
388
+ self.stats["exclusive_sample"] = self.stats.get("exclusive_sample", []) + sample_rows
389
+
390
+ # Only drops if create table succeeded (meaning, the table didn't already exist)
391
+ yield exclusive_rows.drop()
392
+
393
+ # Run as a sequence of thread-local queries (compiled into a ThreadLocalInterpreter)
394
+ db.query(exclusive_rows(exclusive_rows_query), None)
395
+
396
+ def _materialize_diff(self, db, diff_rows, segment_index=None):
397
+ assert self.materialize_to_table
398
+
399
+ append_to_table(db, self.materialize_to_table, diff_rows.limit(self.table_write_limit))
@@ -0,0 +1,240 @@
1
+ """Contains the implementation of two classes:
2
+
3
+ - LexicographicSpace - a lexicographic space of arbitrary dimensions.
4
+ - BoundedLexicographicSpace - a lexicographic space, where the lowest point may be non-zero.
5
+
6
+ A lexicographic space is a space of increasing natural values, ordered by lexicographic order.
7
+ Read more: https://mathworld.wolfram.com/LexicographicOrder.html
8
+
9
+ These abstractions were written to support compound keys in the hashdiff algorithm.
10
+ In the hashdiff algorithm, we rely on the order of the column keys, to segment the table correctly.
11
+ SQL orders the columns of tables based on lexicographic ordering.
12
+ Since we need an evenly spaced "range" function over the space, which has arbitrary dimensions, we have
13
+ to implement it ourself.
14
+
15
+ As a further optimization, since we each time operate on segments of the ordered table, we add support
16
+ for working with a restricted space, which will reduce the likelihood of gaps in our "select", when the
17
+ keys are not evenly distributed.
18
+ """
19
+
20
+ from random import randint, randrange
21
+
22
+ from typing import Tuple
23
+
24
+ import attrs
25
+
26
+ from data_diff.utils import safezip
27
+
28
+ Vector = Tuple[int]
29
+ Interval = Tuple[int]
30
+
31
+
32
+ class Overflow(ValueError):
33
+ pass
34
+
35
+
36
+ def neg_interval(interval):
37
+ return tuple(-i for i in interval)
38
+
39
+
40
+ def neg_v(v: Vector):
41
+ return tuple(-i for i in v)
42
+
43
+
44
+ def sub_v(v1: Vector, v2: Vector):
45
+ return tuple(i1 - i2 for i1, i2 in safezip(v1, v2))
46
+
47
+
48
+ def add_v(v1: Vector, v2: Vector):
49
+ return tuple(i1 + i2 for i1, i2 in safezip(v1, v2))
50
+
51
+
52
+ def rand_v_in_range(v1: Vector, v2: Vector):
53
+ return tuple(irandrange(i1, i2) for i1, i2 in safezip(v1, v2))
54
+
55
+
56
+ def irandrange(start, stop):
57
+ if start == stop:
58
+ return start
59
+ return randrange(start, stop)
60
+
61
+
62
+ @attrs.define(frozen=True)
63
+ class LexicographicSpace:
64
+ """Lexicographic space of arbitrary dimensions.
65
+
66
+ All elements must be of the same length as the number of dimensions. (no rpadding)
67
+ """
68
+
69
+ def __init__(self, dims: Vector) -> None:
70
+ super().__init__()
71
+ self.dims = dims
72
+
73
+ def __contains__(self, v: Vector) -> bool:
74
+ return all(0 <= i < d for i, d in safezip(v, self.dims))
75
+
76
+ def add(self, v1: Vector, v2: Vector) -> Vector:
77
+ # assert v1 in self and v2 in self, (v1, v2)
78
+
79
+ carry = 0
80
+ res = []
81
+ for i1, i2, d in reversed(list(safezip(v1, v2, self.dims))):
82
+ n = i1 + i2 + carry
83
+ carry = n // d
84
+ assert carry <= 1
85
+ n %= d
86
+ res.append(n)
87
+
88
+ if carry:
89
+ raise Overflow("Overflow")
90
+
91
+ new_v = tuple(reversed(res))
92
+ assert new_v in self
93
+ return new_v
94
+
95
+ def sub(self, v1: Vector, v2: Vector):
96
+ return self.add(v1, neg_v(v2))
97
+
98
+ def _divide(self, v: Vector, count: int):
99
+ n = 0
100
+ for x, d in zip(v, self.dims[1:] + (1,), strict=True):
101
+ x += n
102
+ rem = x % count
103
+ n = rem * d
104
+ yield x // count
105
+
106
+ def divide(self, v: Vector, count: int) -> Vector:
107
+ return tuple(self._divide(v, count))
108
+
109
+ def range(self, min_value: Vector, max_value: Vector, count: int):
110
+ assert min_value in self and max_value in self
111
+ count -= 1
112
+ size = self.sub(max_value, min_value)
113
+ interval = self.divide(size, count)
114
+ n = min_value
115
+ for i in range(count):
116
+ yield n
117
+ n = self.add(n, interval)
118
+ yield n
119
+
120
+
121
+ class BoundedLexicographicSpace:
122
+ """Lexicographic space of arbitrary dimensions, where the lowest point may be non-zero.
123
+
124
+ i.e. a space resticted by a "bounding-box" between two arbitrary points.
125
+ """
126
+
127
+ def __init__(self, min_bound: Vector, max_bound: Vector) -> None:
128
+ super().__init__()
129
+
130
+ dims = tuple(mx - mn for mn, mx in safezip(min_bound, max_bound))
131
+ if not all(d >= 0 for d in dims):
132
+ raise ValueError("Error: Negative dimension!")
133
+ if not (dims[0] > 0):
134
+ raise ValueError("First dimension must be non-zero!")
135
+
136
+ self.min_bound = min_bound
137
+ self.max_bound = max_bound
138
+
139
+ self.uspace = LexicographicSpace(dims)
140
+
141
+ def __contains__(self, p: Vector) -> bool:
142
+ return all(mn <= i < mx for i, mn, mx in safezip(p, self.min_bound, self.max_bound))
143
+
144
+ def to_uspace(self, v: Vector) -> Vector:
145
+ assert v in self
146
+ return sub_v(v, self.min_bound)
147
+
148
+ def from_uspace(self, v: Vector) -> Vector:
149
+ res = add_v(v, self.min_bound)
150
+ assert res in self
151
+ return res
152
+
153
+ def add_interval(self, v1: Vector, interval: Interval) -> Vector:
154
+ return self.from_uspace(self.uspace.add(self.to_uspace(v1), interval))
155
+
156
+ def sub_interval(self, v1: Vector, interval: Interval) -> Vector:
157
+ return self.from_uspace(self.uspace.sub(self.to_uspace(v1), interval))
158
+
159
+ def sub(self, v1: Vector, v2: Vector) -> Interval:
160
+ return self.uspace.sub(self.to_uspace(v1), self.to_uspace(v2))
161
+
162
+ def range(self, min_value: Vector, max_value: Vector, count: int):
163
+ return [
164
+ self.from_uspace(v) for v in self.uspace.range(self.to_uspace(min_value), self.to_uspace(max_value), count)
165
+ ]
166
+
167
+
168
+ def test_lex_space():
169
+ # Test add
170
+ binspace = LexicographicSpace((2, 2, 2, 2))
171
+ zero = (0, 0, 0, 0)
172
+ one = (0, 0, 0, 1)
173
+ bin_nums = [zero]
174
+ for i in range(15):
175
+ last = bin_nums[-1]
176
+ bin_nums.append(binspace.add(last, one))
177
+ five = bin_nums[5]
178
+ seven = bin_nums[7]
179
+ eight = bin_nums[8]
180
+ fifteen = bin_nums[15]
181
+
182
+ assert binspace.add(binspace.add(one, five), one) == seven
183
+ assert binspace.add(one, seven) == eight
184
+ assert binspace.add(seven, eight) == fifteen
185
+
186
+ assert binspace.sub(eight, one) == seven
187
+ assert binspace.sub(fifteen, seven) == eight
188
+
189
+ r = list(binspace.range(one, seven, 4))
190
+ assert r == [one, bin_nums[3], five, seven], r
191
+
192
+ decspace = LexicographicSpace((10, 10, 10))
193
+ assert decspace.divide((4, 5, 2), 2) == (2, 2, 6)
194
+ assert decspace.divide((3, 0, 2), 2) == (1, 5, 1)
195
+
196
+ # Restricted space
197
+
198
+ rspace1 = BoundedLexicographicSpace((2, 2), (8, 8))
199
+ assert rspace1.add_interval((2, 2), (0, 0)) == (2, 2)
200
+ assert rspace1.add_interval((2, 2), (0, 1)) == (2, 3)
201
+ assert rspace1.add_interval((2, 2), (0, 6)) == (3, 2)
202
+ assert rspace1.add_interval((2, 2), (0, 7)) == (3, 3)
203
+ # space.add((2,2), (6, 0)) # Overflow
204
+
205
+ rspace2 = BoundedLexicographicSpace((4, 4, 4, 4), (6, 6, 6, 6))
206
+ _one = (4, 4, 4, 5)
207
+ _three = (4, 4, 5, 5)
208
+ _five = (4, 5, 4, 5)
209
+ _seven = (4, 5, 5, 5)
210
+ assert rspace2.add_interval(rspace2.add_interval(_five, one), one) == _seven
211
+ assert rspace2.sub_interval(rspace2.sub_interval(_seven, one), one) == _five
212
+
213
+ r = list(rspace2.range(_one, _seven, 4))
214
+ assert r == [_one, _three, _five, _seven], r
215
+
216
+ # Test range -
217
+ # For random bounds and min/max values, assert that range() generates steps with uniform distances
218
+ MAX_COLUMNS = 16
219
+ MAX_DIM = 10000
220
+ MAX_BISECTION = 128
221
+
222
+ for n in range(1, MAX_COLUMNS):
223
+ min_bound = tuple(randint(0, MAX_DIM) for i in range(n))
224
+ size = tuple(randint(1, MAX_DIM) for i in range(n))
225
+ max_bound = add_v(min_bound, size)
226
+
227
+ sp = BoundedLexicographicSpace(min_bound, max_bound)
228
+
229
+ max_value = rand_v_in_range(min_bound, max_bound)
230
+ min_value = rand_v_in_range(min_bound, max_value)
231
+ for count in range(2, MAX_BISECTION):
232
+ r = sp.range(min_value, max_value, count)
233
+ assert len(r) == count
234
+ diffs = [sp.sub(b, a) for a, b in zip(r[:-1], r[1:])]
235
+ assert len(set(diffs)) == 1 # Uniform!
236
+ # print('.', end='')
237
+
238
+
239
+ if __name__ == "__main__":
240
+ test_lex_space()