Merge pull request #59 from davidanthoff/get_columns_copy_using_missing

Support TableTraits.get_columns_copy_using_missing

Author: malmaud

Project.toml

@@ -19,6 +19,6 @@ DataValues = "≥ 0.4.4"
 IteratorInterfaceExtensions = "^0.1.1, ^1"
 Lazy = "≥ 0.11.4"
 PyCall = "≥ 1.90.0"
-TableTraits = "^0.3.1, ^0.4, ^1"
+TableTraits = "^0.4, ^1"
 TableTraitsUtils = "≥ 0.3.0"
 julia = "≥ 0.7.0"

src/tabletraits.jl

@@ -14,6 +14,16 @@ function TableTraits.getiterator(df::DataFrame)
     return create_tableiterator(column_data, col_names)
 end
 
+TableTraits.supports_get_columns_copy_using_missing(df::DataFrame) = true
+
+function TableTraits.get_columns_copy_using_missing(df::Pandas.DataFrame)
+    # return a named tuple of columns here
+    col_names_raw = [i for i in Pandas.columns(df)]
+    col_names = Symbol.(col_names_raw)
+    cols = (Array(eltype(df[i])==String ? [df[i][j] for j=1:length(df)] : df[i]) for i in col_names_raw)
+    return NamedTuple{tuple(col_names...)}(tuple(cols...))
+end
+
 function _construct_pandas_from_iterabletable(source)
     y = create_columns_from_iterabletable(source, errorhandling=:returnvalue)
     y===nothing && return nothing

test/test_tabletraits.jl

@@ -28,6 +28,10 @@ it_collected = collect(it)
 @test it_collected[1] == (a=1, b="John", c=3.2)
 @test it_collected[2] == (a=2, b="Sally", c=5.8)
 
+@test TableTraits.supports_get_columns_copy_using_missing(df) == true
+cols = TableTraits.get_columns_copy_using_missing(df)
+@test cols == (a=[1,2], b=["John", "Sally"], c=[3.2, 5.8])
+
 table_array2 = [(a=1, b=DataValue("John"), c=3.2), (a=2, b=DataValue("Sally"), c=5.8)]
 
 @test_throws ArgumentError DataFrame(table_array2)
@@ -46,4 +50,7 @@ it3_collected = collect(IteratorInterfaceExtensions.getiterator(df3))
 @test it3_collected[2].b == "Sally"
 @test isnan(it3_collected[2].c)
 
+cols3 = TableTraits.get_columns_copy_using_missing(df3)
+@test isequal(cols3, (a=[NaN,2.], b=["John", "Sally"], c=[3.2, NaN]))
+
 end