test_broadcasting/
test_broadcasting.rs

1//! # Comprehensive Broadcasting Examples
2//!
3//! This example demonstrates Minarrow's broadcasting capabilities across different array types
4//! and operations. Broadcasting allows operations between arrays of different shapes by
5//! automatically replicating smaller arrays to match larger ones.
6//!
7//! ## Broadcasting Rules
8//! - Arrays with matching lengths operate element-wise
9//! - Single-element arrays broadcast to match the length of the other operand
10//! - Type promotion occurs automatically (e.g., Int32 + Float32 → Float32)
11//! - Complex types (Table, Cube, SuperArray) support element-wise broadcasting
12
13use minarrow::{Array, FloatArray, IntegerArray, NumericArray, Table, Value, vec64};
14use std::sync::Arc;
15
16#[cfg(feature = "views")]
17use minarrow::ArrayV;
18
19#[cfg(feature = "chunked")]
20use minarrow::SuperArray;
21
22#[cfg(feature = "cube")]
23use minarrow::Cube;
24
25fn main() {
26    println!("═══════════════════════════════════════════════════════════");
27    println!("  Minarrow Comprehensive Broadcasting Examples");
28    println!("═══════════════════════════════════════════════════════════\n");
29
30    test_integer_broadcasting();
31    test_float_broadcasting();
32    test_mixed_type_promotion();
33    test_scalar_broadcasting();
34    test_division_broadcasting();
35    test_reference_operations();
36    test_subtraction_broadcasting();
37    test_chained_operations();
38    test_table_broadcasting();
39    test_array_view_broadcasting();
40    test_super_array_broadcasting();
41    test_cube_broadcasting();
42
43    println!("\n═══════════════════════════════════════════════════════════");
44    println!("  All broadcasting tests completed successfully!");
45    println!("═══════════════════════════════════════════════════════════");
46}
47
48/// Test integer array broadcasting with multiplication
49fn test_integer_broadcasting() {
50    println!("┌─ Test 1: Integer Broadcasting");
51    println!("│  Operation: [100] * [1, 2, 3, 4, 5]");
52    println!("│  Expected:  [100, 200, 300, 400, 500]");
53
54    let scalar_array = Value::Array(Arc::new(Array::from_int32(IntegerArray::from_slice(
55        &vec64![100],
56    ))));
57    let multi_array = Value::Array(Arc::new(Array::from_int32(IntegerArray::from_slice(
58        &vec64![1, 2, 3, 4, 5],
59    ))));
60
61    match scalar_array * multi_array {
62        Ok(Value::Array(arr_arc)) => {
63            if let Array::NumericArray(NumericArray::Int32(arr)) = arr_arc.as_ref() {
64                println!("│  Result:    {:?}", arr.data.as_slice());
65                println!("└─ ✓ Passed\n");
66            } else {
67                println!("└─ ✗ Error: Unexpected array type\n");
68            }
69        }
70        Ok(_) => println!("└─ ✗ Error: Unexpected result type\n"),
71        Err(e) => println!("└─ ✗ Error: {:?}\n", e),
72    }
73}
74
75/// Test float array broadcasting with addition
76fn test_float_broadcasting() {
77    println!("┌─ Test 2: Float Broadcasting");
78    println!("│  Operation: [2.5] + [1.0, 2.0, 3.0]");
79    println!("│  Expected:  [3.5, 4.5, 5.5]");
80
81    let scalar_float = Value::Array(Arc::new(Array::from_float64(FloatArray::from_slice(
82        &vec64![2.5],
83    ))));
84    let multi_float = Value::Array(Arc::new(Array::from_float64(FloatArray::from_slice(
85        &vec64![1.0, 2.0, 3.0],
86    ))));
87
88    match scalar_float + multi_float {
89        Ok(Value::Array(arr_arc)) => {
90            if let Array::NumericArray(NumericArray::Float64(arr)) = arr_arc.as_ref() {
91                println!("│  Result:    {:?}", arr.data.as_slice());
92                println!("└─ ✓ Passed\n");
93            } else {
94                println!("└─ ✗ Error: Unexpected array type\n");
95            }
96        }
97        Ok(_) => println!("└─ ✗ Error: Unexpected result type\n"),
98        Err(e) => println!("└─ ✗ Error: {:?}\n", e),
99    }
100}
101
102/// Test automatic type promotion from integer to float
103fn test_mixed_type_promotion() {
104    println!("┌─ Test 3: Mixed Type Promotion");
105    println!("│  Operation: Int32[10, 20, 30] + Float32[0.5, 0.5, 0.5]");
106    println!("│  Expected:  Float32[10.5, 20.5, 30.5]");
107
108    let int_array = Value::Array(Arc::new(Array::from_int32(IntegerArray::from_slice(
109        &vec64![10, 20, 30],
110    ))));
111    let float_array = Value::Array(Arc::new(Array::from_float32(FloatArray::from_slice(
112        &vec64![0.5, 0.5, 0.5],
113    ))));
114
115    match int_array + float_array {
116        Ok(Value::Array(arr_arc)) => {
117            if let Array::NumericArray(NumericArray::Float32(arr)) = arr_arc.as_ref() {
118                println!(
119                    "│  Result:    {:?} (promoted to Float32)",
120                    arr.data.as_slice()
121                );
122                println!("└─ ✓ Passed\n");
123            } else {
124                println!("└─ ✗ Error: Unexpected array type\n");
125            }
126        }
127        Ok(_) => println!("└─ ✗ Error: Unexpected result type\n"),
128        Err(e) => println!("└─ ✗ Error: {:?}\n", e),
129    }
130}
131
132/// Test broadcasting with Scalar type (requires scalar_type feature)
133fn test_scalar_broadcasting() {
134    #[cfg(feature = "scalar_type")]
135    {
136        println!("┌─ Test 4: Scalar + Array Broadcasting");
137        println!("│  Operation: Scalar(1000) + [1, 2, 3]");
138        println!("│  Expected:  [1001, 1002, 1003]");
139
140        let scalar = Value::Scalar(minarrow::Scalar::Int64(1000));
141        let array = Value::Array(Arc::new(Array::from_int64(IntegerArray::from_slice(
142            &vec64![1, 2, 3],
143        ))));
144
145        match scalar + array {
146            Ok(Value::Array(arr_arc)) => {
147                if let Array::NumericArray(NumericArray::Int64(arr)) = arr_arc.as_ref() {
148                    println!("│  Result:    {:?}", arr.data.as_slice());
149                    println!("└─ ✓ Passed\n");
150                } else {
151                    println!("└─ ✗ Error: Unexpected array type\n");
152                }
153            }
154            Ok(_) => println!("└─ ✗ Error: Unexpected result type\n"),
155            Err(e) => println!("└─ ✗ Error: {:?}\n", e),
156        }
157    }
158
159    #[cfg(not(feature = "scalar_type"))]
160    {
161        println!("┌─ Test 4: Scalar + Array Broadcasting");
162        println!("└─ ⊘ Skipped (scalar_type feature not enabled)\n");
163    }
164}
165
166/// Test division broadcasting
167fn test_division_broadcasting() {
168    println!("┌─ Test 5: Division Broadcasting");
169    println!("│  Operation: [100.0] / [2.0, 4.0, 5.0, 10.0]");
170    println!("│  Expected:  [50.0, 25.0, 20.0, 10.0]");
171
172    let dividend = Value::Array(Arc::new(Array::from_float64(FloatArray::from_slice(
173        &vec64![100.0],
174    ))));
175    let divisors = Value::Array(Arc::new(Array::from_float64(FloatArray::from_slice(
176        &vec64![2.0, 4.0, 5.0, 10.0],
177    ))));
178
179    match dividend / divisors {
180        Ok(Value::Array(arr_arc)) => {
181            if let Array::NumericArray(NumericArray::Float64(arr)) = arr_arc.as_ref() {
182                println!("│  Result:    {:?}", arr.data.as_slice());
183                println!("└─ ✓ Passed\n");
184            } else {
185                println!("└─ ✗ Error: Unexpected array type\n");
186            }
187        }
188        Ok(_) => println!("└─ ✗ Error: Unexpected result type\n"),
189        Err(e) => println!("└─ ✗ Error: {:?}\n", e),
190    }
191}
192
193/// Test operations using references (non-consuming)
194fn test_reference_operations() {
195    println!("┌─ Test 6: Reference Operations (Non-Consuming)");
196    println!("│  Operation: &[5] * &[10, 20, 30]");
197    println!("│  Expected:  [50, 100, 150]");
198
199    let a = Value::Array(Arc::new(Array::from_int32(IntegerArray::from_slice(
200        &vec64![5],
201    ))));
202    let b = Value::Array(Arc::new(Array::from_int32(IntegerArray::from_slice(
203        &vec64![10, 20, 30],
204    ))));
205
206    match &a * &b {
207        Ok(Value::Array(arr_arc)) => {
208            if let Array::NumericArray(NumericArray::Int32(arr)) = arr_arc.as_ref() {
209                println!("│  Result:    {:?}", arr.data.as_slice());
210                println!("│  Note: Original arrays remain available for reuse");
211                println!("└─ ✓ Passed\n");
212            } else {
213                println!("└─ ✗ Error: Unexpected array type\n");
214            }
215        }
216        Ok(_) => println!("└─ ✗ Error: Unexpected result type\n"),
217        Err(e) => println!("└─ ✗ Error: {:?}\n", e),
218    }
219}
220
221/// Test subtraction with broadcasting
222fn test_subtraction_broadcasting() {
223    println!("┌─ Test 7: Subtraction Broadcasting");
224    println!("│  Operation: [100, 200, 300] - [1]");
225    println!("│  Expected:  [99, 199, 299]");
226
227    let array = Value::Array(Arc::new(Array::from_int32(IntegerArray::from_slice(
228        &vec64![100, 200, 300],
229    ))));
230    let scalar_array = Value::Array(Arc::new(Array::from_int32(IntegerArray::from_slice(
231        &vec64![1],
232    ))));
233
234    match array - scalar_array {
235        Ok(Value::Array(arr_arc)) => {
236            if let Array::NumericArray(NumericArray::Int32(arr)) = arr_arc.as_ref() {
237                println!("│  Result:    {:?}", arr.data.as_slice());
238                println!("└─ ✓ Passed\n");
239            } else {
240                println!("└─ ✗ Error: Unexpected array type\n");
241            }
242        }
243        Ok(_) => println!("└─ ✗ Error: Unexpected result type\n"),
244        Err(e) => println!("└─ ✗ Error: {:?}\n", e),
245    }
246}
247
248/// Test chained operations with broadcasting
249fn test_chained_operations() {
250    println!("┌─ Test 8: Chained Operations");
251    println!("│  Operation: ([2] * [1, 2, 3]) + [10]");
252    println!("│  Expected:  [12, 14, 16]");
253
254    let two = Value::Array(Arc::new(Array::from_int32(IntegerArray::from_slice(
255        &vec64![2],
256    ))));
257    let nums = Value::Array(Arc::new(Array::from_int32(IntegerArray::from_slice(
258        &vec64![1, 2, 3],
259    ))));
260    let ten = Value::Array(Arc::new(Array::from_int32(IntegerArray::from_slice(
261        &vec64![10],
262    ))));
263
264    let step1 = (two * nums).expect("First operation failed");
265    match step1 + ten {
266        Ok(Value::Array(arr_arc)) => {
267            if let Array::NumericArray(NumericArray::Int32(arr)) = arr_arc.as_ref() {
268                println!("│  Result:    {:?}", arr.data.as_slice());
269                println!("└─ ✓ Passed\n");
270            } else {
271                println!("└─ ✗ Error: Unexpected array type\n");
272            }
273        }
274        Ok(_) => println!("└─ ✗ Error: Unexpected result type\n"),
275        Err(e) => println!("└─ ✗ Error: {:?}\n", e),
276    }
277}
278
279/// Test Table broadcasting - operates on each column
280fn test_table_broadcasting() {
281    println!("┌─ Test 9: Table Broadcasting");
282    println!(
283        "│  Operation: Table{{col1:[1,2,3], col2:[4,5,6]}} + Table{{col1:[10,10,10], col2:[20,20,20]}}"
284    );
285    println!("│  Expected:  Table{{col1:[11,12,13], col2:[24,25,26]}}");
286
287    // Create first table with two columns
288    let arr1_col1 = Array::from_int32(IntegerArray::from_slice(&vec64![1, 2, 3]));
289    let arr1_col2 = Array::from_int32(IntegerArray::from_slice(&vec64![4, 5, 6]));
290    let fa1_col1 = minarrow::FieldArray::from_arr("col1", arr1_col1);
291    let fa1_col2 = minarrow::FieldArray::from_arr("col2", arr1_col2);
292    let mut table1 = Table::new("table1".to_string(), None);
293    table1.add_col(fa1_col1);
294    table1.add_col(fa1_col2);
295
296    // Create second table with matching structure
297    let arr2_col1 = Array::from_int32(IntegerArray::from_slice(&vec64![10, 10, 10]));
298    let arr2_col2 = Array::from_int32(IntegerArray::from_slice(&vec64![20, 20, 20]));
299    let fa2_col1 = minarrow::FieldArray::from_arr("col1", arr2_col1);
300    let fa2_col2 = minarrow::FieldArray::from_arr("col2", arr2_col2);
301    let mut table2 = Table::new("table2".to_string(), None);
302    table2.add_col(fa2_col1);
303    table2.add_col(fa2_col2);
304
305    match Value::Table(Arc::new(table1)) + Value::Table(Arc::new(table2)) {
306        Ok(Value::Table(result)) => {
307            if let Array::NumericArray(NumericArray::Int32(col1)) = &result.cols[0].array {
308                println!("│  Result col1: {:?}", col1.data.as_slice());
309            }
310            if let Array::NumericArray(NumericArray::Int32(col2)) = &result.cols[1].array {
311                println!("│  Result col2: {:?}", col2.data.as_slice());
312            }
313            println!("└─ ✓ Passed\n");
314        }
315        Ok(other) => println!("└─ ✗ Error: Unexpected result type {:?}\n", other),
316        Err(e) => println!("└─ ✗ Error: {:?}\n", e),
317    }
318}
319
320/// Test ArrayView broadcasting - efficient windowed operations
321fn test_array_view_broadcasting() {
322    #[cfg(feature = "views")]
323    {
324        println!("┌─ Test 10: ArrayView Broadcasting");
325        println!("│  Operation: ArrayView([2,3,4]) + ArrayView([10,10,10])");
326        println!("│  Expected:  Array([12,13,14])");
327
328        // Create an array and a view into it
329        let arr1 = Array::from_int32(IntegerArray::from_slice(&vec64![1, 2, 3, 4, 5]));
330        let view1 = ArrayV::new(arr1, 1, 3); // View of elements [2,3,4]
331
332        let arr2 = Array::from_int32(IntegerArray::from_slice(&vec64![10, 10, 10]));
333        let view2 = ArrayV::new(arr2, 0, 3);
334
335        match Value::ArrayView(Arc::new(view1)) + Value::ArrayView(Arc::new(view2)) {
336            Ok(Value::Array(arr_arc)) => {
337                if let Array::NumericArray(NumericArray::Int32(result)) = arr_arc.as_ref() {
338                    println!("│  Result:    {:?}", result.data.as_slice());
339                    println!("└─ ✓ Passed\n");
340                } else {
341                    println!("└─ ✗ Error: Unexpected array type\n");
342                }
343            }
344            Ok(_) => println!("└─ ✗ Error: Unexpected result type\n"),
345            Err(e) => println!("└─ ✗ Error: {:?}\n", e),
346        }
347    }
348
349    #[cfg(not(feature = "views"))]
350    {
351        println!("┌─ Test 10: ArrayView Broadcasting");
352        println!("└─ ⊘ Skipped (views feature not enabled)\n");
353    }
354}
355
356/// Test SuperArray broadcasting - chunked array operations
357fn test_super_array_broadcasting() {
358    #[cfg(feature = "chunked")]
359    {
360        println!("┌─ Test 11: SuperArray (Chunked) Broadcasting");
361        println!("│  Operation: SuperArray{{[1,2],[3,4]}} * SuperArray{{[2,2],[2,2]}}");
362        println!("│  Expected:  SuperArray{{[2,4],[6,8]}}");
363
364        // Create chunked arrays (multiple field array chunks)
365        let chunk1_a = Array::from_int32(IntegerArray::from_slice(&vec64![1, 2]));
366        let chunk2_a = Array::from_int32(IntegerArray::from_slice(&vec64![3, 4]));
367        let fa1 = minarrow::FieldArray::from_arr("chunk1", chunk1_a);
368        let fa2 = minarrow::FieldArray::from_arr("chunk1", chunk2_a);
369        let super_arr1 = SuperArray::from_field_array_chunks(vec![fa1, fa2]);
370
371        let chunk1_b = Array::from_int32(IntegerArray::from_slice(&vec64![2, 2]));
372        let chunk2_b = Array::from_int32(IntegerArray::from_slice(&vec64![2, 2]));
373        let fa3 = minarrow::FieldArray::from_arr("chunk1", chunk1_b);
374        let fa4 = minarrow::FieldArray::from_arr("chunk1", chunk2_b);
375        let super_arr2 = SuperArray::from_field_array_chunks(vec![fa3, fa4]);
376
377        match Value::SuperArray(Arc::new(super_arr1)) * Value::SuperArray(Arc::new(super_arr2)) {
378            Ok(Value::SuperArray(result)) => {
379                println!("│  Result with {} chunks:", result.len());
380                for i in 0..result.len() {
381                    if let Some(fa) = result.chunk(i) {
382                        if let Array::NumericArray(NumericArray::Int32(arr)) = &fa.array {
383                            println!("│    Chunk {}: {:?}", i, arr.data.as_slice());
384                        }
385                    }
386                }
387                println!("└─ ✓ Passed\n");
388            }
389            Ok(other) => println!("└─ ✗ Error: Unexpected result type {:?}\n", other),
390            Err(e) => println!("└─ ✗ Error: {:?}\n", e),
391        }
392    }
393
394    #[cfg(not(feature = "chunked"))]
395    {
396        println!("┌─ Test 11: SuperArray (Chunked) Broadcasting");
397        println!("└─ ⊘ Skipped (chunked feature not enabled)\n");
398    }
399}
400
401/// Test Cube broadcasting - 3D tensor operations
402fn test_cube_broadcasting() {
403    #[cfg(feature = "cube")]
404    {
405        println!("┌─ Test 12: Cube (3D) Broadcasting");
406        println!("│  Operation: Cube{{2 tables}} + Cube{{2 tables}}");
407        println!("│  Expected:  Element-wise addition across all tables");
408
409        // Create first cube with 2 tables
410        // First, create columns for table 1
411        let t1_arr1 = Array::from_int32(IntegerArray::from_slice(&vec64![1, 2]));
412        let t1_arr2 = Array::from_int32(IntegerArray::from_slice(&vec64![3, 4]));
413        let t1_fa1 = minarrow::FieldArray::from_arr("col1", t1_arr1);
414        let t1_fa2 = minarrow::FieldArray::from_arr("col2", t1_arr2);
415
416        // Create columns for cube1 via constructor
417        let mut cube1 = Cube::new("cube1".to_string(), Some(vec![t1_fa1, t1_fa2]), None);
418
419        // Add second table to cube1
420        let t2_arr1 = Array::from_int32(IntegerArray::from_slice(&vec64![5, 6]));
421        let t2_arr2 = Array::from_int32(IntegerArray::from_slice(&vec64![7, 8]));
422        let t2_fa1 = minarrow::FieldArray::from_arr("col1", t2_arr1);
423        let t2_fa2 = minarrow::FieldArray::from_arr("col2", t2_arr2);
424        let mut table2 = Table::new("t2".to_string(), None);
425        table2.add_col(t2_fa1);
426        table2.add_col(t2_fa2);
427        cube1.add_table(table2);
428
429        // Create second cube
430        let t3_arr1 = Array::from_int32(IntegerArray::from_slice(&vec64![10, 10]));
431        let t3_arr2 = Array::from_int32(IntegerArray::from_slice(&vec64![20, 20]));
432        let t3_fa1 = minarrow::FieldArray::from_arr("col1", t3_arr1);
433        let t3_fa2 = minarrow::FieldArray::from_arr("col2", t3_arr2);
434        let mut cube2 = Cube::new("cube2".to_string(), Some(vec![t3_fa1, t3_fa2]), None);
435
436        let t4_arr1 = Array::from_int32(IntegerArray::from_slice(&vec64![30, 30]));
437        let t4_arr2 = Array::from_int32(IntegerArray::from_slice(&vec64![40, 40]));
438        let t4_fa1 = minarrow::FieldArray::from_arr("col1", t4_arr1);
439        let t4_fa2 = minarrow::FieldArray::from_arr("col2", t4_arr2);
440        let mut table4 = Table::new("t4".to_string(), None);
441        table4.add_col(t4_fa1);
442        table4.add_col(t4_fa2);
443        cube2.add_table(table4);
444
445        match Value::Cube(Arc::new(cube1)) + Value::Cube(Arc::new(cube2)) {
446            Ok(Value::Cube(result)) => {
447                println!("│  Result cube with {} tables:", result.n_tables());
448                for i in 0..result.n_tables() {
449                    println!("│  Table {}:", i);
450                    if let Some(table) = result.table(i) {
451                        for j in 0..table.n_cols() {
452                            let col = &table.cols[j];
453                            if let Array::NumericArray(NumericArray::Int32(arr)) = &col.array {
454                                println!("│    Column {}: {:?}", j, arr.data.as_slice());
455                            }
456                        }
457                    }
458                }
459                println!("└─ ✓ Passed\n");
460            }
461            Ok(other) => println!("└─ ✗ Error: Unexpected result type {:?}\n", other),
462            Err(e) => println!("└─ ✗ Error: {:?}\n", e),
463        }
464    }
465
466    #[cfg(not(feature = "cube"))]
467    {
468        println!("┌─ Test 12: Cube (3D) Broadcasting");
469        println!("└─ ⊘ Skipped (cube feature not enabled)\n");
470    }
471}