1. import pandas as pd
  2. import numpy as np
  4. def collect_metrics(data_generator, metrics=['size', 'dtype_counts', 'null_counts', 'unique_counts', 'min_max']):
  5. """
  6. Collects metrics of a dataset iteratively to minimize memory usage.
  8. Args:
  9. data_generator: An object that yields batches of data (e.g., a custom iterator).
  10. metrics: A list of metrics to calculate.
  12. Returns:
  13. A dictionary containing the collected metrics.
  14. """
  16. total_size = 0
  17. dtype_counts = {}
  18. null_counts = {}
  19. unique_counts = {}
  20. data_list = [] # Store data batches
  22. for batch in data_generator:
  23. if isinstance(batch, pd.DataFrame):
  24. data_list.append(batch)
  25. elif isinstance(batch, np.ndarray):
  26. data_list.append(batch)
  27. else:
  28. raise TypeError("Batch must be a pandas DataFrame or numpy array.")
  30. if not data_list:
  31. return {
  32. 'size': 0,
  33. 'dtype_counts': {},
  34. 'null_counts': {},
  35. 'unique_counts': {},
  36. 'min_max': {}
  37. }
  39. combined_data = pd.concat(data_list, ignore_index=True)
  40. total_size = combined_data.size
  42. # Calculate dtype counts
  43. for col in combined_data.columns:
  44. dtype_counts[col] = combined_data[col].dtype
  46. # Calculate null counts
  47. for col in combined_data.columns:
  48. null_counts[col] = combined_data[col].isnull().sum()
  50. # Calculate unique counts
  51. for col in combined_data.columns:
  52. unique_counts[col] = combined_data[col].nunique()
  54. # Calculate min/max for numeric columns
  55. for col in combined_data.select_dtypes(include=np.number).columns:
  56. min_val = combined_data[col].min()
  57. max_val = combined_data[col].max()
  58. data_list_col = []
  59. for batch in data_list:
  60. if isinstance(batch, pd.DataFrame):
  61. data_list_col.append(batch[col])
  62. elif isinstance(batch, np.ndarray):
  63. data_list_col.append(batch[:,col])
  65. min_val = np.min(np.concatenate(data_list_col))
  66. max_val = np.max(np.concatenate(data_list_col))
  68. data_list.remove(batch)
  71. min_val = combined_data[col].min()
  72. max_val = combined_data[col].max()
  74. data_list.append(combined_data[col])
  78. return {
  79. 'size': total_size,
  80. 'dtype_counts': dtype_counts,
  81. 'null_counts': null_counts,
  82. 'unique_counts': unique_counts,
  83. 'min_max': {col: (min_val, max_val) for col in combined_data.select_dtypes(include=np.number).columns}
  84. }
  86. if __name__ == '__main__':
  87. # Example Usage:
  88. def dummy_data_generator(num_batches=3):
  89. for i in range(num_batches):
  90. data = {'col1': np.random.rand(100), 'col2': np.random.randint(0, 10, 50), 'col3': ['a', 'b'] * 25}
  91. yield pd.DataFrame(data)
  93. metrics = collect_metrics(dummy_data_generator())
  94. print(metrics)

Add your comment