36 lines
1.3 KiB
Python
36 lines
1.3 KiB
Python
import logging
|
|
|
|
import pandas as pd
|
|
|
|
from .. import constants
|
|
|
|
logging.basicConfig(level=logging.INFO)
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def main():
|
|
# Entrees : dataset source deja prepare (hors git)
|
|
feature_file_path = constants.SOURCE_DIR / constants.FEATURE_FILENAME
|
|
target_file_path = constants.SOURCE_DIR / constants.TARGET_FILENAME
|
|
|
|
logger.info(f"Read dataset from {feature_file_path} and {target_file_path}")
|
|
df_features = pd.read_parquet(feature_file_path)
|
|
df_target = pd.read_parquet(target_file_path)
|
|
|
|
df = df_features.join(df_target)
|
|
timestamps = df.index.get_level_values("timestamp")
|
|
|
|
# Sorties : splits versionnes par DVC dans le depot
|
|
constants.DATASET_DIR.mkdir(parents=True, exist_ok=True)
|
|
logger.info(f"Split strategy: {constants.CHOSEN_SPLIT_STRATEGY.value}")
|
|
for dataset_name, (start_date, end_date) in constants.DATASET_SPLIT_DATES[constants.CHOSEN_SPLIT_STRATEGY].items():
|
|
file_path = constants.DATASET_DIR / f"{dataset_name}.parquet"
|
|
mask = (timestamps.date >= start_date) & (timestamps.date <= end_date)
|
|
df_split = df[mask]
|
|
logger.info(f"Split dataset into {dataset_name} ({start_date} -> {end_date}) with shape: {df_split.shape}")
|
|
df_split.to_parquet(file_path)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|