graphique 1.7__tar.gz → 1.8__tar.gz

{graphique-1.7 → graphique-1.8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: graphique
-Version: 1.7
+Version: 1.8
 Summary: GraphQL service for arrow tables and parquet data sets.
 Author-email: Aric Coady <aric.coady@gmail.com>
 License: Copyright 2022 Aric Coady
@@ -32,6 +32,7 @@ Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
 Classifier: Topic :: Database :: Database Engines/Servers
 Classifier: Topic :: Internet :: WWW/HTTP :: HTTP Servers
 Classifier: Topic :: Internet :: WWW/HTTP :: Indexing/Search
@@ -40,8 +41,10 @@ Classifier: Typing :: Typed
 Requires-Python: >=3.9
 Description-Content-Type: text/markdown
 License-File: LICENSE.txt
-Requires-Dist: pyarrow>=17
+Requires-Dist: pyarrow>=18
 Requires-Dist: strawberry-graphql[asgi,cli]>=0.236
+Requires-Dist: numpy
+Requires-Dist: isodate>=0.7
 Provides-Extra: server
 Requires-Dist: uvicorn[standard]; extra == "server"
 Provides-Extra: cli
@@ -89,8 +92,9 @@ Supply a mapping of names to datasets for multiple roots, and to enable federati
 import pyarrow.dataset as ds
 from graphique import GraphQL
 
-app = GraphQL(ds.dataset(...))  # Table is root query type
-app = GraphQL.federated({<name>: ds.dataset(...), ...}, keys={...})  # Tables on federated fields
+source = ds.dataset(...)
+app = GraphQL(source)  # Table is root query type
+app = GraphQL.federated({<name>: source, ...}, keys={<name>: [], ...})  # Tables on federated fields
 ```
 
 Start like any ASGI app.
@@ -147,6 +151,8 @@ If index columns are detected in the schema metadata, then an initial `filter` w
 ## Dependencies
 * pyarrow
 * strawberry-graphql[asgi,cli]
+* numpy
+* isodate
 * uvicorn (or other [ASGI server](https://asgi.readthedocs.io/en/latest/implementations.html))
 
 ## Tests

{graphique-1.7 → graphique-1.8}/README.md

@@ -40,8 +40,9 @@ Supply a mapping of names to datasets for multiple roots, and to enable federati
 import pyarrow.dataset as ds
 from graphique import GraphQL
 
-app = GraphQL(ds.dataset(...))  # Table is root query type
-app = GraphQL.federated({<name>: ds.dataset(...), ...}, keys={...})  # Tables on federated fields
+source = ds.dataset(...)
+app = GraphQL(source)  # Table is root query type
+app = GraphQL.federated({<name>: source, ...}, keys={<name>: [], ...})  # Tables on federated fields
 ```
 
 Start like any ASGI app.
@@ -98,6 +99,8 @@ If index columns are detected in the schema metadata, then an initial `filter` w
 ## Dependencies
 * pyarrow
 * strawberry-graphql[asgi,cli]
+* numpy
+* isodate
 * uvicorn (or other [ASGI server](https://asgi.readthedocs.io/en/latest/implementations.html))
 
 ## Tests

{graphique-1.7 → graphique-1.8}/graphique/core.py

@@ -12,7 +12,7 @@ import inspect
 import itertools
 import operator
 import json
-from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence
+from collections.abc import Callable, Iterable, Iterator, Mapping
 from dataclasses import dataclass
 from typing import Optional, Union, get_type_hints
 import numpy as np
@@ -24,6 +24,8 @@ from typing_extensions import Self
 
 Array = Union[pa.Array, pa.ChunkedArray]
 Batch = Union[pa.RecordBatch, pa.Table]
+bit_any = functools.partial(functools.reduce, operator.or_)
+bit_all = functools.partial(functools.reduce, operator.and_)
 
 
 class Agg:
@@ -51,28 +53,15 @@ class Agg:
         'tdigest': pc.TDigestOptions,
         'variance': pc.VarianceOptions,
     }
-
-    associatives = {'all', 'any', 'first', 'last', 'max', 'min', 'one', 'product', 'sum'}
-    associatives |= {'count'}  # transformed to be associative
+    ordered = {'first', 'last'}
 
     def __init__(self, name: str, alias: str = '', **options):
         self.name = name
         self.alias = alias or name
         self.options = options
 
-    def astuple(self, func: str) -> tuple:
-        options = self.option_map[func.rpartition('hash_')[-1]](**self.options)
-        return self.name, func, options, self.alias
-
-    one = operator.itemgetter(0)
-    list = staticmethod(lambda array: array)
-
-    @staticmethod
-    def distinct(array: Array, mode: str = 'only_valid') -> pa.Array:
-        values = pc.unique(array)
-        if not values.null_count or mode == 'all':
-            return values
-        return values.drop_null() if mode == 'only_valid' else pa.array([None], array.type)
+    def func_options(self, func: str) -> pc.FunctionOptions:
+        return self.option_map[func.removeprefix('hash_')](**self.options)
 
 
 @dataclass(frozen=True)
@@ -210,8 +199,7 @@ class ListChunk(pa.lib.BaseListArray):
 
     def index(self, **options) -> pa.Array:
         """index for first occurrence of each list scalar"""
-        values = [pc.index(value, **options) for value in ListChunk.scalars(self)]
-        return Column.from_scalars(values)
+        return pa.array(pc.index(value, **options) for value in ListChunk.scalars(self))
 
     @register
     def list_all(ctx, self: pa.list_(pa.bool_())) -> pa.bool_():  # type: ignore
@@ -227,13 +215,6 @@ class ListChunk(pa.lib.BaseListArray):
 class Column(pa.ChunkedArray):
     """Chunked array interface as a namespace of functions."""
 
-    def from_scalars(values: Sequence) -> pa.Array:
-        """Return array from arrow scalars."""
-        return pa.array((value.as_py() for value in values), values[0].type)
-
-    def scalar_type(self):
-        return self.type.value_type if pa.types.is_dictionary(self.type) else self.type
-
     def is_list_type(self):
         funcs = pa.types.is_list, pa.types.is_large_list, pa.types.is_fixed_size_list
         return any(func(self.type) for func in funcs)
@@ -284,26 +265,6 @@ class Column(pa.ChunkedArray):
         mask = predicate(Column.diff(self), *args) if args else Column.diff(self, predicate)
         return pc.indices_nonzero(pa.chunked_array(ends + mask.chunks + ends))
 
-    def first_last(self, **options):
-        if not pa.types.is_dictionary(self.type):
-            return pc.first_last(self, **options).as_py()
-        self = self.combine_chunks()
-        indices = pc.first_last(self.indices, **options).as_py().items()
-        return {key: None if index is None else self.dictionary[index] for key, index in indices}
-
-    def min_max(self, **options):
-        if pa.types.is_dictionary(self.type):
-            self = self.unique().dictionary_decode()
-        return pc.min_max(self, **options).as_py()
-
-    def min(self, **options):
-        """Return min of the values."""
-        return Column.min_max(self, **options)['min']
-
-    def max(self, **options):
-        """Return max of the values."""
-        return Column.min_max(self, **options)['max']
-
     def index(self, value, start=0, end=None) -> int:
         """Return the first index of a value."""
         with contextlib.suppress(NotImplementedError):
@@ -338,10 +299,8 @@ class Column(pa.ChunkedArray):
 class Table(pa.Table):
     """Table interface as a namespace of functions."""
 
-    def map_batch(scanner: ds.Scanner, func: Callable, *rargs, **kwargs) -> pa.Table:
-        # TODO(apache/arrow#31612): replace with user defined function for multiple kernels
-        batches = [func(batch, *rargs, **kwargs) for batch in scanner.to_batches() if batch]
-        return pa.Table.from_batches(batches, None if batches else scanner.projected_schema)
+    def map_batch(self, func: Callable, *args, **kwargs) -> pa.Table:
+        return pa.Table.from_batches(func(batch, *args, **kwargs) for batch in self.to_batches())
 
     def columns(self) -> dict:
         """Return columns as a dictionary."""
@@ -411,39 +370,6 @@ class Table(pa.Table):
         table = Table.union(scalars, Table.from_offsets(lists, offsets))
         return table, Column.diff(offsets)
 
-    def group(self, *names: str, counts: str = '', **funcs: Sequence[Agg]) -> pa.Table:
-        """Group by and aggregate.
-
-        Args:
-            *names: columns to group by
-            counts: alias for optional row counts
-            **funcs: aggregate funcs with columns options
-        """
-        prefix = 'hash_' if names else ''
-        aggs = [agg.astuple(prefix + func) for func in funcs for agg in funcs[func]]
-        columns = list(names) + [agg[0] for agg in aggs]
-        if counts:
-            aggs.append(([], 'hash_count_all', None, counts))
-        if isinstance(self, pa.Table):
-            decl = Declaration('table_source', self.unify_dictionaries())
-        else:
-            decl = Declaration.scan(self, columns=columns)
-        return decl.aggregate(aggs, names).to_table(use_threads=False)
-
-    def aggregate(self, counts: str = '', **funcs: Sequence[Agg]) -> dict:
-        """Return aggregated scalars as a row of data."""
-        row = {counts: len(self)} if counts else {}
-        for key in ('one', 'list', 'distinct'):  # hash only functions
-            func, aggs = getattr(Agg, key), funcs.pop(key, [])
-            row |= {agg.alias: func(self[agg.name], **agg.options) for agg in aggs}
-        if funcs:
-            table = Table.group(self, **funcs)  # type: ignore
-            row |= {name: table[name][0] for name in table.schema.names}
-        for name, value in row.items():
-            if isinstance(value, pa.ChunkedArray):
-                row[name] = value.combine_chunks()
-        return row
-
     def list_fields(self) -> set:
         return {field.name for field in self.schema if Column.is_list_type(field)}
 
@@ -490,7 +416,7 @@ class Table(pa.Table):
             indices: include original indices in the table
         """
         if length == 1 and not indices:
-            return Table.ranked(self, 1, *names, null_placement=null_placement)[:1]
+            return Table.min_max(self, *names)[:1]
         indices_ = Table.sort_indices(self, *names, length=length, null_placement=null_placement)
         table = self.take(indices_)
         if indices:
@@ -510,49 +436,47 @@ class Table(pa.Table):
         table = pa.concat_tables(table for table in tables if table is not None)
         return Table.union(self, Table.from_counts(table, counts))
 
-    def ranked(self, k: int, *names: str, null_placement: str = 'at_end') -> Batch:
-        """Return table filtered by values within dense rank, similar to `select_k_unstable`."""
-        if k == 1:  # optimized for min_max
-            expr = isinstance(self, pa.Table)
-            for name, order in map(sort_key, names):
-                func = Column.min if order == 'ascending' else Column.max
-                value = func(self[name], skip_nulls=null_placement == 'at_end')
-                if value is None:
-                    mask = (pc.field(name) if expr else self[name]).is_null()
-                else:
-                    mask = pc.field(name) == value if expr else pc.equal(self[name], value)
-                self = self.filter(mask)
-            return self
-        values = None
-        kwargs = dict(tiebreaker='dense', null_placement=null_placement)
-        for name, order in map(sort_key, names):
-            if len(self) <= k:
-                break
-            ranks = pc.rank(Column.sort_values(self[name]), order, **kwargs)
-            if values is None:  # optimized to sort only once on first iteration
-                values = ranks
-            else:
-                values = pc.add_checked(pc.multiply_checked(values, pc.max(ranks)), ranks)
-                values = pc.rank(values, 'ascending', tiebreaker='dense')
-            mask = pc.less_equal(values, k)
-            self = self.filter(mask)
-            values = values.filter(mask)
+    def min_max(self, *names: str) -> Self:
+        """Return table filtered by minimum or maximum values."""
+        for key, order in map(sort_key, names):
+            field, asc = pc.field(key), (order == 'ascending')
+            ((value,),) = Nodes.group(self, _=(key, ('min' if asc else 'max'), None)).to_table()
+            self = self.filter(field <= value if asc else field >= value)
         return self
 
-    def get_fragments(self) -> Iterator[ds.Fragment]:
-        """Support filtered datasets if it only references partition keys."""
-        expr = self._scan_options.get('filter')
-        if expr is not None:  # raise ValueError if filter references other fields
-            ds.dataset([], schema=self.partitioning.schema).scanner(filter=expr)
-        return self._get_fragments(expr)
-
-    def fragment_keys(self) -> list:
-        """Filtered partitioned datasets may not have fragments."""
+    def rank(self, k: int, *names: str) -> Self:
+        """Return table filtered by values within dense rank, similar to `select_k_unstable`."""
+        if k == 1:
+            return Table.min_max(self, *names)
+        keys = dict(map(sort_key, names))
+        table = Nodes.group(self, *keys).to_table()
+        table = table.take(pc.select_k_unstable(table, k, keys.items()))
+        exprs = []
+        for key, order in keys.items():
+            field, asc = pc.field(key), (order == 'ascending')
+            exprs.append(field <= pc.max(table[key]) if asc else field >= pc.min(table[key]))
+        return self.filter(bit_all(exprs))
+
+    def fragments(self, *names, counts: str = '') -> pa.Table:
+        """Return selected fragment keys in a table."""
         try:
-            Table.get_fragments(self)
+            expr = self._scan_options.get('filter')
+            if expr is not None:  # raise ValueError if filter references other fields
+                ds.dataset([], schema=self.partitioning.schema).scanner(filter=expr)
         except (AttributeError, ValueError):
-            return []
-        return self.partitioning.schema.names
+            return pa.table({})
+        fragments = self._get_fragments(expr)
+        parts = [ds.get_partition_keys(frag.partition_expression) for frag in fragments]
+        names, table = set(names), pa.Table.from_pylist(parts)  # type: ignore
+        keys = [name for name in table.schema.names if name in names]
+        table = table.group_by(keys, use_threads=False).aggregate([])
+        if not counts:
+            return table
+        if not table.schema:
+            return table.append_column(counts, pa.array([self.count_rows()]))
+        exprs = [bit_all(pc.field(key) == row[key] for key in row) for row in table.to_pylist()]
+        column = [self.filter(expr).count_rows() for expr in exprs]
+        return table.append_column(counts, pa.array(column))
 
     def rank_keys(self, k: int, *names: str, dense: bool = True) -> tuple:
         """Return expression and unmatched fields for partitioned dataset which filters by rank.
@@ -562,39 +486,24 @@ class Table(pa.Table):
             *names: columns to rank by
             dense: use dense rank; false indicates sorting
         """
-        schema = set(Table.fragment_keys(self))
-        keys = dict(itertools.takewhile(lambda key: key[0] in schema, map(sort_key, names)))
+        keys = dict(map(sort_key, names))
+        table = Table.fragments(self, *keys, counts='' if dense else '_')
+        keys = {name: keys[name] for name in table.schema.names if name in keys}
         if not keys:
             return None, names
-        parts = []
-        for fragment in Table.get_fragments(self):
-            parts.append(ds.get_partition_keys(fragment.partition_expression))
-            if not dense:
-                parts[-1][''] = fragment.count_rows()
-        offset, table = len(keys), pa.Table.from_pylist(parts)
         if dense:
-            table = Table.ranked(table, k, *names[:offset])
+            table = table.take(pc.select_k_unstable(table, k, keys.items()))
         else:
-            table = Table.sort(table, *names[:offset])
-            index = pc.index(pc.greater_equal(pc.cumulative_sum(table['']), k), True).as_py()
-            table = table[: index + 1 if index >= 0 else None]
-        remaining = names[offset:]
+            table = table.sort_by(keys.items())
+            totals = itertools.accumulate(table['_'].to_pylist())
+            counts = (count for count, total in enumerate(totals, 1) if total >= k)
+            table = table[: next(counts, None)].remove_column(len(table) - 1)
+        exprs = [bit_all(pc.field(key) == row[key] for key in row) for row in table.to_pylist()]
+        remaining = names[len(keys) :]
         if remaining or not dense:  # fields with a single value are no longer needed
-            selectors = [len(table[name].unique()) > 1 for name in keys]
+            selectors = [len(table[key].unique()) > 1 for key in keys]
             remaining = tuple(itertools.compress(names, selectors)) + remaining
-        fields, prefix = [], []  # type: ignore
-        for index, (name, order) in enumerate(keys.items()):
-            expr, field, values = None, ds.field(name), table[name].unique()
-            value = pc.max(values) if order == 'ascending' else pc.min(values)
-            if index == len(keys) - 1:  # last one
-                expr = field <= value if order == 'ascending' else field >= value
-            elif len(values) > 1:
-                expr = field < value if order == 'ascending' else field > value
-            if expr is not None:
-                fields.append(functools.reduce(operator.and_, prefix + [expr]))
-            prefix.append(field == value)
-            table = table.filter(prefix[-1])
-        return functools.reduce(operator.or_, fields), remaining
+        return bit_any(exprs[: len(table)]), remaining
 
     def flatten(self, indices: str = '') -> Iterator[pa.RecordBatch]:
         """Generate batches with list arrays flattened, optionally with parent indices."""
@@ -632,8 +541,11 @@ class Table(pa.Table):
         return f'{size:n} {prefix}B'
 
 
-class Declaration(ac.Declaration):
-    """[Acero](https://arrow.apache.org/docs/python/api/acero.html) engine declaration."""
+class Nodes(ac.Declaration):
+    """[Acero](https://arrow.apache.org/docs/python/api/acero.html) engine declaration.
+
+    Provides a `Scanner` interface with no "oneshot" limitation.
+    """
 
     option_map = {
         'table_source': ac.TableSourceNodeOptions,
@@ -644,25 +556,65 @@ class Declaration(ac.Declaration):
         'order_by': ac.OrderByNodeOptions,
         'hashjoin': ac.HashJoinNodeOptions,
     }
+    to_batches = ac.Declaration.to_reader  # source compatibility
 
     def __init__(self, name, *args, inputs=None, **options):
         super().__init__(name, self.option_map[name](*args, **options), inputs)
 
-    @classmethod
-    def scan(cls, dataset: ds.Dataset, columns=None) -> Self:
-        """Return source node from a dataset."""
-        expr = dataset._scan_options.get('filter')
-        self = cls('scan', dataset, filter=expr, columns=columns)
-        if expr is not None:
-            self = self.filter(expr)
+    def scan(self, columns: Iterable[str]) -> Self:
+        """Return projected source node, supporting datasets and tables."""
+        if isinstance(self, ds.Dataset):
+            expr = self._scan_options.get('filter')
+            self = Nodes('scan', self, columns=columns)
+            if expr is not None:
+                self = self.apply('filter', expr)
+        elif isinstance(self, pa.Table):
+            self = Nodes('table_source', self)
+        elif isinstance(self, pa.RecordBatch):
+            self = Nodes('table_source', pa.table(self))
         if isinstance(columns, Mapping):
-            self = self.project(columns.values(), columns)
-        elif columns is not None:
-            self = self.project(map(pc.field, columns))
-        return self
+            return self.apply('project', columns.values(), columns)
+        return self.apply('project', map(pc.field, columns))
+
+    @property
+    def schema(self) -> pa.Schema:
+        """projected schema"""
+        with self.to_reader() as reader:
+            return reader.schema
+
+    def scanner(self, **options) -> ds.Scanner:
+        return ds.Scanner.from_batches(self.to_reader(**options))
+
+    def count_rows(self) -> int:
+        """Count matching rows."""
+        return self.scanner().count_rows()
+
+    def head(self, num_rows: int, **options) -> pa.Table:
+        """Load the first N rows."""
+        return self.scanner(**options).head(num_rows)
+
+    def take(self, indices: Iterable[int], **options) -> pa.Table:
+        """Select rows by index."""
+        return self.scanner(**options).take(indices)
 
     def apply(self, name: str, *args, **options) -> Self:
+        """Add a node by name."""
         return type(self)(name, *args, inputs=[self], **options)
 
-    def __getattr__(self, name: str) -> Callable:
-        return functools.partial(self.apply, name)
+    filter = functools.partialmethod(apply, 'filter')
+
+    def group(self, *names, **aggs: tuple) -> Self:
+        """Add `aggregate` node with dictionary support.
+
+        Also supports datasets because aggregation determines the projection.
+        """
+        aggregates, targets = [], set(names)
+        for name, (target, _, _) in aggs.items():
+            aggregates.append(aggs[name] + (name,))
+            targets.update([target] if isinstance(target, str) else target)
+        columns = {name: pc.field(name) for name in targets}
+        for name in columns:
+            field = self.schema.field(name)
+            if pa.types.is_dictionary(field.type):
+                columns[name] = columns[name].cast(field.type.value_type)
+        return Nodes.scan(self, columns).apply('aggregate', aggregates, names)