Enhance Pandas Query Engine Output Processor (#16052)

* Enhance output processor to temporarily adjust display options. * Make format changes. * Update pyproject.toml --------- Co-authored-by: Logan <[email protected]> Co-authored-by: Massimiliano Pippi <[email protected]>
run-llama · Sep 20, 2024 · 2e99701 · 2e99701
1 parent e17a25d
commit 2e99701
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 2 deletions.
diff --git a/llama-index-experimental/llama_index/experimental/query_engine/pandas/output_parser.py b/llama-index-experimental/llama_index/experimental/query_engine/pandas/output_parser.py
@@ -49,10 +49,19 @@ def default_output_processor(
         try:
             # str(pd.dataframe) will truncate output by display.max_colwidth
             # set width temporarily to extract more text
+            current_max_colwidth = pd.get_option("display.max_colwidth")
+            current_max_rows = pd.get_option("display.max_rows")
+            current_max_columns = pd.get_option("display.max_columns")
             if "max_colwidth" in output_kwargs:
                 pd.set_option("display.max_colwidth", output_kwargs["max_colwidth"])
+            if "max_rows" in output_kwargs:
+                pd.set_option("display.max_rows", output_kwargs["max_rows"])
+            if "max_columns" in output_kwargs:
+                pd.set_option("display.max_columns", output_kwargs["max_columns"])
             output_str = str(safe_eval(module_end_str, global_vars, local_vars))
-            pd.reset_option("display.max_colwidth")
+            pd.set_option("display.max_colwidth", current_max_colwidth)
+            pd.set_option("display.max_rows", current_max_rows)
+            pd.set_option("display.max_columns", current_max_columns)
             return output_str
 
         except Exception:

diff --git a/llama-index-experimental/pyproject.toml b/llama-index-experimental/pyproject.toml
@@ -25,7 +25,7 @@ exclude = ["**/BUILD"]
 license = "MIT"
 name = "llama-index-experimental"
 readme = "README.md"
-version = "0.3.2"
+version = "0.3.3"
 
 [tool.poetry.dependencies]
 python = ">=3.10,<4.0"