Alternative schema specification API using variadic generics

[nielsneerhoff]

Hi there!

Kudos! I have just started working with Spark in Python, and this library addresses an issue I've been struggling with from the beginning (especially coming from Java).

Python 3.11.4 was just released, and it includes some interesting new features, among which are variadic generics. In short, these could theoretically allow us to define a DataSet as:

dataset = DataSet['fruit': str, 'sales': int, 'average_price': float]

instead of:

class FruitSales(Schema):
  fruit: str
  sales: int
  average_price: float

dataset = DataSet[FruitSales]

(or equivalently, using PySpark types for the attribute type definitions).

What are your thoughts on adding support for this new type of schema specification to this library? Perhaps I can assist in implementing it.

There are several benefits of using the first schema definition approach:

• It's a bit shorter than specifying the schema line by line
• One can write: dataset = DataSet['fruit': str, ...] if one does not know all attributes of the schema
• It's more consistent with Spark's (currently experimental) type hints API for Pandas

I'm quite new to open-source contributing, please feel free to give feedback on this post and the idea :)

[nanne-aben]

Hi Niels!

Kudos! I have just started working with Spark in Python, and this library addresses an issue I've been struggling with from the beginning (especially coming from Java).

Thank you! Happy to hear that!

dataset = DataSet['fruit': str, 'sales': int, 'average_price': float]

Could you elaborate on when you'd use this form?

When I define my schema classes, I usually refer to them many times, for example:

from pyspark.sql.types import IntegerType
from typedspark import Column, Schema, DataSet, transform_to_schema

class Person(Schema):
    name: Column[StringType]
    age: Column[IntegerType]
    job_id: Column[IntegerType]

class Job(Schema):
    id: Column[IntegerType]
    base_salary: Column[IntegerType]

class PersonWithJob(Person, Job):
    effective_salary: Column[IntegerType]

def foo(persons: DataSet[Person], jobs: DataSet[Job]) -> DataSet[PersonWithJob]:
    return transform_to_schema(
        persons.join(jobs, Person.job_id == Job.id),
        PersonWithJob,
        {
            PersonWithJob.effective_salary: Job.base_salary + Person.age * 100,
        }
    )

With the variadic generics, we can't do this, right? Is there a usecase I'm missing?

[nielsneerhoff]

Hey @nanne-aben , thanks a lot for your response.

I see that the shorter form using variadic generics is less practical when you want to use (multiple) inheritance(s). To still facilitate it using variadic generics, we could perhaps think of something that would allow the DataSet's TypeVar to be both Schema(s) and a variadic generic:

DataSet[Person, Schema]
# or, equivalently
DataSet[name: Column[StringType], ..., base_salary: Column[IntegerType]]

But I think it is less explicit and overly complex.

---

In regards to your question on when I would use the form

dataset = DataSet['fruit': str, 'sales': int, 'average_price': float]

it would mostly be as an in-line type hint. For instance:

def foo(df_with_certain_col: DataSet['fruit': str, ...]) -> DataSet['fruit': str]:
   return df_with_certain_col.select('fruit')

Or, even better:

from constants import FRUIT

def foo(df_with_certain_col: DataSet[**{FRUIT: str, ...}]) -> DataSet[**{FRUIT: str}]:
   return df_with_certain_col.select(FRUIT)

The added benefit is that the column name is not hard-coded. The FRUIT can easily be reused and refactored. For instance in the tests:

df: DataSet[**{FRUIT: str, ...}] = ...  # create some spark DF
assert set(foo(df).collect()) == {Row(**{FRUIT: ...})}

(or, define DfWithFruit = DataSet[**{FRUIT: str, ...}] somewhere and reuse it throughout the above)

Anyway, I agree it would not fit the intention of this library very well. I will close the discussion.

Thanks again and happy coding!

[nanne-aben]

Thanks for elaborating! I agree that the proposed syntax doesn't fit very well :) You do bring up an interesting use-case.

from constants import FRUIT

def foo(df_with_certain_col: DataSet[**{FRUIT: str, ...}]) -> DataSet[**{FRUIT: str}]:
   return df_with_certain_col.select(FRUIT)

So this would be the case where you want to define some of the columns in a DataSet, but not all of them, right? We could actually support something like this using typing.Protocol classes. I'd implement it as something like:

class Fruit(Schema)
    fruit: Column[StringType]

def foo(df_with_certain_col: PartialDataSet[Fruit]) -> DataSet[Fruit]:
   return transform_to_schema(df, Fruit)

I think that would be interesting to include. What do you think about the above syntax?

---

df: DataSet[**{FRUIT: str, ...}] = ...  # create some spark DF
assert set(foo(df).collect()) == {Row(**{FRUIT: ...})}

Just as a quick note: we have support for using schemas to make unit testing easier, check out the example unit test here.

---

def foo(df_with_certain_col: DataSet['fruit': str, ...]) -> DataSet['fruit': str]:
   return df_with_certain_col.select('fruit')

I think this one will be hard to support for linters like mypy or pyright. The return type of df_with_certain_col.select('fruit') will be a DataFrame, or maybe a DataSet, but I don't think we can make it such that it expects a DataSet['fruit': str]. So you'd have to do a casting step:

def foo(df_with_certain_col: DataSet['fruit': str, ...]) -> DataSet['fruit': str]:
   return DataSet['fruit': str](df_with_certain_col.select('fruit'))

And that would make it a lot more verbose.

---

The added benefit is that the column name is not hard-coded. The FRUIT can easily be reused and refactored.

One of the main motivations for typedspark is actually easy refactoring. Consider the following schema again:

class Fruit(Schema):
    fruit: Column[StringType]

We could refactor Fruit.fruit to Fruit.type and your IDE will update it everywhere accordingly.

---

Thanks again and happy coding!

Thank you! Happy coding to you as well!

[nanne-aben]

Edited my post above for tone, my initial response was a bit dry :)