graphique 1.8__tar.gz → 2.0__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 (37) hide show
  1. {graphique-1.8 → graphique-2.0}/PKG-INFO +53 -57
  2. graphique-2.0/README.md +119 -0
  3. graphique-2.0/graphique/__init__.py +3 -0
  4. graphique-2.0/graphique/core.py +83 -0
  5. graphique-2.0/graphique/inputs.py +472 -0
  6. graphique-2.0/graphique/interface.py +384 -0
  7. {graphique-1.8 → graphique-2.0}/graphique/middleware.py +19 -16
  8. graphique-2.0/graphique/models.py +260 -0
  9. {graphique-1.8 → graphique-2.0}/graphique/scalars.py +36 -40
  10. graphique-2.0/graphique/service.py +44 -0
  11. {graphique-1.8 → graphique-2.0}/graphique/shell.py +6 -4
  12. {graphique-1.8 → graphique-2.0}/graphique.egg-info/PKG-INFO +53 -57
  13. {graphique-1.8 → graphique-2.0}/graphique.egg-info/SOURCES.txt +0 -2
  14. {graphique-1.8 → graphique-2.0}/graphique.egg-info/requires.txt +4 -2
  15. {graphique-1.8 → graphique-2.0}/pyproject.toml +26 -16
  16. graphique-2.0/tests/test_core.py +44 -0
  17. graphique-2.0/tests/test_dataset.py +167 -0
  18. graphique-2.0/tests/test_models.py +186 -0
  19. graphique-2.0/tests/test_service.py +342 -0
  20. graphique-1.8/README.md +0 -111
  21. graphique-1.8/graphique/__init__.py +0 -1
  22. graphique-1.8/graphique/core.py +0 -620
  23. graphique-1.8/graphique/inputs.py +0 -703
  24. graphique-1.8/graphique/interface.py +0 -516
  25. graphique-1.8/graphique/models.py +0 -295
  26. graphique-1.8/graphique/service.py +0 -40
  27. graphique-1.8/requirements.in +0 -4
  28. graphique-1.8/tests/test_bench.py +0 -21
  29. graphique-1.8/tests/test_core.py +0 -189
  30. graphique-1.8/tests/test_dataset.py +0 -278
  31. graphique-1.8/tests/test_models.py +0 -390
  32. graphique-1.8/tests/test_service.py +0 -483
  33. {graphique-1.8 → graphique-2.0}/LICENSE.txt +0 -0
  34. {graphique-1.8 → graphique-2.0}/graphique/py.typed +0 -0
  35. {graphique-1.8 → graphique-2.0}/graphique.egg-info/dependency_links.txt +0 -0
  36. {graphique-1.8 → graphique-2.0}/graphique.egg-info/top_level.txt +0 -0
  37. {graphique-1.8 → graphique-2.0}/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.8
4
- Summary: GraphQL service for arrow tables and parquet data sets.
3
+ Version: 2.0
4
+ Summary: GraphQL service for python dataframes and parquet datasets.
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-Expression: 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,10 +12,8 @@ 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
@@ -38,17 +23,20 @@ Classifier: Topic :: Internet :: WWW/HTTP :: HTTP Servers
38
23
  Classifier: Topic :: Internet :: WWW/HTTP :: Indexing/Search
39
24
  Classifier: Topic :: Software Development :: Libraries :: Python Modules
40
25
  Classifier: Typing :: Typed
41
- Requires-Python: >=3.9
26
+ Requires-Python: >=3.10
42
27
  Description-Content-Type: text/markdown
43
28
  License-File: LICENSE.txt
44
- Requires-Dist: pyarrow>=18
29
+ Requires-Dist: pyarrow>=20
45
30
  Requires-Dist: strawberry-graphql[asgi,cli]>=0.236
46
- Requires-Dist: numpy
31
+ Requires-Dist: ibis-framework>=10.7
32
+ Requires-Dist: pyarrow-hotfix
47
33
  Requires-Dist: isodate>=0.7
48
34
  Provides-Extra: server
49
35
  Requires-Dist: uvicorn[standard]; extra == "server"
36
+ Requires-Dist: duckdb; extra == "server"
50
37
  Provides-Extra: cli
51
38
  Requires-Dist: tqdm; extra == "cli"
39
+ Dynamic: license-file
52
40
 
53
41
  [![image](https://img.shields.io/pypi/v/graphique.svg)](https://pypi.org/project/graphique/)
54
42
  ![image](https://img.shields.io/pypi/pyversions/graphique.svg)
@@ -57,42 +45,52 @@ Requires-Dist: tqdm; extra == "cli"
57
45
  [![build](https://github.com/coady/graphique/actions/workflows/build.yml/badge.svg)](https://github.com/coady/graphique/actions/workflows/build.yml)
58
46
  [![image](https://codecov.io/gh/coady/graphique/branch/main/graph/badge.svg)](https://codecov.io/gh/coady/graphique/)
59
47
  [![CodeQL](https://github.com/coady/graphique/actions/workflows/github-code-scanning/codeql/badge.svg)](https://github.com/coady/graphique/actions/workflows/github-code-scanning/codeql)
60
- [![CodSpeed Badge](https://img.shields.io/endpoint?url=https://codspeed.io/badge.json)](https://codspeed.io/coady/graphique)
61
48
  [![image](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json)](https://github.com/astral-sh/ruff)
62
49
  [![image](https://mypy-lang.org/static/mypy_badge.svg)](https://mypy-lang.org/)
63
50
 
64
- [GraphQL](https://graphql.org) service for [arrow](https://arrow.apache.org) tables and [parquet](https://parquet.apache.org) data sets. The schema for a query API is derived automatically.
51
+ [GraphQL](https://graphql.org) service for [ibis](https://ibis-project.org) dataframes, [arrow](https://arrow.apache.org) tables, and [parquet](https://parquet.apache.org) datasets. The schema for a query API is derived automatically.
52
+
53
+ ## Version 2
54
+ When this project started, there was no out-of-core execution engine with performance comparable to [PyArrow](https://arrow.apache.org/docs/python/index.html). So it effectively included one, based on datasets and [Acero](https://arrow.apache.org/docs/python/api/acero.html).
55
+
56
+ Since then the ecosystem has grown considerably: [DuckDB](https://duckdb.org), [DataFusion](https://datafusion.apache.org), and [Ibis](https://ibis-project.org). As of version 2, graphique is based on `ibis`. It provides a common dataframe API for multiple backends, enabling graphique to also have a default but configurable backend.
57
+
58
+ Being a major version upgrade, there are incompatible changes from version 1. However the overall API remains largely the same.
65
59
 
66
60
  ## Usage
61
+ There is an example app which reads a parquet dataset.
67
62
  ```console
68
- % env PARQUET_PATH=... uvicorn graphique.service:app
63
+ env PARQUET_PATH=... uvicorn graphique.service:app
69
64
  ```
70
65
 
71
66
  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`.
72
67
 
73
68
  ```console
74
- % env PARQUET_PATH=... strawberry export-schema graphique.service:app.schema
69
+ env PARQUET_PATH=... strawberry export-schema graphique.service:app.schema
75
70
  ```
76
- outputs the graphql schema for a parquet data set.
71
+ outputs the graphql schema.
77
72
 
78
73
  ### Configuration
79
- Graphique uses [Starlette's config](https://www.starlette.io/config/): in environment variables or a `.env` file. Config variables are used as input to a [parquet dataset](https://arrow.apache.org/docs/python/dataset.html).
74
+ The example app uses [Starlette's config](https://www.starlette.io/config/): in environment variables or a `.env` file.
80
75
 
81
76
  * PARQUET_PATH: path to the parquet directory or file
82
- * FEDERATED = '': field name to extend type `Query` with a federated `Table`
83
- * DEBUG = False: run service in debug mode, which includes metrics
77
+ * FEDERATED = '': field name to extend type `Query` with a federated `Table`
78
+ * METRICS = False: include timings from apollo tracing extension
84
79
  * COLUMNS = None: list of names, or mapping of aliases, of columns to select
85
80
  * FILTERS = None: json `filter` query for which rows to read at startup
86
81
 
87
- For more options create a custom [ASGI](https://asgi.readthedocs.io/en/latest/index.html) app. Call graphique's `GraphQL` on an arrow [Dataset](https://arrow.apache.org/docs/python/api/dataset.html), [Scanner](https://arrow.apache.org/docs/python/generated/pyarrow.dataset.Scanner.html), or [Table](https://arrow.apache.org/docs/python/generated/pyarrow.Table.html). The GraphQL `Table` type will be the root Query type.
82
+ 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.
88
83
 
84
+ ### App
85
+ For more options create a custom [ASGI](https://asgi.readthedocs.io/en/latest/index.html) app. Call graphique's `GraphQL` on an ibis [Table](https://ibis-project.org/reference/expression-tables) or arrow [Dataset](https://arrow.apache.org/docs/python/api/dataset.html).
89
86
  Supply a mapping of names to datasets for multiple roots, and to enable federation.
90
87
 
91
88
  ```python
92
- import pyarrow.dataset as ds
89
+ import ibis
93
90
  from graphique import GraphQL
94
91
 
95
- source = ds.dataset(...)
92
+ source = ibis.read_*(...) # or ibis.connect(...).table(...) or pyarrow.dataset.dataset(...)
93
+ # apply initial projections or filters to `source`
96
94
  app = GraphQL(source) # Table is root query type
97
95
  app = GraphQL.federated({<name>: source, ...}, keys={<name>: [], ...}) # Tables on federated fields
98
96
  ```
@@ -103,55 +101,53 @@ Start like any ASGI app.
103
101
  uvicorn <module>:app
104
102
  ```
105
103
 
106
- 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.
107
-
108
104
  ### API
109
105
  #### types
110
- * `Dataset`: interface for an arrow dataset, scanner, or table.
106
+ * `Dataset`: interface for an ibis table or arrow dataset.
111
107
  * `Table`: implements the `Dataset` interface. Adds typed `row`, `columns`, and `filter` fields from introspecting the schema.
112
- * `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.
113
- * `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.
108
+ * `Column`: interface for an ibis column. Each data type has a corresponding column implementation: Boolean, Int, BigInt, Float, Decimal, Date, Datetime, Time, Duration, Base64, String, Array, Struct. All columns have a `values` field for their list of scalars. Additional fields vary by type.
109
+ * `Row`: scalar fields. 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.
114
110
 
115
111
  #### selection
116
112
  * `slice`: contiguous selection of rows
117
- * `filter`: select rows with simple predicates
118
- * `scan`: select rows and project columns with expressions
113
+ * `filter`: select rows by predicates
114
+ * `join`: join tables by key columns
115
+ * `take`: rows by index
116
+ * `dropNull`: remove rows with nulls
119
117
 
120
118
  #### projection
119
+ * `project`: project columns with expressions
121
120
  * `columns`: provides a field for every `Column` in the schema
122
121
  * `column`: access a column of any type by name
123
122
  * `row`: provides a field for each scalar of a single row
124
- * `apply`: transform columns by applying a function
125
- * `join`: join tables by key columns
123
+ * `cast`: cast column types
124
+ * `fillNull`: fill null values
126
125
 
127
126
  #### aggregation
128
127
  * `group`: group by given columns, and aggregate the others
129
- * `runs`: partition on adjacent values in given columns, transforming the others into list columns
130
- * `tables`: return a list of tables by splitting on the scalars in list columns
131
- * `flatten`: flatten list columns with repeated scalars
128
+ * `distinct`: group with all columns
129
+ * `runs`: provisionally group by adjacency
130
+ * `unnest`: unnest an array column
131
+ * `count`: number of rows
132
132
 
133
133
  #### ordering
134
- * `sort`: sort table by given columns
135
- * `rank`: select rows with smallest or largest values
134
+ * `order`: sort table by given columns
135
+ * options `limit` and `dense`: select rows with smallest or largest values
136
136
 
137
137
  ### Performance
138
- Graphique relies on native [PyArrow](https://arrow.apache.org/docs/python/index.html) routines wherever possible. Otherwise it falls back to using [NumPy](https://numpy.org/doc/stable/) or custom optimizations.
139
-
140
- By default, datasets are read on-demand, with only the necessary rows and columns scanned. Although graphique is a running service, [parquet is performant](https://arrow.apache.org/docs/python/generated/pyarrow.dataset.Dataset.html) at reading a subset of data. Optionally specify `FILTERS` in the json `filter` format to read a subset of rows at startup, trading-off memory for latency. An empty filter (`{}`) will read the whole table.
138
+ Performance is dependent on the [ibis backend](https://ibis-project.org/backends/duckdb), which defaults to [duckdb](https://duckdb.org/). There are no internal Python loops. Scalars do not become Python types until serialized.
141
139
 
142
- Specifying `COLUMNS` will limit memory usage when reading at startup (`FILTERS`). There is little speed difference as unused columns are inherently ignored. Optional aliasing can also be used for camel casing.
143
-
144
- If index columns are detected in the schema metadata, then an initial `filter` will also attempt a binary search on tables.
140
+ [PyArrow](https://arrow.apache.org/docs/python/) is also used for partitioned dataset optimizations, and for any feature which ibis does not support. Table fields are lazily evaluated up until scalars are reached, and automatically cached as needed for multiple fields.
145
141
 
146
142
  ## Installation
147
143
  ```console
148
- % pip install graphique[server]
144
+ pip install graphique[server]
149
145
  ```
150
146
 
151
147
  ## Dependencies
152
- * pyarrow
148
+ * ibis-framework (with duckdb or other backend)
153
149
  * strawberry-graphql[asgi,cli]
154
- * numpy
150
+ * pyarrow
155
151
  * isodate
156
152
  * uvicorn (or other [ASGI server](https://asgi.readthedocs.io/en/latest/implementations.html))
157
153
 
@@ -159,5 +155,5 @@ If index columns are detected in the schema metadata, then an initial `filter` w
159
155
  100% branch coverage.
160
156
 
161
157
  ```console
162
- % pytest [--cov]
158
+ pytest [--cov]
163
159
  ```
@@ -0,0 +1,119 @@
1
+ [![image](https://img.shields.io/pypi/v/graphique.svg)](https://pypi.org/project/graphique/)
2
+ ![image](https://img.shields.io/pypi/pyversions/graphique.svg)
3
+ [![image](https://pepy.tech/badge/graphique)](https://pepy.tech/project/graphique)
4
+ ![image](https://img.shields.io/pypi/status/graphique.svg)
5
+ [![build](https://github.com/coady/graphique/actions/workflows/build.yml/badge.svg)](https://github.com/coady/graphique/actions/workflows/build.yml)
6
+ [![image](https://codecov.io/gh/coady/graphique/branch/main/graph/badge.svg)](https://codecov.io/gh/coady/graphique/)
7
+ [![CodeQL](https://github.com/coady/graphique/actions/workflows/github-code-scanning/codeql/badge.svg)](https://github.com/coady/graphique/actions/workflows/github-code-scanning/codeql)
8
+ [![image](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json)](https://github.com/astral-sh/ruff)
9
+ [![image](https://mypy-lang.org/static/mypy_badge.svg)](https://mypy-lang.org/)
10
+
11
+ [GraphQL](https://graphql.org) service for [ibis](https://ibis-project.org) dataframes, [arrow](https://arrow.apache.org) tables, and [parquet](https://parquet.apache.org) datasets. The schema for a query API is derived automatically.
12
+
13
+ ## Version 2
14
+ When this project started, there was no out-of-core execution engine with performance comparable to [PyArrow](https://arrow.apache.org/docs/python/index.html). So it effectively included one, based on datasets and [Acero](https://arrow.apache.org/docs/python/api/acero.html).
15
+
16
+ Since then the ecosystem has grown considerably: [DuckDB](https://duckdb.org), [DataFusion](https://datafusion.apache.org), and [Ibis](https://ibis-project.org). As of version 2, graphique is based on `ibis`. It provides a common dataframe API for multiple backends, enabling graphique to also have a default but configurable backend.
17
+
18
+ Being a major version upgrade, there are incompatible changes from version 1. However the overall API remains largely the same.
19
+
20
+ ## Usage
21
+ There is an example app which reads a parquet dataset.
22
+ ```console
23
+ env PARQUET_PATH=... uvicorn graphique.service:app
24
+ ```
25
+
26
+ 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`.
27
+
28
+ ```console
29
+ env PARQUET_PATH=... strawberry export-schema graphique.service:app.schema
30
+ ```
31
+ outputs the graphql schema.
32
+
33
+ ### Configuration
34
+ The example app uses [Starlette's config](https://www.starlette.io/config/): in environment variables or a `.env` file.
35
+
36
+ * PARQUET_PATH: path to the parquet directory or file
37
+ * FEDERATED = '': field name to extend type `Query` with a federated `Table`
38
+ * METRICS = False: include timings from apollo tracing extension
39
+ * COLUMNS = None: list of names, or mapping of aliases, of columns to select
40
+ * FILTERS = None: json `filter` query for which rows to read at startup
41
+
42
+ 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.
43
+
44
+ ### App
45
+ For more options create a custom [ASGI](https://asgi.readthedocs.io/en/latest/index.html) app. Call graphique's `GraphQL` on an ibis [Table](https://ibis-project.org/reference/expression-tables) or arrow [Dataset](https://arrow.apache.org/docs/python/api/dataset.html).
46
+ Supply a mapping of names to datasets for multiple roots, and to enable federation.
47
+
48
+ ```python
49
+ import ibis
50
+ from graphique import GraphQL
51
+
52
+ source = ibis.read_*(...) # or ibis.connect(...).table(...) or pyarrow.dataset.dataset(...)
53
+ # apply initial projections or filters to `source`
54
+ app = GraphQL(source) # Table is root query type
55
+ app = GraphQL.federated({<name>: source, ...}, keys={<name>: [], ...}) # Tables on federated fields
56
+ ```
57
+
58
+ Start like any ASGI app.
59
+
60
+ ```console
61
+ uvicorn <module>:app
62
+ ```
63
+
64
+ ### API
65
+ #### types
66
+ * `Dataset`: interface for an ibis table or arrow dataset.
67
+ * `Table`: implements the `Dataset` interface. Adds typed `row`, `columns`, and `filter` fields from introspecting the schema.
68
+ * `Column`: interface for an ibis column. Each data type has a corresponding column implementation: Boolean, Int, BigInt, Float, Decimal, Date, Datetime, Time, Duration, Base64, String, Array, Struct. All columns have a `values` field for their list of scalars. Additional fields vary by type.
69
+ * `Row`: scalar fields. 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.
70
+
71
+ #### selection
72
+ * `slice`: contiguous selection of rows
73
+ * `filter`: select rows by predicates
74
+ * `join`: join tables by key columns
75
+ * `take`: rows by index
76
+ * `dropNull`: remove rows with nulls
77
+
78
+ #### projection
79
+ * `project`: project columns with expressions
80
+ * `columns`: provides a field for every `Column` in the schema
81
+ * `column`: access a column of any type by name
82
+ * `row`: provides a field for each scalar of a single row
83
+ * `cast`: cast column types
84
+ * `fillNull`: fill null values
85
+
86
+ #### aggregation
87
+ * `group`: group by given columns, and aggregate the others
88
+ * `distinct`: group with all columns
89
+ * `runs`: provisionally group by adjacency
90
+ * `unnest`: unnest an array column
91
+ * `count`: number of rows
92
+
93
+ #### ordering
94
+ * `order`: sort table by given columns
95
+ * options `limit` and `dense`: select rows with smallest or largest values
96
+
97
+ ### Performance
98
+ Performance is dependent on the [ibis backend](https://ibis-project.org/backends/duckdb), which defaults to [duckdb](https://duckdb.org/). There are no internal Python loops. Scalars do not become Python types until serialized.
99
+
100
+ [PyArrow](https://arrow.apache.org/docs/python/) is also used for partitioned dataset optimizations, and for any feature which ibis does not support. Table fields are lazily evaluated up until scalars are reached, and automatically cached as needed for multiple fields.
101
+
102
+ ## Installation
103
+ ```console
104
+ pip install graphique[server]
105
+ ```
106
+
107
+ ## Dependencies
108
+ * ibis-framework (with duckdb or other backend)
109
+ * strawberry-graphql[asgi,cli]
110
+ * pyarrow
111
+ * isodate
112
+ * uvicorn (or other [ASGI server](https://asgi.readthedocs.io/en/latest/implementations.html))
113
+
114
+ ## Tests
115
+ 100% branch coverage.
116
+
117
+ ```console
118
+ pytest [--cov]
119
+ ```
@@ -0,0 +1,3 @@
1
+ # ruff: noqa: F401
2
+ from .interface import Dataset
3
+ from .middleware import GraphQL
@@ -0,0 +1,83 @@
1
+ """
2
+ Core arrow utilities.
3
+ """
4
+
5
+ import itertools
6
+
7
+ import ibis
8
+ import pyarrow as pa
9
+ import pyarrow.dataset as ds
10
+
11
+
12
+ class links:
13
+ ref = "https://ibis-project.org/reference"
14
+ types = f"[data type]({ref}/datatypes)"
15
+ schema = f"[table schema]({ref}/schemas#ibis.expr.schema.Schema)"
16
+
17
+
18
+ def getitems(obj, *keys):
19
+ """Nested `getitem`."""
20
+ for key in keys:
21
+ obj = obj[key]
22
+ return obj
23
+
24
+
25
+ def order_key(name: str) -> ibis.Deferred:
26
+ """Parse sort order."""
27
+ return (ibis.desc if name.startswith('-') else ibis.asc)(ibis._[name.lstrip('-')])
28
+
29
+
30
+ class Parquet(ds.Dataset):
31
+ """Partitioned parquet dataset."""
32
+
33
+ def schema(self) -> pa.Schema:
34
+ """partition schema"""
35
+ return self.partitioning.schema if hasattr(self, 'partitioning') else pa.schema([])
36
+
37
+ def keys(self, *names) -> list:
38
+ """Return prefix of matching partition keys."""
39
+ keys = set(Parquet.schema(self).names)
40
+ return list(itertools.takewhile(lambda name: name.lstrip('-') in keys, names))
41
+
42
+ def fragments(self, counts: str = '') -> ibis.Table:
43
+ """Return partition fragments as a table."""
44
+ parts = []
45
+ for frag in self._get_fragments(self._scan_options.get('filter')):
46
+ parts.append(ds.get_partition_keys(frag.partition_expression))
47
+ parts[-1]['__path__'] = frag.path
48
+ if counts:
49
+ parts[-1][counts] = frag.count_rows()
50
+ return ibis.memtable(pa.Table.from_pylist(parts))
51
+
52
+ def group(self, *names, counts: str = '') -> ibis.Table:
53
+ """Return grouped partitions as a table."""
54
+ table = Parquet.fragments(self, counts)
55
+ agg = {counts: table[counts].sum()} if counts else {}
56
+ return table.aggregate(agg, by=names).order_by(*names)
57
+
58
+ def filter(self, expr: ds.Expression) -> ds.Dataset | None:
59
+ """Attempt to apply filter to partition keys."""
60
+ try: # raises ValueError if filter references non-partition keys
61
+ ds.dataset([], schema=self.partitioning.schema).scanner(filter=expr)
62
+ except (AttributeError, ValueError):
63
+ return None
64
+ return self if expr is None else self.filter(expr)
65
+
66
+ def to_table(self) -> ibis.Table:
67
+ """Return ibis `Table` from filtered dataset."""
68
+ paths = [frag.path for frag in self._get_fragments(self._scan_options.get('filter'))]
69
+ hive = isinstance(self.partitioning, ds.HivePartitioning)
70
+ return ibis.read_parquet(paths, hive_partitioning=hive)
71
+
72
+ def rank(self, limit: int, *names: str, dense: bool = False) -> ibis.Table:
73
+ """Return ordered limited partitions as a table."""
74
+ keys = {name.strip('-'): order_key(name) for name in names}
75
+ table = Parquet.fragments(self, counts='_').order_by(*keys.values()).cache()
76
+ groups = table.aggregate(count=ibis._.count(), total=table['_'].sum(), by=list(keys))
77
+ groups = groups.order_by(*keys.values()).cache()
78
+ if not dense:
79
+ totals = itertools.accumulate(groups['total'].to_list())
80
+ limit = next((index for index, total in enumerate(totals, 1) if total >= limit), None) # type: ignore
81
+ limit = groups[:limit]['count'].sum().to_pyarrow().as_py()
82
+ hive = isinstance(self.partitioning, ds.HivePartitioning)
83
+ return ibis.read_parquet(table[:limit]['__path__'].to_list(), hive_partitioning=hive)