nebius/SWE-rebench-V2-PRs huggingface dataset:
Repository: vaexio/vaex
Pull request #199 (2019-04-03 09:59:58)

Adding a unit-test for apply exposing the need to update the lazy functions.


Patch:
diff --git a/packages/vaex-core/vaex/scopes.py b/packages/vaex-core/vaex/scopes.py
index 72fc1d1e..5d508eb0 100644
--- a/packages/vaex-core/vaex/scopes.py
+++ b/packages/vaex-core/vaex/scopes.py
@@ -193,6 +193,8 @@ class _BlockScopeSelection(object):
                     # self._ensure_buffer(variable)
                     return self.evaluate(expression)  # , out=self.buffers[variable])
                     # self.values[variable] = self.buffers[variable]
+                elif variable in self.df.functions:
+                    return self.df.functions[variable].f
                 raise KeyError("Unknown variables or column: %r" % (variable,))
         except:
             import traceback as tb


Test patch:
diff --git a/tests/apply_test.py b/tests/apply_test.py
index 124d93cf..6845bb33 100644
--- a/tests/apply_test.py
+++ b/tests/apply_test.py
@@ -4,3 +4,12 @@ import vaex
 def test_apply():
     ds = vaex.from_arrays(x=[1, 2, 3])
     assert ds.x.apply(lambda x: x + 1).values.tolist() == [2, 3, 4]
+
+
+def test_apply_select():
+    # tests if select supports lambda/functions
+    x = [1, 2, 3, 4, 5]
+    ds = vaex.from_arrays(x=x)
+    ds['x_ind'] = ds.x.apply(lambda w: w > 3)
+    ds.select('x_ind')
+    assert 2 == ds.selected_length()