From ffc87f16731daa1637e23c333932411285f0f6f9 Mon Sep 17 00:00:00 2001
From: Anshul Singhvi <anshulsinghvi@gmail.com>
Date: Sat, 21 Sep 2024 10:29:09 -0700
Subject: [PATCH 1/2] Make flatten work on tables too

Add a test for flatten

Fix eltype from table
---
 src/primitives.jl  |  9 +++++++--
 test/primitives.jl | 19 +++++++++++++++++++
 2 files changed, 26 insertions(+), 2 deletions(-)

diff --git a/src/primitives.jl b/src/primitives.jl
index a100b5cb8..2f81b15ca 100644
--- a/src/primitives.jl
+++ b/src/primitives.jl
@@ -473,8 +473,13 @@ flatten(f, ::Type{Target}, geom) where {Target<:GI.AbstractTrait} = _flatten(f,
 
 _flatten(f, ::Type{Target}, geom) where Target = _flatten(f, Target, GI.trait(geom), geom)
 # Try to flatten over iterables
-_flatten(f, ::Type{Target}, ::Nothing, iterable) where Target =
-    Iterators.flatten(Iterators.map(x -> _flatten(f, Target, x), iterable))
+function _flatten(f, ::Type{Target}, ::Nothing, iterable) where Target
+    if Tables.istable(iterable)
+        Iterators.flatten(Iterators.map(x -> _flatten(f, Target, x), Tables.getcolumn(iterable, first(GI.geometrycolumns(iterable)))))
+    else
+        Iterators.flatten(Iterators.map(x -> _flatten(f, Target, x), iterable))
+    end
+end
 # Flatten feature collections
 function _flatten(f, ::Type{Target}, ::GI.FeatureCollectionTrait, fc) where Target
     Iterators.map(GI.getfeature(fc)) do feature
diff --git a/test/primitives.jl b/test/primitives.jl
index cf57cc2a1..5fe3f8890 100644
--- a/test/primitives.jl
+++ b/test/primitives.jl
@@ -92,6 +92,25 @@ end
     @test GO._tuple_point.(GO.flatten(GI.PointTrait, very_wrapped)) == vcat(pv1, pv2)
     @test collect(GO.flatten(GI.AbstractCurveTrait, [poly])) == [lr1, lr2]
     @test collect(GO.flatten(GI.x, GI.PointTrait, very_wrapped)) == first.(vcat(pv1, pv2))
+    @testset "flatten with tables" begin
+        # Construct a simple table with a geometry column
+        geom_column = [GI.Point(1.0,1.0), GI.Point(2.0,2.0), GI.Point(3.0,3.0)]
+        table = (geometry = geom_column, id = [1, 2, 3])
+        
+        # Test flatten on the table
+        flattened = collect(GO.flatten(GI.PointTrait, table))
+        
+        @test length(flattened) == 3
+        @test all(p isa GI.Point for p in flattened)
+        @test flattened == geom_column
+        
+        # Test flatten with a function
+        flattened_coords = collect(GO.flatten(p -> (GI.x(p), GI.y(p)), GI.PointTrait, table))
+        
+        @test length(flattened_coords) == 3
+        @test all(c isa Tuple{Float64,Float64} for c in flattened_coords)
+        @test flattened_coords == [(1.0,1.0), (2.0,2.0), (3.0,3.0)]
+    end
 end
 
 @testset "reconstruct" begin

From 1daeb2e959ff917500298c8befcecf04a08d4088 Mon Sep 17 00:00:00 2001
From: Anshul Singhvi <anshulsinghvi@gmail.com>
Date: Sun, 22 Sep 2024 12:50:20 -0700
Subject: [PATCH 2/2] Clarify what's going on

Co-authored-by: Rafael Schouten <rafaelschouten@gmail.com>
---
 src/primitives.jl | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/src/primitives.jl b/src/primitives.jl
index 2f81b15ca..34a315fcd 100644
--- a/src/primitives.jl
+++ b/src/primitives.jl
@@ -475,9 +475,10 @@ _flatten(f, ::Type{Target}, geom) where Target = _flatten(f, Target, GI.trait(ge
 # Try to flatten over iterables
 function _flatten(f, ::Type{Target}, ::Nothing, iterable) where Target
     if Tables.istable(iterable)
-        Iterators.flatten(Iterators.map(x -> _flatten(f, Target, x), Tables.getcolumn(iterable, first(GI.geometrycolumns(iterable)))))
+        column = Tables.getcolumn(iterable, first(GI.geometrycolumns(iterable)))
+        Iterators.map(x -> _flatten(f, Target, x), column) |> Iterators.flatten
     else
-        Iterators.flatten(Iterators.map(x -> _flatten(f, Target, x), iterable))
+        Iterators.map(x -> _flatten(f, Target, x), iterable) |> Iterators.flatten
     end
 end
 # Flatten feature collections