test: add unit test for automatic registration of multiple Python tables in SQL queries

Author: kosiew

python/datafusion/context.py

@@ -710,21 +710,25 @@ def _execute_sql() -> DataFrame:
                 self.ctx.sql_with_options(query, options.options_internal)
             )
 
-        try:
-            return _execute_sql()
-        except Exception as err:
-            if not getattr(self, "_auto_python_table_lookup", False):
-                raise
+        auto_lookup_enabled = getattr(self, "_auto_python_table_lookup", False)
+
+        while True:
+            try:
+                return _execute_sql()
+            except Exception as err:
+                if not auto_lookup_enabled:
+                    raise
 
-            missing_tables = self._extract_missing_table_names(err)
-            if not missing_tables:
-                raise
+                missing_tables = self._extract_missing_table_names(err)
+                if not missing_tables:
+                    raise
 
-            registered = self._register_python_tables(missing_tables)
-            if not registered:
-                raise
+                registered = self._register_python_tables(missing_tables)
+                if not registered:
+                    raise
 
-            return _execute_sql()
+                # Retry to allow registering additional tables referenced in the query.
+                continue
 
     def sql_with_options(self, query: str, options: SQLOptions) -> DataFrame:
         """Create a :py:class:`~datafusion.dataframe.DataFrame` from SQL query text.

python/tests/test_context.py

@@ -294,6 +294,35 @@ def test_sql_auto_register_arrow_table():
     assert result[0].column(0).to_pylist()[0] == 6
 
 
+def test_sql_auto_register_multiple_tables_single_query():
+    ctx = SessionContext(auto_register_python_objects=True)
+
+    customers = pa.Table.from_pydict(  # noqa: F841
+        {"customer_id": [1, 2], "name": ["Alice", "Bob"]}
+    )
+    orders = pa.Table.from_pydict(  # noqa: F841
+        {"order_id": [100, 200], "customer_id": [1, 2]}
+    )
+
+    result = ctx.sql(
+        """
+        SELECT c.customer_id, o.order_id
+        FROM customers c
+        JOIN orders o ON c.customer_id = o.customer_id
+        ORDER BY o.order_id
+        """
+    ).collect()
+
+    actual = pa.Table.from_batches(result)
+    expected = pa.Table.from_pydict(
+        {"customer_id": [1, 2], "order_id": [100, 200]}
+    )
+
+    assert actual.equals(expected)
+    assert ctx.table_exist("customers")
+    assert ctx.table_exist("orders")
+
+
 def test_sql_auto_register_arrow_outer_scope():
     ctx = SessionContext()
     ctx.auto_register_python_variables = True