from bukka.utils.files import file_manager
from bukka.utils.bukka_logger import BukkaLogger
from bukka.data_management.dataset_functionality import (
DatasetStatistics,
DatasetManagement,
DatasetIO,
DatasetQuality,
)
from narwhals.typing import FrameT
import narwhals as nw
logger = BukkaLogger(__name__)
[docs]
class Dataset:
"""Dataset class for managing and splitting datasets for expert systems.
This class loads, splits, and manages datasets using Narwhals for dataframe
abstraction. It writes train/test splits as Parquet files and exposes schema
and feature metadata.
Parameters
----------
target_column : str
The name of the target column in the dataset.
file_manager : file_manager.FileManager
An instance of FileManager to handle file paths and dataset storage.
strata : list[str] | str | None, optional
Column(s) to use for stratified splitting. Defaults to None.
stratify : bool, optional
Whether to stratify the split. Defaults to True.
train_size : float, optional
Proportion of the dataset to include in the train split. Defaults to 0.8.
feature_columns : list[str] | None, optional
List of feature column names. If None, all columns except the target are used.
Defaults to None.
Attributes
----------
file_manager : file_manager.FileManager
File manager instance for data paths.
target_column : str
Name of the target column.
feature_columns : list[str]
List of feature column names.
data_schema : dict[str, pyarrow.DataType]
Schema of the training data (column names mapped to PyArrow data types).
train_df : Narwhals DataFrame
Training data split.
test_df : Narwhals DataFrame
Test data split.
Examples
--------
>>> from bukka.utils.files.file_manager import FileManager
>>> from pathlib import Path
>>> fm = FileManager(project_name='my_project', dataset_path=Path('data.csv'))
>>> dataset = Dataset(
... target_column='label',
... file_manager=fm,
... train_size=0.7,
... stratify=True
... )
>>> print(dataset.feature_columns)
['feat1', 'feat2', 'feat3']
>>> print(dataset.data_schema)
{'feat1': int64, 'feat2': double, 'label': int64}
"""
def __init__(
self,
target_column: str,
file_manager: file_manager.FileManager,
strata=None,
stratify=True,
train_size=0.8,
feature_columns: list[str] | None = None,
backend: str = "polars"
):
self.io = DatasetIO()
self.management = DatasetManagement()
self.statistics = DatasetStatistics()
self.quality = DatasetQuality()
logger.debug(f"Initializing Dataset with target_column='{target_column}', train_size={train_size}, stratify={stratify}, backend='{backend}'")
self.file_manager = file_manager
self.target_column = target_column
self.backend = backend
dataset_path = getattr(self.file_manager, 'dataset_path', None)
if dataset_path is not None and dataset_path.exists():
logger.debug(f"Loading dataset from: {dataset_path}")
df = self.io.load_from_file(dataset_path, backend=self.backend)
# Ensure we have a Narwhals DataFrame (for compatibility with tests/mocks)
if not hasattr(df, '__narwhals_dataframe__'):
df = nw.from_native(df)
# Validate target column exists in the dataset
if target_column and target_column not in df.columns:
available_cols = ', '.join(df.columns[:10]) # Show first 10 columns
if len(df.columns) > 10:
available_cols += f", ... ({len(df.columns)} total)"
raise ValueError(
f"Target column '{target_column}' not found in dataset. "
f"Available columns: {available_cols}"
)
else:
logger.debug("No dataset path found or dataset does not exist, skipping dataset loading")
logger.debug(f"Splitting dataset into train/test with train_size={train_size}")
if strata is None and target_column is None:
stratify = False
strata = []
if stratify is None:
strata = []
self.train_df, self.test_df = self.management.split_dataset(
df=df,
target_column=target_column,
strata=strata,
train_size=train_size,
stratify=stratify
)
self.io.save_to_parquet(self.train_df, self.file_manager.train_data_file)
self.io.save_to_parquet(self.test_df, self.file_manager.test_data_file)
logger.debug("Dataset split completed")
if feature_columns == None:
logger.debug("Auto-detecting feature columns from training data")
self.feature_columns = list(self.train_df.columns)
if target_column and target_column in self.feature_columns:
self.feature_columns.remove(target_column)
logger.debug(f"Detected {len(self.feature_columns)} feature columns")
else:
logger.debug(f"Using provided feature columns: {len(feature_columns)} columns")
self.feature_columns = feature_columns
logger.debug(f"Reading schema from: {self.file_manager.train_data_file}")
self.data_schema = dict(self.train_df.schema)
logger.debug(f"Schema loaded with {len(self.data_schema)} columns")
logger.debug("Dataset initialization complete")
[docs]
def identify_multicollinearity_train(self, columns: list[str] = None, threshold: float = 0.8):
"""Identify multicollinear features in the training dataset.
Computes pairwise correlations between numerical columns and returns
pairs with absolute correlation above the specified threshold.
Parameters
----------
columns : list[str], optional
List of column names to check for multicollinearity. If None,
uses all feature columns. Defaults to None.
threshold : float, optional
Correlation threshold above which column pairs are considered
multicollinear. Defaults to 0.8.
Returns
-------
list[tuple[str, str, float]]
List of tuples with correlated column pairs and their correlation
coefficient. Each tuple is (col1, col2, correlation).
Examples
--------
>>> dataset = Dataset(
... target_column='label',
... file_manager=fm
... )
>>> pairs = dataset.identify_multicollinearity_train(['feat1', 'feat2', 'feat3'])
>>> print(pairs)
[('feat1', 'feat2', 0.95), ('feat2', 'feat3', 0.87)]
>>> # Use default feature columns and custom threshold
>>> pairs = dataset.identify_multicollinearity_train(threshold=0.9)
"""
if columns is None:
columns = self.feature_columns
return self.statistics.identify_multicollinearity(self.train_df, columns, threshold)
[docs]
def get_varied_scale_train(self, column_name: str):
"""Calculate the range of a column in the training dataset.
Computes the range (maximum value minus minimum value) for a numerical
column to assess scale variation.
Parameters
----------
column_name : str
Name of the column to analyze. Must be a numerical column.
Returns
-------
float
The range of the column (max - min).
Examples
--------
>>> dataset = Dataset(
... target_column='label',
... file_manager=fm
... )
>>> scale = dataset.get_varied_scale_train('price')
>>> print(scale)
950.75
>>> # Check scale for multiple columns
>>> for col in ['price', 'quantity', 'weight']:
... scale = dataset.get_varied_scale_train(col)
... print(f"{col}: {scale}")
"""
return self.statistics.get_varied_scale(self.train_df, column_name)
[docs]
def check_varied_scale_train(self, column_name: str, threshold: float):
"""Check if a column has varied scale in the training dataset.
Determines whether a numerical column's range exceeds a specified
threshold, indicating that scaling might be beneficial.
Parameters
----------
column_name : str
Name of the column to check. Must be a numerical column.
threshold : float
The threshold value for determining varied scale. If the column's
range (max - min) exceeds this value, returns True.
Returns
-------
bool
True if the column's range exceeds the threshold, False otherwise.
Examples
--------
>>> dataset = Dataset(
... target_column='label',
... file_manager=fm
... )
>>> has_varied = dataset.check_varied_scale_train('price', 100)
>>> print(has_varied)
True
>>> # Check multiple columns for scaling needs
>>> for col in ['price', 'age', 'quantity']:
... needs_scaling = dataset.check_varied_scale_train(col, threshold=100)
... if needs_scaling:
... print(f"{col} needs scaling")
"""
return self.statistics.does_data_have_varied_scale(self.train_df, column_name, threshold)
[docs]
def get_column_null_count(self, column: str) -> int:
"""Get the count of null values in a column from the training dataset.
Parameters
----------
column : str
Name of the column to check.
Returns
-------
int
The count of null values in the column.
Examples
--------
>>> dataset = Dataset(target_column='label', file_manager=fm)
>>> null_count = dataset.get_column_null_count('age')
>>> print(null_count)
5
"""
return self.quality.get_column_null_count(self.train_df, column)
[docs]
def type_of_column(self, column: str) -> str:
"""Get the simplified data type of a column from the training dataset.
Parameters
----------
column : str
Name of the column to check.
Returns
-------
str
The simplified data type: 'int', 'float', 'string', or backend-specific type name.
Examples
--------
>>> dataset = Dataset(target_column='label', file_manager=fm)
>>> dtype = dataset.type_of_column('age')
>>> print(dtype)
'int'
"""
return self.quality.type_of_column(self.train_df, column)
[docs]
def has_outliers(self, column: str, z_threshold: float = 3) -> bool:
"""Check if a column has outliers in the training dataset.
Parameters
----------
column : str
Name of the column to check.
z_threshold : float, optional
Number of standard deviations from mean to consider as outlier.
Defaults to 3.
Returns
-------
bool
True if outliers are detected, False otherwise.
Examples
--------
>>> dataset = Dataset(target_column='label', file_manager=fm)
>>> has_outliers = dataset.has_outliers('price')
>>> print(has_outliers)
True
"""
return self.statistics.has_outliers(self.train_df, column, z_threshold)
[docs]
def get_unq_count(self, column: str) -> int:
"""Get the count of unique values in a column from the training dataset.
Parameters
----------
column : str
Name of the column to analyze.
Returns
-------
int
The number of unique values in the column.
Examples
--------
>>> dataset = Dataset(target_column='label', file_manager=fm)
>>> unique_count = dataset.get_unq_count('category')
>>> print(unique_count)
5
"""
return self.statistics.get_unq_count(self.train_df, column)
[docs]
def has_inconsistent_categorical_data(self, column: str, threshold: float = 0.1) -> bool:
"""Check if a categorical column has inconsistent data in the training dataset.
Parameters
----------
column : str
Name of the column to check.
threshold : float, optional
Threshold for unique value ratio. Defaults to 0.1.
Returns
-------
bool
True if inconsistent categorical data is detected, False otherwise.
Examples
--------
>>> dataset = Dataset(target_column='label', file_manager=fm)
>>> is_inconsistent = dataset.has_inconsistent_categorical_data('category')
>>> print(is_inconsistent)
True
"""
return self.quality.has_inconsistent_categorical_data(self.train_df, column, threshold)
[docs]
def has_multicollinearity(self, columns: list[str] | None = None, threshold: float = 0.8) -> bool:
"""Check if the training dataset has multicollinearity among columns.
Parameters
----------
columns : list[str], optional
List of column names to check. If None, uses feature_columns.
threshold : float, optional
Correlation threshold for detecting multicollinearity. Defaults to 0.8.
Returns
-------
bool
True if any pair of columns has correlation above threshold.
Examples
--------
>>> dataset = Dataset(target_column='label', file_manager=fm)
>>> has_multi = dataset.has_multicollinearity()
>>> print(has_multi)
True
"""
if columns is None:
columns = self.feature_columns
return self.statistics.does_data_have_multicollinearity(self.train_df, columns, threshold)
[docs]
def is_text_column(self, column: str, min_avg_length: int = 50) -> bool:
"""Check if a column contains text data suitable for NLP tasks.
Parameters
----------
column : str
Name of the column to check.
min_avg_length : int, optional
Minimum average string length to be considered text. Defaults to 50.
Returns
-------
bool
True if the column appears to contain text data, False otherwise.
Examples
--------
>>> dataset = Dataset(target_column='label', file_manager=fm)
>>> is_text = dataset.is_text_column('description')
>>> print(is_text)
True
"""
return self.quality.is_text_column(self.train_df, column, min_avg_length)
def __repr__(self):
return f"Dataset(target_column={self.target_column}, feature_columns={self.feature_columns})"