NoahKusaba commented on code in PR #2217: URL: https://github.com/apache/datafusion-ballista/pull/2217#discussion_r3737184055
########## integrations/iceberg/src/lib.rs: ########## @@ -0,0 +1,124 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! # Ballista Iceberg driver +//! +//! Adds distributed Apache Iceberg reads and writes to Ballista. +//! +//! Iceberg's DataFusion integration already produces a complete physical plan; +//! the only thing missing for Ballista is serialization. Ballista ships plans to +//! remote nodes, but the Iceberg plan nodes hold live catalog/storage handles +//! that can't be serialized. This crate's logical and physical extension codecs +//! serialize the minimal config needed to rebuild those handles per node. +//! +//! ## Usage (standalone) +//! +//! ```ignore +//! use std::collections::HashMap; +//! +//! use iceberg_ballista::{register_iceberg_codecs, register_iceberg_table, IcebergCatalogConfig}; +//! use ballista_core::extension::SessionConfigExt; +//! use datafusion::prelude::{SessionConfig, SessionContext}; +//! use iceberg::NamespaceIdent; +//! +//! # async fn run() -> datafusion::error::Result<()> { +//! // 1. Register the Iceberg codecs on the session config, then start standalone Ballista. +//! let config = register_iceberg_codecs(SessionConfig::new_with_ballista()); +//! let ctx = SessionContext::standalone_with_config(config).await?; +//! +//! // 2. Register a catalog-backed Iceberg table for reads and writes. +//! let props = HashMap::from([("uri".to_string(), "http://localhost:8181".to_string())]); +//! let cfg = IcebergCatalogConfig::new("rest", "rest", props); +//! register_iceberg_table(&ctx, "t", cfg, NamespaceIdent::new("ns".into()), "tbl").await?; +//! +//! // 3. INSERT runs distributed across the cluster. +//! ctx.sql("INSERT INTO t SELECT * FROM source").await?.collect().await?; +//! # Ok(()) +//! # } +//! ``` + +mod bridge; +mod logical_codec; +mod physical_codec; + +use std::sync::Arc; + +use ballista_core::extension::SessionConfigExt; +use datafusion::common::DataFusionError; +use datafusion::prelude::{SessionConfig, SessionContext}; +use iceberg::NamespaceIdent; +pub use iceberg_datafusion::IcebergCatalogConfig; +use iceberg_datafusion::to_datafusion_error; + +pub use crate::logical_codec::IcebergLogicalCodec; +pub use crate::physical_codec::IcebergPhysicalCodec; + +/// Installs the Iceberg logical and physical extension codecs onto a +/// [`SessionConfig`]. +/// +/// In a standalone cluster the scheduler and executor both derive their codecs +/// from this config, so one call suffices. For a separately deployed scheduler +/// and executor, set the same codecs on their process configs +/// (`override_logical_codec` / `override_physical_codec`). +pub fn register_iceberg_codecs(config: SessionConfig) -> SessionConfig { + config + .with_ballista_logical_extension_codec(Arc::new(IcebergLogicalCodec::default())) Review Comment: Good catch! Hadn't considered other codec's being used besides Ballista's default. Will add it in a follow up PR. I should also note that currently logical_codec's try_encode falls back to the set inner codec if it's not a icebergTableProvider, so this should get us the desired behavior. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
