diff --git a/CHANGELOG.md b/CHANGELOG.md index 19c02aa..cf2ecc3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,7 @@ - New native writers `SVGWriter`, `PDFWriter`, and `HepWriter`, rendered directly by ggsql (via hephaestus) with no JavaScript round-trip. Each exposes `render(spec)` returning `str` (SVG) or `bytes` (PDF/HEP); render warnings are surfaced via `warnings.warn()`. - New `ggsql.save(spec, file, ...)` helper that writes a spec to disk, picking the writer from the file extension (`.svg`, `.pdf`, `.hep`, or `.json` for Vega-Lite). - Enabled the crate's `builtin-data` feature, making builtin datasets such as `ggsql:world` available. +- New `ggsql.CachingReader` that caches query results to avoid redundant source reads. Wraps a `duckdb://` connection string or any custom Python reader; reads are memoized with a configurable TTL and size budget, and derived `__ggsql_*` tables run on a private in-memory DuckDB cache. ### Changed diff --git a/python/ggsql/__init__.py b/python/ggsql/__init__.py index 76feee7..f237030 100644 --- a/python/ggsql/__init__.py +++ b/python/ggsql/__init__.py @@ -13,6 +13,7 @@ from narwhals.typing import IntoFrame from ggsql._ggsql import ( + CachingReader, DuckDBReader, HepWriter as _RustHepWriter, PDFWriter as _RustPDFWriter, @@ -27,6 +28,7 @@ # PyO3 classes default to __module__ = "builtins"; point them at their real # home so docs tooling (great-docs/griffe) can locate them. for _cls in ( + CachingReader, DuckDBReader, _RustHepWriter, _RustPDFWriter, @@ -40,6 +42,7 @@ __all__ = [ # Classes + "CachingReader", "DuckDBReader", "SVGWriter", "PDFWriter", diff --git a/src/lib.rs b/src/lib.rs index 2de8c59..35c8f6a 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -10,8 +10,9 @@ use arrow::ipc::reader::StreamReader; use arrow::ipc::writer::StreamWriter; use ggsql::dataframe::DataFrame; +use ggsql::reader::cache::{CacheConfig, CacheConfigOverride}; use ggsql::reader::Spec; -use ggsql::reader::{DuckDBReader as RustDuckDBReader, Reader}; +use ggsql::reader::{CachingReader as RustCachingReader, DuckDBReader as RustDuckDBReader, Reader}; use ggsql::validate::{validate as rust_validate, ValidationWarning}; use ggsql::writer::{ HepWriter as RustHepWriter, PdfWriter as RustPdfWriter, SvgWriter as RustSvgWriter, @@ -382,6 +383,215 @@ impl PyDuckDBReader { } } +// ============================================================================ +// PyCachingReader +// ============================================================================ + +/// Caching reader: wraps a primary reader with an in-memory DuckDB cache. +/// +/// Source reads go to the primary and are memoized (TTL + LRU byte budget); +/// all ggsql-generated derived SQL runs on the private in-memory cache. +/// Configuration is read from the environment (GGSQL_CACHE_DISABLED, +/// GGSQL_CACHE_TTL, GGSQL_CACHE_MAX_BYTES) with keyword arguments taking +/// precedence. +/// +/// Examples +/// -------- +/// >>> reader = CachingReader("duckdb://memory", ttl_secs=60) +/// >>> spec = reader.execute("SELECT 1 AS x, 2 AS y VISUALISE x, y DRAW point") +/// +/// >>> # Wrapping a custom Python reader (bridge protocol): +/// >>> class MyReader: +/// ... def execute_sql(self, sql): ... +/// ... def register(self, name, table, replace): ... +/// ... def unregister(self, name): ... +/// >>> reader = CachingReader(MyReader()) +#[pyclass(name = "CachingReader", unsendable)] +struct PyCachingReader { + inner: RustCachingReader, +} + +#[pymethods] +impl PyCachingReader { + /// Create a new caching reader. + /// + /// Parameters + /// ---------- + /// primary : str or object + /// Either a "duckdb://..." connection string (a fresh native + /// DuckDBReader is constructed internally) or a Python object + /// implementing the reader protocol: execute_sql(sql) -> + /// pyarrow.Table, register(name, table, replace), unregister(name). + /// ttl_secs : int, optional + /// Seconds before a memoized read is treated as stale. + /// Default 300 (or GGSQL_CACHE_TTL). + /// max_bytes : int, optional + /// Cumulative byte budget across memoized results before LRU + /// eviction. Default 512 MiB (or GGSQL_CACHE_MAX_BYTES). + /// enabled : bool, optional + /// When False, source reads always hit the primary: the memo is + /// neither consulted nor written. Note this only disables + /// memoization — registered tables still live in the private + /// in-memory cache, reads referencing them still route there, + /// and derived SQL still runs on the cache backend. + /// Default True (unless GGSQL_CACHE_DISABLED is set). + /// + /// Raises + /// ------ + /// ValueError + /// If a connection string has an unsupported scheme, or a + /// DuckDBReader instance is passed (pass its connection + /// string instead). + #[new] + #[pyo3(signature = (primary, *, ttl_secs=None, max_bytes=None, enabled=None))] + fn new( + primary: &Bound<'_, PyAny>, + ttl_secs: Option, + max_bytes: Option, + enabled: Option, + ) -> PyResult { + // Private in-memory DuckDB is the cache backend: derived tables, + // registered data, and memoized reads all live there. + let cache = RustDuckDBReader::from_connection_string("duckdb://memory") + .map_err(|e| PyErr::new::(e.to_string()))?; + + let (primary_reader, primary_uri): (Box, String) = if let Ok(uri) = + primary.extract::<&str>() + { + if !uri.starts_with("duckdb://") { + return Err(PyErr::new::(format!( + "Unsupported connection string {uri:?}: only duckdb:// primaries are currently supported" + ))); + } + let reader = RustDuckDBReader::from_connection_string(uri) + .map_err(|e| PyErr::new::(e.to_string()))?; + (Box::new(reader), uri.to_string()) + } else { + // Moving the inner reader out of a live pyclass isn't + // possible; guide the user to the connection-string form. + if primary.downcast::().is_ok() { + return Err(PyErr::new::( + "CachingReader cannot wrap an existing DuckDBReader; pass its connection string instead (e.g. CachingReader(\"duckdb://memory\"))", + )); + } + // The URI only feeds cache-key derivation, which just needs a + // stable-per-instance string; the type name suffices. + let type_name = primary.get_type().name()?.to_str()?.to_string(); + ( + Box::new(PyReaderBridge { + obj: primary.clone().unbind(), + }), + type_name, + ) + }; + + let config = CacheConfig::from_env().merge(CacheConfigOverride { + enabled, + ttl_secs, + max_bytes, + }); + + Ok(Self { + inner: RustCachingReader::with_config( + primary_reader, + Box::new(cache), + primary_uri, + "duckdb", + config, + ), + }) + } + + /// Register a pyarrow Table as a queryable table. + /// + /// The table lands in the in-memory cache; reads referencing its name + /// are routed to the cache rather than the primary. + /// + /// Parameters + /// ---------- + /// name : str + /// The table name to register under. + /// table : pyarrow.Table + /// The Arrow table to register. + /// replace : bool + /// Whether to replace an existing table of the same name. + #[pyo3(signature = (name, table, replace=false))] + fn register( + &self, + py: Python<'_>, + name: &str, + table: &Bound<'_, PyAny>, + replace: bool, + ) -> PyResult<()> { + let df = py_to_df(py, table)?; + self.inner + .register(name, df, replace) + .map_err(|e| PyErr::new::(e.to_string())) + } + + /// Unregister a previously registered table. + /// + /// Parameters + /// ---------- + /// name : str + /// The table name to unregister. + fn unregister(&self, name: &str) -> PyResult<()> { + self.inner + .unregister(name) + .map_err(|e| PyErr::new::(e.to_string())) + } + + /// Execute a SQL query against the source surface. + /// + /// Row-returning reads are memoized: an identical query within the TTL + /// is served from the cache without hitting the primary. + /// + /// Parameters + /// ---------- + /// sql : str + /// The SQL query to execute. + /// + /// Returns + /// ------- + /// pyarrow.Table + /// The query result as a pyarrow Table. + fn execute_sql(&self, py: Python<'_>, sql: &str) -> PyResult> { + let df = self + .inner + .execute_sql(sql) + .map_err(|e| PyErr::new::(e.to_string()))?; + df_to_py(py, &df) + } + + /// Execute a ggsql query and return the visualization specification. + /// + /// Base reads hit the (memoized) primary; derived tables are + /// materialized into the cache. Unlike a plain bridged reader, queries + /// with layer-level sources (VISUALISE ... FROM ) work, because + /// every layer source is staged through the cache. + /// + /// Parameters + /// ---------- + /// query : str + /// The ggsql query (SQL + VISUALISE clause). + /// + /// Returns + /// ------- + /// Spec + /// The resolved visualization specification ready for rendering. + /// + /// Raises + /// ------ + /// ValueError + /// If the query syntax is invalid or SQL execution fails. + fn execute(&self, query: &str) -> PyResult { + self.inner + .execute(query) + .map(|s| PySpec { inner: s }) + .map_err(|e| PyErr::new::(e.to_string())) + } +} + // ============================================================================ // PyVegaLiteWriter // ============================================================================ @@ -976,7 +1186,7 @@ fn validate(query: &str) -> PyResult { fn execute(query: &str, reader: &Bound<'_, PyAny>) -> PyResult { // Fast path: try all known native reader types // Add new native readers to this list as they're implemented - try_native_readers!(query, reader, PyDuckDBReader); + try_native_readers!(query, reader, PyDuckDBReader, PyCachingReader); // Bridge path: wrap Python object as Reader let bridge = PyReaderBridge { @@ -996,6 +1206,7 @@ fn execute(query: &str, reader: &Bound<'_, PyAny>) -> PyResult { fn _ggsql(m: &Bound<'_, PyModule>) -> PyResult<()> { // Classes m.add_class::()?; + m.add_class::()?; m.add_class::()?; m.add_class::()?; m.add_class::()?; diff --git a/tests/test_ggsql.py b/tests/test_ggsql.py index 22a11da..598b0df 100644 --- a/tests/test_ggsql.py +++ b/tests/test_ggsql.py @@ -10,6 +10,7 @@ import json import re +import time import warnings from pathlib import Path @@ -958,3 +959,130 @@ def test_builtin_dataset_renders_svg(self): ) svg = ggsql.SVGWriter().render(spec) assert " pa.Table: + self.calls += 1 + return self.conn.execute(sql).arrow() + + def register(self, name: str, table: pa.Table, _replace: bool) -> None: + self.conn.register(name, table) + + def unregister(self, name: str) -> None: + self.conn.unregister(name) + + +class TestCachingReader: + def test_create_with_connection_string(self): + reader = ggsql.CachingReader("duckdb://memory") + assert reader is not None + + def test_exported_at_package_level(self): + assert "CachingReader" in ggsql.__all__ + assert ggsql.CachingReader.__module__ == "ggsql._ggsql" + + def test_rejects_non_duckdb_scheme(self): + with pytest.raises(ValueError, match="duckdb://"): + ggsql.CachingReader("sqlite://memory") + + def test_rejects_duckdb_reader_instance(self): + inner = ggsql.DuckDBReader("duckdb://memory") + with pytest.raises(ValueError, match="connection string"): + ggsql.CachingReader(inner) + + def test_execute(self): + reader = ggsql.CachingReader("duckdb://memory") + spec = reader.execute("SELECT 1 AS x, 2 AS y VISUALISE x, y DRAW point") + assert spec.metadata()["rows"] == 1 + + def test_execute_sql(self): + reader = ggsql.CachingReader("duckdb://memory") + table = reader.execute_sql("SELECT 1 AS x") + assert table.num_rows == 1 + + def test_register_and_query(self): + reader = ggsql.CachingReader("duckdb://memory") + reader.register("mytab", pa.table({"x": [1, 2], "y": [3, 4]})) + spec = reader.execute("SELECT * FROM mytab VISUALISE x, y DRAW point") + assert spec.metadata()["rows"] == 2 + + def test_memoizes_identical_reads(self): + primary = _CountingReader() + reader = ggsql.CachingReader(primary) + reader.execute_sql("SELECT * FROM t") + reader.execute_sql("SELECT * FROM t") + assert primary.calls == 1 + + def test_distinct_reads_each_hit_primary(self): + primary = _CountingReader() + reader = ggsql.CachingReader(primary) + reader.execute_sql("SELECT * FROM t") + reader.execute_sql("SELECT v + 1 AS v FROM t") + assert primary.calls == 2 + + def test_ttl_expiry(self): + primary = _CountingReader() + reader = ggsql.CachingReader(primary, ttl_secs=1) + reader.execute_sql("SELECT * FROM t") + time.sleep(1.1) + reader.execute_sql("SELECT * FROM t") + assert primary.calls == 2 + + def test_disabled_kwarg(self): + primary = _CountingReader() + reader = ggsql.CachingReader(primary, enabled=False) + reader.execute_sql("SELECT * FROM t") + reader.execute_sql("SELECT * FROM t") + assert primary.calls == 2 + + def test_disabled_env_var(self, monkeypatch): + monkeypatch.setenv("GGSQL_CACHE_DISABLED", "1") + primary = _CountingReader() + reader = ggsql.CachingReader(primary) + reader.execute_sql("SELECT * FROM t") + reader.execute_sql("SELECT * FROM t") + assert primary.calls == 2 + + +class TestCachingReaderCapabilities: + """The scenarios querychat cannot serve today without CachingReader.""" + + def test_layer_level_source_through_bridge(self): + """VISUALISE ... FROM
works through a bridged primary. + + This is the querychat scenario: today querychat rejects any query + with a layer-level source because it cannot be replayed from one + final SQL result. Through CachingReader the layer source is staged + into the cache, so the query runs end to end. + """ + primary = _CountingReader() + reader = ggsql.CachingReader(primary) + spec = reader.execute("VISUALISE x DRAW histogram MAPPING v AS x FROM t") + assert spec is not None + + def test_registered_table_routes_to_cache(self): + """A registered table is cache-resident; reading it skips the primary.""" + primary = _CountingReader() + reader = ggsql.CachingReader(primary) + reader.register("only_in_cache", pa.table({"v": [1, 2, 3]})) + spec = reader.execute( + "VISUALISE x DRAW histogram MAPPING v AS x FROM only_in_cache" + ) + assert spec is not None + assert primary.calls == 0 + + def test_top_level_execute_fast_path(self): + """ggsql.execute() recognizes CachingReader as native (no bridge).""" + reader = ggsql.CachingReader("duckdb://memory") + spec = ggsql.execute("SELECT 1 AS x, 2 AS y VISUALISE x, y DRAW point", reader) + assert spec.metadata()["rows"] == 1