csit.infra.dash/app/cdash/data/data.py

   1 # Copyright (c) 2023 Cisco and/or its affiliates.
   2 # Licensed under the Apache License, Version 2.0 (the "License");
   3 # you may not use this file except in compliance with the License.
   4 # You may obtain a copy of the License at:
   5 #
   6 #     http://www.apache.org/licenses/LICENSE-2.0
   7 #
   8 # Unless required by applicable law or agreed to in writing, software
   9 # distributed under the License is distributed on an "AS IS" BASIS,
  10 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  11 # See the License for the specific language governing permissions and
  12 # limitations under the License.
  13
  14 """Prepare data for Plotly Dash applications.
  15 """
  16
  17 import logging
  18 import resource
  19 import awswrangler as wr
  20 import pandas as pd
  21 import pyarrow as pa
  22
  23 from yaml import load, FullLoader, YAMLError
  24 from datetime import datetime, timedelta
  25 from time import time
  26 from pytz import UTC
  27 from awswrangler.exceptions import EmptyDataFrame, NoFilesFound
  28 from pyarrow.lib import ArrowInvalid, ArrowNotImplementedError
  29
  30 from ..utils.constants import Constants as C
  31
  32
  33 # If True, pyarrow.Schema is generated. See also condition in the method
  34 # _write_parquet_schema.
  35 # To generate schema, select only one data set in data.yaml file.
  36 GENERATE_SCHEMA = False
  37
  38
  39 class Data:
  40     """Gets the data from parquets and stores it for further use by dash
  41     applications.
  42     """
  43
  44     def __init__(self, data_spec_file: str) -> None:
  45         """Initialize the Data object.
  46
  47         :param data_spec_file: Path to file specifying the data to be read from
  48             parquets.
  49         :type data_spec_file: str
  50         :raises RuntimeError: if it is not possible to open data_spec_file or it
  51             is not a valid yaml file.
  52         """
  53
  54         # Inputs:
  55         self._data_spec_file = data_spec_file
  56
  57         # Specification of data to be read from parquets:
  58         self._data_spec = list()
  59
  60         # Data frame to keep the data:
  61         self._data = {
  62             "statistics": pd.DataFrame(),
  63             "trending": pd.DataFrame(),
  64             "iterative": pd.DataFrame(),
  65             "coverage": pd.DataFrame()
  66         }
  67
  68         # Read from files:
  69         try:
  70             with open(self._data_spec_file, "r") as file_read:
  71                 self._data_spec = load(file_read, Loader=FullLoader)
  72         except IOError as err:
  73             raise RuntimeError(
  74                 f"Not possible to open the file {self._data_spec_file,}\n{err}"
  75             )
  76         except YAMLError as err:
  77             raise RuntimeError(
  78                 f"An error occurred while parsing the specification file "
  79                 f"{self._data_spec_file,}\n"
  80                 f"{err}"
  81             )
  82
  83     @property
  84     def data(self):
  85         return self._data
  86
  87     @staticmethod
  88     def _get_list_of_files(
  89             path,
  90             last_modified_begin=None,
  91             last_modified_end=None,
  92             days=None
  93         ) -> list:
  94         """Get list of interested files stored in S3 compatible storage and
  95         returns it.
  96
  97         :param path: S3 prefix (accepts Unix shell-style wildcards)
  98             (e.g. s3://bucket/prefix) or list of S3 objects paths
  99             (e.g. [s3://bucket/key0, s3://bucket/key1]).
 100         :param last_modified_begin: Filter the s3 files by the Last modified
 101             date of the object. The filter is applied only after list all s3
 102             files.
 103         :param last_modified_end: Filter the s3 files by the Last modified date
 104             of the object. The filter is applied only after list all s3 files.
 105         :param days: Number of days to filter.
 106         :type path: Union[str, List[str]]
 107         :type last_modified_begin: datetime, optional
 108         :type last_modified_end: datetime, optional
 109         :type days: integer, optional
 110         :returns: List of file names.
 111         :rtype: list
 112         """
 113         file_list = list()
 114         if days:
 115             last_modified_begin = datetime.now(tz=UTC) - timedelta(days=days)
 116         try:
 117             file_list = wr.s3.list_objects(
 118                 path=path,
 119                 suffix="parquet",
 120                 last_modified_begin=last_modified_begin,
 121                 last_modified_end=last_modified_end
 122             )
 123             logging.debug("\n".join(file_list))
 124         except NoFilesFound as err:
 125             logging.error(f"No parquets found.\n{err}")
 126         except EmptyDataFrame as err:
 127             logging.error(f"No data.\n{err}")
 128
 129         return file_list
 130
 131     def _validate_columns(self, data_type: str) -> str:
 132         """Check if all columns are present in the dataframe.
 133
 134         :param data_type: The data type defined in data.yaml
 135         :type data_type: str
 136         :returns: Error message if validation fails, otherwise empty string.
 137         :rtype: str
 138         """
 139         defined_columns = set()
 140         for data_set in self._data_spec:
 141             if data_set.get("data_type", str()) == data_type:
 142                 defined_columns.update(data_set.get("columns", set()))
 143
 144         if not defined_columns:
 145             return "No columns defined in the data set(s)."
 146
 147         if self.data[data_type].empty:
 148             return "No data."
 149
 150         ret_msg = str()
 151         for col in defined_columns:
 152             if col not in self.data[data_type].columns:
 153                 if not ret_msg:
 154                     ret_msg = "Missing columns: "
 155                 else:
 156                     ret_msg += ", "
 157                 ret_msg += f"{col}"
 158         return ret_msg
 159
 160     @staticmethod
 161     def _write_parquet_schema(
 162             path,
 163             partition_filter=None,
 164             columns=None,
 165             validate_schema=False,
 166             last_modified_begin=None,
 167             last_modified_end=None,
 168             days=None
 169         ) -> None:
 170         """Auxiliary function to write parquet schemas. Use it instead of
 171         "_create_dataframe_from_parquet" in "read_all_data".
 172
 173         :param path: S3 prefix (accepts Unix shell-style wildcards)
 174             (e.g. s3://bucket/prefix) or list of S3 objects paths
 175             (e.g. [s3://bucket/key0, s3://bucket/key1]).
 176         :param partition_filter: Callback Function filters to apply on PARTITION
 177             columns (PUSH-DOWN filter). This function MUST receive a single
 178             argument (Dict[str, str]) where keys are partitions names and values
 179             are partitions values. Partitions values will be always strings
 180             extracted from S3. This function MUST return a bool, True to read
 181             the partition or False to ignore it. Ignored if dataset=False.
 182         :param columns: Names of columns to read from the file(s).
 183         :param validate_schema: Check that individual file schemas are all the
 184             same / compatible. Schemas within a folder prefix should all be the
 185             same. Disable if you have schemas that are different and want to
 186             disable this check.
 187         :param last_modified_begin: Filter the s3 files by the Last modified
 188             date of the object. The filter is applied only after list all s3
 189             files.
 190         :param last_modified_end: Filter the s3 files by the Last modified date
 191             of the object. The filter is applied only after list all s3 files.
 192         :param days: Number of days to filter.
 193         :type path: Union[str, List[str]]
 194         :type partition_filter: Callable[[Dict[str, str]], bool], optional
 195         :type columns: List[str], optional
 196         :type validate_schema: bool, optional
 197         :type last_modified_begin: datetime, optional
 198         :type last_modified_end: datetime, optional
 199         :type days: integer, optional
 200         """
 201         if days:
 202             last_modified_begin = datetime.now(tz=UTC) - timedelta(days=days)
 203
 204         df = wr.s3.read_parquet(
 205             path=path,
 206             path_suffix="parquet",
 207             ignore_empty=True,
 208             validate_schema=validate_schema,
 209             use_threads=True,
 210             dataset=True,
 211             columns=columns,
 212             partition_filter=partition_filter,
 213             last_modified_begin=last_modified_begin,
 214             last_modified_end=last_modified_end,
 215             chunked=1
 216         )
 217
 218         for itm in df:
 219             try:
 220                 # Specify the condition or remove it:
 221                 if all((
 222                         pd.api.types.is_string_dtype(itm["<column_name>"]),
 223                         pd.api.types.is_string_dtype(itm["telemetry"][0])
 224                     )):
 225                     print(pa.Schema.from_pandas(itm))
 226                     pa.parquet.write_metadata(
 227                         pa.Schema.from_pandas(itm),
 228                         f"{C.PATH_TO_SCHEMAS}_tmp_schema"
 229                     )
 230                     print(itm)
 231                     break
 232             except KeyError:
 233                 pass
 234
 235     @staticmethod
 236     def _create_dataframe_from_parquet(
 237             path,
 238             partition_filter=None,
 239             columns=None,
 240             validate_schema=False,
 241             last_modified_begin=None,
 242             last_modified_end=None,
 243             days=None,
 244             schema=None
 245         ) -> pd.DataFrame:
 246         """Read parquet stored in S3 compatible storage and returns Pandas
 247         Dataframe.
 248
 249         :param path: S3 prefix (accepts Unix shell-style wildcards)
 250             (e.g. s3://bucket/prefix) or list of S3 objects paths
 251             (e.g. [s3://bucket/key0, s3://bucket/key1]).
 252         :param partition_filter: Callback Function filters to apply on PARTITION
 253             columns (PUSH-DOWN filter). This function MUST receive a single
 254             argument (Dict[str, str]) where keys are partitions names and values
 255             are partitions values. Partitions values will be always strings
 256             extracted from S3. This function MUST return a bool, True to read
 257             the partition or False to ignore it. Ignored if dataset=False.
 258         :param columns: Names of columns to read from the file(s).
 259         :param validate_schema: Check that individual file schemas are all the
 260             same / compatible. Schemas within a folder prefix should all be the
 261             same. Disable if you have schemas that are different and want to
 262             disable this check.
 263         :param last_modified_begin: Filter the s3 files by the Last modified
 264             date of the object. The filter is applied only after list all s3
 265             files.
 266         :param last_modified_end: Filter the s3 files by the Last modified date
 267             of the object. The filter is applied only after list all s3 files.
 268         :param days: Number of days to filter.
 269         :param schema: Path to schema to use when reading data from the parquet.
 270         :type path: Union[str, List[str]]
 271         :type partition_filter: Callable[[Dict[str, str]], bool], optional
 272         :type columns: List[str], optional
 273         :type validate_schema: bool, optional
 274         :type last_modified_begin: datetime, optional
 275         :type last_modified_end: datetime, optional
 276         :type days: integer, optional
 277         :type schema: string
 278         :returns: Pandas DataFrame or None if DataFrame cannot be fetched.
 279         :rtype: DataFrame
 280         """
 281         df = pd.DataFrame()
 282         start = time()
 283         if days:
 284             last_modified_begin = datetime.now(tz=UTC) - timedelta(days=days)
 285         try:
 286             df = wr.s3.read_parquet(
 287                 path=path,
 288                 path_suffix="parquet",
 289                 ignore_empty=True,
 290                 schema=schema,
 291                 validate_schema=validate_schema,
 292                 use_threads=True,
 293                 dataset=True,
 294                 columns=columns,
 295                 partition_filter=partition_filter,
 296                 last_modified_begin=last_modified_begin,
 297                 last_modified_end=last_modified_end,
 298                 dtype_backend="pyarrow"
 299             )
 300
 301             df.info(verbose=True, memory_usage="deep")
 302             logging.debug(
 303                 f"\nCreation of dataframe {path} took: {time() - start}\n"
 304             )
 305         except (ArrowInvalid, ArrowNotImplementedError) as err:
 306             logging.error(f"Reading of data from parquets FAILED.\n{repr(err)}")
 307         except NoFilesFound as err:
 308             logging.error(
 309                 f"Reading of data from parquets FAILED.\n"
 310                 f"No parquets found in specified time period.\n"
 311                 f"Nr of days: {days}\n"
 312                 f"last_modified_begin: {last_modified_begin}\n"
 313                 f"{repr(err)}"
 314             )
 315         except EmptyDataFrame as err:
 316             logging.error(
 317                 f"Reading of data from parquets FAILED.\n"
 318                 f"No data in parquets in specified time period.\n"
 319                 f"Nr of days: {days}\n"
 320                 f"last_modified_begin: {last_modified_begin}\n"
 321                 f"{repr(err)}"
 322             )
 323
 324         return df
 325
 326     def read_all_data(self, days: int=None) -> dict:
 327         """Read all data necessary for all applications.
 328
 329         :param days: Number of days to filter. If None, all data will be
 330             downloaded.
 331         :type days: int
 332         :returns: A dictionary where keys are names of parquets and values are
 333             the pandas dataframes with fetched data.
 334         :rtype: dict(str: pandas.DataFrame)
 335         """
 336
 337         data_lists = {
 338             "statistics": list(),
 339             "trending": list(),
 340             "iterative": list(),
 341             "coverage": list()
 342         }
 343
 344         logging.info("\n\nReading data:\n" + "-" * 13 + "\n")
 345         for data_set in self._data_spec:
 346             logging.info(
 347                 f"\n\nReading data for {data_set['data_type']} "
 348                 f"{data_set['partition_name']} {data_set.get('release', '')}\n"
 349             )
 350             schema_file = data_set.get("schema", None)
 351             if schema_file:
 352                 try:
 353                     schema = pa.parquet.read_schema(
 354                         f"{C.PATH_TO_SCHEMAS}{schema_file}"
 355                     )
 356                 except FileNotFoundError as err:
 357                     logging.error(repr(err))
 358                     logging.error("Proceeding without schema.")
 359                     schema = None
 360             else:
 361                 schema = None
 362             partition_filter = lambda part: True \
 363                 if part[data_set["partition"]] == data_set["partition_name"] \
 364                     else False
 365             if data_set["data_type"] in ("trending", "statistics"):
 366                 time_period = days
 367             else:
 368                 time_period = None
 369
 370             if GENERATE_SCHEMA:
 371                 # Generate schema:
 372                 Data._write_parquet_schema(
 373                     path=data_set["path"],
 374                     partition_filter=partition_filter,
 375                     columns=data_set.get("columns", None),
 376                     days=time_period
 377                 )
 378                 return
 379
 380             #  Read data:
 381             data = Data._create_dataframe_from_parquet(
 382                 path=data_set["path"],
 383                 partition_filter=partition_filter,
 384                 columns=data_set.get("columns", None),
 385                 days=time_period,
 386                 schema=schema
 387             )
 388             if data_set["data_type"] in ("iterative", "coverage"):
 389                 data["release"] = data_set["release"]
 390                 data["release"] = data["release"].astype("category")
 391
 392             data_lists[data_set["data_type"]].append(data)
 393
 394         logging.info(
 395             "\n\nData post-processing, validation and summary:\n" +
 396             "-" * 45 + "\n"
 397         )
 398         for key in self._data.keys():
 399             logging.info(f"\n\nDataframe {key}:\n")
 400             self._data[key] = pd.concat(
 401                 data_lists[key],
 402                 ignore_index=True,
 403                 copy=False
 404             )
 405             self._data[key].info(verbose=True, memory_usage="deep")
 406             err_msg = self._validate_columns(key)
 407             if err_msg:
 408                 self._data[key] = pd.DataFrame()
 409                 logging.error(
 410                     f"Data validation FAILED.\n"
 411                     f"{err_msg}\n"
 412                     "Generated dataframe replaced by an empty dataframe."
 413                 )
 414
 415         mem_alloc = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1000
 416         logging.info(f"\n\nMemory allocation: {mem_alloc:.0f}MB\n")
 417
 418         return self._data