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!!