Return Python dataclass with multiple objects from query

Viewed 49

I've created a Python function that returns metadata based on a data class I've instantiated. Here is my code:

from dataclasses import dataclass
from pyspark.sql import SparkSession, DataFrame

@dataclass
class Metadata:
  col1: str
  col2: str
  col3: str

def retrieve_metadata(
    spark: SparkSession, batch_name: str, table_name: str,
) -> Metadata:

    jdbc_url, connection_properties = fetch_db_config()
    query = f"""(SELECT col1, col2 , col3 FROM table
             ) AS q
             """
    # return query
    df_columns = query_jdbc_database(
        spark, query, jdbc_url, connection_properties).collect()

    return Metadata(
        df_columns.col1,
        df_columns.col2,
        df_columns.col3
    )

metadata = retrieve_metadata(spark, x, y)
print(metadata) 

It uses some functions to query a jdbc database, but those are not worth mentioning here. Now, the problem that I have is that my metadata object only contains one row. I expect it to contain as many rows as the query returns. I'm guessing its a really small change, but I can't wrap my head around it.

Edit: Maybe a better question is, is there an easy way to load my spark dataframe into my dataclass?

Any suggestions?

Thanks in advance!!

0 Answers
Related