From 5679f6d40151ec8487d0ee4bbccf6149fb52ef32 Mon Sep 17 00:00:00 2001
From: Jacob Quinn <quinn.jacobd@gmail.com>
Date: Wed, 18 Aug 2021 00:24:43 -0600
Subject: [PATCH] Allow select to work on super-wide tables

Started with code suggested by @OkonSamuel in issue #20. Made some
tweaks and added some tests and this seems to work now.
---
 src/TableOperations.jl | 14 ++++++++++++++
 test/runtests.jl       | 26 ++++++++++++++++++++++++++
 2 files changed, 40 insertions(+)

diff --git a/src/TableOperations.jl b/src/TableOperations.jl
index 9ffb3f4..32822b6 100644
--- a/src/TableOperations.jl
+++ b/src/TableOperations.jl
@@ -106,8 +106,22 @@ end
 
 typesubset(::Tables.Schema{names, types}, ::Tuple{}) where {names, types} = Tuple{}
 
+function typesubset(sch::Tables.Schema{nothing, nothing}, nms::NTuple{N, Symbol}) where {N}
+    names = sch.names
+    types = sch.types
+    return Tuple{types[indexin(collect(nms), names)]...}
+end
+
+function typesubset(sch::Tables.Schema{nothing, nothing}, inds::NTuple{N, Int}) where {N}
+    types = sch.types
+    return Tuple{Any[types[i] for i in inds]...}
+end
+
+typesubset(::Tables.Schema{nothing, nothing}, ::Tuple{}) = Tuple{}
+
 namesubset(::Tables.Schema{names, types}, nms::NTuple{N, Symbol}) where {names, types, N} = nms
 Base.@pure namesubset(::Tables.Schema{names, T}, inds::NTuple{N, Int}) where {names, T, N} = ntuple(i -> names[inds[i]], N)
+Base.@pure namesubset(sch::Tables.Schema{nothing, nothing}, inds::NTuple{N, Int}) where {N} = (names = sch.names; ntuple(i -> names[inds[i]], N))
 namesubset(::Tables.Schema{names, types}, ::Tuple{}) where {names, types} = ()
 namesubset(names, nms::NTuple{N, Symbol}) where {N} = nms
 namesubset(names, inds::NTuple{N, Int}) where {N} = ntuple(i -> names[inds[i]], N)
diff --git a/test/runtests.jl b/test/runtests.jl
index e83f82d..31f5add 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -4,6 +4,16 @@ ctable = (A=[1, missing, 3], B=[1.0, 2.0, 3.0], C=["hey", "there", "sailor"])
 rtable = Tables.rowtable(ctable)
 rtable2 = Iterators.filter(i -> i.a % 2 == 0, [(a=x, b=y) for (x, y) in zip(1:20, 21:40)])
 
+struct ReallyWideTable
+end
+Tables.istable(::Type{ReallyWideTable}) = true
+Tables.columnaccess(::Type{ReallyWideTable}) = true
+Tables.columns(x::ReallyWideTable) = x
+Tables.columnnames(x::ReallyWideTable) = [Symbol(:x, i) for i = 1:100_000]
+Tables.getcolumn(x::ReallyWideTable, i::Int) = rand(10)
+Tables.getcolumn(x::ReallyWideTable, nm::Symbol) = rand(10)
+Tables.schema(x::ReallyWideTable) = Tables.Schema(Tables.columnnames(x), [Float64 for _ = 1:100_000])
+
 @testset "TableOperations.transform" begin
 
 tran = ctable |> TableOperations.transform(C=Symbol)
@@ -93,6 +103,22 @@ end
 
 @testset "TableOperations.select" begin
 
+# 20
+x = ReallyWideTable()
+sel = TableOperations.select(x, :x1, :x2)
+sch = Tables.schema(sel)
+@test sch.names == (:x1, :x2)
+@test sch.types == (Float64, Float64)
+tt = Tables.columntable(sel)
+@test tt.x1 isa Vector{Float64}
+
+sel = TableOperations.select(x, 1, 2)
+sch = Tables.schema(sel)
+@test sch.names == (:x1, :x2)
+@test sch.types == (Float64, Float64)
+tt = Tables.columntable(sel)
+@test tt.x1 isa Vector{Float64}
+
 # 117
 sel = TableOperations.select(ctable)
 @test Tables.istable(typeof(sel))