src.fairreckitlib.data.set.dataset_sampling

  1"""This module contains functionality to create a sample of an existing dataset.
  2
  3Functions:
  4
  5    create_dataset_sample: create a sample of a dataset.
  6    create_dataset_table_samples: create tables samples for a map of key indices.
  7    create_matrix_sample_config: create a sample matrix configuration.
  8    create_matrix_sample: create a sample of a dataset's matrix.
  9
 10This program has been developed by students from the bachelor Computer Science at
 11Utrecht University within the Software Project course.
 12© Copyright Utrecht University (Department of Information and Computing Sciences)
 13"""
 14
 15import os
 16from typing import Dict, List, Optional, Tuple
 17
 18import pandas as pd
 19
 20from ...core.io.io_utility import save_yml
 21from .dataset import Dataset
 22from .dataset_constants import DATASET_CONFIG_FILE
 23from .dataset_config import DatasetMatrixConfig, DatasetIndexConfig, RatingMatrixConfig
 24from .dataset_config import DatasetConfig, DatasetTableConfig, create_dataset_table_config
 25
 26
 27def create_dataset_sample(
 28        output_dir: str,
 29        dataset: Dataset,
 30        num_users: int,
 31        num_items: int) -> Dataset:
 32    """Create a sample of the specified dataset.
 33
 34    Look at the 'create_matrix_sample' function for specifics on how the
 35    matrices of the dataset are sampled. All tables, except the events, that are related
 36    to the user/item keys that are present in the sample matrices are sampled as well.
 37    The generated dataset sample is stored in the output directory before returning it.
 38    This function raises an IOError when the specified output directory already exists.
 39
 40    Args:
 41        output_dir: the path to the directory where the dataset sample will be stored.
 42        dataset: the dataset to create a sample of.
 43        num_users: the number of users in the created sample matrices.
 44        num_items: the number of items in the created sample matrices.
 45
 46    Returns:
 47        the resulting sample dataset.
 48    """
 49    sample_dir = os.path.join(output_dir, dataset.get_name() + '-Sample')
 50    if os.path.isdir(sample_dir):
 51        raise IOError('Failed to create sample, directory already exists.')
 52
 53    os.mkdir(sample_dir)
 54
 55    sample_matrices = {}
 56    key_id_map = {}
 57
 58    for matrix_name in dataset.get_available_matrices():
 59        # create and add matrix sample
 60        sample_matrix_config = create_matrix_sample_config(
 61            sample_dir,
 62            dataset,
 63            matrix_name,
 64            num_users,
 65            num_items
 66        )
 67        sample_matrices[matrix_name] = sample_matrix_config
 68
 69        # append user key indices to the key map
 70        user_keys = key_id_map.get(sample_matrix_config.user.key, [])
 71        user_indices = pd.Series(sample_matrix_config.user.load_indices(sample_dir))
 72        user_keys = pd.Series(user_keys).append(user_indices)
 73        key_id_map[sample_matrix_config.user.key] = user_keys.unique()
 74
 75        # append item key indices to the key map
 76        item_keys = key_id_map.get(sample_matrix_config.item.key, [])
 77        item_indices = pd.Series(sample_matrix_config.item.load_indices(sample_dir))
 78        item_keys = pd.Series(item_keys).append(item_indices)
 79        key_id_map[sample_matrix_config.item.key] = item_keys.unique()
 80
 81    # create sample tables for the key map that contains all the needed indices of all matrices
 82    sample_tables = create_dataset_table_samples(sample_dir, dataset, key_id_map)
 83
 84    # create and save dataset configuration
 85    sample_dataset_config = DatasetConfig(
 86        dataset.get_name() + '-Sample',
 87        {},
 88        sample_matrices,
 89        sample_tables
 90    )
 91    save_yml(os.path.join(sample_dir, DATASET_CONFIG_FILE), sample_dataset_config.to_yml_format())
 92
 93    return Dataset(sample_dir, sample_dataset_config)
 94
 95
 96def create_dataset_table_samples(
 97        output_dir: str,
 98        dataset: Dataset,
 99        key_id_map: Dict[str, List[int]]) -> Dict[str, DatasetTableConfig]:
100    """Create table samples for the specified dataset and key map.
101
102    The key map is used to identify which tables of the dataset are sampled.
103    A table is considered to be a candidate if the key in the map matches the
104    primary key of the table. Any rows that do not contain the needed indices
105    in the key map are filtered.
106
107    Args:
108        output_dir: the path to the directory where the sample tables will be stored.
109        dataset: the dataset to create a sample tables from.
110        key_id_map: a dictionary containing a table key paired with a list of indices
111            that are related to these table keys.
112
113    Returns:
114        a dictionary with the resulting table sample configurations, keyed by table names.
115    """
116    sample_tables = {}
117
118    for table_name in dataset.get_available_tables():
119        table_config = dataset.get_table_config(table_name)
120
121        table = dataset.read_table(table_name)
122        table_modified = False
123
124        for key_id, key_id_list in key_id_map.items():
125            # filter unwanted table rows when the primary key matches
126            if table_config.primary_key == [key_id]:
127                table = table[table[key_id].isin(key_id_list)]
128                table_modified = True
129
130        if table_modified:
131            # store the table sample
132            sample_table_config = create_dataset_table_config(
133                dataset.get_name() + '_' + table_name + '.tsv.bz2',
134                table_config.primary_key,
135                table_config.columns,
136                compression='bz2',
137                encoding=table_config.file.options.encoding,
138                foreign_keys=table_config.foreign_keys,
139                num_records=len(table)
140            )
141            sample_table_config.save_table(table, output_dir)
142            # add sample table configuration
143            sample_tables[table_name] = sample_table_config
144
145    return sample_tables
146
147
148def create_matrix_sample_config(
149        output_dir: str,
150        dataset: Dataset,
151        matrix_name: str,
152        num_users: int,
153        num_items: int) -> Optional[DatasetMatrixConfig]:
154    """Create a dataset matrix sample configuration.
155
156    Look at the 'create_matrix_sample' function for specifics on how the
157    matrix is sampled. The generated matrix and user/item indirection arrays are
158    stored in the output directory and the corresponding configuration is returned.
159
160    Args:
161        output_dir: the path to the directory where the sample matrix will be stored.
162        dataset: the dataset to create a sample matrix from.
163        matrix_name: the name of the matrix to create a sample of.
164        num_users: the number of users in the created sample matrix.
165        num_items: the number of items in the created sample matrix.
166
167    Returns:
168        the sample matrix configuration or None when the specified matrix does not exist.
169    """
170    matrix_config = dataset.get_matrix_config(matrix_name)
171    if matrix_config is None:
172        return None
173
174    sample, users, items = create_matrix_sample(dataset, matrix_name, num_users, num_items)
175
176    # create the user indices config and save the array
177    user_index_config = DatasetIndexConfig(
178        matrix_name + '_user_indices.hdf5',
179        matrix_config.user.key,
180        len(users)
181    )
182    user_index_config.save_indices(output_dir, list(users))
183
184    # create the item indices config and save the array
185    item_index_config = DatasetIndexConfig(
186        matrix_name + '_item_indices.hdf5',
187        matrix_config.item.key,
188        len(items)
189    )
190    item_index_config.save_indices(output_dir, list(items))
191
192    # create the sample matrix table config and save the table
193    sample_table_config = create_dataset_table_config(
194        dataset.get_name() + '_' + matrix_name + '.tsv.bz2',
195        matrix_config.table.primary_key,
196        matrix_config.table.columns,
197        compression='bz2',
198        encoding=matrix_config.table.file.options.encoding,
199        foreign_keys=matrix_config.table.foreign_keys,
200        num_records=len(sample)
201    )
202    sample_table_config.save_table(sample, output_dir)
203
204    return DatasetMatrixConfig(
205        sample_table_config,
206        RatingMatrixConfig(
207            float(sample[matrix_config.table.columns[0]].min()),
208            float(sample[matrix_config.table.columns[0]].max()),
209            matrix_config.ratings.rating_type
210        ),
211        user_index_config,
212        item_index_config
213    )
214
215
216def create_matrix_sample(
217        dataset: Dataset,
218        matrix_name: str,
219        num_users: int,
220        num_items: int) -> Tuple[pd.DataFrame, List[int], List[int]]:
221    """Create a sample for the specified matrix.
222
223    Extracts a sample with the first occurring users and items until the
224    specified amounts are reached, and therefore are only used as an indication.
225    No additional users/items are generated when the dataset matrix has
226    less available amounts than is specified. Moreover, due to the sparsity of the
227    matrix it can turn out that the resulting matrix is very close, but not
228    exactly the specified amounts.
229
230    Args:
231        dataset: the dataset to create a sample matrix from.
232        matrix_name: the name of the matrix to create a sample of.
233        num_users: the number of users in the created sample matrix.
234        num_items: the number of items in the created sample matrix.
235
236    Returns:
237        the sample matrix, the unique user and unique item indices.
238    """
239    matrix_config = dataset.get_matrix_config(matrix_name)
240
241    # clamp num users/items
242    matrix_users = min(matrix_config.user.num_records, num_users)
243    matrix_items = min(matrix_config.item.num_records, num_items)
244
245    # prepare sample dataframe
246    matrix_columns = matrix_config.table.primary_key + matrix_config.table.columns
247    matrix_sample = pd.DataFrame(columns=matrix_columns)
248
249    user_key = matrix_config.user.key
250    item_key = matrix_config.item.key
251
252    # create sample in chunks for very big matrices
253    for _, matrix in enumerate(dataset.read_matrix(matrix_name, chunk_size=50000000)):
254        matrix_sample = pd.concat([matrix_sample, matrix])
255        if len(matrix_sample[user_key].unique()) > matrix_users and \
256                len(matrix_sample[item_key].unique()) > matrix_items:
257            break
258
259    # users may not be number from 0...num_users
260    unique_users = matrix_sample[user_key].unique()
261    matrix_sample = pd.merge(
262        matrix_sample,
263        pd.DataFrame(list(enumerate(unique_users)), columns=['user',user_key]),
264        how='left',
265        on=user_key
266    )
267    # remove any users above the threshold
268    matrix_sample = matrix_sample[matrix_sample['user'] < matrix_users]
269    # recalculate the indirection array
270    unique_users = dataset.resolve_user_ids(matrix_name, matrix_sample[user_key].unique().tolist())
271
272    # items may not be number from 0...num_items
273    unique_items = matrix_sample[item_key].unique()
274    matrix_sample = pd.merge(
275        matrix_sample,
276        pd.DataFrame(list(enumerate(unique_items)), columns=['item',item_key]),
277        how='left',
278        on=item_key
279    )
280    # remove any items above the threshold
281    matrix_sample = matrix_sample[matrix_sample['item'] < matrix_items]
282    # recalculate and resolve the indirection array
283    unique_items = dataset.resolve_item_ids(matrix_name, matrix_sample[item_key].unique().tolist())
284
285    # create sample by removing the extra columns that were added
286    matrix_sample = matrix_sample[['user', 'item'] + matrix_config.table.columns]
287
288    return matrix_sample, unique_users, unique_items