-def test_univariate_filter_corr():
+import pytest
+
+@pytest.mark.parametrize("fs_method, selection_mode, selection_threshold, expected_features", [
+    ("u_corr", None, 0.3, 211),
+    ("anova", "percentile", 0.8, 4),
+    # Add more test cases for other univariate methods
+])
+def test_univariate_filter(fs_method, selection_mode, selection_threshold, expected_features):
     """
-    Test univariate_filter method with 'u_corr' method.
-    :return: None
+    Test univariate_filter method with different parameters.
     """
-
-    # import tsv as pandas DataFrame
     df = pd.read_csv(get_tnbc_data_path(), sep="\t")
-
-    # create FSDataFrame instance
     fs_df = FSDataFrame(df=df, sample_col="Sample", label_col="label")
 
-    # create FSUnivariate instance
-    fs_univariate = FSUnivariate(fs_method="u_corr", selection_threshold=0.3)
-
-    fsdf_filtered = fs_univariate.select_features(fs_df)
-
-    assert fs_df.count_features() == 500
-    assert fsdf_filtered.count_features() == 211
-
-    # Export the filtered DataFrame as Pandas DataFrame
-    df_filtered = fsdf_filtered.to_pandas()
-    df_filtered.to_csv("filtered_tnbc_data.csv", index=False)
-
-
-# test the univariate_filter method with 'anova' method
-def test_univariate_filter_anova():
-    """
-    Test univariate_filter method with 'anova' method.
-    :return: None
-    """
-
-    # import tsv as pandas DataFrame
-    df = pd.read_csv(get_tnbc_data_path(), sep="\t")
-
-    # create FSDataFrame instance
-    fs_df = FSDataFrame(df=df, sample_col="Sample", label_col="label")
-
-    # create FSUnivariate instance
     fs_univariate = FSUnivariate(
-        fs_method="anova", selection_mode="percentile", selection_threshold=0.8
+        fs_method=fs_method,
+        selection_mode=selection_mode,
+        selection_threshold=selection_threshold
     )
 
     fsdf_filtered = fs_univariate.select_features(fs_df)
 
     assert fs_df.count_features() == 500
-    assert fsdf_filtered.count_features() == 4
+    assert fsdf_filtered.count_features() == expected_features
 
-    # Export the filtered DataFrame as Pandas DataFrame
     df_filtered = fsdf_filtered.to_pandas()
-    df_filtered.to_csv("filtered_tnbc_data.csv", index=False)
+    df_filtered.to_csv(f"filtered_tnbc_data_{fs_method}.csv", index=False)
 

-# eam
-# 18.07.22
+"""
+fslite: A Python package for memory-efficient, high-performance feature selection on big and small datasets.
 
+Author: eam
+Date: 18.07.22
+"""
+

 if score_func not in score_func_mapping:
-    raise ValueError(
+    raise InvalidMethodError(
         f"Invalid score_func '{score_func}'. Valid options are: {list(score_func_mapping.keys())}"
     )
 

+from enum import Enum, auto
+
+class FSMethod(Enum):
+    ANOVA = auto()
+    U_CORR = auto()
+    F_REGRESSION = auto()
+    MUTUAL_INFO_REGRESSION = auto()
+    MUTUAL_INFO_CLASSIFICATION = auto()
+
 FS_METHODS = {
     "univariate": {
         "title": "Univariate Feature Selection",
         "description": "Univariate feature selection refers to the process of selecting the most relevant features for "
         "a machine learning model by evaluating each feature individually with respect to the target "
         "variable using univariate statistical tests. It simplifies the feature selection process by "
         "treating each feature independently and assessing its contribution to the predictive "
         "performance of the model.",
         "methods": [
             {
-                "name": "anova",
+                "name": FSMethod.ANOVA,
                 "description": "Univariate ANOVA feature selection (f-classification)",
             },
-            {"name": "u_corr", "description": "Univariate Pearson's correlation"},
-            {"name": "f_regression", "description": "Univariate f-regression"},
+            {"name": FSMethod.U_CORR, "description": "Univariate Pearson's correlation"},
+            {"name": FSMethod.F_REGRESSION, "description": "Univariate f-regression"},
             {
-                "name": "mutual_info_regression",
+                "name": FSMethod.MUTUAL_INFO_REGRESSION,
                 "description": "Univariate mutual information regression",
             },
             {
-                "name": "mutual_info_classification",
+                "name": FSMethod.MUTUAL_INFO_CLASSIFICATION,
                 "description": "Univariate mutual information classification",
             },
         ],
     },
 

+from typing import Union
+import pandas as pd
+import pyarrow as pa
+from pyspark.sql import DataFrame, SparkSession
 
+def import_table(file_path: str, sep: str = '\t', n_partitions: int = 5) -> DataFrame:
+    # Implementation here
+    pass
 
+def import_table_as_psdf(file_path: str, sep: str = '\t', n_partitions: int = 5) -> pd.DataFrame:
+    # Implementation here
+    pass
+

-class InvalidDataError(ValueError):
+class InvalidDataError(Exception):
     """
-    Error raised when an invalid feature selection method is used.
+    Error raised when invalid data is provided for feature selection.
     """
 
     def __init__(self, message):
         super().__init__(f"Invalid data frame: {message}")
 

 ## Experiments and Benchmarks
 
 This document contains the experiments and benchmarks that were conducted to evaluate the performance of fslite. 
 The experiments were conducted on the following datasets:
 
+1. Dataset A: [Brief description, size, characteristics]
+2. Dataset B: [Brief description, size, characteristics]
+3. Dataset C: [Brief description, size, characteristics]
+
+Each dataset was chosen to represent different scenarios and challenges in feature selection tasks.
+

-if sparsity > sparse_threshold:
+if np.isclose(sparsity, sparse_threshold, atol=1e-8) or sparsity > sparse_threshold:
     if dense_matrix_size < memory_threshold * available_memory:
         # Use dense matrix if enough memory is available
         logging.info(
             f"Data is sparse (sparsity={sparsity:.2f}) but enough memory available. "
             f"Using a dense matrix."
         )
         self.__matrix = numerical_df.to_numpy(dtype=np.float32)
         self.__is_sparse = False
     else:
         # Use sparse matrix due to memory constraints
         logging.info(
             f"Data is sparse (sparsity={sparsity:.2f}), memory insufficient for dense matrix. "
             f"Using a sparse matrix representation."
         )
         self.__matrix = sparse.csr_matrix(
             numerical_df.to_numpy(dtype=np.float32)
         )
         self.__is_sparse = True
 else:
     # Use dense matrix since it's not sparse
     logging.info(
         f"Data is not sparse (sparsity={sparsity:.2f}), using a dense matrix."
     )
     self.__matrix = numerical_df.to_numpy(dtype=np.float32)
     self.__is_sparse = False
 

+TESTDATA_DIR = Path(__file__).parent.parent / "testdata"
+
 def get_tnbc_data_path() -> str:
     """
     Return path to example dataset (TNBC) with 44 samples and 500 features.
-
     """
-    tnbc_path = Path(__file__).parent.parent / "testdata/TNBC.tsv.gz"
-    return tnbc_path.__str__()
-
+    return str(TESTDATA_DIR / "TNBC.tsv.gz")
 
 def get_tnbc_data_missing_values_path() -> str:
     """
     Return path to example dataset (TNBC) with missing values.
+    """
+    return str(TESTDATA_DIR / "TNBC_missing.tsv")
 
-    """
-    tnbc_path = Path(__file__).parent.parent / "testdata/TNBC_missing.tsv"
-    return tnbc_path.__str__()
-

-sdf = import_table('data.tsv.bgz',
+sdf = import_table(file_path='data.tsv.bgz',
                    sep='\t',
                    n_partitions=5)
 

-if univariate_method == "anova":
-    selected_features = self.univariate_feature_selector(
-        df, score_func="f_classif", **kwargs
-    )
-elif univariate_method == "f_regression":
-    selected_features = self.univariate_feature_selector(
-        df, score_func="f_regression", **kwargs
-    )
-elif univariate_method == "u_corr":
-    selected_features = univariate_correlation_selector(df, **kwargs)
-elif univariate_method == "mutual_info_classification":
-    selected_features = self.univariate_feature_selector(
-        df, score_func="mutual_info_classif", **kwargs
-    )
-elif univariate_method == "mutual_info_regression":
-    selected_features = self.univariate_feature_selector(
-        df, score_func="mutual_info_regression", **kwargs
-    )
+method_mapping = {
+    "anova": lambda: self.univariate_feature_selector(df, score_func="f_classif", **kwargs),
+    "f_regression": lambda: self.univariate_feature_selector(df, score_func="f_regression", **kwargs),
+    "u_corr": lambda: univariate_correlation_selector(df, **kwargs),
+    "mutual_info_classification": lambda: self.univariate_feature_selector(df, score_func="mutual_info_classif", **kwargs),
+    "mutual_info_regression": lambda: self.univariate_feature_selector(df, score_func="mutual_info_regression", **kwargs),
+}
 
+selected_features = method_mapping.get(univariate_method, lambda: [])()
+

 # Compute correlation matrix
 if corr_method == "pearson":
     corr_matrix = np.corrcoef(f_matrix, rowvar=False)
 elif corr_method == "spearman":
     corr_matrix, _ = spearmanr(f_matrix)
 else:
     raise ValueError(
         f"Unsupported correlation method '{corr_method}'. Use 'pearson' or 'spearman'."
     )
 
+# For large datasets, consider using a more efficient method
+# For example, you could use pandas' corr() method with a custom function
+# that computes correlation for chunks of the data at a time
+