diff --git a/Architecture.md b/Architecture.md index 49f2dc8..8073c2c 100644 --- a/Architecture.md +++ b/Architecture.md @@ -994,8 +994,8 @@ On a bare single-file local Parquet scan, `physical_minmax` reads the footer sta `SQLSource` in `sql.py` is the second kind of source, next to `LFQueryBuilder`. (Elsewhere in this document, "both source kinds" means the two forms of a builder: a resident frame and a scan.) The data stays in a database -(Postgres, DuckDB or ClickHouse), and each trace aggregation runs there as one -SQL query. It offers the surface the engine and the server use of an +(Postgres, DuckDB, ClickHouse or Trino), and each trace aggregation runs there +as one SQL query. It offers the surface the engine and the server use of an `LFQueryBuilder`: `schema`, `cache` (set by the registrar), `is_scan` (always true: a plan reads rows in no order), `physical_minmax`, `compile_filter` and `aggregate`. The server stores it in `_sources` like a builder, and `Figure`, @@ -1036,17 +1036,37 @@ SQLSource domain-minimum row to bucket -1. It is also cast to a non-nullable double: Postgres reads an exponent literal as `numeric`, so a `numeric` column would compare and bucket in decimal, not as the double that Polars uses. A - nullable cast made a ClickHouse bucket query 2.5 times slower. + nullable cast made a ClickHouse bucket query 2.5 times slower. A plan + groups by the bucket expression, not by its alias: Trino cannot see a + select alias in `GROUP BY`, and Postgres and ClickHouse make the same plan + for both forms. - **Dialect meaning lives in `SQLFrame`.** SQLGlot quotes identifiers and renders literals per dialect. The functions whose meaning differs per dialect (`phys`, `nan`, `num`, `int_div`, `extreme_by`, `lit`) are defined - there. + there. The median in `_agg_sql` also switches on the dialect. +- **Trino.** `date_diff` counts epoch units exactly, but in milliseconds at + most. A microsecond column adds the microseconds within the millisecond from + `to_unixtime`, a double. That is exact while one ulp of the epoch seconds is + below 0.5 µs: from the year 1834 to 2106, which the test checks. Outside it + Trino 484 was 1 µs off for about a fifth of the values. Shifting the value + by its whole milliseconds first (`date_add`) is exact everywhere but made a + timestamp histogram 1.6 times slower; the text form was 2.6 times slower. + The whole double rounded wrong from 2100 on. + Trino has no exact percentile, so the median is the mean of the two middle + values of `array_sort(array_agg(v))` (each group in memory), each halved + first so the sum cannot overflow. The `trino` driver returns rows, so the dtypes come from the type + names of the cursor description. A column of another type (uuid, time, + varbinary, array, map, row) is left out of the schema: Trino compares it with + no string literal, so a selection on it could not run. A timestamp finer than + microseconds is left out too: the driver rounds it to a Python datetime, the + epoch in SQL does not, so a bin would disagree with the value. - **Identifiers.** Column names come from the spec, which the browser sends, so every column reaches the SQL through `SQLSource.col`, which refuses a name - that is not in the schema (and, on ClickHouse, a name with a backslash, which - ClickHouse reads as an escape inside a quoted identifier). Every other - identifier is a FlexViz constant: a trace uid is never an SQL alias. The - user's `query=` goes into a `WITH` clause as written, never re-rendered. + that is not in the schema. On ClickHouse it also refuses a name with a + backslash, which ClickHouse reads as an escape inside a quoted identifier. + Postgres, DuckDB and Trino double a quote and have no backslash escape. Every + other identifier is a FlexViz constant: a trace uid is never an SQL alias. + The user's `query=` goes into a `WITH` clause as written, never re-rendered. --- diff --git a/docs/guides/databases.md b/docs/guides/databases.md index 7c27c2b..51b363f 100644 --- a/docs/guides/databases.md +++ b/docs/guides/databases.md @@ -27,13 +27,17 @@ detail table with a join: |---|---|---|---|---|---| | Live ClickHouse | same machine | 0.6 | 1.0 | 0.55 | 1.1 | | Live ClickHouse | 20 Mbit/s, 30 ms | 0.7 | 2.0 | 1.3 | 1.5 | +| Live Trino, Parquet files | same machine | 0.1 | 4.3 | 1.25 | 3.6 | +| Live Trino, Parquet files | 20 Mbit/s, 30 ms | 0.4 | 5.1 | 2.4 | 4.3 | +| Live DuckDB, the same Parquet files | same machine | 0.1 | 1.5 | 0.6 | 1.7 | | Live Postgres | same machine | 0.1 | 8.8 | 0.9 | 7.2 | | Live Postgres | 20 Mbit/s, 30 ms | 0.8 | 15.2 | 3.9 | 9.5 | | Extract, ADBC | same machine | 2.0 | 0.5 | 0.35 | 0.4 | | Extract, ADBC | 20 Mbit/s, 30 ms | 105 | 0.4 | 0.3 | 0.3 | | Extract, SQLAlchemy | same machine | 9.7 | 0.5 | 0.35 | 0.4 | -On the full 100M-row table, live ClickHouse renders in 1.5 s and live +On the full 100M-row table, live ClickHouse renders in 1.5 s, live DuckDB on +Parquet files in 3.0 s, live Trino on the same files in 4.6 s, and live Postgres in 45 s on the same machine. ### Short, wide runs @@ -54,9 +58,9 @@ machine, because every chart query read the whole table. Use these rules: -- **A columnar engine** (ClickHouse or DuckDB): use a live `SQLSource`. The - engine aggregates millions of rows in milliseconds, so every zoom and - selection stays interactive. +- **A columnar engine** (ClickHouse, DuckDB or Trino): use a live + `SQLSource`. The engine aggregates millions of rows in milliseconds to + seconds, so every zoom and selection stays interactive. - **Postgres with a few hundred thousand rows behind each chart**: a live `SQLSource` is interactive, if the query can use an index. Postgres aggregates 10 to 30 times more slowly than a columnar engine, so the row @@ -135,6 +139,35 @@ src = fv.SQLSource( ) ``` +Trino example (`pip install trino`): + +```python +import trino + +src = fv.SQLSource( + lambda: trino.dbapi.connect( + host="trino.example", port=8080, user="reader", catalog="iceberg", schema="fab" + ), + table="log_detail", + dialect="trino", +) +``` + +On Trino, a column of a type such as `uuid`, `json`, `time`, `varbinary`, +`char`, `array`, `map` or `row` is not available to charts. Cast it in the +query, for example `CAST(run_id AS varchar) AS run_id`. A `char` value comes +back padded with spaces, which a selection on it cannot match. A timestamp +finer than microseconds (`timestamp(9)`) is not available either; read it as +`CAST(ts AS timestamp(6)) AS ts`. Keep a `timestamp with time zone` column in +UTC: `ts AT TIME ZONE 'UTC' AS ts`. In a zone with daylight saving time, the +driver reads a value in the repeated hour one hour early. + +Use Trino through its catalogs for files and lakehouse tables (Hive, Iceberg, +Delta Lake). Do not put Trino in front of Postgres to speed it up: Trino's +Postgres connector does not push a binned `GROUP BY` down, so Trino reads every +row from Postgres. A 60-bin histogram over 100M rows took 63 s that way and +11 s in Postgres directly. + ### What runs in the database Each chart is one `GROUP BY` query. A line becomes one row per x bucket (per @@ -146,21 +179,27 @@ binned columns, so that the bins line up across charts. The results are equal to the results on a Polars frame with the same rows, with these exceptions: -- A sum or a mean can be different in the last digits, because the database - adds the values in its own order. +- A sum, a mean or a median can be different in the last digits, because the + database adds and halves the values in its own order. - If more rows have the same y minimum or y maximum in a line bucket, the database can show the x of a different one of these rows. - On Postgres, a sum of 64-bit integers is exact only up to 2**53. - On Postgres, the minimum and the maximum of a text column follow the collation of the database. +- On Postgres and Trino, a microsecond timestamp bins exactly only from the + year 1834 to 2106 (see Limits). -The test suite compares the two for every supported chart on DuckDB, Postgres -and ClickHouse. +The test suite compares the two for every supported chart on DuckDB, Postgres, +ClickHouse and Trino. ### Supported dialects and charts -Dialects: Postgres, DuckDB and ClickHouse. Other dialects give an error at -construction. +Dialects: Postgres, DuckDB, ClickHouse and Trino. Other dialects give an error +at construction. + +The test suite runs on these four engines only. FlexViz does not test engines +that are compatible with one of them, for example TimescaleDB (Postgres), +MotherDuck (DuckDB) or Starburst (Trino). | Chart | Live on a `SQLSource` | |---|---| @@ -187,6 +226,13 @@ have, gives an error when you add the chart, in your own code. in memory instead of sorting on disk. A line over 5M rows took 3.1 s with 64 MB and 1.9 s with 1 GB: `ALTER ROLE flexviz_reader SET work_mem = '1GB'`. +- On Trino, a median (`agg="median"`) sorts the values of each group in + memory, because Trino has no exact percentile function. +- On Postgres and Trino, the epoch microseconds of a microsecond timestamp + are exact from the year 1834 to 2106. Outside that range (for example a + `9999-12-31` sentinel date) a value can be 1 µs or more off, so it can fall + in the bin next to its own at a bin edge. The value that a chart shows stays + exact. - `Dashboard(src, cache=True)` declares that the data does not change while the server runs. FlexViz then keeps the column bounds and the first render. Without it, every request first asks the database for the column bounds @@ -241,3 +287,20 @@ If the rows do not fit in memory, write them to Parquet once and scan the file: df.write_parquet("runs.parquet") fv.Dashboard(pl.scan_parquet("runs.parquet")) ``` + +## Files on object storage: no warehouse + +If the data is already Parquet or Iceberg files, for example on S3, one +machine can read it without a warehouse: + +- Polars: pass `pl.scan_parquet("s3://bucket/runs/*.parquet")` or + `pl.scan_iceberg(...)` to `Dashboard`. FlexViz keeps the scan lazy. +- DuckDB: pass a DuckDB connection to `SQLSource`, with a query that reads the + files, for example + `fv.SQLSource(duckdb.connect(), query="SELECT * FROM read_parquet('s3://bucket/runs/*.parquet')")`. + For a private bucket, run `CREATE SECRET` on the connection first, and for + Iceberg, attach the catalog with DuckDB's `iceberg` extension first. + +On one machine and the same Parquet files, DuckDB rendered the dashboard of +the first table in 1.5 s and Trino in 4.3 s. Trino is the better choice when +your team already runs it, or when the data needs more than one machine. diff --git a/flexviz/sql.py b/flexviz/sql.py index 1ae3edc..17f0ffd 100644 --- a/flexviz/sql.py +++ b/flexviz/sql.py @@ -31,6 +31,7 @@ import polars as pl try: + import sqlglot from sqlglot import exp from sqlglot.dialects.dialect import Dialect from sqlglot.tokens import TokenType @@ -45,7 +46,7 @@ from .trace.base import _typed_range_bounds #: The dialects whose SQL has run against a real database in the test suite. -_DIALECTS = ("postgres", "duckdb", "clickhouse") +_DIALECTS = ("postgres", "duckdb", "clickhouse", "trino") _DIALECT_ALIASES = {"postgresql": "postgres", "clickhousedb": "clickhouse"} # Every column a plan adds carries this prefix, like the Polars plans. @@ -85,6 +86,40 @@ } +def _trino_dtype(name: str) -> pl.DataType | None: + """The Polars dtype of a Trino type name, as the ``trino`` driver's rows load. + + ``decimal`` reads as Float64, as Postgres ``numeric``. None for other + types (uuid, time, varbinary, array, map, row, ...): the driver returns + Python objects for them, and Trino compares them with no string. None also + for a timestamp finer than microseconds: the driver rounds it to a Python + datetime, but the epoch in SQL does not, so a bin would disagree with the + value. ``CAST(c AS timestamp(6))`` in the query reads it. + """ + base, _, args = name.partition("(") + simple = { + "boolean": pl.Boolean(), + "tinyint": pl.Int8(), + "smallint": pl.Int16(), + "integer": pl.Int32(), + "bigint": pl.Int64(), + "real": pl.Float32(), + "double": pl.Float64(), + "decimal": pl.Float64(), + "date": pl.Date(), + "varchar": pl.String(), + } + if base in simple: + return simple[base] + if base == "timestamp": + precision = int(args.split(")")[0]) if args else 3 + if precision > 6: + return None + zone = "UTC" if name.endswith("with time zone") else None + return pl.Datetime("ms" if precision <= 3 else "us", zone) + return None + + def _clickhouse_dtype(name: str) -> pl.DataType: """The Polars dtype of a ClickHouse type name, as its DB-API rows load.""" # In any order: ClickHouse writes LowCardinality(Nullable(T)). @@ -197,6 +232,22 @@ def _fetch(conn: Any, sql: str) -> pl.DataFrame: cur.execute(sql) schema = {d[0]: _clickhouse_dtype(str(d[1])) for d in cur.description} return pl.DataFrame(cur.fetchall(), schema=schema, orient="row") + if module == "trino": + # Rows without types: the cursor names each column's Trino type. A + # column of a type without a dtype is left out, so it is not in the + # schema, and a chart on it fails when it is added. + cur = conn.cursor() + try: + cur.execute(sql) + rows = cur.fetchall() + types = [(d[0], _trino_dtype(d[1])) for d in cur.description] + finally: + cur.close() + keep = [i for i, (_, t) in enumerate(types) if t is not None] + if len(keep) < len(types): + rows = [[r[i] for i in keep] for r in rows] + schema = {types[i][0]: types[i][1] for i in keep} + return pl.DataFrame(rows, schema=schema, orient="row") return pl.read_database(sql, conn) @@ -219,7 +270,9 @@ def _alias(e: exp.Expression, name: str) -> exp.Alias: def _num(v: float) -> exp.Expression: """A numeric literal that parses back to the same double or integer.""" if isinstance(v, float) and not math.isfinite(v): - return exp.cast(exp.Literal.string(repr(v)), exp.DataType.Type.DOUBLE) + # Trino casts only the spelled-out name, which every dialect reads. + text = "NaN" if math.isnan(v) else ("Infinity" if v > 0 else "-Infinity") + return exp.cast(exp.Literal.string(text), exp.DataType.Type.DOUBLE) if v < 0: # Parenthesized, so ``a - -1`` can never print as the comment ``a --1``. return exp.Paren(this=exp.Neg(this=_num(-v))) @@ -237,6 +290,14 @@ def _and(*conds: exp.Expression) -> exp.Expression: return exp.and_(*conds) if conds else exp.true() +def _trino(template: str, **kw: exp.Expression) -> exp.Expression: + """A Trino expression from SQL text: each ``:name`` is replaced by the + expression ``name``, so no identifier is formatted into the text.""" + return sqlglot.parse_one(template, read="trino").transform( + lambda n: kw[n.name].copy() if isinstance(n, exp.Placeholder) else n + ) + + @dataclass(frozen=True) class SQLFrame: """The source's rows under the cross-filter: what a SQL plan reads. @@ -256,9 +317,13 @@ def col(self, name: str) -> exp.Column: def num(self, name: str) -> exp.Expression: """A value column as a number. Postgres neither sums, takes the min of, - nor casts to a double a Boolean, so it becomes an integer there.""" + nor casts to a double a Boolean, and Trino does not sum one, so it + becomes an integer there.""" c = self.col(name) - if self.dialect == "postgres" and self.source.schema[name] == pl.Boolean: + if ( + self.dialect in ("postgres", "trino") + and self.source.schema[name] == pl.Boolean + ): return exp.cast(c, exp.DataType.Type.INT) return c @@ -283,6 +348,8 @@ def phys(self, name: str, dtype: pl.DataType | None) -> exp.Expression: if dtype == pl.Date: if d == "clickhouse": return exp.func("toInt32", c) + if d == "trino": + return _trino("date_diff('day', DATE '1970-01-01', :c)", c=c) epoch = exp.cast(exp.Literal.string("1970-01-01"), exp.DataType.Type.DATE) return exp.cast(c - epoch, exp.DataType.Type.BIGINT) if not isinstance(dtype, pl.Datetime): @@ -299,6 +366,28 @@ def phys(self, name: str, dtype: pl.DataType | None) -> exp.Expression: f"toUnixTimestamp64{fn}", exp.func("toDateTime64", c, exp.Literal.number(precision)), ) + if d == "trino": + # date_diff counts exactly, but in milliseconds at most. The + # microseconds within the millisecond come from to_unixtime, a + # double, as on Postgres: exact from 1834 to 2106. shortcut: + # outside that range a value can be 1 us off (Trino 484); shifting + # the value by its whole milliseconds first (date_add) is exact + # everywhere, but made a timestamp histogram 1.6 times slower. + epoch = "TIMESTAMP '1970-01-01 00:00:00" + ( + " UTC'" if dtype.time_zone else "'" + ) + ms = f"date_diff('millisecond', {epoch}, :c)" + if unit == "ms": + return _trino(ms, c=c) + seconds = ( + "to_unixtime(:c)" + if dtype.time_zone + else "to_unixtime(with_timezone(:c, 'UTC'))" + ) + return _trino( + f"{ms} * 1000 + CAST(round(({seconds} - {ms} / 1e3) * 1e6) AS bigint)", + c=c, + ) # Postgres keeps microseconds. date_part returns a double (2x faster # than the numeric EXTRACT), which rounds back to the exact count while # one ulp of the epoch seconds is below 0.5 us: from 1834 to 2106. @@ -322,7 +411,7 @@ def nan(self, name: str) -> exp.Expression: if self.dialect == "clickhouse": # isNaN refuses a Decimal, which the schema reads as Float64. return exp.func("isNaN", exp.func("toFloat64", c)) - return exp.func("isnan", c) + return exp.func("is_nan" if self.dialect == "trino" else "isnan", c) def usable(self, name: str, dtype: pl.DataType | None) -> exp.Expression: """Not null, and not NaN for a float: the rows a bin or a bucket takes.""" @@ -340,7 +429,8 @@ def int_div(self, a: exp.Expression, b: int) -> exp.Expression: return exp.func("intDiv", a, _num(b)) if d == "duckdb": return exp.IntDiv(this=a, expression=_num(b)) - # Postgres: ``/`` on two integers truncates, which is floor for a >= 0. + # Postgres and Trino: ``/`` on two integers truncates, which is floor + # for a >= 0. SQLGlot's Trino IntDiv divides as a double. return exp.Div(this=a, expression=_num(b), typed=True) def lit(self, value: Any, dtype: pl.DataType | None = None) -> exp.Expression: @@ -353,8 +443,6 @@ def lit(self, value: Any, dtype: pl.DataType | None = None) -> exp.Expression: return self._ns_lit(value) value = pl.Series([value], dtype=pl.Int64).cast(dtype).item() if isinstance(value, float): - if math.isnan(value): - return exp.cast(exp.Literal.string("NaN"), exp.DataType.Type.DOUBLE) return _num(value) if isinstance(value, int) and not isinstance(value, bool): return _num(value) @@ -872,7 +960,8 @@ def bucket_extrema( _alias(hi_y, names["hi_y"]), ) .where(*conds) - .group_by(*[g.copy() for g in groups], exp.column(_BUCKET, quoted=True)) + # By the expression, not its alias, which Trino cannot see in GROUP BY. + .group_by(*[g.copy() for g in groups], bucket.copy()) ) df = frame.collect(select) @@ -941,7 +1030,7 @@ def hist_counts( _alias(exp.Count(this=exp.Star()), _P + "count"), ) .where(*conds) - .group_by(*[g.copy() for g in groups], exp.column(_P + "b", quoted=True)) + .group_by(*[g.copy() for g in groups], b.copy()) ) df = frame.collect(select) dtypes: dict[str, pl.DataType] = {_P + "b": pl.Int32(), _P + "count": pl.UInt32()} @@ -1026,7 +1115,7 @@ def run(frame: SQLFrame) -> pl.DataFrame: select = ( exp.select(_alias(cell, _P + "c"), _alias(value, _P + "v")) .where(*conds) - .group_by(exp.column(_P + "c", quoted=True)) + .group_by(cell.copy()) ) df = frame.collect(select) idx = df[_P + "c"].cast(pl.Int64).to_numpy() @@ -1081,6 +1170,18 @@ def _agg_sql( # Exact and interpolating, like Polars' median; plain ``median`` # in ClickHouse samples. return exp.Anonymous(this="quantileExactInclusive(0.5)", expressions=[v]) + if frame.dialect == "trino": + # Trino has only approx_percentile: the mean of the middle values of + # the group, sorted once. The array holds the group in memory. Each + # value is halved first, so the sum cannot overflow; halving a + # double is exact unless the result is subnormal. + a = "element_at(s, (cardinality(s) + 1) / 2)" + b = "element_at(s, cardinality(s) / 2 + 1)" + return _trino( + "transform(ARRAY[array_sort(array_agg(CAST(:v AS double))" + f" FILTER (WHERE :v IS NOT NULL))], s -> {a} / 2 + {b} / 2)[1]", + v=v, + ) return exp.PercentileCont(this=v, expression=_num(0.5)) raise ValueError(f"agg {agg!r} is not supported on a {frame.dialect} source") diff --git a/flexviz/trace/base.py b/flexviz/trace/base.py index 5f8bd64..58eb173 100644 --- a/flexviz/trace/base.py +++ b/flexviz/trace/base.py @@ -259,7 +259,10 @@ def check_source(self, source: LFQueryBuilder | SQLSource) -> None: ] missing = [c for c in dict.fromkeys(cols) if c not in source.schema] if missing: - raise ValueError(f"column(s) {missing} not in the SQL source") + raise ValueError( + f"column(s) {missing} not in the SQL source, or of a type FlexViz " + "cannot read (cast it in the query)" + ) def domain_cols(self, update_range: dict[str, Any]) -> tuple[str, ...]: """Columns whose ``(min, max)`` this trace's spec needs. diff --git a/tests/test_sql.py b/tests/test_sql.py index 3673b99..f5c2822 100644 --- a/tests/test_sql.py +++ b/tests/test_sql.py @@ -1,9 +1,11 @@ """SQLSource: a database source gives the deltas a Polars source gives. The core matrix runs on an in-process DuckDB. A reduced matrix runs on Postgres -when ``FLEXVIZ_TEST_POSTGRES_URI`` is set, and on ClickHouse when +when ``FLEXVIZ_TEST_POSTGRES_URI`` is set, on ClickHouse when ``FLEXVIZ_TEST_CLICKHOUSE_URI`` is set -(``clickhouse://user:password@host:port/database``). +(``clickhouse://user:password@host:port/database``), and on Trino when +``FLEXVIZ_TEST_TRINO_URI`` is set (``trino://user@host:port/catalog/schema``, a +catalog that can create tables, such as ``memory``). """ from __future__ import annotations @@ -32,7 +34,7 @@ from flexviz.LF import AggregationSpec, as_source from flexviz.server import app, register_source from flexviz.spec import ClauseFilter, SelectionPredicate, decode_spec -from flexviz.sql import SQLFrame, SQLSource, _clickhouse_dtype +from flexviz.sql import SQLFrame, SQLSource, _clickhouse_dtype, _trino_dtype pytestmark = pytest.mark.integration @@ -42,6 +44,7 @@ PG_ENV = "FLEXVIZ_TEST_POSTGRES_URI" CH_ENV = "FLEXVIZ_TEST_CLICKHOUSE_URI" +TRINO_ENV = "FLEXVIZ_TEST_TRINO_URI" N = 12_000 T0 = dt.datetime(2026, 1, 1) @@ -719,7 +722,7 @@ def test_literals_round_trip_exactly(duck): @pytest.mark.parametrize("value", [-1.5, -3, -(2**62), -1e-300]) def test_negative_literal_subtracts(duck, value): """``2 - -1`` printed as ``2 --1`` would read as a comment.""" - for dialect in ("duckdb", "postgres", "clickhouse"): + for dialect in ("duckdb", "postgres", "clickhouse", "trino"): frame = SQLFrame(SQLSource(duck, table="src", dialect=dialect)) sql = exp.select( exp.Sub(this=exp.Literal.number(2), expression=frame.lit(value)) @@ -729,7 +732,7 @@ def test_negative_literal_subtracts(duck, value): assert duck.execute(sql).fetchone()[0] == 2 - value -@pytest.mark.parametrize("dialect", ["duckdb", "postgres", "clickhouse"]) +@pytest.mark.parametrize("dialect", ["duckdb", "postgres", "clickhouse", "trino"]) def test_identifiers_are_quoted(dialect): con = duckdb.connect() con.execute('CREATE TABLE t ("we""ird name" INT, "select" INT)') @@ -768,6 +771,8 @@ def test_clickhouse_refuses_a_backslash_identifier(): "postgres", "E'it''s a \\\\ back\\\\slash \"dq\" -- not a comment; DROP TABLE src'", ), + # Trino has no backslash escape in a string literal. + ("trino", "'it''s a \\ back\\slash \"dq\" -- not a comment; DROP TABLE src'"), # ClickHouse reads a backslash as an escape. ( "clickhouse", @@ -926,7 +931,8 @@ def test_dialect_alias_is_accepted(duck): @pytest.mark.parametrize( - ("name", "dialect"), [("postgresql", "postgres"), ("clickhousedb", "clickhouse")] + ("name", "dialect"), + [("postgresql", "postgres"), ("clickhousedb", "clickhouse"), ("trino", "trino")], ) def test_sqlalchemy_dialect_names(name, dialect): """The dialect of a SQLAlchemy engine comes from its dialect name.""" @@ -1610,7 +1616,11 @@ def _real_matrix() -> list[tuple[str, str, str]]: # Generated only for a configured database: hundreds of skips say nothing. return [ (backend, case, event) - for backend, env in (("postgres", PG_ENV), ("clickhouse", CH_ENV)) + for backend, env in ( + ("postgres", PG_ENV), + ("clickhouse", CH_ENV), + ("trino", TRINO_ENV), + ) if os.environ.get(env) for case, event in _matrix(REAL_CASES, REAL_EVENTS) ] @@ -1624,7 +1634,7 @@ def test_real_database_matches_polars(request, backend, case, event): _assert_same(client, name + "_ref", name, case, event) -@pytest.mark.parametrize("backend", ["postgres", "clickhouse"]) +@pytest.mark.parametrize("backend", ["postgres", "clickhouse", "trino"]) def test_real_database_literals_round_trip(request, backend): src = request.getfixturevalue(f"{backend}_source") values = [*_random_doubles(30), math.inf, -math.inf, math.nan, 2**62 + 1, -7] @@ -1659,6 +1669,173 @@ def test_clickhouse_minmax_over_no_rows_is_null(clickhouse_client): clickhouse_client.command(f'DROP TABLE IF EXISTS "{table}"') +_TRINO_TYPES = { + pl.Float64: "double", + pl.Int16: "smallint", + pl.Int32: "integer", + pl.Int64: "bigint", + pl.String: "varchar", + pl.Boolean: "boolean", + pl.Date: "date", +} + + +def _trino_type(name: str, dtype: pl.DataType) -> str: + if name == "dec": + return "decimal(12,3)" + if isinstance(dtype, pl.Datetime): + return "timestamp(6)" + (" with time zone" if dtype.time_zone else "") + return _TRINO_TYPES[dtype.base_type()] + + +def _trino_value(v: Any) -> str: + """A Trino literal for one test value. Test data only, never user input.""" + if v is None: + return "NULL" + if isinstance(v, bool): + return "TRUE" if v else "FALSE" + if isinstance(v, float): + if math.isnan(v): + return "nan()" + if math.isinf(v): + return "infinity()" if v > 0 else "-infinity()" + return f"DOUBLE '{v!r}'" + if isinstance(v, int): + return str(v) + if isinstance(v, dt.datetime): + zone = " UTC" if v.tzinfo else "" + return ( + f"TIMESTAMP '{v.replace(tzinfo=None).isoformat(' ', 'microseconds')}{zone}'" + ) + if isinstance(v, dt.date): + return f"DATE '{v.isoformat()}'" + return "'" + str(v).replace("'", "''") + "'" + + +def _trino_connect() -> Any: + trino = pytest.importorskip("trino") + parts = urlsplit(os.environ[TRINO_ENV]) + catalog, schema = parts.path.strip("/").split("/") + return trino.dbapi.connect( + host=parts.hostname, + port=parts.port or 8080, + user=unquote(parts.username or "flexviz"), + catalog=catalog, + schema=schema, + ) + + +def _trino_execute(sql: str) -> None: + conn = _trino_connect() + try: + cur = conn.cursor() + cur.execute(sql) + cur.fetchall() + finally: + conn.close() + + +@pytest.fixture(scope="module") +def trino_source() -> Any: + if not os.environ.get(TRINO_ENV): + pytest.skip(f"set {TRINO_ENV} to run against Trino") + pytest.importorskip("trino") + table = f"fv_test_{secrets.token_hex(6)}" + df = DF.select(_REAL_COLUMNS) + cols = ", ".join(f'"{c}" {_trino_type(c, t)}' for c, t in df.schema.items()) + _trino_execute(f'CREATE TABLE "{table}" ({cols})') + try: + for chunk in df.iter_slices(2_000): + rows = ", ".join( + "(" + ", ".join(_trino_value(v) for v in row) + ")" + for row in chunk.iter_rows() + ) + _trino_execute(f'INSERT INTO "{table}" VALUES {rows}') + src = SQLSource(_trino_connect, dialect="trino", table=table) + ref = src._run(exp.select("*").order_by(exp.column("rid", quoted=True))) + register_source(_real_name("trino") + "_ref", ref) + register_source(_real_name("trino"), src) + yield src + finally: + _trino_execute(f'DROP TABLE IF EXISTS "{table}"') + + +def test_trino_schema_matches_the_frame(trino_source): + """The type names map to the frame's dtypes.""" + assert trino_source.schema == DF.select(_REAL_COLUMNS).schema + + +def test_trino_quotes_identifiers_and_strings(trino_source): + """A column name with a double quote and a comment marker, and a string + value with quotes, a backslash and SQL, reach Trino as data.""" + name = 'we"ird -- name; DROP TABLE x' + query = ( + f'SELECT "w" AS "{name.replace(chr(34), chr(34) * 2)}", "cat" ' + f'FROM "{trino_source._table}"' + ) + src = SQLSource(_trino_connect, dialect="trino", query=query) + frame = SQLFrame(src) + cond = exp.or_( + exp.EQ(this=frame.col("cat"), expression=frame.lit(AWKWARD)), + frame.range_cond(name, -1.0, 1.0), + ) + got = src._run(exp.select(exp.Count(this=exp.Star())).where(cond)).item() + want = DF.filter(pl.col("w").is_between(-1.0, 1.0)).height + assert got == want + + +def test_trino_millisecond_timestamp(trino_source): + """A timestamp(3) column reads as milliseconds, and bins and buckets in + milliseconds: the exact ``date_diff`` path, not the text form.""" + cols = ["rid", "ts_ms", "w"] + query = ( + 'SELECT "rid", CAST("ts" AS timestamp(3)) AS "ts_ms", "w" ' + f'FROM "{trino_source._table}"' + ) + src = SQLSource(_trino_connect, dialect="trino", query=query) + assert src.schema == DF.select(cols).schema + register_source("_sql_test_trino_ms", src) + register_source("_sql_test_trino_ms_ref", DF.select(cols)) + client = TestClient(app) + dash = Dashboard() + hist = dash.add_figure() + CASES["hist_ts_ms"][0](hist) + line = dash.add_figure() + line.add_line(x="ts_ms", y="w", n_points=300) + for event in ("init", "zoom"): + state, ev = _event( + event, {hist._uid: CASES["hist_ts_ms"][1], line._uid: LINE_TS_ZOOM} + ) + want = _post(client, dash, "_sql_test_trino_ms_ref", state, ev) + got = _post(client, dash, "_sql_test_trino_ms", state, ev) + assert want[hist._uid] and want[line._uid] + assert not _diff(want, got) + + +@pytest.mark.parametrize( + "cast", + [ + "AT TIME ZONE 'UTC' AS timestamp(6)", + "AT TIME ZONE 'America/New_York' AS timestamp(6) with time zone", + ], +) +def test_trino_epoch_microseconds_are_exact(trino_source, cast): + """``phys`` gives the exact epoch microseconds from 1834 to 2106, where one + ulp of the epoch seconds is below 0.5 us, and where a plain double rounds + wrong from about 2100 on.""" + query = f"""SELECT us, CAST(from_unixtime_nanos(us * 1000) {cast}) AS c + FROM (SELECT b + n * 7919 + (n * n) % 997 AS us + FROM UNNEST(sequence(0, 9999)) AS t(n) + CROSS JOIN UNNEST(ARRAY[BIGINT '-4294000000000000', + BIGINT '-1500000000123457', BIGINT '0', BIGINT '1767225600000000', + BIGINT '4200000000000000', BIGINT '4294000000000000']) AS s(b))""" + src = SQLSource(_trino_connect, dialect="trino", query=query) + frame = SQLFrame(src) + phys = frame.phys("c", src.schema["c"]) + wrong = exp.CountIf(this=exp.NEQ(this=phys, expression=frame.col("us"))) + assert src._run(exp.select(wrong, exp.Count(this=exp.Star()))).row(0) == (0, 60_000) + + @pytest.mark.parametrize( ("name", "dtype"), [ @@ -1674,3 +1851,73 @@ def test_clickhouse_minmax_over_no_rows_is_null(clickhouse_client): ) def test_clickhouse_type_names(name, dtype): assert _clickhouse_dtype(name) == dtype + + +@pytest.mark.parametrize( + ("name", "dtype"), + [ + ("integer", pl.Int32()), + ("decimal(10, 2)", pl.Float64()), + ("varchar(3)", pl.String()), + # Padded values that a filter cannot match: Trino strips the spaces. + ("char(3)", None), + ("timestamp(3)", pl.Datetime("ms")), + ("timestamp(6)", pl.Datetime("us")), + ("timestamp(9)", None), + ("timestamp(7) with time zone", None), + ("timestamp(6) with time zone", pl.Datetime("us", "UTC")), + ("uuid", None), + ("time(3)", None), + ("array(integer)", None), + ], +) +def test_trino_type_names(name, dtype): + assert _trino_dtype(name) == dtype + + +def test_trino_median_of_large_values(): + """The median of two middle values whose sum overflows a double halves each + one first; a group of one value is that value.""" + if not os.environ.get(TRINO_ENV): + pytest.skip(f"set {TRINO_ENV} to run against Trino") + df = pl.DataFrame( + { + "g": ["one", "two", "two", "even", "even", "even", "even"], + "v": [1e308, 1e308, 1.5e308, 1.0, 2.0, 4.0, 1e300], + } + ) + rows = ", ".join(f"('{g}', DOUBLE '{v!r}')" for g, v in df.iter_rows()) + query = f"SELECT * FROM (VALUES {rows}) AS t(g, v)" + register_source("_sql_test_trino_median_ref", df) + register_source( + "_sql_test_trino_median", + SQLSource(_trino_connect, dialect="trino", query=query), + ) + client = TestClient(app) + dash = Dashboard() + dash.add_figure().add_bar(labels="g", values="v", agg="median") + state, ev = _event("init", {}) + want = _post(client, dash, "_sql_test_trino_median_ref", state, ev) + got = _post(client, dash, "_sql_test_trino_median", state, ev) + assert not _diff(want, got) + + +@pytest.mark.parametrize( + "column", + [ + # Trino compares a uuid with no string, so a selection could not run. + "uuid '12151fd2-7586-11e9-8f9e-2a86e4085a59'", + # The driver rounds to microseconds, which the epoch in SQL does not. + "CAST(\"ts\" + INTERVAL '0.000000501' SECOND AS timestamp(9))", + ], + ids=["uuid", "timestamp9"], +) +def test_trino_column_without_dtype_fails_at_add_time(trino_source, column): + """A column of a type FlexViz cannot read is not in the schema, so a chart + on it fails when it is added.""" + query = f'SELECT {column} AS u, "w" FROM "{trino_source._table}"' + src = SQLSource(_trino_connect, dialect="trino", query=query) + assert list(src.schema) == ["w"] + fig = Dashboard(src).add_figure() + with pytest.raises(ValueError, match="of a type FlexViz cannot read"): + fig.add_histogram(x="u")