chore: import upstream snapshot with attribution
This commit is contained in:
@@ -0,0 +1,142 @@
|
||||
# ┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓
|
||||
# ┃ ██████ ██████ ██████ █ █ █ █ █ █▄ ▀███ █ ┃
|
||||
# ┃ ▄▄▄▄▄█ █▄▄▄▄▄ ▄▄▄▄▄█ ▀▀▀▀▀█▀▀▀▀▀ █ ▀▀▀▀▀█ ████████▌▐███ ███▄ ▀█ █ ▀▀▀▀▀ ┃
|
||||
# ┃ █▀▀▀▀▀ █▀▀▀▀▀ █▀██▀▀ ▄▄▄▄▄ █ ▄▄▄▄▄█ ▄▄▄▄▄█ ████████▌▐███ █████▄ █ ▄▄▄▄▄ ┃
|
||||
# ┃ █ ██████ █ ▀█▄ █ ██████ █ ███▌▐███ ███████▄ █ ┃
|
||||
# ┣━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┫
|
||||
# ┃ Copyright (c) 2017, the Perspective Authors. ┃
|
||||
# ┃ ╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌ ┃
|
||||
# ┃ This file is part of the Perspective library, distributed under the terms ┃
|
||||
# ┃ of the [Apache License 2.0](https://www.apache.org/licenses/LICENSE-2.0). ┃
|
||||
# ┗━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┛
|
||||
|
||||
|
||||
class VirtualServerHandler:
|
||||
"""
|
||||
An interface for implementing a Perspective `VirtualServer`. It operates
|
||||
thusly:
|
||||
|
||||
- A table is selected by name (validated via `get_hosted_tables`).
|
||||
|
||||
- The UI will ask the model to create a temporary table with the results
|
||||
of querying this table with a specific query `config`, a simple struct
|
||||
which reflects the UI configurable fields (see `get_features`).
|
||||
|
||||
- The UI will query slices of the temporary table as it needs them to
|
||||
render. This may be a rectangular slice, a whole column or the entire
|
||||
set, and it is returned from teh model via a custom push-only
|
||||
struct `PerspectiveColumn` for now, though in the future we will support
|
||||
e.g. Polars and other arrow-native formats directly.
|
||||
|
||||
- The UI will delete its own temporary tables via `view_delete` but it is
|
||||
ok for them to die intermittently, the UI will recover automatically.
|
||||
"""
|
||||
|
||||
def get_features(self):
|
||||
"""
|
||||
[OPTIONAL] Toggle UI features through data model support. For example,
|
||||
setting `"group_by": False` would hide the "Group By" UI control, as
|
||||
well as prevent this field from appearing in `config` dicts later
|
||||
provided to `table_make_view`.
|
||||
|
||||
This API defaults to just "columns", e.g. a simple flat datagrid in
|
||||
which you can just scroll, select and format columns.
|
||||
|
||||
# Example
|
||||
|
||||
```python
|
||||
return {
|
||||
"group_by": True,
|
||||
"split_by": True,
|
||||
"sort": True,
|
||||
"expressions": True,
|
||||
"filter_ops": {
|
||||
"integer": ["==", "<"],
|
||||
},
|
||||
"aggregates": {
|
||||
"string": ["count"],
|
||||
"float": ["count", "sum"],
|
||||
},
|
||||
}
|
||||
```
|
||||
"""
|
||||
|
||||
pass
|
||||
|
||||
def get_hosted_tables(self) -> list[str]:
|
||||
"""
|
||||
List of `Table` names available to query from.
|
||||
"""
|
||||
|
||||
pass
|
||||
|
||||
def table_schema(self, table_name):
|
||||
"""
|
||||
Get the _Perspective Schema_ for a `Table`, a mapping of column name to
|
||||
Perspective column types, a simplified set of six visually-relevant
|
||||
types mapped from DuckDB's much richer type system. Optionally,
|
||||
a model may also implement `view_schema` which describes temporary
|
||||
tables, but for DuckDB this method is identical.
|
||||
"""
|
||||
|
||||
pass
|
||||
|
||||
def table_size(self, table_name):
|
||||
"""
|
||||
Get a table's row count. Optionally, a model may also implement the
|
||||
`view_size` method to get the row count for temporary tables, but for
|
||||
DuckDB this method is identical.
|
||||
"""
|
||||
|
||||
pass
|
||||
|
||||
def view_schema(self, view_name, config):
|
||||
return self.table_schema(view_name)
|
||||
|
||||
def view_size(self, view_name):
|
||||
return self.table_size(view_name)
|
||||
|
||||
def table_make_view(self, table_name, view_name, config):
|
||||
"""
|
||||
Create a temporary table `view_name` from the results of querying
|
||||
`table_name` with a query configuration `config`.
|
||||
"""
|
||||
|
||||
pass
|
||||
|
||||
def table_validate_expression(self, view_name, expression):
|
||||
"""
|
||||
[OPTIONAL] Given a temporary table `view_name`, validate the type of
|
||||
a column expression string `expression`, or raise an error if the
|
||||
expression is invalid. This is enabeld by `"expressions"` via
|
||||
`get_features` and defaults to allow all expressions.
|
||||
"""
|
||||
|
||||
pass
|
||||
|
||||
def view_delete(self, view_name):
|
||||
"""
|
||||
Delete a temporary table. The UI will do this automatically, and it
|
||||
can recover.
|
||||
"""
|
||||
|
||||
pass
|
||||
|
||||
def view_get_min_max(self, view_name, column_name, config):
|
||||
"""
|
||||
[OPTIONAL] Get the min and max values of a column in a view.
|
||||
Returns a tuple of (min, max) as native Python values.
|
||||
"""
|
||||
|
||||
pass
|
||||
|
||||
def view_get_data(self, view_name, config, viewport, data):
|
||||
"""
|
||||
Serialize a rectangular slice `viewport` from temporary table
|
||||
`view_name`, into the `PerspectiveColumn` serialization API injected
|
||||
via `data`. The push-only `PerspectiveColumn` type can handle casting
|
||||
Python types as input, but once a type is pushed to a column name it
|
||||
must not be changed.
|
||||
"""
|
||||
|
||||
pass
|
||||
@@ -0,0 +1,240 @@
|
||||
# ┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓
|
||||
# ┃ ██████ ██████ ██████ █ █ █ █ █ █▄ ▀███ █ ┃
|
||||
# ┃ ▄▄▄▄▄█ █▄▄▄▄▄ ▄▄▄▄▄█ ▀▀▀▀▀█▀▀▀▀▀ █ ▀▀▀▀▀█ ████████▌▐███ ███▄ ▀█ █ ▀▀▀▀▀ ┃
|
||||
# ┃ █▀▀▀▀▀ █▀▀▀▀▀ █▀██▀▀ ▄▄▄▄▄ █ ▄▄▄▄▄█ ▄▄▄▄▄█ ████████▌▐███ █████▄ █ ▄▄▄▄▄ ┃
|
||||
# ┃ █ ██████ █ ▀█▄ █ ██████ █ ███▌▐███ ███████▄ █ ┃
|
||||
# ┣━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┫
|
||||
# ┃ Copyright (c) 2017, the Perspective Authors. ┃
|
||||
# ┃ ╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌ ┃
|
||||
# ┃ This file is part of the Perspective library, distributed under the terms ┃
|
||||
# ┃ of the [Apache License 2.0](https://www.apache.org/licenses/LICENSE-2.0). ┃
|
||||
# ┗━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┛
|
||||
|
||||
import perspective
|
||||
|
||||
from datetime import datetime
|
||||
import logging
|
||||
|
||||
from perspective.virtual_servers import VirtualServerHandler
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
NUMBER_AGGS = [
|
||||
"sum",
|
||||
"count",
|
||||
"any_value",
|
||||
"arbitrary",
|
||||
"array_agg",
|
||||
"avg",
|
||||
"bit_and",
|
||||
"bit_or",
|
||||
"bit_xor",
|
||||
"bitstring_agg",
|
||||
"bool_and",
|
||||
"bool_or",
|
||||
"countif",
|
||||
"favg",
|
||||
"fsum",
|
||||
"geomean",
|
||||
"kahan_sum",
|
||||
"last",
|
||||
"max",
|
||||
"min",
|
||||
"product",
|
||||
"string_agg",
|
||||
"sumkahan",
|
||||
]
|
||||
|
||||
STRING_AGGS = [
|
||||
"count",
|
||||
"any_value",
|
||||
"arbitrary",
|
||||
"first",
|
||||
"countif",
|
||||
"last",
|
||||
"string_agg",
|
||||
]
|
||||
|
||||
FILTER_OPS = [
|
||||
"==",
|
||||
"!=",
|
||||
"LIKE",
|
||||
"IS DISTINCT FROM",
|
||||
"IS NOT DISTINCT FROM",
|
||||
">=",
|
||||
"<=",
|
||||
">",
|
||||
"<",
|
||||
]
|
||||
|
||||
|
||||
class ClickhouseVirtualSession:
|
||||
def __init__(self, callback, db):
|
||||
self.session = perspective.VirtualServer(ClickhouseVirtualServerHandler(db))
|
||||
self.callback = callback
|
||||
|
||||
def handle_request(self, msg):
|
||||
self.callback(self.session.handle_request(msg))
|
||||
|
||||
|
||||
class ClickhouseVirtualServer:
|
||||
def __init__(self, db):
|
||||
self.db = db
|
||||
|
||||
def new_session(self, callback):
|
||||
return ClickhouseVirtualSession(callback, self.db)
|
||||
|
||||
|
||||
class ClickhouseVirtualServerHandler(VirtualServerHandler):
|
||||
"""
|
||||
An implementation of a `perspective.VirtualServerHandler` for ClickHouse.
|
||||
"""
|
||||
|
||||
def __init__(self, db):
|
||||
self.db = db
|
||||
self.sql_builder = perspective.GenericSQLVirtualServerModel(
|
||||
{"create_entity": "VIEW", "grouping_fn": "GROUPING"}
|
||||
)
|
||||
|
||||
def get_features(self):
|
||||
return {
|
||||
"group_by": True,
|
||||
"split_by": False,
|
||||
"sort": True,
|
||||
"expressions": True,
|
||||
"group_rollup_mode": ["rollup", "flat", "total"],
|
||||
"filter_ops": {
|
||||
"integer": FILTER_OPS,
|
||||
"float": FILTER_OPS,
|
||||
"string": FILTER_OPS,
|
||||
"boolean": FILTER_OPS,
|
||||
"date": FILTER_OPS,
|
||||
"datetime": FILTER_OPS,
|
||||
},
|
||||
"aggregates": {
|
||||
"integer": NUMBER_AGGS,
|
||||
"float": NUMBER_AGGS,
|
||||
"string": STRING_AGGS,
|
||||
"boolean": STRING_AGGS,
|
||||
"date": STRING_AGGS,
|
||||
"datetime": STRING_AGGS,
|
||||
},
|
||||
}
|
||||
|
||||
def get_hosted_tables(self):
|
||||
query = "SHOW TABLES"
|
||||
results = run_query(self.db, query)
|
||||
return [result[0] for result in results]
|
||||
|
||||
def table_schema(self, table_name, config=None):
|
||||
query = self.sql_builder.table_schema(table_name)
|
||||
results = run_query(self.db, query)
|
||||
schema = {}
|
||||
for result in results:
|
||||
col_name = result[0]
|
||||
if not col_name.startswith("__"):
|
||||
schema[col_name] = clickhouse_type_to_psp(result[1])
|
||||
|
||||
return schema
|
||||
|
||||
def view_column_size(self, view_name, config):
|
||||
query = f"SELECT COUNT() FROM system.columns WHERE table = '{view_name}'"
|
||||
results = run_query(self.db, query)
|
||||
gs = len(config["group_by"])
|
||||
return results[0][0] - (
|
||||
0 if gs == 0 else gs + (1 if len(config["split_by"]) == 0 else 0)
|
||||
)
|
||||
|
||||
def table_size(self, table_name):
|
||||
query = self.sql_builder.table_size(table_name)
|
||||
results = run_query(self.db, query)
|
||||
return results[0][0]
|
||||
|
||||
def table_make_view(self, table_name, view_name, config):
|
||||
query = self.sql_builder.table_make_view(table_name, view_name, config)
|
||||
run_query(self.db, query, execute=True)
|
||||
|
||||
def table_validate_expression(self, view_name, expression):
|
||||
query = self.sql_builder.table_validate_expression(view_name, expression)
|
||||
results = run_query(self.db, query)
|
||||
return clickhouse_type_to_psp(results[0][1])
|
||||
|
||||
def view_delete(self, view_name):
|
||||
query = self.sql_builder.view_delete(view_name)
|
||||
run_query(self.db, query, execute=True)
|
||||
|
||||
def view_get_min_max(self, view_name, column_name, config):
|
||||
query = self.sql_builder.view_get_min_max(view_name, column_name, config)
|
||||
results = run_query(self.db, query)
|
||||
row = results[0]
|
||||
return (row[0], row[1])
|
||||
|
||||
def view_get_data(self, view_name, config, schema, viewport, data):
|
||||
group_by = config["group_by"]
|
||||
query = self.sql_builder.view_get_data(view_name, config, viewport, schema)
|
||||
results, columns, dtypes = run_query(self.db, query, columns=True)
|
||||
for cidx, col in enumerate(columns):
|
||||
dtype = clickhouse_type_to_psp(str(dtypes[cidx]))
|
||||
for ridx, row in enumerate(results):
|
||||
grouping_id = row[0] if len(group_by) > 0 else None
|
||||
|
||||
value = row[cidx]
|
||||
if dtype == "string" and not isinstance(value, str):
|
||||
value = str(value)
|
||||
|
||||
data.set_col(dtype, col, ridx, value, grouping_id)
|
||||
|
||||
|
||||
################################################################################
|
||||
#
|
||||
# ClickHouse Utils
|
||||
|
||||
|
||||
def clickhouse_type_to_psp(name):
|
||||
"""Convert a ClickHouse `dtype` to a Perspective `ColumnType`."""
|
||||
if name.startswith("Nullable(") and name.endswith(")"):
|
||||
name = name[9:-1]
|
||||
|
||||
if name.startswith("Array"):
|
||||
return "string"
|
||||
|
||||
if name in ("Int64", "UInt64", "Float64"):
|
||||
return "float"
|
||||
|
||||
if name == "String":
|
||||
return "string"
|
||||
|
||||
if name == "DateTime":
|
||||
return "datetime"
|
||||
|
||||
if name == "Date":
|
||||
return "date"
|
||||
|
||||
msg = f"Unknown type '{name}'"
|
||||
raise ValueError(msg)
|
||||
|
||||
|
||||
def run_query(db, query, execute=False, columns=False):
|
||||
query = " ".join(query.split())
|
||||
start = datetime.now()
|
||||
result = None
|
||||
try:
|
||||
if execute:
|
||||
db.command(query)
|
||||
else:
|
||||
req = db.query(query)
|
||||
result = req.result_rows
|
||||
except Exception as e:
|
||||
logger.error(e)
|
||||
logger.error(f"{query}")
|
||||
raise e
|
||||
else:
|
||||
logger.debug(f"{datetime.now() - start} {query}")
|
||||
if columns:
|
||||
return (
|
||||
result,
|
||||
req.column_names,
|
||||
[(x.name if hasattr(x, "name") else str(x)) for x in req.column_types],
|
||||
)
|
||||
else:
|
||||
return result
|
||||
@@ -0,0 +1,236 @@
|
||||
# ┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓
|
||||
# ┃ ██████ ██████ ██████ █ █ █ █ █ █▄ ▀███ █ ┃
|
||||
# ┃ ▄▄▄▄▄█ █▄▄▄▄▄ ▄▄▄▄▄█ ▀▀▀▀▀█▀▀▀▀▀ █ ▀▀▀▀▀█ ████████▌▐███ ███▄ ▀█ █ ▀▀▀▀▀ ┃
|
||||
# ┃ █▀▀▀▀▀ █▀▀▀▀▀ █▀██▀▀ ▄▄▄▄▄ █ ▄▄▄▄▄█ ▄▄▄▄▄█ ████████▌▐███ █████▄ █ ▄▄▄▄▄ ┃
|
||||
# ┃ █ ██████ █ ▀█▄ █ ██████ █ ███▌▐███ ███████▄ █ ┃
|
||||
# ┣━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┫
|
||||
# ┃ Copyright (c) 2017, the Perspective Authors. ┃
|
||||
# ┃ ╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌ ┃
|
||||
# ┃ This file is part of the Perspective library, distributed under the terms ┃
|
||||
# ┃ of the [Apache License 2.0](https://www.apache.org/licenses/LICENSE-2.0). ┃
|
||||
# ┗━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┛
|
||||
|
||||
import io
|
||||
import duckdb
|
||||
import perspective
|
||||
import pyarrow.ipc as ipc
|
||||
|
||||
from datetime import datetime
|
||||
import logging
|
||||
|
||||
from perspective.virtual_servers import VirtualServerHandler
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
NUMBER_AGGS = [
|
||||
"sum",
|
||||
"count",
|
||||
"any_value",
|
||||
"arbitrary",
|
||||
# "arg_max",
|
||||
# "arg_max_null",
|
||||
# "arg_min",
|
||||
# "arg_min_null",
|
||||
"array_agg",
|
||||
"avg",
|
||||
"bit_and",
|
||||
"bit_or",
|
||||
"bit_xor",
|
||||
"bitstring_agg",
|
||||
"bool_and",
|
||||
"bool_or",
|
||||
"countif",
|
||||
"favg",
|
||||
"fsum",
|
||||
"geomean",
|
||||
# "histogram",
|
||||
# "histogram_values",
|
||||
"kahan_sum",
|
||||
"last",
|
||||
# "list"
|
||||
"max",
|
||||
# "max_by"
|
||||
"min",
|
||||
# "min_by"
|
||||
"product",
|
||||
"string_agg",
|
||||
"sumkahan",
|
||||
# "weighted_avg",
|
||||
]
|
||||
|
||||
STRING_AGGS = [
|
||||
"count",
|
||||
"any_value",
|
||||
"arbitrary",
|
||||
"first",
|
||||
"countif",
|
||||
"last",
|
||||
"string_agg",
|
||||
]
|
||||
|
||||
FILTER_OPS = [
|
||||
"==",
|
||||
"!=",
|
||||
"LIKE",
|
||||
"IS DISTINCT FROM",
|
||||
"IS NOT DISTINCT FROM",
|
||||
">=",
|
||||
"<=",
|
||||
">",
|
||||
"<",
|
||||
]
|
||||
|
||||
|
||||
class DuckDBVirtualSession:
|
||||
def __init__(self, callback, db):
|
||||
self.session = perspective.VirtualServer(DuckDBVirtualServerHandler(db))
|
||||
self.callback = callback
|
||||
|
||||
def handle_request(self, msg):
|
||||
self.callback(self.session.handle_request(msg))
|
||||
|
||||
|
||||
class DuckDBVirtualServer:
|
||||
def __init__(self, db):
|
||||
self.db = db
|
||||
|
||||
def new_session(self, callback):
|
||||
return DuckDBVirtualSession(callback, self.db)
|
||||
|
||||
|
||||
class DuckDBVirtualServerHandler(VirtualServerHandler):
|
||||
"""
|
||||
An implementation of a `perspective.VirtualServerHandler` for DuckDB.
|
||||
"""
|
||||
|
||||
def __init__(self, db):
|
||||
self.db = db
|
||||
self.sql_builder = perspective.GenericSQLVirtualServerModel({})
|
||||
|
||||
def get_features(self):
|
||||
return {
|
||||
"group_by": True,
|
||||
"split_by": True,
|
||||
"sort": True,
|
||||
"expressions": True,
|
||||
"group_rollup_mode": ["rollup", "flat", "total"],
|
||||
"filter_ops": {
|
||||
"integer": FILTER_OPS,
|
||||
"float": FILTER_OPS,
|
||||
"string": FILTER_OPS,
|
||||
"boolean": FILTER_OPS,
|
||||
"date": FILTER_OPS,
|
||||
"datetime": FILTER_OPS,
|
||||
},
|
||||
"aggregates": {
|
||||
"integer": NUMBER_AGGS,
|
||||
"float": NUMBER_AGGS,
|
||||
"string": STRING_AGGS,
|
||||
"boolean": STRING_AGGS,
|
||||
"date": STRING_AGGS,
|
||||
"datetime": STRING_AGGS,
|
||||
},
|
||||
}
|
||||
|
||||
def get_hosted_tables(self):
|
||||
query = self.sql_builder.get_hosted_tables()
|
||||
results = run_query(self.db, query)
|
||||
return [f"{result[0]}.{result[2]}" for result in results]
|
||||
|
||||
def table_schema(self, table_name, config=None):
|
||||
query = self.sql_builder.table_schema(table_name)
|
||||
results = run_query(self.db, query)
|
||||
schema = {}
|
||||
for result in results:
|
||||
col_name = result[0]
|
||||
if not col_name.startswith("__"):
|
||||
schema[col_name] = duckdb_type_to_psp(result[1])
|
||||
|
||||
return schema
|
||||
|
||||
def view_column_size(self, table_name, config):
|
||||
query = self.sql_builder.view_column_size(table_name)
|
||||
results = run_query(self.db, query)
|
||||
gs = len(config["group_by"])
|
||||
return results[0][0] - (
|
||||
0 if gs == 0 else gs + (1 if len(config["split_by"]) == 0 else 0)
|
||||
)
|
||||
|
||||
def table_size(self, table_name):
|
||||
query = self.sql_builder.table_size(table_name)
|
||||
results = run_query(self.db, query)
|
||||
return results[0][0]
|
||||
|
||||
def table_make_view(self, table_name, view_name, config):
|
||||
query = self.sql_builder.table_make_view(table_name, view_name, config)
|
||||
run_query(self.db, query, execute=True)
|
||||
|
||||
def table_validate_expression(self, view_name, expression):
|
||||
query = self.sql_builder.table_validate_expression(view_name, expression)
|
||||
results = run_query(self.db, query)
|
||||
return duckdb_type_to_psp(results[0][1])
|
||||
|
||||
def view_delete(self, view_name):
|
||||
query = self.sql_builder.view_delete(view_name)
|
||||
run_query(self.db, query, execute=True)
|
||||
|
||||
def view_get_min_max(self, view_name, column_name, config):
|
||||
query = self.sql_builder.view_get_min_max(view_name, column_name, config)
|
||||
results = run_query(self.db, query)
|
||||
row = results[0]
|
||||
return (row[0], row[1])
|
||||
|
||||
def view_get_data(self, view_name, config, schema, viewport, data):
|
||||
query = self.sql_builder.view_get_data(view_name, config, viewport, schema)
|
||||
result = self.db.sql(query)
|
||||
arrow_table = result.fetch_arrow_table()
|
||||
buf = io.BytesIO()
|
||||
with ipc.new_stream(buf, arrow_table.schema) as writer:
|
||||
writer.write_table(arrow_table)
|
||||
data.from_arrow_ipc(buf.getvalue())
|
||||
|
||||
|
||||
################################################################################
|
||||
#
|
||||
# DuckDB Utils
|
||||
|
||||
|
||||
def duckdb_type_to_psp(name):
|
||||
"""Convert a DuckDB `dtype` to a Perspective `ColumnType`."""
|
||||
if name == "VARCHAR":
|
||||
return "string"
|
||||
if name in ("DOUBLE", "BIGINT", "HUGEINT"):
|
||||
return "float"
|
||||
if name == "INTEGER":
|
||||
return "integer"
|
||||
if name == "DATE":
|
||||
return "date"
|
||||
if name == "BOOLEAN":
|
||||
return "boolean"
|
||||
if name == "TIMESTAMP":
|
||||
return "datetime"
|
||||
|
||||
msg = f"Unknown type '{name}'"
|
||||
raise ValueError(msg)
|
||||
|
||||
|
||||
def run_query(db, query, execute=False, columns=False):
|
||||
query = " ".join(query.split())
|
||||
start = datetime.now()
|
||||
result = None
|
||||
try:
|
||||
if execute:
|
||||
db.execute(query)
|
||||
else:
|
||||
req = db.sql(query)
|
||||
result = req.fetchall()
|
||||
except (duckdb.ParserException, duckdb.BinderException) as e:
|
||||
logger.error(e)
|
||||
logger.error(f"{query}")
|
||||
raise e
|
||||
else:
|
||||
logger.debug(f"{datetime.now() - start} {query}")
|
||||
if columns:
|
||||
return (result, req.columns, req.dtypes)
|
||||
else:
|
||||
return result
|
||||
@@ -0,0 +1,710 @@
|
||||
# ┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓
|
||||
# ┃ ██████ ██████ ██████ █ █ █ █ █ █▄ ▀███ █ ┃
|
||||
# ┃ ▄▄▄▄▄█ █▄▄▄▄▄ ▄▄▄▄▄█ ▀▀▀▀▀█▀▀▀▀▀ █ ▀▀▀▀▀█ ████████▌▐███ ███▄ ▀█ █ ▀▀▀▀▀ ┃
|
||||
# ┃ █▀▀▀▀▀ █▀▀▀▀▀ █▀██▀▀ ▄▄▄▄▄ █ ▄▄▄▄▄█ ▄▄▄▄▄█ ████████▌▐███ █████▄ █ ▄▄▄▄▄ ┃
|
||||
# ┃ █ ██████ █ ▀█▄ █ ██████ █ ███▌▐███ ███████▄ █ ┃
|
||||
# ┣━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┫
|
||||
# ┃ Copyright (c) 2017, the Perspective Authors. ┃
|
||||
# ┃ ╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌ ┃
|
||||
# ┃ This file is part of the Perspective library, distributed under the terms ┃
|
||||
# ┃ of the [Apache License 2.0](https://www.apache.org/licenses/LICENSE-2.0). ┃
|
||||
# ┗━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┛
|
||||
|
||||
import logging
|
||||
import polars as pl
|
||||
import perspective
|
||||
|
||||
from datetime import datetime
|
||||
import re
|
||||
|
||||
from perspective.virtual_servers import VirtualServerHandler
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
NUMBER_AGGS = [
|
||||
"sum",
|
||||
"count",
|
||||
"any_value",
|
||||
"avg",
|
||||
"mean",
|
||||
"max",
|
||||
"min",
|
||||
"first",
|
||||
"last",
|
||||
]
|
||||
|
||||
STRING_AGGS = [
|
||||
"count",
|
||||
"any_value",
|
||||
"first",
|
||||
"last",
|
||||
]
|
||||
|
||||
FILTER_OPS = [
|
||||
"==",
|
||||
"!=",
|
||||
">=",
|
||||
"<=",
|
||||
">",
|
||||
"<",
|
||||
]
|
||||
|
||||
AGG_MAP = {
|
||||
"sum": lambda e: e.sum(),
|
||||
"count": lambda e: e.count(),
|
||||
"avg": lambda e: e.mean(),
|
||||
"mean": lambda e: e.mean(),
|
||||
"min": lambda e: e.min(),
|
||||
"max": lambda e: e.max(),
|
||||
"first": lambda e: e.first(),
|
||||
"last": lambda e: e.last(),
|
||||
"any_value": lambda e: e.first(),
|
||||
"arbitrary": lambda e: e.first(),
|
||||
}
|
||||
|
||||
|
||||
class PolarsVirtualSession:
|
||||
def __init__(self, callback, tables):
|
||||
self.session = perspective.VirtualServer(PolarsVirtualServerHandler(tables))
|
||||
self.callback = callback
|
||||
|
||||
def handle_request(self, msg):
|
||||
self.callback(self.session.handle_request(msg))
|
||||
|
||||
|
||||
class PolarsVirtualServer:
|
||||
def __init__(self, tables):
|
||||
self.tables = tables
|
||||
|
||||
def new_session(self, callback):
|
||||
return PolarsVirtualSession(callback, self.tables)
|
||||
|
||||
|
||||
class PolarsVirtualServerHandler(VirtualServerHandler):
|
||||
"""
|
||||
An implementation of a `perspective.VirtualServerHandler` for Polars.
|
||||
"""
|
||||
|
||||
def __init__(self, tables):
|
||||
self.tables = tables
|
||||
self.views = {}
|
||||
self.view_schemas = {}
|
||||
|
||||
def get_features(self):
|
||||
return {
|
||||
"group_by": True,
|
||||
"split_by": True,
|
||||
"sort": True,
|
||||
"expressions": True,
|
||||
"group_rollup_mode": ["rollup", "flat", "total"],
|
||||
"filter_ops": {
|
||||
"integer": FILTER_OPS,
|
||||
"float": FILTER_OPS,
|
||||
"string": FILTER_OPS,
|
||||
"boolean": ["==", "!="],
|
||||
"date": FILTER_OPS,
|
||||
"datetime": FILTER_OPS,
|
||||
},
|
||||
"aggregates": {
|
||||
"integer": NUMBER_AGGS,
|
||||
"float": NUMBER_AGGS,
|
||||
"string": STRING_AGGS,
|
||||
"boolean": STRING_AGGS,
|
||||
"date": STRING_AGGS,
|
||||
"datetime": STRING_AGGS,
|
||||
},
|
||||
}
|
||||
|
||||
def get_hosted_tables(self):
|
||||
return list(self.tables.keys())
|
||||
|
||||
def table_schema(self, table_name, config=None):
|
||||
df = self.tables[table_name]
|
||||
schema = {}
|
||||
for col_name, dtype in df.schema.items():
|
||||
if not col_name.startswith("__"):
|
||||
schema[col_name] = polars_type_to_psp(dtype)
|
||||
return schema
|
||||
|
||||
def table_size(self, table_name):
|
||||
return self.tables[table_name].height
|
||||
|
||||
def view_schema(self, view_name, config):
|
||||
if view_name in self.view_schemas:
|
||||
return self.view_schemas[view_name]
|
||||
return self.table_schema(view_name)
|
||||
|
||||
def view_size(self, view_name):
|
||||
if view_name in self.views:
|
||||
return self.views[view_name].height
|
||||
return self.table_size(view_name)
|
||||
|
||||
def table_validate_expression(self, table_name, expression):
|
||||
df = self.tables.get(table_name)
|
||||
if df is None:
|
||||
return None
|
||||
expr = parse_expression(expression)
|
||||
result = df.select(expr.alias("__expr__"))
|
||||
return polars_type_to_psp(result["__expr__"].dtype)
|
||||
|
||||
def table_make_view(self, table_name, view_name, config):
|
||||
start = datetime.now()
|
||||
df = self.tables[table_name]
|
||||
group_by = config.get("group_by", [])
|
||||
columns = [c for c in config.get("columns", []) if c is not None]
|
||||
aggregates = config.get("aggregates", {})
|
||||
sort = config.get("sort", [])
|
||||
filters = config.get("filter", [])
|
||||
split_by = config.get("split_by", [])
|
||||
expressions = config.get("expressions", {})
|
||||
group_rollup_mode = config.get("group_rollup_mode", "rollup")
|
||||
|
||||
if expressions:
|
||||
for expr_name, expr_str in expressions.items():
|
||||
expr = parse_expression(expr_str)
|
||||
df = df.with_columns(expr.alias(expr_name))
|
||||
|
||||
df = apply_filters(df, filters)
|
||||
|
||||
col_alias = lambda c: c.replace("_", "-")
|
||||
is_flat = group_rollup_mode == "flat"
|
||||
is_total = group_rollup_mode == "total"
|
||||
|
||||
if is_total:
|
||||
if split_by:
|
||||
result = build_split_by_total(
|
||||
df, split_by, columns, aggregates, col_alias
|
||||
)
|
||||
else:
|
||||
result = build_total(df, columns, aggregates, col_alias)
|
||||
elif split_by and group_by:
|
||||
if is_flat:
|
||||
result = build_split_by_grouped_flat(
|
||||
df, group_by, split_by, columns, aggregates, col_alias
|
||||
)
|
||||
result = apply_sort_split_by_flat(
|
||||
result, sort, columns, group_by, split_by
|
||||
)
|
||||
else:
|
||||
result = build_split_by_grouped(
|
||||
df, group_by, split_by, columns, aggregates, col_alias
|
||||
)
|
||||
result = apply_sort_grouped(result, sort, group_by, col_alias)
|
||||
elif split_by:
|
||||
result = build_split_by_flat(df, split_by, columns, col_alias)
|
||||
result = apply_sort_flat(result, sort, col_alias)
|
||||
elif group_by:
|
||||
if is_flat:
|
||||
result = build_flat_group_by(
|
||||
df, group_by, columns, aggregates, col_alias
|
||||
)
|
||||
result = apply_sort_flat(result, sort, col_alias)
|
||||
else:
|
||||
result = build_rollup(df, group_by, columns, aggregates, col_alias)
|
||||
result = apply_sort_grouped(result, sort, group_by, col_alias)
|
||||
else:
|
||||
select_exprs = [pl.col(c).alias(col_alias(c)) for c in columns]
|
||||
result = df.select(select_exprs)
|
||||
result = apply_sort_flat(result, sort, col_alias)
|
||||
|
||||
self.views[view_name] = result
|
||||
self.view_schemas[view_name] = compute_view_schema(result)
|
||||
logger.debug(
|
||||
f"{datetime.now() - start} table_make_view {table_name} -> {view_name}"
|
||||
)
|
||||
|
||||
def view_delete(self, view_name):
|
||||
self.views.pop(view_name, None)
|
||||
self.view_schemas.pop(view_name, None)
|
||||
|
||||
def view_get_min_max(self, view_name, column_name, config):
|
||||
df = self.views[view_name]
|
||||
col = df[column_name]
|
||||
min_val = col.min()
|
||||
max_val = col.max()
|
||||
return (min_val, max_val)
|
||||
|
||||
def view_get_data(self, view_name, config, schema, viewport, data):
|
||||
df = self.views.get(view_name)
|
||||
if df is None:
|
||||
return
|
||||
|
||||
group_by = config.get("group_by", [])
|
||||
split_by = config.get("split_by", [])
|
||||
group_rollup_mode = config.get("group_rollup_mode", "rollup")
|
||||
is_split_by = len(split_by) > 0
|
||||
is_flat = group_rollup_mode == "flat"
|
||||
|
||||
start_row = viewport.get("start_row", 0) or 0
|
||||
end_row = viewport.get("end_row") or df.height
|
||||
start_col = viewport.get("start_col", 0) or 0
|
||||
end_col = viewport.get("end_col")
|
||||
|
||||
length = min(end_row, df.height) - start_row
|
||||
if length <= 0:
|
||||
return
|
||||
df_slice = df.slice(start_row, length)
|
||||
|
||||
data_columns = [c for c in schema.keys() if not c.startswith("__")]
|
||||
if end_col is not None:
|
||||
data_columns = data_columns[start_col:end_col]
|
||||
else:
|
||||
data_columns = data_columns[start_col:]
|
||||
|
||||
has_group_by = len(group_by) > 0
|
||||
has_grouping_id = has_group_by and not is_flat
|
||||
|
||||
all_cols = []
|
||||
if has_grouping_id:
|
||||
all_cols.append("__GROUPING_ID__")
|
||||
for idx in range(len(group_by)):
|
||||
all_cols.append(f"__ROW_PATH_{idx}__")
|
||||
all_cols.extend(data_columns)
|
||||
|
||||
grouping_ids = None
|
||||
if has_grouping_id:
|
||||
grouping_ids = df_slice["__GROUPING_ID__"].to_list()
|
||||
|
||||
for cidx, col in enumerate(all_cols):
|
||||
if cidx == 0 and has_grouping_id:
|
||||
continue
|
||||
|
||||
series = df_slice[col]
|
||||
dtype = polars_type_to_psp(series.dtype)
|
||||
values = series.to_list()
|
||||
|
||||
push_col = col
|
||||
if is_split_by and not col.startswith("__"):
|
||||
push_col = col.replace("_", "|")
|
||||
|
||||
for ridx, value in enumerate(values):
|
||||
if grouping_ids:
|
||||
grouping_id = grouping_ids[ridx]
|
||||
elif has_group_by:
|
||||
grouping_id = 0
|
||||
else:
|
||||
grouping_id = None
|
||||
|
||||
if value is not None and isinstance(value, float) and value != value:
|
||||
value = None
|
||||
|
||||
data.set_col(dtype, push_col, ridx, value, grouping_id)
|
||||
|
||||
|
||||
################################################################################
|
||||
#
|
||||
# Polars Utils
|
||||
|
||||
|
||||
def polars_type_to_psp(dtype):
|
||||
"""Convert a Polars `dtype` to a Perspective `ColumnType`."""
|
||||
if dtype in (pl.Utf8, pl.String):
|
||||
return "string"
|
||||
if dtype == pl.Categorical:
|
||||
return "string"
|
||||
if dtype in (pl.Int8, pl.Int16, pl.Int32, pl.UInt8, pl.UInt16):
|
||||
return "integer"
|
||||
if dtype in (pl.Int64, pl.UInt64, pl.UInt32, pl.Float32, pl.Float64):
|
||||
return "float"
|
||||
if dtype == pl.Date:
|
||||
return "date"
|
||||
if dtype == pl.Boolean:
|
||||
return "boolean"
|
||||
if isinstance(dtype, pl.Datetime) or dtype == pl.Datetime:
|
||||
return "datetime"
|
||||
|
||||
msg = f"Unknown Polars type '{dtype}'"
|
||||
raise ValueError(msg)
|
||||
|
||||
|
||||
def apply_filters(df, filters):
|
||||
"""Apply a list of filter configs to a DataFrame."""
|
||||
if not filters:
|
||||
return df
|
||||
|
||||
mask = pl.lit(True)
|
||||
for filt in filters:
|
||||
col_name = filt[0]
|
||||
op = filt[1]
|
||||
value = filt[2] if len(filt) > 2 else None
|
||||
|
||||
if value is None:
|
||||
continue
|
||||
|
||||
col_expr = pl.col(col_name)
|
||||
if op == "==":
|
||||
mask = mask & (col_expr == value)
|
||||
elif op == "!=":
|
||||
mask = mask & (col_expr != value)
|
||||
elif op == ">":
|
||||
mask = mask & (col_expr > value)
|
||||
elif op == "<":
|
||||
mask = mask & (col_expr < value)
|
||||
elif op == ">=":
|
||||
mask = mask & (col_expr >= value)
|
||||
elif op == "<=":
|
||||
mask = mask & (col_expr <= value)
|
||||
|
||||
return df.filter(mask)
|
||||
|
||||
|
||||
def get_polars_agg_expr(col, agg_name, filter_expr=None):
|
||||
"""Convert an aggregate name to a Polars expression."""
|
||||
if isinstance(agg_name, list):
|
||||
agg_name = agg_name[0]
|
||||
if isinstance(agg_name, dict):
|
||||
agg_name = "first"
|
||||
expr = pl.col(col)
|
||||
if filter_expr is not None:
|
||||
expr = expr.filter(filter_expr)
|
||||
if agg_name in AGG_MAP:
|
||||
return AGG_MAP[agg_name](expr)
|
||||
|
||||
msg = f"Unknown aggregate '{agg_name}'"
|
||||
raise ValueError(msg)
|
||||
|
||||
|
||||
def default_aggregate(col_name, df):
|
||||
"""Return the default aggregate for a column based on its type."""
|
||||
dtype = df[col_name].dtype
|
||||
psp_type = polars_type_to_psp(dtype)
|
||||
if psp_type in ("integer", "float"):
|
||||
return "sum"
|
||||
return "count"
|
||||
|
||||
|
||||
def build_rollup(df, group_by, columns, aggregates, col_alias):
|
||||
"""Emulate GROUP BY ROLLUP using multiple group_by operations."""
|
||||
n = len(group_by)
|
||||
frames = []
|
||||
data_columns = [c for c in columns if c not in group_by]
|
||||
|
||||
for level in range(n + 1):
|
||||
num_groups = n - level
|
||||
active_groups = group_by[:num_groups]
|
||||
|
||||
agg_exprs = []
|
||||
for col in data_columns:
|
||||
agg_name = aggregates.get(col, default_aggregate(col, df))
|
||||
agg_exprs.append(get_polars_agg_expr(col, agg_name).alias(col_alias(col)))
|
||||
|
||||
if active_groups:
|
||||
grouped = df.group_by(active_groups, maintain_order=True).agg(agg_exprs)
|
||||
else:
|
||||
grouped = df.select(agg_exprs)
|
||||
|
||||
for idx in range(n):
|
||||
if idx < num_groups:
|
||||
grouped = grouped.with_columns(
|
||||
pl.col(group_by[idx]).alias(f"__ROW_PATH_{idx}__")
|
||||
)
|
||||
else:
|
||||
src_dtype = df[group_by[idx]].dtype
|
||||
grouped = grouped.with_columns(
|
||||
pl.lit(None).cast(src_dtype).alias(f"__ROW_PATH_{idx}__")
|
||||
)
|
||||
|
||||
grouping_id = sum(1 << i for i in range(num_groups, n))
|
||||
grouped = grouped.with_columns(
|
||||
pl.lit(grouping_id).cast(pl.Int64).alias("__GROUPING_ID__")
|
||||
)
|
||||
|
||||
for gb_col in active_groups:
|
||||
if gb_col in grouped.columns:
|
||||
grouped = grouped.drop(gb_col)
|
||||
|
||||
frames.append(grouped)
|
||||
|
||||
result = pl.concat(frames, how="diagonal")
|
||||
path_cols = [f"__ROW_PATH_{i}__" for i in range(n)]
|
||||
data_col_aliases = [col_alias(c) for c in data_columns]
|
||||
final_order = ["__GROUPING_ID__"] + path_cols + data_col_aliases
|
||||
result = result.select([c for c in final_order if c in result.columns])
|
||||
return result
|
||||
|
||||
|
||||
def build_flat_group_by(df, group_by, columns, aggregates, col_alias):
|
||||
"""Build a simple GROUP BY (no rollup) - only leaf-level rows."""
|
||||
n = len(group_by)
|
||||
data_columns = [c for c in columns if c not in group_by]
|
||||
|
||||
agg_exprs = []
|
||||
for col in data_columns:
|
||||
agg_name = aggregates.get(col, default_aggregate(col, df))
|
||||
agg_exprs.append(get_polars_agg_expr(col, agg_name).alias(col_alias(col)))
|
||||
|
||||
grouped = df.group_by(group_by, maintain_order=True).agg(agg_exprs)
|
||||
|
||||
for idx in range(n):
|
||||
grouped = grouped.with_columns(
|
||||
pl.col(group_by[idx]).alias(f"__ROW_PATH_{idx}__")
|
||||
)
|
||||
|
||||
for gb_col in group_by:
|
||||
if gb_col in grouped.columns:
|
||||
grouped = grouped.drop(gb_col)
|
||||
|
||||
path_cols = [f"__ROW_PATH_{i}__" for i in range(n)]
|
||||
data_col_aliases = [col_alias(c) for c in data_columns]
|
||||
final_order = path_cols + data_col_aliases
|
||||
result = grouped.select([c for c in final_order if c in grouped.columns])
|
||||
return result.sort(path_cols)
|
||||
|
||||
|
||||
def build_total(df, columns, aggregates, col_alias):
|
||||
"""Build a single total row aggregating the entire dataset."""
|
||||
agg_exprs = []
|
||||
for col in columns:
|
||||
agg_name = aggregates.get(col, default_aggregate(col, df))
|
||||
agg_exprs.append(get_polars_agg_expr(col, agg_name).alias(col_alias(col)))
|
||||
return df.select(agg_exprs)
|
||||
|
||||
|
||||
def build_split_by_total(df, split_by, columns, aggregates, col_alias):
|
||||
"""Build a single total row with split_by (pivot) columns."""
|
||||
split_col = split_by[0]
|
||||
data_columns = [c for c in columns if c not in split_by]
|
||||
split_values = sorted(df[split_col].unique().to_list())
|
||||
|
||||
agg_exprs = []
|
||||
for sv in split_values:
|
||||
filter_expr = pl.col(split_col) == sv
|
||||
for dc in data_columns:
|
||||
agg_name = aggregates.get(dc, default_aggregate(dc, df))
|
||||
col_name = f"{sv}_{col_alias(dc)}"
|
||||
agg_exprs.append(
|
||||
get_polars_agg_expr(dc, agg_name, filter_expr=filter_expr).alias(
|
||||
col_name
|
||||
)
|
||||
)
|
||||
|
||||
return df.select(agg_exprs)
|
||||
|
||||
|
||||
def build_split_by_grouped_flat(df, group_by, split_by, columns, aggregates, col_alias):
|
||||
"""Build a flat grouped view with split_by (pivot) columns - no rollup rows."""
|
||||
n = len(group_by)
|
||||
split_col = split_by[0]
|
||||
data_columns = [c for c in columns if c not in group_by and c not in split_by]
|
||||
split_values = sorted(df[split_col].unique().to_list())
|
||||
|
||||
agg_exprs = []
|
||||
for sv in split_values:
|
||||
filter_expr = pl.col(split_col) == sv
|
||||
for dc in data_columns:
|
||||
agg_name = aggregates.get(dc, default_aggregate(dc, df))
|
||||
col_name = f"{sv}_{col_alias(dc)}"
|
||||
agg_exprs.append(
|
||||
get_polars_agg_expr(dc, agg_name, filter_expr=filter_expr).alias(
|
||||
col_name
|
||||
)
|
||||
)
|
||||
|
||||
for dc in data_columns:
|
||||
agg_name = aggregates.get(dc, default_aggregate(dc, df))
|
||||
agg_exprs.append(get_polars_agg_expr(dc, agg_name).alias(f"__SORT_{dc}__"))
|
||||
|
||||
grouped = df.group_by(group_by, maintain_order=True).agg(agg_exprs)
|
||||
|
||||
for idx in range(n):
|
||||
grouped = grouped.with_columns(
|
||||
pl.col(group_by[idx]).alias(f"__ROW_PATH_{idx}__")
|
||||
)
|
||||
|
||||
for gb_col in group_by:
|
||||
if gb_col in grouped.columns:
|
||||
grouped = grouped.drop(gb_col)
|
||||
|
||||
path_cols = [f"__ROW_PATH_{i}__" for i in range(n)]
|
||||
data_col_names = []
|
||||
for sv in split_values:
|
||||
for dc in data_columns:
|
||||
data_col_names.append(f"{sv}_{col_alias(dc)}")
|
||||
sort_col_names = [f"__SORT_{dc}__" for dc in data_columns]
|
||||
final_order = path_cols + data_col_names + sort_col_names
|
||||
result = grouped.select([c for c in final_order if c in grouped.columns])
|
||||
return result.sort(path_cols)
|
||||
|
||||
|
||||
def apply_sort_grouped(df, sort_config, group_by, col_alias):
|
||||
"""Apply sort to a ROLLUP result DataFrame."""
|
||||
n = len(group_by)
|
||||
|
||||
sort_cols = []
|
||||
sort_desc = []
|
||||
for entry in sort_config:
|
||||
col = entry[0]
|
||||
direction = entry[1]
|
||||
if direction != "none":
|
||||
aliased = col_alias(col)
|
||||
if aliased in df.columns:
|
||||
sort_cols.append(aliased)
|
||||
sort_desc.append(direction in ("desc", "col desc"))
|
||||
|
||||
if not sort_cols:
|
||||
# Default: tree order by row path, nulls first
|
||||
path_cols = [f"__ROW_PATH_{i}__" for i in range(n)]
|
||||
return df.sort(path_cols, descending=[False] * n, nulls_last=False)
|
||||
|
||||
# With explicit sort: grand total first, then rest sorted
|
||||
is_total = pl.lit(True)
|
||||
for i in range(n):
|
||||
is_total = is_total & pl.col(f"__ROW_PATH_{i}__").is_null()
|
||||
|
||||
total_row = df.filter(is_total)
|
||||
rest = df.filter(~is_total)
|
||||
rest = rest.sort(sort_cols, descending=sort_desc)
|
||||
return pl.concat([total_row, rest])
|
||||
|
||||
|
||||
def apply_sort_split_by_flat(df, sort_config, columns, group_by, split_by):
|
||||
"""Apply sort to a flat split_by grouped DataFrame using __SORT__ columns."""
|
||||
data_columns = [c for c in columns if c not in group_by and c not in split_by]
|
||||
sort_cols = []
|
||||
sort_desc = []
|
||||
for entry in sort_config:
|
||||
col = entry[0]
|
||||
direction = entry[1]
|
||||
if direction != "none":
|
||||
sort_name = f"__SORT_{col}__"
|
||||
if sort_name in df.columns:
|
||||
sort_cols.append(sort_name)
|
||||
sort_desc.append(direction in ("desc", "col desc"))
|
||||
|
||||
if sort_cols:
|
||||
df = df.sort(sort_cols, descending=sort_desc)
|
||||
|
||||
drop_cols = [
|
||||
f"__SORT_{dc}__" for dc in data_columns if f"__SORT_{dc}__" in df.columns
|
||||
]
|
||||
if drop_cols:
|
||||
df = df.drop(drop_cols)
|
||||
return df
|
||||
|
||||
|
||||
def apply_sort_flat(df, sort_config, col_alias):
|
||||
"""Apply sort to a flat (non-grouped) DataFrame."""
|
||||
if not sort_config:
|
||||
return df
|
||||
|
||||
sort_cols = []
|
||||
sort_descending = []
|
||||
for sort_entry in sort_config:
|
||||
col = sort_entry[0]
|
||||
direction = sort_entry[1]
|
||||
if direction != "none":
|
||||
aliased = col_alias(col)
|
||||
if aliased in df.columns:
|
||||
sort_cols.append(aliased)
|
||||
sort_descending.append(direction in ("desc", "col desc"))
|
||||
|
||||
if sort_cols:
|
||||
return df.sort(sort_cols, descending=sort_descending)
|
||||
return df
|
||||
|
||||
|
||||
def compute_view_schema(df):
|
||||
"""Compute the Perspective schema for a view DataFrame."""
|
||||
schema = {}
|
||||
for col_name, dtype in df.schema.items():
|
||||
if col_name.startswith("__"):
|
||||
continue
|
||||
schema[col_name] = polars_type_to_psp(dtype)
|
||||
return schema
|
||||
|
||||
|
||||
def build_split_by_grouped(df, group_by, split_by, columns, aggregates, col_alias):
|
||||
"""Build a grouped rollup with split_by (pivot) columns."""
|
||||
n = len(group_by)
|
||||
split_col = split_by[0]
|
||||
data_columns = [c for c in columns if c not in group_by and c not in split_by]
|
||||
split_values = sorted(df[split_col].unique().to_list())
|
||||
|
||||
frames = []
|
||||
for level in range(n + 1):
|
||||
num_groups = n - level
|
||||
active_groups = group_by[:num_groups]
|
||||
|
||||
agg_exprs = []
|
||||
for sv in split_values:
|
||||
filter_expr = pl.col(split_col) == sv
|
||||
for dc in data_columns:
|
||||
agg_name = aggregates.get(dc, default_aggregate(dc, df))
|
||||
col_name = f"{sv}_{col_alias(dc)}"
|
||||
agg_exprs.append(
|
||||
get_polars_agg_expr(dc, agg_name, filter_expr=filter_expr).alias(
|
||||
col_name
|
||||
)
|
||||
)
|
||||
|
||||
if active_groups:
|
||||
grouped = df.group_by(active_groups, maintain_order=True).agg(agg_exprs)
|
||||
else:
|
||||
grouped = df.select(agg_exprs)
|
||||
|
||||
for idx in range(n):
|
||||
if idx < num_groups:
|
||||
grouped = grouped.with_columns(
|
||||
pl.col(group_by[idx]).alias(f"__ROW_PATH_{idx}__")
|
||||
)
|
||||
else:
|
||||
src_dtype = df[group_by[idx]].dtype
|
||||
grouped = grouped.with_columns(
|
||||
pl.lit(None).cast(src_dtype).alias(f"__ROW_PATH_{idx}__")
|
||||
)
|
||||
|
||||
grouping_id = sum(1 << i for i in range(num_groups, n))
|
||||
grouped = grouped.with_columns(
|
||||
pl.lit(grouping_id).cast(pl.Int64).alias("__GROUPING_ID__")
|
||||
)
|
||||
|
||||
for gb_col in active_groups:
|
||||
if gb_col in grouped.columns:
|
||||
grouped = grouped.drop(gb_col)
|
||||
|
||||
frames.append(grouped)
|
||||
|
||||
result = pl.concat(frames, how="diagonal")
|
||||
path_cols = [f"__ROW_PATH_{i}__" for i in range(n)]
|
||||
data_col_names = []
|
||||
for sv in split_values:
|
||||
for dc in data_columns:
|
||||
data_col_names.append(f"{sv}_{col_alias(dc)}")
|
||||
final_order = ["__GROUPING_ID__"] + path_cols + data_col_names
|
||||
result = result.select([c for c in final_order if c in result.columns])
|
||||
return result
|
||||
|
||||
|
||||
def build_split_by_flat(df, split_by, columns, col_alias):
|
||||
"""Build a flat (non-grouped) split_by view."""
|
||||
split_col = split_by[0]
|
||||
data_columns = [c for c in columns if c not in split_by]
|
||||
split_values = sorted(df[split_col].unique().to_list())
|
||||
|
||||
exprs = []
|
||||
for sv in split_values:
|
||||
for dc in data_columns:
|
||||
col_name = f"{sv}_{col_alias(dc)}"
|
||||
exprs.append(
|
||||
pl.when(pl.col(split_col) == sv)
|
||||
.then(pl.col(dc))
|
||||
.otherwise(None)
|
||||
.alias(col_name)
|
||||
)
|
||||
|
||||
return df.select(exprs)
|
||||
|
||||
|
||||
def parse_expression(expr_str):
|
||||
"""Parse a Perspective expression string into a Polars expression."""
|
||||
pattern = r'"([^"]*)"'
|
||||
parts = []
|
||||
last_end = 0
|
||||
for match in re.finditer(pattern, expr_str):
|
||||
parts.append(expr_str[last_end : match.start()])
|
||||
col_name = match.group(1)
|
||||
parts.append(f'pl.col("{col_name}")')
|
||||
last_end = match.end()
|
||||
parts.append(expr_str[last_end:])
|
||||
polars_expr_str = "".join(parts)
|
||||
return eval(polars_expr_str, {"pl": pl, "__builtins__": {}})
|
||||
Reference in New Issue
Block a user