Re: [D] Performant streaming of Parquet to Postgres [arrow-rs]

["adriangb (via GitHub)" <gi...@apache.org>]

GitHub user adriangb closed a discussion: Performant streaming of Parquet to Postgres

I couldn't find anything out there so I'm looking into writing something to move data from parquet files (possibly stored in an object store or HTTP) into Postgres as fast as possible. For my use case I have a couple of requirements:
- Able to work with larger than memory files without blowing up memory or taking forever
- Can be hooked into from Python

My initial attempt was using Polars but unfortunately it is not able to read Parquet files in batches efficiently and converting from Arrow -> Python types (via Polars) and then Python types -> Postgres (via asyncpg) is slow.

So naturally I'm looking at doing this in Rust to speed things up.

Here's my plan, I want to see if it seems viable. I'll have a Rust library that goes from `RecordBatch` to Postgres' binary format. It'll look something like this:

```rust
pub struct PostgresBinaryCopyEncoder {}

impl PostgresBinaryCopyEncoder {
    // Initialize with arrow schema for these RecordBatches
    pub fn new(types: Schema) -> BinaryCopyInWriter {}

    // If this is the first call inlcude headers in the returned data
    // Gather all of the columns via RecordBatch.columns
    // Zip the columns and iterate over them row-wise
    // For each value use the column's type to downcast into a Rust type
    // Use the ToSql trait which is already implemented for Rust values to encode into pg's binary format
    pub fn encode(&mut self, values: &RecordBatch) -> Result<&[u8], Error> {}

   // Get the footers
   pub fn finalize(&mut self) -> Result<&[u8, Error]> {}
```

And then a Python side that manages the IO so that it can decide to stream the bytes into Postgres directly, write to a file somewhere, read the parquet file from disk, in threads for an async environment, etc. That would look something like this:

```python
from thislib,_rust_stuff import PosgresBinaryCopyEncoder as PosgresBinaryCopyEncoderRust

class PosgresBinaryCopyEncoder:
    def __init__(self, stream: RecordBatchStreamReader) -> None:  ...

   def stream(self) -> Iterable[bytes]:
      schema = self._stream.schema
      encoder = RecordBatchStreamReader(schema)
      while True:
          try:
              batch = self._stream.read_next_batch()
              yield encoder.encoder(batch)
          except StopIteration:
              yield encoder.finalize()       
```

I've already written something to encode Rust types into Postgres types based on [rust-postgres-binary-copy](https://github.com/sfackler/rust-postgres-binary-copy). Where I'm a bit confused is on how to iterate over a RecordBatch and get Rust types out of it. Alternatively I guess we could iterate over a RecordBatch and get arrow data out of it but then I'd need to re-implement the `postgres_types::ToSql` trait for all arrow types. And I expect an extra hop from Arrow -> Rust native type will be quite cheap.

GitHub link: https://github.com/apache/arrow-rs/discussions/3641

----
This is an automatically sent email for github@arrow.apache.org.
To unsubscribe, please send an email to: github-unsubscribe@arrow.apache.org