diff --git a/python/datafusion/tests/test_context.py b/python/datafusion/tests/test_context.py index 0184280c2..4af00a3b4 100644 --- a/python/datafusion/tests/test_context.py +++ b/python/datafusion/tests/test_context.py @@ -466,6 +466,13 @@ def test_table_exist(ctx): assert ctx.table_exist("t") is True +def test_table_not_found(ctx): + from uuid import uuid4 + + with pytest.raises(KeyError): + ctx.table(f"not-found-{uuid4()}") + + def test_read_json(ctx): path = os.path.dirname(os.path.abspath(__file__)) diff --git a/src/context.rs b/src/context.rs index 4433d94c2..3ab783495 100644 --- a/src/context.rs +++ b/src/context.rs @@ -765,7 +765,8 @@ impl PySessionContext { } pub fn table(&self, name: &str, py: Python) -> PyResult { - let x = wait_for_future(py, self.ctx.table(name)).map_err(DataFusionError::from)?; + let x = wait_for_future(py, self.ctx.table(name)) + .map_err(|e| PyKeyError::new_err(e.to_string()))?; Ok(PyDataFrame::new(x)) }