gitlabds 2.4.2__tar.gz → 2.4.3__tar.gz

{gitlabds-2.4.2 → gitlabds-2.4.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: gitlabds
-Version: 2.4.2
+Version: 2.4.3
 Summary: GitLab Data Science Tools
 Author-email: Kevin Dietz <kdietz@gitlab.com>
 License-Expression: MIT

{gitlabds-2.4.2 → gitlabds-2.4.3}/gitlabds/split_data.py

@@ -140,7 +140,6 @@ def split_data(
         train_df, test_df = train_test_split(
             df_copy,
             train_size=train_pct,
-            test_size=1-train_pct,
             random_state=random_state,
             shuffle=shuffle,
             stratify=stratify_param
@@ -241,7 +240,6 @@ def split_data(
         train_groups, test_groups = train_test_split(
             group_target.index.values,
             train_size=train_pct,
-            test_size=1 - train_pct,
             random_state=random_state,
             shuffle=shuffle,
             stratify=group_stratify,
@@ -263,7 +261,6 @@ def split_data(
         X_train, X_test, y_train, y_test = train_test_split(
             X, y,
             train_size=train_pct,
-            test_size=1-train_pct,
             random_state=random_state,
             shuffle=shuffle,
             stratify=stratify_param

{gitlabds-2.4.2 → gitlabds-2.4.3}/gitlabds.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: gitlabds
-Version: 2.4.2
+Version: 2.4.3
 Summary: GitLab Data Science Tools
 Author-email: Kevin Dietz <kdietz@gitlab.com>
 License-Expression: MIT

{gitlabds-2.4.2 → gitlabds-2.4.3}/gitlabds.egg-info/SOURCES.txt

@@ -26,4 +26,5 @@ gitlabds.egg-info/requires.txt
 gitlabds.egg-info/top_level.txt
 tests/test_feature_reduction.py
 tests/test_record_benchmarking.py
-tests/test_snowflake_connectors.py
+tests/test_snowflake_connectors.py
+tests/test_split_data.py

gitlabds-2.4.3/tests/test_split_data.py

@@ -0,0 +1,59 @@
+import numpy as np
+import pandas as pd
+import pytest
+
+from gitlabds import split_data
+
+
+def _frame(n: int, groups: int = 0) -> pd.DataFrame:
+    rng = np.random.default_rng(0)
+    data = {
+        "x1": rng.random(n),
+        "x2": rng.random(n),
+        "y": (np.arange(n) % 2).astype(int),
+    }
+    if groups:
+        data["grp"] = np.repeat(np.arange(groups), n // groups)
+    return pd.DataFrame(data)
+
+
+@pytest.mark.parametrize("train_pct", [0.7, 0.85])
+@pytest.mark.parametrize("n", [100, 1000])
+def test_no_dv_split_on_exact_multiples(n, train_pct):
+    """Row counts where (1 - train_pct) * n is a whole number must not raise."""
+    df = _frame(n)
+    train_df, y_train, test_df, y_test, _ = split_data(
+        df, train_pct=train_pct, stratify=False
+    )
+    assert len(train_df) + len(test_df) == n
+    assert y_train.empty and y_test.empty
+
+
+@pytest.mark.parametrize("train_pct", [0.7, 0.85])
+@pytest.mark.parametrize("n", [100, 1000])
+def test_row_level_split_on_exact_multiples(n, train_pct):
+    df = _frame(n)
+    X_train, y_train, X_test, y_test, _ = split_data(
+        df, dv="y", train_pct=train_pct, stratify=False
+    )
+    assert len(X_train) + len(X_test) == n
+    assert len(y_train) == len(X_train)
+    assert len(y_test) == len(X_test)
+
+
+@pytest.mark.parametrize("train_pct", [0.7, 0.85])
+def test_group_split_on_exact_multiple_group_count(train_pct):
+    """The group count, not the row count, drives the group-aware split size."""
+    n, groups = 2000, 100
+    df = _frame(n, groups=groups)
+    X_train, y_train, X_test, y_test, _ = split_data(
+        df,
+        dv="y",
+        groupby="grp",
+        train_pct=train_pct,
+        stratify=False,
+    )
+    assert len(X_train) + len(X_test) == n
+    train_groups = set(X_train["grp"])
+    test_groups = set(X_test["grp"])
+    assert not (train_groups & test_groups)