Depricate Arrow2, update Polars version and keep all functionaliy

[EricFecteau]

Issue & Roadblock

I have been trying to use ConnectorX to move data from PostgreSQL (or any other database) to Polars, in Rust. I have hit two specific roadblocks:

• To move data from Databases to Polars, you have to use the Arrow2, a deprecated crate
• If you use Arrow2 to move data from a Database to Polars, you get Polars 0.32, a fairly old version of Polars (August 2023) compared to the most recent version of Polars 0.45.

Since I need Polars version 0.45 (or near this version), I had to come up with another solution.

Attempted solution

One solution is to use Polars to convert the Arrow-rs ArrayData object (by iterating through the RecordBatch and the Columns) and build a Polars DataFrame (using the current version of Polars). This is a popular solution online for moving Arrow-rs to Polars. But unfortunately, the Arrow-rs support for Polars was removed in 0.44.

Solution found

Since both the arrow-rs and the polars_arrow implement exporting and importing Arrow Array through the FFI bindings via Arrow’s C Data Interface, it's possible to Zero-Copy Arrow data between them. Ultimately, I was able to build this with this code:

// Connect to PostgreSQL through the ConnectorX
let source_conn =
    SourceConn::try_from("postgresql://postgres:postgres@localhost:5432").unwrap();

// Prepare query
let queries = &[CXQuery::from("SELECT * FROM table")];

// ConnectorX query PostgreSQL and return arrow object
let arrow_obj = get_arrow(&source_conn, None, queries)
    .unwrap()
    .arrow()
    .unwrap();

// Get column names as Polars PlSmallStr
let names = arrow_obj[0]
    .schema()
    .fields()
    .iter()
    .map(|f| PlSmallStr::from(f.name()))
    .collect::<Vec<PlSmallStr>>();

// Ready LazyFrame vector for the chunks
let mut lf_vec = vec![];

// The received arrow is chunked (for parallel processing) by ConnectorX
for batch in arrow_obj.into_iter() {
    // Bach column vector
    let mut columns = Vec::with_capacity(batch.num_columns());

    // Arrow stores data by columns, therefore need to be Zero-copied by column
    for (i, col) in batch.columns().iter().enumerate() {
        // Convert to arrow_data::data::ArrayData (arrow-rs)
        let array = col.to_data();

        // Convert to ffi with arrow-rs
        let (out_array, out_schema) = arrow::ffi::to_ffi(&array).unwrap();

        // Import field from ffi with polars
        let field = unsafe {
            polars_arrow::ffi::import_field_from_c(transmute::<
                &arrow::ffi::FFI_ArrowSchema,
                &polars_arrow::ffi::ArrowSchema,
            >(&out_schema))
        }
        .unwrap();

        // Import data from ffi with polars
        let data = unsafe {
            polars_arrow::ffi::import_array_from_c(
                transmute::<arrow::ffi::FFI_ArrowArray, polars_arrow::ffi::ArrowArray>(
                    out_array,
                ),
                field.dtype().clone(),
            )
        }
        .unwrap();

        // Create Polars series from arrow column
        columns.push(Series::from_arrow(names[i].clone(), data).unwrap());
    }

    // Create DataFrame from the columns
    lf_vec.push(DataFrame::from_iter(columns).lazy());
}

// Concat the chunks
let union_args = UnionArgs::default();
let lf = concat(lf_vec, union_args).unwrap();

// Print table
println!("{}", lf.collect().unwrap())

Conclusion

From my understanding, the main reason the Arrow2 crate is supported is to move data to Polars in Rust exclusively. So my idea is that everything Arrow2 related can be dropped from this crate, and some interoperability with arrow-rs and polars_arrow can be built in as a "destination" using the FFI Arrow methods (and therefore .polars() can be built using that). This would significantly simplify the create as arrow2 would no longer need to be implemented as transport.

[wangxiaoying]

Hi @EricFecteau , thanks for opening the discussion and proposing the solution!

Indeed, arrow2 was introduced in connectorx for exporting data to polars in the beginning. Since arrow2 is replaced with polars_arrow for polars, we need to add a new destination for polars, which can leverage the existing arrow destination for minimizing the implementation effort.

The solution looks reasonable to me. I'm currently kind of busy with other works and may not be able to start working on it lately. Please feel free to open a PR for it and let me know if you have any questions during implementation! I would suggest to create a new destination polars in connectorx following the above and keep the arrow2 solution for now (for compatibility) and drop it later.

[EricFecteau]

Hello again @wangxiaoying ! Thank you for the reply and recommendation! I have implemented the solution on my fork (here in the update_polars branch for arrow2 and arrow_rs_polars branch for arrow-rs), but have two design questions and one technical question before doing a PR.

Design 1: Polars version

Since the reason I encountered this issue was that connectorx had a very old version of Polars (0.32) and was incompatible with the version I needed (0.45), should all Rust "destinations" be set up as workspace = true in Cargo.toml, to allow the user the specify the version they need? This is how arrow-rs and arrow2 are done. This is how I currently implemented it in the Crates.toml for Polars and polars-arrow.

Any concerns there - specially with minimal compatible version? Should I instead of `polars = { version = ">=0.43" } as the minimal working version (tested). Polars change frequently so the minimal working version (0.43) is fairly recent. I would not be surprised if a future version also breaks this in the near future.

Design 2: No changes to python version

If I understand correctly, this change has zero impact on the python binding, because Polars is not a destination for ConnectorX-python (only Arrow-rs, Arrow2 and Pandas). Instead, Polars uses the arrow2 return type directly when they use ConnectorX. Long term, if arrow2 gets deprecated from this crate (as arrow2 is generally deprecated), the ConnectorX-python would need to also deprecate the arrow2.

I think the Polars could already move to arrow-rs for this, since they use the FFI anyway. Maybe something I could explore more in the future.

Technical 1: Arrow2 - Polars test failures

When importing FFI with polars_arrow (exported with arrow2) instead of the previous method (not using FFI), the test_pg_pl_varchar_array test and the test_pg_pl_text_array test fail. They fail with the "An ArrowArray of type LargeUtf8 must have buffer 2." error (found here in polars-arrow). I am not familiar enough with Arrow memory formats to know why that happens (see technical details). The tests for arrow2 use PostgreSQL, but with arrow-rs, "text arrays", "varchar arrays" and "bool arrays" are not implemented at all for PostgreSQL.

Issue technical details

Here is way more info in case it can help solve the issue. I am more than happy to try implementing something with some guidance.

For arrow2, the VarcharArray and the TextArray both convert themselves to Utf8Array, passing through this "impl" to make it a LargeUtf8. Arrow2 successfully exports it to a Utf8Array.

When using import_array_from_c from polars_arrow, I then get an error. It tries to create the polars_arrow LargeUtf8 from the FFI, and when requesting a u8 buffer (index = 2) (and creating it), when the buffer pointer is fetched, it fails because index == array_n_buffers. This makes the error ("An ArrowArray of type LargeUtf8 must have buffer 2.") weird, because array.n_buffers does indeed have a buffer == 2.

Not sure if this is a arrow2 export issue, a arrow_polars import issue or a connectorx issue.

For arrow-rs, I can't test it within connectorx, because postgres_arrow.rs does not implement arrays.

Pull request resolution

Here is a decision that needs to be made before I make this pull request - what should be the way forward:

1. Disable these two tests (for arrow2 and don't implement them for arrow-rs) and keep the new FFI method for both.
2. Keep Polars 0.32 for arrow2 (and the previous code) and implement the FFI for arrow-rs (using a newer version of Polars) - I think this should be possible in the Cargo.toml by renaming the dependency to something like polars_old. Do not port the "text arrays", "varchar arrays" and "bool arrays" since they are not implement with PostgreSQL.
3. Figure out bug, fix bug with FFI-arrow2 and implement the arrays with PostgreSQL-arrow-rs and implement for FFI-arrow-rs.

I suspect that there is not a ton of users using the arrays, if they were never implemented for arrow-rs.

[EricFecteau]

Awesome @wangxiaoying! Thanks.

Yeah for #1 I could add test maybe with the dummy source? (would have to check if it implements arrays) because PostgreSQL does not implement arrays in arrow (it does for arrow2). I suspect it will work perfectly with FFI between arrow and polars-arrow no problem as you mentioned!

For #2 I can create polars_old instead! I could look into arrays and see if I could get rid of Arrow2! Maybe do a second PR for this.

I will do these changes to the PR this weekend.

[wangxiaoying]

Thank you!

Yeah for #1 I could add test maybe with the dummy source? (would have to check if it implements arrays) because PostgreSQL does not implement arrays in arrow (it does for arrow2).

Currently postgres has the array types already. After adding array types to arrow, you can just add the corresponding type mapping in this transport file. For example, if we want to support varchar array, we can add something like this line in arrow2. Doing this should enable converting array from postgres to arrow. Then we can use the current select test_textarray from test_types directly in the test.

[EricFecteau]

Thanks for the guidance! I implemented arrays for arrow and reverted to using polars_old = 0.32 for arrow2. I am interested in doing another PR to remove arrow2 from the crate, but I think it deserves it’s own PR. Now all the test (plus the tests I added) pass.

A few comments / questions:

• I implemented all the types as Vec<Option<Type>> for LargeList, as it seems to be necessary for arrow to make them all Option<Type>.
• I had to up the Arrow version (I used the lastest) because the one that was currently used as the workspace version did not have Feild::new_list_field(). This makes the minimum version of Arrow 50.0.0.
• Is there a reason why PostgreSQL makes all the int types i64 and float types f64? This is true for Arrow and Arrow2 (and arrow2 arrays). This makes it so that floats like -1e-37 become -9.99999991097579e-38 and makes tests hard. I did the same thing for arrays, but could not write proper tests for them.
• Is there a reason why bool arrays in Postgres can’t be nullable? It seems intentional.

[wangxiaoying]

I implemented all the types as Vec<Option> for LargeList, as it seems to be necessary for arrow to make them all Option.

I think it makes sense!

I had to up the Arrow version

I made the arrow version from 54 to 53 since it is the latest version that our dependency datafusion supports. Otherwise it would result in compile error.

Is there a reason why PostgreSQL makes all the int types i64 and float types f64? This is true for Arrow and Arrow2 (and arrow2 arrays).

It's more like a technical debt in the previous implementation. I think we can update it in another PR.

Is there a reason why bool arrays in Postgres can’t be nullable?

Not really. I have updated the test to remove this filter.

I am interested in doing another PR to remove arrow2 from the crate, but I think it deserves it’s own PR.

Sure, I think we can definitely have a separate PR for doing this. I have invited you to be a collaborator of connectorx. Hopefully it will make it easier for you to contribute!

[EricFecteau]

Amazing, thank you very much for all the help!

[deanm0000]

As an outside observer who is usually a big polars fan, I'd suggest you just leave it to polars to make their arrow interoperable with the official one. The other libraries they use (fastexcel, delta, iceberg, etc) do their arrow to arrow via pyarrow.

[EricFecteau]

Unfortunately, Polars has decided, as of 0.44 to no longer support interoperability with the official Arrow crate: pola-rs/polars#19312

[deanm0000]

I knew they took out their Arrow2Arrow trait (or whatever it was called). I'm really only intending to say that if the error referenced above An ArrowArray of type LargeUtf8 must have buffer 2. is a big hold up then don't feel the need to support the rust interop yourself. All the other polars dependencies (AFAIK) do interop in python. To be honest, in rereading the thread, I'm a little confused about whether the problem is just with the old arrow2 crate or if it's the current polars-arrow, if it's only the deprecated one and cffi works for the current version then I'm not trying to talk you out of how you're doing it. I think it's great.