mloda 0.3.2__py3-none-any.whl → 0.4.0__py3-none-any.whl

mloda_plugins/feature_group/experimental/time_window/base.py

@@ -7,16 +7,19 @@ from __future__ import annotations
 from abc import abstractmethod
 from typing import Any, List, Optional, Set, Type, Union
 
-from mloda_core.abstract_plugins.abstract_feature_group import AbstractFeatureGroup
-from mloda_core.abstract_plugins.components.feature import Feature
-from mloda_core.abstract_plugins.components.feature_chainer.feature_chain_parser import FeatureChainParser
-from mloda_core.abstract_plugins.components.feature_name import FeatureName
-from mloda_core.abstract_plugins.components.feature_set import FeatureSet
-from mloda_core.abstract_plugins.components.options import Options
+from mloda import FeatureGroup
+from mloda import Feature
+from mloda.provider import FeatureChainParser
+from mloda.provider import (
+    FeatureChainParserMixin,
+)
+from mloda.user import FeatureName
+from mloda.provider import FeatureSet
+from mloda import Options
 from mloda_plugins.feature_group.experimental.default_options_key import DefaultOptionKeys
 
 
-class TimeWindowFeatureGroup(AbstractFeatureGroup):
+class TimeWindowFeatureGroup(FeatureChainParserMixin, FeatureGroup):
     # Option keys for time window configuration
     WINDOW_FUNCTION = "window_function"
     WINDOW_SIZE = "window_size"
@@ -66,21 +69,21 @@ class TimeWindowFeatureGroup(AbstractFeatureGroup):
 
     ## Requirements
     - The input data must have a datetime column that can be used for time-based operations
-    - By default, the feature group will use DefaultOptionKeys.reference_time (default: "time_filter")
+    - By default, the feature group will use DefaultOptionKeys.reference_time (default: "reference_time")
     - You can specify a custom time column by setting the reference_time option in the feature group options
 
     """
 
     @classmethod
-    def get_time_filter_feature(cls, options: Optional[Options] = None) -> str:
+    def get_reference_time_column(cls, options: Optional[Options] = None) -> str:
         """
-        Get the time filter feature name from options or use the default.
+        Get the reference time column name from options or use the default.
 
         Args:
-            options: Optional Options object that may contain a custom time filter feature name
+            options: Optional Options object that may contain a custom reference time column name
 
         Returns:
-            The time filter feature name to use
+            The reference time column name to use
         """
         reference_time_key = DefaultOptionKeys.reference_time.value
         if options and options.get(reference_time_key):
@@ -90,7 +93,7 @@ class TimeWindowFeatureGroup(AbstractFeatureGroup):
                     f"Invalid reference_time option: {reference_time}. Must be string. Is: {type(reference_time)}."
                 )
             return reference_time
-        return reference_time_key
+        return DefaultOptionKeys.reference_time.value
 
     # Define supported window functions
     WINDOW_FUNCTIONS = {
@@ -123,28 +126,28 @@ class TimeWindowFeatureGroup(AbstractFeatureGroup):
         # Window function parameter (context parameter)
         WINDOW_FUNCTION: {
             **WINDOW_FUNCTIONS,  # Reference existing WINDOW_FUNCTIONS dict
-            DefaultOptionKeys.mloda_context: True,  # Mark as context parameter
-            DefaultOptionKeys.mloda_strict_validation: True,  # Enable strict validation
+            DefaultOptionKeys.context: True,  # Mark as context parameter
+            DefaultOptionKeys.strict_validation: True,  # Enable strict validation
         },
         # Window size parameter (context parameter)
         WINDOW_SIZE: {
             "explanation": "Size of the time window (must be positive integer)",
-            DefaultOptionKeys.mloda_context: True,  # Mark as context parameter
-            DefaultOptionKeys.mloda_strict_validation: True,  # Enable strict validation
-            DefaultOptionKeys.mloda_validation_function: lambda x: (isinstance(x, int) and x > 0)
+            DefaultOptionKeys.context: True,  # Mark as context parameter
+            DefaultOptionKeys.strict_validation: True,  # Enable strict validation
+            DefaultOptionKeys.validation_function: lambda x: (isinstance(x, int) and x > 0)
             or (isinstance(x, str) and x.isdigit() and int(x) > 0),
         },
         # Time unit parameter (context parameter)
         TIME_UNIT: {
             **TIME_UNITS,  # Reference existing TIME_UNITS dict
-            DefaultOptionKeys.mloda_context: True,  # Mark as context parameter
-            DefaultOptionKeys.mloda_strict_validation: True,  # Enable strict validation
+            DefaultOptionKeys.context: True,  # Mark as context parameter
+            DefaultOptionKeys.strict_validation: True,  # Enable strict validation
         },
         # Source feature parameter (context parameter)
         DefaultOptionKeys.in_features: {
             "explanation": "Source feature to apply time window operation to",
-            DefaultOptionKeys.mloda_context: True,  # Mark as context parameter
-            DefaultOptionKeys.mloda_strict_validation: False,  # Flexible validation
+            DefaultOptionKeys.context: True,  # Mark as context parameter
+            DefaultOptionKeys.strict_validation: False,  # Flexible validation
         },
     }
 
@@ -152,6 +155,11 @@ class TimeWindowFeatureGroup(AbstractFeatureGroup):
     PATTERN = "__"
     PREFIX_PATTERN = r".*__([\w]+)_(\d+)_([\w]+)_window$"
 
+    # In-feature configuration for FeatureChainParserMixin
+    MIN_IN_FEATURES = 1
+    MAX_IN_FEATURES = 1
+
+    # Custom input_features needed to add time_filter_feature
     def input_features(self, options: Options, feature_name: FeatureName) -> Optional[Set[Feature]]:
         """Extract source feature from either configuration-based options or string parsing."""
 
@@ -160,7 +168,7 @@ class TimeWindowFeatureGroup(AbstractFeatureGroup):
         # Try string-based parsing first
         _, source_feature = FeatureChainParser.parse_feature_name(feature_name.name, [self.PREFIX_PATTERN])
         if source_feature is not None:
-            time_filter_feature = Feature(self.get_time_filter_feature(options))
+            time_filter_feature = Feature(self.get_reference_time_column(options))
             return {Feature(source_feature), time_filter_feature}
 
         # Fall back to configuration-based approach
@@ -170,9 +178,66 @@ class TimeWindowFeatureGroup(AbstractFeatureGroup):
                 f"Expected exactly one source feature, but found {len(source_features)}: {source_features}"
             )
 
-        time_filter_feature = Feature(self.get_time_filter_feature(options))
+        time_filter_feature = Feature(self.get_reference_time_column(options))
         return set(source_features) | {time_filter_feature}
 
+    @classmethod
+    def _extract_time_window_params(cls, feature: Feature) -> tuple[Optional[str], Optional[int], Optional[str]]:
+        """
+        Extract time window parameters (window_function, window_size, time_unit) from a feature.
+
+        Tries string-based parsing first using parse_time_window_prefix, falls back to configuration.
+
+        Args:
+            feature: The feature to extract parameters from
+
+        Returns:
+            Tuple of (window_function, window_size, time_unit), where any value may be None if not found
+        """
+        feature_name = feature.get_name()
+
+        # Try string-based parsing first
+        try:
+            window_function, window_size, time_unit = cls.parse_time_window_prefix(feature_name)
+            return window_function, window_size, time_unit
+        except ValueError:
+            pass
+
+        # Fall back to configuration
+        window_function = feature.options.get(cls.WINDOW_FUNCTION)
+        window_size = feature.options.get(cls.WINDOW_SIZE)
+        time_unit = feature.options.get(cls.TIME_UNIT)
+
+        # Convert window_size to int if it's a string
+        if window_size is not None and isinstance(window_size, str):
+            window_size = int(window_size)
+
+        return window_function, window_size, time_unit
+
+    @classmethod
+    def _extract_time_window_params_and_source_features(cls, feature: Feature) -> tuple[str, int, str, str]:
+        """
+        Extract time window parameters and source feature from a feature.
+
+        Tries string-based parsing first, falls back to configuration-based approach.
+
+        Args:
+            feature: The feature to extract parameters from
+
+        Returns:
+            Tuple of (window_function, window_size, time_unit, source_feature_name)
+
+        Raises:
+            ValueError: If parameters cannot be extracted
+        """
+        source_features = cls._extract_source_features(feature)
+        window_function, window_size, time_unit = cls._extract_time_window_params(feature)
+
+        if window_function is None or window_size is None or time_unit is None:
+            raise ValueError(f"Could not extract time window parameters from: {feature.name}")
+
+        return window_function, window_size, time_unit, source_features[0]
+
     @classmethod
     def parse_time_window_prefix(cls, feature_name: str) -> tuple[str, int, str]:
         """
@@ -243,24 +308,7 @@ class TimeWindowFeatureGroup(AbstractFeatureGroup):
         """Extract the time unit from the feature name."""
         return cls.parse_time_window_prefix(feature_name)[2]
 
-    @classmethod
-    def match_feature_group_criteria(
-        cls,
-        feature_name: Union[FeatureName, str],
-        options: Options,
-        data_access_collection: Optional[Any] = None,
-    ) -> bool:
-        """Check if feature name matches the expected pattern for time window features."""
-        if isinstance(feature_name, FeatureName):
-            feature_name = feature_name.name
-
-        # Use unified parser approach with PROPERTY_MAPPING
-        return FeatureChainParser.match_configuration_feature_chain_parser(
-            feature_name,
-            options,
-            property_mapping=cls.PROPERTY_MAPPING,
-            prefix_patterns=[cls.PREFIX_PATTERN],
-        )
+    # match_feature_group_criteria() inherited from FeatureChainParserMixin
 
     @classmethod
     def calculate_feature(cls, data: Any, features: FeatureSet) -> Any:
@@ -283,52 +331,17 @@ class TimeWindowFeatureGroup(AbstractFeatureGroup):
                     raise ValueError("All features must have the same options.")
             _options = feature.options
 
-        time_filter_feature = cls.get_time_filter_feature(_options)
+        reference_time_column = cls.get_reference_time_column(_options)
 
-        cls._check_time_filter_feature_exists(data, time_filter_feature)
+        cls._check_reference_time_column_exists(data, reference_time_column)
 
-        cls._check_time_filter_feature_is_datetime(data, time_filter_feature)
+        cls._check_reference_time_column_is_datetime(data, reference_time_column)
 
         # Process each requested feature
         for feature in features.features:
-            feature_name = feature.get_name()
-
-            # Try string-based parsing first (for legacy features)
-            parsed_params, in_features = FeatureChainParser.parse_feature_name(feature_name, [cls.PREFIX_PATTERN])
-
-            if in_features is not None:
-                # String-based approach succeeded
-                window_function, window_size, time_unit = cls.parse_time_window_prefix(feature_name)
-            else:
-                # Fall back to configuration-based approach
-                has_config_params = (
-                    feature.options.get(cls.WINDOW_FUNCTION) is not None
-                    and feature.options.get(cls.WINDOW_SIZE) is not None
-                    and feature.options.get(cls.TIME_UNIT) is not None
-                )
-
-                if not has_config_params:
-                    raise ValueError(
-                        f"Feature '{feature_name}' does not match string pattern and lacks configuration parameters"
-                    )
-
-                # Configuration-based approach
-                source_features = feature.options.get_in_features()
-                if len(source_features) != 1:
-                    raise ValueError(
-                        f"Expected exactly one source feature, but found {len(source_features)}: {source_features}"
-                    )
-                source_feature = next(iter(source_features))
-                in_features = source_feature.get_name()
-
-                # Extract parameters from options
-                window_function = feature.options.get(cls.WINDOW_FUNCTION)
-                window_size = feature.options.get(cls.WINDOW_SIZE)
-                time_unit = feature.options.get(cls.TIME_UNIT)
-
-                # Convert window_size to int if it's a string
-                if isinstance(window_size, str):
-                    window_size = int(window_size)
+            window_function, window_size, time_unit, in_features = cls._extract_time_window_params_and_source_features(
+                feature
+            )
 
             # Resolve multi-column features automatically
             # If in_features is "onehot_encoded__product", this discovers
@@ -340,7 +353,7 @@ class TimeWindowFeatureGroup(AbstractFeatureGroup):
             cls._check_source_features_exist(data, resolved_columns)
 
             result = cls._perform_window_operation(
-                data, window_function, window_size, time_unit, resolved_columns, time_filter_feature
+                data, window_function, window_size, time_unit, resolved_columns, reference_time_column
             )
 
             data = cls._add_result_to_data(data, feature.get_name(), result)
@@ -349,31 +362,31 @@ class TimeWindowFeatureGroup(AbstractFeatureGroup):
 
     @classmethod
     @abstractmethod
-    def _check_time_filter_feature_exists(cls, data: Any, time_filter_feature: str) -> None:
+    def _check_reference_time_column_exists(cls, data: Any, reference_time_column: str) -> None:
         """
-        Check if the time filter feature exists in the data.
+        Check if the reference time column exists in the data.
 
         Args:
             data: The input data
-            time_filter_feature: The name of the time filter feature
+            reference_time_column: The name of the reference time column
 
         Raises:
-            ValueError: If the time filter feature does not exist in the data
+            ValueError: If the reference time column does not exist in the data
         """
         ...
 
     @classmethod
     @abstractmethod
-    def _check_time_filter_feature_is_datetime(cls, data: Any, time_filter_feature: str) -> None:
+    def _check_reference_time_column_is_datetime(cls, data: Any, reference_time_column: str) -> None:
         """
-        Check if the time filter feature is a datetime column.
+        Check if the reference time column is a datetime column.
 
         Args:
             data: The input data
-            time_filter_feature: The name of the time filter feature
+            reference_time_column: The name of the reference time column
 
         Raises:
-            ValueError: If the time filter feature is not a datetime column
+            ValueError: If the reference time column is not a datetime column
         """
         ...
 
@@ -447,7 +460,7 @@ class TimeWindowFeatureGroup(AbstractFeatureGroup):
             time_unit: The time unit for the window
             in_features: List of resolved source feature names to perform window operation on
             time_filter_feature: The name of the time filter feature to use for time-based operations.
-                                If None, uses the value from get_time_filter_feature().
+                                If None, uses the value from get_reference_time_column().
 
         Returns:
             The result of the window operation

mloda_plugins/feature_group/experimental/time_window/pandas.py

@@ -6,7 +6,7 @@ from __future__ import annotations
 
 from typing import Any, List, Optional, Set, Type, Union
 
-from mloda_core.abstract_plugins.compute_frame_work import ComputeFrameWork
+from mloda import ComputeFramework
 from mloda_plugins.compute_framework.base_implementations.pandas.dataframe import PandasDataFrame
 from mloda_plugins.feature_group.experimental.time_window.base import TimeWindowFeatureGroup
 
@@ -19,25 +19,25 @@ except ImportError:
 
 class PandasTimeWindowFeatureGroup(TimeWindowFeatureGroup):
     @classmethod
-    def compute_framework_rule(cls) -> Union[bool, Set[Type[ComputeFrameWork]]]:
+    def compute_framework_rule(cls) -> Union[bool, Set[Type[ComputeFramework]]]:
         return {PandasDataFrame}
 
     @classmethod
-    def _check_time_filter_feature_exists(cls, data: pd.DataFrame, time_filter_feature: str) -> None:
-        """Check if the time filter feature exists in the DataFrame."""
-        if time_filter_feature not in data.columns:
+    def _check_reference_time_column_exists(cls, data: pd.DataFrame, reference_time_column: str) -> None:
+        """Check if the reference time column exists in the DataFrame."""
+        if reference_time_column not in data.columns:
             raise ValueError(
-                f"Time filter feature '{time_filter_feature}' not found in data. "
+                f"Reference time column '{reference_time_column}' not found in data. "
                 f"Please ensure the DataFrame contains this column."
             )
 
     @classmethod
-    def _check_time_filter_feature_is_datetime(cls, data: pd.DataFrame, time_filter_feature: str) -> None:
-        """Check if the time filter feature is a datetime column."""
-        if not pd.api.types.is_datetime64_any_dtype(data[time_filter_feature]):
+    def _check_reference_time_column_is_datetime(cls, data: pd.DataFrame, reference_time_column: str) -> None:
+        """Check if the reference time column is a datetime column."""
+        if not pd.api.types.is_datetime64_any_dtype(data[reference_time_column]):
             raise ValueError(
-                f"Time filter feature '{time_filter_feature}' must be a datetime column. "
-                f"Current dtype: {data[time_filter_feature].dtype}"
+                f"Reference time column '{reference_time_column}' must be a datetime column. "
+                f"Current dtype: {data[reference_time_column].dtype}"
             )
 
     @classmethod
@@ -93,14 +93,14 @@ class PandasTimeWindowFeatureGroup(TimeWindowFeatureGroup):
             time_unit: The time unit for the window
             in_features: List of source feature names (may be single or multiple columns)
             time_filter_feature: The name of the time filter feature to use for time-based operations.
-                                If None, uses the value from get_time_filter_feature().
+                                If None, uses the value from get_reference_time_column().
 
         Returns:
             The result of the window operation
         """
         # Use the default time filter feature if none is provided
         if time_filter_feature is None:
-            time_filter_feature = cls.get_time_filter_feature()
+            time_filter_feature = cls.get_reference_time_column()
 
         # Create a copy of the DataFrame with the time filter feature as the index
         # This is necessary for time-based rolling operations

mloda_plugins/feature_group/experimental/time_window/pyarrow.py

@@ -10,7 +10,7 @@ import datetime
 import pyarrow as pa
 import pyarrow.compute as pc
 
-from mloda_core.abstract_plugins.compute_frame_work import ComputeFrameWork
+from mloda import ComputeFramework
 
 from mloda_plugins.compute_framework.base_implementations.pyarrow.table import PyArrowTable
 from mloda_plugins.feature_group.experimental.time_window.base import TimeWindowFeatureGroup
@@ -18,25 +18,25 @@ from mloda_plugins.feature_group.experimental.time_window.base import TimeWindow
 
 class PyArrowTimeWindowFeatureGroup(TimeWindowFeatureGroup):
     @classmethod
-    def compute_framework_rule(cls) -> Union[bool, Set[Type[ComputeFrameWork]]]:
+    def compute_framework_rule(cls) -> Union[bool, Set[Type[ComputeFramework]]]:
         return {PyArrowTable}
 
     @classmethod
-    def _check_time_filter_feature_exists(cls, data: pa.Table, time_filter_feature: str) -> None:
-        """Check if the time filter feature exists in the Table."""
-        if time_filter_feature not in data.schema.names:
+    def _check_reference_time_column_exists(cls, data: pa.Table, reference_time_column: str) -> None:
+        """Check if the reference time column exists in the Table."""
+        if reference_time_column not in data.schema.names:
             raise ValueError(
-                f"Time filter feature '{time_filter_feature}' not found in data. "
+                f"Reference time column '{reference_time_column}' not found in data. "
                 f"Please ensure the Table contains this column."
             )
 
     @classmethod
-    def _check_time_filter_feature_is_datetime(cls, data: pa.Table, time_filter_feature: str) -> None:
-        """Check if the time filter feature is a datetime column."""
-        time_column = data.column(time_filter_feature)
+    def _check_reference_time_column_is_datetime(cls, data: pa.Table, reference_time_column: str) -> None:
+        """Check if the reference time column is a datetime column."""
+        time_column = data.column(reference_time_column)
         if not pa.types.is_timestamp(time_column.type):
             raise ValueError(
-                f"Time filter feature '{time_filter_feature}' must be a timestamp column. "
+                f"Reference time column '{reference_time_column}' must be a timestamp column. "
                 f"Current type: {time_column.type}"
             )
 
@@ -103,14 +103,14 @@ class PyArrowTimeWindowFeatureGroup(TimeWindowFeatureGroup):
             time_unit: The time unit for the window
             in_features: List of source feature names (may be single or multiple columns)
             time_filter_feature: The name of the time filter feature to use for time-based operations.
-                                If None, uses the value from get_time_filter_feature().
+                                If None, uses the value from get_reference_time_column().
 
         Returns:
             The result of the window operation as a PyArrow Array
         """
         # Use the default time filter feature if none is provided
         if time_filter_feature is None:
-            time_filter_feature = cls.get_time_filter_feature()
+            time_filter_feature = cls.get_reference_time_column()
 
         # Get the time column
         time_column = data.column(time_filter_feature)

mloda_plugins/feature_group/input_data/api_data/api_data.py

@@ -1,11 +1,9 @@
 from typing import Any, Optional
-from mloda_core.abstract_plugins.abstract_feature_group import AbstractFeatureGroup
-from mloda_core.abstract_plugins.components.feature_set import FeatureSet
-from mloda_core.abstract_plugins.components.input_data.api.api_input_data import ApiInputData
-from mloda_core.abstract_plugins.components.input_data.base_input_data import BaseInputData
+from mloda import FeatureGroup
+from mloda.provider import FeatureSet, ApiData as ApiInputData, BaseInputData
 
 
-class ApiInputDataFeature(AbstractFeatureGroup):
+class ApiInputDataFeature(FeatureGroup):
     """
     Base class for API-based input data feature groups.
 
@@ -39,9 +37,9 @@ class ApiInputDataFeature(AbstractFeatureGroup):
     Uses Options to specify API data access configuration:
 
     ```python
-    from mloda_core.abstract_plugins.components.feature import Feature
-    from mloda_core.abstract_plugins.components.options import Options
-    from mloda_core.abstract_plugins.components.input_data.api.api_input_data import ApiInputData
+    from mloda import Feature
+    from mloda import Options
+    from mloda.core.abstract_plugins.components.input_data.api.api_input_data import ApiInputData
 
     feature = Feature(
         name="user_profile",
@@ -60,8 +58,8 @@ class ApiInputDataFeature(AbstractFeatureGroup):
     ### Basic API Data Access
 
     ```python
-    from mloda_core.abstract_plugins.components.feature import Feature
-    from mloda_core.abstract_plugins.components.options import Options
+    from mloda import Feature
+    from mloda import Options
 
     # Simple API feature reference
     feature = Feature(name="api_user_score")
@@ -70,7 +68,7 @@ class ApiInputDataFeature(AbstractFeatureGroup):
     ### Configuration-Based with Endpoint Mapping
 
     ```python
-    from mloda_core.abstract_plugins.components.input_data.api.api_input_data import ApiInputData
+    from mloda.core.abstract_plugins.components.input_data.api.api_input_data import ApiInputData
 
     # Map multiple API response fields
     feature = Feature(

mloda_plugins/feature_group/input_data/read_context_files.py

@@ -3,12 +3,12 @@ from pathlib import Path
 from typing import Any, List, Set
 
 
-from mloda_core.abstract_plugins.abstract_feature_group import AbstractFeatureGroup
-from mloda_core.abstract_plugins.components.feature import Feature
-from mloda_core.abstract_plugins.components.feature_name import FeatureName
-from mloda_core.abstract_plugins.components.feature_set import FeatureSet
-from mloda_core.abstract_plugins.components.link import JoinType
-from mloda_core.abstract_plugins.components.options import Options
+from mloda import FeatureGroup
+from mloda import Feature
+from mloda.user import FeatureName
+from mloda.provider import FeatureSet
+from mloda.user import JoinType
+from mloda import Options
 from mloda_plugins.feature_group.experimental.default_options_key import DefaultOptionKeys
 from mloda_plugins.feature_group.experimental.dynamic_feature_group_factory.dynamic_feature_group_factory import (
     DynamicFeatureGroupCreator,
@@ -27,7 +27,7 @@ except ImportError:
     pd = None
 
 
-class ConcatenatedFileContent(AbstractFeatureGroup):
+class ConcatenatedFileContent(FeatureGroup):
     """
     A feature group that reads and combines content from files within a directory (default: python files).
 

mloda_plugins/feature_group/input_data/read_db.py

@@ -1,9 +1,7 @@
 from typing import Any, Dict, List, Optional, Tuple, Union
-from mloda_core.abstract_plugins.components.data_access_collection import DataAccessCollection
-from mloda_core.abstract_plugins.components.feature_set import FeatureSet
-from mloda_core.abstract_plugins.components.hashable_dict import HashableDict
-from mloda_core.abstract_plugins.components.input_data.base_input_data import BaseInputData
-from mloda_core.abstract_plugins.components.options import Options
+from mloda.user import DataAccessCollection
+from mloda.provider import FeatureSet, HashableDict, BaseInputData
+from mloda import Options
 
 
 class ReadDB(BaseInputData):
@@ -85,11 +83,11 @@ class ReadDB(BaseInputData):
             data_accesses.append(data_access)
 
         if not data_accesses:
-            return False
+            return None
 
         matched_data_access = cls.match_read_db_data_access(data_accesses, feature_names)
-        if matched_data_access is False:
-            return False
+        if matched_data_access is None:
+            return None
         return matched_data_access
 
     @classmethod
@@ -110,7 +108,7 @@ class ReadDB(BaseInputData):
                     return data_access
             except NotImplementedError:
                 continue
-        return False
+        return None
 
     @classmethod
     def get_connection(cls, credentials: Any) -> Any:

mloda_plugins/feature_group/input_data/read_db_feature.py

@@ -1,12 +1,12 @@
 from typing import Any, Optional
 
-from mloda_core.abstract_plugins.abstract_feature_group import AbstractFeatureGroup
-from mloda_core.abstract_plugins.components.feature_set import FeatureSet
-from mloda_core.abstract_plugins.components.input_data.base_input_data import BaseInputData
+from mloda import FeatureGroup
+from mloda.provider import FeatureSet
+from mloda.provider import BaseInputData
 from mloda_plugins.feature_group.input_data.read_db import ReadDB
 
 
-class ReadDBFeature(AbstractFeatureGroup):
+class ReadDBFeature(FeatureGroup):
     @classmethod
     def input_data(cls) -> Optional[BaseInputData]:
         return ReadDB()

mloda_plugins/feature_group/input_data/read_dbs/sqlite.py

@@ -4,10 +4,9 @@ from typing import Any
 import pyarrow as pa
 import sqlite3
 
-from mloda_core.abstract_plugins.components.feature_set import FeatureSet
-from mloda_core.abstract_plugins.components.hashable_dict import HashableDict
-from mloda_core.abstract_plugins.components.data_types import DataType
-from mloda_core.abstract_plugins.components.options import Options
+from mloda.provider import FeatureSet, HashableDict
+from mloda.user import DataType
+from mloda import Options
 from mloda_plugins.feature_group.input_data.read_db import ReadDB
 
 
@@ -46,9 +45,9 @@ class SQLITEReader(ReadDB):
     Uses Options with database credentials and configuration:
 
     ```python
-    from mloda_core.abstract_plugins.components.feature import Feature
-    from mloda_core.abstract_plugins.components.options import Options
-    from mloda_core.abstract_plugins.components.hashable_dict import HashableDict
+    from mloda import Feature
+    from mloda import Options
+    from mloda.core.abstract_plugins.components.hashable_dict import HashableDict
 
     feature = Feature(
         name="customer_name",
@@ -68,8 +67,8 @@ class SQLITEReader(ReadDB):
     ### Basic SQLite Feature Access
 
     ```python
-    from mloda_core.abstract_plugins.components.feature import Feature
-    from mloda_core.abstract_plugins.components.hashable_dict import HashableDict
+    from mloda import Feature
+    from mloda.core.abstract_plugins.components.hashable_dict import HashableDict
 
     # Simple column reference from SQLite database
     feature = Feature(
@@ -117,7 +116,7 @@ class SQLITEReader(ReadDB):
     ### Using DataAccessCollection
 
     ```python
-    from mloda_core.abstract_plugins.components.data_access_collection import DataAccessCollection
+    from mloda.user import DataAccessCollection
 
     # Configure database access at the collection level
     data_access = DataAccessCollection(
@@ -214,11 +213,22 @@ class SQLITEReader(ReadDB):
     def load_data(cls, data_access: Any, features: FeatureSet) -> Any:
         query = cls.build_query(features)
         result, column_names = cls.read_db(data_access, query)
-        return cls.read_as_pa_data(result, column_names)
+        return cls.read_as_pa_data(result, column_names, features)
 
     @classmethod
-    def read_as_pa_data(cls, result: Any, column_names: Any) -> Any:
-        schema = pa.schema([(column_names[i], DataType.infer_arrow_type(result[0][i])) for i in range(len(result[0]))])
+    def read_as_pa_data(cls, result: Any, column_names: Any, features: Any) -> Any:
+        feature_map = {f.get_name(): f for f in features.features}
+
+        schema_fields = []
+        for i, col_name in enumerate(column_names):
+            feature = feature_map.get(col_name)
+            if feature and feature.data_type:
+                arrow_type = DataType.to_arrow_type(feature.data_type)
+            else:
+                arrow_type = DataType.infer_arrow_type(result[0][i])
+            schema_fields.append((col_name, arrow_type))
+
+        schema = pa.schema(schema_fields)
         data_dicts = [{column_names[i]: row[i] for i in range(len(row))} for row in result]
         table = pa.Table.from_pylist(data_dicts, schema=schema)
         return table