oxigeo-distributed 0.2.3

Distributed processing capabilities for OxiGeo using Apache Arrow Flight
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151

//! End-to-end integration test for the Coordinator → Flight → Worker pipeline.
//!
//! Starts a real tonic Flight server backed by a [`Worker`], registers it with a
//! [`Coordinator`], and dispatches a task over gRPC via
//! [`Coordinator::dispatch_task_to_worker`]. This exercises the wiring that
//! previously did not exist: task serialization, Arrow IPC data transfer, remote
//! execution, and result collection, all across a real socket.

#![allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)]

use std::sync::Arc;
use std::time::Duration;

use arrow::array::Int32Array;
use arrow::datatypes::{DataType, Field, Schema};
use arrow::record_batch::RecordBatch;
use oxigeo_distributed::coordinator::{Coordinator, CoordinatorConfig};
use oxigeo_distributed::flight::FlightClient;
use oxigeo_distributed::flight::server::FlightServer;
use oxigeo_distributed::task::{PartitionId, Task, TaskId, TaskOperation};
use oxigeo_distributed::worker::{Worker, WorkerConfig};

fn test_batch() -> Arc<RecordBatch> {
    let schema = Arc::new(Schema::new(vec![Field::new(
        "value",
        DataType::Int32,
        false,
    )]));
    Arc::new(
        RecordBatch::try_new(
            schema,
            vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5]))],
        )
        .expect("batch"),
    )
}

async fn wait_until_reachable(address: &str) {
    for _ in 0..100 {
        if FlightClient::new(address.to_string()).await.is_ok() {
            return;
        }
        tokio::time::sleep(Duration::from_millis(50)).await;
    }
    panic!("flight server at {address} never became reachable");
}

#[tokio::test]
async fn coordinator_dispatches_task_to_remote_worker_over_flight() {
    // Reserve an ephemeral port, then hand it to the tonic server.
    let listener = std::net::TcpListener::bind("127.0.0.1:0").expect("bind");
    let port = listener.local_addr().expect("addr").port();
    drop(listener);
    let socket_addr: std::net::SocketAddr = format!("127.0.0.1:{port}").parse().expect("sockaddr");
    let http_addr = format!("http://127.0.0.1:{port}");

    // A real worker served over Flight.
    let worker = Arc::new(Worker::new(WorkerConfig::new("worker-remote".to_string())));
    let server = FlightServer::new().with_worker(worker);
    let service = server.into_service();

    let server_handle = tokio::spawn(async move {
        tonic::transport::Server::builder()
            .add_service(service)
            .serve(socket_addr)
            .await
            .expect("server");
    });

    wait_until_reachable(&http_addr).await;

    // A coordinator that knows the worker by its Flight address.
    let coordinator = Coordinator::new(CoordinatorConfig::new(http_addr.clone()));
    coordinator
        .add_worker("worker-remote".to_string(), http_addr.clone())
        .expect("add worker");

    // Dispatch a Filter task (value > 2) end-to-end; 3 of 5 rows survive.
    let task = Task::new(
        TaskId(1),
        PartitionId(0),
        TaskOperation::Filter {
            expression: "value > 2".to_string(),
        },
    );

    let result = coordinator
        .dispatch_task_to_worker(task, "worker-remote", test_batch())
        .await
        .expect("dispatch");

    assert!(result.is_success(), "remote task must succeed: {result:?}");
    let output = result.data.expect("output batch");
    assert_eq!(
        output.num_rows(),
        3,
        "the filter must have really executed on the remote worker"
    );

    // The coordinator recorded the completion in its own bookkeeping.
    let progress = coordinator.get_progress().expect("progress");
    assert_eq!(progress.completed_tasks, 1);

    server_handle.abort();
}

#[tokio::test]
async fn coordinator_dispatch_reports_failure_for_bad_expression() {
    let listener = std::net::TcpListener::bind("127.0.0.1:0").expect("bind");
    let port = listener.local_addr().expect("addr").port();
    drop(listener);
    let socket_addr: std::net::SocketAddr = format!("127.0.0.1:{port}").parse().expect("sockaddr");
    let http_addr = format!("http://127.0.0.1:{port}");

    let worker = Arc::new(Worker::new(WorkerConfig::new("worker-remote".to_string())));
    let service = FlightServer::new().with_worker(worker).into_service();
    let server_handle = tokio::spawn(async move {
        tonic::transport::Server::builder()
            .add_service(service)
            .serve(socket_addr)
            .await
            .expect("server");
    });

    wait_until_reachable(&http_addr).await;

    let coordinator = Coordinator::new(CoordinatorConfig::new(http_addr.clone()));
    coordinator
        .add_worker("worker-remote".to_string(), http_addr.clone())
        .expect("add worker");

    // A filter on a non-existent column must surface as a task failure, not a
    // crash or a silent success.
    let task = Task::new(
        TaskId(2),
        PartitionId(0),
        TaskOperation::Filter {
            expression: "no_such_column > 2".to_string(),
        },
    );

    let result = coordinator
        .dispatch_task_to_worker(task, "worker-remote", test_batch())
        .await
        .expect("dispatch returns a result even on task failure");

    assert!(result.is_failure(), "bad expression must fail the task");
    assert!(result.error.is_some());

    server_handle.abort();
}