graphique 1.7__tar.gz → 1.9__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.
- {graphique-1.7 → graphique-1.9}/PKG-INFO +14 -22
- {graphique-1.7 → graphique-1.9}/README.md +5 -2
- {graphique-1.7 → graphique-1.9}/graphique/core.py +126 -175
- {graphique-1.7 → graphique-1.9}/graphique/inputs.py +166 -165
- {graphique-1.7 → graphique-1.9}/graphique/interface.py +119 -168
- {graphique-1.7 → graphique-1.9}/graphique/middleware.py +11 -26
- {graphique-1.7 → graphique-1.9}/graphique/models.py +37 -48
- {graphique-1.7 → graphique-1.9}/graphique/scalars.py +19 -38
- {graphique-1.7 → graphique-1.9}/graphique/service.py +4 -0
- {graphique-1.7 → graphique-1.9}/graphique/shell.py +9 -6
- {graphique-1.7 → graphique-1.9}/graphique.egg-info/PKG-INFO +14 -22
- {graphique-1.7 → graphique-1.9}/graphique.egg-info/requires.txt +3 -1
- {graphique-1.7 → graphique-1.9}/pyproject.toml +5 -7
- {graphique-1.7 → graphique-1.9}/requirements.in +3 -1
- {graphique-1.7 → graphique-1.9}/tests/test_bench.py +4 -4
- {graphique-1.7 → graphique-1.9}/tests/test_core.py +25 -88
- {graphique-1.7 → graphique-1.9}/tests/test_dataset.py +13 -15
- {graphique-1.7 → graphique-1.9}/tests/test_models.py +13 -11
- {graphique-1.7 → graphique-1.9}/tests/test_service.py +1 -1
- {graphique-1.7 → graphique-1.9}/LICENSE.txt +0 -0
- {graphique-1.7 → graphique-1.9}/graphique/__init__.py +0 -0
- {graphique-1.7 → graphique-1.9}/graphique/py.typed +0 -0
- {graphique-1.7 → graphique-1.9}/graphique.egg-info/SOURCES.txt +0 -0
- {graphique-1.7 → graphique-1.9}/graphique.egg-info/dependency_links.txt +0 -0
- {graphique-1.7 → graphique-1.9}/graphique.egg-info/top_level.txt +0 -0
- {graphique-1.7 → graphique-1.9}/setup.cfg +0 -0
|
@@ -1,22 +1,9 @@
|
|
|
1
|
-
Metadata-Version: 2.
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
2
|
Name: graphique
|
|
3
|
-
Version: 1.
|
|
3
|
+
Version: 1.9
|
|
4
4
|
Summary: GraphQL service for arrow tables and parquet data sets.
|
|
5
5
|
Author-email: Aric Coady <aric.coady@gmail.com>
|
|
6
|
-
License:
|
|
7
|
-
|
|
8
|
-
Licensed under the Apache License, Version 2.0 (the "License");
|
|
9
|
-
you may not use this file except in compliance with the License.
|
|
10
|
-
You may obtain a copy of the License at
|
|
11
|
-
|
|
12
|
-
http://www.apache.org/licenses/LICENSE-2.0
|
|
13
|
-
|
|
14
|
-
Unless required by applicable law or agreed to in writing, software
|
|
15
|
-
distributed under the License is distributed on an "AS IS" BASIS,
|
|
16
|
-
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
17
|
-
See the License for the specific language governing permissions and
|
|
18
|
-
limitations under the License.
|
|
19
|
-
|
|
6
|
+
License: Apache-2.0
|
|
20
7
|
Project-URL: Homepage, https://github.com/coady/graphique
|
|
21
8
|
Project-URL: Documentation, https://coady.github.io/graphique
|
|
22
9
|
Project-URL: Changelog, https://github.com/coady/graphique/blob/main/CHANGELOG.md
|
|
@@ -25,27 +12,29 @@ Keywords: graphql,arrow,parquet
|
|
|
25
12
|
Classifier: Development Status :: 5 - Production/Stable
|
|
26
13
|
Classifier: Intended Audience :: Developers
|
|
27
14
|
Classifier: Intended Audience :: Science/Research
|
|
28
|
-
Classifier: License :: OSI Approved :: Apache Software License
|
|
29
15
|
Classifier: Operating System :: OS Independent
|
|
30
16
|
Classifier: Programming Language :: Python :: 3
|
|
31
|
-
Classifier: Programming Language :: Python :: 3.9
|
|
32
17
|
Classifier: Programming Language :: Python :: 3.10
|
|
33
18
|
Classifier: Programming Language :: Python :: 3.11
|
|
34
19
|
Classifier: Programming Language :: Python :: 3.12
|
|
20
|
+
Classifier: Programming Language :: Python :: 3.13
|
|
35
21
|
Classifier: Topic :: Database :: Database Engines/Servers
|
|
36
22
|
Classifier: Topic :: Internet :: WWW/HTTP :: HTTP Servers
|
|
37
23
|
Classifier: Topic :: Internet :: WWW/HTTP :: Indexing/Search
|
|
38
24
|
Classifier: Topic :: Software Development :: Libraries :: Python Modules
|
|
39
25
|
Classifier: Typing :: Typed
|
|
40
|
-
Requires-Python: >=3.
|
|
26
|
+
Requires-Python: >=3.10
|
|
41
27
|
Description-Content-Type: text/markdown
|
|
42
28
|
License-File: LICENSE.txt
|
|
43
|
-
Requires-Dist: pyarrow>=
|
|
29
|
+
Requires-Dist: pyarrow>=20
|
|
44
30
|
Requires-Dist: strawberry-graphql[asgi,cli]>=0.236
|
|
31
|
+
Requires-Dist: numpy
|
|
32
|
+
Requires-Dist: isodate>=0.7
|
|
45
33
|
Provides-Extra: server
|
|
46
34
|
Requires-Dist: uvicorn[standard]; extra == "server"
|
|
47
35
|
Provides-Extra: cli
|
|
48
36
|
Requires-Dist: tqdm; extra == "cli"
|
|
37
|
+
Dynamic: license-file
|
|
49
38
|
|
|
50
39
|
[](https://pypi.org/project/graphique/)
|
|
51
40
|

|
|
@@ -89,8 +78,9 @@ Supply a mapping of names to datasets for multiple roots, and to enable federati
|
|
|
89
78
|
import pyarrow.dataset as ds
|
|
90
79
|
from graphique import GraphQL
|
|
91
80
|
|
|
92
|
-
|
|
93
|
-
app = GraphQL
|
|
81
|
+
source = ds.dataset(...)
|
|
82
|
+
app = GraphQL(source) # Table is root query type
|
|
83
|
+
app = GraphQL.federated({<name>: source, ...}, keys={<name>: [], ...}) # Tables on federated fields
|
|
94
84
|
```
|
|
95
85
|
|
|
96
86
|
Start like any ASGI app.
|
|
@@ -147,6 +137,8 @@ If index columns are detected in the schema metadata, then an initial `filter` w
|
|
|
147
137
|
## Dependencies
|
|
148
138
|
* pyarrow
|
|
149
139
|
* strawberry-graphql[asgi,cli]
|
|
140
|
+
* numpy
|
|
141
|
+
* isodate
|
|
150
142
|
* uvicorn (or other [ASGI server](https://asgi.readthedocs.io/en/latest/implementations.html))
|
|
151
143
|
|
|
152
144
|
## Tests
|
|
@@ -40,8 +40,9 @@ Supply a mapping of names to datasets for multiple roots, and to enable federati
|
|
|
40
40
|
import pyarrow.dataset as ds
|
|
41
41
|
from graphique import GraphQL
|
|
42
42
|
|
|
43
|
-
|
|
44
|
-
app = GraphQL
|
|
43
|
+
source = ds.dataset(...)
|
|
44
|
+
app = GraphQL(source) # Table is root query type
|
|
45
|
+
app = GraphQL.federated({<name>: source, ...}, keys={<name>: [], ...}) # Tables on federated fields
|
|
45
46
|
```
|
|
46
47
|
|
|
47
48
|
Start like any ASGI app.
|
|
@@ -98,6 +99,8 @@ If index columns are detected in the schema metadata, then an initial `filter` w
|
|
|
98
99
|
## Dependencies
|
|
99
100
|
* pyarrow
|
|
100
101
|
* strawberry-graphql[asgi,cli]
|
|
102
|
+
* numpy
|
|
103
|
+
* isodate
|
|
101
104
|
* uvicorn (or other [ASGI server](https://asgi.readthedocs.io/en/latest/implementations.html))
|
|
102
105
|
|
|
103
106
|
## Tests
|
|
@@ -12,9 +12,9 @@ import inspect
|
|
|
12
12
|
import itertools
|
|
13
13
|
import operator
|
|
14
14
|
import json
|
|
15
|
-
from collections.abc import Callable, Iterable, Iterator, Mapping
|
|
15
|
+
from collections.abc import Callable, Iterable, Iterator, Mapping
|
|
16
16
|
from dataclasses import dataclass
|
|
17
|
-
from typing import
|
|
17
|
+
from typing import TypeAlias, get_type_hints
|
|
18
18
|
import numpy as np
|
|
19
19
|
import pyarrow as pa
|
|
20
20
|
import pyarrow.acero as ac
|
|
@@ -22,8 +22,10 @@ import pyarrow.compute as pc
|
|
|
22
22
|
import pyarrow.dataset as ds
|
|
23
23
|
from typing_extensions import Self
|
|
24
24
|
|
|
25
|
-
Array =
|
|
26
|
-
Batch =
|
|
25
|
+
Array: TypeAlias = pa.Array | pa.ChunkedArray
|
|
26
|
+
Batch: TypeAlias = pa.RecordBatch | pa.Table
|
|
27
|
+
bit_any = functools.partial(functools.reduce, operator.or_)
|
|
28
|
+
bit_all = functools.partial(functools.reduce, operator.and_)
|
|
27
29
|
|
|
28
30
|
|
|
29
31
|
class Agg:
|
|
@@ -51,36 +53,22 @@ class Agg:
|
|
|
51
53
|
'tdigest': pc.TDigestOptions,
|
|
52
54
|
'variance': pc.VarianceOptions,
|
|
53
55
|
}
|
|
54
|
-
|
|
55
|
-
associatives = {'all', 'any', 'first', 'last', 'max', 'min', 'one', 'product', 'sum'}
|
|
56
|
-
associatives |= {'count'} # transformed to be associative
|
|
56
|
+
ordered = {'first', 'last'}
|
|
57
57
|
|
|
58
58
|
def __init__(self, name: str, alias: str = '', **options):
|
|
59
59
|
self.name = name
|
|
60
60
|
self.alias = alias or name
|
|
61
61
|
self.options = options
|
|
62
62
|
|
|
63
|
-
def
|
|
64
|
-
|
|
65
|
-
return self.name, func, options, self.alias
|
|
66
|
-
|
|
67
|
-
one = operator.itemgetter(0)
|
|
68
|
-
list = staticmethod(lambda array: array)
|
|
69
|
-
|
|
70
|
-
@staticmethod
|
|
71
|
-
def distinct(array: Array, mode: str = 'only_valid') -> pa.Array:
|
|
72
|
-
values = pc.unique(array)
|
|
73
|
-
if not values.null_count or mode == 'all':
|
|
74
|
-
return values
|
|
75
|
-
return values.drop_null() if mode == 'only_valid' else pa.array([None], array.type)
|
|
63
|
+
def func_options(self, func: str) -> pc.FunctionOptions:
|
|
64
|
+
return self.option_map[func.removeprefix('hash_')](**self.options)
|
|
76
65
|
|
|
77
66
|
|
|
78
|
-
@dataclass(frozen=True)
|
|
67
|
+
@dataclass(frozen=True, slots=True)
|
|
79
68
|
class Compare:
|
|
80
69
|
"""Comparable wrapper for bisection search."""
|
|
81
70
|
|
|
82
71
|
value: object
|
|
83
|
-
__slots__ = ('value',) # slots keyword in Python >=3.10
|
|
84
72
|
|
|
85
73
|
def __lt__(self, other):
|
|
86
74
|
return self.value < other.as_py()
|
|
@@ -163,7 +151,7 @@ class ListChunk(pa.lib.BaseListArray):
|
|
|
163
151
|
offsets = self.values.offsets.take(self.offsets)
|
|
164
152
|
return type(self).from_arrays(offsets, self.values.values)
|
|
165
153
|
|
|
166
|
-
def aggregate(self, **funcs:
|
|
154
|
+
def aggregate(self, **funcs: pc.FunctionOptions | None) -> pa.RecordBatch:
|
|
167
155
|
"""Return aggregated scalars by grouping each hash function on the parent indices.
|
|
168
156
|
|
|
169
157
|
If there are empty or null scalars, then the result must be padded with null defaults and
|
|
@@ -210,8 +198,7 @@ class ListChunk(pa.lib.BaseListArray):
|
|
|
210
198
|
|
|
211
199
|
def index(self, **options) -> pa.Array:
|
|
212
200
|
"""index for first occurrence of each list scalar"""
|
|
213
|
-
|
|
214
|
-
return Column.from_scalars(values)
|
|
201
|
+
return pa.array(pc.index(value, **options) for value in ListChunk.scalars(self))
|
|
215
202
|
|
|
216
203
|
@register
|
|
217
204
|
def list_all(ctx, self: pa.list_(pa.bool_())) -> pa.bool_(): # type: ignore
|
|
@@ -227,13 +214,6 @@ class ListChunk(pa.lib.BaseListArray):
|
|
|
227
214
|
class Column(pa.ChunkedArray):
|
|
228
215
|
"""Chunked array interface as a namespace of functions."""
|
|
229
216
|
|
|
230
|
-
def from_scalars(values: Sequence) -> pa.Array:
|
|
231
|
-
"""Return array from arrow scalars."""
|
|
232
|
-
return pa.array((value.as_py() for value in values), values[0].type)
|
|
233
|
-
|
|
234
|
-
def scalar_type(self):
|
|
235
|
-
return self.type.value_type if pa.types.is_dictionary(self.type) else self.type
|
|
236
|
-
|
|
237
217
|
def is_list_type(self):
|
|
238
218
|
funcs = pa.types.is_list, pa.types.is_large_list, pa.types.is_fixed_size_list
|
|
239
219
|
return any(func(self.type) for func in funcs)
|
|
@@ -260,7 +240,7 @@ class Column(pa.ChunkedArray):
|
|
|
260
240
|
if not pa.types.is_dictionary(self.type):
|
|
261
241
|
return self
|
|
262
242
|
array = self if isinstance(self, pa.Array) else self.combine_chunks()
|
|
263
|
-
return pc.rank(array.dictionary
|
|
243
|
+
return pc.rank(array.dictionary).take(array.indices)
|
|
264
244
|
|
|
265
245
|
def pairwise_diff(self, period: int = 1) -> Array:
|
|
266
246
|
"""`pairwise_diff` with chunked array support."""
|
|
@@ -284,26 +264,6 @@ class Column(pa.ChunkedArray):
|
|
|
284
264
|
mask = predicate(Column.diff(self), *args) if args else Column.diff(self, predicate)
|
|
285
265
|
return pc.indices_nonzero(pa.chunked_array(ends + mask.chunks + ends))
|
|
286
266
|
|
|
287
|
-
def first_last(self, **options):
|
|
288
|
-
if not pa.types.is_dictionary(self.type):
|
|
289
|
-
return pc.first_last(self, **options).as_py()
|
|
290
|
-
self = self.combine_chunks()
|
|
291
|
-
indices = pc.first_last(self.indices, **options).as_py().items()
|
|
292
|
-
return {key: None if index is None else self.dictionary[index] for key, index in indices}
|
|
293
|
-
|
|
294
|
-
def min_max(self, **options):
|
|
295
|
-
if pa.types.is_dictionary(self.type):
|
|
296
|
-
self = self.unique().dictionary_decode()
|
|
297
|
-
return pc.min_max(self, **options).as_py()
|
|
298
|
-
|
|
299
|
-
def min(self, **options):
|
|
300
|
-
"""Return min of the values."""
|
|
301
|
-
return Column.min_max(self, **options)['min']
|
|
302
|
-
|
|
303
|
-
def max(self, **options):
|
|
304
|
-
"""Return max of the values."""
|
|
305
|
-
return Column.min_max(self, **options)['max']
|
|
306
|
-
|
|
307
267
|
def index(self, value, start=0, end=None) -> int:
|
|
308
268
|
"""Return the first index of a value."""
|
|
309
269
|
with contextlib.suppress(NotImplementedError):
|
|
@@ -338,10 +298,8 @@ class Column(pa.ChunkedArray):
|
|
|
338
298
|
class Table(pa.Table):
|
|
339
299
|
"""Table interface as a namespace of functions."""
|
|
340
300
|
|
|
341
|
-
def map_batch(
|
|
342
|
-
|
|
343
|
-
batches = [func(batch, *rargs, **kwargs) for batch in scanner.to_batches() if batch]
|
|
344
|
-
return pa.Table.from_batches(batches, None if batches else scanner.projected_schema)
|
|
301
|
+
def map_batch(self, func: Callable, *args, **kwargs) -> pa.Table:
|
|
302
|
+
return pa.Table.from_batches(func(batch, *args, **kwargs) for batch in self.to_batches())
|
|
345
303
|
|
|
346
304
|
def columns(self) -> dict:
|
|
347
305
|
"""Return columns as a dictionary."""
|
|
@@ -411,39 +369,6 @@ class Table(pa.Table):
|
|
|
411
369
|
table = Table.union(scalars, Table.from_offsets(lists, offsets))
|
|
412
370
|
return table, Column.diff(offsets)
|
|
413
371
|
|
|
414
|
-
def group(self, *names: str, counts: str = '', **funcs: Sequence[Agg]) -> pa.Table:
|
|
415
|
-
"""Group by and aggregate.
|
|
416
|
-
|
|
417
|
-
Args:
|
|
418
|
-
*names: columns to group by
|
|
419
|
-
counts: alias for optional row counts
|
|
420
|
-
**funcs: aggregate funcs with columns options
|
|
421
|
-
"""
|
|
422
|
-
prefix = 'hash_' if names else ''
|
|
423
|
-
aggs = [agg.astuple(prefix + func) for func in funcs for agg in funcs[func]]
|
|
424
|
-
columns = list(names) + [agg[0] for agg in aggs]
|
|
425
|
-
if counts:
|
|
426
|
-
aggs.append(([], 'hash_count_all', None, counts))
|
|
427
|
-
if isinstance(self, pa.Table):
|
|
428
|
-
decl = Declaration('table_source', self.unify_dictionaries())
|
|
429
|
-
else:
|
|
430
|
-
decl = Declaration.scan(self, columns=columns)
|
|
431
|
-
return decl.aggregate(aggs, names).to_table(use_threads=False)
|
|
432
|
-
|
|
433
|
-
def aggregate(self, counts: str = '', **funcs: Sequence[Agg]) -> dict:
|
|
434
|
-
"""Return aggregated scalars as a row of data."""
|
|
435
|
-
row = {counts: len(self)} if counts else {}
|
|
436
|
-
for key in ('one', 'list', 'distinct'): # hash only functions
|
|
437
|
-
func, aggs = getattr(Agg, key), funcs.pop(key, [])
|
|
438
|
-
row |= {agg.alias: func(self[agg.name], **agg.options) for agg in aggs}
|
|
439
|
-
if funcs:
|
|
440
|
-
table = Table.group(self, **funcs) # type: ignore
|
|
441
|
-
row |= {name: table[name][0] for name in table.schema.names}
|
|
442
|
-
for name, value in row.items():
|
|
443
|
-
if isinstance(value, pa.ChunkedArray):
|
|
444
|
-
row[name] = value.combine_chunks()
|
|
445
|
-
return row
|
|
446
|
-
|
|
447
372
|
def list_fields(self) -> set:
|
|
448
373
|
return {field.name for field in self.schema if Column.is_list_type(field)}
|
|
449
374
|
|
|
@@ -465,7 +390,7 @@ class Table(pa.Table):
|
|
|
465
390
|
return Table.union(self, Table.from_counts(table, counts))
|
|
466
391
|
|
|
467
392
|
def sort_indices(
|
|
468
|
-
self, *names: str, length:
|
|
393
|
+
self, *names: str, length: int | None = None, null_placement: str = 'at_end'
|
|
469
394
|
) -> pa.Array:
|
|
470
395
|
"""Return indices which would sort the table by columns, optimized for fixed length."""
|
|
471
396
|
func = functools.partial(pc.sort_indices, null_placement=null_placement)
|
|
@@ -478,7 +403,7 @@ class Table(pa.Table):
|
|
|
478
403
|
def sort(
|
|
479
404
|
self,
|
|
480
405
|
*names: str,
|
|
481
|
-
length:
|
|
406
|
+
length: int | None = None,
|
|
482
407
|
indices: str = '',
|
|
483
408
|
null_placement: str = 'at_end',
|
|
484
409
|
) -> Batch:
|
|
@@ -490,7 +415,7 @@ class Table(pa.Table):
|
|
|
490
415
|
indices: include original indices in the table
|
|
491
416
|
"""
|
|
492
417
|
if length == 1 and not indices:
|
|
493
|
-
return Table.
|
|
418
|
+
return Table.min_max(self, *names)[:1]
|
|
494
419
|
indices_ = Table.sort_indices(self, *names, length=length, null_placement=null_placement)
|
|
495
420
|
table = self.take(indices_)
|
|
496
421
|
if indices:
|
|
@@ -510,49 +435,47 @@ class Table(pa.Table):
|
|
|
510
435
|
table = pa.concat_tables(table for table in tables if table is not None)
|
|
511
436
|
return Table.union(self, Table.from_counts(table, counts))
|
|
512
437
|
|
|
513
|
-
def
|
|
514
|
-
"""Return table filtered by
|
|
515
|
-
|
|
516
|
-
|
|
517
|
-
|
|
518
|
-
|
|
519
|
-
value = func(self[name], skip_nulls=null_placement == 'at_end')
|
|
520
|
-
if value is None:
|
|
521
|
-
mask = (pc.field(name) if expr else self[name]).is_null()
|
|
522
|
-
else:
|
|
523
|
-
mask = pc.field(name) == value if expr else pc.equal(self[name], value)
|
|
524
|
-
self = self.filter(mask)
|
|
525
|
-
return self
|
|
526
|
-
values = None
|
|
527
|
-
kwargs = dict(tiebreaker='dense', null_placement=null_placement)
|
|
528
|
-
for name, order in map(sort_key, names):
|
|
529
|
-
if len(self) <= k:
|
|
530
|
-
break
|
|
531
|
-
ranks = pc.rank(Column.sort_values(self[name]), order, **kwargs)
|
|
532
|
-
if values is None: # optimized to sort only once on first iteration
|
|
533
|
-
values = ranks
|
|
534
|
-
else:
|
|
535
|
-
values = pc.add_checked(pc.multiply_checked(values, pc.max(ranks)), ranks)
|
|
536
|
-
values = pc.rank(values, 'ascending', tiebreaker='dense')
|
|
537
|
-
mask = pc.less_equal(values, k)
|
|
538
|
-
self = self.filter(mask)
|
|
539
|
-
values = values.filter(mask)
|
|
438
|
+
def min_max(self, *names: str) -> Self:
|
|
439
|
+
"""Return table filtered by minimum or maximum values."""
|
|
440
|
+
for key, order in map(sort_key, names):
|
|
441
|
+
field, asc = pc.field(key), (order == 'ascending')
|
|
442
|
+
((value,),) = Nodes.group(self, _=(key, ('min' if asc else 'max'), None)).to_table()
|
|
443
|
+
self = self.filter(field <= value if asc else field >= value)
|
|
540
444
|
return self
|
|
541
445
|
|
|
542
|
-
def
|
|
543
|
-
"""
|
|
544
|
-
|
|
545
|
-
|
|
546
|
-
|
|
547
|
-
|
|
548
|
-
|
|
549
|
-
|
|
550
|
-
|
|
446
|
+
def rank(self, k: int, *names: str) -> Self:
|
|
447
|
+
"""Return table filtered by values within dense rank, similar to `select_k_unstable`."""
|
|
448
|
+
if k == 1:
|
|
449
|
+
return Table.min_max(self, *names)
|
|
450
|
+
keys = dict(map(sort_key, names))
|
|
451
|
+
table = Nodes.group(self, *keys).to_table()
|
|
452
|
+
table = table.take(pc.select_k_unstable(table, k, keys.items()))
|
|
453
|
+
exprs = []
|
|
454
|
+
for key, order in keys.items():
|
|
455
|
+
field, asc = pc.field(key), (order == 'ascending')
|
|
456
|
+
exprs.append(field <= pc.max(table[key]) if asc else field >= pc.min(table[key]))
|
|
457
|
+
return self.filter(bit_all(exprs))
|
|
458
|
+
|
|
459
|
+
def fragments(self, *names, counts: str = '') -> pa.Table:
|
|
460
|
+
"""Return selected fragment keys in a table."""
|
|
551
461
|
try:
|
|
552
|
-
|
|
462
|
+
expr = self._scan_options.get('filter')
|
|
463
|
+
if expr is not None: # raise ValueError if filter references other fields
|
|
464
|
+
ds.dataset([], schema=self.partitioning.schema).scanner(filter=expr)
|
|
553
465
|
except (AttributeError, ValueError):
|
|
554
|
-
return
|
|
555
|
-
|
|
466
|
+
return pa.table({})
|
|
467
|
+
fragments = self._get_fragments(expr)
|
|
468
|
+
parts = [ds.get_partition_keys(frag.partition_expression) for frag in fragments]
|
|
469
|
+
names, table = set(names), pa.Table.from_pylist(parts) # type: ignore
|
|
470
|
+
keys = [name for name in table.schema.names if name in names]
|
|
471
|
+
table = table.group_by(keys, use_threads=False).aggregate([])
|
|
472
|
+
if not counts:
|
|
473
|
+
return table
|
|
474
|
+
if not table.schema:
|
|
475
|
+
return table.append_column(counts, pa.array([self.count_rows()]))
|
|
476
|
+
exprs = [bit_all(pc.field(key) == row[key] for key in row) for row in table.to_pylist()]
|
|
477
|
+
column = [self.filter(expr).count_rows() for expr in exprs]
|
|
478
|
+
return table.append_column(counts, pa.array(column))
|
|
556
479
|
|
|
557
480
|
def rank_keys(self, k: int, *names: str, dense: bool = True) -> tuple:
|
|
558
481
|
"""Return expression and unmatched fields for partitioned dataset which filters by rank.
|
|
@@ -562,39 +485,24 @@ class Table(pa.Table):
|
|
|
562
485
|
*names: columns to rank by
|
|
563
486
|
dense: use dense rank; false indicates sorting
|
|
564
487
|
"""
|
|
565
|
-
|
|
566
|
-
|
|
488
|
+
keys = dict(map(sort_key, names))
|
|
489
|
+
table = Table.fragments(self, *keys, counts='' if dense else '_')
|
|
490
|
+
keys = {name: keys[name] for name in table.schema.names if name in keys}
|
|
567
491
|
if not keys:
|
|
568
492
|
return None, names
|
|
569
|
-
parts = []
|
|
570
|
-
for fragment in Table.get_fragments(self):
|
|
571
|
-
parts.append(ds.get_partition_keys(fragment.partition_expression))
|
|
572
|
-
if not dense:
|
|
573
|
-
parts[-1][''] = fragment.count_rows()
|
|
574
|
-
offset, table = len(keys), pa.Table.from_pylist(parts)
|
|
575
493
|
if dense:
|
|
576
|
-
table =
|
|
494
|
+
table = table.take(pc.select_k_unstable(table, k, keys.items()))
|
|
577
495
|
else:
|
|
578
|
-
table =
|
|
579
|
-
|
|
580
|
-
|
|
581
|
-
|
|
496
|
+
table = table.sort_by(keys.items())
|
|
497
|
+
totals = itertools.accumulate(table['_'].to_pylist())
|
|
498
|
+
counts = (count for count, total in enumerate(totals, 1) if total >= k)
|
|
499
|
+
table = table[: next(counts, None)].remove_column(len(table) - 1)
|
|
500
|
+
exprs = [bit_all(pc.field(key) == row[key] for key in row) for row in table.to_pylist()]
|
|
501
|
+
remaining = names[len(keys) :]
|
|
582
502
|
if remaining or not dense: # fields with a single value are no longer needed
|
|
583
|
-
selectors = [len(table[
|
|
503
|
+
selectors = [len(table[key].unique()) > 1 for key in keys]
|
|
584
504
|
remaining = tuple(itertools.compress(names, selectors)) + remaining
|
|
585
|
-
|
|
586
|
-
for index, (name, order) in enumerate(keys.items()):
|
|
587
|
-
expr, field, values = None, ds.field(name), table[name].unique()
|
|
588
|
-
value = pc.max(values) if order == 'ascending' else pc.min(values)
|
|
589
|
-
if index == len(keys) - 1: # last one
|
|
590
|
-
expr = field <= value if order == 'ascending' else field >= value
|
|
591
|
-
elif len(values) > 1:
|
|
592
|
-
expr = field < value if order == 'ascending' else field > value
|
|
593
|
-
if expr is not None:
|
|
594
|
-
fields.append(functools.reduce(operator.and_, prefix + [expr]))
|
|
595
|
-
prefix.append(field == value)
|
|
596
|
-
table = table.filter(prefix[-1])
|
|
597
|
-
return functools.reduce(operator.or_, fields), remaining
|
|
505
|
+
return bit_any(exprs[: len(table)]), remaining
|
|
598
506
|
|
|
599
507
|
def flatten(self, indices: str = '') -> Iterator[pa.RecordBatch]:
|
|
600
508
|
"""Generate batches with list arrays flattened, optionally with parent indices."""
|
|
@@ -612,7 +520,7 @@ class Table(pa.Table):
|
|
|
612
520
|
offset += len(batch)
|
|
613
521
|
yield pa.RecordBatch.from_pydict(columns)
|
|
614
522
|
|
|
615
|
-
def split(self) -> Iterator[
|
|
523
|
+
def split(self) -> Iterator[pa.RecordBatch | None]:
|
|
616
524
|
"""Generate tables from splitting list scalars."""
|
|
617
525
|
lists = Table.list_fields(self)
|
|
618
526
|
scalars = set(self.schema.names) - lists
|
|
@@ -632,8 +540,11 @@ class Table(pa.Table):
|
|
|
632
540
|
return f'{size:n} {prefix}B'
|
|
633
541
|
|
|
634
542
|
|
|
635
|
-
class
|
|
636
|
-
"""[Acero](https://arrow.apache.org/docs/python/api/acero.html) engine declaration.
|
|
543
|
+
class Nodes(ac.Declaration):
|
|
544
|
+
"""[Acero](https://arrow.apache.org/docs/python/api/acero.html) engine declaration.
|
|
545
|
+
|
|
546
|
+
Provides a `Scanner` interface with no "oneshot" limitation.
|
|
547
|
+
"""
|
|
637
548
|
|
|
638
549
|
option_map = {
|
|
639
550
|
'table_source': ac.TableSourceNodeOptions,
|
|
@@ -644,25 +555,65 @@ class Declaration(ac.Declaration):
|
|
|
644
555
|
'order_by': ac.OrderByNodeOptions,
|
|
645
556
|
'hashjoin': ac.HashJoinNodeOptions,
|
|
646
557
|
}
|
|
558
|
+
to_batches = ac.Declaration.to_reader # source compatibility
|
|
647
559
|
|
|
648
560
|
def __init__(self, name, *args, inputs=None, **options):
|
|
649
561
|
super().__init__(name, self.option_map[name](*args, **options), inputs)
|
|
650
562
|
|
|
651
|
-
|
|
652
|
-
|
|
653
|
-
|
|
654
|
-
|
|
655
|
-
|
|
656
|
-
|
|
657
|
-
|
|
563
|
+
def scan(self, columns: Iterable[str]) -> Self:
|
|
564
|
+
"""Return projected source node, supporting datasets and tables."""
|
|
565
|
+
if isinstance(self, ds.Dataset):
|
|
566
|
+
expr = self._scan_options.get('filter')
|
|
567
|
+
self = Nodes('scan', self, columns=columns)
|
|
568
|
+
if expr is not None:
|
|
569
|
+
self = self.apply('filter', expr)
|
|
570
|
+
elif isinstance(self, pa.Table):
|
|
571
|
+
self = Nodes('table_source', self)
|
|
572
|
+
elif isinstance(self, pa.RecordBatch):
|
|
573
|
+
self = Nodes('table_source', pa.table(self))
|
|
658
574
|
if isinstance(columns, Mapping):
|
|
659
|
-
|
|
660
|
-
|
|
661
|
-
|
|
662
|
-
|
|
575
|
+
return self.apply('project', columns.values(), columns)
|
|
576
|
+
return self.apply('project', map(pc.field, columns))
|
|
577
|
+
|
|
578
|
+
@property
|
|
579
|
+
def schema(self) -> pa.Schema:
|
|
580
|
+
"""projected schema"""
|
|
581
|
+
with self.to_reader() as reader:
|
|
582
|
+
return reader.schema
|
|
583
|
+
|
|
584
|
+
def scanner(self, **options) -> ds.Scanner:
|
|
585
|
+
return ds.Scanner.from_batches(self.to_reader(**options))
|
|
586
|
+
|
|
587
|
+
def count_rows(self) -> int:
|
|
588
|
+
"""Count matching rows."""
|
|
589
|
+
return self.scanner().count_rows()
|
|
590
|
+
|
|
591
|
+
def head(self, num_rows: int, **options) -> pa.Table:
|
|
592
|
+
"""Load the first N rows."""
|
|
593
|
+
return self.scanner(**options).head(num_rows)
|
|
594
|
+
|
|
595
|
+
def take(self, indices: Iterable[int], **options) -> pa.Table:
|
|
596
|
+
"""Select rows by index."""
|
|
597
|
+
return self.scanner(**options).take(indices)
|
|
663
598
|
|
|
664
599
|
def apply(self, name: str, *args, **options) -> Self:
|
|
600
|
+
"""Add a node by name."""
|
|
665
601
|
return type(self)(name, *args, inputs=[self], **options)
|
|
666
602
|
|
|
667
|
-
|
|
668
|
-
|
|
603
|
+
filter = functools.partialmethod(apply, 'filter')
|
|
604
|
+
|
|
605
|
+
def group(self, *names, **aggs: tuple) -> Self:
|
|
606
|
+
"""Add `aggregate` node with dictionary support.
|
|
607
|
+
|
|
608
|
+
Also supports datasets because aggregation determines the projection.
|
|
609
|
+
"""
|
|
610
|
+
aggregates, targets = [], set(names)
|
|
611
|
+
for name, (target, _, _) in aggs.items():
|
|
612
|
+
aggregates.append(aggs[name] + (name,))
|
|
613
|
+
targets.update([target] if isinstance(target, str) else target)
|
|
614
|
+
columns = {name: pc.field(name) for name in targets}
|
|
615
|
+
for name in columns:
|
|
616
|
+
field = self.schema.field(name)
|
|
617
|
+
if pa.types.is_dictionary(field.type):
|
|
618
|
+
columns[name] = columns[name].cast(field.type.value_type)
|
|
619
|
+
return Nodes.scan(self, columns).apply('aggregate', aggregates, names)
|