gdb-client 0.3.0__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.
@@ -0,0 +1,8 @@
1
+ __pycache__/
2
+ *.pyc
3
+ *.egg-info/
4
+ dist/
5
+ build/
6
+ .pytest_cache/
7
+ .venv/
8
+ env/
@@ -0,0 +1,176 @@
1
+ Apache License
2
+ Version 2.0, January 2004
3
+ http://www.apache.org/licenses/
4
+
5
+ TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
6
+
7
+ 1. Definitions.
8
+
9
+ "License" shall mean the terms and conditions for use, reproduction,
10
+ and distribution as defined by Sections 1 through 9 of this document.
11
+
12
+ "Licensor" shall mean the copyright owner or entity authorized by
13
+ the copyright owner that is granting the License.
14
+
15
+ "Legal Entity" shall mean the union of the acting entity and all
16
+ other entities that control, are controlled by, or are under common
17
+ control with that entity. For the purposes of this definition,
18
+ "control" means (i) the power, direct or indirect, to cause the
19
+ direction or management of such entity, whether by contract or
20
+ otherwise, or (ii) ownership of fifty percent (50%) or more of the
21
+ outstanding shares, or (iii) beneficial ownership of such entity.
22
+
23
+ "You" (or "Your") shall mean an individual or Legal Entity
24
+ exercising permissions granted by this License.
25
+
26
+ "Source" form shall mean the preferred form for making modifications,
27
+ including but not limited to software source code, documentation
28
+ source, and configuration files.
29
+
30
+ "Object" form shall mean any form resulting from mechanical
31
+ transformation or translation of a Source form, including but
32
+ not limited to compiled object code, generated documentation,
33
+ and conversions to other media types.
34
+
35
+ "Work" shall mean the work of authorship, whether in Source or
36
+ Object form, made available under the License, as indicated by a
37
+ copyright notice that is included in or attached to the work
38
+ (an example is provided in the Appendix below).
39
+
40
+ "Derivative Works" shall mean any work, whether in Source or Object
41
+ form, that is based on (or derived from) the Work and for which the
42
+ editorial revisions, annotations, elaborations, or other modifications
43
+ represent, as a whole, an original work of authorship. For the purposes
44
+ of this License, Derivative Works shall not include works that remain
45
+ separable from, or merely link (or bind by name) to the interfaces of,
46
+ the Work and Derivative Works thereof.
47
+
48
+ "Contribution" shall mean any work of authorship, including
49
+ the original version of the Work and any modifications or additions
50
+ to that Work or Derivative Works thereof, that is intentionally
51
+ submitted to Licensor for inclusion in the Work by the copyright owner
52
+ or by an individual or Legal Entity authorized to submit on behalf of
53
+ the copyright owner. For the purposes of this definition, "submitted"
54
+ means any form of electronic, verbal, or written communication sent
55
+ to the Licensor or its representatives, including but not limited to
56
+ communication on electronic mailing lists, source code control systems,
57
+ and issue tracking systems that are managed by, or on behalf of, the
58
+ Licensor for the purpose of discussing and improving the Work, but
59
+ excluding communication that is conspicuously marked or otherwise
60
+ designated in writing by the copyright owner as "Not a Contribution."
61
+
62
+ "Contributor" shall mean Licensor and any individual or Legal Entity
63
+ on behalf of whom a Contribution has been received by Licensor and
64
+ subsequently incorporated within the Work.
65
+
66
+ 2. Grant of Copyright License. Subject to the terms and conditions of
67
+ this License, each Contributor hereby grants to You a perpetual,
68
+ worldwide, non-exclusive, no-charge, royalty-free, irrevocable
69
+ copyright license to reproduce, prepare Derivative Works of,
70
+ publicly display, publicly perform, sublicense, and distribute the
71
+ Work and such Derivative Works in Source or Object form.
72
+
73
+ 3. Grant of Patent License. Subject to the terms and conditions of
74
+ this License, each Contributor hereby grants to You a perpetual,
75
+ worldwide, non-exclusive, no-charge, royalty-free, irrevocable
76
+ (except as stated in this section) patent license to make, have made,
77
+ use, offer to sell, sell, import, and otherwise transfer the Work,
78
+ where such license applies only to those patent claims licensable
79
+ by such Contributor that are necessarily infringed by their
80
+ Contribution(s) alone or by combination of their Contribution(s)
81
+ with the Work to which such Contribution(s) was submitted. If You
82
+ institute patent litigation against any entity (including a
83
+ cross-claim or counterclaim in a lawsuit) alleging that the Work
84
+ or a Contribution incorporated within the Work constitutes direct
85
+ or contributory patent infringement, then any patent licenses
86
+ granted to You under this License for that Work shall terminate
87
+ as of the date such litigation is filed.
88
+
89
+ 4. Redistribution. You may reproduce and distribute copies of the
90
+ Work or Derivative Works thereof in any medium, with or without
91
+ modifications, and in Source or Object form, provided that You
92
+ meet the following conditions:
93
+
94
+ (a) You must give any other recipients of the Work or
95
+ Derivative Works a copy of this License; and
96
+
97
+ (b) You must cause any modified files to carry prominent notices
98
+ stating that You changed the files; and
99
+
100
+ (c) You must retain, in the Source form of any Derivative Works
101
+ that You distribute, all copyright, patent, trademark, and
102
+ attribution notices from the Source form of the Work,
103
+ excluding those notices that do not pertain to any part of
104
+ the Derivative Works; and
105
+
106
+ (d) If the Work includes a "NOTICE" text file as part of its
107
+ distribution, then any Derivative Works that You distribute must
108
+ include a readable copy of the attribution notices contained
109
+ within such NOTICE file, excluding those notices that do not
110
+ pertain to any part of the Derivative Works, in at least one
111
+ of the following places: within a NOTICE text file distributed
112
+ as part of the Derivative Works; within the Source form or
113
+ documentation, if provided along with the Derivative Works; or,
114
+ within a display generated by the Derivative Works, if and
115
+ wherever such third-party notices normally appear. The contents
116
+ of the NOTICE file are for informational purposes only and
117
+ do not modify the License. You may add Your own attribution
118
+ notices within Derivative Works that You distribute, alongside
119
+ or as an addendum to the NOTICE text from the Work, provided
120
+ that such additional attribution notices cannot be construed
121
+ as modifying the License.
122
+
123
+ You may add Your own copyright statement to Your modifications and
124
+ may provide additional or different license terms and conditions
125
+ for use, reproduction, or distribution of Your modifications, or
126
+ for any such Derivative Works as a whole, provided Your use,
127
+ reproduction, and distribution of the Work otherwise complies with
128
+ the conditions stated in this License.
129
+
130
+ 5. Submission of Contributions. Unless You explicitly state otherwise,
131
+ any Contribution intentionally submitted for inclusion in the Work
132
+ by You to the Licensor shall be under the terms and conditions of
133
+ this License, without any additional terms or conditions.
134
+ Notwithstanding the above, nothing herein shall supersede or modify
135
+ the terms of any separate license agreement you may have executed
136
+ with Licensor regarding such Contributions.
137
+
138
+ 6. Trademarks. This License does not grant permission to use the trade
139
+ names, trademarks, service marks, or product names of the Licensor,
140
+ except as required for reasonable and customary use in describing the
141
+ origin of the Work and reproducing the content of the NOTICE file.
142
+
143
+ 7. Disclaimer of Warranty. Unless required by applicable law or
144
+ agreed to in writing, Licensor provides the Work (and each
145
+ Contributor provides its Contributions) on an "AS IS" BASIS,
146
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
147
+ implied, including, without limitation, any warranties or conditions
148
+ of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
149
+ PARTICULAR PURPOSE. You are solely responsible for determining the
150
+ appropriateness of using or redistributing the Work and assume any
151
+ risks associated with Your exercise of permissions under this License.
152
+
153
+ 8. Limitation of Liability. In no event and under no legal theory,
154
+ whether in tort (including negligence), contract, or otherwise,
155
+ unless required by applicable law (such as deliberate and grossly
156
+ negligent acts) or agreed to in writing, shall any Contributor be
157
+ liable to You for damages, including any direct, indirect, special,
158
+ incidental, or exemplary damages of any character arising as a
159
+ result of this License or out of the use or inability to use the
160
+ Work (including but not limited to damages for loss of goodwill,
161
+ work stoppage, computer failure or malfunction, or any and all
162
+ other commercial damages or losses), even if such Contributor
163
+ has been advised of the possibility of such damages.
164
+
165
+ 9. Accepting Warranty or Additional Liability. While redistributing
166
+ the Work or Derivative Works thereof, You may choose to offer,
167
+ and charge a fee for, acceptance of support, warranty, indemnity,
168
+ or other liability obligations and/or rights consistent with this
169
+ License. However, in accepting such obligations, You may act only
170
+ on Your own behalf and on Your sole responsibility, not on behalf
171
+ of any other Contributor, and only if You agree to indemnify,
172
+ defend, and hold each Contributor harmless for any liability
173
+ incurred by, or claims asserted against, such Contributor by reason
174
+ of your accepting any such warranty or additional liability.
175
+
176
+ END OF TERMS AND CONDITIONS
@@ -0,0 +1,231 @@
1
+ Metadata-Version: 2.5
2
+ Name: gdb-client
3
+ Version: 0.3.0
4
+ Summary: High-performance Python client for GDB graph database with Polars parallel scatter-ingest via Apache Arrow Flight
5
+ Project-URL: Homepage, https://github.com/1orgar/gdb
6
+ Project-URL: Repository, https://github.com/1orgar/gdb-py-client
7
+ Project-URL: Documentation, https://github.com/1orgar/gdb/blob/main/docs/SERVER_AND_QUERY_GUIDE.md
8
+ Project-URL: Issues, https://github.com/1orgar/gdb-py-client/issues
9
+ Author-email: GDB Team <dev@gdb.io>
10
+ License-Expression: Apache-2.0
11
+ License-File: LICENSE
12
+ Keywords: apache-arrow,arrow-flight,cypher,gdb,graph-database,htap,mpp,opencypher,polars,scatter-ingest
13
+ Classifier: Development Status :: 4 - Beta
14
+ Classifier: Intended Audience :: Developers
15
+ Classifier: Intended Audience :: Science/Research
16
+ Classifier: License :: OSI Approved :: Apache Software License
17
+ Classifier: Operating System :: OS Independent
18
+ Classifier: Programming Language :: Python :: 3
19
+ Classifier: Programming Language :: Python :: 3.9
20
+ Classifier: Programming Language :: Python :: 3.10
21
+ Classifier: Programming Language :: Python :: 3.11
22
+ Classifier: Programming Language :: Python :: 3.12
23
+ Classifier: Programming Language :: Python :: 3.13
24
+ Classifier: Topic :: Database
25
+ Classifier: Topic :: Database :: Database Engines/Servers
26
+ Classifier: Topic :: Scientific/Engineering :: Information Analysis
27
+ Requires-Python: >=3.9
28
+ Requires-Dist: polars>=1.0.0
29
+ Requires-Dist: pyarrow>=14.0.0
30
+ Requires-Dist: requests>=2.28.0
31
+ Provides-Extra: dev
32
+ Requires-Dist: build>=1.0.0; extra == 'dev'
33
+ Requires-Dist: pytest>=7.0.0; extra == 'dev'
34
+ Requires-Dist: twine>=5.0.0; extra == 'dev'
35
+ Description-Content-Type: text/markdown
36
+
37
+ # GDB Python Client (`gdb-client`)
38
+
39
+ [![PyPI version](https://img.shields.io/pypi/v/gdb-client.svg)](https://pypi.org/project/gdb-client/)
40
+ [![Python versions](https://img.shields.io/pypi/pyversions/gdb-client.svg)](https://pypi.org/project/gdb-client/)
41
+ [![License](https://img.shields.io/badge/License-Apache_2.0-blue.svg)](https://opensource.org/licenses/Apache-2.0)
42
+ [![CI](https://github.com/1orgar/gdb-py-client/actions/workflows/ci.yml/badge.svg)](https://github.com/1orgar/gdb-py-client/actions/workflows/ci.yml)
43
+
44
+ Official high-performance Python client for **[GDB (Distributed HTAP In-Memory Graph Database)](https://github.com/1orgar/gdb)**.
45
+
46
+ Powered by **Polars** SIMD vectorized expressions and **Apache Arrow Flight** high-speed TCP streaming, `gdb-client` delivers ultra-high-throughput parallel scatter-ingest directly into in-memory partition stores of cluster nodes.
47
+
48
+ ---
49
+
50
+ ## ⚡ Key Features
51
+
52
+ - **Polars-Powered Scatter-Ingest**: Automatically partitions massive DataFrames using Polars SIMD expressions and streams partitions in parallel directly to primary replica nodes in the cluster hash ring.
53
+ - **Dedicated Arrow Flight Ingestion Ports (`:8860+`)**: Bypasses HTTP JSON serialization overhead, streaming binary Arrow `RecordBatch` payloads directly into GDB's in-memory `DeltaMemTable`.
54
+ - **Automatic Cluster Hash Ring Discovery**: Queries cluster coordinator topology (`/cluster`) to resolve ring tokens ($u \pmod N$), partition mapping, and per-node Arrow Flight endpoints.
55
+ - **Native Polars Query Results**: Runs openCypher / GQL queries and yields results directly as `polars.DataFrame` via Arrow Flight `do_get` with zero-copy deserialization.
56
+ - **Graceful Fallbacks**: Works out of the box with standard library fallbacks if optional dependencies are missing.
57
+
58
+ ---
59
+
60
+ ## 📦 Installation
61
+
62
+ Install from PyPI:
63
+
64
+ ```bash
65
+ pip install gdb-client
66
+ ```
67
+
68
+ Or install the latest development version directly from GitHub:
69
+
70
+ ```bash
71
+ pip install git+https://github.com/1orgar/gdb-py-client.git
72
+ ```
73
+
74
+ ### Requirements
75
+
76
+ - Python >= 3.9
77
+ - `polars >= 1.0.0`
78
+ - `pyarrow >= 14.0.0`
79
+ - `requests >= 2.28.0`
80
+
81
+ ---
82
+
83
+ ## 🚀 Quickstart
84
+
85
+ ```python
86
+ import polars as pl
87
+ from gdb_client import GdbClient
88
+
89
+ # 1. Connect to any cluster node (auto-discovers all ring peers and Flight ports)
90
+ client = GdbClient(seed_url="http://localhost:8847")
91
+
92
+ # 2. Ingest 1,000,000 vertices in parallel directly into target nodes
93
+ users_df = pl.DataFrame({
94
+ "id": range(1, 1_000_001),
95
+ "name": [f"User_{i}" for i in range(1, 1_000_001)],
96
+ "age": [20 + (i % 50) for i in range(1, 1_000_001)],
97
+ })
98
+
99
+ res_v = client.scatter_ingest_vertices(users_df, label="User", id_col="id")
100
+ print(f"Vertices: {res_v['rows_ingested']} rows in {res_v['elapsed_seconds']}s "
101
+ f"({res_v['rows_per_second']:,.0f} rows/s)")
102
+
103
+ # 3. Ingest edges in parallel (automatically hashed by source vertex)
104
+ edges_df = pl.DataFrame({
105
+ "src": range(1, 500_001),
106
+ "dst": range(500_001, 1_000_001),
107
+ })
108
+
109
+ res_e = client.scatter_ingest_edges(edges_df, edge_type="FOLLOWS")
110
+ print(f"Edges: {res_e['rows_ingested']} rows in {res_e['elapsed_seconds']}s "
111
+ f"({res_e['rows_per_second']:,.0f} rows/s)")
112
+
113
+ # 4. Trigger background compaction into immutable Chunked-CSR
114
+ client.compact()
115
+
116
+ # 5. Query graph using openCypher directly into a Polars DataFrame
117
+ df = client.query("MATCH (a:User)-[:FOLLOWS]->(b:User) RETURN a.name, b.name LIMIT 10;")
118
+ print(df)
119
+
120
+ # 6. Close Flight TCP connections
121
+ client.close()
122
+ ```
123
+
124
+ ---
125
+
126
+ ## 🏗 Architecture: How Scatter-Ingest Works
127
+
128
+ In a traditional database setup, client libraries send bulk requests to a single coordinator or load balancer, which must unpack the data, hash every key, and re-transmit records over the network to the correct storage nodes.
129
+
130
+ `gdb-client` moves partition routing to the client side:
131
+
132
+ ```
133
+ ┌─────────────────────────────────────────┐
134
+ │ polars.DataFrame │
135
+ │ (10,000,000 Rows in RAM) │
136
+ └────────────────────┬────────────────────┘
137
+ │
138
+ Polars SIMD Hash: (col % TotalNodes)
139
+ │
140
+ ┌──────────────────────────┼──────────────────────────┐
141
+ ▼ ▼ ▼
142
+ [Partition Chunk 0] [Partition Chunk 1] [Partition Chunk 2]
143
+ │ │ │
144
+ PyArrow RecordBatch PyArrow RecordBatch PyArrow RecordBatch
145
+ │ │ │
146
+ Arrow Flight do_put Arrow Flight do_put Arrow Flight do_put
147
+ (TCP Port :8860) (TCP Port :8861) (TCP Port :8862)
148
+ ▼ ▼ ▼
149
+ ┌─────────────────┐ ┌─────────────────┐ ┌─────────────────┐
150
+ │ GDB Node #1 │ │ GDB Node #2 │ │ GDB Node #3 │
151
+ │ DeltaMemTable │ │ DeltaMemTable │ │ DeltaMemTable │
152
+ └─────────────────┘ └─────────────────┘ └─────────────────┘
153
+ ```
154
+
155
+ 1. **Client-Side Vectorized Hashing**: Polars computes target partitions for all rows simultaneously using native CPU SIMD instructions.
156
+ 2. **Direct Multi-Node Streaming**: Partitions are streamed over concurrent TCP Flight channels directly to the responsible nodes' `--client-flight-port` (`:8860`, `:8861`, `:8862`, etc.).
157
+ 3. **Zero Coordinator Bottleneck**: Cluster inter-node network bandwidth is preserved for MPP shuffle exchange and graph traversal.
158
+
159
+ ---
160
+
161
+ ## 📖 API Reference
162
+
163
+ ### `GdbClient`
164
+
165
+ ```python
166
+ GdbClient(seed_url: str = "http://localhost:8847", auth_token: Optional[str] = None)
167
+ ```
168
+
169
+ - **`seed_url`**: HTTP address of any healthy node in the GDB cluster.
170
+ - **`auth_token`**: Optional Bearer token for clusters with authorization enabled.
171
+
172
+ #### Methods
173
+
174
+ | Method | Parameters | Return Type | Description |
175
+ | :--- | :--- | :--- | :--- |
176
+ | `scatter_ingest_vertices` | `df: pl.DataFrame, label: str = "User", id_col: str = "id", max_workers: int = 8` | `dict` | Partitions vertices and streams Arrow RecordBatches directly to target nodes. |
177
+ | `scatter_ingest_edges` | `df: pl.DataFrame, edge_type: str = "KNOWS", src_col: str = "src", dst_col: str = "dst", rank_col: Optional[str] = None, max_workers: int = 8` | `dict` | Partitions edges by source ID and streams Arrow RecordBatches to target nodes. |
178
+ | `query` | `cypher: str, node_id: Optional[int] = None` | `pl.DataFrame` | Executes an openCypher query and returns the tabular result as a Polars DataFrame. |
179
+ | `compact` | — | `dict` | Triggers asynchronous CSR compaction across cluster nodes. |
180
+ | `refresh_topology` | — | `None` | Queries `/cluster` and reconnects Flight clients to active nodes. |
181
+ | `close` | — | `None` | Closes open Arrow Flight TCP client connections. |
182
+
183
+ ---
184
+
185
+ ### Ingestion Performance Metrics
186
+
187
+ `scatter_ingest_vertices` and `scatter_ingest_edges` return a detailed performance dictionary:
188
+
189
+ ```python
190
+ {
191
+ "rows_ingested": 1000000,
192
+ "elapsed_seconds": 0.4281,
193
+ "rows_per_second": 2335902.8,
194
+ "partitions": 3
195
+ }
196
+ ```
197
+
198
+ ---
199
+
200
+ ## 🛠 Local Development & Building Wheel
201
+
202
+ To build the wheel and source distribution locally:
203
+
204
+ ```bash
205
+ # Clone the repository
206
+ git clone https://github.com/1orgar/gdb-py-client.git
207
+ cd gdb-py-client
208
+
209
+ # Install build dependencies
210
+ pip install -e ".[dev]"
211
+
212
+ # Run tests
213
+ python -m unittest discover tests
214
+
215
+ # Build .whl and .tar.gz distributions
216
+ python -m build
217
+
218
+ # Validate package metadata with twine
219
+ twine check dist/*
220
+ ```
221
+
222
+ Built artifacts will be located in `dist/`:
223
+ - `dist/gdb_client-0.3.0-py3-none-any.whl`
224
+ - `dist/gdb_client-0.3.0.tar.gz`
225
+
226
+
227
+ ---
228
+
229
+ ## 📄 License
230
+
231
+ Licensed under the Apache License, Version 2.0. See the [LICENSE](LICENSE) file for details.
@@ -0,0 +1,195 @@
1
+ # GDB Python Client (`gdb-client`)
2
+
3
+ [![PyPI version](https://img.shields.io/pypi/v/gdb-client.svg)](https://pypi.org/project/gdb-client/)
4
+ [![Python versions](https://img.shields.io/pypi/pyversions/gdb-client.svg)](https://pypi.org/project/gdb-client/)
5
+ [![License](https://img.shields.io/badge/License-Apache_2.0-blue.svg)](https://opensource.org/licenses/Apache-2.0)
6
+ [![CI](https://github.com/1orgar/gdb-py-client/actions/workflows/ci.yml/badge.svg)](https://github.com/1orgar/gdb-py-client/actions/workflows/ci.yml)
7
+
8
+ Official high-performance Python client for **[GDB (Distributed HTAP In-Memory Graph Database)](https://github.com/1orgar/gdb)**.
9
+
10
+ Powered by **Polars** SIMD vectorized expressions and **Apache Arrow Flight** high-speed TCP streaming, `gdb-client` delivers ultra-high-throughput parallel scatter-ingest directly into in-memory partition stores of cluster nodes.
11
+
12
+ ---
13
+
14
+ ## ⚡ Key Features
15
+
16
+ - **Polars-Powered Scatter-Ingest**: Automatically partitions massive DataFrames using Polars SIMD expressions and streams partitions in parallel directly to primary replica nodes in the cluster hash ring.
17
+ - **Dedicated Arrow Flight Ingestion Ports (`:8860+`)**: Bypasses HTTP JSON serialization overhead, streaming binary Arrow `RecordBatch` payloads directly into GDB's in-memory `DeltaMemTable`.
18
+ - **Automatic Cluster Hash Ring Discovery**: Queries cluster coordinator topology (`/cluster`) to resolve ring tokens ($u \pmod N$), partition mapping, and per-node Arrow Flight endpoints.
19
+ - **Native Polars Query Results**: Runs openCypher / GQL queries and yields results directly as `polars.DataFrame` via Arrow Flight `do_get` with zero-copy deserialization.
20
+ - **Graceful Fallbacks**: Works out of the box with standard library fallbacks if optional dependencies are missing.
21
+
22
+ ---
23
+
24
+ ## 📦 Installation
25
+
26
+ Install from PyPI:
27
+
28
+ ```bash
29
+ pip install gdb-client
30
+ ```
31
+
32
+ Or install the latest development version directly from GitHub:
33
+
34
+ ```bash
35
+ pip install git+https://github.com/1orgar/gdb-py-client.git
36
+ ```
37
+
38
+ ### Requirements
39
+
40
+ - Python >= 3.9
41
+ - `polars >= 1.0.0`
42
+ - `pyarrow >= 14.0.0`
43
+ - `requests >= 2.28.0`
44
+
45
+ ---
46
+
47
+ ## 🚀 Quickstart
48
+
49
+ ```python
50
+ import polars as pl
51
+ from gdb_client import GdbClient
52
+
53
+ # 1. Connect to any cluster node (auto-discovers all ring peers and Flight ports)
54
+ client = GdbClient(seed_url="http://localhost:8847")
55
+
56
+ # 2. Ingest 1,000,000 vertices in parallel directly into target nodes
57
+ users_df = pl.DataFrame({
58
+ "id": range(1, 1_000_001),
59
+ "name": [f"User_{i}" for i in range(1, 1_000_001)],
60
+ "age": [20 + (i % 50) for i in range(1, 1_000_001)],
61
+ })
62
+
63
+ res_v = client.scatter_ingest_vertices(users_df, label="User", id_col="id")
64
+ print(f"Vertices: {res_v['rows_ingested']} rows in {res_v['elapsed_seconds']}s "
65
+ f"({res_v['rows_per_second']:,.0f} rows/s)")
66
+
67
+ # 3. Ingest edges in parallel (automatically hashed by source vertex)
68
+ edges_df = pl.DataFrame({
69
+ "src": range(1, 500_001),
70
+ "dst": range(500_001, 1_000_001),
71
+ })
72
+
73
+ res_e = client.scatter_ingest_edges(edges_df, edge_type="FOLLOWS")
74
+ print(f"Edges: {res_e['rows_ingested']} rows in {res_e['elapsed_seconds']}s "
75
+ f"({res_e['rows_per_second']:,.0f} rows/s)")
76
+
77
+ # 4. Trigger background compaction into immutable Chunked-CSR
78
+ client.compact()
79
+
80
+ # 5. Query graph using openCypher directly into a Polars DataFrame
81
+ df = client.query("MATCH (a:User)-[:FOLLOWS]->(b:User) RETURN a.name, b.name LIMIT 10;")
82
+ print(df)
83
+
84
+ # 6. Close Flight TCP connections
85
+ client.close()
86
+ ```
87
+
88
+ ---
89
+
90
+ ## 🏗 Architecture: How Scatter-Ingest Works
91
+
92
+ In a traditional database setup, client libraries send bulk requests to a single coordinator or load balancer, which must unpack the data, hash every key, and re-transmit records over the network to the correct storage nodes.
93
+
94
+ `gdb-client` moves partition routing to the client side:
95
+
96
+ ```
97
+ ┌─────────────────────────────────────────┐
98
+ │ polars.DataFrame │
99
+ │ (10,000,000 Rows in RAM) │
100
+ └────────────────────┬────────────────────┘
101
+ │
102
+ Polars SIMD Hash: (col % TotalNodes)
103
+ │
104
+ ┌──────────────────────────┼──────────────────────────┐
105
+ ▼ ▼ ▼
106
+ [Partition Chunk 0] [Partition Chunk 1] [Partition Chunk 2]
107
+ │ │ │
108
+ PyArrow RecordBatch PyArrow RecordBatch PyArrow RecordBatch
109
+ │ │ │
110
+ Arrow Flight do_put Arrow Flight do_put Arrow Flight do_put
111
+ (TCP Port :8860) (TCP Port :8861) (TCP Port :8862)
112
+ ▼ ▼ ▼
113
+ ┌─────────────────┐ ┌─────────────────┐ ┌─────────────────┐
114
+ │ GDB Node #1 │ │ GDB Node #2 │ │ GDB Node #3 │
115
+ │ DeltaMemTable │ │ DeltaMemTable │ │ DeltaMemTable │
116
+ └─────────────────┘ └─────────────────┘ └─────────────────┘
117
+ ```
118
+
119
+ 1. **Client-Side Vectorized Hashing**: Polars computes target partitions for all rows simultaneously using native CPU SIMD instructions.
120
+ 2. **Direct Multi-Node Streaming**: Partitions are streamed over concurrent TCP Flight channels directly to the responsible nodes' `--client-flight-port` (`:8860`, `:8861`, `:8862`, etc.).
121
+ 3. **Zero Coordinator Bottleneck**: Cluster inter-node network bandwidth is preserved for MPP shuffle exchange and graph traversal.
122
+
123
+ ---
124
+
125
+ ## 📖 API Reference
126
+
127
+ ### `GdbClient`
128
+
129
+ ```python
130
+ GdbClient(seed_url: str = "http://localhost:8847", auth_token: Optional[str] = None)
131
+ ```
132
+
133
+ - **`seed_url`**: HTTP address of any healthy node in the GDB cluster.
134
+ - **`auth_token`**: Optional Bearer token for clusters with authorization enabled.
135
+
136
+ #### Methods
137
+
138
+ | Method | Parameters | Return Type | Description |
139
+ | :--- | :--- | :--- | :--- |
140
+ | `scatter_ingest_vertices` | `df: pl.DataFrame, label: str = "User", id_col: str = "id", max_workers: int = 8` | `dict` | Partitions vertices and streams Arrow RecordBatches directly to target nodes. |
141
+ | `scatter_ingest_edges` | `df: pl.DataFrame, edge_type: str = "KNOWS", src_col: str = "src", dst_col: str = "dst", rank_col: Optional[str] = None, max_workers: int = 8` | `dict` | Partitions edges by source ID and streams Arrow RecordBatches to target nodes. |
142
+ | `query` | `cypher: str, node_id: Optional[int] = None` | `pl.DataFrame` | Executes an openCypher query and returns the tabular result as a Polars DataFrame. |
143
+ | `compact` | — | `dict` | Triggers asynchronous CSR compaction across cluster nodes. |
144
+ | `refresh_topology` | — | `None` | Queries `/cluster` and reconnects Flight clients to active nodes. |
145
+ | `close` | — | `None` | Closes open Arrow Flight TCP client connections. |
146
+
147
+ ---
148
+
149
+ ### Ingestion Performance Metrics
150
+
151
+ `scatter_ingest_vertices` and `scatter_ingest_edges` return a detailed performance dictionary:
152
+
153
+ ```python
154
+ {
155
+ "rows_ingested": 1000000,
156
+ "elapsed_seconds": 0.4281,
157
+ "rows_per_second": 2335902.8,
158
+ "partitions": 3
159
+ }
160
+ ```
161
+
162
+ ---
163
+
164
+ ## 🛠 Local Development & Building Wheel
165
+
166
+ To build the wheel and source distribution locally:
167
+
168
+ ```bash
169
+ # Clone the repository
170
+ git clone https://github.com/1orgar/gdb-py-client.git
171
+ cd gdb-py-client
172
+
173
+ # Install build dependencies
174
+ pip install -e ".[dev]"
175
+
176
+ # Run tests
177
+ python -m unittest discover tests
178
+
179
+ # Build .whl and .tar.gz distributions
180
+ python -m build
181
+
182
+ # Validate package metadata with twine
183
+ twine check dist/*
184
+ ```
185
+
186
+ Built artifacts will be located in `dist/`:
187
+ - `dist/gdb_client-0.3.0-py3-none-any.whl`
188
+ - `dist/gdb_client-0.3.0.tar.gz`
189
+
190
+
191
+ ---
192
+
193
+ ## 📄 License
194
+
195
+ Licensed under the Apache License, Version 2.0. See the [LICENSE](LICENSE) file for details.
@@ -0,0 +1,71 @@
1
+ [build-system]
2
+ requires = ["hatchling"]
3
+ build-backend = "hatchling.build"
4
+
5
+ [project]
6
+ name = "gdb-client"
7
+ version = "0.3.0"
8
+ description = "High-performance Python client for GDB graph database with Polars parallel scatter-ingest via Apache Arrow Flight"
9
+ readme = "README.md"
10
+ requires-python = ">=3.9"
11
+ license = "Apache-2.0"
12
+ authors = [
13
+ { name = "GDB Team", email = "dev@gdb.io" }
14
+ ]
15
+ keywords = [
16
+ "graph-database",
17
+ "gdb",
18
+ "apache-arrow",
19
+ "arrow-flight",
20
+ "polars",
21
+ "cypher",
22
+ "opencypher",
23
+ "scatter-ingest",
24
+ "htap",
25
+ "mpp",
26
+ ]
27
+ classifiers = [
28
+ "Development Status :: 4 - Beta",
29
+ "Intended Audience :: Developers",
30
+ "Intended Audience :: Science/Research",
31
+ "License :: OSI Approved :: Apache Software License",
32
+ "Programming Language :: Python :: 3",
33
+ "Programming Language :: Python :: 3.9",
34
+ "Programming Language :: Python :: 3.10",
35
+ "Programming Language :: Python :: 3.11",
36
+ "Programming Language :: Python :: 3.12",
37
+ "Programming Language :: Python :: 3.13",
38
+ "Topic :: Database",
39
+ "Topic :: Database :: Database Engines/Servers",
40
+ "Topic :: Scientific/Engineering :: Information Analysis",
41
+ "Operating System :: OS Independent",
42
+ ]
43
+ dependencies = [
44
+ "polars>=1.0.0",
45
+ "pyarrow>=14.0.0",
46
+ "requests>=2.28.0",
47
+ ]
48
+
49
+ [project.optional-dependencies]
50
+ dev = [
51
+ "pytest>=7.0.0",
52
+ "build>=1.0.0",
53
+ "twine>=5.0.0",
54
+ ]
55
+
56
+ [project.urls]
57
+ Homepage = "https://github.com/1orgar/gdb"
58
+ Repository = "https://github.com/1orgar/gdb-py-client"
59
+ Documentation = "https://github.com/1orgar/gdb/blob/main/docs/SERVER_AND_QUERY_GUIDE.md"
60
+ Issues = "https://github.com/1orgar/gdb-py-client/issues"
61
+
62
+ [tool.hatch.build.targets.wheel]
63
+ packages = ["src/gdb_client"]
64
+
65
+ [tool.hatch.build.targets.sdist]
66
+ include = [
67
+ "/src",
68
+ "/tests",
69
+ "/README.md",
70
+ "/pyproject.toml",
71
+ ]
@@ -0,0 +1,7 @@
1
+ """GDB Python Client: High-Performance Distributed Graph Database Connector."""
2
+
3
+ from .client import GdbClient
4
+ from .topology import ClusterTopology, NodeInfo
5
+
6
+ __version__ = "0.3.0"
7
+ __all__ = ["GdbClient", "ClusterTopology", "NodeInfo"]
@@ -0,0 +1,261 @@
1
+ from __future__ import annotations
2
+
3
+ import time
4
+ from concurrent.futures import ThreadPoolExecutor, as_completed
5
+ from typing import Dict, List, Optional, Union
6
+ import json
7
+
8
+ try:
9
+ import polars as pl
10
+ except ImportError:
11
+ pl = None
12
+
13
+ try:
14
+ import pyarrow as pa
15
+ import pyarrow.flight as flight
16
+ except ImportError:
17
+ pa = None
18
+ flight = None
19
+
20
+ try:
21
+ import requests
22
+ except ImportError:
23
+ requests = None
24
+
25
+ from .topology import ClusterTopology, NodeInfo
26
+
27
+
28
+ class GdbClient:
29
+ """High-performance Python client for GDB.
30
+
31
+ Features:
32
+ - Polars-powered parallel scatter-ingest directly to target cluster nodes via Arrow Flight.
33
+ - Zero-copy Flight DoGet queries returning polars.DataFrame directly.
34
+ - Automatic cluster hash ring discovery and partition routing.
35
+ """
36
+
37
+ def __init__(self, seed_url: str = "http://localhost:8847", auth_token: Optional[str] = None):
38
+ self.seed_url = seed_url.rstrip("/")
39
+ self.auth_token = auth_token
40
+ self.topology = ClusterTopology(seed_url=self.seed_url)
41
+ self._flight_clients: Dict[int, flight.FlightClient] = {}
42
+ self._connect_flight_clients()
43
+
44
+ def _connect_flight_clients(self):
45
+ self._flight_clients.clear()
46
+ for node in self.topology.nodes:
47
+ location = f"grpc://{node.host}:{node.client_flight_port}"
48
+ try:
49
+ client = flight.FlightClient(location)
50
+ self._flight_clients[node.node_id] = client
51
+ except Exception as e:
52
+ # Retry or log warning
53
+ pass
54
+
55
+ def refresh_topology(self):
56
+ """Refreshes cluster nodes and recreates Flight connections."""
57
+ self.topology.refresh()
58
+ self._connect_flight_clients()
59
+
60
+ def scatter_ingest_vertices(
61
+ self,
62
+ df: pl.DataFrame,
63
+ label: str = "User",
64
+ id_col: str = "id",
65
+ max_workers: int = 8,
66
+ ) -> Dict[str, Union[int, float]]:
67
+ """Partitions vertex DataFrame in parallel using Polars expressions,
68
+
69
+ and streams each partition directly to the primary replica node via Arrow Flight.
70
+ """
71
+ start_time = time.perf_counter()
72
+ total_nodes = self.topology.total_nodes
73
+ if total_nodes == 0:
74
+ raise RuntimeError("No active cluster nodes discovered in topology")
75
+
76
+ # 1. Compute target node partition in Polars (vectorized zero-copy)
77
+ if df[id_col].dtype in (pl.Int64, pl.Int32, pl.UInt64, pl.UInt32):
78
+ part_expr = (pl.col(id_col) % total_nodes).cast(pl.UInt32).alias("_partition")
79
+ else:
80
+ part_expr = (pl.col(id_col).hash() % total_nodes).cast(pl.UInt32).alias("_partition")
81
+
82
+ partitioned_df = df.with_columns(part_expr)
83
+ partitions = partitioned_df.partition_by("_partition", as_dict=True)
84
+
85
+ # 2. Dispatch partitions to target nodes concurrently
86
+ total_rows = 0
87
+
88
+ def _send_partition(part_idx_tuple, sub_df: pl.DataFrame) -> int:
89
+ part_idx = part_idx_tuple[0] if isinstance(part_idx_tuple, tuple) else part_idx_tuple
90
+ target_node = self.topology.nodes[part_idx % total_nodes]
91
+ client = self._flight_clients.get(target_node.node_id)
92
+ if client is None:
93
+ client = flight.FlightClient(f"grpc://{target_node.host}:{target_node.client_flight_port}")
94
+
95
+ clean_df = sub_df.drop("_partition")
96
+ arrow_table = clean_df.to_arrow()
97
+
98
+ # Attach ingest descriptor command
99
+ descriptor_cmd = json.dumps({"type": "vertex", "label": label})
100
+ descriptor = flight.FlightDescriptor.for_command(descriptor_cmd.encode("utf-8"))
101
+
102
+ writer, reader = client.do_put(descriptor, arrow_table.schema)
103
+ writer.write_table(arrow_table)
104
+ writer.close()
105
+
106
+ # Drain result
107
+ for _ in reader:
108
+ pass
109
+
110
+ return len(sub_df)
111
+
112
+ with ThreadPoolExecutor(max_workers=min(max_workers, len(partitions) or 1)) as executor:
113
+ futures = [
114
+ executor.submit(_send_partition, part_key, sub_df)
115
+ for part_key, sub_df in partitions.items()
116
+ ]
117
+ for f in as_completed(futures):
118
+ total_rows += f.result()
119
+
120
+ elapsed = time.perf_counter() - start_time
121
+ throughput = total_rows / elapsed if elapsed > 0 else 0
122
+
123
+ return {
124
+ "rows_ingested": total_rows,
125
+ "elapsed_seconds": round(elapsed, 4),
126
+ "rows_per_second": round(throughput, 2),
127
+ "partitions": len(partitions),
128
+ }
129
+
130
+ def scatter_ingest_edges(
131
+ self,
132
+ df: pl.DataFrame,
133
+ edge_type: str = "KNOWS",
134
+ src_col: str = "src",
135
+ dst_col: str = "dst",
136
+ rank_col: Optional[str] = None,
137
+ max_workers: int = 8,
138
+ ) -> Dict[str, Union[int, float]]:
139
+ """Partitions edge DataFrame by source vertex ID in parallel using Polars,
140
+
141
+ and streams directly to primary replica nodes via Arrow Flight.
142
+ """
143
+ start_time = time.perf_counter()
144
+ total_nodes = self.topology.total_nodes
145
+ if total_nodes == 0:
146
+ raise RuntimeError("No active cluster nodes discovered in topology")
147
+
148
+ if df[src_col].dtype in (pl.Int64, pl.Int32, pl.UInt64, pl.UInt32):
149
+ part_expr = (pl.col(src_col) % total_nodes).cast(pl.UInt32).alias("_partition")
150
+ else:
151
+ part_expr = (pl.col(src_col).hash() % total_nodes).cast(pl.UInt32).alias("_partition")
152
+
153
+ partitioned_df = df.with_columns(part_expr)
154
+ partitions = partitioned_df.partition_by("_partition", as_dict=True)
155
+
156
+ total_rows = 0
157
+
158
+ def _send_edge_partition(part_idx_tuple, sub_df: pl.DataFrame) -> int:
159
+ part_idx = part_idx_tuple[0] if isinstance(part_idx_tuple, tuple) else part_idx_tuple
160
+ target_node = self.topology.nodes[part_idx % total_nodes]
161
+ client = self._flight_clients.get(target_node.node_id)
162
+ if client is None:
163
+ client = flight.FlightClient(f"grpc://{target_node.host}:{target_node.client_flight_port}")
164
+
165
+ clean_df = sub_df.drop("_partition")
166
+ arrow_table = clean_df.to_arrow()
167
+
168
+ descriptor_cmd = json.dumps({"type": "edge", "edge_type": edge_type})
169
+ descriptor = flight.FlightDescriptor.for_command(descriptor_cmd.encode("utf-8"))
170
+
171
+ writer, reader = client.do_put(descriptor, arrow_table.schema)
172
+ writer.write_table(arrow_table)
173
+ writer.close()
174
+
175
+ for _ in reader:
176
+ pass
177
+
178
+ return len(sub_df)
179
+
180
+ with ThreadPoolExecutor(max_workers=min(max_workers, len(partitions) or 1)) as executor:
181
+ futures = [
182
+ executor.submit(_send_edge_partition, part_key, sub_df)
183
+ for part_key, sub_df in partitions.items()
184
+ ]
185
+ for f in as_completed(futures):
186
+ total_rows += f.result()
187
+
188
+ elapsed = time.perf_counter() - start_time
189
+ throughput = total_rows / elapsed if elapsed > 0 else 0
190
+
191
+ return {
192
+ "rows_ingested": total_rows,
193
+ "elapsed_seconds": round(elapsed, 4),
194
+ "rows_per_second": round(throughput, 2),
195
+ "partitions": len(partitions),
196
+ }
197
+
198
+ def query(self, cypher: str, node_id: Optional[int] = None) -> pl.DataFrame:
199
+ """Executes an openCypher query and returns the results directly as a Polars DataFrame.
200
+
201
+ Prefers Arrow Flight DoGet for maximum throughput; falls back to HTTP REST API.
202
+ """
203
+ # Choose client node
204
+ target_client = None
205
+ if node_id is not None and node_id in self._flight_clients:
206
+ target_client = self._flight_clients[node_id]
207
+ elif self._flight_clients:
208
+ target_client = next(iter(self._flight_clients.values()))
209
+
210
+ # Attempt Flight DoGet first
211
+ if target_client is not None:
212
+ try:
213
+ ticket = flight.Ticket(cypher.encode("utf-8"))
214
+ reader = target_client.do_get(ticket)
215
+ table = reader.read_all()
216
+ return pl.from_arrow(table)
217
+ except Exception:
218
+ pass
219
+
220
+ # Fallback to HTTP API
221
+ target_url = f"{self.seed_url}/query"
222
+ if requests is not None:
223
+ resp = requests.post(target_url, json={"query": cypher}, timeout=30)
224
+ resp.raise_for_status()
225
+ data = resp.json()
226
+ else:
227
+ import urllib.request
228
+ req = urllib.request.Request(
229
+ target_url,
230
+ data=json.dumps({"query": cypher}).encode("utf-8"),
231
+ headers={"Content-Type": "application/json"},
232
+ )
233
+ with urllib.request.urlopen(req, timeout=30) as resp:
234
+ data = json.loads(resp.read().decode("utf-8"))
235
+
236
+ columns = data.get("columns", [])
237
+ rows = data.get("rows", [])
238
+
239
+ if not columns or not rows:
240
+ return pl.DataFrame() if pl is not None else {}
241
+
242
+ if pl is not None:
243
+ col_dict = {col: [row[i] for row in rows] for i, col in enumerate(columns)}
244
+ return pl.DataFrame(col_dict)
245
+ return {"columns": columns, "rows": rows}
246
+
247
+ def compact(self) -> dict:
248
+ """Triggers asynchronous compaction across the cluster."""
249
+ target_url = f"{self.seed_url}/compact"
250
+ if requests is not None:
251
+ resp = requests.post(target_url, timeout=10)
252
+ return resp.json()
253
+ else:
254
+ import urllib.request
255
+ req = urllib.request.Request(target_url, data=b"{}", headers={"Content-Type": "application/json"})
256
+ with urllib.request.urlopen(req, timeout=10) as resp:
257
+ return json.loads(resp.read().decode("utf-8"))
258
+
259
+ def close(self):
260
+ """Closes all active Flight client connections."""
261
+ self._flight_clients.clear()
@@ -0,0 +1,91 @@
1
+ from dataclasses import dataclass
2
+ from typing import List, Optional
3
+ import json
4
+ import urllib.request
5
+ import urllib.error
6
+
7
+ try:
8
+ import requests
9
+ except ImportError:
10
+ requests = None
11
+
12
+
13
+ @dataclass
14
+ class NodeInfo:
15
+ node_id: int
16
+ http_url: str
17
+ flight_port: int
18
+ client_flight_port: int
19
+ host: str = "127.0.0.1"
20
+
21
+ @classmethod
22
+ def from_dict(cls, data: dict) -> "NodeInfo":
23
+ http_url = data.get("http_url", "http://127.0.0.1:8847")
24
+ host = "127.0.0.1"
25
+ if "//" in http_url:
26
+ host_part = http_url.split("//")[1].split(":")[0]
27
+ if host_part:
28
+ host = host_part
29
+
30
+ return cls(
31
+ node_id=int(data.get("node_id", 1)),
32
+ http_url=http_url,
33
+ flight_port=int(data.get("flight_port", 8848)),
34
+ client_flight_port=int(data.get("client_flight_port", 8860)),
35
+ host=host,
36
+ )
37
+
38
+
39
+ class ClusterTopology:
40
+ """Discovers and caches the cluster hash ring and node ports."""
41
+
42
+ def __init__(self, seed_url: str = "http://localhost:8847"):
43
+ self.seed_url = seed_url.rstrip("/")
44
+ self.nodes: List[NodeInfo] = []
45
+ self.replication_factor: int = 1
46
+ self.partitions: int = 4
47
+ self.refresh()
48
+
49
+ def refresh(self):
50
+ url = f"{self.seed_url}/cluster"
51
+ try:
52
+ if requests is not None:
53
+ resp = requests.get(url, timeout=5)
54
+ resp.raise_for_status()
55
+ data = resp.json()
56
+ else:
57
+ req = urllib.request.Request(url, headers={"User-Agent": "gdb-client/0.2.0"})
58
+ with urllib.request.urlopen(req, timeout=5) as resp:
59
+ data = json.loads(resp.read().decode("utf-8"))
60
+ raw_nodes = data.get("ring_nodes", [])
61
+ if raw_nodes:
62
+ self.nodes = [NodeInfo.from_dict(n) for n in raw_nodes]
63
+ else:
64
+ self.nodes = [
65
+ NodeInfo(
66
+ node_id=int(data.get("node_id", 1)),
67
+ http_url=self.seed_url,
68
+ flight_port=int(data.get("flight_port", 8848)),
69
+ client_flight_port=int(data.get("client_flight_port", 8860)),
70
+ )
71
+ ]
72
+ self.replication_factor = int(data.get("replication_factor", 1))
73
+ self.partitions = int(data.get("partitions", 4))
74
+ except Exception as e:
75
+ # Fallback single node
76
+ self.nodes = [
77
+ NodeInfo(
78
+ node_id=1,
79
+ http_url=self.seed_url,
80
+ flight_port=8848,
81
+ client_flight_port=8860,
82
+ )
83
+ ]
84
+
85
+ @property
86
+ def total_nodes(self) -> int:
87
+ return len(self.nodes)
88
+
89
+ def target_node_for_key(self, key: int) -> NodeInfo:
90
+ idx = key % len(self.nodes)
91
+ return self.nodes[idx]
@@ -0,0 +1,55 @@
1
+ import unittest
2
+ from gdb_client.topology import ClusterTopology, NodeInfo
3
+
4
+ try:
5
+ import polars as pl
6
+ HAS_POLARS = True
7
+ except ImportError:
8
+ HAS_POLARS = False
9
+
10
+
11
+ class TestTopologyAndClient(unittest.TestCase):
12
+ def test_topology_routing(self):
13
+ topology = ClusterTopology.__new__(ClusterTopology)
14
+ topology.nodes = [
15
+ NodeInfo(node_id=1, http_url="http://127.0.0.1:8847", flight_port=8848, client_flight_port=8860),
16
+ NodeInfo(node_id=2, http_url="http://127.0.0.1:8846", flight_port=8849, client_flight_port=8861),
17
+ NodeInfo(node_id=3, http_url="http://127.0.0.1:8845", flight_port=8850, client_flight_port=8862),
18
+ ]
19
+ topology.replication_factor = 3
20
+ topology.partitions = 8
21
+
22
+ self.assertEqual(topology.total_nodes, 3)
23
+
24
+ # Key modulo routing tests
25
+ target_0 = topology.target_node_for_key(0)
26
+ self.assertEqual(target_0.node_id, 1)
27
+ self.assertEqual(target_0.client_flight_port, 8860)
28
+
29
+ target_1 = topology.target_node_for_key(1)
30
+ self.assertEqual(target_1.node_id, 2)
31
+ self.assertEqual(target_1.client_flight_port, 8861)
32
+
33
+ target_2 = topology.target_node_for_key(2)
34
+ self.assertEqual(target_2.node_id, 3)
35
+ self.assertEqual(target_2.client_flight_port, 8862)
36
+
37
+ target_3 = topology.target_node_for_key(3)
38
+ self.assertEqual(target_3.node_id, 1)
39
+
40
+ @unittest.skipUnless(HAS_POLARS, "polars is required for this test")
41
+ def test_polars_partitioning_logic(self):
42
+ df = pl.DataFrame({
43
+ "id": [1, 2, 3, 4, 5, 6],
44
+ "name": ["A", "B", "C", "D", "E", "F"],
45
+ })
46
+ total_nodes = 3
47
+ df_part = df.with_columns((pl.col("id") % total_nodes).cast(pl.UInt32).alias("_partition"))
48
+ partitions = df_part.partition_by("_partition", as_dict=True)
49
+
50
+ self.assertEqual(len(partitions), 3)
51
+ self.assertEqual(sum(len(sub) for sub in partitions.values()), 6)
52
+
53
+
54
+ if __name__ == "__main__":
55
+ unittest.main()