address selectcols issue #991

JuliaAI · Jan 5, 2025 · 6b8b899 · 6b8b899
1 parent 82b4e7f
commit 6b8b899
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 4 deletions.
diff --git a/src/interface/data_utils.jl b/src/interface/data_utils.jl
@@ -101,7 +101,7 @@ function MMI.selectcols(::FI, ::Val{:table}, X, c::Union{Symbol, Integer})
     return Tables.getcolumn(cols, c)
 end
 
-function MMI.selectcols(::FI, ::Val{:table}, X, c::Union{Colon, AbstractArray})
+function MMI.selectcols(::FI, ::Val{:table}, X, c)
     if isdataframe(X)
         return X[!, c]
     else
@@ -115,18 +115,28 @@ end
 # utils for `select`*
 
 # project named tuple onto a tuple with only specified `labels` or indices:
-function project(t::NamedTuple, labels::AbstractArray{Symbol})
+function project(t::NamedTuple, labels::Union{AbstractArray{Symbol},NTuple{<:Any,Symbol}})
     return NamedTuple{tuple(labels...)}(t)
 end
 
 project(t::NamedTuple, label::Colon) = t
 project(t::NamedTuple, label::Symbol) = project(t, [label,])
 project(t::NamedTuple, i::Integer) = project(t, [i,])
 
-function project(t::NamedTuple, indices::AbstractArray{<:Integer})
+function project(
+    t::NamedTuple,
+    indices::AbstractArray{<:Integer},
+    )
     return NamedTuple{tuple(keys(t)[indices]...)}(tuple([t[i] for i in indices]...))
 end
 
+function project(
+    t::NamedTuple,
+    indices::NTuple{<:Any,<:Integer},
+    )
+    return NamedTuple{tuple(keys(t)[[indices...]]...)}(tuple([t[i] for i in indices]...))
+end
+
 # utils for selectrows
 typename(X) = split(string(supertype(typeof(X))), '.')[end]
 isdataframe(X) = typename(X) == "AbstractDataFrame"

diff --git a/test/interface/data_utils.jl b/test/interface/data_utils.jl
@@ -149,7 +149,9 @@ end
     s = schema(tt)
     @test nrows(tt) == N
 
-    @test selectcols(tt, 4:6) ==
+    @test selectcols(tt, 4:6) == selectcols(tt, (4, 5, 6)) ==
+        selectcols(tt, (:x4, :x5, :z)) ==
+        selectcols(tt, [:x4, :x5, :z]) ==
         selectcols(TypedTables.Table(x4=tt.x4, x5=tt.x5, z=tt.z), :)
     @test selectcols(tt, [:x1, :z]) ==
         selectcols(TypedTables.Table(x1=tt.x1, z=tt.z), :)