Skip to content

Commit

Permalink
doc: demostrate retrieving query result in arrow format
Browse files Browse the repository at this point in the history
  • Loading branch information
fanyang01 committed Dec 6, 2024
1 parent d2b3863 commit cb6cb7d
Show file tree
Hide file tree
Showing 2 changed files with 35 additions and 0 deletions.
13 changes: 13 additions & 0 deletions compatibility/pg-pytools/pyarrow_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -48,3 +48,16 @@
df_from_pg = df_from_pg.astype({'id': 'int64', 'num': 'int64'})
# Compare the original DataFrame with the DataFrame from PostgreSQL
assert df.equals(df_from_pg), "DataFrames are not equal"

# Copy query result to a pandas DataFrame
arrow_data = io.BytesIO()
with cur.copy("COPY (SELECT id, num * num AS num FROM test.tb1) TO STDOUT (FORMAT arrow)") as copy:
for block in copy:
arrow_data.write(block)

with pa.ipc.open_stream(arrow_data.getvalue()) as reader:
df_from_pg = reader.read_pandas().astype({'id': 'int64', 'num': 'int64'})
df['num'] = df['num'] ** 2
df = df.drop('data', axis='columns')
# Compare the original DataFrame with the DataFrame from PostgreSQL
assert df.equals(df_from_pg), "DataFrames are not equal"
22 changes: 22 additions & 0 deletions docs/tutorial/pg-python-data-tools.md
Original file line number Diff line number Diff line change
Expand Up @@ -119,3 +119,25 @@ polars_df = pl.from_arrow(arrow_df)
```python
polars_df = pl.from_pandas(pandas_df)
```

## 4. Retrieving Query Results as DataFrames

You can also retrieve query results from MyDuck Server as DataFrames using Arrow format. Here is an example:

```python
import io
import polars as pl
import pyarrow as pa
import psycopg

# Copy query result to a Polars DataFrame
arrow_data = io.BytesIO()
with cur.copy("COPY (SELECT id, num * num AS num FROM test.tb1) TO STDOUT (FORMAT arrow)") as copy:
for block in copy:
arrow_data.write(block)

with pa.ipc.open_stream(arrow_data.getvalue()) as reader:
arrow_table = reader.read_all()
polars_df = pl.from_arrow(arrow_table)
print(polars_df)
```

0 comments on commit cb6cb7d

Please sign in to comment.