sqllocks-shape-simulation 0.9.1__tar.gz

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 (44) hide show
  1. sqllocks_shape_simulation-0.9.1/LICENSE +21 -0
  2. sqllocks_shape_simulation-0.9.1/PKG-INFO +61 -0
  3. sqllocks_shape_simulation-0.9.1/README.md +38 -0
  4. sqllocks_shape_simulation-0.9.1/pyproject.toml +38 -0
  5. sqllocks_shape_simulation-0.9.1/setup.cfg +4 -0
  6. sqllocks_shape_simulation-0.9.1/src/shape_simulation/__init__.py +157 -0
  7. sqllocks_shape_simulation-0.9.1/src/shape_simulation/_patterns.py +321 -0
  8. sqllocks_shape_simulation-0.9.1/src/shape_simulation/_tables.py +156 -0
  9. sqllocks_shape_simulation-0.9.1/src/shape_simulation/cli.py +57 -0
  10. sqllocks_shape_simulation-0.9.1/src/shape_simulation/cli_batch.py +421 -0
  11. sqllocks_shape_simulation-0.9.1/src/shape_simulation/cli_patterns.py +221 -0
  12. sqllocks_shape_simulation-0.9.1/src/shape_simulation/clickstream_patterns.py +442 -0
  13. sqllocks_shape_simulation-0.9.1/src/shape_simulation/file_drop.py +453 -0
  14. sqllocks_shape_simulation-0.9.1/src/shape_simulation/financial_patterns.py +445 -0
  15. sqllocks_shape_simulation-0.9.1/src/shape_simulation/hybrid.py +192 -0
  16. sqllocks_shape_simulation-0.9.1/src/shape_simulation/iot_patterns.py +430 -0
  17. sqllocks_shape_simulation-0.9.1/src/shape_simulation/operational_log_patterns.py +542 -0
  18. sqllocks_shape_simulation-0.9.1/src/shape_simulation/pulse_patterns.py +629 -0
  19. sqllocks_shape_simulation-0.9.1/src/shape_simulation/scd2_file_drops.py +446 -0
  20. sqllocks_shape_simulation-0.9.1/src/shape_simulation/state_machine.py +496 -0
  21. sqllocks_shape_simulation-0.9.1/src/shape_simulation/stream_emit.py +495 -0
  22. sqllocks_shape_simulation-0.9.1/src/sqllocks_shape_simulation.egg-info/PKG-INFO +61 -0
  23. sqllocks_shape_simulation-0.9.1/src/sqllocks_shape_simulation.egg-info/SOURCES.txt +42 -0
  24. sqllocks_shape_simulation-0.9.1/src/sqllocks_shape_simulation.egg-info/dependency_links.txt +1 -0
  25. sqllocks_shape_simulation-0.9.1/src/sqllocks_shape_simulation.egg-info/entry_points.txt +2 -0
  26. sqllocks_shape_simulation-0.9.1/src/sqllocks_shape_simulation.egg-info/requires.txt +1 -0
  27. sqllocks_shape_simulation-0.9.1/src/sqllocks_shape_simulation.egg-info/top_level.txt +1 -0
  28. sqllocks_shape_simulation-0.9.1/tests/test_clickstream.py +187 -0
  29. sqllocks_shape_simulation-0.9.1/tests/test_common.py +124 -0
  30. sqllocks_shape_simulation-0.9.1/tests/test_config_checks.py +74 -0
  31. sqllocks_shape_simulation-0.9.1/tests/test_e2e.py +350 -0
  32. sqllocks_shape_simulation-0.9.1/tests/test_file_drop.py +233 -0
  33. sqllocks_shape_simulation-0.9.1/tests/test_financial.py +365 -0
  34. sqllocks_shape_simulation-0.9.1/tests/test_hybrid.py +99 -0
  35. sqllocks_shape_simulation-0.9.1/tests/test_iot.py +218 -0
  36. sqllocks_shape_simulation-0.9.1/tests/test_operational_log.py +189 -0
  37. sqllocks_shape_simulation-0.9.1/tests/test_patterns_e2e.py +266 -0
  38. sqllocks_shape_simulation-0.9.1/tests/test_portability.py +53 -0
  39. sqllocks_shape_simulation-0.9.1/tests/test_pulse.py +288 -0
  40. sqllocks_shape_simulation-0.9.1/tests/test_scd2_file_drops.py +191 -0
  41. sqllocks_shape_simulation-0.9.1/tests/test_skeleton.py +9 -0
  42. sqllocks_shape_simulation-0.9.1/tests/test_state_machine.py +140 -0
  43. sqllocks_shape_simulation-0.9.1/tests/test_stream_emit.py +319 -0
  44. sqllocks_shape_simulation-0.9.1/tests/test_tables.py +62 -0
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2025-2026 SQLLocks (Jonathan Stewart)
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
@@ -0,0 +1,61 @@
1
+ Metadata-Version: 2.4
2
+ Name: sqllocks-shape-simulation
3
+ Version: 0.9.1
4
+ Summary: Shape plugin: Simulation scenarios
5
+ License-Expression: MIT
6
+ Project-URL: Homepage, https://github.com/sqllocks/shape
7
+ Project-URL: Repository, https://github.com/sqllocks/shape
8
+ Project-URL: Issues, https://github.com/sqllocks/shape/issues
9
+ Project-URL: Documentation, https://github.com/sqllocks/shape/tree/main/plugins/shape-simulation
10
+ Classifier: Programming Language :: Python :: 3
11
+ Classifier: Programming Language :: Python :: 3 :: Only
12
+ Classifier: Programming Language :: Python :: 3.11
13
+ Classifier: Programming Language :: Python :: 3.12
14
+ Classifier: Programming Language :: Python :: 3.13
15
+ Classifier: Programming Language :: Python :: 3.14
16
+ Classifier: Operating System :: OS Independent
17
+ Classifier: Topic :: Database
18
+ Requires-Python: >=3.11
19
+ Description-Content-Type: text/markdown
20
+ License-File: LICENSE
21
+ Requires-Dist: sqllocks-shape==0.9.1
22
+ Dynamic: license-file
23
+
24
+ # sqllocks-shape-simulation
25
+
26
+ Shape plugin: simulation scenarios. Turns generated tables into what an upstream system would
27
+ produce over time, and generates the telemetry of systems that have no table behind them.
28
+
29
+ | Simulator | Module | `shape simulate` |
30
+ |---|---|---|
31
+ | `FileDropSimulator` | `shape_simulation.file_drop` | `file-drop` |
32
+ | `SCD2FileDropSimulator` | `shape_simulation.scd2_file_drops` | `scd2` |
33
+ | `StreamEmitter` | `shape_simulation.stream_emit` | `stream` |
34
+ | `HybridSimulator` | `shape_simulation.hybrid` | `hybrid` |
35
+ | `WorkflowSimulator` | `shape_simulation.state_machine` | `workflow` |
36
+ | `ClickstreamSimulator` | `shape_simulation.clickstream_patterns` | `clickstream` |
37
+ | `FinancialStreamSimulator` | `shape_simulation.financial_patterns` | `financial` |
38
+ | `IoTTelemetrySimulator` | `shape_simulation.iot_patterns` | `iot` |
39
+ | `OperationalLogSimulator` | `shape_simulation.operational_log_patterns` | `operational-log` |
40
+ | `PulseDemandSimulator` | `shape_simulation.pulse_patterns` | `pulse` |
41
+
42
+ Each takes a configuration and, where it works on existing data, Arrow tables (or a generation
43
+ result), and is deterministic for a seed. The first five are described in
44
+ `docs/SIMULATION_FILES_EVENTS.md`, the pattern simulators (the last five) in
45
+ `docs/plugins/simulation.md`.
46
+
47
+ ```bash
48
+ pip install 'sqllocks-shape[simulation]'
49
+ shape simulate clickstream --set users=500 -o out/
50
+ ```
51
+
52
+ ```python
53
+ from shape_simulation.clickstream_patterns import ClickstreamConfig, ClickstreamSimulator
54
+
55
+ result = ClickstreamSimulator(ClickstreamConfig(users=500, seed=7)).run()
56
+ result.sessions, result.page_views, result.funnels, result.stats
57
+ ```
58
+
59
+ It registers one plugin command, `simulate` (`shape.commands`); the simulators are importable from
60
+ `shape_simulation` without it. Its version always equals core's (`sqllocks-shape`), and it is
61
+ released together with core. How plugins are written: `docs/plugins/authoring.md` in the repository.
@@ -0,0 +1,38 @@
1
+ # sqllocks-shape-simulation
2
+
3
+ Shape plugin: simulation scenarios. Turns generated tables into what an upstream system would
4
+ produce over time, and generates the telemetry of systems that have no table behind them.
5
+
6
+ | Simulator | Module | `shape simulate` |
7
+ |---|---|---|
8
+ | `FileDropSimulator` | `shape_simulation.file_drop` | `file-drop` |
9
+ | `SCD2FileDropSimulator` | `shape_simulation.scd2_file_drops` | `scd2` |
10
+ | `StreamEmitter` | `shape_simulation.stream_emit` | `stream` |
11
+ | `HybridSimulator` | `shape_simulation.hybrid` | `hybrid` |
12
+ | `WorkflowSimulator` | `shape_simulation.state_machine` | `workflow` |
13
+ | `ClickstreamSimulator` | `shape_simulation.clickstream_patterns` | `clickstream` |
14
+ | `FinancialStreamSimulator` | `shape_simulation.financial_patterns` | `financial` |
15
+ | `IoTTelemetrySimulator` | `shape_simulation.iot_patterns` | `iot` |
16
+ | `OperationalLogSimulator` | `shape_simulation.operational_log_patterns` | `operational-log` |
17
+ | `PulseDemandSimulator` | `shape_simulation.pulse_patterns` | `pulse` |
18
+
19
+ Each takes a configuration and, where it works on existing data, Arrow tables (or a generation
20
+ result), and is deterministic for a seed. The first five are described in
21
+ `docs/SIMULATION_FILES_EVENTS.md`, the pattern simulators (the last five) in
22
+ `docs/plugins/simulation.md`.
23
+
24
+ ```bash
25
+ pip install 'sqllocks-shape[simulation]'
26
+ shape simulate clickstream --set users=500 -o out/
27
+ ```
28
+
29
+ ```python
30
+ from shape_simulation.clickstream_patterns import ClickstreamConfig, ClickstreamSimulator
31
+
32
+ result = ClickstreamSimulator(ClickstreamConfig(users=500, seed=7)).run()
33
+ result.sessions, result.page_views, result.funnels, result.stats
34
+ ```
35
+
36
+ It registers one plugin command, `simulate` (`shape.commands`); the simulators are importable from
37
+ `shape_simulation` without it. Its version always equals core's (`sqllocks-shape`), and it is
38
+ released together with core. How plugins are written: `docs/plugins/authoring.md` in the repository.
@@ -0,0 +1,38 @@
1
+ [build-system]
2
+ requires = ["setuptools>=77"]
3
+ build-backend = "setuptools.build_meta"
4
+
5
+ [project]
6
+ name = "sqllocks-shape-simulation"
7
+ # Lockstep with core (T-09): scripts/check_plugin_skeletons.py fails when this drifts.
8
+ version = "0.9.1"
9
+ description = "Shape plugin: Simulation scenarios"
10
+ readme = "README.md"
11
+ requires-python = ">=3.11"
12
+ license = "MIT"
13
+ license-files = ["LICENSE"]
14
+ classifiers = [
15
+ "Programming Language :: Python :: 3",
16
+ "Programming Language :: Python :: 3 :: Only",
17
+ "Programming Language :: Python :: 3.11",
18
+ "Programming Language :: Python :: 3.12",
19
+ "Programming Language :: Python :: 3.13",
20
+ "Programming Language :: Python :: 3.14",
21
+ "Operating System :: OS Independent",
22
+ "Topic :: Database",
23
+ ]
24
+ dependencies = ["sqllocks-shape==0.9.1"]
25
+
26
+ # Entry points (shape.* groups, plugin API v1); see docs/plugins/authoring.md.
27
+
28
+ [project.urls]
29
+ Homepage = "https://github.com/sqllocks/shape"
30
+ Repository = "https://github.com/sqllocks/shape"
31
+ Issues = "https://github.com/sqllocks/shape/issues"
32
+ Documentation = "https://github.com/sqllocks/shape/tree/main/plugins/shape-simulation"
33
+
34
+ [project.entry-points."shape.commands"]
35
+ simulate = "shape_simulation.cli:SimulateCommand"
36
+
37
+ [tool.setuptools.packages.find]
38
+ where = ["src"]
@@ -0,0 +1,4 @@
1
+ [egg_info]
2
+ tag_build =
3
+ tag_date = 0
4
+
@@ -0,0 +1,157 @@
1
+ """Shape plugin: simulation scenarios (P6-04).
2
+
3
+ Each simulator lives in its own module and is also importable from here, lazily (importing the
4
+ package loads no simulator, and none of Arrow or NumPy):
5
+
6
+ * ``file_drop``: :class:`FileDropSimulator`, an upstream source landing files over a date range;
7
+ * ``scd2_file_drops``: :class:`SCD2FileDropSimulator`, a full load and daily versioned deltas;
8
+ * ``stream_emit``: :class:`StreamEmitter`, tables as enveloped events on the emit runtime;
9
+ * ``hybrid``: :class:`HybridSimulator`, a file drop and a stream of the same tables, linked;
10
+ * ``state_machine``: :class:`WorkflowSimulator`, business-process events with dwell times;
11
+ * ``clickstream_patterns``, ``financial_patterns``, ``iot_patterns``,
12
+ ``operational_log_patterns``, ``pulse_patterns``: the pattern simulators (sessions and funnels,
13
+ reversals and fraud bursts, sensor drift and alert storms, service logs and traces, rideshare
14
+ telemetry and marts), which take a configuration and, where they layer on existing data, Arrow
15
+ tables, and return Arrow tables and summary statistics;
16
+ * ``cli``, ``cli_batch``, ``cli_patterns``: ``shape simulate``, the command that runs them.
17
+ """
18
+
19
+ from __future__ import annotations
20
+
21
+ import importlib
22
+ from typing import TYPE_CHECKING, Any
23
+
24
+ SHAPE_API = "1.0"
25
+ """The plugin targets the plugin API major version this Shape provides."""
26
+
27
+ # name -> module that defines it (one line per name; each simulator adds its own lines)
28
+ _EXPORTS: dict[str, str] = {
29
+ "ClickstreamConfig": "shape_simulation.clickstream_patterns",
30
+ "ClickstreamResult": "shape_simulation.clickstream_patterns",
31
+ "ClickstreamSimulator": "shape_simulation.clickstream_patterns",
32
+ "FileDropConfig": "shape_simulation.file_drop",
33
+ "FileDropResult": "shape_simulation.file_drop",
34
+ "FileDropSimulator": "shape_simulation.file_drop",
35
+ "FinancialStreamConfig": "shape_simulation.financial_patterns",
36
+ "FinancialStreamResult": "shape_simulation.financial_patterns",
37
+ "FinancialStreamSimulator": "shape_simulation.financial_patterns",
38
+ "HybridConfig": "shape_simulation.hybrid",
39
+ "HybridResult": "shape_simulation.hybrid",
40
+ "HybridSimulator": "shape_simulation.hybrid",
41
+ "IoTTelemetryConfig": "shape_simulation.iot_patterns",
42
+ "IoTTelemetryResult": "shape_simulation.iot_patterns",
43
+ "IoTTelemetrySimulator": "shape_simulation.iot_patterns",
44
+ "OperationalLogConfig": "shape_simulation.operational_log_patterns",
45
+ "OperationalLogResult": "shape_simulation.operational_log_patterns",
46
+ "OperationalLogSimulator": "shape_simulation.operational_log_patterns",
47
+ "PulseDemandConfig": "shape_simulation.pulse_patterns",
48
+ "PulseDemandSimulator": "shape_simulation.pulse_patterns",
49
+ "PulseSimResult": "shape_simulation.pulse_patterns",
50
+ "SCD2FileDropConfig": "shape_simulation.scd2_file_drops",
51
+ "SCD2FileDropResult": "shape_simulation.scd2_file_drops",
52
+ "SCD2FileDropSimulator": "shape_simulation.scd2_file_drops",
53
+ "SimulateCommand": "shape_simulation.cli",
54
+ "StateDefinition": "shape_simulation.state_machine",
55
+ "StreamEmitConfig": "shape_simulation.stream_emit",
56
+ "StreamEmitResult": "shape_simulation.stream_emit",
57
+ "StreamEmitter": "shape_simulation.stream_emit",
58
+ "TransitionRule": "shape_simulation.state_machine",
59
+ "WorkflowConfig": "shape_simulation.state_machine",
60
+ "WorkflowResult": "shape_simulation.state_machine",
61
+ "WorkflowSimulator": "shape_simulation.state_machine",
62
+ "get_preset_workflow": "shape_simulation.state_machine",
63
+ }
64
+
65
+ __all__ = [
66
+ "ClickstreamConfig",
67
+ "ClickstreamResult",
68
+ "ClickstreamSimulator",
69
+ "FileDropConfig",
70
+ "FileDropResult",
71
+ "FileDropSimulator",
72
+ "FinancialStreamConfig",
73
+ "FinancialStreamResult",
74
+ "FinancialStreamSimulator",
75
+ "HybridConfig",
76
+ "HybridResult",
77
+ "HybridSimulator",
78
+ "IoTTelemetryConfig",
79
+ "IoTTelemetryResult",
80
+ "IoTTelemetrySimulator",
81
+ "OperationalLogConfig",
82
+ "OperationalLogResult",
83
+ "OperationalLogSimulator",
84
+ "PulseDemandConfig",
85
+ "PulseDemandSimulator",
86
+ "PulseSimResult",
87
+ "SCD2FileDropConfig",
88
+ "SCD2FileDropResult",
89
+ "SCD2FileDropSimulator",
90
+ "SimulateCommand",
91
+ "StateDefinition",
92
+ "StreamEmitConfig",
93
+ "StreamEmitResult",
94
+ "StreamEmitter",
95
+ "TransitionRule",
96
+ "WorkflowConfig",
97
+ "WorkflowResult",
98
+ "WorkflowSimulator",
99
+ "get_preset_workflow",
100
+ ]
101
+
102
+ if TYPE_CHECKING:
103
+ from shape_simulation.cli import SimulateCommand
104
+ from shape_simulation.clickstream_patterns import (
105
+ ClickstreamConfig,
106
+ ClickstreamResult,
107
+ ClickstreamSimulator,
108
+ )
109
+ from shape_simulation.file_drop import FileDropConfig, FileDropResult, FileDropSimulator
110
+ from shape_simulation.financial_patterns import (
111
+ FinancialStreamConfig,
112
+ FinancialStreamResult,
113
+ FinancialStreamSimulator,
114
+ )
115
+ from shape_simulation.hybrid import HybridConfig, HybridResult, HybridSimulator
116
+ from shape_simulation.iot_patterns import (
117
+ IoTTelemetryConfig,
118
+ IoTTelemetryResult,
119
+ IoTTelemetrySimulator,
120
+ )
121
+ from shape_simulation.operational_log_patterns import (
122
+ OperationalLogConfig,
123
+ OperationalLogResult,
124
+ OperationalLogSimulator,
125
+ )
126
+ from shape_simulation.pulse_patterns import (
127
+ PulseDemandConfig,
128
+ PulseDemandSimulator,
129
+ PulseSimResult,
130
+ )
131
+ from shape_simulation.scd2_file_drops import (
132
+ SCD2FileDropConfig,
133
+ SCD2FileDropResult,
134
+ SCD2FileDropSimulator,
135
+ )
136
+ from shape_simulation.state_machine import (
137
+ StateDefinition,
138
+ TransitionRule,
139
+ WorkflowConfig,
140
+ WorkflowResult,
141
+ WorkflowSimulator,
142
+ get_preset_workflow,
143
+ )
144
+ from shape_simulation.stream_emit import StreamEmitConfig, StreamEmitResult, StreamEmitter
145
+
146
+
147
+ def __getattr__(name: str) -> Any:
148
+ module = _EXPORTS.get(name)
149
+ if module is None:
150
+ raise AttributeError(f"module 'shape_simulation' has no attribute {name!r}")
151
+ value = getattr(importlib.import_module(module), name)
152
+ globals()[name] = value
153
+ return value
154
+
155
+
156
+ def __dir__() -> list[str]:
157
+ return sorted([*globals(), *_EXPORTS])
@@ -0,0 +1,321 @@
1
+ """Shared helpers of the pattern simulators (clickstream, financial, IoT, operational log, pulse).
2
+
3
+ The simulators work on Arrow tables (inputs and outputs) with numpy, so the plugin needs nothing
4
+ beyond Shape's own dependencies. Everything random comes from one ``numpy`` generator seeded by
5
+ the configuration, ids included: the same configuration and seed give the same tables.
6
+ """
7
+
8
+ from __future__ import annotations
9
+
10
+ import json
11
+ from collections.abc import Mapping
12
+ from datetime import UTC, datetime
13
+ from pathlib import Path
14
+ from typing import Any, ClassVar
15
+
16
+ import numpy as np
17
+ import pyarrow as pa # type: ignore[import-untyped]
18
+ import pyarrow.compute as pc # type: ignore[import-untyped]
19
+
20
+ FORMATS = ("parquet", "csv", "jsonl")
21
+
22
+
23
+ # ---- ids ------------------------------------------------------------------------------------
24
+
25
+
26
+ def uuid_strings(rng: np.random.Generator, n: int) -> list[str]:
27
+ """``n`` version-4 style UUID strings drawn from ``rng`` (so a seed reproduces them)."""
28
+ if n <= 0:
29
+ return []
30
+ raw = rng.integers(0, 256, size=(n, 16), dtype=np.uint8)
31
+ raw[:, 6] = (raw[:, 6] & 0x0F) | 0x40
32
+ raw[:, 8] = (raw[:, 8] & 0x3F) | 0x80
33
+ hexa = raw.tobytes().hex()
34
+ return [
35
+ f"{h[0:8]}-{h[8:12]}-{h[12:16]}-{h[16:20]}-{h[20:32]}"
36
+ for h in (hexa[i * 32 : (i + 1) * 32] for i in range(n))
37
+ ]
38
+
39
+
40
+ def hex_strings(rng: np.random.Generator, n: int, length: int) -> list[str]:
41
+ """``n`` random lowercase hex strings of ``length`` characters."""
42
+ if n <= 0:
43
+ return []
44
+ nbytes = (length + 1) // 2
45
+ raw = rng.integers(0, 256, size=(n, nbytes), dtype=np.uint8).tobytes().hex()
46
+ width = nbytes * 2
47
+ return [raw[i * width : i * width + length] for i in range(n)]
48
+
49
+
50
+ # ---- weighted draws -------------------------------------------------------------------------
51
+
52
+
53
+ def pick(
54
+ rng: np.random.Generator,
55
+ values: list[Any] | np.ndarray,
56
+ size: int,
57
+ p: list[float] | np.ndarray | None = None,
58
+ ) -> np.ndarray:
59
+ """``size`` draws from ``values`` (uniform, or with probabilities ``p``, normalised)."""
60
+ if p is not None:
61
+ weights = np.asarray(p, dtype=np.float64)
62
+ weights = weights / weights.sum()
63
+ return np.asarray(values)[rng.choice(len(values), size=size, p=weights)]
64
+ return np.asarray(values)[rng.integers(0, len(values), size=size)]
65
+
66
+
67
+ # ---- Arrow in -------------------------------------------------------------------------------
68
+
69
+
70
+ def as_table(obj: Any) -> pa.Table:
71
+ """An Arrow table from a table, a record batch, a mapping of columns or a pandas frame."""
72
+ if isinstance(obj, pa.Table):
73
+ return obj
74
+ if isinstance(obj, pa.RecordBatch):
75
+ return pa.Table.from_batches([obj])
76
+ if isinstance(obj, Mapping):
77
+ return pa.table(dict(obj))
78
+ to_arrow = getattr(obj, "to_arrow", None)
79
+ if callable(to_arrow):
80
+ out = to_arrow()
81
+ if isinstance(out, pa.Table):
82
+ return out
83
+ if hasattr(obj, "__arrow_c_stream__"):
84
+ return pa.table(obj)
85
+ if hasattr(obj, "columns") and hasattr(obj, "dtypes"): # a pandas frame
86
+ return pa.Table.from_pandas(obj, preserve_index=False)
87
+ raise TypeError(f"cannot read {type(obj).__name__} as an Arrow table")
88
+
89
+
90
+ def table_mapping(tables: Any) -> dict[str, pa.Table]:
91
+ """``tables`` (a mapping of tables, or a generation result with ``.tables``) as a dict."""
92
+ inner = getattr(tables, "tables", tables)
93
+ if not isinstance(inner, Mapping):
94
+ raise TypeError("tables must map table names to tables")
95
+ return {str(k): as_table(v) for k, v in inner.items()}
96
+
97
+
98
+ def _is_number(value: Any) -> bool:
99
+ return not isinstance(value, bool) and isinstance(value, int | float | np.integer | np.floating)
100
+
101
+
102
+ def check_settings(
103
+ config: Any,
104
+ *,
105
+ positive: tuple[str, ...] = (),
106
+ non_negative: tuple[str, ...] = (),
107
+ probabilities: tuple[str, ...] = (),
108
+ non_empty: tuple[str, ...] = (),
109
+ ) -> None:
110
+ """Refuse settings a simulator cannot use, with a ``ValueError`` that names the setting.
111
+
112
+ Every field whose default is a number must be a number (``--set users=abc`` gives text);
113
+ ``positive`` ones are above 0 (they divide), ``non_negative`` ones 0 or more,
114
+ ``probabilities`` in [0, 1] and ``non_empty`` lists hold at least one item. ``None`` is
115
+ accepted where it is the default (an unset optional)."""
116
+ import dataclasses
117
+
118
+ for f in dataclasses.fields(config):
119
+ default = f.default
120
+ if isinstance(default, bool) or not isinstance(default, int | float):
121
+ continue
122
+ value = getattr(config, f.name)
123
+ if not _is_number(value):
124
+ raise ValueError(f"{f.name} must be a number, got {value!r}")
125
+ for names, strict, high, what in (
126
+ (positive, True, None, "above 0"),
127
+ (non_negative, False, None, "0 or more"),
128
+ (probabilities, False, 1.0, "between 0 and 1"),
129
+ ):
130
+ for name in names:
131
+ value = getattr(config, name)
132
+ if value is None:
133
+ continue
134
+ if not _is_number(value):
135
+ raise ValueError(f"{name} must be a number, got {value!r}")
136
+ number = float(value)
137
+ # written so that NaN fails every rule
138
+ inside = number > 0 if strict else number >= 0
139
+ if high is not None:
140
+ inside = inside and number <= high
141
+ if not inside:
142
+ raise ValueError(f"{name} must be {what}, got {value!r}")
143
+ for name in non_empty:
144
+ if not getattr(config, name):
145
+ raise ValueError(f"{name} must hold at least one item")
146
+
147
+
148
+ def timestamp_us(col: pa.ChunkedArray | pa.Array) -> tuple[np.ndarray, np.ndarray, str | None]:
149
+ """A date, timestamp or ISO-8601 string column as ``(microseconds, valid, timezone)``.
150
+
151
+ ``microseconds`` is int64 since the epoch (0 where null), ``valid`` marks the non-null
152
+ cells and ``timezone`` is the column's time zone, if it had one.
153
+ """
154
+ if isinstance(col, pa.ChunkedArray):
155
+ col = col.combine_chunks() if col.num_chunks else pa.array([], type=col.type)
156
+ tz: str | None = None
157
+ if pa.types.is_timestamp(col.type):
158
+ tz = col.type.tz
159
+ col = col.cast(pa.timestamp("us", tz))
160
+ elif pa.types.is_date(col.type):
161
+ col = col.cast(pa.timestamp("us"))
162
+ elif pa.types.is_string(col.type) or pa.types.is_large_string(col.type):
163
+ col = _text_times(col)
164
+ else:
165
+ raise TypeError(f"{col.type} is not a date or timestamp type")
166
+ valid = ~np.asarray(col.is_null().to_numpy(zero_copy_only=False), dtype=bool)
167
+ ints = col.cast(pa.int64()).fill_null(0).to_numpy(zero_copy_only=False)
168
+ return ints.astype(np.int64), valid, tz
169
+
170
+
171
+ def _text_times(col: pa.Array) -> pa.Array:
172
+ """ISO-8601 text as naive ``timestamp[us]``: a time with ``Z`` or an offset becomes its UTC
173
+ wall time, one without is taken as it is. Text that is not a time raises ``ValueError``."""
174
+ try:
175
+ return pc.cast(col, pa.timestamp("us"))
176
+ except pa.ArrowInvalid:
177
+ pass
178
+ try: # every time carries a zone
179
+ return pc.cast(col, pa.timestamp("us", "UTC")).cast(pa.timestamp("us"))
180
+ except pa.ArrowInvalid:
181
+ pass
182
+ values: list[datetime | None] = []
183
+ for text in col.to_pylist(): # some with a zone, some without
184
+ if text is None:
185
+ values.append(None)
186
+ continue
187
+ try:
188
+ value = datetime.fromisoformat(text)
189
+ except ValueError:
190
+ raise ValueError(f"{text!r} is not an ISO-8601 date or time") from None
191
+ if value.tzinfo is not None:
192
+ value = value.astimezone(UTC).replace(tzinfo=None)
193
+ values.append(value)
194
+ return pa.array(values, pa.timestamp("us"))
195
+
196
+
197
+ def timestamps(us: np.ndarray, tz: str | None = None, valid: np.ndarray | None = None) -> pa.Array:
198
+ """A ``timestamp[us]`` array (``tz`` if given) from microseconds, null where not ``valid``."""
199
+ mask = None if valid is None else ~np.asarray(valid, dtype=bool)
200
+ return pa.array(
201
+ np.asarray(us, dtype=np.int64).astype("datetime64[us]"),
202
+ type=pa.timestamp("us", tz),
203
+ mask=mask,
204
+ )
205
+
206
+
207
+ def parse_start(text: str) -> int:
208
+ """An ISO-8601 instant as microseconds since the epoch (a missing zone means UTC)."""
209
+ moment = datetime.fromisoformat(text)
210
+ if moment.tzinfo is None:
211
+ moment = moment.replace(tzinfo=UTC)
212
+ return int(moment.timestamp() * 1_000_000)
213
+
214
+
215
+ def float_values(col: pa.ChunkedArray | pa.Array) -> tuple[np.ndarray, np.ndarray]:
216
+ """A numeric column as ``(float64 values, valid)``; null cells are NaN."""
217
+ arr = pc.cast(col, pa.float64())
218
+ if isinstance(arr, pa.ChunkedArray):
219
+ arr = arr.combine_chunks() if arr.num_chunks else pa.array([], type=pa.float64())
220
+ valid = ~np.asarray(arr.is_null().to_numpy(zero_copy_only=False), dtype=bool)
221
+ values = arr.fill_null(float("nan")).to_numpy(zero_copy_only=False).astype(np.float64)
222
+ return values, valid
223
+
224
+
225
+ def float_array(values: np.ndarray) -> pa.Array:
226
+ """A float64 Arrow array from ``values``, null where the value is NaN."""
227
+ values = np.asarray(values, dtype=np.float64)
228
+ return pa.array(values, type=pa.float64(), mask=np.isnan(values))
229
+
230
+
231
+ def empty_table(schema: pa.Schema) -> pa.Table:
232
+ return schema.empty_table()
233
+
234
+
235
+ def combine(tables: list[pa.Table]) -> pa.Table:
236
+ """The rows of ``tables`` stacked, with the union of their columns (missing cells are null).
237
+
238
+ Where the same column has different types (decimal and float amounts, say) numeric
239
+ columns become float64 and anything else text.
240
+ """
241
+ names: list[str] = []
242
+ types: dict[str, pa.DataType] = {}
243
+ for table in tables:
244
+ for field in table.schema:
245
+ if field.name not in types:
246
+ names.append(field.name)
247
+ types[field.name] = field.type
248
+ elif types[field.name] != field.type:
249
+ a, b = types[field.name], field.type
250
+ numeric = all(
251
+ pa.types.is_integer(t) or pa.types.is_floating(t) or pa.types.is_decimal(t)
252
+ for t in (a, b)
253
+ )
254
+ types[field.name] = pa.float64() if numeric else pa.string()
255
+ schema = pa.schema([pa.field(n, types[n]) for n in names])
256
+ parts = []
257
+ for table in tables:
258
+ cols = []
259
+ for name in names:
260
+ if name in table.column_names:
261
+ cols.append(pc.cast(table.column(name), types[name]))
262
+ else:
263
+ cols.append(pa.nulls(table.num_rows, type=types[name]))
264
+ parts.append(pa.Table.from_arrays(cols, schema=schema))
265
+ return pa.concat_tables(parts)
266
+
267
+
268
+ def value_counts(values: np.ndarray) -> dict[str, int]:
269
+ """``{value: count}`` ordered by count (largest first; ties by first appearance)."""
270
+ if len(values) == 0:
271
+ return {}
272
+ uniq, first, counts = np.unique(values, return_index=True, return_counts=True)
273
+ order = np.lexsort((first, -counts))
274
+ return {str(uniq[i]): int(counts[i]) for i in order}
275
+
276
+
277
+ # ---- results --------------------------------------------------------------------------------
278
+
279
+
280
+ class TablesResult:
281
+ """What every simulator's result offers: its tables, as flat events, and on disk."""
282
+
283
+ TABLES: ClassVar[tuple[str, ...]] = ()
284
+ stats: dict[str, Any]
285
+
286
+ def table_map(self) -> dict[str, pa.Table]:
287
+ """Every output table by name."""
288
+ return {name: getattr(self, name) for name in self.TABLES}
289
+
290
+ def events(self, table: str, *, seq_start: int = 0) -> pa.RecordBatch:
291
+ """The rows of ``table`` as flat stream events: its columns plus ``_shape_table``,
292
+ ``_shape_seq`` and (when it has a date or timestamp column) ``_shape_event_time``.
293
+ Feed them to ``shape.streaming.emit`` sinks or an emitter."""
294
+ from shape.streaming.emit.formats import with_event_fields
295
+
296
+ source = self.table_map().get(table)
297
+ if source is None:
298
+ raise KeyError(f"no table {table!r}; the result has {', '.join(self.table_map())}")
299
+ batch = source.combine_chunks().to_batches()
300
+ merged = batch[0] if batch else pa.RecordBatch.from_pylist([], schema=source.schema)
301
+ events: pa.RecordBatch = with_event_fields(merged, table, seq_start)
302
+ return events
303
+
304
+ def write(self, directory: str | Path, fmt: str = "parquet") -> dict[str, Path]:
305
+ """Write every table to ``directory`` as ``<table>.<fmt>`` (parquet, csv or jsonl)
306
+ and return the paths."""
307
+ if fmt not in FORMATS:
308
+ raise ValueError(f"unknown format {fmt!r}; choose from {', '.join(FORMATS)}")
309
+ from shape.builtins.sinks.files import CsvSink, JsonlSink, ParquetSink
310
+
311
+ sink: Any = {"parquet": ParquetSink, "csv": CsvSink, "jsonl": JsonlSink}[fmt]()
312
+ out = Path(directory)
313
+ out.mkdir(parents=True, exist_ok=True)
314
+ written: dict[str, Path] = {}
315
+ for name, table in self.table_map().items():
316
+ sink.write(str(out) + "/", name, table.to_batches(), schema=table.schema)
317
+ written[name] = out / f"{name}.{fmt}"
318
+ (out / "stats.json").write_text(
319
+ json.dumps(self.stats, indent=2, default=str) + "\n", encoding="utf-8", newline="\n"
320
+ )
321
+ return written