# /// script # requires-python = ">=3.11" # dependencies = [ # "marimo", # "duckdb==1.3.2", # "pyarrow==19.0.1", # "polars[pyarrow]==1.25.2", # "pandas==2.2.3", # "sqlglot==27.0.0", # "psutil==7.0.0", # "altair", # ] # /// import marimo __generated_with = "0.14.12" app = marimo.App(width="medium") @app.cell(hide_code=True) def _(mo): mo.md( r""" # Working with Apache Arrow *By [Thomas Liang](https://github.com/thliang01)* # """ ) return @app.cell(hide_code=True) def _(mo): mo.md( r""" [Apache Arrow](https://arrow.apache.org/) is a multi-language toolbox for building high performance applications that process and transport large data sets. It is designed to both improve the performance of analytical algorithms and the efficiency of moving data from one system or programming language to another. A critical component of Apache Arrow is its in-memory columnar format, a standardized, language-agnostic specification for representing structured, table-like datasets in-memory. This data format has a rich data type system (included nested and user-defined data types) designed to support the needs of analytic database systems, data frame libraries, and more. DuckDB has native support for Apache Arrow, which is an in-memory columnar data format. This allows for efficient data transfer between DuckDB and other Arrow-compatible systems, such as Polars and Pandas (via PyArrow). In this notebook, we'll explore how to: - Create an Arrow table from a DuckDB query. - Load an Arrow table into DuckDB. - Convert between DuckDB, Arrow, and Polars/Pandas DataFrames. - Combining data from multiple sources - Performance benefits """ ) return @app.cell def _(mo): mo.sql( """ CREATE TABLE IF NOT EXISTS users ( id INTEGER, name VARCHAR, age INTEGER, city VARCHAR ); INSERT INTO users VALUES (1, 'Alice', 30, 'New York'), (2, 'Bob', 24, 'London'), (3, 'Charlie', 35, 'Paris'), (4, 'David', 29, 'New York'), (5, 'Eve', 40, 'London'); """ ) return (users,) @app.cell(hide_code=True) def _(mo): mo.md( r""" ## 1. Creating an Arrow Table from a DuckDB Query You can directly fetch the results of a DuckDB query as an Apache Arrow table using the `.arrow()` method on the query result. """ ) return @app.cell def _(mo, users): users_arrow_table = mo.sql( # type: ignore """ SELECT * FROM users WHERE age > 30; """ ).to_arrow() return (users_arrow_table,) @app.cell(hide_code=True) def _(mo): mo.md(r"""The `.arrow()` method returns a `pyarrow.Table` object. We can inspect its schema:""") return @app.cell def _(users_arrow_table): users_arrow_table.schema return @app.cell(hide_code=True) def _(mo): mo.md( r""" ## 2. Loading an Arrow Table into DuckDB You can also register an existing Arrow table (or a Polars/Pandas DataFrame, which uses Arrow under the hood) directly with DuckDB. This allows you to query the in-memory data without any copying, which is highly efficient. """ ) return @app.cell def _(pa): # Create an Arrow table in Python new_data = pa.table({ 'id': [6, 7], 'name': ['Fiona', 'George'], 'age': [22, 45], 'city': ['Berlin', 'Tokyo'] }) return (new_data,) @app.cell(hide_code=True) def _(mo): mo.md(r"""Now, we can query this Arrow table `new_data` directly from SQL by embedding it in the query.""") return @app.cell def _(mo, new_data): mo.sql( f""" SELECT name, age, city FROM new_data WHERE age > 30; """ ) return @app.cell(hide_code=True) def _(mo): mo.md( r""" ## 3. Convert between DuckDB, Arrow, and Polars/Pandas DataFrames. The real power of DuckDB's Arrow integration comes from its seamless interoperability with data frame libraries like Polars and Pandas. Because they all share the Arrow in-memory format, conversions are often zero-copy and extremely fast. """ ) return @app.cell(hide_code=True) def _(mo): mo.md(r"""### From DuckDB to Polars/Pandas""") return @app.cell def _(pl, users_arrow_table): # Convert the Arrow table to a Polars DataFrame users_polars_df = pl.from_arrow(users_arrow_table) users_polars_df return @app.cell def _(users_arrow_table): # Convert the Arrow table to a Pandas DataFrame users_pandas_df = users_arrow_table.to_pandas() users_pandas_df return @app.cell(hide_code=True) def _(mo): mo.md(r"""### From Polars/Pandas to DuckDB""") return @app.cell def _(pl): # Create a Polars DataFrame polars_df = pl.DataFrame({ "product_id": [101, 102, 103], "product_name": ["Laptop", "Mouse", "Keyboard"], "price": [1200.00, 25.50, 75.00] }) polars_df return (polars_df,) @app.cell(hide_code=True) def _(mo): mo.md(r"""Now we can query this Polars DataFrame directly in DuckDB:""") return @app.cell def _(mo, polars_df): # Query the Polars DataFrame directly in DuckDB mo.sql( f""" SELECT product_name, price FROM polars_df WHERE price > 50 ORDER BY price DESC; """ ) return @app.cell(hide_code=True) def _(mo): mo.md(r"""Similarly, we can query a Pandas DataFrame:""") return @app.cell def _(pd): # Create a Pandas DataFrame pandas_df = pd.DataFrame({ "order_id": [1001, 1002, 1003, 1004], "product_id": [101, 102, 103, 101], "quantity": [1, 2, 1, 3], "order_date": pd.to_datetime(['2024-01-15', '2024-01-16', '2024-01-16', '2024-01-17']) }) pandas_df return (pandas_df,) @app.cell def _(mo, pandas_df): # Query the Pandas DataFrame in DuckDB mo.sql( f""" SELECT order_date, SUM(quantity) as total_quantity FROM pandas_df GROUP BY order_date ORDER BY order_date; """ ) return @app.cell(hide_code=True) def _(mo): mo.md( r""" ## 4. Advanced Example: Combining Multiple Data Sources One of the most powerful features is the ability to join data from different sources (DuckDB tables, Arrow tables, Polars/Pandas DataFrames) in a single query: """ ) return @app.cell def _(mo, pandas_df, polars_df, users): # Join the DuckDB users table with the Polars products DataFrame and Pandas orders DataFrame result = mo.sql( f""" SELECT u.name as customer_name, p.product_name, o.quantity, p.price, (o.quantity * p.price) as total_amount FROM users u CROSS JOIN pandas_df o JOIN polars_df p ON o.product_id = p.product_id WHERE u.id = 1 -- Just for Alice ORDER BY o.order_date; """ ) result return @app.cell(hide_code=True) def _(mo): mo.md( r""" ## 5. Performance Benefits of Arrow Integration The zero-copy integration between DuckDB and Apache Arrow delivers significant performance and memory benefits. This seamless integration enables: ### Key Benefits: - **Memory Efficiency**: Arrow's columnar format uses 20-40% less memory than traditional DataFrames through compact columnar representation and better compression ratios - **Zero-Copy Operations**: Data can be shared between DuckDB and Arrow-compatible systems (Polars, Pandas) without any data copying, eliminating redundant memory usage - **Query Performance**: 2-10x faster queries compared to traditional approaches that require data copying - **Larger-than-Memory Analysis**: Both DuckDB and Arrow-compatible libraries support streaming query results, allowing you to execute queries on data larger than available memory by processing data in batches. - **Advanced Query Optimization**: DuckDB's optimizer can push down filters and projections directly into Arrow scans, reading only relevant columns and partitions Let's demonstrate these benefits with concrete examples: """ ) return @app.cell(hide_code=True) def _(mo): mo.md(r"""### Memory Efficiency Demonstration""") return @app.cell def _(pd, pl): import sys import time # Create identical datasets in different formats n_rows = 1_000_000 # Pandas DataFrame (traditional approach) pandas_data = pd.DataFrame({ "id": range(n_rows), "value": [i * 2.5 for i in range(n_rows)], "category": [f"cat_{i % 100}" for i in range(n_rows)], "description": [f"This is a longer text description for row {i}" for i in range(n_rows)] }) # Polars DataFrame (Arrow-based) polars_data = pl.DataFrame({ "id": range(n_rows), "value": pl.Series([i * 2.5 for i in range(n_rows)]), "category": pl.Series([f"cat_{i % 100}" for i in range(n_rows)]), "description": pl.Series([f"This is a longer text description for row {i}" for i in range(n_rows)]) }) # Get memory usage pandas_memory = pandas_data.memory_usage(deep=True).sum() / 1024 / 1024 # MB polars_memory = polars_data.estimated_size() / 1024 / 1024 # MB print(f"Dataset size: {n_rows:,} rows") print(f"Pandas memory usage: {pandas_memory:.2f} MB") print(f"Polars (Arrow) memory usage: {polars_memory:.2f} MB") print(f"Memory savings: {((pandas_memory - polars_memory) / pandas_memory * 100):.1f}%") return pandas_data, polars_data, time @app.cell(hide_code=True) def _(mo): mo.md(r"""### Performance Comparison: Arrow vs Non-Arrow Approaches""") return @app.cell(hide_code=True) def _(mo): mo.md(r"""Let's compare three approaches for the same analytical query:""") return @app.cell def _(duckdb, mo, pandas_data, polars_data, time): # Test query: group by category and calculate aggregations query = """ SELECT category, COUNT(*) as count, AVG(value) as avg_value, MIN(value) as min_value, MAX(value) as max_value, SUM(value) as sum_value FROM data_source GROUP BY category ORDER BY count DESC """ # Approach 1: Traditional - Copy data to DuckDB table start_time = time.time() conn = duckdb.connect(':memory:') conn.execute("CREATE TABLE pandas_table AS SELECT * FROM pandas_data") result1 = conn.execute(query.replace("data_source", "pandas_table")).fetchall() # conn.close() approach1_time = time.time() - start_time # Approach 2: Direct Pandas query (no DuckDB) start_time = time.time() result2 = pandas_data.groupby('category').agg({ 'id': 'count', 'value': ['mean', 'min', 'max', 'sum'] }).sort_values(('id', 'count'), ascending=False) approach2_time = time.time() - start_time # Approach 3: Arrow-based - Zero-copy with Polars start_time = time.time() result3 = mo.sql( f""" SELECT category, COUNT(*) as count, AVG(value) as avg_value, MIN(value) as min_value, MAX(value) as max_value, SUM(value) as sum_value FROM polars_data GROUP BY category ORDER BY count DESC """ ) approach3_time = time.time() - start_time print("Performance Comparison:") print(f"1. Traditional (copy to DuckDB): {approach1_time:.3f} seconds") print(f"2. Pandas groupby: {approach2_time:.3f} seconds") print(f"3. Arrow-based (zero-copy): {approach3_time:.3f} seconds") print(f"\nSpeedup vs traditional: {approach1_time/approach3_time:.1f}x") print(f"Speedup vs pandas: {approach2_time/approach3_time:.1f}x") # Return timing variables but not the closed connection return approach1_time, approach2_time, approach3_time @app.cell(hide_code=True) def _(mo): mo.md(r"""### Visualizing the Performance Difference""") return @app.cell def _(approach1_time, approach2_time, approach3_time, mo, pl): import altair as alt # Create a bar chart showing the performance comparison performance_data = pl.DataFrame({ "Approach": ["Traditional\n(Copy to DuckDB)", "Pandas\nGroupBy", "Arrow-based\n(Zero-copy)"], "Time (seconds)": [approach1_time, approach2_time, approach3_time] }) # Create the Altair chart chart = alt.Chart(performance_data.to_pandas()).mark_bar().encode( x=alt.X("Approach", type="nominal", sort="-y"), y=alt.Y("Time (seconds)", type="quantitative"), color=alt.Color("Approach", type="nominal", scale=alt.Scale(range=["#ff6b6b", "#ffd93d", "#6bcf7f"])) ).properties( title="Query Performance Comparison", width=400, height=300 ) # Display using marimo's altair_chart UI element mo.ui.altair_chart(chart) return alt, chart, performance_data @app.cell(hide_code=True) def _(mo): mo.md(r"""### Complex Query Performance""") return @app.cell(hide_code=True) def _(mo): mo.md(r"""Let's test a more complex query with joins and window functions:""") return @app.cell def _(mo, pl, polars_data, time): # Create additional datasets for join operations categories_df = pl.DataFrame({ "category": [f"cat_{i}" for i in range(100)], "category_group": [f"group_{i // 10}" for i in range(100)], "priority": [i % 5 + 1 for i in range(100)] }) # Complex query with join and window functions new_start_time = time.time() complex_result = mo.sql( f""" WITH ranked_data AS ( SELECT d.*, c.category_group, c.priority, ROW_NUMBER() OVER (PARTITION BY c.category_group ORDER BY d.value DESC) as rank_in_group, AVG(d.value) OVER (PARTITION BY c.category_group) as group_avg_value FROM polars_data d JOIN categories_df c ON d.category = c.category ) SELECT category_group, COUNT(DISTINCT category) as unique_categories, AVG(value) as avg_value, MAX(value) as max_value, AVG(group_avg_value) as avg_group_value, COUNT(CASE WHEN rank_in_group <= 10 THEN 1 END) as top_10_count FROM ranked_data GROUP BY category_group ORDER BY avg_value DESC """ ) complex_query_time = time.time() - new_start_time print(f"Complex query with joins and window functions completed in {complex_query_time:.3f} seconds") complex_result return (categories_df,) @app.cell(hide_code=True) def _(mo): mo.md( r""" ### Memory Efficiency During Operations Let's demonstrate how Arrow's zero-copy operations save memory during data transformations: """ ) return @app.cell def _(polars_data, time): import os import pyarrow.compute as pc # Add this import # Get current process process = psutil.Process(os.getpid()) # Measure memory before operations memory_before = process.memory_info().rss / 1024 / 1024 # MB # Perform multiple Arrow-based operations (zero-copy) latest_start_time = time.time() # These operations use Arrow's zero-copy capabilities arrow_table = polars_data.to_arrow() arrow_sliced = arrow_table.slice(0, 100000) # Use PyArrow compute functions for filtering arrow_filtered = arrow_table.filter(pc.greater(arrow_table['value'], 500000)) arrow_ops_time = time.time() - latest_start_time memory_after_arrow = process.memory_info().rss / 1024 / 1024 # MB # Compare with traditional copy-based operations latest_start_time = time.time() # These operations may create copies depending on Pandas' Copy-on-Write (CoW) behavior pandas_copy = polars_data.to_pandas() pandas_sliced = pandas_copy.iloc[:100000].copy() pandas_filtered = pandas_copy[pandas_copy['value'] > 500000].copy() copy_ops_time = time.time() - latest_start_time memory_after_copy = process.memory_info().rss / 1024 / 1024 # MB print("Memory Usage Comparison:") print(f"Initial memory: {memory_before:.2f} MB") print(f"After Arrow operations: {memory_after_arrow:.2f} MB (diff: +{memory_after_arrow - memory_before:.2f} MB)") print(f"After copy operations: {memory_after_copy:.2f} MB (diff: +{memory_after_copy - memory_before:.2f} MB)") print(f"\nTime comparison:") print(f"Arrow operations: {arrow_ops_time:.3f} seconds") print(f"Copy operations: {copy_ops_time:.3f} seconds") print(f"Speedup: {copy_ops_time/arrow_ops_time:.1f}x") return pc @app.cell(hide_code=True) def _(mo): mo.md( r""" ## Summary In this notebook, we've explored: 1. **Creating Arrow tables from DuckDB queries** using `.to_arrow()` 2. **Loading Arrow tables into DuckDB** and querying them directly 3. **Converting between DuckDB, Arrow, Polars, and Pandas** with zero-copy operations 4. **Combining data from multiple sources** in a single SQL query 5. **Performance and memory benefits** including: - **Memory efficiency**: Arrow format uses 20-40% less memory than traditional DataFrames - **Query performance**: 2-10x faster queries through zero-copy operations - **Reduced memory overhead**: Operations on Arrow data avoid creating copies - **Better scalability**: Can handle larger datasets within the same memory constraints The seamless integration between DuckDB and Arrow-compatible systems makes it easy to work with data across different tools while maintaining high performance and memory efficiency. """ ) return @app.cell def _(): import marimo as mo import pyarrow as pa import polars as pl import pandas as pd import duckdb import sqlglot import psutil return duckdb, mo, pa, pd, pl if __name__ == "__main__": app.run()