csv_batcher.csv_pooler

 1from enum import StrEnum, auto
 2from multiprocessing import Pool
 3from typing import Callable
 4
 5import pandas as pd
 6
 7from csv_batcher.csv_splitter import CSVSplitter
 8from csv_batcher.utils.logger import logging
 9
10class CallbackWith(StrEnum):
11    """
12    CallbackWith Enum, used to control what is passed to callback function
13    """
14    # Pass the chunked CSV filename:
15    CSV_FILENAME = auto()
16
17    # Pass a Dataframe created from the chunked CSV file:
18    DATAFRAME = auto()
19
20    # First creates a DataFrame, then calls `apply()` with the callback
21    #   function. This causes the callback function to be called with
22    #   each row as a `pd.Series` object:
23    DATAFRAME_ROW = auto()
24class CSVPooler:
25    def __init__(
26        self,
27        csv_filename: str,
28        process_fn: Callable,
29        callback_with: CallbackWith = CallbackWith.CSV_FILENAME,
30        pool_size: int = 5,
31        chunk_lines: int = 10000,
32    ):
33        """
34        Construct `Pooler` with given `csv_filename`, `process_fn`, `as_dataframe`, `pool_size', 'chunk_size'
35
36        Args:
37            csv_filename (str): Name of CSV file
38            process_fn (Callable): A function that accepts a single argument
39                By default, this is the path to a chunked CSV file
40                If `as_dataframe` is True, then the argument sent is a dataframe of the chunked CSV
41            callback_with (CallbackWith): Controls what is sent to callback function.
42                @see CallbackWith enumeration for details
43                Defaults to CallbackWith.CSV_FILENAME.
44            as_dataframe_rows (bool): When True, a dataframe is created as with as_dataframe.
45                that is sent to `process_fn` instead.  Defaults to False.
46            pool_size (int, optional): Number of workers to uses. Defaults to 8.
47            chunk_lines (int, optional): Target row count for each chunked CSV. Last chunk may
48                have fewer rows.  Defaults to 10000.
49        """
50        self.csv_filename = csv_filename
51        self.process_fn = process_fn
52        self.callback_with = callback_with
53        self.pool_size = pool_size
54        self.chunk_lines = chunk_lines
55
56    def process(self):
57        """
58        Processes `self.csv_filename` by using `CSVSplitter` to split it
59            into multiple temporary files defined by `self.chunk_lines`.
60        Use `multiprocessing.Pool` to use multiple process workers to process
61        the group of CSVs.
62        """
63        processed_count = 0
64        csv_splitter = CSVSplitter(self.csv_filename, self.chunk_lines)
65        try:
66            csv_file_cnt = len(csv_splitter.csv_files())
67            logging.info(f"Pooling against {csv_file_cnt} files")
68            with Pool(self.pool_size) as p:
69                for result, count in p.imap(self._process_csv, csv_splitter.csv_files()):
70                    yield(result)
71                    processed_count += count
72        finally:
73            csv_splitter.cleanup()
74
75        logging.info(f"Processed {processed_count} rows from {csv_file_cnt} CSV Files")
76
77    def _process_csv(self, csv_chunk_filename):
78        if self.callback_with == CallbackWith.CSV_FILENAME:
79            result = self.process_fn(csv_chunk_filename)
80            with open(csv_chunk_filename) as f:
81                # Get total lines and subtract for header:
82                count = sum(1 for line in f) - 1
83        elif self.callback_with == CallbackWith.DATAFRAME:
84            df = pd.read_csv(csv_chunk_filename, skipinitialspace=True, index_col=None)
85            count = df.shape[0]
86            result = self.process_fn(df)
87        elif self.callback_with == CallbackWith.DATAFRAME_ROW:
88            df = pd.read_csv(csv_chunk_filename, skipinitialspace=True, index_col=None)
89            count = df.shape[0]
90            result = df.apply(self.process_fn, axis=1)
91
92        return result, count