remove "_v2" from stats files

Author: kylebd99

src/TensorStats/TensorStats.jl

@@ -3,6 +3,6 @@
 #include("tensor-stats.jl")
 #include("propagate-stats.jl")
 include("StaticBitset.jl")
-include("tensor-stats_v2.jl")
-include("propagate-stats_v2.jl")
+include("tensor-stats.jl")
+include("propagate-stats.jl")
 include("cost-estimates.jl")

src/TensorStats/propagate-stats.jl

@@ -9,7 +9,7 @@ function merge_tensor_stats_union(op,  all_stats::Vararg{TensorStats})
     throw(error("merge_tensor_stats_union not implemented for: ", typeof(all_stats[1])))
 end
 
-function reduce_tensor_stats(op, reduce_indices::Set{IndexExpr}, stats::TensorStats)
+function reduce_tensor_stats(op, init, reduce_indices::Set{IndexExpr}, stats::TensorStats)
     throw(error("reduce_tensor_stats not implemented for: ", typeof(stats)))
 end
 
@@ -22,48 +22,53 @@ end
 function merge_tensor_def(op, all_defs::Vararg{TensorDef})
     new_default_value = op([def.default_value for def in all_defs]...)
     new_index_set = union([def.index_set for def in all_defs]...)
-    new_dim_sizes = Dict()
+    new_dim_sizes = Dict{IndexExpr, UInt128}()
     for index in new_index_set
         for def in all_defs
             if index in def.index_set
                 new_dim_sizes[index] = def.dim_sizes[index]
             end
         end
     end
-    @assert new_index_set ⊆ keys(new_dim_sizes)
+#    @assert new_index_set ⊆ keys(new_dim_sizes)
     return TensorDef(new_index_set, new_dim_sizes, new_default_value, nothing, nothing, nothing)
 end
 
-function reduce_tensor_def(op, reduce_indices::Set{IndexExpr}, def::TensorDef)
+function reduce_tensor_def(op, init, reduce_indices::Set{IndexExpr}, def::TensorDef)
     op = op isa PlanNode ? op.val : op
-    new_default_value = nothing
-    if isidentity(op, def.default_value) || isidempotent(op)
-        new_default_value = op(def.default_value, def.default_value)
-    elseif op == +
-        new_default_value = def.default_value * prod([def.dim_sizes[x] for x in reduce_indices])
-    elseif op == *
-        new_default_value = def.default_value ^ prod([def.dim_sizes[x] for x in reduce_indices])
-    else
-        # This is going to be VERY SLOW. Should raise a warning about reductions over non-identity default values.
-        # Depending on the semantics of reductions, we might be able to do this faster.
-        println("Warning: A reduction can take place over a tensor whose default value is not the reduction operator's identity. \\
-                         This can result in a large slowdown as the new default is calculated.")
-        new_default_value = op([def.default_value for _ in prod([def.dim_sizes[x] for x in reduce_indices])]...)
+    init = init isa PlanNode ? init.val : init
+    if isnothing(init)
+        if isnothing(op) && isnothing(init)
+            init = def.default_value
+        elseif isidentity(op, def.default_value) || isidempotent(op)
+            init = op(def.default_value, def.default_value)
+        elseif op == +
+            init = def.default_value * prod([def.dim_sizes[x] for x in reduce_indices])
+        elseif op == *
+            init = def.default_value ^ prod([def.dim_sizes[x] for x in reduce_indices])
+        else
+            # This is going to be VERY SLOW. Should raise a warning about reductions over non-identity default values.
+            # Depending on the semantics of reductions, we might be able to do this faster.
+            println("Warning: A reduction can take place over a tensor whose default value is not the reduction operator's identity. \\
+                            This can result in a large slowdown as the new default is calculated.")
+            init = op([def.default_value for _ in prod([def.dim_sizes[x] for x in reduce_indices])]...)
+        end
     end
+    @assert !isnothing(init)
     new_index_set = setdiff(def.index_set, reduce_indices)
-    new_dim_sizes = Dict()
+    new_dim_sizes = Dict{IndexExpr, UInt128}()
     for index in new_index_set
         new_dim_sizes[index] = def.dim_sizes[index]
     end
-    return TensorDef(new_index_set, new_dim_sizes, new_default_value, nothing, nothing, nothing)
+    return TensorDef(new_index_set, new_dim_sizes, init, nothing, nothing, nothing)
 end
 
 # This function determines whether a binary operation is union-like or join-like and creates
 # new statistics objects accordingly.
 function merge_tensor_stats(op, all_stats::Vararg{ST}) where ST <: TensorStats
-    new_def = merge_tensor_def(op, [get_def(stats) for stats in all_stats]...)
-    join_like_args = []
-    union_like_args = []
+    new_def::TensorDef = merge_tensor_def(op, [get_def(stats) for stats in all_stats]...)
+    join_like_args = ST[]
+    union_like_args = ST[]
     for stats in all_stats
         if length(get_index_set(stats)) == 0
             continue
@@ -90,8 +95,8 @@ function merge_tensor_stats(op::PlanNode, all_stats::Vararg{ST}) where ST <:Tens
     return merge_tensor_stats(op.val, all_stats...)
 end
 
-function reduce_tensor_stats(op, reduce_indices::Union{Vector{PlanNode}, Set{PlanNode}}, stats::ST) where ST <:TensorStats
-    return reduce_tensor_stats(op, Set{IndexExpr}([idx.name for idx in reduce_indices]), stats)
+function reduce_tensor_stats(op, init, reduce_indices::Union{Vector{PlanNode}, Set{PlanNode}}, stats::ST) where ST <:TensorStats
+    return reduce_tensor_stats(op, init, Set{IndexExpr}([idx.name for idx in reduce_indices]), stats)
 end
 
 function transpose_tensor_def(index_order::Vector{IndexExpr}, def::TensorDef)
@@ -101,25 +106,25 @@ end
 
 ################# NaiveStats Propagation ##################################################
  # We do everything in log for numerical stability
-function merge_tensor_stats_join(op, new_def, all_stats::Vararg{NaiveStats})
+function merge_tensor_stats_join(op, new_def::TensorDef, all_stats::Vararg{NaiveStats})
     new_dim_space_size = sum([log2(get_dim_size(new_def, idx)) for idx in new_def.index_set])
     prob_non_default = sum([log2(stats.cardinality) - sum([log2(get_dim_size(stats, idx)) for idx in get_index_set(stats)]) for stats in all_stats])
     new_cardinality = 2^(prob_non_default + new_dim_space_size)
     return NaiveStats(new_def, new_cardinality)
 end
 
-function merge_tensor_stats_union(op, new_def, all_stats::Vararg{NaiveStats})
+function merge_tensor_stats_union(op, new_def::TensorDef, all_stats::Vararg{NaiveStats})
     new_dim_space_size = sum([log2(get_dim_size(new_def, idx)) for idx in new_def.index_set])
     prob_default = sum([log2(1 - 2^(log2(stats.cardinality) - sum([log2(get_dim_size(stats, idx)) for idx in get_index_set(stats)]))) for stats in all_stats])
     new_cardinality = 2^(log2(1 - 2^prob_default) + new_dim_space_size)
     return NaiveStats(new_def, new_cardinality)
 end
 
-function reduce_tensor_stats(op, reduce_indices::Set{IndexExpr}, stats::NaiveStats)
+function reduce_tensor_stats(op, init, reduce_indices::Set{IndexExpr}, stats::NaiveStats)
     if length(reduce_indices) == 0
         return copy_stats(stats)
     end
-    new_def = reduce_tensor_def(op, reduce_indices, get_def(stats))
+    new_def = reduce_tensor_def(op, init, reduce_indices, get_def(stats))
     new_dim_space_size = sum([log2(get_dim_size(new_def, idx)) for idx in new_def.index_set])
     old_dim_space_size = sum([log2(get_dim_size(stats, idx)) for idx in get_index_set(stats)])
     prob_default_value = 1 - 2^(log2(stats.cardinality)-old_dim_space_size)
@@ -135,33 +140,66 @@ function transpose_tensor_stats(index_order::Vector{IndexExpr}, stats::NaiveStat
 end
 
 ################# DCStats Propagation ##################################################
-function merge_tensor_stats_join(op, new_def, all_stats::Vararg{DCStats})
-    new_dc_dict = Dict()
-    for dc in ∪([stats.dcs for stats in all_stats]...)
-        dc_key = get_dc_key(dc)
-        current_dc = get(new_dc_dict, dc_key, Inf)
-        if dc.d < current_dc
-            new_dc_dict[dc_key] = dc.d
+
+function unify_dc_ints(all_stats, new_def)
+    final_idx_2_int = Dict{IndexExpr, Int}()
+    final_int_2_idx = Dict{Int, IndexExpr}()
+    max_int = 1
+    for (i, idx) in enumerate(union([keys(stat.idx_2_int) for stat in all_stats]...))
+        final_idx_2_int[idx] = max_int
+        final_int_2_idx[max_int] = idx
+        max_int += 1
+    end
+    for idx in get_index_set(new_def)
+        if !haskey(final_idx_2_int, idx)
+            final_idx_2_int[idx] = max_int
+            final_int_2_idx[max_int] = idx
+            max_int += 1
         end
-        @assert new_dc_dict[dc_key] != 0
     end
-    new_stats = DCStats(new_def, Set{DC}(DC(key.X, key.Y, d) for (key, d) in new_dc_dict))
+    final_idx_2_int, final_int_2_idx
+end
+
+convert_bitset(int_to_int, b) = SmallBitSet([int_to_int[x] for x in b])
+
+function merge_tensor_stats_join(op, new_def::TensorDef, all_stats::Vararg{DCStats})
+    if length(all_stats) == 1
+        return DCStats(new_def, copy(all_stats[1].idx_2_int), copy(all_stats[1].int_2_idx), copy(all_stats[1].dcs))
+    end
+    final_idx_2_int, final_int_2_idx = unify_dc_ints(all_stats, new_def)
+    new_dc_dict = Dict{DCKey, Float64}()
+    for stats in all_stats
+        for dc in stats.dcs
+            dc_key = (X= SmallBitSet(Int[final_idx_2_int[stats.int_2_idx[x]] for x in dc.X]),
+                        Y= SmallBitSet(Int[final_idx_2_int[stats.int_2_idx[y]] for y in dc.Y]))
+            current_dc = get(new_dc_dict, dc_key, Inf)
+            if dc.d < current_dc
+                new_dc_dict[dc_key] = dc.d
+            end
+        end
+    end
+    new_stats = DCStats(new_def, final_idx_2_int, final_int_2_idx, Set{DC}(DC(key.X, key.Y, d) for (key, d) in new_dc_dict))
     return new_stats
 end
 
-function merge_tensor_stats_union(op, new_def, all_stats::Vararg{DCStats})
-    dc_keys = counter(Any)
+function merge_tensor_stats_union(op, new_def::TensorDef, all_stats::Vararg{DCStats})
+    if length(all_stats) == 1
+        return DCStats(new_def, copy(all_stats[1].idx_2_int), copy(all_stats[1].int_2_idx), copy(all_stats[1].dcs))
+    end
+    final_idx_2_int, final_int_2_idx = unify_dc_ints(all_stats, new_def)
+    dc_keys = counter(DCKey)
     stats_dcs = []
     # We start by extending all arguments' dcs to the new dimensions and infer dcs as needed
     for stats in all_stats
-#        condense_stats!(stats, timeout=1000)
-        dcs = Dict()
+        dcs = Dict{DCKey, Float64}()
         Z = setdiff(get_index_set(new_def), get_index_set(stats))
         Z_dimension_space_size = get_dim_space_size(new_def, Z)
         for dc in stats.dcs
-            dcs[(X=dc.X, Y=dc.Y)] = dc.d
-            inc!(dc_keys, (X=dc.X, Y=dc.Y))
-            ext_dc_key = (X=dc.X, Y=∪(dc.Y, Z))
+            new_key::DCKey = (X= SmallBitSet(Int[final_idx_2_int[stats.int_2_idx[x]] for x in dc.X]),
+                            Y= SmallBitSet(Int[final_idx_2_int[stats.int_2_idx[y]] for y in dc.Y]))
+            dcs[new_key] = dc.d
+            inc!(dc_keys, new_key)
+            ext_dc_key = (X=new_key.X, Y=∪(new_key.Y, idxs_to_bitset(final_idx_2_int, Z)))
             if !haskey(dcs, ext_dc_key)
                 inc!(dc_keys, ext_dc_key)
             end
@@ -172,32 +210,34 @@ function merge_tensor_stats_union(op, new_def, all_stats::Vararg{DCStats})
 
     # We only keep DCs which can be inferred from all inputs. Otherwise, we might miss
     # important information which simply wasn't inferred
-    new_dcs = Dict{Any, UInt128}()
+    new_dcs = Dict{DCKey, UInt128}()
     for (key, count) in dc_keys
         if count == length(all_stats)
             new_dcs[key] = min(typemax(UInt64), sum([get(dcs, key, UInt128(0)) for dcs in stats_dcs]))
-            if key.Y ⊆ get_index_set(new_def)
-                new_dcs[key] = min(new_dcs[key], get_dim_space_size(new_def, key.Y))
+            if  key.Y ⊆ idxs_to_bitset(final_idx_2_int, get_index_set(new_def))
+                new_dcs[key] = min(new_dcs[key], get_dim_space_size(new_def, bitset_to_idxs(final_int_2_idx, key.Y)))
             end
         end
     end
 
+#=
     for Y in subsets(collect(get_index_set(new_def)))
-        Y = Set{IndexExpr}(Y)
-        proj_dc_key = (X=Set{IndexExpr}(), Y=Y)
-        new_dcs[proj_dc_key] = min(get(new_dcs, proj_dc_key, typemax(UInt64)/2), get_dim_space_size(new_def, Y))
+        proj_dc_key = (X=SmallBitSet(), Y=idxs_to_bitset(final_idx_2_int, Y))
+        new_dcs[proj_dc_key] = min(get(new_dcs, proj_dc_key, typemax(UInt64)/2), get_dim_space_size(new_def, Set(Y)))
     end
-
-    return DCStats(new_def, Set{DC}(DC(key.X, key.Y, d) for (key, d) in new_dcs))
+ =#
+    return DCStats(new_def, final_idx_2_int, final_int_2_idx, Set{DC}(DC(key.X, key.Y, d) for (key, d) in new_dcs))
 end
 
-function reduce_tensor_stats(op, reduce_indices::Set{IndexExpr}, stats::DCStats)
+function reduce_tensor_stats(op, init, reduce_indices::Set{IndexExpr}, stats::DCStats)
     if length(reduce_indices) == 0
         return copy_stats(stats)
     end
-    new_def = reduce_tensor_def(op, reduce_indices, get_def(stats))
+    new_def = reduce_tensor_def(op, init, reduce_indices, get_def(stats))
     new_dcs = copy(stats.dcs)
-    new_stats = DCStats(new_def, new_dcs)
+    new_idx_2_int = copy(stats.idx_2_int)
+    new_int_2_idx = copy(stats.int_2_idx)
+    new_stats = DCStats(new_def, new_idx_2_int, new_int_2_idx, new_dcs)
     return new_stats
 end