Skip to content

Commit cacba9b

Browse files
committed
Collect was occuring twice on repr
1 parent 3ba4d9d commit cacba9b

File tree

1 file changed

+1
-25
lines changed

1 file changed

+1
-25
lines changed

src/dataframe.rs

Lines changed: 1 addition & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -125,8 +125,6 @@ impl PyDataFrame {
125125
return Ok("No data to display".to_string());
126126
}
127127

128-
let df = self.df.as_ref().clone().limit(0, Some(10))?;
129-
let batches = wait_for_future(py, df.collect())?;
130128
let batches_as_displ =
131129
pretty::pretty_format_batches(&batches).map_err(py_datafusion_err)?;
132130

@@ -139,7 +137,7 @@ impl PyDataFrame {
139137
}
140138

141139
fn _repr_html_(&self, py: Python) -> PyDataFusionResult<String> {
142-
let (batches, mut has_more) = wait_for_future(
140+
let (batches, has_more) = wait_for_future(
143141
py,
144142
collect_record_batches_to_display(
145143
self.df.as_ref().clone(),
@@ -209,35 +207,13 @@ impl PyDataFrame {
209207
})
210208
.collect::<Result<Vec<_>, _>>()?;
211209

212-
let total_memory: usize = batches
213-
.iter()
214-
.map(|batch| batch.get_array_memory_size())
215-
.sum();
216210
let rows_per_batch = batches.iter().map(|batch| batch.num_rows());
217-
let total_rows = rows_per_batch.clone().sum();
218-
219-
let num_rows_to_display = match total_memory > MAX_TABLE_BYTES_TO_DISPLAY {
220-
true => {
221-
let ratio = MAX_TABLE_BYTES_TO_DISPLAY as f32 / total_memory as f32;
222-
let mut reduced_row_num = (total_rows as f32 * ratio).round() as usize;
223-
if reduced_row_num < MIN_TABLE_ROWS_TO_DISPLAY {
224-
reduced_row_num = MIN_TABLE_ROWS_TO_DISPLAY.min(total_rows);
225-
}
226-
227-
has_more = has_more || reduced_row_num < total_rows;
228-
reduced_row_num
229-
}
230-
false => total_rows,
231-
};
232211

233212
// We need to build up row by row for html
234213
let mut table_row = 0;
235214
for (batch_formatter, num_rows_in_batch) in batch_formatters.iter().zip(rows_per_batch) {
236215
for batch_row in 0..num_rows_in_batch {
237216
table_row += 1;
238-
if table_row > num_rows_to_display {
239-
break;
240-
}
241217
let mut cells = Vec::new();
242218
for (col, formatter) in batch_formatter.iter().enumerate() {
243219
let cell_data = formatter.value(batch_row).to_string();

0 commit comments

Comments
 (0)