Skip to content

Commit 065fa40

Browse files
committed
feat: Enhance PyDataFrame to support customizable display options
- Updated `PyDataFrame` constructor to accept a `PyDataframeDisplayConfig` parameter for improved DataFrame display customization. - Modified multiple methods in `PySessionContext` to pass the display configuration when creating `PyDataFrame` instances, ensuring consistent display settings across different DataFrame operations.
1 parent 5dfb9ce commit 065fa40

File tree

2 files changed

+34
-14
lines changed

2 files changed

+34
-14
lines changed

src/context.rs

Lines changed: 27 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -457,7 +457,7 @@ impl PySessionContext {
457457
pub fn sql(&mut self, query: &str, py: Python) -> PyDataFusionResult<PyDataFrame> {
458458
let result = self.ctx.sql(query);
459459
let df = wait_for_future(py, result)?;
460-
Ok(PyDataFrame::new(df))
460+
Ok(PyDataFrame::new(df, self.ctx.display_config.clone()))
461461
}
462462

463463
#[pyo3(signature = (query, options=None))]
@@ -474,7 +474,7 @@ impl PySessionContext {
474474
};
475475
let result = self.ctx.sql_with_options(query, options);
476476
let df = wait_for_future(py, result)?;
477-
Ok(PyDataFrame::new(df))
477+
Ok(PyDataFrame::new(df, self.ctx.display_config.clone()))
478478
}
479479

480480
#[pyo3(signature = (partitions, name=None, schema=None))]
@@ -509,13 +509,16 @@ impl PySessionContext {
509509

510510
let table = wait_for_future(py, self._table(&table_name))?;
511511

512-
let df = PyDataFrame::new(table);
512+
let df = PyDataFrame::new(table, self.ctx.display_config.clone());
513513
Ok(df)
514514
}
515515

516516
/// Create a DataFrame from an existing logical plan
517517
pub fn create_dataframe_from_logical_plan(&mut self, plan: PyLogicalPlan) -> PyDataFrame {
518-
PyDataFrame::new(DataFrame::new(self.ctx.state(), plan.plan.as_ref().clone()))
518+
PyDataFrame::new(
519+
DataFrame::new(self.ctx.state(), plan.plan.as_ref().clone()),
520+
self.ctx.display_config.clone(),
521+
)
519522
}
520523

521524
/// Construct datafusion dataframe from Python list
@@ -883,15 +886,18 @@ impl PySessionContext {
883886
pub fn table(&self, name: &str, py: Python) -> PyResult<PyDataFrame> {
884887
let x = wait_for_future(py, self.ctx.table(name))
885888
.map_err(|e| PyKeyError::new_err(e.to_string()))?;
886-
Ok(PyDataFrame::new(x))
889+
Ok(PyDataFrame::new(x, self.ctx.display_config.clone()))
887890
}
888891

889892
pub fn table_exist(&self, name: &str) -> PyDataFusionResult<bool> {
890893
Ok(self.ctx.table_exist(name)?)
891894
}
892895

893896
pub fn empty_table(&self) -> PyDataFusionResult<PyDataFrame> {
894-
Ok(PyDataFrame::new(self.ctx.read_empty()?))
897+
Ok(
898+
PyDataFrame::new(self.ctx.read_empty()?),
899+
self.ctx.display_config.clone(),
900+
)
895901
}
896902

897903
pub fn session_id(&self) -> String {
@@ -926,7 +932,7 @@ impl PySessionContext {
926932
let result = self.ctx.read_json(path, options);
927933
wait_for_future(py, result)?
928934
};
929-
Ok(PyDataFrame::new(df))
935+
Ok(PyDataFrame::new(df, self.ctx.display_config.clone()))
930936
}
931937

932938
#[allow(clippy::too_many_arguments)]
@@ -971,12 +977,18 @@ impl PySessionContext {
971977
let paths = path.extract::<Vec<String>>()?;
972978
let paths = paths.iter().map(|p| p as &str).collect::<Vec<&str>>();
973979
let result = self.ctx.read_csv(paths, options);
974-
let df = PyDataFrame::new(wait_for_future(py, result)?);
980+
let df = PyDataFrame::new(
981+
wait_for_future(py, result)?,
982+
self.ctx.display_config.clone(),
983+
);
975984
Ok(df)
976985
} else {
977986
let path = path.extract::<String>()?;
978987
let result = self.ctx.read_csv(path, options);
979-
let df = PyDataFrame::new(wait_for_future(py, result)?);
988+
let df = PyDataFrame::new(
989+
wait_for_future(py, result)?,
990+
self.ctx.display_config.clone(),
991+
);
980992
Ok(df)
981993
}
982994
}
@@ -1014,7 +1026,10 @@ impl PySessionContext {
10141026
.collect();
10151027

10161028
let result = self.ctx.read_parquet(path, options);
1017-
let df = PyDataFrame::new(wait_for_future(py, result)?);
1029+
let df = PyDataFrame::new(
1030+
wait_for_future(py, result)?,
1031+
self.ctx.display_config.clone(),
1032+
);
10181033
Ok(df)
10191034
}
10201035

@@ -1039,12 +1054,12 @@ impl PySessionContext {
10391054
let read_future = self.ctx.read_avro(path, options);
10401055
wait_for_future(py, read_future)?
10411056
};
1042-
Ok(PyDataFrame::new(df))
1057+
Ok(PyDataFrame::new(df, self.ctx.display_config.clone()))
10431058
}
10441059

10451060
pub fn read_table(&self, table: &PyTable) -> PyDataFusionResult<PyDataFrame> {
10461061
let df = self.ctx.read_table(table.table())?;
1047-
Ok(PyDataFrame::new(df))
1062+
Ok(PyDataFrame::new(df, self.ctx.display_config.clone()))
10481063
}
10491064

10501065
fn __repr__(&self) -> PyResult<String> {

src/dataframe.rs

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,7 @@ use pyo3::types::{PyCapsule, PyTuple, PyTupleMethods};
4343
use tokio::task::JoinHandle;
4444

4545
use crate::catalog::PyTable;
46+
use crate::context::PyDataframeDisplayConfig;
4647
use crate::errors::{py_datafusion_err, PyDataFusionError};
4748
use crate::expr::sort_expr::to_sort_expressions;
4849
use crate::physical_plan::PyExecutionPlan;
@@ -83,12 +84,16 @@ const MAX_LENGTH_CELL_WITHOUT_MINIMIZE: usize = 25;
8384
#[derive(Clone)]
8485
pub struct PyDataFrame {
8586
df: Arc<DataFrame>,
87+
display_config: PyDataframeDisplayConfig,
8688
}
8789

8890
impl PyDataFrame {
8991
/// creates a new PyDataFrame
90-
pub fn new(df: DataFrame) -> Self {
91-
Self { df: Arc::new(df) }
92+
pub fn new(df: DataFrame, display_config: PyDataframeDisplayConfig) -> Self {
93+
Self {
94+
df: Arc::new(df),
95+
display_config,
96+
}
9297
}
9398
}
9499

0 commit comments

Comments
 (0)