DataFusion Integration¶

use rocklake_datafusion::RockLakeCatalog;
use datafusion::prelude::*;
use std::sync::Arc;

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    // Create a DataFusion session
    let ctx = SessionContext::new();

    // Open a RockLake catalog from object storage
    let catalog = RockLakeCatalog::open("s3://my-bucket/lakehouse/catalog/").await?;

    // Register it with DataFusion under the name "lake"
    ctx.register_catalog("lake", Arc::new(catalog));

    // Query tables through standard SQL
    let df = ctx.sql("SELECT event_type, count(*) as cnt 
                      FROM lake.analytics.events 
                      WHERE timestamp > '2024-03-01' 
                      GROUP BY event_type 
                      ORDER BY cnt DESC")
        .await?;

    df.show().await?;
    Ok(())
}

use rocklake_datafusion::{RockLakeCatalog, CatalogConfig};
use object_store::aws::AmazonS3Builder;

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    // Configure object storage explicitly
    let config = CatalogConfig::builder()
        .storage_url("s3://my-bucket/catalog/")
        .aws_region("us-east-1")
        .aws_endpoint("https://s3.us-east-1.amazonaws.com")
        .build();

    let catalog = RockLakeCatalog::open_with_config(config).await?;

    let ctx = SessionContext::new();
    ctx.register_catalog("lake", Arc::new(catalog));

    // Register the same S3 object store for DataFusion to read Parquet files
    let s3 = AmazonS3Builder::from_env()
        .with_bucket_name("my-bucket")
        .build()?;
    ctx.register_object_store("s3://my-bucket", Arc::new(s3));

    let df = ctx.sql("SELECT * FROM lake.analytics.events LIMIT 10").await?;
    df.show().await?;
    Ok(())
}

use rocklake_datafusion::RockLakeCatalog;

// Open catalog at a specific point in time
let catalog = RockLakeCatalog::open_at_snapshot("s3://bucket/catalog/", 1000).await?;
ctx.register_catalog("lake_snapshot", Arc::new(catalog));

// All queries through this catalog see the state at snapshot 1000
let df = ctx.sql("SELECT count(*) FROM lake_snapshot.analytics.events").await?;

// Register multiple catalogs for cross-catalog queries
let prod = RockLakeCatalog::open("s3://prod-bucket/catalog/").await?;
let staging = RockLakeCatalog::open("s3://staging-bucket/catalog/").await?;

ctx.register_catalog("prod", Arc::new(prod));
ctx.register_catalog("staging", Arc::new(staging));

// Compare data across environments
let df = ctx.sql("
    SELECT 'prod' as env, count(*) as rows FROM prod.analytics.events
    UNION ALL
    SELECT 'staging' as env, count(*) as rows FROM staging.analytics.events
").await?;

use rocklake_catalog::store::CatalogStore;
use rocklake_datafusion::RockLakeCatalogProvider;
use datafusion::prelude::*;
use std::sync::Arc;

// Assume `store` is an Arc<Mutex<CatalogStore>> already open.
let provider = RockLakeCatalogProvider::from_catalog_store(store.clone()).await?;

let ctx = SessionContext::new();
ctx.register_catalog("lake", Arc::new(provider));

// Example: REST API that queries a DuckLake catalog
async fn query_handler(req: QueryRequest) -> Result<QueryResponse, Error> {
    let ctx = get_session_context(); // Pre-configured with RockLake catalog
    let df = ctx.sql(&req.sql).await?;
    let batches = df.collect().await?;
    Ok(QueryResponse::from_arrow(batches))
}

// Streaming analytics that reads catalog metadata from RockLake
// and processes new Parquet files as they arrive
async fn process_new_files(catalog: &RockLakeCatalog) {
    let table = catalog.get_table("analytics", "events").await?;
    let new_files = table.files_since_snapshot(last_processed_snapshot).await?;
    for file in new_files {
        process_parquet_file(&file.path).await?;
    }
}

// Integration test: verify RockLake returns same results through both paths
#[tokio::test]
async fn test_catalog_consistency() {
    let catalog = RockLakeCatalog::open("./test-catalog/").await.unwrap();
    let ctx = SessionContext::new();
    ctx.register_catalog("lake", Arc::new(catalog));

    let df = ctx.sql("SELECT count(*) FROM lake.public.events").await.unwrap();
    let batches = df.collect().await.unwrap();
    let count: i64 = batches[0].column(0).as_primitive::<Int64Type>().value(0);

    assert_eq!(count, expected_count_from_duckdb);
}

// Self-contained analytics binary
fn main() {
    let rt = tokio::runtime::Runtime::new().unwrap();
    rt.block_on(async {
        let catalog = RockLakeCatalog::open("./data/catalog/").await.unwrap();
        let ctx = SessionContext::new();
        ctx.register_catalog("local", Arc::new(catalog));

        // Run analytics queries
        let report = ctx.sql("SELECT ...").await.unwrap();
        export_to_csv(report, "output.csv").await;
    });
}

use datafusion::error::DataFusionError;

match ctx.sql("SELECT * FROM lake.nonexistent.table").await {
    Ok(df) => { /* success */ },
    Err(DataFusionError::Plan(msg)) => {
        // Schema or table not found
        eprintln!("Planning error: {}", msg);
    },
    Err(DataFusionError::External(e)) => {
        // Storage or catalog error
        eprintln!("External error: {}", e);
    },
    Err(e) => {
        eprintln!("Other error: {}", e);
    }
}

// Example: REST API that queries the lakehouse
async fn handle_query(query: String) -> Result<Vec<RecordBatch>> {
    let catalog = RockLakeCatalog::open("s3://bucket/catalog/").await?;
    let ctx = SessionContext::new();
    ctx.register_catalog("lake", Arc::new(catalog));

    let df = ctx.sql(&query).await?;
    let batches = df.collect().await?;
    Ok(batches)
}

// Validate that a table has expected columns before loading data
let catalog = RockLakeCatalog::open(&storage_url).await?;
let ctx = SessionContext::new();
ctx.register_catalog("lake", Arc::new(catalog));

let schema = ctx.sql("DESCRIBE lake.analytics.events").await?.collect().await?;
let columns: Vec<String> = schema.iter()
    .flat_map(|b| b.column(0).as_any().downcast_ref::<StringArray>())
    .flat_map(|a| (0..a.len()).map(|i| a.value(i).to_string()))
    .collect();

assert!(columns.contains(&"event_id".to_string()));
assert!(columns.contains(&"created_at".to_string()));

// Use the catalog for metadata but execute differently
let catalog = RockLakeCatalog::open(&storage_url).await?;
let ctx = SessionContext::new();
ctx.register_catalog("lake", Arc::new(catalog));

// Get the logical plan (not execution)
let plan = ctx.sql("SELECT * FROM lake.analytics.events WHERE date > '2024-01-01'")
    .await?
    .into_optimized_plan()?;

// Extract file list from the plan for custom processing
let files = extract_data_files_from_plan(&plan);
for file in files {
    custom_streaming_processor.ingest(file).await?;
}