add a DataFrames extension for better table reconstruction (#312)

Author: asinghvi17

GeometryOpsCore/src/apply.jl

@@ -175,9 +175,11 @@ with the same schema, but with the new geometry column.
 This new table may be of the same type as the old one iff `Tables.materializer` is defined for 
 that table.  If not, then a `NamedTuple` is returned.
 =#
-function _apply_table(f::F, target, iterable::IterableType; geometrycolumn = nothing, preserve_default_metadata = false, threaded, kw...) where {F, IterableType}
-    _get_col_pair(colname) = colname => Tables.getcolumn(iterable, colname)
-    # We extract the geometry column and run `apply` on it.
+function _apply_table(f::F, target, iterable::IterableType; geometrycolumn = nothing, preserve_default_metadata = false, threaded, kw...) where {F, IterableType}    # We extract the geometry column and run `apply` on it.
+    # First, we need the table schema:
+    input_schema = Tables.schema(iterable)
+    input_colnames = input_schema.names
+    # then, we find the geometry column(s)
     geometry_columns = if isnothing(geometrycolumn)
         GI.geometrycolumns(iterable)
     elseif geometrycolumn isa NTuple{N, <: Symbol} where N
@@ -187,31 +189,31 @@ function _apply_table(f::F, target, iterable::IterableType; geometrycolumn = not
     else
         throw(ArgumentError("geometrycolumn must be a Symbol or a tuple of Symbols, got a $(typeof(geometrycolumn))"))
     end
-    if !all(Base.Fix2(in, Tables.columnnames(iterable)), geometry_columns)
+    if !Base.issubset(geometry_columns, input_colnames)
         throw(ArgumentError(
             """
             `apply`: the `geometrycolumn` kwarg must be a subset of the column names of the table, 
             got $(geometry_columns)
             but the table has columns 
-            $(Tables.columnnames(iterable))
+            $(input_colnames)
             """
             ))
     end
+    # here we apply the function to the geometry column(s).
+    apply_kw = if isempty(used_reconstruct_table_kwargs(iterable))
+        kw
+    else
+        Base.structdiff(values(kw), NamedTuple{used_reconstruct_table_kwargs(iterable)})
+    end
     new_geometry_vecs = map(geometry_columns) do colname
-        _apply(f, target, Tables.getcolumn(iterable, colname); threaded, kw...)
+        _apply(f, target, Tables.getcolumn(iterable, colname); threaded, apply_kw...)
     end
-    # Then, we obtain the schema of the table,
-    old_schema = Tables.schema(iterable)
-    # filter the geometry column out,
-    new_names = filter(x -> !(x in geometry_columns), old_schema.names)
+    # Then, we filter the geometry column(s) out,
+    new_names = filter(x -> !(x in geometry_columns), input_colnames)
     # and try to rebuild the same table as the best type - either the original type of `iterable`,
     # or a named tuple which is the default fallback.
-    result = Tables.materializer(iterable)(
-        merge(
-            NamedTuple{geometry_columns, Base.Tuple{typeof.(new_geometry_vecs)...}}(new_geometry_vecs),
-            NamedTuple(Iterators.map(_get_col_pair, new_names))
-        )
-    )
+    # See the function directly below this one for the actual fallback implementation.
+    result = reconstruct_table(iterable, geometry_columns, new_geometry_vecs, new_names; kw...)
     # Finally, we ensure that metadata is propagated correctly.
     # This can only happen if the original table supports metadata reads,
     # and the result supports metadata writes.
@@ -246,6 +248,51 @@ function _apply_table(f::F, target, iterable::IterableType; geometrycolumn = not
     return result
 end
 
+
+"""
+    used_reconstruct_table_kwargs(input)
+
+Return a tuple of the kwargs that should be passed to `reconstruct_table` for the given input.
+
+This is "semi-public" API, and required for any input type that defines `reconstruct_table`.
+"""
+function used_reconstruct_table_kwargs(input) 
+    ()
+end
+
+"""
+    reconstruct_table(input, geometry_column_names, geometry_columns, other_column_names, args...; kwargs...)
+
+Reconstruct a table from the given input, geometry column names,
+geometry columns, and other column names.
+
+Any function that defines `reconstruct_table` must also define `used_reconstruct_table_kwargs`.
+
+The input must be a table.
+
+The function should return a best-effort attempt at a table of the same type as the input,
+with the new geometry column(s) and other columns.
+
+The fallback implementation invokes `Tables.materializer`.  But if you want to be efficient 
+and pass e.g. arbitrary kwargs to the materializer, or materialize in a different way, you 
+can do so by overloading this function for your desired input type.
+
+This is "semi-public" API and while it may add optional arguments, it will not add new required 
+positional arguments. All implementations must allow arbitrary kwargs to pass through and harvest 
+what they need.
+"""
+function reconstruct_table(input, geometry_column_names, geometry_columns, other_column_names, args...; kwargs...)
+    @assert Tables.istable(input)
+    _get_col_pair(colname) = colname => Tables.getcolumn(input, colname)
+
+    return Tables.materializer(input)(
+        merge(
+            NamedTuple{geometry_column_names, Base.Tuple{typeof.(geometry_columns)...}}(geometry_columns),
+            NamedTuple(Iterators.map(_get_col_pair, other_column_names))
+        )
+    )
+end
+
 # Rewrap all FeatureCollectionTrait feature collections as GI.FeatureCollection
 # Maybe use threads to call _apply on component features
 @inline function _apply(f::F, target, ::GI.FeatureCollectionTrait, fc;

Project.toml

@@ -21,12 +21,14 @@ Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 Tables = "bd369af6-aec1-5ad0-b16a-f7cc5008161c"
 
 [weakdeps]
+DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
 FlexiJoins = "e37f2e79-19fa-4eb7-8510-b63b51fe0a37"
 LibGEOS = "a90b1aa1-3769-5649-ba7e-abc5a9d163eb"
 Proj = "c94c279d-25a6-4763-9509-64d165bea63e"
 TGGeometry = "d7e755d2-3c95-4bcf-9b3c-79ab1a78647b"
 
 [extensions]
+GeometryOpsDataFramesExt = "DataFrames"
 GeometryOpsFlexiJoinsExt = "FlexiJoins"
 GeometryOpsLibGEOSExt = "LibGEOS"
 GeometryOpsProjExt = "Proj"
@@ -37,6 +39,7 @@ AbstractTrees = "0.4"
 AdaptivePredicates = "1.2"
 CoordinateTransformations = "0.5, 0.6"
 DataAPI = "1"
+DataFrames = "1"
 DelaunayTriangulation = "1.0.4"
 ExactPredicates = "2.2.8"
 Extents = "0.1.5"

ext/GeometryOpsDataFramesExt/GeometryOpsDataFramesExt.jl

@@ -0,0 +1,45 @@
+#=
+# DataFrames extension
+
+This module simply extends Core's `reconstruct_table` method to:
+- work with DataFrames
+- not copy columns unless it is necessary to do so
+- allow passing through known kwargs to the constructor
+
+
+In the future, if we ever end up defining `ApplyToFeatures` on a table,
+then we will need to add some form of method for that...
+which will likely entail adding an extra positional argument to the
+reconstruct_table method, and checking whether `other_column_names`
+is equal to `setdiff(Tables.columnnames(input), geometry_column_names)`.
+
+If it is not then we will have to reconstruct the whole DataFrame from the 
+GI.Feature named-tuple row table representation.
+=#
+module GeometryOpsDataFramesExt
+
+import GeometryOpsCore
+using DataFrames
+
+GeometryOpsCore.used_reconstruct_table_kwargs(::DataFrames.DataFrame) = (:copycols,)
+
+function GeometryOpsCore.reconstruct_table(
+        input::DataFrames.DataFrame, geometry_column_names, geometry_columns, 
+        other_column_names, args...; 
+        copycols = true, kwargs...
+    )
+    # Create a new dataframe, let the rest be the same
+    new_df = DataFrame(input; copycols)
+
+    # Copy over the geometry columns
+    for (colname, col) in zip(geometry_column_names, geometry_columns)
+        new_df[!, colname] = col
+    end
+
+    # The other columns were already copied over by the constructor,
+    # so we're done.
+    # Metadata is set in the `_apply_table` method.
+    return new_df
+end
+
+end

test/extensions/dataframes.jl

@@ -0,0 +1,15 @@
+using Test
+using DataFrames
+import GeometryOps as GO, GeoInterface as GI
+
+using NaturalEarth
+
+@testset "DataFrames extension: can we use copycols=false?" begin
+    df = DataFrame(naturalearth("admin_0_countries", 110); copycols = true)
+    transformed = GO.transform(identity, df; copycols = true)
+    transformed_lazy = GO.transform(identity, df; copycols = false)
+
+    @test transformed.NAME == transformed_lazy.NAME
+    @test transformed.NAME !== df.NAME # test that they are different arrays
+    @test transformed_lazy.NAME === df.NAME # test that they are the same array
+end

test/runtests.jl

@@ -47,4 +47,5 @@ end
 # Extensions
 @safetestset "FlexiJoins" begin include("extensions/flexijoins.jl") end
 @safetestset "LibGEOS" begin include("extensions/libgeos.jl") end
-@safetestset "TGGeometry" begin include("extensions/tggeometry.jl") end
+@safetestset "TGGeometry" begin include("extensions/tggeometry.jl") end
+@safetestset "DataFrames" begin include("extensions/dataframes.jl") end