NoahKusaba commented on code in PR #2217: URL: https://github.com/apache/datafusion-ballista/pull/2217#discussion_r3717268513
########## integrations/iceberg/tests/distributed_read_write.rs: ########## @@ -0,0 +1,712 @@ +// 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. + +//! End-to-end distributed Iceberg write/read test. +//! +//! Each test starts its own Iceberg REST catalog + MinIO with testcontainers +//! (see [`fixture`]), so this needs a working docker daemon. That is why it +//! sits behind the `integration-tests` feature and is not part of a plain +//! `cargo test`: +//! +//! ```bash +//! cargo test -p iceberg-ballista --features integration-tests --test distributed_read_write +//! ``` + +#![cfg(feature = "integration-tests")] + +mod fixture; + +use std::collections::HashMap; +use std::sync::Arc; +use std::time::Duration; + +use arrow::array::RecordBatch; +use ballista::datafusion::assert_batches_eq; +use ballista::datafusion::execution::{SessionState, SessionStateBuilder}; +use ballista::datafusion::prelude::{SessionConfig, SessionContext}; +use ballista::prelude::{SessionConfigExt, SessionContextExt}; +use ballista_core::serde::protobuf::scheduler_grpc_client::SchedulerGrpcClient; +use ballista_executor::new_standalone_executor_from_state; +use ballista_scheduler::standalone::new_standalone_scheduler_from_state; +use iceberg::spec::{ + NestedField, PrimitiveType, Schema, Transform, Type, UnboundPartitionField, + UnboundPartitionSpec, +}; +use iceberg::transaction::{AddColumn, ApplyTransactionAction, Transaction}; +use iceberg::{Catalog, CatalogBuilder, NamespaceIdent, TableCreation, TableIdent}; +use iceberg_ballista::{ + IcebergCatalogConfig, register_iceberg_catalog, register_iceberg_codecs, + register_iceberg_table, +}; +use iceberg_catalog_rest::RestCatalogBuilder; +use iceberg_datafusion::IcebergTableProvider; +use iceberg_storage_opendal::OpenDalStorageFactory; + +use crate::fixture::IcebergFixture; + +/// Session state with the Iceberg codecs installed. +fn iceberg_session_state(config: SessionConfig) -> SessionState { + SessionStateBuilder::new() + .with_config(register_iceberg_codecs(config)) + .with_default_features() + .build() +} + +/// Runs a SQL statement to completion and returns its batches. +async fn run_sql(ctx: &SessionContext, sql: &str) -> Vec<RecordBatch> { + ctx.sql(sql) + .await + .expect("plan sql") + .collect() + .await + .expect("run sql") +} + +async fn build_rest_catalog(props: &HashMap<String, String>) -> impl Catalog + use<> { + RestCatalogBuilder::default() + .with_storage_factory(Arc::new(OpenDalStorageFactory::S3 { + customized_credential_load: None, + })) + .load("rest", props.clone()) + .await + .expect("build rest catalog") +} + +/// Creates the test namespace. Each test owns its catalog, so this never races +/// another creator. +async fn create_namespace(catalog: &impl Catalog) -> NamespaceIdent { + let namespace = NamespaceIdent::new("ballista_it".to_string()); + catalog + .create_namespace(&namespace, HashMap::new()) + .await + .expect("create namespace"); + namespace +} + +/// Creates `table_name` with `schema` (optionally partitioned) in the test +/// namespace and returns that namespace. +async fn create_table_with( + props: &HashMap<String, String>, + table_name: &str, + schema: Schema, + partition_spec: Option<UnboundPartitionSpec>, +) -> NamespaceIdent { + let catalog = build_rest_catalog(props).await; + let namespace = create_namespace(&catalog).await; + + let builder = TableCreation::builder() + .name(table_name.to_string()) + .schema(schema) + .properties(HashMap::new()); + let creation = match partition_spec { + Some(spec) => builder.partition_spec(spec).build(), + None => builder.build(), + }; + catalog + .create_table(&namespace, creation) + .await + .expect("create table"); + + namespace +} + +async fn create_table( + props: &HashMap<String, String>, + table_name: &str, +) -> NamespaceIdent { + let schema = Schema::builder() + .with_schema_id(0) + .with_fields(vec![ + NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(), + NestedField::required(2, "name", Type::Primitive(PrimitiveType::String)) + .into(), + ]) + .build() + .unwrap(); + create_table_with(props, table_name, schema, None).await +} + +/// Loads the table from the catalog and returns its current snapshot id. +async fn current_snapshot_id( + props: &HashMap<String, String>, + namespace: &NamespaceIdent, + table_name: &str, +) -> i64 { + build_rest_catalog(props) + .await + .load_table(&TableIdent::new(namespace.clone(), table_name.to_string())) + .await + .expect("load table") + .metadata() + .current_snapshot_id() + .expect("table has a snapshot") +} + +/// Creates a table partitioned by `region` (identity). A distributed INSERT then +/// fans the rows out to one writer per region — exercising the partition-value +/// expression (`PartitionExpr`) serialization across the cluster. +async fn create_partitioned_table( + props: &HashMap<String, String>, + table_name: &str, +) -> NamespaceIdent { + // Optional (nullable) fields so the schema matches the nullable columns a + // `VALUES` source produces; the partitioned-write path checks nullability. + let schema = Schema::builder() + .with_schema_id(0) + .with_fields(vec![ + NestedField::optional(1, "id", Type::Primitive(PrimitiveType::Int)).into(), + NestedField::optional(2, "region", Type::Primitive(PrimitiveType::String)) + .into(), + ]) + .build() + .unwrap(); + let partition_spec = UnboundPartitionSpec::builder() + .with_spec_id(0) + // The REST catalog requires an explicit partition field-id (some + // catalogs assign one automatically; REST does not). + .add_partition_fields([UnboundPartitionField { + source_id: 2, + field_id: Some(1000), + name: "region".to_string(), + transform: Transform::Identity, + }]) + .unwrap() + .build(); + create_table_with(props, table_name, schema, Some(partition_spec)).await +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn distributed_insert_and_read() { + let _ = env_logger::builder().is_test(true).try_init(); + + let fixture = IcebergFixture::start().await; + let props = fixture.props(); + let table_name = "events".to_string(); + let namespace = create_table(&props, &table_name).await; + + let state = iceberg_session_state( + SessionConfig::new_with_ballista() + .with_target_partitions(2) + .with_ballista_standalone_parallelism(2), + ); + let ctx = SessionContext::standalone_with_state(state) + .await + .expect("start standalone ballista"); + + let catalog_config = IcebergCatalogConfig::new("rest", "rest", props.clone()); + register_iceberg_table( + &ctx, + "events", + catalog_config, + namespace, + table_name.clone(), + ) + .await + .expect("register iceberg table"); + + // Read *before* any insert. This first read is what makes the re-reads below + // meaningful: it forces the registered provider through a full plan/scan + // cycle while the table is still empty, so a provider that froze its table + // metadata here would keep returning this empty snapshot forever. + assert_batches_eq!( + ["+---+", "| n |", "+---+", "| 0 |", "+---+"], + &run_sql(&ctx, "SELECT count(*) AS n FROM events").await + ); + + run_sql( + &ctx, + "INSERT INTO events VALUES (1, 'alice'), (2, 'bob'), (3, 'carol')", + ) + .await; + + assert_batches_eq!( + ["+---+", "| n |", "+---+", "| 3 |", "+---+"], + &run_sql(&ctx, "SELECT count(*) AS n FROM events").await + ); + + assert_batches_eq!( + [ + "+----+-------+", + "| id | name |", + "+----+-------+", + "| 1 | alice |", + "| 2 | bob |", + "| 3 | carol |", + "+----+-------+", + ], + &run_sql(&ctx, "SELECT id, name FROM events ORDER BY id").await + ); + + // Insert again through the same registration and re-read. `scan` reloads + // table metadata from the catalog every time, so each query plans against + // the snapshot current at *that* moment — the codec then pins it for the + // executors. A provider (or a cached catalog/table) that went stale after + // the first read would still report 3 rows here. + run_sql(&ctx, "INSERT INTO events VALUES (4, 'dave'), (5, 'erin')").await; + + // The re-read must observe the interposed insert: both the original and the + // newly inserted rows. + assert_batches_eq!( + ["+---+", "| n |", "+---+", "| 5 |", "+---+"], + &run_sql(&ctx, "SELECT count(*) AS n FROM events").await + ); + + assert_batches_eq!( + [ + "+----+-------+", + "| id | name |", + "+----+-------+", + "| 1 | alice |", + "| 2 | bob |", + "| 3 | carol |", + "| 4 | dave |", + "| 5 | erin |", + "+----+-------+", + ], + &run_sql(&ctx, "SELECT id, name FROM events ORDER BY id").await + ); + + // Catalog-level registration: mount the whole Iceberg catalog and read the + // same table as `<catalog>.<namespace>.<table>`. The providers built through + // the catalog carry the config too, so this distributed read exercises the + // catalog/schema config-threading path end to end. + register_iceberg_catalog( + &ctx, + "ice", + IcebergCatalogConfig::new("rest", "rest", props), + ) + .await + .expect("register iceberg catalog"); + assert_batches_eq!( + ["+---+", "| n |", "+---+", "| 5 |", "+---+"], + &run_sql( + &ctx, + &format!("SELECT count(*) AS n FROM ice.ballista_it.{table_name}"), + ) + .await + ); +} + +/// Distributed correctness on a real multi-executor cluster, writing a +/// **partitioned** table. +/// +/// Where [`distributed_insert_and_read`] uses standalone Ballista (one in-process +/// executor) and an unpartitioned table, this stands up a single scheduler with +/// **several in-process executors** and writes a table partitioned by `region`. A +/// partitioned write injects a partition-value expression (`PartitionExpr`) into +/// the physical plan, so this exercises that expression's serialization through +/// the codec on top of the plan-node serialization — and fans the rows out to one +/// writer per region across the executors. +/// +/// The assertions target the correctness properties of a distributed, multi-writer +/// write: +/// 1. the write commits exactly **one** atomic snapshot (not one per task), +/// 2. the parallel writers contributed multiple data files (one per region), and +/// 3. every input row lands exactly once (no loss, no duplication). +#[tokio::test(flavor = "multi_thread", worker_threads = 8)] Review Comment: @andygrove This PR does cover a true multi-executor cluster deployment. This is the test that covers that, by making a standalone scheduler + executors. I'll add in validation to ensure both executors did work (not just 1), and add an example for use. -- 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]
