#466 feat: Python Client Supports Multiple File Output (#478)

Updates the Python client to support multiple file on output. Writing the tests for this proved difficult so I went with manually testing: a single result ✅ and multiple results ✅
kaskada-ai · Jul 7, 2023 · 741bfe2 · 741bfe2
2 parents 8e80987 + 41f7555
commit 741bfe2
Showing 1 changed file with 19 additions and 23 deletions.
diff --git a/clients/python/src/fenlmagic/__init__.py b/clients/python/src/fenlmagic/__init__.py
@@ -145,30 +145,26 @@ def fenl(self, arg, cell=None):
                 client=self.client,
             )
             query_result = QueryResult(expression, resp)
-
             if not dry_run and render_dataframe:
-                # TODO: figure out how to support reading from multiple parquet paths
-                if (
-                    len(
-                        query_result.query_response.destination.object_store.output_paths.paths
-                    )
-                    > 0
-                ):
-                    if response_as == query.ResponseType.FILE_TYPE_PARQUET:
-                        df = pandas.read_parquet(
-                            query_result.query_response.destination.object_store.output_paths.paths[
-                                0
-                            ],
-                            engine="pyarrow",
-                        )
-                        query_result.set_dataframe(df)
-                    elif response_as == query.ResponseType.FILE_TYPE_CSV:
-                        df = pandas.read_csv(
-                            query_result.query_response.destination.object_store.output_paths.paths[
-                                0
-                            ],
-                        )
-                        query_result.set_dataframe(df)
+                output_paths = (
+                    query_result.query_response.destination.object_store.output_paths.paths
+                )
+                if len(output_paths) > 0:
+                    dataframes = []
+                    for parquet_output_path in output_paths:
+                        if response_as == query.ResponseType.FILE_TYPE_PARQUET:
+                            df = pandas.read_parquet(
+                                parquet_output_path, engine="pyarrow"
+                            )
+                            dataframes.append(df)
+                        elif response_as == query.ResponseType.FILE_TYPE_CSV:
+                            df = pandas.read_csv(parquet_output_path, engine="pyarrow")
+                            dataframes.append(df)
+                        else:
+                            raise NotImplementedError(
+                                f"unknown response type: {response_as}", response_as
+                            )
+                    query_result.set_dataframe(pandas.concat(dataframes))
 
             if var is not None:
                 IPython.get_ipython().push({var: query_result})