Provide more control over the Parquet schema when writing to Parquet

[adamreeve]

Description

Sometimes we want to generate Parquet files to be read by tools that expect a specific Parquet schema, and at the moment this requires exporting to Arrow and then using PyArrow to cast to the desired format first. It would be great if we could write Parquet directly from Polars.

In particular, it would be useful to be able to specify the nullability of columns, and also of any types nested within other columns like array elements. Currently Polars assumes everything is nullable. This causes problems when reading Parquet files in C#/.NET with ParquetSharp or Parquet.Net for example, as non-nullable integers will be read as the int type, but nullable integers need to be read as Nullable<int>.

One potential solution would be to add a new arrow_schema parameter to write_parquet and sink_parquet, which can be None or an instance of pyarrow.Schema. There are some edge cases where the Parquet schema can't be fully constrained by the Arrow schema, eg. I think writing decimals with integer Parquet primitives rather than fixed length byte arrays is one thing that can't be defined in the Arrow schema, but PyArrow doesn't expose this level of customisation either so I think that's fine.

If write_parquet(use_pyarrow=True) is used, the implementation can simply cast the Arrow table to the desired schema before writing it. This won't unnecessarily copy any buffers and will check that any non-nullable columns don't contain any nulls. Handling the Rust ParquetWriter is a bit trickier. We could maybe enforce that the schema matches the result of schema_to_arrow_checked besides any nullability differences, and then validate that any non-nullable columns do not contain nulls when writing.

I'm happy to implement this, but want to check that this feature would be accepted first.

[choucavalier]

@adamreeve that would be great!

how are you currently casting the pa.Table data types? could you share some code?

in my case, i need to cast the (by-default) nullable data types of polars to nullable=False in pyarrow but it does not work yet (it's been implemented in arrow 2 weeks ago and has not been released yet)

i think it would be great if the schema was also an extra arg of .to_arrow, especially because .to_arrow is also used by methods like .write_delta (e.g. when merge a polars dataframe into a delta table, right now it can fail because of the difference in data types in-between the arrow table generated by polars and the data types of the already-existing delta table..)

[adamreeve]

Hi @choucavalier

I haven't made any progress on this, I discussed it elsewhere with the Polars team and there was a concern about adding a dependency on PyArrow to Polars for Parquet writing, so we weren't sure what the best approach would be to solve this problem.

in my case, i need to cast the (by-default) nullable data types of polars to nullable=False in pyarrow but it does not work yet (it's been implemented in arrow 2 weeks ago and has not been released yet)

It looks like that change is specifically about nested fields, eg. within a struct typed column. Fortunately I don't think this has been an issue for us, as casting works with non-nested and array typed columns, eg:

import polars as pl
import pyarrow as pa
import pyarrow.parquet as pq

df = pl.DataFrame({
    "x": [[0, 1, 2], [3, 4]]
})

schema = pa.schema([
    pa.field("x", type=pa.list_(pa.field("item", type=pa.int32(), nullable=False)), nullable=False)
])
table = df.to_arrow().cast(schema)
pq.write_table(table, "data.parquet")

[choucavalier]

Yes, indeed in my case I'm using struct typed columns, and the casting simply does not work when going from nullable (assumed by polars) to non-nullable struct fields.

I've currently circumvented the issue by converting to numpy and re-constructing the pa.Table myself