多进程从Cassandra拉取数据遇PicklingError问题求助

That pickling error you're hitting happens because Cassandra's ExecutionResult and ResultSet objects can't be serialized (pickled) to pass between processes. These objects hold live connection state and aren't designed to be transferred across process boundaries. Let's fix this by converting query results to a serializable format inside child processes before sending them back to the main process.

Step 1: Define the Missing pandas_factory #

You referenced pandas_factory but didn't include its definition. This function converts Cassandra ResultSets directly to Pandas DataFrames (which are pickle-friendly):

def pandas_factory(colnames, rows):
    return pd.DataFrame(rows, columns=colnames)

Step 2: Modified Cassandra DB Class #

We'll adjust the class to extract serializable data from execution results before returning them from child processes:

from cassandra.cluster import Cluster
import pandas as pd
import sys
import os
from multiprocessing import Pool
from cassandra.concurrent import execute_concurrent_with_args

# Add the missing pandas_factory here
def pandas_factory(colnames, rows):
    return pd.DataFrame(rows, columns=colnames)

ip_address = '127.0.0.1'

class cassandra_db(object):
    concurrency = 2 # Matches default execute_concurrent_with_args limit
    def __init__(self, process_count=None):
        self.pool = Pool(processes=process_count, initializer=self._setup)

    @classmethod
    def _setup(cls):
        # Each child process gets its own Cassandra connection (correct approach!)
        cls.session = Cluster([ip_address]).connect(keyspace='test')
        cls.session.row_factory = pandas_factory
        cls.prepared = cls.session.prepare('SELECT * FROM tr_test WHERE key=?')

    def close_pool(self):
        self.pool.close()
        self.pool.join()

    def get_results(self, params):
        try:
            xrange
        except NameError:
            xrange = range
        params = list(params)
        print("-----> ", params)
        print("-----+>", self.concurrency)
        
        # Split params into chunks for multiprocessing
        chunks = [params[n:n + self.concurrency] for n in xrange(0, len(params), self.concurrency)]
        # Collect results from all child processes
        all_chunk_results = self.pool.map(_multiprocess_get, chunks)
        
        # Flatten and concatenate all DataFrames into one
        df_list = [df for chunk in all_chunk_results for df in chunk]
        final_df = pd.concat(df_list, ignore_index=True)
        return final_df

    @classmethod
    def _results_from_concurrent(cls, params):
        execution_results = execute_concurrent_with_args(cls.session, cls.prepared, params)
        data_frames = []
        for res in execution_results:
            if res.success:
                # Extract the DataFrame from the execution result
                data_frames.append(res.result_or_exc)
            else:
                # Add custom error handling here if needed
                print(f"Query failed: {res.result_or_exc}")
        return data_frames

def _multiprocess_get(params):
    return cassandra_db._results_from_concurrent(params)

Step 3: Updated Calling Code #

Your original code tried to call df.head() on a method that returned nothing. The modified get_results now returns a concatenated DataFrame, so we can use it correctly:

import os
import pandas as pd
import sys
relative_path='/home/anji'
sys.path.append(os.path.join(relative_path ,'commons','Database Operations'))
from cassandra_db import cassandra_db

processes = 2
con_db = cassandra_db(processes)
keys=[(1,),(2,)]
df = con_db.get_results(keys)
print("Result", df.head())
# Always clean up the multiprocessing pool when done
con_db.close_pool()

Key Fixes Explained #

1. Serializable Data: We convert ExecutionResult/ResultSet to Pandas DataFrames inside child processes, so only pickle-friendly data is sent back to the main process.
2. Result Aggregation: The get_results method now collects all chunk results, flattens them, and combines them into a single DataFrame for easy use.
3. Basic Error Handling: Added checks for failed queries—you can expand this to log errors or retry failed requests as needed.
4. Pool Cleanup: Added con_db.close_pool() to properly shut down the multiprocessing pool after use.

Additional Notes #

• Keep your process_count reasonable (start with 2-4) to avoid overwhelming Cassandra or your system's resources.
• If you don't need Pandas, you can convert ResultSets to lists of dictionaries instead (e.g., res.result_or_exc.to_dict('records')), which are also serializable.

内容的提问来源于stack exchange，提问作者Anji