timsaucer commented on code in PR #1036:
URL:
https://github.com/apache/datafusion-python/pull/1036#discussion_r2008764438
##########
src/dataframe.rs:
##########
@@ -771,3 +871,82 @@ fn record_batch_into_schema(
RecordBatch::try_new(schema, data_arrays)
}
+
+/// This is a helper function to return the first non-empty record batch from
executing a DataFrame.
+/// It additionally returns a bool, which indicates if there are more record
batches available.
+/// We do this so we can determine if we should indicate to the user that the
data has been
+/// truncated. This collects until we have achived both of these two conditions
+///
+/// - We have collected our minimum number of rows
+/// - We have reached our limit, either data size or maximum number of rows
+///
+/// Otherwise it will return when the stream has exhausted. If you want a
specific number of
+/// rows, set min_rows == max_rows.
+async fn collect_record_batches_to_display(
+ df: DataFrame,
+ min_rows: usize,
+ max_rows: usize,
+) -> Result<(Vec<RecordBatch>, bool), DataFusionError> {
+ let mut stream = df.execute_stream().await?;
+ let mut size_estimate_so_far = 0;
+ let mut rows_so_far = 0;
+ let mut record_batches = Vec::default();
+ let mut has_more = false;
+
+ while (size_estimate_so_far < MAX_TABLE_BYTES_TO_DISPLAY && rows_so_far <
max_rows)
+ || rows_so_far < min_rows
+ {
+ let mut rb = match stream.next().await {
+ None => {
+ break;
+ }
+ Some(Ok(r)) => r,
+ Some(Err(e)) => return Err(e),
+ };
+
+ let mut rows_in_rb = rb.num_rows();
+ if rows_in_rb > 0 {
+ size_estimate_so_far += rb.get_array_memory_size();
+
+ if size_estimate_so_far > MAX_TABLE_BYTES_TO_DISPLAY {
+ let ratio = MAX_TABLE_BYTES_TO_DISPLAY as f32 /
size_estimate_so_far as f32;
+ let total_rows = rows_in_rb + rows_so_far;
+
+ let mut reduced_row_num = (total_rows as f32 * ratio).round()
as usize;
Review Comment:
Yes. And the data size is an estimate as well. The point is to get a general
ball park and not necessarily an exact measure. It should still indicate that
that data have been truncated.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscr...@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscr...@datafusion.apache.org
For additional commands, e-mail: github-h...@datafusion.apache.org
