diff --git a/src/dataframe.rs b/src/dataframe.rs index 605e1161..fcce90d9 100644 --- a/src/dataframe.rs +++ b/src/dataframe.rs @@ -172,8 +172,7 @@ impl PyDataFrame { #[pyo3(signature = (num=20))] fn show(&self, py: Python, num: usize) -> PyResult<()> { let df = self.df.as_ref().clone().limit(0, Some(num))?; - let batches = wait_for_future(py, df.collect())?; - pretty::print_batches(&batches).map_err(|err| PyArrowException::new_err(err.to_string())) + print_dataframe(py, df) } /// Filter out duplicate rows @@ -217,8 +216,7 @@ impl PyDataFrame { #[pyo3(signature = (verbose=false, analyze=false))] fn explain(&self, py: Python, verbose: bool, analyze: bool) -> PyResult<()> { let df = self.df.as_ref().clone().explain(verbose, analyze)?; - let batches = wait_for_future(py, df.collect())?; - pretty::print_batches(&batches).map_err(|err| PyArrowException::new_err(err.to_string())) + print_dataframe(py, df) } /// Get the logical plan for this `DataFrame` @@ -389,3 +387,20 @@ impl PyDataFrame { Ok(wait_for_future(py, self.df.as_ref().clone().count())?) } } + +/// Print DataFrame +fn print_dataframe(py: Python, df: DataFrame) -> PyResult<()> { + // Get string representation of record batches + let batches = wait_for_future(py, df.collect())?; + let batches_as_string = pretty::pretty_format_batches(&batches); + let result = match batches_as_string { + Ok(batch) => format!("DataFrame()\n{batch}"), + Err(err) => format!("Error: {:?}", err.to_string()), + }; + + // Import the Python 'builtins' module to access the print function + // Note that println! does not print to the Python debug console and is not visible in notebooks for instance + let print = py.import("builtins")?.getattr("print")?; + print.call1((result,))?; + Ok(()) +}