Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions a5_rust/Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 1 addition & 1 deletion a5_rust/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -9,4 +9,4 @@ name = "a5_rust"
crate-type = ["staticlib"]

[dependencies]
a5 = "0.8.0"
a5 = "0.9.0"
61 changes: 27 additions & 34 deletions a5_rust/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -14,13 +14,6 @@ pub struct ResultLonLat {
pub error: *mut std::os::raw::c_char, // null if no error
}

#[repr(C)]
pub struct ResultSpherical {
pub theta: f64,
pub phi: f64,
pub error: *mut std::os::raw::c_char,
}

#[repr(C)]
pub struct CellBoundaryOptions {
pub closed_ring: bool,
Expand Down Expand Up @@ -251,17 +244,6 @@ pub extern "C" fn a5_get_num_children(parent_res: i32, child_res: i32) -> usize
a5::get_num_children(parent_res, child_res)
}

#[no_mangle]
pub extern "C" fn a5_cell_to_spherical(cell: u64) -> ResultSpherical {
match a5::cell_to_spherical(cell) {
Ok(sph) => ResultSpherical { theta: sph.theta.get(), phi: sph.phi.get(), error: std::ptr::null_mut() },
Err(e) => {
let err_msg = CString::new(e.to_string()).unwrap();
ResultSpherical { theta: 0.0, phi: 0.0, error: err_msg.into_raw() }
}
}
}

#[no_mangle]
pub extern "C" fn a5_spherical_cap(cell_id: u64, radius: f64) -> CellArray {
cell_vec_result_to_c(a5::spherical_cap(cell_id, radius))
Expand Down Expand Up @@ -294,18 +276,6 @@ pub extern "C" fn a5_is_valid_cell(index: u64) -> bool {
}
}

#[no_mangle]
pub extern "C" fn a5_spherical_to_cell(theta: f64, phi: f64, resolution: i32) -> ResultU64 {
let spherical = a5::coordinate_systems::Spherical::new(a5::Radians::new(theta), a5::Radians::new(phi));
match a5::spherical_to_cell(spherical, resolution) {
Ok(cell) => ResultU64 { value: cell, error: std::ptr::null_mut() },
Err(e) => {
let err_msg = CString::new(e.to_string()).unwrap();
ResultU64 { value: 0, error: err_msg.into_raw() }
}
}
}

// Build a Vec<a5::LonLat> from a C array of LonLatDegrees (lon, lat) input points.
fn lonlat_slice_to_vec(points: *const LonLatDegrees, len: usize) -> Vec<a5::LonLat> {
let slice = unsafe { std::slice::from_raw_parts(points, len) };
Expand All @@ -321,12 +291,35 @@ pub extern "C" fn a5_line_string_to_cells(points: *const LonLatDegrees, len: usi
cell_vec_result_to_c(a5::line_string_to_cells(&lonlats, resolution))
}

// GeoJSON-style polygon: ring 0 is the outer ring, rings 1.. are holes. The
// rings are passed flattened into a single `points` buffer, with `ring_lengths`
// giving the vertex count of each ring (so the offsets can be reconstructed).
// Holes are excluded by the a5 crate itself - the caller does no hole handling.
#[no_mangle]
pub extern "C" fn a5_polygon_to_cells(points: *const LonLatDegrees, len: usize, resolution: i32) -> CellArray {
if points.is_null() || len == 0 {
pub extern "C" fn a5_polygon_to_cells(
points: *const LonLatDegrees,
ring_lengths: *const usize,
ring_count: usize,
resolution: i32,
) -> CellArray {
if points.is_null() || ring_lengths.is_null() || ring_count == 0 {
return CellArray { data: std::ptr::null_mut(), len: 0, error: std::ptr::null_mut() };
}
let lonlats = lonlat_slice_to_vec(points, len);
cell_vec_result_to_c(a5::polygon_to_cells(&lonlats, resolution))
let lengths = unsafe { std::slice::from_raw_parts(ring_lengths, ring_count) };
let total: usize = lengths.iter().sum();
let flat = unsafe { std::slice::from_raw_parts(points, total) };

let mut rings: Vec<Vec<a5::LonLat>> = Vec::with_capacity(ring_count);
let mut offset = 0;
for &len in lengths {
rings.push(
flat[offset..offset + len]
.iter()
.map(|p| a5::LonLat::new(p.lon, p.lat))
.collect(),
);
offset += len;
}
cell_vec_result_to_c(a5::polygon_to_cells(&rings, resolution))
}

34 changes: 0 additions & 34 deletions docs/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -155,8 +155,6 @@ Visualizing that A5 cell shows:
|----------|---------|-------------|
| `a5_lonlat_to_cell(lon, lat, res)` | `UBIGINT` | Cell containing a coordinate |
| `a5_cell_to_lonlat(cell)` | `DOUBLE[2]` | Cell center `[lon, lat]` |
| `a5_cell_to_spherical(cell)` | `DOUBLE[2]` | Cell center `[theta, phi]` (radians) |
| `a5_spherical_to_cell(theta, phi, res)` | `UBIGINT` | Cell from spherical coords (inverse of above) |
| `a5_cell_to_boundary(cell [, closed, segments])` | `DOUBLE[2][]` | Boundary vertices |
| `a5_cell_area(res)` | `DOUBLE` | Cell area (m²) at a resolution |
| `a5_get_resolution(cell)` | `INTEGER` | Resolution of a cell |
Expand Down Expand Up @@ -347,38 +345,6 @@ SELECT unnest(a5_cell_to_boundary(207618739568, false, 5)) as boundary_points;
└───────────────────────────────────────────┘
```



#### `a5_cell_to_spherical(cell_id) -> DOUBLE[2]`

Returns the spherical coordinates [theta, phi] in radians of an A5 cell center, where theta is the azimuthal angle and phi is the polar angle.

**Example:**
```sql
SELECT a5_cell_to_spherical(a5_lonlat_to_cell(-74.0060, 40.7128, 15)) as spherical_coords;
```

#### `a5_spherical_to_cell(theta, phi, resolution) -> UBIGINT`

Returns the A5 cell at the given resolution containing the spherical coordinates [theta, phi] (in radians). This is the inverse of `a5_cell_to_spherical`.

**Parameters:**

- `theta` (DOUBLE): Azimuthal angle in radians
- `phi` (DOUBLE): Polar angle in radians
- `resolution` (INTEGER): Resolution level (0-30)

**Example:**
```sql
SELECT a5_spherical_to_cell(-0.512679, 0.913528, 10) as cell;
┌─────────────────────┐
│ cell │
│ uint64 │
├─────────────────────┤
│ 1937278465245970432 │
└─────────────────────┘
```

### Region Functions

#### `a5_geometry_to_cells(geom, resolution) -> UBIGINT[]`
Expand Down
152 changes: 18 additions & 134 deletions src/a5_extension.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@
namespace duckdb {

#define MAX_RESOLUTION 30
#define A5_EXTENSION_VERSION "2026060904"
#define A5_EXTENSION_VERSION "2026061701"

// Helper function to validate resolution and throw with a clear error message
inline void ValidateResolution(int32_t resolution, const char *function_name) {
Expand Down Expand Up @@ -369,35 +369,6 @@ inline void A5GetNumChildrenFun(DataChunk &args, ExpressionState &state, Vector
});
}

inline void A5CellToSphericalFun(DataChunk &args, ExpressionState &state, Vector &result) {
auto &cell_vector = args.data[0];

auto &result_data_children = ArrayVector::GetEntry(result);
double *data_ptr = FlatVector::GetData<double>(result_data_children);

UnifiedVectorFormat cell_id_format;
cell_vector.ToUnifiedFormat(args.size(), cell_id_format);
uint64_t *input_data_ptr = FlatVector::GetData<uint64_t>(cell_vector);

for (idx_t i = 0; i < args.size(); i++) {
auto cell_idx = cell_id_format.sel->get_index(i);
if (!cell_id_format.validity.RowIsValid(cell_idx)) {
FlatVector::SetNull(result, i, true);
continue;
}

struct ResultSpherical res = a5_cell_to_spherical(input_data_ptr[cell_idx]);
ThrowRustError(res.error, "a5_cell_to_spherical");

data_ptr[i * 2] = res.theta;
data_ptr[i * 2 + 1] = res.phi;
}

if (args.size() == 1) {
result.SetVectorType(VectorType::CONSTANT_VECTOR);
}
}

inline void A5SphericalCapFun(DataChunk &args, ExpressionState &state, Vector &result) {
ListVector::Reserve(result, args.size() * 4);
uint64_t offset = 0;
Expand Down Expand Up @@ -492,21 +463,6 @@ inline void A5IsValidCellFun(DataChunk &args, ExpressionState &state, Vector &re
[&](uint64_t cell) { return a5_is_valid_cell(cell); });
}

inline void A5SphericalToCellFun(DataChunk &args, ExpressionState &state, Vector &result) {
auto &theta_vector = args.data[0];
auto &phi_vector = args.data[1];
auto &resolution_vector = args.data[2];

TernaryExecutor::Execute<double, double, int32_t, uint64_t>(
theta_vector, phi_vector, resolution_vector, result, args.size(),
[&](double theta, double phi, int32_t resolution) {
ValidateResolution(resolution, "a5_spherical_to_cell");
struct ResultU64 res = a5_spherical_to_cell(theta, phi, resolution);
ThrowRustError(res.error, "a5_spherical_to_cell");
return res.value;
});
}

// ---------------------------------------------------------------------------
// GEOMETRY (WKB) writers
//
Expand Down Expand Up @@ -641,74 +597,34 @@ static vector<LonLatDegrees> WkbReadRing(WkbCursor &cur, idx_t dims) {

// Fill a polygon (outer ring minus any holes) into the accumulator.
//
// A5's polygon_to_cells returns a *compacted* (mixed-resolution) covering, so the outer
// covering and a hole's covering generally share no cell IDs and cannot be differenced
// directly. To subtract holes we uncompact both to the target resolution, take the set
// difference at that uniform resolution, then re-compact for output. The (common) no-hole
// case skips all of this and passes the crate's compacted covering through unchanged.
// Holes are excluded by the a5 crate itself: we flatten all rings (outer first, then
// holes) into a single point buffer plus a per-ring length array and hand them to
// a5_polygon_to_cells, which returns the compacted covering of the outer ring with the
// holes already removed. Empty rings are dropped so a degenerate ring never shifts the
// outer-ring-is-first convention.
static void PolygonRingsToCells(const vector<vector<LonLatDegrees>> &rings, int32_t resolution, CellAccumulator &acc,
const char *function_name) {
if (rings.empty() || rings[0].empty()) {
return;
}
auto outer = a5_polygon_to_cells(rings[0].data(), rings[0].size(), resolution);
ThrowCellArrayError(outer, function_name);

bool has_holes = false;
for (size_t r = 1; r < rings.size(); r++) {
if (!rings[r].empty()) {
has_holes = true;
break;
}
}
if (!has_holes) {
for (size_t i = 0; i < outer.len; i++) {
acc.Add(outer.data[i]);
}
a5_free_cell_array(outer);
return;
}

// Expand the outer covering to a uniform resolution.
auto outer_uniform = a5_uncompact(outer.data, outer.len, resolution);
ThrowCellArrayError(outer_uniform, function_name);
a5_free_cell_array(outer);

// Collect the uniform-resolution cells of every hole.
std::unordered_set<uint64_t> holes;
for (size_t r = 1; r < rings.size(); r++) {
if (rings[r].empty()) {
vector<LonLatDegrees> points;
vector<uintptr_t> ring_lengths;
ring_lengths.reserve(rings.size());
for (const auto &ring : rings) {
if (ring.empty()) {
continue;
}
auto hole = a5_polygon_to_cells(rings[r].data(), rings[r].size(), resolution);
ThrowCellArrayError(hole, function_name);
auto hole_uniform = a5_uncompact(hole.data, hole.len, resolution);
ThrowCellArrayError(hole_uniform, function_name);
a5_free_cell_array(hole);
for (size_t i = 0; i < hole_uniform.len; i++) {
holes.insert(hole_uniform.data[i]);
}
a5_free_cell_array(hole_uniform);
ring_lengths.push_back(ring.size());
points.insert(points.end(), ring.begin(), ring.end());
}

// Difference, then re-compact so the output matches the no-hole convention.
vector<uint64_t> kept;
kept.reserve(outer_uniform.len);
for (size_t i = 0; i < outer_uniform.len; i++) {
if (holes.find(outer_uniform.data[i]) == holes.end()) {
kept.push_back(outer_uniform.data[i]);
}
}
a5_free_cell_array(outer_uniform);
if (kept.empty()) {
return;
}
auto compacted = a5_compact(kept.data(), kept.size());
ThrowCellArrayError(compacted, function_name);
for (size_t i = 0; i < compacted.len; i++) {
acc.Add(compacted.data[i]);
auto cells = a5_polygon_to_cells(points.data(), ring_lengths.data(), ring_lengths.size(), resolution);
ThrowCellArrayError(cells, function_name);
for (size_t i = 0; i < cells.len; i++) {
acc.Add(cells.data[i]);
}
a5_free_cell_array(compacted);
a5_free_cell_array(cells);
}

// Recursively read a (possibly multi-part) geometry and accumulate its A5 cells.
Expand Down Expand Up @@ -1063,21 +979,6 @@ static void LoadInternal(ExtensionLoader &loader) {
loader.RegisterFunction(std::move(info));
}

// a5_cell_to_spherical: Returns the spherical coordinates of a cell center
{
auto func = ScalarFunction("a5_cell_to_spherical", {LogicalType::UBIGINT},
LogicalType::ARRAY(LogicalType::DOUBLE, 2), A5CellToSphericalFun);
CreateScalarFunctionInfo info(func);
FunctionDescription desc;
desc.description = "Returns the spherical coordinates [theta, phi] in radians of an A5 cell center";
desc.parameter_names = {"cell"};
desc.parameter_types = {LogicalType::UBIGINT};
desc.examples = {"a5_cell_to_spherical(a5_lonlat_to_cell(-122.4, 37.8, 10))"};
desc.categories = {"a5", "geospatial"};
info.descriptions.push_back(std::move(desc));
loader.RegisterFunction(std::move(info));
}

// a5_spherical_cap: Returns cells within a spherical cap radius
{
auto func = ScalarFunction("a5_spherical_cap", {LogicalType::UBIGINT, LogicalType::DOUBLE},
Expand Down Expand Up @@ -1153,23 +1054,6 @@ static void LoadInternal(ExtensionLoader &loader) {
loader.RegisterFunction(std::move(info));
}

// a5_spherical_to_cell: Returns the cell containing the given spherical coordinates
{
auto func =
ScalarFunction("a5_spherical_to_cell", {LogicalType::DOUBLE, LogicalType::DOUBLE, LogicalType::INTEGER},
LogicalType::UBIGINT, A5SphericalToCellFun);
CreateScalarFunctionInfo info(func);
FunctionDescription desc;
desc.description = "Returns the A5 cell at the given resolution containing the spherical coordinates "
"[theta, phi] (in radians); the inverse of a5_cell_to_spherical";
desc.parameter_names = {"theta", "phi", "resolution"};
desc.parameter_types = {LogicalType::DOUBLE, LogicalType::DOUBLE, LogicalType::INTEGER};
desc.examples = {"a5_spherical_to_cell(2.14, 0.92, 10)"};
desc.categories = {"a5", "geospatial"};
info.descriptions.push_back(std::move(desc));
loader.RegisterFunction(std::move(info));
}

// a5_geometry_to_cells: Returns the cells covering any geometry
{
auto func = ScalarFunction("a5_geometry_to_cells", {LogicalType::GEOMETRY(), LogicalType::INTEGER},
Expand Down
Loading
Loading