feat(python): support data= dict parameter on execute() for inline DataFrames

Allows passing a dict of DataFrames to reader.execute() and the module-level
execute() function, which are registered before query execution and unregistered
afterward (cleanup happens even on error).

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: cpsievert

ggsql-python/python/ggsql/_ggsql.pyi

@@ -94,7 +94,12 @@ class DuckDBReader:
         """
         ...
 
-    def execute(self, query: str) -> Spec:
+    def execute(
+        self,
+        query: str,
+        *,
+        data: dict[str, pl.DataFrame] | None = None,
+    ) -> Spec:
         """Execute a ggsql query and return the visualization specification.
 
         This is the main entry point for creating visualizations. It parses
@@ -105,6 +110,10 @@ class DuckDBReader:
         ----------
         query
             The ggsql query (SQL + VISUALISE clause).
+        data
+            Optional dictionary mapping table names to DataFrames. Tables are
+            registered before execution and unregistered afterward (even on
+            error).
 
         Returns
         -------
@@ -385,7 +394,12 @@ def validate(query: str) -> Validated:
     """
     ...
 
-def execute(query: str, reader: object) -> Spec:
+def execute(
+    query: str,
+    reader: object,
+    *,
+    data: dict[str, pl.DataFrame] | None = None,
+) -> Spec:
     """Execute a ggsql query with a reader (native or custom Python object).
 
     This is a convenience function for custom readers. For native readers,
@@ -399,6 +413,10 @@ def execute(query: str, reader: object) -> Spec:
         The database reader to execute SQL against. Can be a native
         ``DuckDBReader`` for optimal performance, or any Python object with
         an ``execute_sql(sql: str) -> polars.DataFrame`` method.
+    data
+        Optional dictionary mapping table names to DataFrames. Tables are
+        registered before execution and unregistered afterward (even on
+        error).
 
     Returns
     -------

ggsql-python/src/lib.rs

@@ -337,6 +337,9 @@ impl PyDuckDBReader {
     /// ----------
     /// query : str
     ///     The ggsql query (SQL + VISUALISE clause).
+    /// data : dict[str, polars.DataFrame] | None
+    ///     Optional dictionary mapping table names to DataFrames. Tables are
+    ///     registered before execution and unregistered afterward (even on error).
     ///
     /// Returns
     /// -------
@@ -354,11 +357,48 @@ impl PyDuckDBReader {
     /// >>> spec = reader.execute("SELECT 1 AS x, 2 AS y VISUALISE x, y DRAW point")
     /// >>> writer = VegaLiteWriter()
     /// >>> json_output = writer.render(spec)
-    fn execute(&self, query: &str) -> PyResult<PySpec> {
-        self.inner
+    #[pyo3(signature = (query, *, data=None))]
+    fn execute(&self, py: Python<'_>, query: &str, data: Option<&Bound<'_, PyDict>>) -> PyResult<PySpec> {
+        // Register DataFrames from data dict
+        let registered_names = if let Some(data_dict) = data {
+            self.register_data_dict(py, data_dict)?
+        } else {
+            vec![]
+        };
+
+        // Execute query (capture result, don't return early)
+        let result = self.inner
             .execute(query)
             .map(|s| PySpec { inner: s })
-            .map_err(ggsql_err_to_py)
+            .map_err(ggsql_err_to_py);
+
+        // Cleanup: unregister temporary tables (even on error)
+        for name in &registered_names {
+            let _ = self.inner.unregister(name);
+        }
+
+        result
+    }
+}
+
+impl PyDuckDBReader {
+    /// Register DataFrames from a Python dict. Returns list of registered names for cleanup.
+    /// This is a private Rust helper, not exposed to Python.
+    fn register_data_dict(
+        &self,
+        py: Python<'_>,
+        data: &Bound<'_, PyDict>,
+    ) -> PyResult<Vec<String>> {
+        let mut names = Vec::new();
+        for (key, value) in data.iter() {
+            let name: String = key.extract()?;
+            let df = py_to_polars(py, &value)?;
+            self.inner
+                .register(&name, df, true)
+                .map_err(ggsql_err_to_py)?;
+            names.push(name);
+        }
+        Ok(names)
     }
 }
 
@@ -729,6 +769,9 @@ fn validate(query: &str) -> PyResult<PyValidated> {
 ///     The database reader to execute SQL against. Can be a native Reader
 ///     for optimal performance, or any Python object with an
 ///     `execute_sql(sql: str) -> polars.DataFrame` method.
+/// data : dict[str, polars.DataFrame] | None
+///     Optional dictionary mapping table names to DataFrames. Tables are
+///     registered before execution and unregistered afterward (even on error).
 ///
 /// Returns
 /// -------
@@ -755,19 +798,80 @@ fn validate(query: &str) -> PyResult<PyValidated> {
 /// >>> reader = MyReader()
 /// >>> spec = execute("SELECT * FROM data VISUALISE x, y DRAW point", reader)
 #[pyfunction]
-fn execute(query: &str, reader: &Bound<'_, PyAny>) -> PyResult<PySpec> {
-    // 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);
+#[pyo3(signature = (query, reader, *, data=None))]
+fn execute(py: Python<'_>, query: &str, reader: &Bound<'_, PyAny>, data: Option<&Bound<'_, PyDict>>) -> PyResult<PySpec> {
+    // Native reader fast path: DuckDBReader
+    // Note: we can't use the try_native_readers! macro here because it uses `return`
+    // which would skip cleanup of registered tables.
+    if let Ok(native) = reader.downcast::<PyDuckDBReader>() {
+        // Register DataFrames if provided
+        let registered_names = if let Some(data_dict) = data {
+            native.borrow().register_data_dict(py, data_dict)?
+        } else {
+            vec![]
+        };
+
+        // Execute (capture result for cleanup)
+        let result = native.borrow().inner.execute(query)
+            .map(|s| PySpec { inner: s })
+            .map_err(ggsql_err_to_py);
+
+        // Cleanup: unregister temporary tables (even on error)
+        for name in &registered_names {
+            let _ = native.borrow().inner.unregister(name);
+        }
+
+        return result;
+    }
 
     // Bridge path: wrap Python object as Reader
+    // Register DataFrames if provided
+    let registered_names = if let Some(data_dict) = data {
+        register_data_on_reader(py, reader, data_dict)?
+    } else {
+        vec![]
+    };
+
     let bridge = PyReaderBridge {
         obj: reader.clone().unbind(),
     };
-    bridge
+    let result = bridge
         .execute(query)
         .map(|s| PySpec { inner: s })
-        .map_err(ggsql_err_to_py)
+        .map_err(ggsql_err_to_py);
+
+    // Cleanup for bridge path
+    for name in &registered_names {
+        let _ = call_unregister(py, reader, name);
+    }
+
+    result
+}
+
+/// Register DataFrames from a Python dict onto a Python reader object.
+/// Returns list of registered names for cleanup.
+fn register_data_on_reader(
+    py: Python<'_>,
+    reader: &Bound<'_, PyAny>,
+    data: &Bound<'_, PyDict>,
+) -> PyResult<Vec<String>> {
+    let mut names = Vec::new();
+    for (key, value) in data.iter() {
+        let name: String = key.extract()?;
+        let df = py_to_polars(py, &value)?;
+        let py_df = polars_to_py(py, &df)?;
+        reader.call_method("register", (&name, py_df, true), None)?;
+        names.push(name);
+    }
+    Ok(names)
+}
+
+/// Call unregister on a reader if the method exists.
+fn call_unregister(_py: Python<'_>, reader: &Bound<'_, PyAny>, name: &str) -> PyResult<()> {
+    if reader.hasattr("unregister")? {
+        reader.call_method1("unregister", (name,))?;
+    }
+    Ok(())
 }
 
 // ============================================================================

ggsql-python/tests/test_ggsql.py

@@ -656,6 +656,83 @@ def test_render_chart_facet(self):
         assert isinstance(chart, altair.FacetChart)
 
 
+class TestExecuteWithData:
+    """Tests for reader.execute() with data= parameter."""
+
+    def test_execute_with_single_dataframe(self):
+        """Can pass a single DataFrame via data dict."""
+        reader = ggsql.DuckDBReader("duckdb://memory")
+        df = pl.DataFrame({"x": [1, 2, 3], "y": [10, 20, 30]})
+        spec = reader.execute(
+            "SELECT * FROM mydata VISUALISE x, y DRAW point",
+            data={"mydata": df},
+        )
+        assert spec.metadata()["rows"] == 3
+
+    def test_execute_with_multiple_dataframes(self):
+        """Can pass multiple DataFrames via data dict."""
+        reader = ggsql.DuckDBReader("duckdb://memory")
+        df1 = pl.DataFrame({"id": [1, 2, 3], "y": [10, 20, 30]})
+        df2 = pl.DataFrame({"id": [2, 3], "category": ["A", "B"]})
+        spec = reader.execute(
+            "SELECT t1.id AS x, t1.y FROM t1 JOIN t2 ON t1.id = t2.id "
+            "VISUALISE x, y DRAW point",
+            data={"t1": df1, "t2": df2},
+        )
+        assert spec.metadata()["rows"] == 2
+
+    def test_execute_with_data_cleans_up(self):
+        """DataFrames passed via data= are unregistered after execution."""
+        reader = ggsql.DuckDBReader("duckdb://memory")
+        df = pl.DataFrame({"x": [1, 2, 3], "y": [10, 20, 30]})
+        reader.execute(
+            "SELECT * FROM temp VISUALISE x, y DRAW point",
+            data={"temp": df},
+        )
+        # Table should be cleaned up — querying it should fail
+        with pytest.raises((ggsql.ReaderError, ValueError)):
+            reader.execute_sql("SELECT * FROM temp")
+
+    def test_execute_with_data_cleans_up_on_error(self):
+        """DataFrames are unregistered even if execution fails."""
+        reader = ggsql.DuckDBReader("duckdb://memory")
+        df = pl.DataFrame({"x": [1, 2, 3], "y": [10, 20, 30]})
+        with pytest.raises((ggsql.ParseError, ggsql.ValidationError, ValueError)):
+            reader.execute(
+                "SELECT * FROM temp VISUALISE DRAW not_a_geom",
+                data={"temp": df},
+            )
+        # Table should still be cleaned up
+        with pytest.raises((ggsql.ReaderError, ValueError)):
+            reader.execute_sql("SELECT * FROM temp")
+
+    def test_execute_without_data_still_works(self):
+        """Calling execute() without data= still works as before."""
+        reader = ggsql.DuckDBReader("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_with_empty_data(self):
+        """Passing empty data= dict works fine."""
+        reader = ggsql.DuckDBReader("duckdb://memory")
+        spec = reader.execute(
+            "SELECT 1 AS x, 2 AS y VISUALISE x, y DRAW point",
+            data={},
+        )
+        assert spec.metadata()["rows"] == 1
+
+    def test_module_execute_with_data(self):
+        """Module-level execute() also supports data= parameter."""
+        reader = ggsql.DuckDBReader("duckdb://memory")
+        df = pl.DataFrame({"x": [1, 2, 3], "y": [10, 20, 30]})
+        spec = ggsql.execute(
+            "SELECT * FROM mydata VISUALISE x, y DRAW point",
+            reader,
+            data={"mydata": df},
+        )
+        assert spec.metadata()["rows"] == 3
+
+
 class TestTypeStubs:
     """Tests for type stub presence and correctness."""