Add parts argument to load_dataset function

[ostreech1997]

Before submitting (must do checklist)

• Did you read the contribution guide?
• Did you update the docs? We use Numpy format for all the methods and classes.
• Did you write any new necessary tests?
• Did you update the CHANGELOG?

Proposed Changes

Closing issues

Closes #74.

[codecov]

Codecov Report

❗ No coverage uploaded for pull request base (internal_datasets@959b3a0). Click here to learn what that means.
Patch has no changes to coverable lines.

Additional details and impacted files

@@                 Coverage Diff                  @@
##             internal_datasets      #79   +/-   ##
====================================================
  Coverage                     ?   88.61%           
====================================================
  Files                        ?      194           
  Lines                        ?    12488           
  Branches                     ?        0           
====================================================
  Hits                         ?    11066           
  Misses                       ?     1422           
  Partials                     ?        0           

☔ View full report in Codecov by Sentry.

📢 Have feedback on the report? Share it here.

[github-actions]

🚀 Deployed on https://deploy-preview-79--etna-docs.netlify.app

[ostreech1997]

Script for time measurements for saving data in wide and long formats.

import tempfile
import urllib.request
import zipfile
from pathlib import Path
import time

import pandas as pd

from etna.datasets.tsdataset import TSDataset


dataset_dir = Path.home() / ".etna" / "electricity_15T"


def _download_dataset_zip(url: str, file_name: str, **kwargs) -> pd.DataFrame:
    try:
        with tempfile.TemporaryDirectory() as td:
            temp_path = Path(td) / "temp.zip"
            urllib.request.urlretrieve(url, temp_path)
            with zipfile.ZipFile(temp_path) as f:
                f.extractall(td)
                df = pd.read_csv(Path(td) / file_name, **kwargs)
    except Exception as err:
        raise Exception(f"Error during downloading and reading dataset. Reason: {repr(err)}")
    return df


def prepare_data():
    url = "https://archive.ics.uci.edu/static/public/321/electricityloaddiagrams20112014.zip"
    dataset_dir.mkdir(exist_ok=True, parents=True)
    data = _download_dataset_zip(url=url, file_name="LD2011_2014.txt", sep=";", dtype=str)

    data = data.rename({"Unnamed: 0": "timestamp"}, axis=1)
    data["timestamp"] = pd.to_datetime(data["timestamp"])
    dt_list = sorted(data["timestamp"].unique())
    data = data.melt("timestamp", var_name="segment", value_name="target")
    data["target"] = data["target"].str.replace(",", ".").astype(float)
    data_train = data[data["timestamp"].isin(dt_list[: -15 * 24])]
    data_test = data[data["timestamp"].isin(dt_list[-15 * 24:])]
    return data, data_train, data_test


def save_wide():
    data, data_train, data_test = prepare_data()
    TSDataset.to_dataset(data).to_csv(dataset_dir / "electricity_15T_full.csv.gz", index=True, compression="gzip")
    TSDataset.to_dataset(data_train).to_csv(
        dataset_dir / "electricity_15T_train.csv.gz", index=True, compression="gzip"
    )
    TSDataset.to_dataset(data_test).to_csv(dataset_dir / "electricity_15T_test.csv.gz", index=True, compression="gzip")


def load_wide():
    data = pd.read_csv(
        dataset_dir / f"electricity_15T_full.csv.gz",
        compression="gzip",
        header=[0, 1],
        index_col=[0],
        parse_dates=[0]
    )
    _ = TSDataset(data, freq="15T")


def save_long():
    data, data_train, data_test = prepare_data()
    data.to_csv(dataset_dir / "electricity_15T_full.csv.gz", index=False, compression="gzip")
    data_train.to_csv(
        dataset_dir / "electricity_15T_train.csv.gz", index=False, compression="gzip"
    )
    data_test.to_csv(dataset_dir / "electricity_15T_test.csv.gz", index=False, compression="gzip")


def load_long():
    data = pd.read_csv(
        dataset_dir / f"electricity_15T_full.csv.gz",
        compression="gzip",
        parse_dates=[0]
    )
    _ = TSDataset(TSDataset.to_dataset(data), freq="15T")


def main():
    time_start = time.time()
    save_wide()
    time_end = time.time()
    print("Time for saving data in wide format:", (time_end - time_start) / 60)

    time_start = time.time()
    load_wide()
    time_end = time.time()
    print("Time for loading data in wide format:", (time_end - time_start) / 60)

    time_start = time.time()
    save_long()
    time_end = time.time()
    print("Time for saving data in long format:", (time_end - time_start) / 60)

    time_start = time.time()
    load_long()
    time_end = time.time()
    print("Time for loading data in long format:", (time_end - time_start) / 60)


if __name__ == "__main__":
    main()
    

Results:

Time for saving data in wide format: 6.355555299917857
Time for loading data in wide format: 0.18998863299687704
Time for saving data in long format: 11.959305040041606
Time for loading data in long format: 1.4660529494285583