Performance tuning joinAcross (my version)

Question

I have written what I think is a more or less complete version of JoinAcross with code below. I've done this to try to avoid the misbehavior I've observed in the function, namely the memory leak in JoinAcross that I've documented in my other post linked here. This function is critical to several a package I've written for backtesting financial trading strategies and so performance is key. The memory leak has crashed my wolframscipts and so I set out to write my own version but I'm running out of ideas of how to shave those last few microseconds off the performance to put it on par with the built-in JoinAcross. Here's the code for my version which I've called joinAcross.

Can someone please help me to improve the performance of this function or help make this better in any way?

For large lists of associations this approach also becomes super memory hungry due to the use of Outer. Is there a way to avoid that?

I believe the starting place to investigate is my use of Select in getting the left or right matches. Is there a faster way of getting those without first getting the inner matches?

Thank you in advance for your help and hope this function helps others as well who often use Association's or JoinAcross.

ClearAll[joinAcross];
Options[joinAcross] = {KeyCollisionFunction -> Left};
joinAcross::badjspec = 
  "Invalid join specification. Enter either \"Inner\", \"Outer\", \
\"Left\", or \"Right\".";
joinAcross::badkcspec = 
  "Invalid KeyCollisionFunction specification. Enter None, Left, \
Right, or a valid function of one paramter returning a list of two \
strings such as {\"First\"->#,\"Second\"->#}.";
joinAcross::keynx = "Key `1` does not exist in Association `2`.";
joinAcross[alist_, blist_, spec : (_String | {__String}), 
   jspec : ("Inner" | "Outer" | "Left" | "Right"), 
   opts : OptionsPattern[]] :=
  Module[{
    kcspec = OptionValue[KeyCollisionFunction], keyfunc, valfunc,
    speclist = Flatten[{spec}],
    akdata, akdata1, akdata2,
    keys,
    nonspeclist,
    kcfunc,
    assocCombos,
    innerMatches, innerSpecs, innerMergedMatches, innerResults = {},
    patt,
    leftMergedMatches, leftResults = {},
    rightMergedMatches, rightResults = {},
    result = {}
    },
   (*Handle key collision function to get key and value functions.*)
 \
  {keyfunc, valfunc} = Switch[kcspec,
     None, {Identity, (Missing["Unmatched"] &)},
     Left, {Identity, First},
     Right, {Identity, Last},
     Function[{_, _}], {kcspec, ({First@#, Last@#} &)},
     _, Message[joinAcross::badkcspec]; Return[$Failed]
     ];
   
   akdata = 
    KeyUnion[Catenate[{alist, blist}], (Missing["Unmatched"] &)];
   {akdata1, akdata2} = TakeDrop[akdata, Length@alist];
   assocCombos = Flatten[Outer[List, akdata1, akdata2], 1];
   
   keys = 
    Flatten[
     Intersection[DeleteDuplicates@Keys@data1, 
      DeleteDuplicates@Keys@data2], 1];
   nonspeclist = Complement[keys, speclist];
   
   kcfunc = Which[
      MemberQ[speclist, #1], #1 -> First@#2,
      MemberQ[nonspeclist, #1] && kcspec =!= None, 
      Thread[Rule[keyfunc[#1], valfunc[#2]]],
      MemberQ[nonspeclist, #1] && kcspec === None, Nothing,
      True, #1 -> 
       SelectFirst[#2, ! MissingQ[#] &, Missing["Unmatched"]]
      ] &;
   
   innerMatches = 
    Select[assocCombos, #[[1, speclist]] === #[[2, speclist]] &];
   innerSpecs = Values@innerMatches[[All, 1, speclist]];
   innerMergedMatches = Merge[#, Identity] & /@ innerMatches;
   innerResults = 
    Table[
     Association @@ 
      KeyValueMap[kcfunc, innerMergedMatch], {innerMergedMatch, 
      innerMergedMatches}];
   
   If[jspec == "Inner", Return[innerResults]];
   
   (*Get inner matches key pattern.*)
   
   patt = 
    Alternatives @@ 
     KeyValuePattern /@ Normal@innerMatches[[All, 1, speclist]];
   
   (*Right join.*)
   If[jspec == "Right",
    rightMergedMatches = 
     Merge[#, {Missing["Unmatched"], Last@#} &] & /@ 
       Select[
        assocCombos, #[[1, speclist]] =!= #[[2, speclist]] && 
          FreeQ[#[[2]], patt] &] // DeleteDuplicates;
    rightResults = 
     Table[
      Association @@ 
       KeyValueMap[kcfunc, rightMergedMatch], {rightMergedMatch, 
       rightMergedMatches}];
    Return[Catenate[{innerResults, rightResults}]];
    ];
   
   (*Left join.*)
   If[jspec == "Left",
    leftMergedMatches = 
     Merge[#, {First@#, Missing["Unmatched"]} &] & /@ 
       Select[
        assocCombos, #[[1, speclist]] =!= #[[2, speclist]] && 
          FreeQ[#[[1]], patt] &] // DeleteDuplicates;
    leftResults = 
     Table[
      Association @@ 
       KeyValueMap[kcfunc, leftMergedMatch], {leftMergedMatch, 
       leftMergedMatches}];
    Return[Catenate[{innerResults, leftResults}]];
    ];
   
   (*Outer join.*)
   If[jspec == "Outer",
    rightMergedMatches = 
     Merge[#, {Missing["Unmatched"], Last@#} &] & /@ 
       Select[
        assocCombos, #[[1, speclist]] =!= #[[2, speclist]] && 
          FreeQ[#[[2]], patt] &] // DeleteDuplicates;
    rightResults = 
     Table[
      Association @@ 
       KeyValueMap[kcfunc, rightMergedMatch], {rightMergedMatch, 
       rightMergedMatches}];
    leftMergedMatches = 
     Merge[#, {First@#, Missing["Unmatched"]} &] & /@ 
       Select[
        assocCombos, #[[1, speclist]] =!= #[[2, speclist]] && 
          FreeQ[#[[1]], patt] &] // DeleteDuplicates;
    leftResults = 
     Table[
      Association @@ 
       KeyValueMap[kcfunc, leftMergedMatch], {leftMergedMatch, 
       leftMergedMatches}];
    Return[Catenate[{innerResults, rightResults, leftResults}]]
    ];
   ];

I just use some test data to compare the outputs from this version to the built-in JoinAcross with the following random data:

SeedRandom[1234];
data1 = Table[<|"Date" -> DatePlus[Today, t], 
    "Value" -> RandomReal[]|>, {t, -5, 15}];
data2 = Table[<|"Date" -> DatePlus[Today, t], 
     "Value" -> RandomReal[]|>, {t, 1, 5}];

Answer 1

Hm. I think the root of all evil in your code is the line assocCombos = Flatten[Outer[List, akdata1, akdata2], 1];. Say alist has length m and blist has length n. Then this line makes the complexity of your code go up to at least $O(m \cdot n)$. And for me it feels like JoinAcross should have complexity at most $O(m + n)$-ish (maybe wit some $\log$s).

I cannot provide a full-fletched implementation. But I can make the basic use case (one key and the "Left" strategy) working.

joinAcross2[a_, b_, key_, mergefun_] := Module[{overlap, i, j},
   overlap = Select[
     Merge[
      {PositionIndex[a[[All, key]]], PositionIndex[b[[All, key]]]},
      Identity
      ],
     Length[#] == 2 &
     ];
   {i, j} = Transpose[Values[overlap]];
   MapThread[
    Merge[{#1, #2}, mergefun] &, {a[[Flatten[i]]], b[[Flatten[j]]]}]
   ];

Now let's compare:

SeedRandom[1234];
m = 100 2 2;
n = 200 2 2;
data1 = Table[<|"Date" -> DatePlus[Today, t], 
    "Value" -> RandomReal[]|>, {t, 1, m}];
data2 = Table[<|"Date" -> DatePlus[Today, t], 
    "Value" -> RandomReal[]|>, {t, 1, n}];



aa = JoinAcross[data1, data2, "Date", "Left"]; // AbsoluteTiming // First
bb = joinAcross[data1, data2, "Date", "Left"]; // AbsoluteTiming // First
cc = joinAcross2[data1, data2, "Date", First]; // AbsoluteTiming // First

aa == bb == cc 

0.000854

18.4975

0.002572

True

Not quite as fast as the built-in variant, but at least in the same ball park.