geo_polygonize_core/
arrow_api.rs

1use crate::error::PolygonizeError;
2use crate::types::{Coord3D, Line3D};
3use crate::Polygonizer;
4use arrow::array::{Array, AsArray, GenericListArray};
5use arrow::datatypes::{DataType, Field, Float64Type};
6use geo_traits::to_geo::ToGeoLineString;
7use geoarrow::array::{GeoArrowArray, GeoArrowArrayAccessor, LineStringArray, PolygonBuilder};
8use geoarrow::datatypes::{Dimension, GeoArrowType, Metadata, PolygonType};
9use std::convert::TryFrom;
10use std::sync::Arc;
11
12pub use crate::options::PolygonizerOptions;
13
14pub fn polygonize_arrow(
15    array: &dyn Array,
16    field: &Field,
17    options: PolygonizerOptions,
18) -> Result<geoarrow::array::PolygonArray, PolygonizeError> {
19    let mut polygonizer = Polygonizer::with_options(options);
20    let metadata = match GeoArrowType::from_extension_field(field)
21        .map_err(|e| PolygonizeError::ArrowError(e.to_string()))?
22    {
23        Some(GeoArrowType::LineString(typ)) => {
24            if typ.dimension() != Dimension::XY {
25                return Err(PolygonizeError::UnsupportedOptionCombination {
26                    reason: format!(
27                        "GeoArrow polygonization currently supports XY coordinates only, got {:?}",
28                        typ.dimension()
29                    ),
30                });
31            }
32            typ.metadata().clone()
33        }
34        Some(other) => {
35            return Err(PolygonizeError::InvalidArgumentType {
36                field: field.name().to_string(),
37                expected: "geoarrow.linestring".to_string(),
38                actual: format!("{other:?}"),
39            });
40        }
41        None => Arc::new(
42            Metadata::try_from(field).map_err(|e| PolygonizeError::ArrowError(e.to_string()))?,
43        ),
44    };
45
46    if let Some(edges) = metadata.edges() {
47        return Err(PolygonizeError::UnsupportedOptionCombination {
48            reason: format!("GeoArrow polygonization supports planar edges only, got {edges:?}"),
49        });
50    }
51
52    let mut lines = Vec::new();
53
54    match LineStringArray::try_from((array, field)) {
55        Ok(arr) => process_linestring_array(&arr, &mut lines)?,
56        Err(error)
57            if field
58                .extension_type_name()
59                .is_some_and(|name| name != "ogc.geoarrow.linestring") =>
60        {
61            return Err(PolygonizeError::ArrowError(error.to_string()));
62        }
63        Err(_) => match array.data_type() {
64            DataType::List(_) => process_list_array(array.as_list::<i32>(), &mut lines)?,
65            DataType::LargeList(_) => process_list_array(array.as_list::<i64>(), &mut lines)?,
66            _ => {
67                return Err(PolygonizeError::ArrowError(format!(
68                    "Failed to convert input array to LineStringArray and fallback failed. DataType: {:?}, Field: {:?}.",
69                    array.data_type(),
70                    field
71                )));
72            }
73        },
74    }
75
76    let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
77        polygonizer.add_lines(lines);
78        polygonizer.polygonize()
79    }))
80    .unwrap_or_else(|_| {
81        Err(PolygonizeError::Panic(
82            "Panic occurred in Rust core".to_string(),
83        ))
84    })
85    .map_err(|e| PolygonizeError::TopologyFailure {
86        reason: format!("Polygonization error: {:?}", e),
87    })?;
88
89    let geo_polygons: Vec<geo::Polygon> = result
90        .polygons
91        .into_iter()
92        .map(|p| {
93            let exterior = geo::LineString::from(
94                p.exterior
95                    .into_iter()
96                    .map(|c| (c.x, c.y))
97                    .collect::<Vec<_>>(),
98            );
99            let interiors = p
100                .interiors
101                .into_iter()
102                .map(|ring: Vec<Coord3D>| {
103                    geo::LineString::from(ring.into_iter().map(|c| (c.x, c.y)).collect::<Vec<_>>())
104                })
105                .collect();
106            geo::Polygon::new(exterior, interiors)
107        })
108        .collect();
109
110    let mut builder = PolygonBuilder::new(PolygonType::new(Dimension::XY, metadata));
111    for poly in geo_polygons {
112        builder
113            .push_polygon(Some(&poly))
114            .map_err(|e| PolygonizeError::ArrowError(format!("Failed to push polygon: {}", e)))?;
115    }
116    Ok(builder.finish())
117}
118
119fn process_linestring_array(
120    arr: &LineStringArray,
121    lines: &mut Vec<Line3D>,
122) -> Result<(), PolygonizeError> {
123    for i in 0..arr.len() {
124        if let Ok(Some(geom)) = arr.get(i) {
125            let ls = geom.to_line_string();
126            for line in ls.lines() {
127                if !line.start.x.is_finite()
128                    || !line.start.y.is_finite()
129                    || !line.end.x.is_finite()
130                    || !line.end.y.is_finite()
131                {
132                    return Err(PolygonizeError::InvalidGeometry {
133                        reason: "NaN or Inf coordinates detected in LineStringArray".to_string(),
134                    });
135                }
136                let p1 = Coord3D::new(line.start.x, line.start.y, 0.0);
137                let p2 = Coord3D::new(line.end.x, line.end.y, 0.0);
138                lines.push(Line3D::new(p1, p2, 0));
139            }
140        }
141    }
142    Ok(())
143}
144
145// Manual fallback for GenericListArray<Offset>
146fn process_list_array<O: arrow::array::OffsetSizeTrait>(
147    list_arr: &GenericListArray<O>,
148    lines: &mut Vec<Line3D>,
149) -> Result<(), PolygonizeError> {
150    // Values should be StructArray with x, y
151    let values = list_arr.values();
152    let struct_arr = values
153        .as_struct_opt()
154        .ok_or_else(|| PolygonizeError::InvalidBufferShape {
155            reason: "List values must be Struct".to_string(),
156        })?;
157
158    if struct_arr.num_columns() != 2 {
159        return Err(PolygonizeError::UnsupportedOptionCombination {
160            reason: format!(
161                "GeoArrow polygonization currently supports XY coordinates only, got {} coordinate columns",
162                struct_arr.num_columns()
163            ),
164        });
165    }
166
167    // Get x and y columns
168    // We assume field names "x" and "y" or indices 0 and 1
169    // GeoArrow spec uses "x", "y".
170    let x_arr = struct_arr
171        .column_by_name("x")
172        .or_else(|| {
173            if struct_arr.num_columns() > 0 {
174                Some(struct_arr.column(0))
175            } else {
176                None
177            }
178        })
179        .ok_or_else(|| PolygonizeError::InvalidBufferShape {
180            reason: "Struct missing 'x' column".to_string(),
181        })?;
182
183    let y_arr = struct_arr
184        .column_by_name("y")
185        .or_else(|| {
186            if struct_arr.num_columns() > 1 {
187                Some(struct_arr.column(1))
188            } else {
189                None
190            }
191        })
192        .ok_or_else(|| PolygonizeError::InvalidBufferShape {
193            reason: "Struct missing 'y' column".to_string(),
194        })?;
195
196    let x_vals = x_arr.as_primitive_opt::<Float64Type>().ok_or_else(|| {
197        PolygonizeError::InvalidArgumentType {
198            field: "x".to_string(),
199            expected: "Float64".to_string(),
200            actual: format!("{:?}", x_arr.data_type()),
201        }
202    })?;
203    let y_vals = y_arr.as_primitive_opt::<Float64Type>().ok_or_else(|| {
204        PolygonizeError::InvalidArgumentType {
205            field: "y".to_string(),
206            expected: "Float64".to_string(),
207            actual: format!("{:?}", y_arr.data_type()),
208        }
209    })?;
210
211    for i in 0..list_arr.len() {
212        if list_arr.is_null(i) {
213            continue;
214        }
215        let start = list_arr.value_offsets()[i].to_usize().ok_or_else(|| {
216            PolygonizeError::InvalidBufferShape {
217                reason: "Invalid start offset: cannot convert to usize".to_string(),
218            }
219        })?;
220        let end = list_arr.value_offsets()[i + 1].to_usize().ok_or_else(|| {
221            PolygonizeError::InvalidBufferShape {
222                reason: "Invalid end offset: cannot convert to usize".to_string(),
223            }
224        })?;
225
226        if end <= start {
227            continue;
228        }
229
230        if end > x_vals.len() || end > y_vals.len() {
231            return Err(PolygonizeError::InvalidBufferShape {
232                reason: "Offset out of bounds for x/y coordinate arrays".to_string(),
233            });
234        }
235
236        // Iterate points in the linestring
237        for j in start..end - 1 {
238            let x1 = x_vals.value(j);
239            let y1 = y_vals.value(j);
240            let x2 = x_vals.value(j + 1);
241            let y2 = y_vals.value(j + 1);
242
243            if !x1.is_finite() || !y1.is_finite() || !x2.is_finite() || !y2.is_finite() {
244                return Err(PolygonizeError::InvalidGeometry {
245                    reason: "NaN or Inf coordinates detected in list array".to_string(),
246                });
247            }
248
249            let p1 = Coord3D::new(x1, y1, 0.0);
250            let p2 = Coord3D::new(x2, y2, 0.0);
251            lines.push(Line3D::new(p1, p2, 0));
252        }
253    }
254
255    Ok(())
256}
geo_polygonize_core/arrow_api.rs

geo_polygonize_core/
arrow_api.rs
