update to trigger dask (#1184)

* update to trigger dask

* update maxdate, NaN filtering

* update special case check

* update deepcopy

* update inequality

* added series to class check

* restore deleted operators

* fixed equal_to indexing

* present_on_multiple_rows_within fix

* rework empty

* updated empty check

* empty

* equality

* restored USDM test suite

* restore missing relationship_data

* Remove cartesian products, fix ContentsDefineVLMDatasetBuilder, fix test_present_on_multiple_rows_within

* Fix get_codelist_attributes dask series

* fixes for record_count and unique

* Fix indexes for dask match datasets

* reenumerate index instead of preserving old index

* compute before boolean masking

* better index reset

* try range instead of rangeindex

* revert to set_index

* Revert most of distinct and record_count test changes

* better record_count example

* Remove comments

---------

Co-authored-by: gerrycampion <85252124+gerrycampion@users.noreply.github.com>
Co-authored-by: Gerry Campion <gcampion@cdisc.org>

Author: SFJohnson24

.github/workflows/test_suite.yml

@@ -83,8 +83,9 @@ jobs:
       - name: Run validation with Dask
         id: dask_run
         continue-on-error: true
+        env:
+          DATASET_SIZE_THRESHOLD: 0
         run: |
-          echo "DATASET_SIZE_THRESHOLD=0" > .env
           python core.py validate -s sdtmig -v 3-3 ${{ env.RULE_LIST }} -d CORE_Test_Suite/data -ct sdtmct-2020-03-27 -of json -o CORE_Test_Suite/dask-results -l info || true
 
           if [ -f "CORE_Test_Suite/dask-results.json" ]; then

cdisc_rules_engine/check_operators/dataframe_operators.py

@@ -163,14 +163,10 @@ def _check_equality(
         comparison_data = (
             comparator if comparator not in row or value_is_literal else row[comparator]
         )
-        target_values = row[target]
-        target_is_empty = pd.isna(target_values)
-        if not target_is_empty and isinstance(row[target], str):
-            target_is_empty = row[target] == ""
-        comp_is_empty = pd.isna(comparison_data)
-        if not comp_is_empty and isinstance(comparison_data, str):
-            comp_is_empty = comparison_data == ""
-        if target_is_empty or comp_is_empty:
+        both_null = (comparison_data == "" or comparison_data is None) & (
+            row[target] == "" or row[target] is None
+        )
+        if both_null:
             return False
         if case_insensitive:
             target_val = row[target].lower() if row[target] else None
@@ -198,32 +194,11 @@ def _check_inequality(
         comparison_data = (
             comparator if comparator not in row or value_is_literal else row[comparator]
         )
-        target_is_empty = pd.isna(row[target])
-        if isinstance(row[target], str) and (
-            not target_is_empty.any()
-            if hasattr(target_is_empty, "any")
-            else not target_is_empty
-        ):
-            target_is_empty = row[target] == ""
-        comp_is_empty = pd.isna(comparison_data)
-        if isinstance(comparison_data, str) and (
-            not comp_is_empty.any()
-            if hasattr(comp_is_empty, "any")
-            else not comp_is_empty
-        ):
-            comp_is_empty = comparison_data == ""
-        target_is_empty_scalar = (
-            target_is_empty.any()
-            if hasattr(target_is_empty, "any")
-            else target_is_empty
+        both_null = (comparison_data == "" or comparison_data is None) & (
+            row[target] == "" or row[target] is None
         )
-        comp_is_empty_scalar = (
-            comp_is_empty.any() if hasattr(comp_is_empty, "any") else comp_is_empty
-        )
-        if target_is_empty_scalar and comp_is_empty_scalar:
+        if both_null:
             return False
-        if target_is_empty_scalar or comp_is_empty_scalar:
-            return True
         if case_insensitive:
             target_val = row[target].lower() if row[target] else None
             comparison_val = comparison_data.lower() if comparison_data else None
@@ -243,7 +218,8 @@ def equal_to(self, other_value):
         return self.value.apply(
             lambda row: self._check_equality(row, target, comparator, value_is_literal),
             axis=1,
-        ).astype(bool)
+            meta=(None, "bool"),
+        ).reset_index(drop=True)
 
     @log_operator_execution
     @type_operator(FIELD_DATAFRAME)
@@ -294,7 +270,8 @@ def not_equal_to(self, other_value):
                 row, target, comparator, value_is_literal
             ),
             axis=1,
-        )
+            meta=(None, "bool"),
+        ).reset_index(drop=True)
 
     @log_operator_execution
     @type_operator(FIELD_DATAFRAME)
@@ -1290,7 +1267,9 @@ def present_on_multiple_rows_within(self, other_value: dict):
             lambda x: self.validate_series_length(x, target, min_count), meta=meta
         )
         uuid = str(uuid4())
-        return self.value.merge(results.rename(uuid), on=target)[uuid]
+        return self.value.merge(
+            results.rename(uuid).reset_index(), on=[group_by_column, target]
+        )[uuid]
 
     def validate_series_length(
         self, data: DatasetInterface, target: str, min_length: int

cdisc_rules_engine/dataset_builders/contents_define_vlm_dataset_builder.py

@@ -1,3 +1,4 @@
+from pandas import Series
 from cdisc_rules_engine.dataset_builders.values_dataset_builder import (
     ValuesDatasetBuilder,
 )
@@ -74,16 +75,14 @@ def build(self):
 
     @staticmethod
     def apply_filters(
-        vlm_row: dict, data_contents_df: DatasetInterface
+        vlm_row: Series, data_contents_df: DatasetInterface
     ) -> DatasetInterface:
         filter_results = data_contents_df.apply(
             lambda data_contents_row: vlm_row["filter"](data_contents_row), axis=1
         )
         row_numbers = data_contents_df.data["row_number"].copy()
         rows = row_numbers.where(filter_results)
-        filtered_df = rows.dropna().to_frame()
-        lut_subset = data_contents_df.__class__.cartesian_product(
-            filtered_df,
-            vlm_row.to_frame().T[["define_variable_name", "define_vlm_name"]],
-        )
+        filtered_df = data_contents_df.__class__(rows.dropna().to_frame())
+        vlm_row_df = vlm_row.to_frame().T[["define_variable_name", "define_vlm_name"]]
+        lut_subset = filtered_df.assign(**vlm_row_df.iloc[0])
         return lut_subset

cdisc_rules_engine/models/dataset/dask_dataset.py

@@ -44,6 +44,9 @@ def data(self, data):
 
     def __getitem__(self, item):
         try:
+            if hasattr(item, "dtype") and pd.api.types.is_bool_dtype(item.dtype):
+                # Handle boolean indexing
+                return self._data.compute()[item]
             return self._data[item].compute().reset_index(drop=True)
         except ValueError as e:
             # Handle boolean indexing length mismatch which occurs when filtering
@@ -93,6 +96,14 @@ def __len__(self):
 
         return self.length
 
+    def __deepcopy__(self, memo):
+        pandas_df = self._data.compute()
+        fresh_dask_df = dd.from_pandas(pandas_df, npartitions=DEFAULT_NUM_PARTITIONS)
+        new_instance = self.__class__(fresh_dask_df)
+        new_instance.length = self.length
+        memo[id(self)] = new_instance
+        return new_instance
+
     @classmethod
     def from_dict(cls, data: dict, **kwargs):
         dataframe = dd.from_dict(data, npartitions=DEFAULT_NUM_PARTITIONS, **kwargs)
@@ -232,7 +243,7 @@ def melt(
         return self.__class__(new_data)
 
     def assign(self, **kwargs):
-        return self.data.assign(**kwargs)
+        return self.__class__(self.data.assign(**kwargs))
 
     def copy(self):
         new_data = self._data.copy()
@@ -265,18 +276,6 @@ def get_error_rows(self, results) -> "pd.Dataframe":
             1000, npartitions=-1
         )
 
-    @classmethod
-    def cartesian_product(cls, left, right):
-        """
-        Return the cartesian product of two dataframes
-        """
-        return cls(
-            dd.from_pandas(
-                left.compute().merge(right, how="cross"),
-                npartitions=DEFAULT_NUM_PARTITIONS,
-            )
-        )
-
     def dropna(self, inplace=False, **kwargs):
         result = self._data.dropna(**kwargs)
         if inplace:
@@ -371,7 +370,9 @@ def partition_isin(partition):
         return result
 
     def filter_by_value(self, column, values):
-        computed_data = self._data.compute()
-        mask = computed_data[column].isin(values)
-        filtered_df = computed_data[mask]
-        return filtered_df
+        mask = self._data[column].isin(values)
+        return self.__class__(self._data[mask])
+
+    def max(self, *args, **kwargs):
+        result = self._data.max(*args, **kwargs)
+        return self.__class__(result)

cdisc_rules_engine/models/dataset/dataset_interface.py

@@ -170,6 +170,13 @@ def len(self) -> int:
         Return the length of the dataset
         """
 
+    @abstractmethod
+    def assign(self, **kwargs):
+        """
+        Assign new columns to the dataset.
+        This method should return a new instance of the dataset with the new columns added.
+        """
+
     @abstractmethod
     def copy(self) -> "DatasetInterface":
         """
@@ -194,12 +201,6 @@ def where(cond, other, **kwargs):
         Wrapper for dataframe where function
         """
 
-    @abstractmethod
-    def cartesian_product(cls, left, right):
-        """
-        Return the cartesian product of two dataframes
-        """
-
     @abstractmethod
     def sort_values(self, by, **kwargs):
         """

cdisc_rules_engine/models/dataset/pandas_dataset.py

@@ -187,6 +187,9 @@ def len(self) -> int:
     def size(self) -> int:
         return self._data.memory_usage().sum()
 
+    def assign(self, **kwargs):
+        return self.__class__(self.data.assign(**kwargs))
+
     def copy(self):
         new_data = self._data.copy()
         return self.__class__(new_data)
@@ -206,13 +209,6 @@ def where(self, cond, other, **kwargs):
         new_data = self._data.where(cond, other, **kwargs)
         return self.__class__(new_data)
 
-    @classmethod
-    def cartesian_product(cls, left, right):
-        """
-        Return the cartesian product of two dataframes
-        """
-        return cls(left.merge(right, how="cross"))
-
     def sort_values(self, by: Union[str, list[str]], **kwargs) -> "pd.Dataframe":
         """
         Sort the underlying dataframe and return a raw dataframe.
@@ -247,6 +243,9 @@ def astype(self, dtype, **kwargs):
     def min(self, *args, **kwargs):
         return self.__class__(self._data.min(*args, **kwargs))
 
+    def max(self, *args, **kwargs):
+        return self.__class__(self._data.max(*args, **kwargs))
+
     def reset_index(self, drop=False, **kwargs):
         """
         Reset the index of the dataset.

cdisc_rules_engine/operations/get_codelist_attributes.py

@@ -55,7 +55,7 @@ def _get_codelist_attributes(self):
 
         # 4.0 merge the two datasets by CC
         # -------------------------------------------------------------------
-        cc_key = ct_data[ct_name].to_list()
+        cc_key = ct_data[ct_name]
         ct_list = ct_cache[(ct_cache[ct_name].isin(cc_key))]
         ds_len = self.params.dataframe.len()
         result = pd.Series([ct_list[ct_attribute].values[0] for _ in range(ds_len)])

cdisc_rules_engine/operations/max_date.py

@@ -12,7 +12,13 @@ def _execute_operation(self):
             else:
                 result = max_date.isoformat()
         else:
-            result = self.params.dataframe.groupby(
-                self.params.grouping, as_index=False
-            ).data.max()
+            result = self.params.dataframe.groupby(self.params.grouping).max()
+        if isinstance(result, pd.Series):
+            result = result.apply(lambda x: x.isoformat() if pd.notna(x) else "")
+        elif isinstance(result, pd.DataFrame):
+            for col in result.columns:
+                if pd.api.types.is_datetime64_any_dtype(result[col]):
+                    result[col] = result[col].apply(
+                        lambda x: x.isoformat() if pd.notna(x) else ""
+                    )
         return result

cdisc_rules_engine/operations/record_count.py

@@ -1,4 +1,5 @@
 import pandas as pd
+from numpy import int64
 
 from cdisc_rules_engine.operations.base_operation import BaseOperation
 
@@ -33,6 +34,7 @@ def _execute_operation(self) -> pd.Series:
                         how="left",
                     )
                     .fillna(0)
+                    .astype({self.params.target: int64})
                 )
             return group_df
         return result

cdisc_rules_engine/services/data_services/base_data_service.py

@@ -5,6 +5,7 @@
 from concurrent.futures import ThreadPoolExecutor
 import os
 import numpy as np
+import dask.dataframe as dd
 
 from cdisc_rules_engine.interfaces import (
     CacheServiceInterface,
@@ -29,6 +30,7 @@
     get_directory_path,
     search_in_list_of_dicts,
     tag_source,
+    replace_nan_values_in_df,
 )
 from cdisc_rules_engine.utilities.sdtm_utilities import get_class_and_domain_metadata
 from cdisc_rules_engine.models.dataset.dataset_interface import DatasetInterface
@@ -266,8 +268,22 @@ def _contains_topic_variable(
         """
 
         def check_presence(key):
-            in_dataset = key in dataset
-            in_values = key in self.dataset_implementation.get_series_values(dataset)
+            if hasattr(dataset, "columns"):
+                columns = dataset.columns
+                if hasattr(columns, "tolist"):
+                    columns = columns.tolist()
+                in_dataset = key in columns
+                in_values = key in self.dataset_implementation.get_series_values(
+                    dataset
+                )
+            else:
+                series_values = dataset.values
+                if hasattr(series_values, "tolist"):
+                    series_values = series_values.tolist()
+                in_dataset = key in series_values
+                in_values = key in self.dataset_implementation.get_series_values(
+                    dataset
+                )
             return in_dataset or in_values
 
         if not check_presence("DOMAIN") and not check_presence("RDOMAIN"):
@@ -300,9 +316,16 @@ def _replace_nans_in_specified_cols_with_none(
         """
         Replaces NaN in specified columns with None.
         """
-        if isinstance(column_names, List):
-            column_names = dataset.data[column_names].columns
-        dataset.data = dataset.data.replace(np.nan, {col: None for col in column_names})
+        valid_columns = [col for col in column_names if col in dataset.data.columns]
+        if not valid_columns:
+            return dataset
+        if isinstance(dataset.data, dd.DataFrame):
+            dataset.data = dataset.data.map_partitions(
+                replace_nan_values_in_df, valid_columns
+            )
+        else:
+            dataset.data = replace_nan_values_in_df(dataset.data, valid_columns)
+        return dataset
 
     async def _async_get_dataset(
         self, function_to_call: Callable, dataset_name: str, **kwargs