moved logic to return original dataframe sooner for more efficiency

pola-rs · Oct 25, 2023 · a8dac99 · a8dac99
1 parent 824c28d
commit a8dac99
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 7 deletions.
diff --git a/crates/polars-core/src/frame/explode.rs b/crates/polars-core/src/frame/explode.rs
@@ -279,6 +279,10 @@ impl DataFrame {
                     }
                 })
                 .collect();
+            // if value_vars is still empty after attempting to populate it, return original dataframe
+            if value_vars.is_empty() {
+                return Ok(self.clone());
+            }
         }
 
         // values will all be placed in single column, so we must find their supertype
@@ -288,10 +292,6 @@ impl DataFrame {
                 .get(v)
                 .ok_or_else(|| polars_err!(ColumnNotFound: "{}", v))
         });
-        // if there only is a single column, we return the original DataFrame
-        if value_vars.len() < 1 {
-            return Ok(self.clone());
-        }
 
         let mut st = iter.next().unwrap()?.clone();
         for dt in iter {

diff --git a/py-polars/tests/unit/operations/test_melt.py b/py-polars/tests/unit/operations/test_melt.py
@@ -70,7 +70,12 @@ def test_melt_projection_pd_7747() -> None:
     )
     assert_frame_equal(result, expected)
 
+
 def test_melt_single_column() -> None:
-    df = pl.DataFrame({'single_column': [1,2,3],})
-    result = df.melt('single_column')
-    assert_frame_equal(result, df)
+    df = pl.DataFrame(
+        {
+            "single_column": [1, 2, 3],
+        }
+    )
+    result = df.melt("single_column")
+    assert_frame_equal(result, df)