SQL Pipe

🚀 The feature

Allows for sourcing datasets from SQL queries. Should allow to substitute in different backends. eg Athena, Presto, Postgres. Should do smart batching to minimize number of query requests and maximize data returned per query.

Batching will probably be backend dependent. Eg Athena/Presto do not support LIMIT .. SKIP ... queries

Motivation, pitch

Data is often stored in large datalakes on which SQL queries can be run to transform and load data. At some point this data must be turned into tensors however loading of this takes quite a lot of care and each query can take multiples of seconds meaning it cannot be done

Alternatives

Could not find any existing implementations. I have written my own solutions that allows user to index on a query but it is very hacky.

Current solution looks something like this

from itertools import groupby
class SQLTAble():

    @cached
    def get_patient(self, id) -> pd.DataFrame:
        
        return self.get_patients([id])

    def get_patients(self, ids) -> pd.DataFrame:

        ids = '(' + ','.join(ids) + ')'

        query = f"""
        SELECT patientid, offset, parameterid, value from icu_data
        where patientid in {ids}
        order by patientid, offset, parameterid
        """

        return client.query(query)

    def prefill_cache(self, all_patients):

        for patients in batch(all_patients, batch_size):

            results = self.get_patients(patients)

            patient_iter, patient_results_iter = groupby(results, lambda x: x['patientid'])

            for patient, patient_results in zip(patient_iter, patient_results_iter):

                self.add_cache(patient, patient_results)

the @cached decorator check if the patient is already in an lmdb database and shortcuts having to perform the query if it is.

Naievely fetching patient by patient is far too slow but it's also not possible to fetch all patients at once as this would result in memory issues.

Not specifically related to the SQLPipe ... but for this specific use case that I am working on at the moment there would also be a need to transform the data from columnar format to timeseries format. This is also quite expensive at scale and might benefit from a Pipe function.

eg above example has (patientid, offset, parameterid, value), this should be turned into a tensor of shape (patient x seq_len x n_parameters) and the values filled from value. Basically the data is queried in sparse COO format and needs to be densified. (There may be an SQL way of doing this, if so would love to hear it. My SQL is admittedly weak)

Additional context

No response

No files, tests, or concrete entry points are named. Start by locating the repository's data-loading and Pipe interfaces, then determine how SQL backends, batching, caching, and tensor conversion should fit together; done requires an agreed scope and implementation plan for the SQL source feature.