feat: add normalization for missing table names extraction in SessionContext

Author: kosiew

python/datafusion/context.py

@@ -743,17 +743,33 @@ def sql_with_options(self, query: str, options: SQLOptions) -> DataFrame:
 
     @staticmethod
     def _extract_missing_table_names(err: Exception) -> list[str]:
+        def _normalize(names: list[Any]) -> list[str]:
+            tables: list[str] = []
+            for raw_name in names:
+                if not raw_name:
+                    continue
+                raw_str = str(raw_name)
+                tables.append(raw_str.rsplit(".", 1)[-1])
+            return tables
+
+        missing_tables = getattr(err, "missing_table_names", None)
+        if missing_tables is not None:
+            if isinstance(missing_tables, str):
+                candidates: list[Any] = [missing_tables]
+            else:
+                try:
+                    candidates = list(missing_tables)
+                except TypeError:
+                    candidates = [missing_tables]
+
+            return _normalize(candidates)
+
         message = str(err)
         matches = set()
         for pattern in (r"table '([^']+)' not found", r"No table named '([^']+)'"):
             matches.update(re.findall(pattern, message))
 
-        tables: list[str] = []
-        for raw_name in matches:
-            if not raw_name:
-                continue
-            tables.append(raw_name.rsplit(".", 1)[-1])
-        return tables
+        return _normalize(list(matches))
 
     def _register_python_tables(self, tables: list[str]) -> bool:
         registered_any = False

python/tests/test_context.py

@@ -268,6 +268,20 @@ def test_sql_missing_table_without_auto_register(ctx):
     assert "arrow_table" in missing_tables
 
 
+def test_extract_missing_table_names_from_attribute():
+    class MissingTablesError(Exception):
+        def __init__(self) -> None:
+            super().__init__("custom error")
+            self.missing_table_names = (
+                "catalog.schema.arrow_table",
+                "plain_table",
+            )
+
+    err = MissingTablesError()
+    missing_tables = SessionContext._extract_missing_table_names(err)
+    assert missing_tables == ["arrow_table", "plain_table"]
+
+
 def test_sql_auto_register_arrow_table():
     ctx = SessionContext(auto_register_python_variables=True)
     arrow_table = pa.Table.from_pydict({"value": [1, 2, 3]})  # noqa: F841