feat: add SQL tests for local Arrow, Pandas, and Polars dataframes

Author: kosiew

python/tests/test_context.py

@@ -348,6 +348,45 @@ def test_from_pandas(ctx):
     assert df.collect()[0].num_rows == 3
 
 
+def test_sql_from_local_arrow_table(ctx):
+    arrow_table = pa.Table.from_pydict({"a": [1, 2], "b": ["x", "y"]})
+
+    result = ctx.sql("SELECT * FROM arrow_table ORDER BY a").collect()
+    actual = pa.Table.from_batches(result)
+    expected = pa.Table.from_pydict({"a": [1, 2], "b": ["x", "y"]})
+
+    assert actual.equals(expected)
+
+
+def test_sql_from_local_pandas_dataframe(ctx):
+    pd = pytest.importorskip("pandas")
+    pandas_df = pd.DataFrame({"a": [3, 1], "b": ["z", "y"]})
+
+    result = ctx.sql("SELECT * FROM pandas_df ORDER BY a").collect()
+    actual = pa.Table.from_batches(result)
+    expected = pa.Table.from_pydict({"a": [1, 3], "b": ["y", "z"]})
+
+    assert actual.equals(expected)
+
+
+def test_sql_from_local_polars_dataframe(ctx):
+    pl = pytest.importorskip("polars")
+    polars_df = pl.DataFrame({"a": [2, 1], "b": ["beta", "alpha"]})
+
+    result = ctx.sql("SELECT * FROM polars_df ORDER BY a").collect()
+    actual = pa.Table.from_batches(result)
+    expected = pa.Table.from_pydict({"a": [1, 2], "b": ["alpha", "beta"]})
+
+    assert actual.equals(expected)
+
+
+def test_sql_from_local_unsupported_object(ctx):
+    unsupported = object()
+
+    with pytest.raises(Exception, match="table 'unsupported' not found"):
+        ctx.sql("SELECT * FROM unsupported").collect()
+
+
 def test_from_polars(ctx):
     # create a dataframe from Polars dataframe
     pd = pytest.importorskip("polars")