graphique 1.0__tar.gz → 1.2__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.
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.1
2
2
  Name: graphique
3
- Version: 1.0
3
+ Version: 1.2
4
4
  Summary: GraphQL service for arrow tables and parquet data sets.
5
5
  Author-email: Aric Coady <aric.coady@gmail.com>
6
6
  License: Copyright 2022 Aric Coady
@@ -26,16 +26,16 @@ Classifier: Intended Audience :: Science/Research
26
26
  Classifier: License :: OSI Approved :: Apache Software License
27
27
  Classifier: Operating System :: OS Independent
28
28
  Classifier: Programming Language :: Python :: 3
29
- Classifier: Programming Language :: Python :: 3.7
30
29
  Classifier: Programming Language :: Python :: 3.8
31
30
  Classifier: Programming Language :: Python :: 3.9
32
31
  Classifier: Programming Language :: Python :: 3.10
32
+ Classifier: Programming Language :: Python :: 3.11
33
33
  Classifier: Topic :: Database :: Database Engines/Servers
34
34
  Classifier: Topic :: Internet :: WWW/HTTP :: HTTP Servers
35
35
  Classifier: Topic :: Internet :: WWW/HTTP :: Indexing/Search
36
36
  Classifier: Topic :: Software Development :: Libraries :: Python Modules
37
37
  Classifier: Typing :: Typed
38
- Requires-Python: >=3.7
38
+ Requires-Python: >=3.8
39
39
  Description-Content-Type: text/markdown
40
40
  Provides-Extra: server
41
41
  License-File: LICENSE.txt
@@ -57,7 +57,7 @@ License-File: LICENSE.txt
57
57
  % env PARQUET_PATH=... uvicorn graphique.service:app
58
58
  ```
59
59
 
60
- Open http://localhost:8000/graphql to try out the API in [GraphiQL](https://github.com/graphql/graphiql/tree/main/packages/graphiql#readme). There is a test fixture at `./tests/fixtures/zipcodes.parquet`.
60
+ Open http://localhost:8000/ to try out the API in [GraphiQL](https://github.com/graphql/graphiql/tree/main/packages/graphiql#readme). There is a test fixture at `./tests/fixtures/zipcodes.parquet`.
61
61
 
62
62
  ```console
63
63
  % env PARQUET_PATH=... strawberry export-schema graphique.service:app.schema
@@ -88,16 +88,16 @@ app = GraphQL.federated({<name>: ds.dataset(...), ...}, keys={...}) # Tables on
88
88
  Start like any ASGI app.
89
89
 
90
90
  ```console
91
- uvicorn <package>:app
91
+ uvicorn <module>:app
92
92
  ```
93
93
 
94
94
  Configuration options exist to provide a convenient no-code solution, but are subject to change in the future. Using a custom app is recommended for production usage.
95
95
 
96
96
  ### API
97
97
  #### types
98
- * `Dataset`: an interface for an arrow dataset or table.
99
- * `Table`: implements the `Dataset` interface with a known schema.
100
- * `Column`: an arrow Column (a.k.a. ChunkedArray). Each arrow data type has a corresponding column implementation: Boolean, Int, Long, Float, Decimal, Date, DateTime, Time, Duration, Base64, String, List, Struct. All columns have a `values` field for their list of scalars. Additional fields vary by type.
98
+ * `Dataset`: interface for an arrow dataset, scanner, or table.
99
+ * `Table`: implements the `Dataset` interface. Adds typed `row`, `columns`, and `filter` fields from introspecting the schema.
100
+ * `Column`: interface for an arrow column (a.k.a. ChunkedArray). Each arrow data type has a corresponding column implementation: Boolean, Int, Long, Float, Decimal, Date, Datetime, Time, Duration, Base64, String, List, Struct. All columns have a `values` field for their list of scalars. Additional fields vary by type.
101
101
  * `Row`: scalar fields. Arrow tables are column-oriented, and graphique encourages that usage for performance. A single `row` field is provided for convenience, but a field for a list of rows is not. Requesting parallel columns is far more efficient.
102
102
 
103
103
  #### selection
@@ -138,8 +138,8 @@ If index columns are detected in the schema metadata, then an initial `filter` w
138
138
  ```
139
139
 
140
140
  ## Dependencies
141
- * pyarrow >=10
142
- * strawberry-graphql[asgi] >=0.134.4
141
+ * pyarrow >=12
142
+ * strawberry-graphql[asgi,cli]
143
143
  * uvicorn (or other [ASGI server](https://asgi.readthedocs.io/en/latest/implementations.html))
144
144
 
145
145
  ## Tests
@@ -150,6 +150,23 @@ If index columns are detected in the schema metadata, then an initial `filter` w
150
150
  ```
151
151
 
152
152
  ## Changes
153
+ 1.2
154
+
155
+ * Pyarrow >=12 required
156
+ * Grouping fragments optimized
157
+ * Group by empty columns
158
+ * Batch sorting and grouping into lists
159
+
160
+ 1.1
161
+
162
+ * Pyarrow >=11 required
163
+ * Python >=3.8 required
164
+ * Scannable functions added
165
+ * List aggregations deprecated
166
+ * Group by fragments
167
+ * Month day nano interval array
168
+ * `min` and `max` fields memory optimized
169
+
153
170
  1.0
154
171
 
155
172
  * Pyarrow >=10 required
@@ -216,9 +233,9 @@ If index columns are detected in the schema metadata, then an initial `filter` w
216
233
 
217
234
  0.2
218
235
 
236
+ * Pyarrow >= 2 required
219
237
  * `ListColumn` and `StructColumn` types
220
238
  * `Groups` type with `aggregate` field
221
239
  * `group` and `unique` optimized
222
- * pyarrow >= 2 required
223
240
  * Statistical fields: `mode`, `stddev`, `variance`
224
241
  * `is_in`, `min`, and `max` optimized
@@ -15,7 +15,7 @@
15
15
  % env PARQUET_PATH=... uvicorn graphique.service:app
16
16
  ```
17
17
 
18
- Open http://localhost:8000/graphql to try out the API in [GraphiQL](https://github.com/graphql/graphiql/tree/main/packages/graphiql#readme). There is a test fixture at `./tests/fixtures/zipcodes.parquet`.
18
+ Open http://localhost:8000/ to try out the API in [GraphiQL](https://github.com/graphql/graphiql/tree/main/packages/graphiql#readme). There is a test fixture at `./tests/fixtures/zipcodes.parquet`.
19
19
 
20
20
  ```console
21
21
  % env PARQUET_PATH=... strawberry export-schema graphique.service:app.schema
@@ -46,16 +46,16 @@ app = GraphQL.federated({<name>: ds.dataset(...), ...}, keys={...}) # Tables on
46
46
  Start like any ASGI app.
47
47
 
48
48
  ```console
49
- uvicorn <package>:app
49
+ uvicorn <module>:app
50
50
  ```
51
51
 
52
52
  Configuration options exist to provide a convenient no-code solution, but are subject to change in the future. Using a custom app is recommended for production usage.
53
53
 
54
54
  ### API
55
55
  #### types
56
- * `Dataset`: an interface for an arrow dataset or table.
57
- * `Table`: implements the `Dataset` interface with a known schema.
58
- * `Column`: an arrow Column (a.k.a. ChunkedArray). Each arrow data type has a corresponding column implementation: Boolean, Int, Long, Float, Decimal, Date, DateTime, Time, Duration, Base64, String, List, Struct. All columns have a `values` field for their list of scalars. Additional fields vary by type.
56
+ * `Dataset`: interface for an arrow dataset, scanner, or table.
57
+ * `Table`: implements the `Dataset` interface. Adds typed `row`, `columns`, and `filter` fields from introspecting the schema.
58
+ * `Column`: interface for an arrow column (a.k.a. ChunkedArray). Each arrow data type has a corresponding column implementation: Boolean, Int, Long, Float, Decimal, Date, Datetime, Time, Duration, Base64, String, List, Struct. All columns have a `values` field for their list of scalars. Additional fields vary by type.
59
59
  * `Row`: scalar fields. Arrow tables are column-oriented, and graphique encourages that usage for performance. A single `row` field is provided for convenience, but a field for a list of rows is not. Requesting parallel columns is far more efficient.
60
60
 
61
61
  #### selection
@@ -96,8 +96,8 @@ If index columns are detected in the schema metadata, then an initial `filter` w
96
96
  ```
97
97
 
98
98
  ## Dependencies
99
- * pyarrow >=10
100
- * strawberry-graphql[asgi] >=0.134.4
99
+ * pyarrow >=12
100
+ * strawberry-graphql[asgi,cli]
101
101
  * uvicorn (or other [ASGI server](https://asgi.readthedocs.io/en/latest/implementations.html))
102
102
 
103
103
  ## Tests
@@ -108,6 +108,23 @@ If index columns are detected in the schema metadata, then an initial `filter` w
108
108
  ```
109
109
 
110
110
  ## Changes
111
+ 1.2
112
+
113
+ * Pyarrow >=12 required
114
+ * Grouping fragments optimized
115
+ * Group by empty columns
116
+ * Batch sorting and grouping into lists
117
+
118
+ 1.1
119
+
120
+ * Pyarrow >=11 required
121
+ * Python >=3.8 required
122
+ * Scannable functions added
123
+ * List aggregations deprecated
124
+ * Group by fragments
125
+ * Month day nano interval array
126
+ * `min` and `max` fields memory optimized
127
+
111
128
  1.0
112
129
 
113
130
  * Pyarrow >=10 required
@@ -174,9 +191,9 @@ If index columns are detected in the schema metadata, then an initial `filter` w
174
191
 
175
192
  0.2
176
193
 
194
+ * Pyarrow >= 2 required
177
195
  * `ListColumn` and `StructColumn` types
178
196
  * `Groups` type with `aggregate` field
179
197
  * `group` and `unique` optimized
180
- * pyarrow >= 2 required
181
198
  * Statistical fields: `mode`, `stddev`, `variance`
182
199
  * `is_in`, `min`, and `max` optimized
@@ -1,3 +1,3 @@
1
1
  from .middleware import GraphQL # noqa: F401
2
2
 
3
- __version__ = '1.0'
3
+ __version__ = '1.2'
@@ -7,17 +7,18 @@ Their methods are called as functions.
7
7
  import bisect
8
8
  import contextlib
9
9
  import functools
10
+ import inspect
10
11
  import itertools
12
+ import operator
11
13
  import json
12
- from concurrent import futures
13
14
  from dataclasses import dataclass
14
- from typing import Callable, Iterable, Iterator, Optional, Sequence, Union
15
+ from typing import Callable, Iterable, Iterator, Optional, Sequence, Union, get_type_hints
15
16
  import numpy as np # type: ignore
16
17
  import pyarrow as pa
17
18
  import pyarrow.compute as pc
18
19
  import pyarrow.dataset as ds
19
20
 
20
- threader = futures.ThreadPoolExecutor(pa.cpu_count())
21
+ Array = Union[pa.Array, pa.ChunkedArray]
21
22
 
22
23
 
23
24
  class Agg:
@@ -51,7 +52,19 @@ class Agg:
51
52
  self.options = options
52
53
 
53
54
  def astuple(self, func: str) -> tuple:
54
- return f'hash_{func}', self.option_map[func](**self.options)
55
+ options = self.option_map[func](**self.options)
56
+ return self.name, func, options
57
+
58
+ @classmethod
59
+ def getfunc(cls, name: str) -> Callable:
60
+ """Return callable with named parameters for keyword arguments."""
61
+ if name in ('first', 'last', 'one'):
62
+ return operator.itemgetter(-1 if name == 'last' else 0)
63
+ if name == 'element':
64
+ return lambda array, index: array[index]
65
+ if name == 'slice':
66
+ return lambda array, start, stop, step: array[start:stop:step]
67
+ return getattr(pc, name)
55
68
 
56
69
 
57
70
  @dataclass(frozen=True)
@@ -78,15 +91,34 @@ def decode(array: pa.Array) -> pa.Array:
78
91
  return array.dictionary_decode() if isinstance(array, pa.DictionaryArray) else array
79
92
 
80
93
 
81
- class ListChunk(pa.lib.BaseListArray):
82
- value_length = pc.list_value_length
94
+ def register(func: Callable) -> pc.Function:
95
+ """Register user defined scalar function."""
96
+ doc = inspect.getdoc(func)
97
+ doc = {'summary': doc.splitlines()[0], 'description': doc} # type: ignore
98
+ annotations = dict(get_type_hints(func))
99
+ pc.register_scalar_function(func, func.__name__, doc, annotations, annotations.pop('return'))
100
+ return pc.get_function(func.__name__)
101
+
83
102
 
103
+ @register
104
+ def digitize(
105
+ ctx, array: pa.float64(), bins: pa.list_(pa.float64()), right: pa.bool_() # type: ignore
106
+ ) -> pa.int64(): # type: ignore
107
+ """Return the indices of the bins to which each value in input array belongs."""
108
+ return pa.array(np.digitize(array, bins.values, right.as_py()))
109
+
110
+
111
+ class ListChunk(pa.lib.BaseListArray):
84
112
  def from_counts(counts: pa.IntegerArray, values: pa.Array) -> pa.LargeListArray:
85
113
  """Return list array by converting counts into offsets."""
86
114
  offsets = pa.concat_arrays([pa.array([0], counts.type), pc.cumulative_sum_checked(counts)])
87
115
  cls = pa.LargeListArray if offsets.type == 'int64' else pa.ListArray
88
116
  return cls.from_arrays(offsets, values)
89
117
 
118
+ def from_scalars(values: Iterable) -> pa.LargeListArray:
119
+ """Return list array from array scalars."""
120
+ return ListChunk.from_counts(pa.array(map(len, values)), pa.concat_arrays(values))
121
+
90
122
  def element(self, index: int) -> pa.Array:
91
123
  """element at index of each list scalar; defaults to null"""
92
124
  with contextlib.suppress(ValueError):
@@ -113,31 +145,41 @@ class ListChunk(pa.lib.BaseListArray):
113
145
  def map_list(self, func: Callable, **kwargs) -> pa.lib.BaseListArray:
114
146
  """Return list array by mapping function across scalars, with null handling."""
115
147
  values = [func(value, **kwargs) for value in ListChunk.scalars(self)]
116
- return ListChunk.from_counts(pa.array(map(len, values)), pa.concat_arrays(values))
148
+ return ListChunk.from_scalars(values)
117
149
 
118
- def aggregate(self, **funcs: Optional[pc.FunctionOptions]) -> pa.StructArray:
150
+ def inner_flatten(self) -> pa.lib.BaseListArray:
151
+ """Return flattened inner lists from a nested list array."""
152
+ offsets = self.values.offsets.take(self.offsets)
153
+ return type(self).from_arrays(offsets, self.values.values)
154
+
155
+ def aggregate(self, **funcs: Optional[pc.FunctionOptions]) -> pa.RecordBatch:
119
156
  """Return aggregated scalars by grouping each hash function on the parent indices.
120
157
 
121
158
  If there are empty or null scalars, then the result must be padded with null defaults and
122
159
  reordered. If the function is a `count`, then the default is 0.
123
160
  """
124
- items = {f'hash_{name}': funcs[name] for name in funcs}.items()
125
- indices = pc.list_parent_indices(self)
126
- groups = pc._group_by([pc.list_flatten(self)] * len(funcs), [indices], items)
127
- if len(groups) == len(self): # no empty or null scalars
128
- return groups
161
+ columns = {'key': pc.list_parent_indices(self), '': pc.list_flatten(self)}
162
+ items = [('', name, funcs[name]) for name in funcs]
163
+ table = pa.table(columns).group_by(['key']).aggregate(items)
164
+ indices, table = table['key'], table.remove_column(table.schema.get_field_index('key'))
165
+ (batch,) = table.to_batches()
166
+ if len(batch) == len(self): # no empty or null scalars
167
+ return batch
129
168
  mask = pc.equal(pc.list_value_length(self), 0)
130
169
  empties = pc.indices_nonzero(Column.fill_null(mask, True))
131
- indices = pa.concat_arrays([groups.field(-1).cast('uint64'), empties])
132
- counters = [field.name for field in groups.type if 'count' in field.name]
133
- empties = pa.repeat(pa.scalar(dict.fromkeys(counters, 0), groups.type), len(empties))
134
- return pa.concat_arrays([groups, empties]).take(pc.sort_indices(indices))
170
+ indices = pa.chunked_array(indices.chunks + [empties.cast(indices.type)])
171
+ columns = {}
172
+ for field in batch.schema:
173
+ scalar = pa.scalar(0 if 'count' in field.name else None, field.type)
174
+ columns[field.name] = pa.repeat(scalar, len(empties))
175
+ table = pa.concat_tables([table, pa.table(columns)]).combine_chunks()
176
+ return table.to_batches()[0].take(pc.sort_indices(indices))
135
177
 
136
178
  def min_max(self, **options) -> pa.Array:
137
179
  if pa.types.is_dictionary(self.type.value_type):
138
- self = ListChunk.aggregate(self, distinct=None).field(0)
180
+ (self,) = ListChunk.aggregate(self, distinct=None)
139
181
  self = type(self).from_arrays(self.offsets, self.values.dictionary_decode())
140
- return ListChunk.aggregate(self, min_max=pc.ScalarAggregateOptions(**options)).field(0)
182
+ return ListChunk.aggregate(self, min_max=pc.ScalarAggregateOptions(**options))[0]
141
183
 
142
184
  def min(self, **options) -> pa.Array:
143
185
  """min value of each list scalar"""
@@ -157,51 +199,45 @@ class ListChunk(pa.lib.BaseListArray):
157
199
 
158
200
  def index(self, **options) -> pa.Array:
159
201
  """index for first occurrence of each list scalar"""
160
- return pa.array(pc.index(value, **options).as_py() for value in ListChunk.scalars(self))
202
+ values = [pc.index(value, **options) for value in ListChunk.scalars(self)]
203
+ return Column.from_scalars(values)
204
+
205
+ @register
206
+ def list_all(ctx, self: pa.list_(pa.bool_())) -> pa.bool_(): # type: ignore
207
+ """Test whether all elements in a boolean array evaluate to true."""
208
+ return ListChunk.aggregate(self, all=None)[0]
209
+
210
+ @register
211
+ def list_any(ctx, self: pa.list_(pa.bool_())) -> pa.bool_(): # type: ignore
212
+ """Test whether any element in a boolean array evaluates to true."""
213
+ return ListChunk.aggregate(self, any=None)[0]
161
214
 
162
215
 
163
216
  class Column(pa.ChunkedArray):
164
217
  """Chunked array interface as a namespace of functions."""
165
218
 
166
- def map(self, func: Callable, **kwargs) -> pa.ChunkedArray:
167
- map_ = threader.map if self.num_chunks > 1 else map
168
- return pa.chunked_array(map_(functools.partial(func, **kwargs), self.iterchunks())) # type: ignore
219
+ def from_scalars(values: Sequence) -> pa.Array:
220
+ """Return array from arrow scalars."""
221
+ return pa.array((value.as_py() for value in values), values[0].type)
169
222
 
170
223
  def scalar_type(self):
171
224
  return self.type.value_type if pa.types.is_dictionary(self.type) else self.type
172
225
 
173
226
  def is_list_type(self):
174
- return pa.types.is_list(self.type) or pa.types.is_large_list(self.type)
227
+ funcs = pa.types.is_list, pa.types.is_large_list, pa.types.is_fixed_size_list
228
+ return any(func(self.type) for func in funcs)
175
229
 
176
- def combine_dictionaries(self) -> tuple:
230
+ def call_indices(self, func: Callable) -> Array:
177
231
  if not pa.types.is_dictionary(self.type):
178
- return self, None
179
- if not self or len(self) <= max(len(chunk.dictionary) for chunk in self.iterchunks()):
180
- return self.cast(self.type.value_type), None
181
- self = self.unify_dictionaries()
182
- indices = pa.chunked_array(chunk.indices for chunk in self.iterchunks())
183
- return self.chunk(0).dictionary, indices
184
-
185
- def indices(self) -> Union[pa.Array, pa.ChunkedArray]:
186
- """Return chunked indices suitable for aggregation."""
187
- if isinstance(self, pa.Array):
188
- return pa.array(np.arange(len(self)))
189
- offsets = list(itertools.accumulate(map(len, self.iterchunks())))
190
- return pa.chunked_array(map(np.arange, [0] + offsets, offsets))
191
-
192
- def call_indices(self, func: Callable) -> pa.ChunkedArray:
193
- dictionary, indices = Column.combine_dictionaries(self)
194
- if indices is None:
195
- return func(dictionary)
196
- return pa.chunked_array(
197
- pa.DictionaryArray.from_arrays(chunk, dictionary) for chunk in func(indices).chunks
198
- )
232
+ return func(self)
233
+ array = self.combine_chunks()
234
+ return pa.DictionaryArray.from_arrays(func(array.indices), array.dictionary)
199
235
 
200
- def fill_null_backward(self) -> pa.ChunkedArray:
236
+ def fill_null_backward(self) -> Array:
201
237
  """`fill_null_backward` with dictionary support."""
202
238
  return Column.call_indices(self, pc.fill_null_backward)
203
239
 
204
- def fill_null_forward(self) -> pa.ChunkedArray:
240
+ def fill_null_forward(self) -> Array:
205
241
  """`fill_null_forward` with dictionary support."""
206
242
  return Column.call_indices(self, pc.fill_null_forward)
207
243
 
@@ -209,11 +245,11 @@ class Column(pa.ChunkedArray):
209
245
  """Optimized `fill_null` to check `null_count`."""
210
246
  return self.fill_null(value) if self.null_count else self
211
247
 
212
- def sort_values(self) -> pa.Array:
213
- dictionary, indices = Column.combine_dictionaries(self)
214
- if indices is None:
215
- return dictionary
216
- return pc.sort_indices(pc.sort_indices(dictionary)).take(indices)
248
+ def sort_values(self) -> Array:
249
+ if not pa.types.is_dictionary(self.type):
250
+ return self
251
+ array = self if isinstance(self, pa.Array) else self.combine_chunks()
252
+ return pc.sort_indices(pc.sort_indices(array.dictionary)).take(array.indices)
217
253
 
218
254
  def diff(self, func: Callable = pc.subtract) -> pa.ChunkedArray:
219
255
  """Return discrete differences between adjacent values."""
@@ -224,7 +260,7 @@ class Column(pa.ChunkedArray):
224
260
 
225
261
  Args:
226
262
  predicate: binary function applied to adjacent values
227
- args: apply binary function to scalar, using `subtract` as the difference function
263
+ *args: apply binary function to scalar, using `subtract` as the difference function
228
264
  """
229
265
  ends = [pa.array([True])]
230
266
  mask = predicate(Column.diff(self), *args) if args else Column.diff(self, predicate)
@@ -243,12 +279,6 @@ class Column(pa.ChunkedArray):
243
279
  """Return max of the values."""
244
280
  return Column.min_max(self, **options)['max']
245
281
 
246
- def digitize(self, bins: Iterable, right=False) -> pa.ChunkedArray:
247
- """Return the indices of the bins to which each value in input array belongs."""
248
- if not isinstance(bins, (pa.Array, np.ndarray)):
249
- bins = pa.array(bins, self.type)
250
- return Column.map(self, np.digitize, bins=bins, right=bool(right))
251
-
252
282
  def index(self, value, start=0, end=None) -> int:
253
283
  """Return the first index of a value."""
254
284
  with contextlib.suppress(NotImplementedError):
@@ -283,10 +313,10 @@ class Column(pa.ChunkedArray):
283
313
  class Table(pa.Table):
284
314
  """Table interface as a namespace of functions."""
285
315
 
286
- def map_batch(self, func: Callable, *rargs, **kwargs) -> pa.Table:
287
- return pa.Table.from_batches(
288
- threader.map(lambda batch: func(batch, *rargs, **kwargs), self.to_batches())
289
- )
316
+ def map_batch(scanner: ds.Scanner, func: Callable, *rargs, **kwargs) -> pa.Table:
317
+ # TODO(apache/arrow#31612): replace with user defined function for multiple kernels
318
+ batches = [func(batch, *rargs, **kwargs) for batch in scanner.to_batches() if batch]
319
+ return pa.Table.from_batches(batches, None if batches else scanner.projected_schema)
290
320
 
291
321
  def union(*tables: pa.Table) -> pa.Table:
292
322
  """Return table with union of columns."""
@@ -316,7 +346,7 @@ class Table(pa.Table):
316
346
  Assumes the table is sorted by the column name, i.e., indexed.
317
347
  """
318
348
  (slc,) = Column.find(self[name], value)
319
- return pa.concat_tables([self[: slc.start], self[slc.stop :]]) # noqa: E203
349
+ return pa.concat_tables([self[: slc.start], self[slc.stop :]])
320
350
 
321
351
  def from_offsets(self, offsets: pa.IntegerArray) -> pa.Table:
322
352
  """Return table with columns converted into list columns."""
@@ -333,8 +363,8 @@ class Table(pa.Table):
333
363
  """Return table partitioned by discrete differences and corresponding counts.
334
364
 
335
365
  Args:
336
- names: columns to partition by `not_equal` which will return scalars
337
- predicates: inequality predicates with optional args which will return list arrays;
366
+ *names: columns to partition by `not_equal` which will return scalars
367
+ **predicates: inequality predicates with optional args which will return list arrays;
338
368
  if the predicate has args, it will be called on the differences
339
369
  """
340
370
  offsets = pa.chunked_array(
@@ -358,65 +388,89 @@ class Table(pa.Table):
358
388
  """Group by and aggregate.
359
389
 
360
390
  Args:
361
- names: columns to group by
391
+ *names: columns to group by
362
392
  counts: alias for optional row counts
363
393
  first: columns to take first value
364
394
  last: columns to take last value
365
- funcs: aggregate funcs with columns options
395
+ **funcs: aggregate funcs with columns options
366
396
  """
367
- none = not (counts or first or last or any(funcs.values()))
368
- column = self[names[0]]
369
- args = [(column, 'hash_count', pc.CountOptions(mode='all'))] if none or counts else []
397
+ if isinstance(self, pa.Table):
398
+ self = self.unify_dictionaries()
399
+ columns = {name: self[name] for name in self.schema.names}
400
+ aliases, args = {}, [] # type: ignore
401
+ if counts:
402
+ aliases['count_all'] = counts
403
+ args.append(([], 'count_all'))
370
404
  if first or last:
371
- args.append((Column.indices(column), 'hash_min_max', None))
405
+ columns[''] = np.arange(len(self))
406
+ args.append(('', 'min_max'))
372
407
  lists: Sequence[Agg] = []
373
408
  list_cols = [self[agg.name] for agg in funcs.get('list', [])]
374
409
  if any(pa.types.is_dictionary(col.type) or Column.is_list_type(col) for col in list_cols):
375
- args.append((Column.indices(column), 'hash_list', None))
410
+ columns[''] = np.arange(len(self))
411
+ args.append(('', 'list'))
376
412
  lists = funcs.pop('list')
377
- for func in funcs:
378
- args += [(self[agg.name], *agg.astuple(func)) for agg in funcs[func]] # type: ignore
379
- values, hashes, options = zip(*args)
380
- keys = map(self.column, names)
381
- if isinstance(self, pa.Table):
382
- keys = (key.unify_dictionaries() for key in keys) # type: ignore
383
- values = (
384
- value.unify_dictionaries() if 'distinct' in func else value
385
- for value, func in zip(values, hashes)
386
- )
387
- arrays = iter(pc._group_by(values, keys, zip(hashes, options)).flatten())
388
- if none:
389
- next(arrays)
390
- columns = {counts: next(arrays)} if counts else {}
413
+ for func, aggs in funcs.items():
414
+ aliases.update({f'{agg.name}_{func}': agg.alias for agg in aggs})
415
+ args += (agg.astuple(func) for agg in aggs) # type: ignore
416
+ table = pa.table(columns).group_by(list(names)).aggregate(args)
417
+ columns = {name: table[name].combine_chunks() for name in table.schema.names}
391
418
  if first or last:
392
- for aggs, indices in zip([first, last], next(arrays).flatten()):
419
+ for aggs, indices in zip([first, last], columns.pop('_min_max').flatten()):
393
420
  columns.update({agg.alias: self[agg.name].take(indices) for agg in aggs})
394
421
  if lists:
395
- indices = next(arrays)
422
+ indices = columns.pop('_list')
396
423
  table = self.select({agg.name for agg in lists}).take(indices.values)
397
424
  table = Table.from_offsets(table, indices.offsets)
398
- columns.update({agg.alias: table[agg.name] for agg in lists})
399
- for aggs in funcs.values():
400
- columns.update(zip([agg.alias for agg in aggs], arrays))
401
- columns.update(zip(names, map(decode, arrays)))
425
+ columns.update({agg.alias: table[agg.name].combine_chunks() for agg in lists})
426
+ columns.update({name: decode(columns[name]) for name in names})
427
+ columns = {aliases.get(name, name): columns[name] for name in columns}
402
428
  return type(self).from_pydict(columns)
403
429
 
430
+ def aggregate(self, counts: str = '', **funcs: Sequence[Agg]) -> dict:
431
+ """Return aggregated scalars as a row of data."""
432
+ row = {name: self[name] for name in self.column_names}
433
+ if counts:
434
+ row[counts] = len(self)
435
+ aliases, args = {}, [] # type: ignore
436
+ for key in funcs:
437
+ if key in Agg.option_map:
438
+ aliases.update({f'{agg.name}_{key}': agg.alias for agg in funcs[key]})
439
+ args += (agg.astuple(key) for agg in funcs[key])
440
+ else:
441
+ func = Agg.getfunc(key)
442
+ row.update({agg.alias: func(self[agg.name], **agg.options) for agg in funcs[key]})
443
+ if args:
444
+ table = self.group_by([]).aggregate(args)
445
+ row.update({aliases[name]: table[name][0] for name in table.column_names})
446
+ for name, value in row.items():
447
+ if isinstance(value, pa.ChunkedArray):
448
+ row[name] = value.combine_chunks()
449
+ return row
450
+
451
+ def list_fields(self) -> set:
452
+ return {field.name for field in self.schema if Column.is_list_type(field)}
453
+
404
454
  def list_value_length(self) -> pa.Array:
405
- lists = {name for name in self.column_names if Column.is_list_type(self[name])}
455
+ lists = Table.list_fields(self)
406
456
  if not lists:
407
- raise ValueError(f"no list columns available: {self.column_names}")
457
+ raise ValueError(f"no list columns available: {self.schema.names}")
408
458
  counts, *others = (pc.list_value_length(self[name]) for name in lists)
409
459
  if any(counts != other for other in others):
410
460
  raise ValueError(f"list columns have different value lengths: {lists}")
411
- return counts.chunk(0)
461
+ return counts if isinstance(counts, pa.Array) else counts.chunk(0)
412
462
 
413
- def sort_list(self, *names, length: int = None) -> pa.Table:
463
+ def sort_list(
464
+ self, *names: str, length: Optional[int] = None, null_placement: str = 'at_end'
465
+ ) -> pa.Table:
414
466
  """Return table with list columns sorted within scalars."""
415
467
  keys = dict(map(sort_key, names)) # type: ignore
416
468
  columns = {name: Column.sort_values(pc.list_flatten(self[name])) for name in keys}
417
469
  columns[''] = pc.list_parent_indices(self[next(iter(keys))])
418
470
  keys = dict({'': 'ascending'}, **keys)
419
- indices = pc.sort_indices(pa.table(columns), sort_keys=keys.items())
471
+ indices = pc.sort_indices(
472
+ pa.table(columns), sort_keys=keys.items(), null_placement=null_placement
473
+ )
420
474
  counts = Table.list_value_length(self)
421
475
  if length is not None:
422
476
  indices = pa.concat_arrays(
@@ -428,27 +482,35 @@ class Table(pa.Table):
428
482
 
429
483
  def select_list(self, apply: Callable = lambda c: c) -> pa.Table:
430
484
  """Return table with only the list columns."""
431
- names = [name for name in self.column_names if Column.is_list_type(self[name])]
485
+ names = list(Table.list_fields(self))
432
486
  return pa.table(list(map(apply, self.select(names))), names)
433
487
 
434
- def sort_indices(self, *names, length: int = None) -> pa.Table:
488
+ def sort_indices(
489
+ self, *names: str, length: Optional[int] = None, null_placement: str = 'at_end'
490
+ ) -> pa.Array:
435
491
  """Return indices which would sort the table by columns, optimized for fixed length."""
436
- func = pc.sort_indices
492
+ func = functools.partial(pc.sort_indices, null_placement=null_placement)
437
493
  if length is not None:
438
494
  func = functools.partial(pc.select_k_unstable, k=length)
439
495
  keys = dict(map(sort_key, names)) # type: ignore
440
496
  table = pa.table({name: Column.sort_values(self[name]) for name in keys})
441
497
  return func(table, sort_keys=keys.items())
442
498
 
443
- def sort(self, *names, length: int = None, indices: str = '') -> pa.Table:
499
+ def sort(
500
+ self,
501
+ *names: str,
502
+ length: Optional[int] = None,
503
+ indices: str = '',
504
+ null_placement: str = 'at_end',
505
+ ) -> pa.Table:
444
506
  """Return table sorted by columns, optimized for fixed length.
445
507
 
446
508
  Args:
447
- names: columns to sort by
509
+ *names: columns to sort by
448
510
  length: maximum number of rows to return
449
511
  indices: include original indices in the table
450
512
  """
451
- indices_ = Table.sort_indices(self, *names, length=length)
513
+ indices_ = Table.sort_indices(self, *names, length=length, null_placement=null_placement)
452
514
  table = self.take(indices_)
453
515
  if indices:
454
516
  table = table.append_column(indices, indices_)
@@ -456,7 +518,7 @@ class Table(pa.Table):
456
518
  metadata = {'index_columns': list(itertools.takewhile(func, names))}
457
519
  return table.replace_schema_metadata({'pandas': json.dumps(metadata)})
458
520
 
459
- def filter_list(self, expr: ds.Expression) -> 'Table':
521
+ def filter_list(self, expr: ds.Expression) -> pa.Table:
460
522
  """Return table with list columns filtered within scalars."""
461
523
  table = Table.select_list(self)
462
524
  counts = Table.list_value_length(table)
@@ -469,15 +531,16 @@ class Table(pa.Table):
469
531
  flattened = pa.table(columns, self.column_names)
470
532
  mask = ds.dataset(flattened).to_table(columns={'mask': expr})['mask'].combine_chunks()
471
533
  masks = type(first).from_arrays(first.offsets, mask)
472
- counts = Column.fill_null(ListChunk.aggregate(masks, sum=None).field(0), 0)
534
+ counts = Column.fill_null(ListChunk.aggregate(masks, sum=None)[0], 0)
473
535
  flattened = flattened.select(table.column_names).filter(mask)
474
536
  return Table.union(self, Table.from_counts(flattened, counts))
475
537
 
476
- def matched(self, func: Callable, *names: str) -> pa.Table:
477
- for name in names:
538
+ def min_max(self, *names: str) -> pa.Table:
539
+ for name, order in map(sort_key, names):
540
+ func = Column.min if order == 'ascending' else Column.max
478
541
  if Column.is_list_type(self[name]):
479
542
  self = self.append_column('', getattr(ListChunk, func.__name__)(self[name]))
480
543
  self = Table.filter_list(self, pc.field(name) == pc.field('')).drop([''])
481
544
  else:
482
- self = ds.dataset(self).to_table(filter=pc.field(name) == func(self[name]))
545
+ self = self.filter(pc.field(name) == func(self[name]))
483
546
  return self