From 2245043bbd15fa4e0daf8c0bd24c56c4c4eb3c6b Mon Sep 17 00:00:00 2001
From: Bishwajit <147722855+bishwajit-db@users.noreply.github.com>
Date: Wed, 8 May 2024 14:47:34 +0530
Subject: [PATCH] Maintain PySpark compatibility for
 databricks.labs.lsql.core.Row (#99)

Add asDict to databricks.labs.lsql.core.Row to maintain PySpark
compatibility
---
 src/databricks/labs/lsql/backends.py | 1 +
 src/databricks/labs/lsql/core.py     | 4 ++++
 2 files changed, 5 insertions(+)

diff --git a/src/databricks/labs/lsql/backends.py b/src/databricks/labs/lsql/backends.py
index cad73f1d..25e28fcc 100644
--- a/src/databricks/labs/lsql/backends.py
+++ b/src/databricks/labs/lsql/backends.py
@@ -207,6 +207,7 @@ def fetch(self, sql: str, *, catalog: str | None = None, schema: str | None = No
                 self._spark.sql(f"USE CATALOG {catalog}")
             if schema:
                 self._spark.sql(f"USE SCHEMA {schema}")
+            # TODO: pyspark.sql.Row is being returned instead of databricks.labs.lsql.core.Row
             return iter(self._spark.sql(sql).collect())
         except Exception as e:
             error_message = str(e)
diff --git a/src/databricks/labs/lsql/core.py b/src/databricks/labs/lsql/core.py
index 541c411d..38e27dd7 100644
--- a/src/databricks/labs/lsql/core.py
+++ b/src/databricks/labs/lsql/core.py
@@ -66,6 +66,10 @@ def as_dict(self) -> dict[str, Any]:
         """Convert the row to a dictionary with the same conventions as Databricks SDK."""
         return dict(zip(self.__columns__, self, strict=True))
 
+    # PySpark's compatibility
+    def asDict(self, recursive: bool = False) -> dict[str, Any]:
+        return self.as_dict()
+
     def __eq__(self, other):
         """Check if the rows are equal."""
         if not isinstance(other, Row):