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.
Files changed (26) hide show
  1. {graphique-1.7 → graphique-1.9}/PKG-INFO +14 -22
  2. {graphique-1.7 → graphique-1.9}/README.md +5 -2
  3. {graphique-1.7 → graphique-1.9}/graphique/core.py +126 -175
  4. {graphique-1.7 → graphique-1.9}/graphique/inputs.py +166 -165
  5. {graphique-1.7 → graphique-1.9}/graphique/interface.py +119 -168
  6. {graphique-1.7 → graphique-1.9}/graphique/middleware.py +11 -26
  7. {graphique-1.7 → graphique-1.9}/graphique/models.py +37 -48
  8. {graphique-1.7 → graphique-1.9}/graphique/scalars.py +19 -38
  9. {graphique-1.7 → graphique-1.9}/graphique/service.py +4 -0
  10. {graphique-1.7 → graphique-1.9}/graphique/shell.py +9 -6
  11. {graphique-1.7 → graphique-1.9}/graphique.egg-info/PKG-INFO +14 -22
  12. {graphique-1.7 → graphique-1.9}/graphique.egg-info/requires.txt +3 -1
  13. {graphique-1.7 → graphique-1.9}/pyproject.toml +5 -7
  14. {graphique-1.7 → graphique-1.9}/requirements.in +3 -1
  15. {graphique-1.7 → graphique-1.9}/tests/test_bench.py +4 -4
  16. {graphique-1.7 → graphique-1.9}/tests/test_core.py +25 -88
  17. {graphique-1.7 → graphique-1.9}/tests/test_dataset.py +13 -15
  18. {graphique-1.7 → graphique-1.9}/tests/test_models.py +13 -11
  19. {graphique-1.7 → graphique-1.9}/tests/test_service.py +1 -1
  20. {graphique-1.7 → graphique-1.9}/LICENSE.txt +0 -0
  21. {graphique-1.7 → graphique-1.9}/graphique/__init__.py +0 -0
  22. {graphique-1.7 → graphique-1.9}/graphique/py.typed +0 -0
  23. {graphique-1.7 → graphique-1.9}/graphique.egg-info/SOURCES.txt +0 -0
  24. {graphique-1.7 → graphique-1.9}/graphique.egg-info/dependency_links.txt +0 -0
  25. {graphique-1.7 → graphique-1.9}/graphique.egg-info/top_level.txt +0 -0
  26. {graphique-1.7 → graphique-1.9}/setup.cfg +0 -0
@@ -1,22 +1,9 @@
1
- Metadata-Version: 2.1
1
+ Metadata-Version: 2.4
2
2
  Name: graphique
3
- Version: 1.7
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: Copyright 2022 Aric Coady
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.9
26
+ Requires-Python: >=3.10
41
27
  Description-Content-Type: text/markdown
42
28
  License-File: LICENSE.txt
43
- Requires-Dist: pyarrow>=17
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
  [![image](https://img.shields.io/pypi/v/graphique.svg)](https://pypi.org/project/graphique/)
51
40
  ![image](https://img.shields.io/pypi/pyversions/graphique.svg)
@@ -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
- app = GraphQL(ds.dataset(...)) # Table is root query type
93
- app = GraphQL.federated({<name>: ds.dataset(...), ...}, keys={...}) # Tables on federated fields
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
- app = GraphQL(ds.dataset(...)) # Table is root query type
44
- app = GraphQL.federated({<name>: ds.dataset(...), ...}, keys={...}) # Tables on federated fields
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, Sequence
15
+ from collections.abc import Callable, Iterable, Iterator, Mapping
16
16
  from dataclasses import dataclass
17
- from typing import Optional, Union, get_type_hints
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 = Union[pa.Array, pa.ChunkedArray]
26
- Batch = Union[pa.RecordBatch, pa.Table]
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 astuple(self, func: str) -> tuple:
64
- options = self.option_map[func.rpartition('hash_')[-1]](**self.options)
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: Optional[pc.FunctionOptions]) -> pa.RecordBatch:
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
- values = [pc.index(value, **options) for value in ListChunk.scalars(self)]
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, 'ascending').take(array.indices)
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(scanner: ds.Scanner, func: Callable, *rargs, **kwargs) -> pa.Table:
342
- # TODO(apache/arrow#31612): replace with user defined function for multiple kernels
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: Optional[int] = None, null_placement: str = 'at_end'
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: Optional[int] = None,
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.ranked(self, 1, *names, null_placement=null_placement)[:1]
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 ranked(self, k: int, *names: str, null_placement: str = 'at_end') -> Batch:
514
- """Return table filtered by values within dense rank, similar to `select_k_unstable`."""
515
- if k == 1: # optimized for min_max
516
- expr = isinstance(self, pa.Table)
517
- for name, order in map(sort_key, names):
518
- func = Column.min if order == 'ascending' else Column.max
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 get_fragments(self) -> Iterator[ds.Fragment]:
543
- """Support filtered datasets if it only references partition keys."""
544
- expr = self._scan_options.get('filter')
545
- if expr is not None: # raise ValueError if filter references other fields
546
- ds.dataset([], schema=self.partitioning.schema).scanner(filter=expr)
547
- return self._get_fragments(expr)
548
-
549
- def fragment_keys(self) -> list:
550
- """Filtered partitioned datasets may not have fragments."""
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
- Table.get_fragments(self)
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
- return self.partitioning.schema.names
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
- schema = set(Table.fragment_keys(self))
566
- keys = dict(itertools.takewhile(lambda key: key[0] in schema, map(sort_key, names)))
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 = Table.ranked(table, k, *names[:offset])
494
+ table = table.take(pc.select_k_unstable(table, k, keys.items()))
577
495
  else:
578
- table = Table.sort(table, *names[:offset])
579
- index = pc.index(pc.greater_equal(pc.cumulative_sum(table['']), k), True).as_py()
580
- table = table[: index + 1 if index >= 0 else None]
581
- remaining = names[offset:]
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[name].unique()) > 1 for name in keys]
503
+ selectors = [len(table[key].unique()) > 1 for key in keys]
584
504
  remaining = tuple(itertools.compress(names, selectors)) + remaining
585
- fields, prefix = [], [] # type: ignore
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[Optional[pa.RecordBatch]]:
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 Declaration(ac.Declaration):
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
- @classmethod
652
- def scan(cls, dataset: ds.Dataset, columns=None) -> Self:
653
- """Return source node from a dataset."""
654
- expr = dataset._scan_options.get('filter')
655
- self = cls('scan', dataset, filter=expr, columns=columns)
656
- if expr is not None:
657
- self = self.filter(expr)
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
- self = self.project(columns.values(), columns)
660
- elif columns is not None:
661
- self = self.project(map(pc.field, columns))
662
- return self
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
- def __getattr__(self, name: str) -> Callable:
668
- return functools.partial(self.apply, name)
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)