Support Polars DataFrame export from QueryResult (#2985)

kuzudb · Mar 4, 2024 · 0d84c73 · 0d84c73
1 parent 04c554b
commit 0d84c73
Show file tree

Hide file tree

Showing 3 changed files with 444 additions and 247 deletions.
diff --git a/tools/python_api/requirements_dev.txt b/tools/python_api/requirements_dev.txt
@@ -1,6 +1,7 @@
 pybind11>=2.6.0
 pytest
 pandas
+polars
 networkx~=3.0.0
 numpy
 pyarrow==14.0.1

diff --git a/tools/python_api/src_py/query_result.py b/tools/python_api/src_py/query_result.py
@@ -98,6 +98,19 @@ def get_as_df(self):
 
         return self._query_result.getAsDF()
 
+    def get_as_pl(self):
+        """
+        Get the query result as a Polars DataFrame.
+
+        Returns
+        -------
+        polars.DataFrame
+            Query result as a Polars DataFrame.
+        """
+
+        import polars as pl
+        return pl.from_arrow(data=self.get_as_arrow(10_000))
+
     def get_as_arrow(self, chunk_size):
         """
         Get the query result as a PyArrow Table.