Convert `req_ch` and `req_athena` to use structs/adbc

[josevalim]

Recently we moved most drivers in kino_db that already relied on ADBC to use only ADBC, and no longer depend on Explorer. The exceptions are req_ch and req_athena, which still use Explorer. We have three options:

1. Continue using Explorer
2. Use ADBC + DuckDB with read_parquet (we probably start a database and connection on the fly)
3. Introduce a %ReqCh.Result{} kind of structure that implements Table.Reader (so we probably decode JSON and then parse it)

[josevalim]

Some additional thoughts on this. Unfortunately, using ADBC + DuckDB comes with two issues:

First, you can't DuckDB's read_parquet from an in-memory blob. So we would need to write the file to disk and then parse it. I guess that's at least better than loading the whole file into memory anyway.

However, the biggest issue is that, once we get a result set from ADBC, we really can't do much with it. Adbc.Connection.query! returns an Adbc.Result, and we can use Elixir functions to manipulate it, but I can only feed it back into DuckDB by sending the whole data back to the database. This means we are doing Parquet -> DuckDB -> Elixir -> DuckDB while the best option would be to keep the Parquet data inside DuckDB. Here is what we would need to do in practice.

In the setup cell:

Adbc.download_driver!(:duckdb)
db = Kino.start_child!({Adbc.Database, driver: :duckdb})
conn = Kino.start_child!({Adbc.Connection, database: db})

Then when executing code:

data = ReqCh.query!(...).body
File.write!(tmp_file, data)

# Process data and keep in DuckDB
Adbc.Connection.query!(conn, "CREATE TEMPORARY TABLE parquet_data AS (SELECT * FROM read_parquet($1)", [tmp_file])

# Do another query to display results
Adbc.Connection.query!(conn, "SELECT * FROM parquet_data")

While there is a huge amount of boilerplate here, it is worth noting those issues are not really specific to DuckDB. If I get data out of PostgreSQL and then I want to use it in a query again, I have to send it back, unless I stored them tables. So the important bit is that, we want to keep data in DuckDB, and then when displaying it, we get what we need. In a way, that's how something like Polars/Explorer works too. Data is kept in Rust and we convert it to Elixir when needed.

However, the other issue is that, once we create temporary tables, the data will be there in the table forever, until you reset the connection. They are effectively side-effects that are not visibile in Elixir and Livebook can't change track them. One could think of using random table names and pass them around:

table_name = :cryto.strong_rand_bytes |> Base.encode32()

# Process data and keep in DuckDB
Adbc.Connection.query!(conn, "CREATE TEMPORARY #{table} parquet_data AS (SELECT * FROM read_parquet($1)", [tmp_file])

# Do another query to display results
Adbc.Connection.query!(conn, "SELECT * FROM #{table}")

While the above is enough to add change tracking to Livebook, it now comes with memory leaks, as each evaluation creates a new temporary table that is never removed (until the connection is reset). Luckily, we can solve this problem by returning a data structure that removes the temporary table when GCed.

View and tables APIs

We could introduce two new functions to Adbc.Connection, called view and table, that works like this:

%Adbc.Relation{type: :table, name: "random_value"} = Adbc.Connection.table!(conn, "SELECT * FROM read_parquet($1)", [tmp_file])

Internally this emits a "CREATE TEMPORARY TABLE" which we know the reference of. When Adbc.Relation is garbage collected, we emit a command to delete it.

Note this is a general pattern, no longer related to req_ch/req_athena. This pattern can be used everywhere in notebooks to write composable queries, where each query we run stores its results in the database, which we then read when necessary (allowing us, for example, to implement a custom Kino.DataTable with LIMIT+OFFSET and exports). I believe https://hex.tech/ does something similar.

We can also add a view! command that does the same but creating a view (so in this case it effectively works like a lazy dataframe):

%Adbc.Relation{type: :view, name: "random_value"} = Adbc.Connection.table!(conn, "SELECT * FROM #{relation.name} WHERE ...", [tmp_file])

And we could have a collect function that converts it into a result (effectively it is "SELECT * FROM #{table}").

This is not perfect though as it has the following issues:

1. I can see the above being useful in notebooks, but I can't see the benefit in using it in real code. In real code, you'd rather compose queries or use CTEs, instead of going the GCed tables route. So this could be a notebook only abstraction, but that hurts portability. Unless we have a mechanism to convert a series of relations into a long CTE?

2. This pattern effectively requires folks to use string interpolation in queries to interpolate the table names, which I am not sure we want to encourage (we could perhaps solve it with custom sigils, but I am not sure if they would play well with Kino.DB)

3. This pattern of tracking temporary tables could also be used in Adbc.Connection.bulk_ingest

4. While we are adding new functions to Adbc.Connection, the behaviour is database specific, as we need to know the exact CREATE VIEW and CREATE TABLE statements to emit, and currently Adbc.Connection knows nothing about query building, this may make it hard to generalize)

[billylanchantin]

There appears to be a relevant option:

ADBC_INGEST_OPTION_TEMPORARY

Example use in docs:

https://arrow.apache.org/adbc/current/python/recipe/postgresql.html#create-append-to-a-temporary-table

From those docs:

After closing the connection, the temporary table is implicitly dropped. If we reconnect, the table won’t exist

So some of what you're outlining may be supported without Adbc.Connection needing to be aware of the need to GC upon closure. But I don't think it addresses the main problem on its own. I'll continue to think on if there's a way to leverage the existing native support for temporary tables for what Livebook needs.

[josevalim]

Thanks! We do support the ingest option on Adbc.Connection.bulk_ingest on main (added yesterday) but it only deals with adding external data to DuckDB, not about transforming existing data. :)

[billylanchantin]

Unsurprisingly, you're way ahead of me! What I was envisioning in fact already exists and is documented:

# Efficiently insert from a query (within query_pointer callback)
# This is most useful for transferring across databases.
# Within the same database, you most likely have custom SQL commands,
# such as COPY, CREATE TEMPORARY TABLE, etc.
Adbc.Connection.query_pointer(source_conn, "SELECT * FROM source_table", fn stream ->
  Adbc.Connection.bulk_insert(dest_conn, stream, table: "dest_table")
end)

---

Regarding the table/view API, I wonder if Adbc.Connection or even Adbc is the right place? Despite it being a useful general pattern, I'm not sure it has anything to do with Adbc specifically. Maybe a better home is KinoDB? That way the DB-specific nature of CREATE ... parts isn't so much of an issue.

Running with this, one potential source of inspiration is Polars. They have immutable transformation semantics which pairs very nicely with Elixir's own immutability. It's what makes DataFrame operations feel so natural:

df1 = DF.new(...)           # These variables don't share
df2 = df1 |> DF.mutate(...) # memory by design.

So perhaps what we want is an API whose goal is to enable those same semantics on an arbitrary DB(s). This is effectively the same mechanics you've already proposed, but with a slightly different purpose. Something like (names are intentionally terrible):

remote_db = Kino.start_child!({Adbc.Database, driver: :postgresql})
local_db = Kino.start_child!({Adbc.Database, driver: :duckdb})

result1 =
  KinoDB.NewAPI.temp_table_from("select * from postgres_table",
    table_name: "result1",
    # Since source_db != target_db, the API uses
    # the new query_pointer/bulk_insert approach.
    source_db: remote_db,
    target_db: local_db
  )

result2 =
  KinoDB.NewAPI.temp_table_from("select * from result1 where ...",
    table_name: "result2",
    # Since source_db == target_db, the query call
    # instead gets wrapped in `CREATE TEMPORARY TABLE`
    db: local_db
  )

This way, result can be whatever it needs to be to facilitate the needs of Kino (reference that gets pagination queries for display, etc.) without feeling out of place.

We could also use a macro to remove the need to specify the table name:

# Macro
result = KinoDB.NewAPI.temp_table_from_macro("...", db: db)
# Rewritten as
result = KinoDB.NewAPI.temp_table_from_function("...", db: db, table_name: "result")

I think this approach will facilitate a Livebook experience that feels like what you get from Explorer out of the box but without being tied to Polars or even Arrow.

EDIT: Thinking more, KinoDB might be a weird dep for ReqCh/ReqAthena? The same NewAPI could live in Adbc too, though the CREATE ... commands will still feel a little out of place.