singlestoredb 1.14.2__cp38-abi3-win_amd64.whl → 1.15.1__cp38-abi3-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 singlestoredb might be problematic. Click here for more details.

Files changed (55) hide show
  1. _singlestoredb_accel.pyd +0 -0
  2. singlestoredb/__init__.py +2 -2
  3. singlestoredb/ai/chat.py +14 -0
  4. singlestoredb/apps/_python_udfs.py +3 -3
  5. singlestoredb/config.py +11 -0
  6. singlestoredb/docstring/__init__.py +33 -0
  7. singlestoredb/docstring/attrdoc.py +126 -0
  8. singlestoredb/docstring/common.py +230 -0
  9. singlestoredb/docstring/epydoc.py +267 -0
  10. singlestoredb/docstring/google.py +412 -0
  11. singlestoredb/docstring/numpydoc.py +562 -0
  12. singlestoredb/docstring/parser.py +100 -0
  13. singlestoredb/docstring/py.typed +1 -0
  14. singlestoredb/docstring/rest.py +256 -0
  15. singlestoredb/docstring/tests/__init__.py +1 -0
  16. singlestoredb/docstring/tests/_pydoctor.py +21 -0
  17. singlestoredb/docstring/tests/test_epydoc.py +729 -0
  18. singlestoredb/docstring/tests/test_google.py +1007 -0
  19. singlestoredb/docstring/tests/test_numpydoc.py +1100 -0
  20. singlestoredb/docstring/tests/test_parse_from_object.py +109 -0
  21. singlestoredb/docstring/tests/test_parser.py +248 -0
  22. singlestoredb/docstring/tests/test_rest.py +547 -0
  23. singlestoredb/docstring/tests/test_util.py +70 -0
  24. singlestoredb/docstring/util.py +141 -0
  25. singlestoredb/functions/decorator.py +51 -31
  26. singlestoredb/functions/ext/asgi.py +381 -35
  27. singlestoredb/functions/ext/timer.py +98 -0
  28. singlestoredb/functions/signature.py +374 -241
  29. singlestoredb/functions/typing/numpy.py +20 -0
  30. singlestoredb/functions/typing/pandas.py +2 -0
  31. singlestoredb/functions/typing/polars.py +2 -0
  32. singlestoredb/functions/typing/pyarrow.py +2 -0
  33. singlestoredb/fusion/handlers/files.py +4 -4
  34. singlestoredb/fusion/handlers/models.py +1 -1
  35. singlestoredb/fusion/handlers/stage.py +4 -4
  36. singlestoredb/magics/run_personal.py +82 -1
  37. singlestoredb/magics/run_shared.py +82 -1
  38. singlestoredb/management/__init__.py +1 -0
  39. singlestoredb/management/cluster.py +1 -1
  40. singlestoredb/management/manager.py +15 -5
  41. singlestoredb/management/region.py +104 -2
  42. singlestoredb/management/workspace.py +174 -3
  43. singlestoredb/tests/ext_funcs/__init__.py +133 -55
  44. singlestoredb/tests/test.sql +22 -0
  45. singlestoredb/tests/test_connection.py +18 -8
  46. singlestoredb/tests/test_ext_func.py +90 -0
  47. singlestoredb/tests/test_management.py +190 -0
  48. singlestoredb/tests/test_udf.py +43 -15
  49. {singlestoredb-1.14.2.dist-info → singlestoredb-1.15.1.dist-info}/METADATA +1 -1
  50. {singlestoredb-1.14.2.dist-info → singlestoredb-1.15.1.dist-info}/RECORD +55 -31
  51. /singlestoredb/functions/{typing.py → typing/__init__.py} +0 -0
  52. {singlestoredb-1.14.2.dist-info → singlestoredb-1.15.1.dist-info}/LICENSE +0 -0
  53. {singlestoredb-1.14.2.dist-info → singlestoredb-1.15.1.dist-info}/WHEEL +0 -0
  54. {singlestoredb-1.14.2.dist-info → singlestoredb-1.15.1.dist-info}/entry_points.txt +0 -0
  55. {singlestoredb-1.14.2.dist-info → singlestoredb-1.15.1.dist-info}/top_level.txt +0 -0
@@ -1,5 +1,7 @@
1
1
  #!/usr/bin/env python3
2
2
  # mypy: disable-error-code="type-arg"
3
+ import asyncio
4
+ import time
3
5
  import typing
4
6
  from typing import List
5
7
  from typing import NamedTuple
@@ -7,10 +9,6 @@ from typing import Optional
7
9
  from typing import Tuple
8
10
 
9
11
  import numpy as np
10
- import numpy.typing as npt
11
- import pandas as pd
12
- import polars as pl
13
- import pyarrow as pa
14
12
 
15
13
  import singlestoredb.functions.dtypes as dt
16
14
  from singlestoredb.functions import Masked
@@ -24,6 +22,49 @@ from singlestoredb.functions.dtypes import MEDIUMINT
24
22
  from singlestoredb.functions.dtypes import SMALLINT
25
23
  from singlestoredb.functions.dtypes import TEXT
26
24
  from singlestoredb.functions.dtypes import TINYINT
25
+ from singlestoredb.functions.typing import numpy as npt
26
+ from singlestoredb.functions.typing import pandas as pdt
27
+ from singlestoredb.functions.typing import polars as plt
28
+ from singlestoredb.functions.typing import pyarrow as pat
29
+
30
+
31
+ @udf
32
+ def doc_test(x: int, y: float) -> int:
33
+ """
34
+ A simple function to test the decorator and documentation.
35
+
36
+ Parameters
37
+ ----------
38
+ x : int
39
+ An integer to be multiplied by 2.
40
+ y : float
41
+ A float that is not used in the computation.
42
+
43
+ Examples
44
+ --------
45
+ Basic usage of the function:
46
+ >>> doc_test(
47
+ ... 3, 4.5,
48
+ ... )
49
+ 6
50
+
51
+ Another example with different values:
52
+ >>> doc_test(5, 2.0)
53
+ 10
54
+
55
+ SQL Example
56
+ sql> SELECT doc_test(3, 4.5);
57
+ 6
58
+
59
+ Final text
60
+
61
+ Returns
62
+ -------
63
+ int
64
+ The input integer multiplied by 2.
65
+
66
+ """
67
+ return x * 2
27
68
 
28
69
 
29
70
  @udf
@@ -36,19 +77,44 @@ def double_mult(x: float, y: float) -> float:
36
77
  return x * y
37
78
 
38
79
 
80
+ @udf(timeout=2)
81
+ def timeout_double_mult(x: float, y: float) -> float:
82
+ time.sleep(5)
83
+ return x * y
84
+
85
+
86
+ @udf
87
+ async def async_double_mult(x: float, y: float) -> float:
88
+ return x * y
89
+
90
+
91
+ @udf(timeout=2)
92
+ async def async_timeout_double_mult(x: float, y: float) -> float:
93
+ await asyncio.sleep(5)
94
+ return x * y
95
+
96
+
39
97
  @udf(
40
98
  args=[DOUBLE(nullable=False), DOUBLE(nullable=False)],
41
99
  returns=DOUBLE(nullable=False),
42
100
  )
43
- def pandas_double_mult(x: pd.Series, y: pd.Series) -> pd.Series:
101
+ def pandas_double_mult(x: pdt.Series, y: pdt.Series) -> pdt.Series:
44
102
  return x * y
45
103
 
46
104
 
47
105
  @udf
48
106
  def numpy_double_mult(
49
- x: npt.NDArray[np.float64],
50
- y: npt.NDArray[np.float64],
51
- ) -> npt.NDArray[np.float64]:
107
+ x: npt.Float64Array,
108
+ y: npt.Float64Array,
109
+ ) -> npt.Float64Array:
110
+ return x * y
111
+
112
+
113
+ @udf
114
+ async def async_numpy_double_mult(
115
+ x: npt.Float64Array,
116
+ y: npt.Float64Array,
117
+ ) -> npt.Float64Array:
52
118
  return x * y
53
119
 
54
120
 
@@ -56,7 +122,7 @@ def numpy_double_mult(
56
122
  args=[DOUBLE(nullable=False), DOUBLE(nullable=False)],
57
123
  returns=DOUBLE(nullable=False),
58
124
  )
59
- def arrow_double_mult(x: pa.Array, y: pa.Array) -> pa.Array:
125
+ def arrow_double_mult(x: pat.Array, y: pat.Array) -> pat.Array:
60
126
  import pyarrow.compute as pc
61
127
  return pc.multiply(x, y)
62
128
 
@@ -65,7 +131,7 @@ def arrow_double_mult(x: pa.Array, y: pa.Array) -> pa.Array:
65
131
  args=[DOUBLE(nullable=False), DOUBLE(nullable=False)],
66
132
  returns=DOUBLE(nullable=False),
67
133
  )
68
- def polars_double_mult(x: pl.Series, y: pl.Series) -> pl.Series:
134
+ def polars_double_mult(x: plt.Series, y: plt.Series) -> plt.Series:
69
135
  return x * y
70
136
 
71
137
 
@@ -106,12 +172,12 @@ def tinyint_mult(x: Optional[int], y: Optional[int]) -> Optional[int]:
106
172
 
107
173
 
108
174
  @tinyint_udf
109
- def pandas_tinyint_mult(x: pd.Series, y: pd.Series) -> pd.Series:
175
+ def pandas_tinyint_mult(x: pdt.Series, y: pdt.Series) -> pdt.Series:
110
176
  return x * y
111
177
 
112
178
 
113
179
  @tinyint_udf
114
- def polars_tinyint_mult(x: pl.Series, y: pl.Series) -> pl.Series:
180
+ def polars_tinyint_mult(x: plt.Series, y: plt.Series) -> plt.Series:
115
181
  return x * y
116
182
 
117
183
 
@@ -121,7 +187,7 @@ def numpy_tinyint_mult(x: np.ndarray, y: np.ndarray) -> np.ndarray:
121
187
 
122
188
 
123
189
  @tinyint_udf
124
- def arrow_tinyint_mult(x: pa.Array, y: pa.Array) -> pa.Array:
190
+ def arrow_tinyint_mult(x: pat.Array, y: pat.Array) -> pat.Array:
125
191
  import pyarrow.compute as pc
126
192
  return pc.multiply(x, y)
127
193
 
@@ -144,12 +210,12 @@ def smallint_mult(x: Optional[int], y: Optional[int]) -> Optional[int]:
144
210
 
145
211
 
146
212
  @smallint_udf
147
- def pandas_smallint_mult(x: pd.Series, y: pd.Series) -> pd.Series:
213
+ def pandas_smallint_mult(x: pdt.Series, y: pdt.Series) -> pdt.Series:
148
214
  return x * y
149
215
 
150
216
 
151
217
  @smallint_udf
152
- def polars_smallint_mult(x: pl.Series, y: pl.Series) -> pl.Series:
218
+ def polars_smallint_mult(x: plt.Series, y: plt.Series) -> plt.Series:
153
219
  return x * y
154
220
 
155
221
 
@@ -159,7 +225,7 @@ def numpy_smallint_mult(x: np.ndarray, y: np.ndarray) -> np.ndarray:
159
225
 
160
226
 
161
227
  @smallint_udf
162
- def arrow_smallint_mult(x: pa.Array, y: pa.Array) -> pa.Array:
228
+ def arrow_smallint_mult(x: pat.Array, y: pat.Array) -> pat.Array:
163
229
  import pyarrow.compute as pc
164
230
  return pc.multiply(x, y)
165
231
 
@@ -183,12 +249,12 @@ def mediumint_mult(x: Optional[int], y: Optional[int]) -> Optional[int]:
183
249
 
184
250
 
185
251
  @mediumint_udf
186
- def pandas_mediumint_mult(x: pd.Series, y: pd.Series) -> pd.Series:
252
+ def pandas_mediumint_mult(x: pdt.Series, y: pdt.Series) -> pdt.Series:
187
253
  return x * y
188
254
 
189
255
 
190
256
  @mediumint_udf
191
- def polars_mediumint_mult(x: pl.Series, y: pl.Series) -> pl.Series:
257
+ def polars_mediumint_mult(x: plt.Series, y: plt.Series) -> plt.Series:
192
258
  return x * y
193
259
 
194
260
 
@@ -198,7 +264,7 @@ def numpy_mediumint_mult(x: np.ndarray, y: np.ndarray) -> np.ndarray:
198
264
 
199
265
 
200
266
  @mediumint_udf
201
- def arrow_mediumint_mult(x: pa.Array, y: pa.Array) -> pa.Array:
267
+ def arrow_mediumint_mult(x: pat.Array, y: pat.Array) -> pat.Array:
202
268
  import pyarrow.compute as pc
203
269
  return pc.multiply(x, y)
204
270
 
@@ -222,12 +288,12 @@ def bigint_mult(x: Optional[int], y: Optional[int]) -> Optional[int]:
222
288
 
223
289
 
224
290
  @bigint_udf
225
- def pandas_bigint_mult(x: pd.Series, y: pd.Series) -> pd.Series:
291
+ def pandas_bigint_mult(x: pdt.Series, y: pdt.Series) -> pdt.Series:
226
292
  return x * y
227
293
 
228
294
 
229
295
  @bigint_udf
230
- def polars_bigint_mult(x: pl.Series, y: pl.Series) -> pl.Series:
296
+ def polars_bigint_mult(x: plt.Series, y: plt.Series) -> plt.Series:
231
297
  return x * y
232
298
 
233
299
 
@@ -237,7 +303,7 @@ def numpy_bigint_mult(x: np.ndarray, y: np.ndarray) -> np.ndarray:
237
303
 
238
304
 
239
305
  @bigint_udf
240
- def arrow_bigint_mult(x: pa.Array, y: pa.Array) -> pa.Array:
306
+ def arrow_bigint_mult(x: pat.Array, y: pat.Array) -> pat.Array:
241
307
  import pyarrow.compute as pc
242
308
  return pc.multiply(x, y)
243
309
 
@@ -261,12 +327,12 @@ def nullable_tinyint_mult(x: Optional[int], y: Optional[int]) -> Optional[int]:
261
327
 
262
328
 
263
329
  @nullable_tinyint_udf
264
- def pandas_nullable_tinyint_mult(x: pd.Series, y: pd.Series) -> pd.Series:
330
+ def pandas_nullable_tinyint_mult(x: pdt.Series, y: pdt.Series) -> pdt.Series:
265
331
  return x * y
266
332
 
267
333
 
268
334
  @nullable_tinyint_udf
269
- def polars_nullable_tinyint_mult(x: pl.Series, y: pl.Series) -> pl.Series:
335
+ def polars_nullable_tinyint_mult(x: plt.Series, y: plt.Series) -> plt.Series:
270
336
  return x * y
271
337
 
272
338
 
@@ -276,7 +342,7 @@ def numpy_nullable_tinyint_mult(x: np.ndarray, y: np.ndarray) -> np.ndarray:
276
342
 
277
343
 
278
344
  @nullable_tinyint_udf
279
- def arrow_nullable_tinyint_mult(x: pa.Array, y: pa.Array) -> pa.Array:
345
+ def arrow_nullable_tinyint_mult(x: pat.Array, y: pat.Array) -> pat.Array:
280
346
  import pyarrow.compute as pc
281
347
  return pc.multiply(x, y)
282
348
 
@@ -299,12 +365,12 @@ def nullable_smallint_mult(x: Optional[int], y: Optional[int]) -> Optional[int]:
299
365
 
300
366
 
301
367
  @nullable_smallint_udf
302
- def pandas_nullable_smallint_mult(x: pd.Series, y: pd.Series) -> pd.Series:
368
+ def pandas_nullable_smallint_mult(x: pdt.Series, y: pdt.Series) -> pdt.Series:
303
369
  return x * y
304
370
 
305
371
 
306
372
  @nullable_smallint_udf
307
- def polars_nullable_smallint_mult(x: pl.Series, y: pl.Series) -> pl.Series:
373
+ def polars_nullable_smallint_mult(x: plt.Series, y: plt.Series) -> plt.Series:
308
374
  return x * y
309
375
 
310
376
 
@@ -314,7 +380,7 @@ def numpy_nullable_smallint_mult(x: np.ndarray, y: np.ndarray) -> np.ndarray:
314
380
 
315
381
 
316
382
  @nullable_smallint_udf
317
- def arrow_nullable_smallint_mult(x: pa.Array, y: pa.Array) -> pa.Array:
383
+ def arrow_nullable_smallint_mult(x: pat.Array, y: pat.Array) -> pat.Array:
318
384
  import pyarrow.compute as pc
319
385
  return pc.multiply(x, y)
320
386
 
@@ -338,12 +404,12 @@ def nullable_mediumint_mult(x: Optional[int], y: Optional[int]) -> Optional[int]
338
404
 
339
405
 
340
406
  @nullable_mediumint_udf
341
- def pandas_nullable_mediumint_mult(x: pd.Series, y: pd.Series) -> pd.Series:
407
+ def pandas_nullable_mediumint_mult(x: pdt.Series, y: pdt.Series) -> pdt.Series:
342
408
  return x * y
343
409
 
344
410
 
345
411
  @nullable_mediumint_udf
346
- def polars_nullable_mediumint_mult(x: pl.Series, y: pl.Series) -> pl.Series:
412
+ def polars_nullable_mediumint_mult(x: plt.Series, y: plt.Series) -> plt.Series:
347
413
  return x * y
348
414
 
349
415
 
@@ -353,7 +419,7 @@ def numpy_nullable_mediumint_mult(x: np.ndarray, y: np.ndarray) -> np.ndarray:
353
419
 
354
420
 
355
421
  @nullable_mediumint_udf
356
- def arrow_nullable_mediumint_mult(x: pa.Array, y: pa.Array) -> pa.Array:
422
+ def arrow_nullable_mediumint_mult(x: pat.Array, y: pat.Array) -> pat.Array:
357
423
  import pyarrow.compute as pc
358
424
  return pc.multiply(x, y)
359
425
 
@@ -377,12 +443,12 @@ def nullable_bigint_mult(x: Optional[int], y: Optional[int]) -> Optional[int]:
377
443
 
378
444
 
379
445
  @nullable_bigint_udf
380
- def pandas_nullable_bigint_mult(x: pd.Series, y: pd.Series) -> pd.Series:
446
+ def pandas_nullable_bigint_mult(x: pdt.Series, y: pdt.Series) -> pdt.Series:
381
447
  return x * y
382
448
 
383
449
 
384
450
  @nullable_bigint_udf
385
- def polars_nullable_bigint_mult(x: pl.Series, y: pl.Series) -> pl.Series:
451
+ def polars_nullable_bigint_mult(x: plt.Series, y: plt.Series) -> plt.Series:
386
452
  return x * y
387
453
 
388
454
 
@@ -392,7 +458,7 @@ def numpy_nullable_bigint_mult(x: np.ndarray, y: np.ndarray) -> np.ndarray:
392
458
 
393
459
 
394
460
  @nullable_bigint_udf
395
- def arrow_nullable_bigint_mult(x: pa.Array, y: pa.Array) -> pa.Array:
461
+ def arrow_nullable_bigint_mult(x: pat.Array, y: pat.Array) -> pat.Array:
396
462
  import pyarrow.compute as pc
397
463
  return pc.multiply(x, y)
398
464
 
@@ -410,7 +476,7 @@ def string_mult(x: str, times: int) -> str:
410
476
 
411
477
 
412
478
  @udf(args=[TEXT(nullable=False), BIGINT(nullable=False)], returns=TEXT(nullable=False))
413
- def pandas_string_mult(x: pd.Series, times: pd.Series) -> pd.Series:
479
+ def pandas_string_mult(x: pdt.Series, times: pdt.Series) -> pdt.Series:
414
480
  return x * times
415
481
 
416
482
 
@@ -447,8 +513,8 @@ def nullable_string_mult(x: Optional[str], times: Optional[int]) -> Optional[str
447
513
  returns=TINYINT(nullable=True),
448
514
  )
449
515
  def pandas_nullable_tinyint_mult_with_masks(
450
- x: Masked[pd.Series], y: Masked[pd.Series],
451
- ) -> Masked[pd.Series]:
516
+ x: Masked[pdt.Series], y: Masked[pdt.Series],
517
+ ) -> Masked[pdt.Series]:
452
518
  x_data, x_nulls = x
453
519
  y_data, y_nulls = y
454
520
  return Masked(x_data * y_data, x_nulls | y_nulls)
@@ -468,8 +534,8 @@ def numpy_nullable_tinyint_mult_with_masks(
468
534
  returns=TINYINT(nullable=True),
469
535
  )
470
536
  def polars_nullable_tinyint_mult_with_masks(
471
- x: Masked[pl.Series], y: Masked[pl.Series],
472
- ) -> Masked[pl.Series]:
537
+ x: Masked[plt.Series], y: Masked[plt.Series],
538
+ ) -> Masked[plt.Series]:
473
539
  x_data, x_nulls = x
474
540
  y_data, y_nulls = y
475
541
  return Masked(x_data * y_data, x_nulls | y_nulls)
@@ -480,8 +546,8 @@ def polars_nullable_tinyint_mult_with_masks(
480
546
  returns=TINYINT(nullable=True),
481
547
  )
482
548
  def arrow_nullable_tinyint_mult_with_masks(
483
- x: Masked[pa.Array], y: Masked[pa.Array],
484
- ) -> Masked[pa.Array]:
549
+ x: Masked[pat.Array], y: Masked[pat.Array],
550
+ ) -> Masked[pat.Array]:
485
551
  import pyarrow.compute as pc
486
552
  x_data, x_nulls = x
487
553
  y_data, y_nulls = y
@@ -489,7 +555,7 @@ def arrow_nullable_tinyint_mult_with_masks(
489
555
 
490
556
 
491
557
  @udf(returns=[TEXT(nullable=False, name='res')])
492
- def numpy_fixed_strings() -> Table[npt.NDArray[np.str_]]:
558
+ def numpy_fixed_strings() -> Table[npt.StrArray]:
493
559
  out = np.array(
494
560
  [
495
561
  'hello',
@@ -502,7 +568,7 @@ def numpy_fixed_strings() -> Table[npt.NDArray[np.str_]]:
502
568
 
503
569
 
504
570
  @udf(returns=[TEXT(nullable=False, name='res'), TINYINT(nullable=False, name='res2')])
505
- def numpy_fixed_strings_2() -> Table[npt.NDArray[np.str_], npt.NDArray[np.int8]]:
571
+ def numpy_fixed_strings_2() -> Table[npt.StrArray, npt.Int8Array]:
506
572
  out = np.array(
507
573
  [
508
574
  'hello',
@@ -515,7 +581,7 @@ def numpy_fixed_strings_2() -> Table[npt.NDArray[np.str_], npt.NDArray[np.int8]]
515
581
 
516
582
 
517
583
  @udf(returns=[BLOB(nullable=False, name='res')])
518
- def numpy_fixed_binary() -> Table[npt.NDArray[np.bytes_]]:
584
+ def numpy_fixed_binary() -> Table[npt.BytesArray]:
519
585
  out = np.array(
520
586
  [
521
587
  'hello'.encode('utf8'),
@@ -537,6 +603,11 @@ def table_function(n: int) -> Table[List[int]]:
537
603
  return Table([10] * n)
538
604
 
539
605
 
606
+ @udf
607
+ async def async_table_function(n: int) -> Table[List[int]]:
608
+ return Table([10] * n)
609
+
610
+
540
611
  @udf(
541
612
  returns=[
542
613
  dt.INT(name='c_int', nullable=False),
@@ -561,8 +632,8 @@ def table_function_struct(n: int) -> Table[List[MyTable]]:
561
632
 
562
633
  @udf
563
634
  def vec_function(
564
- x: npt.NDArray[np.float64], y: npt.NDArray[np.float64],
565
- ) -> npt.NDArray[np.float64]:
635
+ x: npt.Float64Array, y: npt.Float64Array,
636
+ ) -> npt.Float64Array:
566
637
  return x * y
567
638
 
568
639
 
@@ -577,8 +648,8 @@ class VecOutputs(typing.NamedTuple):
577
648
 
578
649
  @udf(args=VecInputs, returns=VecOutputs)
579
650
  def vec_function_ints(
580
- x: npt.NDArray[np.int_], y: npt.NDArray[np.int_],
581
- ) -> npt.NDArray[np.int_]:
651
+ x: npt.IntArray, y: npt.IntArray,
652
+ ) -> npt.IntArray:
582
653
  return x * y
583
654
 
584
655
 
@@ -589,9 +660,16 @@ class DFOutputs(typing.NamedTuple):
589
660
 
590
661
  @udf(args=VecInputs, returns=DFOutputs)
591
662
  def vec_function_df(
592
- x: npt.NDArray[np.int_], y: npt.NDArray[np.int_],
593
- ) -> Table[pd.DataFrame]:
594
- return pd.DataFrame(dict(res=[1, 2, 3], res2=[1.1, 2.2, 3.3]))
663
+ x: npt.IntArray, y: npt.IntArray,
664
+ ) -> Table[pdt.DataFrame]:
665
+ return pdt.DataFrame(dict(res=[1, 2, 3], res2=[1.1, 2.2, 3.3]))
666
+
667
+
668
+ @udf(args=VecInputs, returns=DFOutputs)
669
+ async def async_vec_function_df(
670
+ x: npt.IntArray, y: npt.IntArray,
671
+ ) -> Table[pdt.DataFrame]:
672
+ return pdt.DataFrame(dict(res=[1, 2, 3], res2=[1.1, 2.2, 3.3]))
595
673
 
596
674
 
597
675
  class MaskOutputs(typing.NamedTuple):
@@ -600,8 +678,8 @@ class MaskOutputs(typing.NamedTuple):
600
678
 
601
679
  @udf(args=VecInputs, returns=MaskOutputs)
602
680
  def vec_function_ints_masked(
603
- x: Masked[npt.NDArray[np.int_]], y: Masked[npt.NDArray[np.int_]],
604
- ) -> Table[Masked[npt.NDArray[np.int_]]]:
681
+ x: Masked[npt.IntArray], y: Masked[npt.IntArray],
682
+ ) -> Table[Masked[npt.IntArray]]:
605
683
  x_data, x_nulls = x
606
684
  y_data, y_nulls = y
607
685
  return Table(Masked(x_data * y_data, x_nulls | y_nulls))
@@ -614,8 +692,8 @@ class MaskOutputs2(typing.NamedTuple):
614
692
 
615
693
  @udf(args=VecInputs, returns=MaskOutputs2)
616
694
  def vec_function_ints_masked2(
617
- x: Masked[npt.NDArray[np.int_]], y: Masked[npt.NDArray[np.int_]],
618
- ) -> Table[Masked[npt.NDArray[np.int_]], Masked[npt.NDArray[np.int_]]]:
695
+ x: Masked[npt.IntArray], y: Masked[npt.IntArray],
696
+ ) -> Table[Masked[npt.IntArray], Masked[npt.IntArray]]:
619
697
  x_data, x_nulls = x
620
698
  y_data, y_nulls = y
621
699
  return Table(
@@ -14,6 +14,28 @@ INSERT INTO data SET id='e', name='elephants', value=0;
14
14
 
15
15
  COMMIT;
16
16
 
17
+ CREATE ROWSTORE TABLE IF NOT EXISTS longer_data (
18
+ id VARCHAR(255) NOT NULL,
19
+ name VARCHAR(255) NOT NULL,
20
+ value BIGINT NOT NULL,
21
+ PRIMARY KEY (id) USING HASH
22
+ ) DEFAULT CHARSET = utf8 COLLATE = utf8_unicode_ci;
23
+
24
+ INSERT INTO longer_data SET id='a', name='antelopes', value=2;
25
+ INSERT INTO longer_data SET id='b', name='bears', value=2;
26
+ INSERT INTO longer_data SET id='c', name='cats', value=5;
27
+ INSERT INTO longer_data SET id='d', name='dogs', value=4;
28
+ INSERT INTO longer_data SET id='e', name='elephants', value=0;
29
+ INSERT INTO longer_data SET id='f', name='ferrets', value=2;
30
+ INSERT INTO longer_data SET id='g', name='gorillas', value=4;
31
+ INSERT INTO longer_data SET id='h', name='horses', value=6;
32
+ INSERT INTO longer_data SET id='i', name='iguanas', value=2;
33
+ INSERT INTO longer_data SET id='j', name='jaguars', value=0;
34
+ INSERT INTO longer_data SET id='k', name='kiwis', value=0;
35
+ INSERT INTO longer_data SET id='l', name='leopards', value=1;
36
+
37
+ COMMIT;
38
+
17
39
  CREATE ROWSTORE TABLE IF NOT EXISTS data_with_nulls (
18
40
  id VARCHAR(255) NOT NULL,
19
41
  name VARCHAR(255),
@@ -1443,6 +1443,11 @@ class TestConnection(unittest.TestCase):
1443
1443
  out = cur.fetchone()
1444
1444
  row = dict(zip(names, out.row(0)))
1445
1445
 
1446
+ # Recent versions of polars have a problem with decimals
1447
+ class FixCompare(str):
1448
+ def __eq__(self, other):
1449
+ return super().__eq__(other.replace('precision=None', 'precision=22'))
1450
+
1446
1451
  dtypes = [
1447
1452
  ('id', 'Int32'),
1448
1453
  ('tinyint', 'Int8'),
@@ -1464,10 +1469,10 @@ class TestConnection(unittest.TestCase):
1464
1469
  ('float', 'Float32'),
1465
1470
  ('double', 'Float64'),
1466
1471
  ('real', 'Float64'),
1467
- ('decimal', 'Decimal(precision=22, scale=6)'),
1468
- ('dec', 'Decimal(precision=22, scale=6)'),
1469
- ('fixed', 'Decimal(precision=22, scale=6)'),
1470
- ('numeric', 'Decimal(precision=22, scale=6)'),
1472
+ ('decimal', FixCompare('Decimal(precision=22, scale=6)')),
1473
+ ('dec', FixCompare('Decimal(precision=22, scale=6)')),
1474
+ ('fixed', FixCompare('Decimal(precision=22, scale=6)')),
1475
+ ('numeric', FixCompare('Decimal(precision=22, scale=6)')),
1471
1476
  ('date', 'Date'),
1472
1477
  ('time', "Duration(time_unit='us')"),
1473
1478
  ('time_6', "Duration(time_unit='us')"),
@@ -1585,6 +1590,11 @@ class TestConnection(unittest.TestCase):
1585
1590
  out = cur.fetchone()
1586
1591
  row = dict(zip(names, out.row(0)))
1587
1592
 
1593
+ # Recent versions of polars have a problem with decimals
1594
+ class FixCompare(str):
1595
+ def __eq__(self, other):
1596
+ return super().__eq__(other.replace('precision=None', 'precision=22'))
1597
+
1588
1598
  dtypes = [
1589
1599
  ('id', 'Int32'),
1590
1600
  ('tinyint', 'Int8'),
@@ -1606,10 +1616,10 @@ class TestConnection(unittest.TestCase):
1606
1616
  ('float', 'Float32'),
1607
1617
  ('double', 'Float64'),
1608
1618
  ('real', 'Float64'),
1609
- ('decimal', 'Decimal(precision=22, scale=6)'),
1610
- ('dec', 'Decimal(precision=22, scale=6)'),
1611
- ('fixed', 'Decimal(precision=22, scale=6)'),
1612
- ('numeric', 'Decimal(precision=22, scale=6)'),
1619
+ ('decimal', FixCompare('Decimal(precision=22, scale=6)')),
1620
+ ('dec', FixCompare('Decimal(precision=22, scale=6)')),
1621
+ ('fixed', FixCompare('Decimal(precision=22, scale=6)')),
1622
+ ('numeric', FixCompare('Decimal(precision=22, scale=6)')),
1613
1623
  ('date', 'Date'),
1614
1624
  ('time', "Duration(time_unit='us')"),
1615
1625
  ('time_6', "Duration(time_unit='us')"),
@@ -162,6 +162,43 @@ class TestExtFunc(unittest.TestCase):
162
162
  'from data order by id',
163
163
  )
164
164
 
165
+ def test_timeout_double_mult(self):
166
+ with self.assertRaises(self.conn.OperationalError) as exc:
167
+ self.cur.execute(
168
+ 'select timeout_double_mult(value, 100) as res '
169
+ 'from longer_data order by id',
170
+ )
171
+ assert 'timeout' in str(exc.exception).lower()
172
+
173
+ def test_async_double_mult(self):
174
+ self.cur.execute(
175
+ 'select async_double_mult(value, 100) as res from data order by id',
176
+ )
177
+
178
+ assert [tuple(x) for x in self.cur] == \
179
+ [(200.0,), (200.0,), (500.0,), (400.0,), (0.0,)]
180
+
181
+ desc = self.cur.description
182
+ assert len(desc) == 1
183
+ assert desc[0].name == 'res'
184
+ assert desc[0].type_code == ft.DOUBLE
185
+ assert desc[0].null_ok is False
186
+
187
+ # NULL is not valid
188
+ with self.assertRaises(self.conn.OperationalError):
189
+ self.cur.execute(
190
+ 'select async_double_mult(value, NULL) as res '
191
+ 'from data order by id',
192
+ )
193
+
194
+ def test_async_timeout_double_mult(self):
195
+ with self.assertRaises(self.conn.OperationalError) as exc:
196
+ self.cur.execute(
197
+ 'select async_timeout_double_mult(value, 100) as res '
198
+ 'from longer_data order by id',
199
+ )
200
+ assert 'timeout' in str(exc.exception).lower()
201
+
165
202
  def test_pandas_double_mult(self):
166
203
  self.cur.execute(
167
204
  'select pandas_double_mult(value, 100) as res '
@@ -206,6 +243,28 @@ class TestExtFunc(unittest.TestCase):
206
243
  'from data order by id',
207
244
  )
208
245
 
246
+ def test_async_numpy_double_mult(self):
247
+ self.cur.execute(
248
+ 'select async_numpy_double_mult(value, 100) as res '
249
+ 'from data order by id',
250
+ )
251
+
252
+ assert [tuple(x) for x in self.cur] == \
253
+ [(200.0,), (200.0,), (500.0,), (400.0,), (0.0,)]
254
+
255
+ desc = self.cur.description
256
+ assert len(desc) == 1
257
+ assert desc[0].name == 'res'
258
+ assert desc[0].type_code == ft.DOUBLE
259
+ assert desc[0].null_ok is False
260
+
261
+ # NULL is not valid
262
+ with self.assertRaises(self.conn.OperationalError):
263
+ self.cur.execute(
264
+ 'select async_numpy_double_mult(value, NULL) as res '
265
+ 'from data order by id',
266
+ )
267
+
209
268
  def test_arrow_double_mult(self):
210
269
  self.cur.execute(
211
270
  'select arrow_double_mult(value, 100) as res '
@@ -1246,6 +1305,17 @@ class TestExtFunc(unittest.TestCase):
1246
1305
  assert desc[0].type_code == ft.LONGLONG
1247
1306
  assert desc[0].null_ok is False
1248
1307
 
1308
+ def test_async_table_function(self):
1309
+ self.cur.execute('select * from async_table_function(5)')
1310
+
1311
+ assert [x[0] for x in self.cur] == [10, 10, 10, 10, 10]
1312
+
1313
+ desc = self.cur.description
1314
+ assert len(desc) == 1
1315
+ assert desc[0].name == 'a'
1316
+ assert desc[0].type_code == ft.LONGLONG
1317
+ assert desc[0].null_ok is False
1318
+
1249
1319
  def test_table_function_tuple(self):
1250
1320
  self.cur.execute('select * from table_function_tuple(3)')
1251
1321
 
@@ -1310,6 +1380,26 @@ class TestExtFunc(unittest.TestCase):
1310
1380
  assert desc[1].type_code == ft.DOUBLE
1311
1381
  assert desc[1].null_ok is False
1312
1382
 
1383
+ def test_async_vec_function_df(self):
1384
+ self.cur.execute('select * from async_vec_function_df(5, 10)')
1385
+
1386
+ out = list(self.cur)
1387
+
1388
+ assert out == [
1389
+ (1, 1.1),
1390
+ (2, 2.2),
1391
+ (3, 3.3),
1392
+ ]
1393
+
1394
+ desc = self.cur.description
1395
+ assert len(desc) == 2
1396
+ assert desc[0].name == 'res'
1397
+ assert desc[0].type_code == ft.SHORT
1398
+ assert desc[0].null_ok is False
1399
+ assert desc[1].name == 'res2'
1400
+ assert desc[1].type_code == ft.DOUBLE
1401
+ assert desc[1].null_ok is False
1402
+
1313
1403
  def test_vec_function_ints_masked(self):
1314
1404
  self.cur.execute('select * from vec_function_ints_masked(5, 10)')
1315
1405