cleanup

Eventual-Inc · Feb 15, 2024 · e8ab252 · e8ab252
1 parent 5dfa3d2
commit e8ab252
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 9 deletions.
diff --git a/daft/dataframe/dataframe.py b/daft/dataframe/dataframe.py
@@ -242,9 +242,7 @@ def _populate_preview(self) -> None:
             preview_results = preview_df._result
             assert preview_results is not None
 
-            print("here")
             preview_partition = preview_results._get_merged_vpartition()
-            print("not here")
             self._preview = DataFramePreview(
                 preview_partition=preview_partition,
                 dataframe_num_rows=len(self),

diff --git a/tests/cookbook/test_write.py b/tests/cookbook/test_write.py
@@ -22,7 +22,7 @@ def test_parquet_write(tmp_path):
 
     assert len(pd_df) == 1
     assert pd_df._preview.preview_partition is None
-    pd_df._populate_preview()
+    pd_df.__repr__()
     assert len(pd_df._preview.preview_partition) == 1
 
 
@@ -36,7 +36,7 @@ def test_parquet_write_with_partitioning(tmp_path):
 
     assert len(pd_df) == 5
     assert pd_df._preview.preview_partition is None
-    pd_df._populate_preview()
+    pd_df.__repr__()
     assert len(pd_df._preview.preview_partition) == 5
 
 
@@ -46,7 +46,7 @@ def test_empty_parquet_write_without_partitioning(tmp_path):
     output_files = df.write_parquet(tmp_path)
     assert len(output_files) == 0
     assert output_files._preview.preview_partition is None
-    output_files._populate_preview()
+    output_files.__repr__()
     assert len(output_files._preview.preview_partition) == 0
 
 
@@ -56,7 +56,7 @@ def test_empty_parquet_write_with_partitioning(tmp_path):
     output_files = df.write_parquet(tmp_path, partition_cols=["Borough"])
     assert len(output_files) == 0
     assert output_files._preview.preview_partition is None
-    output_files._populate_preview()
+    output_files.__repr__()
     assert len(output_files._preview.preview_partition) == 0
 
 
@@ -78,7 +78,7 @@ def test_parquet_write_with_partitioning_readback_values(tmp_path):
 
     assert len(output_files) == 5
     assert output_files._preview.preview_partition is None
-    output_files._populate_preview()
+    output_files.__repr__()
     assert len(output_files._preview.preview_partition) == 5
 
 
@@ -204,7 +204,7 @@ def test_csv_write(tmp_path):
 
     assert len(pd_df) == 1
     assert pd_df._preview.preview_partition is None
-    pd_df._populate_preview()
+    pd_df.__repr__()
     assert len(pd_df._preview.preview_partition) == 1
 
 

diff --git a/tests/dataframe/test_repr.py b/tests/dataframe/test_repr.py
@@ -87,7 +87,7 @@ def test_empty_repr(make_df):
     assert df._repr_html_() == "<small>(No data to display: Dataframe has no columns)</small>"
 
 
-@pytest.mark.parametrize("num_preview_rows", [None])
+@pytest.mark.parametrize("num_preview_rows", [9, 10, None])
 def test_repr_with_non_default_preview_rows(make_df, num_preview_rows):
     df = make_df({"A": [i for i in range(10)], "B": [i for i in range(10)]})
     df.collect(num_preview_rows=num_preview_rows)