real-ladybug 0.0.1.dev1__cp311-cp311-win_amd64.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.

Potentially problematic release.


This version of real-ladybug might be problematic. Click here for more details.

Files changed (114) hide show
  1. real_ladybug/__init__.py +83 -0
  2. real_ladybug/_lbug.cp311-win_amd64.pyd +0 -0
  3. real_ladybug/_lbug.exp +0 -0
  4. real_ladybug/_lbug.lib +0 -0
  5. real_ladybug/async_connection.py +226 -0
  6. real_ladybug/connection.py +323 -0
  7. real_ladybug/constants.py +7 -0
  8. real_ladybug/database.py +307 -0
  9. real_ladybug/prepared_statement.py +51 -0
  10. real_ladybug/py.typed +0 -0
  11. real_ladybug/query_result.py +511 -0
  12. real_ladybug/torch_geometric_feature_store.py +185 -0
  13. real_ladybug/torch_geometric_graph_store.py +131 -0
  14. real_ladybug/torch_geometric_result_converter.py +282 -0
  15. real_ladybug/types.py +39 -0
  16. real_ladybug-0.0.1.dev1.dist-info/METADATA +88 -0
  17. real_ladybug-0.0.1.dev1.dist-info/RECORD +114 -0
  18. real_ladybug-0.0.1.dev1.dist-info/WHEEL +5 -0
  19. real_ladybug-0.0.1.dev1.dist-info/licenses/LICENSE +21 -0
  20. real_ladybug-0.0.1.dev1.dist-info/top_level.txt +3 -0
  21. real_ladybug-0.0.1.dev1.dist-info/zip-safe +1 -0
  22. real_ladybug-source/scripts/antlr4/hash.py +2 -0
  23. real_ladybug-source/scripts/antlr4/keywordhandler.py +47 -0
  24. real_ladybug-source/scripts/collect-extensions.py +68 -0
  25. real_ladybug-source/scripts/collect-single-file-header.py +126 -0
  26. real_ladybug-source/scripts/export-dbs.py +101 -0
  27. real_ladybug-source/scripts/export-import-test.py +345 -0
  28. real_ladybug-source/scripts/extension/purge-beta.py +34 -0
  29. real_ladybug-source/scripts/generate-cpp-docs/collect_files.py +122 -0
  30. real_ladybug-source/scripts/generate-tinysnb.py +34 -0
  31. real_ladybug-source/scripts/get-clangd-diagnostics.py +233 -0
  32. real_ladybug-source/scripts/migrate-lbug-db.py +308 -0
  33. real_ladybug-source/scripts/multiplatform-test-helper/collect-results.py +71 -0
  34. real_ladybug-source/scripts/multiplatform-test-helper/notify-discord.py +68 -0
  35. real_ladybug-source/scripts/pip-package/package_tar.py +90 -0
  36. real_ladybug-source/scripts/pip-package/setup.py +130 -0
  37. real_ladybug-source/scripts/run-clang-format.py +408 -0
  38. real_ladybug-source/scripts/setup-extension-repo.py +67 -0
  39. real_ladybug-source/scripts/test-simsimd-dispatch.py +45 -0
  40. real_ladybug-source/scripts/update-nightly-build-version.py +81 -0
  41. real_ladybug-source/third_party/brotli/scripts/dictionary/step-01-download-rfc.py +16 -0
  42. real_ladybug-source/third_party/brotli/scripts/dictionary/step-02-rfc-to-bin.py +34 -0
  43. real_ladybug-source/third_party/brotli/scripts/dictionary/step-03-validate-bin.py +35 -0
  44. real_ladybug-source/third_party/brotli/scripts/dictionary/step-04-generate-java-literals.py +85 -0
  45. real_ladybug-source/third_party/pybind11/tools/codespell_ignore_lines_from_errors.py +35 -0
  46. real_ladybug-source/third_party/pybind11/tools/libsize.py +36 -0
  47. real_ladybug-source/third_party/pybind11/tools/make_changelog.py +63 -0
  48. real_ladybug-source/tools/python_api/build/real_ladybug/__init__.py +83 -0
  49. real_ladybug-source/tools/python_api/build/real_ladybug/async_connection.py +226 -0
  50. real_ladybug-source/tools/python_api/build/real_ladybug/connection.py +323 -0
  51. real_ladybug-source/tools/python_api/build/real_ladybug/constants.py +7 -0
  52. real_ladybug-source/tools/python_api/build/real_ladybug/database.py +307 -0
  53. real_ladybug-source/tools/python_api/build/real_ladybug/prepared_statement.py +51 -0
  54. real_ladybug-source/tools/python_api/build/real_ladybug/py.typed +0 -0
  55. real_ladybug-source/tools/python_api/build/real_ladybug/query_result.py +511 -0
  56. real_ladybug-source/tools/python_api/build/real_ladybug/torch_geometric_feature_store.py +185 -0
  57. real_ladybug-source/tools/python_api/build/real_ladybug/torch_geometric_graph_store.py +131 -0
  58. real_ladybug-source/tools/python_api/build/real_ladybug/torch_geometric_result_converter.py +282 -0
  59. real_ladybug-source/tools/python_api/build/real_ladybug/types.py +39 -0
  60. real_ladybug-source/tools/python_api/src_py/__init__.py +83 -0
  61. real_ladybug-source/tools/python_api/src_py/async_connection.py +226 -0
  62. real_ladybug-source/tools/python_api/src_py/connection.py +323 -0
  63. real_ladybug-source/tools/python_api/src_py/constants.py +7 -0
  64. real_ladybug-source/tools/python_api/src_py/database.py +307 -0
  65. real_ladybug-source/tools/python_api/src_py/prepared_statement.py +51 -0
  66. real_ladybug-source/tools/python_api/src_py/py.typed +0 -0
  67. real_ladybug-source/tools/python_api/src_py/query_result.py +511 -0
  68. real_ladybug-source/tools/python_api/src_py/torch_geometric_feature_store.py +185 -0
  69. real_ladybug-source/tools/python_api/src_py/torch_geometric_graph_store.py +131 -0
  70. real_ladybug-source/tools/python_api/src_py/torch_geometric_result_converter.py +282 -0
  71. real_ladybug-source/tools/python_api/src_py/types.py +39 -0
  72. real_ladybug-source/tools/python_api/test/conftest.py +230 -0
  73. real_ladybug-source/tools/python_api/test/disabled_test_extension.py +73 -0
  74. real_ladybug-source/tools/python_api/test/ground_truth.py +430 -0
  75. real_ladybug-source/tools/python_api/test/test_arrow.py +694 -0
  76. real_ladybug-source/tools/python_api/test/test_async_connection.py +159 -0
  77. real_ladybug-source/tools/python_api/test/test_blob_parameter.py +145 -0
  78. real_ladybug-source/tools/python_api/test/test_connection.py +49 -0
  79. real_ladybug-source/tools/python_api/test/test_database.py +234 -0
  80. real_ladybug-source/tools/python_api/test/test_datatype.py +372 -0
  81. real_ladybug-source/tools/python_api/test/test_df.py +564 -0
  82. real_ladybug-source/tools/python_api/test/test_dict.py +112 -0
  83. real_ladybug-source/tools/python_api/test/test_exception.py +54 -0
  84. real_ladybug-source/tools/python_api/test/test_fsm.py +227 -0
  85. real_ladybug-source/tools/python_api/test/test_get_header.py +49 -0
  86. real_ladybug-source/tools/python_api/test/test_helper.py +8 -0
  87. real_ladybug-source/tools/python_api/test/test_issue.py +147 -0
  88. real_ladybug-source/tools/python_api/test/test_iteration.py +96 -0
  89. real_ladybug-source/tools/python_api/test/test_networkx.py +437 -0
  90. real_ladybug-source/tools/python_api/test/test_parameter.py +340 -0
  91. real_ladybug-source/tools/python_api/test/test_prepared_statement.py +117 -0
  92. real_ladybug-source/tools/python_api/test/test_query_result.py +54 -0
  93. real_ladybug-source/tools/python_api/test/test_query_result_close.py +44 -0
  94. real_ladybug-source/tools/python_api/test/test_scan_pandas.py +676 -0
  95. real_ladybug-source/tools/python_api/test/test_scan_pandas_pyarrow.py +714 -0
  96. real_ladybug-source/tools/python_api/test/test_scan_polars.py +165 -0
  97. real_ladybug-source/tools/python_api/test/test_scan_pyarrow.py +167 -0
  98. real_ladybug-source/tools/python_api/test/test_timeout.py +11 -0
  99. real_ladybug-source/tools/python_api/test/test_torch_geometric.py +640 -0
  100. real_ladybug-source/tools/python_api/test/test_torch_geometric_remote_backend.py +111 -0
  101. real_ladybug-source/tools/python_api/test/test_udf.py +207 -0
  102. real_ladybug-source/tools/python_api/test/test_version.py +6 -0
  103. real_ladybug-source/tools/python_api/test/test_wal.py +80 -0
  104. real_ladybug-source/tools/python_api/test/type_aliases.py +10 -0
  105. real_ladybug-source/tools/rust_api/update_version.py +47 -0
  106. real_ladybug-source/tools/shell/test/conftest.py +218 -0
  107. real_ladybug-source/tools/shell/test/test_helper.py +60 -0
  108. real_ladybug-source/tools/shell/test/test_shell_basics.py +325 -0
  109. real_ladybug-source/tools/shell/test/test_shell_commands.py +656 -0
  110. real_ladybug-source/tools/shell/test/test_shell_control_edit.py +438 -0
  111. real_ladybug-source/tools/shell/test/test_shell_control_search.py +468 -0
  112. real_ladybug-source/tools/shell/test/test_shell_esc_edit.py +232 -0
  113. real_ladybug-source/tools/shell/test/test_shell_esc_search.py +162 -0
  114. real_ladybug-source/tools/shell/test/test_shell_flags.py +645 -0
@@ -0,0 +1,694 @@
1
+ from __future__ import annotations
2
+
3
+ import math
4
+ from datetime import date, datetime, timedelta
5
+ from decimal import Decimal
6
+ from typing import Any
7
+ from uuid import UUID
8
+
9
+ import pytest
10
+
11
+ import ground_truth
12
+ import real_ladybug as lb
13
+ import polars as pl
14
+ import pyarrow as pa
15
+ import pytz
16
+ from pandas import Timestamp
17
+ from type_aliases import ConnDB
18
+ from real_ladybug.constants import ID, LABEL, SRC, DST, NODES
19
+
20
+ _expected_dtypes = {
21
+ # ------------------------------------------------
22
+ # person
23
+ # ------------------------------------------------
24
+ "a.ID": {"arrow": pa.int64(), "pl": pl.Int64},
25
+ "a.fName": {"arrow": pa.string(), "pl": pl.String},
26
+ "a.gender": {"arrow": pa.int64(), "pl": pl.Int64},
27
+ "a.isStudent": {"arrow": pa.bool_(), "pl": pl.Boolean},
28
+ "a.isWorker": {"arrow": pa.bool_(), "pl": pl.Boolean},
29
+ "a.age": {"arrow": pa.int64(), "pl": pl.Int64},
30
+ "a.eyeSight": {"arrow": pa.float64(), "pl": pl.Float64},
31
+ "a.birthdate": {"arrow": pa.date32(), "pl": pl.Date},
32
+ "a.registerTime": {"arrow": pa.timestamp("us"), "pl": pl.Datetime("us")},
33
+ "a.lastJobDuration": {"arrow": pa.duration("us"), "pl": pl.Duration("us")},
34
+ "a.workedHours": {"arrow": pa.list_(pa.int64()), "pl": pl.List(pl.Int64)},
35
+ "a.usedNames": {"arrow": pa.list_(pa.string()), "pl": pl.List(pl.String)},
36
+ "a.courseScoresPerTerm": {"arrow": pa.list_(pa.list_(pa.int64())), "pl": pl.List(pl.List(pl.Int64))},
37
+ "a.grades": {"arrow": pa.list_(pa.int64(), 4), "pl": pl.Array(pl.Int64, shape=(4,))},
38
+ "a.height": {"arrow": pa.float32(), "pl": pl.Float32},
39
+ "a.u": {"arrow": pa.uuid(), "pl": pl.String},
40
+ # ------------------------------------------------
41
+ # movies
42
+ # ------------------------------------------------
43
+ "a.length": {"arrow": pa.int32(), "pl": pl.Int32},
44
+ "m.name": {"arrow": pa.string(), "pl": pl.String},
45
+ "a.description": {
46
+ "arrow": pa.struct([
47
+ ("rating", pa.float64()),
48
+ ("stars", pa.int8()),
49
+ ("views", pa.int64()),
50
+ ("release", pa.timestamp("us")),
51
+ ("release_ns", pa.timestamp("ns")),
52
+ ("release_ms", pa.timestamp("ms")),
53
+ ("release_sec", pa.timestamp("s")),
54
+ ("release_tz", pa.timestamp("us", tz="UTC")),
55
+ ("film", pa.date32()),
56
+ ("u8", pa.uint8()),
57
+ ("u16", pa.uint16()),
58
+ ("u32", pa.uint32()),
59
+ ("u64", pa.uint64()),
60
+ ("hugedata", pa.decimal128(38, 0)),
61
+ ]),
62
+ "pl": pl.Struct({
63
+ "rating": pl.Float64,
64
+ "stars": pl.Int8,
65
+ "views": pl.Int64,
66
+ "release": pl.Datetime(time_unit="us"),
67
+ "release_ns": pl.Datetime(time_unit="ns"),
68
+ "release_ms": pl.Datetime(time_unit="ms"),
69
+ "release_sec": pl.Datetime(time_unit="ms"),
70
+ "release_tz": pl.Datetime(time_unit="us", time_zone="UTC"),
71
+ "film": pl.Date,
72
+ "u8": pl.UInt8,
73
+ "u16": pl.UInt16,
74
+ "u32": pl.UInt32,
75
+ "u64": pl.UInt64,
76
+ "hugedata": pl.Decimal(precision=38, scale=0),
77
+ }),
78
+ },
79
+ # ------------------------------------------------
80
+ # miscellaneous
81
+ # ------------------------------------------------
82
+ "a.lbl": {"arrow": pa.string(), "pl": pl.String},
83
+ "a.orgCode": {"arrow": pa.int64(), "pl": pl.Int64},
84
+ }
85
+
86
+
87
+ def get_result(query_result: lb.QueryResult, result_type: str, chunk_size: int | None) -> Any:
88
+ sz = [] if (chunk_size is None or result_type == "pl") else [chunk_size]
89
+ res = getattr(query_result, f"get_as_{result_type}")(*sz)
90
+ if result_type == "arrow" and chunk_size:
91
+ assert res[0].num_chunks == max((len(res) // chunk_size), 1)
92
+ return res
93
+
94
+
95
+ def assert_column_equals(data: Any, col_name: str, return_type: str, expected_values: list[Any]) -> None:
96
+ col = data[col_name]
97
+ col_dtype = col.dtype if hasattr(col, "dtype") else col.type
98
+ values = col.to_pylist() if return_type == "arrow" else col.to_list()
99
+
100
+ assert len(col) == len(expected_values)
101
+ assert values == expected_values, f"Unexpected values for {col_name} ({return_type!r})"
102
+ assert col_dtype == _expected_dtypes[col_name][return_type], f"Unexpected dtype for {col_name} ({return_type!r})"
103
+
104
+
105
+ def assert_col_names(data: Any, expected_col_names: list[str]) -> None:
106
+ col_names = [(c._name if hasattr(c, "_name") else c.name) for c in data]
107
+ assert col_names == expected_col_names, f"Unexpected column names: {col_names!r}"
108
+
109
+
110
+ def test_to_arrow(conn_db_readonly: ConnDB) -> None:
111
+ conn, _ = conn_db_readonly
112
+
113
+ def _test_person_table(_conn: lb.Connection, return_type: str, chunk_size: int | None = None) -> None:
114
+ query = "MATCH (a:person) RETURN a.* ORDER BY a.ID"
115
+ data = get_result(_conn.execute(query), return_type, chunk_size)
116
+ assert len(data.columns) == 16
117
+
118
+ assert_column_equals(
119
+ data=data,
120
+ col_name="a.ID",
121
+ return_type=return_type,
122
+ expected_values=[0, 2, 3, 5, 7, 8, 9, 10],
123
+ )
124
+
125
+ assert_column_equals(
126
+ data=data,
127
+ col_name="a.fName",
128
+ return_type=return_type,
129
+ expected_values=[
130
+ "Alice",
131
+ "Bob",
132
+ "Carol",
133
+ "Dan",
134
+ "Elizabeth",
135
+ "Farooq",
136
+ "Greg",
137
+ "Hubert Blaine Wolfeschlegelsteinhausenbergerdorff",
138
+ ],
139
+ )
140
+
141
+ assert_column_equals(
142
+ data=data,
143
+ col_name="a.gender",
144
+ return_type=return_type,
145
+ expected_values=[1, 2, 1, 2, 1, 2, 2, 2],
146
+ )
147
+
148
+ assert_column_equals(
149
+ data=data,
150
+ col_name="a.isStudent",
151
+ return_type=return_type,
152
+ expected_values=[True, True, False, False, False, True, False, False],
153
+ )
154
+
155
+ assert_column_equals(
156
+ data=data,
157
+ col_name="a.isWorker",
158
+ return_type=return_type,
159
+ expected_values=[False, False, True, True, True, False, False, True],
160
+ )
161
+
162
+ assert_column_equals(
163
+ data=data,
164
+ col_name="a.age",
165
+ return_type=return_type,
166
+ expected_values=[35, 30, 45, 20, 20, 25, 40, 83],
167
+ )
168
+
169
+ assert_column_equals(
170
+ data=data,
171
+ col_name="a.eyeSight",
172
+ return_type=return_type,
173
+ expected_values=[5.0, 5.1, 5.0, 4.8, 4.7, 4.5, 4.9, 4.9],
174
+ )
175
+
176
+ assert_column_equals(
177
+ data=data,
178
+ col_name="a.birthdate",
179
+ return_type=return_type,
180
+ expected_values=[
181
+ date(1900, 1, 1),
182
+ date(1900, 1, 1),
183
+ date(1940, 6, 22),
184
+ date(1950, 7, 23),
185
+ date(1980, 10, 26),
186
+ date(1980, 10, 26),
187
+ date(1980, 10, 26),
188
+ date(1990, 11, 27),
189
+ ],
190
+ )
191
+
192
+ assert_column_equals(
193
+ data=data,
194
+ col_name="a.registerTime",
195
+ return_type=return_type,
196
+ expected_values=[
197
+ datetime(2011, 8, 20, 11, 25, 30),
198
+ datetime(2008, 11, 3, 15, 25, 30, 526),
199
+ datetime(1911, 8, 20, 2, 32, 21),
200
+ datetime(2031, 11, 30, 12, 25, 30),
201
+ datetime(1976, 12, 23, 11, 21, 42),
202
+ datetime(1972, 7, 31, 13, 22, 30, 678559),
203
+ datetime(1976, 12, 23, 4, 41, 42),
204
+ datetime(2023, 2, 21, 13, 25, 30),
205
+ ],
206
+ )
207
+
208
+ assert_column_equals(
209
+ data=data,
210
+ col_name="a.lastJobDuration",
211
+ return_type=return_type,
212
+ expected_values=[
213
+ timedelta(days=1082, seconds=46920),
214
+ timedelta(days=3750, seconds=46800, microseconds=24),
215
+ timedelta(days=2, seconds=1451),
216
+ timedelta(days=3750, seconds=46800, microseconds=24),
217
+ timedelta(days=2, seconds=1451),
218
+ timedelta(seconds=1080, microseconds=24000),
219
+ timedelta(days=3750, seconds=46800, microseconds=24),
220
+ timedelta(days=1082, seconds=46920),
221
+ ],
222
+ )
223
+
224
+ assert_column_equals(
225
+ data=data,
226
+ col_name="a.workedHours",
227
+ return_type=return_type,
228
+ expected_values=[[10, 5], [12, 8], [4, 5], [1, 9], [2], [3, 4, 5, 6, 7], [1], [10, 11, 12, 3, 4, 5, 6, 7]],
229
+ )
230
+
231
+ assert_column_equals(
232
+ data=data,
233
+ col_name="a.usedNames",
234
+ return_type=return_type,
235
+ expected_values=[
236
+ ["Aida"],
237
+ ["Bobby"],
238
+ ["Carmen", "Fred"],
239
+ ["Wolfeschlegelstein", "Daniel"],
240
+ ["Ein"],
241
+ ["Fesdwe"],
242
+ ["Grad"],
243
+ ["Ad", "De", "Hi", "Kye", "Orlan"],
244
+ ],
245
+ )
246
+
247
+ assert_column_equals(
248
+ data=data,
249
+ col_name="a.courseScoresPerTerm",
250
+ return_type=return_type,
251
+ expected_values=[
252
+ [[10, 8], [6, 7, 8]],
253
+ [[8, 9], [9, 10]],
254
+ [[8, 10]],
255
+ [[7, 4], [8, 8], [9]],
256
+ [[6], [7], [8]],
257
+ [[8]],
258
+ [[10]],
259
+ [[7], [10], [6, 7]],
260
+ ],
261
+ )
262
+
263
+ assert_column_equals(
264
+ data=data,
265
+ col_name="a.grades",
266
+ return_type=return_type,
267
+ expected_values=[
268
+ [96, 54, 86, 92],
269
+ [98, 42, 93, 88],
270
+ [91, 75, 21, 95],
271
+ [76, 88, 99, 89],
272
+ [96, 59, 65, 88],
273
+ [80, 78, 34, 83],
274
+ [43, 83, 67, 43],
275
+ [77, 64, 100, 54],
276
+ ],
277
+ )
278
+ assert_column_equals(
279
+ data=data,
280
+ col_name="a.u",
281
+ return_type=return_type,
282
+ expected_values=[
283
+ UUID(u) if return_type != "pl" else u
284
+ for u in [
285
+ "a0eebc99-9c0b-4ef8-bb6d-6bb9bd380a11",
286
+ "a0eebc99-9c0b-4ef8-bb6d-6bb9bd380a12",
287
+ "a0eebc99-9c0b-4ef8-bb6d-6bb9bd380a13",
288
+ "a0eebc99-9c0b-4ef8-bb6d-6bb9bd380a14",
289
+ "a0eebc99-9c0b-4ef8-bb6d-6bb9bd380a15",
290
+ "a0eebc99-9c0b-4ef8-bb6d-6bb9bd380a16",
291
+ "a0eebc99-9c0b-4ef8-bb6d-6bb9bd380a17",
292
+ "a0eebc99-9c0b-4ef8-bb6d-6bb9bd380a18",
293
+ ]
294
+ ],
295
+ )
296
+
297
+ def _test_movies_table(_conn: lb.Connection, return_type: str, chunk_size: int | None = None) -> None:
298
+ query = "MATCH (a:movies) RETURN a.length, a.description ORDER BY a.length"
299
+ data = get_result(_conn.execute(query), return_type, chunk_size)
300
+
301
+ assert_col_names(data, ["a.length", "a.description"])
302
+ assert_column_equals(
303
+ data=data,
304
+ col_name="a.length",
305
+ return_type=return_type,
306
+ expected_values=[126, 298, 2544],
307
+ )
308
+
309
+ assert_column_equals(
310
+ data=data,
311
+ col_name="a.description",
312
+ return_type=return_type,
313
+ expected_values=[
314
+ {
315
+ "rating": 5.3,
316
+ "stars": 2,
317
+ "views": 152,
318
+ "release": datetime(2011, 8, 20, 11, 25, 30),
319
+ "release_ns": datetime(2011, 8, 20, 11, 25, 30, 123456)
320
+ if return_type == "pl"
321
+ else Timestamp("2011-08-20 11:25:30.123456"),
322
+ "release_ms": datetime(2011, 8, 20, 11, 25, 30, 123000),
323
+ "release_sec": datetime(2011, 8, 20, 11, 25, 30),
324
+ "release_tz": datetime(2011, 8, 20, 11, 25, 30, 123456, pytz.UTC),
325
+ "film": date(2012, 5, 11),
326
+ "u8": 220,
327
+ "u16": 20,
328
+ "u32": 1,
329
+ "u64": 180,
330
+ "hugedata": Decimal(1844674407370955161811111111),
331
+ },
332
+ {
333
+ "rating": 1223.0,
334
+ "stars": 100,
335
+ "views": 10003,
336
+ "release": datetime(2011, 2, 11, 16, 44, 22),
337
+ "release_ns": datetime(2011, 2, 11, 16, 44, 22, 123456)
338
+ if return_type == "pl"
339
+ else Timestamp("2011-02-11 16:44:22.123456"),
340
+ "release_ms": datetime(2011, 2, 11, 16, 44, 22, 123000),
341
+ "release_sec": datetime(2011, 2, 11, 16, 44, 22),
342
+ "release_tz": datetime(2011, 2, 11, 16, 44, 22, 123456, pytz.UTC),
343
+ "film": date(2013, 2, 22),
344
+ "u8": 1,
345
+ "u16": 15,
346
+ "u32": 200,
347
+ "u64": 4,
348
+ "hugedata": Decimal(-15),
349
+ },
350
+ {
351
+ "rating": 7.0,
352
+ "stars": 10,
353
+ "views": 982,
354
+ "release": datetime(2018, 11, 13, 13, 33, 11),
355
+ "release_ns": datetime(2018, 11, 13, 13, 33, 11, 123456)
356
+ if return_type == "pl"
357
+ else Timestamp("2018-11-13 13:33:11.123456"),
358
+ "release_ms": datetime(2018, 11, 13, 13, 33, 11, 123000),
359
+ "release_sec": datetime(2018, 11, 13, 13, 33, 11),
360
+ "release_tz": datetime(2018, 11, 13, 13, 33, 11, 123456, pytz.UTC),
361
+ "film": date(2014, 9, 12),
362
+ "u8": 12,
363
+ "u16": 120,
364
+ "u32": 55,
365
+ "u64": 1,
366
+ "hugedata": Decimal(-1844674407370955161511),
367
+ },
368
+ ],
369
+ )
370
+
371
+ def _test_utf8_string(_conn: lb.Connection, return_type: str, chunk_size: int | None = None) -> None:
372
+ query = "MATCH (m:movies) RETURN m.name"
373
+ data = get_result(_conn.execute(query), return_type, chunk_size)
374
+
375
+ assert_col_names(data, ["m.name"])
376
+ assert_column_equals(
377
+ data=data,
378
+ col_name="m.name",
379
+ return_type=return_type,
380
+ expected_values=["Sóló cón tu párejâ", "The 😂😃🧘🏻‍♂️🌍🌦️🍞🚗 movie", "Roma"],
381
+ )
382
+
383
+ def _test_in_small_chunk_size(_conn: lb.Connection, return_type: str, chunk_size: int | None = None) -> None:
384
+ query = "MATCH (a:person) RETURN a.age, a.fName ORDER BY a.ID"
385
+ data = get_result(_conn.execute(query), return_type, chunk_size)
386
+
387
+ assert_col_names(data, ["a.age", "a.fName"])
388
+ assert_column_equals(
389
+ data=data,
390
+ col_name="a.age",
391
+ return_type=return_type,
392
+ expected_values=[35, 30, 45, 20, 20, 25, 40, 83],
393
+ )
394
+
395
+ assert_column_equals(
396
+ data=data,
397
+ col_name="a.fName",
398
+ return_type=return_type,
399
+ expected_values=[
400
+ "Alice",
401
+ "Bob",
402
+ "Carol",
403
+ "Dan",
404
+ "Elizabeth",
405
+ "Farooq",
406
+ "Greg",
407
+ "Hubert Blaine Wolfeschlegelsteinhausenbergerdorff",
408
+ ],
409
+ )
410
+
411
+ def _test_with_nulls(_conn: lb.Connection, return_type: str, chunk_size: int | None = None) -> None:
412
+ query = "MATCH (a:person:organisation) RETURN label(a) AS `a.lbl`, a.fName, a.orgCode ORDER BY a.ID"
413
+ data = get_result(_conn.execute(query), return_type, chunk_size)
414
+
415
+ assert_col_names(data, ["a.lbl", "a.fName", "a.orgCode"])
416
+ assert_column_equals(
417
+ data=data,
418
+ col_name="a.lbl",
419
+ return_type=return_type,
420
+ expected_values=[
421
+ "person",
422
+ "organisation",
423
+ "person",
424
+ "person",
425
+ "organisation",
426
+ "person",
427
+ "organisation",
428
+ "person",
429
+ "person",
430
+ "person",
431
+ "person",
432
+ ],
433
+ )
434
+
435
+ assert_column_equals(
436
+ data=data,
437
+ col_name="a.fName",
438
+ return_type=return_type,
439
+ expected_values=[
440
+ "Alice",
441
+ None,
442
+ "Bob",
443
+ "Carol",
444
+ None,
445
+ "Dan",
446
+ None,
447
+ "Elizabeth",
448
+ "Farooq",
449
+ "Greg",
450
+ "Hubert Blaine Wolfeschlegelsteinhausenbergerdorff",
451
+ ],
452
+ )
453
+
454
+ assert_column_equals(
455
+ data=data,
456
+ col_name="a.orgCode",
457
+ return_type=return_type,
458
+ expected_values=[None, 325, None, None, 934, None, 824, None, None, None, None],
459
+ )
460
+
461
+ _test_person_table(conn, "arrow", 9)
462
+ _test_person_table(conn, "pl")
463
+ _test_movies_table(conn, "arrow", 8)
464
+ _test_movies_table(conn, "pl")
465
+ _test_utf8_string(conn, "arrow", 3)
466
+ _test_utf8_string(conn, "pl")
467
+ _test_in_small_chunk_size(conn, "arrow", 4)
468
+ _test_in_small_chunk_size(conn, "pl", 4)
469
+ _test_with_nulls(conn, "arrow", 12)
470
+ _test_with_nulls(conn, "pl")
471
+
472
+
473
+ def test_to_arrow_map(conn_db_readonly: ConnDB) -> None:
474
+ conn = conn_db_readonly[0]
475
+ results = (
476
+ conn.execute("RETURN map([1, 2, 3], [{a: 1, b: 2, c: '3'}, {a: 2, b: 3, c: '4'}, {a: 3, b: 4, c: '5'}])")
477
+ .get_as_arrow(8)[0]
478
+ .to_pylist()
479
+ )
480
+ assert results == [
481
+ [(1, {"a": 1, "b": 2, "c": "3"}), (2, {"a": 2, "b": 3, "c": "4"}), (3, {"a": 3, "b": 4, "c": "5"})]
482
+ ]
483
+
484
+ result = conn.execute("RETURN map(['abc', NULL, 'qwe'], [123, 456, 781527])")
485
+ error = "Cannot convert map with null key to Arrow: {abc=123, =456, qwe=781527}"
486
+ with pytest.raises(RuntimeError, match=error):
487
+ result.get_as_arrow(1)
488
+
489
+ result = conn.execute("RETURN map(['abc', 'xyz', 'qwe'], [123, 456, 781527])").get_as_arrow(1)
490
+ assert result[0].to_pylist(maps_as_pydicts="strict") == [{"abc": 123, "xyz": 456, "qwe": 781527}]
491
+
492
+
493
+ def test_to_arrow_array(conn_db_readwrite: ConnDB) -> None:
494
+ conn = conn_db_readwrite[0]
495
+ conn.execute("CREATE NODE TABLE school (id int64 primary key, prop1 int64[3])")
496
+ conn.execute("CREATE (p:school {id: 5})")
497
+ conn.execute("CREATE (p:school {id: 8, prop1: [2,3,1]})")
498
+ conn.execute("CREATE (p:school {id: 18, prop1: null})")
499
+ results = conn.execute("match (p:school) return p.*").get_as_arrow(8)
500
+ assert results[0].to_pylist() == [5, 8, 18]
501
+ assert results[1].to_pylist() == [None, [2, 3, 1], None]
502
+
503
+
504
+ def test_to_arrow_complex(conn_db_readonly: ConnDB) -> None:
505
+ conn, _ = conn_db_readonly
506
+
507
+ def _test_node_helper(srcStruct, dstStruct):
508
+ assert set(srcStruct.keys()) == set(dstStruct.keys())
509
+ for key in srcStruct:
510
+ if key == ID: # there isn't any guarantee on the value of _ID, so ignore it
511
+ continue
512
+ if type(srcStruct[key]) is float:
513
+ assert math.fabs(srcStruct[key] - dstStruct[key]) < 1e-5
514
+ else:
515
+ assert srcStruct[key] == dstStruct[key]
516
+
517
+ def _test_node(_conn: lb.Connection) -> None:
518
+ query = "MATCH (p:person) RETURN p ORDER BY p.ID"
519
+ query_result = _conn.execute(query)
520
+ arrow_tbl = query_result.get_as_arrow()
521
+ p_col = arrow_tbl.column(0)
522
+
523
+ for a, b in zip(
524
+ p_col.to_pylist(),
525
+ [
526
+ ground_truth.TINY_SNB_PERSONS_GROUND_TRUTH[0],
527
+ ground_truth.TINY_SNB_PERSONS_GROUND_TRUTH[2],
528
+ ground_truth.TINY_SNB_PERSONS_GROUND_TRUTH[3],
529
+ ground_truth.TINY_SNB_PERSONS_GROUND_TRUTH[5],
530
+ ground_truth.TINY_SNB_PERSONS_GROUND_TRUTH[7],
531
+ ground_truth.TINY_SNB_PERSONS_GROUND_TRUTH[8],
532
+ ground_truth.TINY_SNB_PERSONS_GROUND_TRUTH[9],
533
+ ground_truth.TINY_SNB_PERSONS_GROUND_TRUTH[10],
534
+ ],
535
+ ):
536
+ _test_node_helper(a, b)
537
+
538
+ def _test_node_rel(_conn: lb.Connection) -> None:
539
+ query = "MATCH (a:person)-[e:workAt]->(b:organisation) RETURN a, e, b ORDER BY a.ID, b.ID"
540
+ query_result = _conn.execute(query)
541
+ arrow_tbl = query_result.get_as_arrow(3)
542
+ assert arrow_tbl.num_columns == 3
543
+ a_col = arrow_tbl.column(0)
544
+ assert len(a_col) == 3
545
+ e_col = arrow_tbl.column(1)
546
+ assert len(e_col) == 3
547
+ b_col = arrow_tbl.column(2)
548
+ assert len(b_col) == 3
549
+ for a, b in zip(
550
+ a_col.to_pylist(),
551
+ [
552
+ ground_truth.TINY_SNB_PERSONS_GROUND_TRUTH[3],
553
+ ground_truth.TINY_SNB_PERSONS_GROUND_TRUTH[5],
554
+ ground_truth.TINY_SNB_PERSONS_GROUND_TRUTH[7],
555
+ ],
556
+ ):
557
+ _test_node_helper(a, b)
558
+ for a, b in zip(
559
+ e_col.to_pylist(),
560
+ [
561
+ {
562
+ SRC: {"offset": 2, "table": 0},
563
+ DST: {"offset": 1, "table": 1},
564
+ ID: {"offset": 0, "table": 5},
565
+ LABEL: "workAt",
566
+ "grading": [3.8, 2.5],
567
+ "rating": 8.2,
568
+ "year": 2015,
569
+ },
570
+ {
571
+ SRC: {"offset": 3, "table": 0},
572
+ DST: {"offset": 2, "table": 1},
573
+ ID: {"offset": 1, "table": 5},
574
+ LABEL: "workAt",
575
+ "grading": [2.1, 4.4],
576
+ "rating": 7.6,
577
+ "year": 2010,
578
+ },
579
+ {
580
+ SRC: {"offset": 4, "table": 0},
581
+ DST: {"offset": 2, "table": 1},
582
+ ID: {"offset": 2, "table": 5},
583
+ LABEL: "workAt",
584
+ "grading": [9.2, 3.1],
585
+ "rating": 9.2,
586
+ "year": 2015,
587
+ },
588
+ ],
589
+ ):
590
+ _test_node_helper(a, b)
591
+
592
+ for a, b in zip(
593
+ b_col.to_pylist(),
594
+ [
595
+ ground_truth.TINY_SNB_ORGANISATIONS_GROUND_TRUTH[4],
596
+ ground_truth.TINY_SNB_ORGANISATIONS_GROUND_TRUTH[6],
597
+ ground_truth.TINY_SNB_ORGANISATIONS_GROUND_TRUTH[6],
598
+ ],
599
+ ):
600
+ _test_node_helper(a, b)
601
+
602
+ def _test_marries_table(_conn: lb.Connection) -> None:
603
+ query = "MATCH (a:person)-[e:marries]->(b:person) RETURN e.* ORDER BY a.ID, b.ID"
604
+ arrow_tbl = _conn.execute(query).get_as_arrow(0)
605
+ assert arrow_tbl.num_columns == 3
606
+
607
+ used_addr_col = arrow_tbl.column(0)
608
+ assert used_addr_col.type == pa.list_(pa.string())
609
+ assert len(used_addr_col) == 3
610
+ assert used_addr_col.to_pylist() == [["toronto"], None, []]
611
+
612
+ addr_col = arrow_tbl.column(1)
613
+ assert addr_col.type == pa.list_(pa.int16(), 2)
614
+ assert len(addr_col) == 3
615
+ assert addr_col.to_pylist() == [[4, 5], [2, 5], [3, 9]]
616
+
617
+ note_col = arrow_tbl.column(2)
618
+ assert note_col.type == pa.string()
619
+ assert len(note_col) == 3
620
+ assert note_col.to_pylist() == [None, "long long long string", "short str"]
621
+
622
+ def _test_recursive_rel(_conn: lb.Connection) -> None:
623
+ query = "MATCH path=(a:person)-[e:knows*2..2]->(b:person) return path order by (nodes(path)[1]).ID, (nodes(path)[2]).ID, (nodes(path)[3]).ID"
624
+ arrow_tbl = _conn.execute(query).get_as_arrow(0)
625
+ # paths should be:
626
+ expected_nodes = [
627
+ [0, 2, 0],
628
+ [0, 2, 3],
629
+ [0, 2, 5],
630
+ [0, 3, 0],
631
+ [0, 3, 2],
632
+ [0, 3, 5],
633
+ [0, 5, 0],
634
+ [0, 5, 2],
635
+ [0, 5, 3],
636
+ [2, 0, 2],
637
+ [2, 0, 3],
638
+ [2, 0, 5],
639
+ [2, 3, 0],
640
+ [2, 3, 2],
641
+ [2, 3, 5],
642
+ [2, 5, 0],
643
+ [2, 5, 2],
644
+ [2, 5, 3],
645
+ [3, 0, 2],
646
+ [3, 0, 3],
647
+ [3, 0, 5],
648
+ [3, 2, 0],
649
+ [3, 2, 3],
650
+ [3, 2, 5],
651
+ [3, 5, 0],
652
+ [3, 5, 2],
653
+ [3, 5, 3],
654
+ [5, 0, 2],
655
+ [5, 0, 3],
656
+ [5, 0, 5],
657
+ [5, 2, 0],
658
+ [5, 2, 3],
659
+ [5, 2, 5],
660
+ [5, 3, 0],
661
+ [5, 3, 2],
662
+ [5, 3, 5],
663
+ ]
664
+ for row, expected in zip(arrow_tbl["path"], expected_nodes):
665
+ cur_ids = []
666
+ rel_ids = []
667
+ for node in row[NODES]:
668
+ cur_ids += [node["ID"].as_py()]
669
+ assert expected == cur_ids
670
+
671
+ def _test_serial(_conn: lb.Connection) -> None:
672
+ arrow_tbl = _conn.execute("MATCH (a:moviesSerial) RETURN a.ID AS id").get_as_arrow(0)
673
+ assert arrow_tbl.num_columns == 1
674
+ assert len(arrow_tbl) == 3
675
+ assert arrow_tbl["id"].to_pylist() == [0, 1, 2]
676
+
677
+ def _test_blob(_conn: lb.Connection) -> None:
678
+ arrow_tbl = _conn.execute("RETURN BLOB('\\\\xBC\\\\xBD\\\\xBA\\\\xAA') as result").get_as_arrow(1)
679
+ assert arrow_tbl.num_columns == 1
680
+ assert len(arrow_tbl) == 1
681
+ assert arrow_tbl["result"][0].as_py() == b"\xbc\xbd\xba\xaa"
682
+
683
+ _test_node(conn)
684
+ _test_node_rel(conn)
685
+ _test_marries_table(conn)
686
+ _test_recursive_rel(conn)
687
+ _test_serial(conn)
688
+ _test_blob(conn)
689
+
690
+ def test_to_arrow1(conn: lb.Connection) -> None:
691
+ query = "MATCH (a:person)-[e:knows]->(:person) RETURN e.summary"
692
+ res = conn.execute(query)
693
+ arrow_tbl = conn.execute(query).get_as_arrow(-1) # what is a chunk size of -1 even supposed to mean?
694
+ assert arrow_tbl == []