Using DataStructures in FunctionCompile

Question

This is a follow up to this StackExchange question

I am trying to do the equivalent of Select on a "FixedArray" data structure.

Here is what I have tried so far:

fixedArrayDS = CreateDataStructure["FixedArray", 2^10];
vanillaArray = RandomReal[{0, 1}, fixedArrayDS["Length"]];
Do[fixedArrayDS["SetPart", i, vanillaArray[[i]]], {i, Length[vanillaArray]}]

I'd like to do something like this:

Select[vanillaArray, # < 0.5 &]

without extracting all the elements (i.e, this: Select[fixedArrayDS["Elements"], # < 0.5 &] is considered "cheating")

Let's try to compile a function to do this (I am trying to teach myself how to use FunctionCompile with data structures). Sow and Reap seems a plausible way to do this.

select = FunctionCompile[
  Function[
   {Typed[fa, "FixedArray"], Typed[bound, "Real64"]},
   Last[
    Reap[
     Do[
      If[fa["Part", i] < bound, Sow[fa["Part", i]]],
      {i, 1, fa["Length"]}
      ]
     ]
    ]
   ]
  ]

doesn't work and comes back with a TypeError stating that a version of Reap can't be found that takes one argument. I'll come back to why Reap isn't going to work later and a hypothesis about using KernelFunction later.

As sanity checks. Test and see if this is the correct syntax for FunctionCompile on this data structure.

first = FunctionCompile[
  Function[
   {Typed[fa, "FixedArray"], Typed[multiplier, "Real64"]},
   fa["Part", 1] multiplier
   ]
  ]

this works

first[fixedArrayDS, 1]

so the FunctionCompile syntax seems to be correct.

Another sanity check. See if Compile works with Reap and Sow:

cf = Compile[{{l, _Real, 1}, {b, _Real}},
  Last[Reap[
    Do[
     If[l[[i]] < b, Sow[l[[i]]]],
     {i, 1, Length[l]}
     ]
    ]
   ]
  ]

which works, but is slower than the built-in Select:

RepeatedTiming[cf[vanillaArray, .5]]

Back to the original problem. Getting clues from Tom Wickham Jones' talk in 2020. Reap is not going to work because it isn't in the following list:

Needs["Compile`Utilities`Support`"]
Compile`Utilities`SupportedSymbols[]

Perhaps using KernelFunction in the FunctionCompile would work--but how would I use Type for Reap?

Another solution might be to use another data structure that has a rapid Append:

select = FunctionCompile[
  Function[
   {Typed[fa, "FixedArray"], Typed[bound, "Real64"]},
   Block[{collection = CreateDataStructure["DynamicArray"]},
    Do[
     If[fa["Part", i] < bound, collection["Append", fa["Part", i]]],
     {i, 1, fa["Length"]}
     ];
    collection
    ]
   ]
  ]

But, this doesn't compile either.

Finally, there is a hint that there is a solution for this task that might be in development. See around minute 51:30 in the TWJ talk linked above; it appears there is a package for using Iterators that would be a natural implement select.

Needs["IteratorTools`"]

But it looks like it didn't ship with 12.3.

Answer 1

I give a simplified version below that is slightly faster than the solution from illan.

select=FunctionCompile[Function[{Typed[pa,"PackedArray"["Real64",1]],Typed[bound,"Real64"]},
Module[{da,elem},
  da=CreateDataStructure@"DynamicArray";
  Do[
    elem=Part[pa,i];
    If[elem<bound,da["Append",elem]],{i,1,Length@pa}
  ];
  Developer`ToPackedArray@da]
]];

Then we use it like so

array = RandomReal[{0, 1}, 100];
select[array, 0.5];

Update

However, I find the implementation below that uses Compile is almost twice as fast and the implementation above. Considering that and this example, it seems using Compile can usually give a faster program than using FunctionCompile.

selectCompiled=Compile[{{array,_Real,1},{bound,_Real}},Module[{elem,bag=Internal`Bag[]},
Do[
  elem=Compile`GetElement[array,n];
  If[elem<bound,Internal`StuffBag[bag,elem]],{n,Length@array}
];
Internal`BagPart[bag,All]],RuntimeOptions->"Speed",CompilationTarget->"C"];

Answer 2

This is a bit too long for a comment, but below is a slightly modified version that uses the "DynamicArray" approach.

The basic issue is that all the precompiled versions of data structures that can be used in top-level WL via CreateDataStructure are only available for the "Expression" type.

The underlying implementations are actually polymorphic and can be instantiated for different types, but only in compiled code.

However, the results (e.g. a "FixedArray" or a "DynamicArray" with "Real64" elements) cannot be used directly in top-level, although they can be passed in to compiled code. This is why the result of select is converted prior to returning it.

select = 
  FunctionCompile[
   Function[{Typed[fa, "FixedArray"["Real64"]], 
     Typed[bound, "Real64"]}, 
    Block[{collection = CreateDataStructure["DynamicArray"]}, 
     Do[If[fa["Part", i] < bound, 
       collection["Append", fa["Part", i]]], {i, 1, fa["Length"]}];
     Developer`ToPackedArray[collection]]]];

toFixedArrayReal64 = 
  FunctionCompile[
   Function[{Typed[pa, "PackedArray"["Real64", 1]]},
    Module[{len, res},
     len = Length[pa];
     res = CreateDataStructure["FixedArray", 0., len];
     Do[res[[i]] = pa[[i]], {i, len}];
     res]]];

len = 2^20;
vanillaArray = RandomReal[{0, 1}, len];
fixedArray = toFixedArrayReal64[vanillaArray];

RepeatedTiming[res1 = Select[vanillaArray, # < 0.5 &];]

(* {0.395801, Null} *)

RepeatedTiming[res2 = select[fixedArray, 0.5];]

(* {0.0183779, Null} *)

res1 === res2

(* True *)