fix: support extra column properties when creating column object (#67)

eakmanrq · Jun 12, 2024 · ea6733e · ea6733e
1 parent fe4c367
commit ea6733e
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/sqlframe/spark/catalog.py b/sqlframe/spark/catalog.py
@@ -519,7 +519,10 @@ def listColumns(
                 )
                 for col in df.columns
             ]
-        return [Column(*x) for x in self._spark_catalog.listColumns(tableName, dbName)]
+        return [
+            Column(**{name: x._asdict()[name] for name in Column._fields})
+            for x in self._spark_catalog.listColumns(tableName, dbName)
+        ]
 
     def listFunctions(
         self, dbName: t.Optional[str] = None, pattern: t.Optional[str] = None

diff --git a/tests/integration/engines/bigquery/test_bigquery_session.py b/tests/integration/engines/bigquery/test_bigquery_session.py
@@ -16,5 +16,5 @@ def test_session_from_config():
     conn.cursor().execute("CREATE SCHEMA IF NOT EXISTS db1")
     conn.cursor().execute("CREATE TABLE IF NOT EXISTS db1.test_table (cola INT, colb STRING)")
     session = BigQuerySession.builder.config("default_dataset", "sqlframe.db1").getOrCreate()
-    columns = session.catalog.get_columns("test_table")
+    columns = session.catalog.get_columns("db1.test_table")
     assert columns == {"`cola`": exp.DataType.build("BIGINT"), "`colb`": exp.DataType.build("TEXT")}