Repository: dask/dask

Issue:
DataFrame.sample function signature differs from pandas
- `pandas.DataFrame.sample(n, frac, ...)`
- `dask.dataframe.DataFrame.sample(frac, ...)`

We can't reliably do `n`, so we should raise if it's anything but `None`. But it'd still be nice to match the signature.

Discussion:

Patch:
diff --git a/dask/dataframe/core.py b/dask/dataframe/core.py
index 146b1eb9f..be22829a1 100644
--- a/dask/dataframe/core.py
+++ b/dask/dataframe/core.py
@@ -1030,12 +1030,15 @@ Dask Name: {name}, {task} tasks""".format(klass=self.__class__.__name__,
     def bfill(self, axis=None, limit=None):
         return self.fillna(method='bfill', limit=limit, axis=axis)
 
-    def sample(self, frac, replace=False, random_state=None):
+    def sample(self, n=None, frac=None, replace=False, random_state=None):
         """ Random sample of items
 
         Parameters
         ----------
-        frac : float
+        n : int, optional
+            Number of items to return is not supported by dask. Use frac
+            instead.
+        frac : float, optional
             Fraction of axis items to return.
         replace : boolean, optional
             Sample with or without replacement. Default = False.
@@ -1048,6 +1051,17 @@ Dask Name: {name}, {task} tasks""".format(klass=self.__class__.__name__,
         DataFrame.random_split
         pandas.DataFrame.sample
         """
+        if n is not None:
+            msg = ("sample does not support the number of sampled items "
+                   "parameter, 'n'. Please use the 'frac' parameter instead.")
+            if isinstance(n, Number) and 0 <= n <= 1:
+                warnings.warn(msg)
+                frac = n
+            else:
+                raise ValueError(msg)
+
+        if frac is None:
+            raise ValueError("frac must not be None")
 
         if random_state is None:
             random_state = np.random.RandomState()


Test patch:
diff --git a/dask/dataframe/tests/test_dataframe.py b/dask/dataframe/tests/test_dataframe.py
index b155de6c9..06d1dae90 100644
--- a/dask/dataframe/tests/test_dataframe.py
+++ b/dask/dataframe/tests/test_dataframe.py
@@ -1379,7 +1379,7 @@ def test_embarrassingly_parallel_operations():
     assert_eq(a.notnull(), df.notnull())
     assert_eq(a.isnull(), df.isnull())
 
-    assert len(a.sample(0.5).compute()) < len(df)
+    assert len(a.sample(frac=0.5).compute()) < len(df)
 
 
 def test_fillna():
@@ -1452,26 +1452,46 @@ def test_sample():
                       index=[10, 20, 30, 40, 50, 60])
     a = dd.from_pandas(df, 2)
 
-    b = a.sample(0.5)
+    b = a.sample(frac=0.5)
 
     assert_eq(b, b)
 
-    c = a.sample(0.5, random_state=1234)
-    d = a.sample(0.5, random_state=1234)
+    c = a.sample(frac=0.5, random_state=1234)
+    d = a.sample(frac=0.5, random_state=1234)
     assert_eq(c, d)
 
-    assert a.sample(0.5)._name != a.sample(0.5)._name
+    assert a.sample(frac=0.5)._name != a.sample(frac=0.5)._name
 
 
 def test_sample_without_replacement():
     df = pd.DataFrame({'x': [1, 2, 3, 4, None, 6], 'y': list('abdabd')},
                       index=[10, 20, 30, 40, 50, 60])
     a = dd.from_pandas(df, 2)
-    b = a.sample(0.7, replace=False)
+    b = a.sample(frac=0.7, replace=False)
     bb = b.index.compute()
     assert len(bb) == len(set(bb))
 
 
+def test_sample_raises():
+    df = pd.DataFrame({'x': [1, 2, 3, 4, None, 6], 'y': list('abdabd')},
+                      index=[10, 20, 30, 40, 50, 60])
+    a = dd.from_pandas(df, 2)
+
+    # Make sure frac is replaced with n when 0 <= n <= 1
+    # This is so existing code (i.e. ddf.sample(0.5)) won't break
+    with pytest.warns(UserWarning):
+        b = a.sample(0.5, random_state=1234)
+    c = a.sample(frac=0.5, random_state=1234)
+    assert_eq(b, c)
+
+    with pytest.raises(ValueError):
+        a.sample(n=10)
+
+    # Make sure frac is provided
+    with pytest.raises(ValueError):
+        a.sample(frac=None)
+
+
 def test_datetime_accessor():
     df = pd.DataFrame({'x': [1, 2, 3, 4]})
     df['x'] = df.x.astype('M8[us]')