Select only first n elements in operator form?

Question

Is there any way to make Select take a second argument (to select only the first n elems) while in operator form?

It's too slow to perform the whole selection and then take the first n, for example:

dataset = 
  Dataset[Table[<|"a" -> RandomChoice[{1, 2}], "b" -> RandomReal[]|>, 
    1000000]];
Length@dataset[Select[#a == 1 & (*only want the first 1k found*)]]    
AbsoluteTiming[dataset[Select[#a == 1 &]][;; 1000];]    
AbsoluteTiming[Select[dataset, #a == 1 &, 1000];]

I don't want to put the dataset inside the Select like above, I'd like to use Dataset's single bracket query notation.

Answer 1

Alas, no -- at least not as long as we wish to use Select as a descending Dataset/Query operator. The reason is that the query compiler will only recognize Select as a descending query operator when it has exactly one argument:

Needs["Dataset`"]

DescendingQ[Select[#==1&]]

(* True *)

Any attempt to curry other arguments will cause the expression to be interpreted as an ascending operator instead, as is the case for all functions that are not on the blessed list of descending operators:

op1 = Select[#, #==1&, 1000]&;
op2[c_, n_][l_] := Select[l, c, n]
op3[c_, n_] := Select[#, c, n] &
op4 = Curry[Select, {2, 3}];
op5 = Curry[{2, 3}][Select];

AscendingQ /@
  { op1
  , op2[#==1&, 1000]
  , op3[#==1&, 1000]
  , op4[#==1&, 1000]
  , op5[#==1&, 1000]
  }

(* {True, True, True, True, True} *)

We can see the difference using a small dataset:

ds = Range[10] // Dataset;

Contrast the action of the descending version of Select...

ds[Select[# < 4&], # + 100&]

... with that of the ascending version:

ds[Select[#, # < 4 &, 3] &, # + 100 &]

In the descending version, the elements are first filtered and then added to 100. In the ascending version, the elements are first added to 100 and then filtered.

We can often work around this situation by issuing consecutive queries:

ds[Select[#, # < 4 &, 3]&][All, # + 100&]

or by using subqueries:

ds[Select[#, # < 4 &, 3] & /* Query[All, # + 100 &]]

(A pedantic note: the All operators in these last queries are not strictly necessary given the listability of Plus, but they illustrate the general principle.)

It is a shame that the query compiler does not have some special treatment for the Curry operator that was introduced in version 11.3. It could be used to supplement and/or re-order the arguments to the various specially-recognized descending operators (especially Select, SelectFirst, and GroupBy). Perhaps in some future version...

Answer 2

If you want to keep the operator form, you could define a helper function to do this:

SelectSubset[crit_, n_][expr_] := Select[expr, crit, n]

Then:

r1 = dataset[SelectSubset[#a == 1&, 1000]]; //RepeatedTiming
r2 = Select[dataset, #a == 1&, 1000]; //RepeatedTiming

Normal@r1 === Normal@r2

{0.0029, Null}

{0.0030, Null}

True

Answer 3

Of course:

dataset[Select[#a == 1 &]][;; 1000]; // RepeatedTiming // First
Select[dataset, #a == 1 &, 1000]; // RepeatedTiming // First
dataset[data \[Function] Select[data, #a == 1 &, 1000]]; // RepeatedTiming // First

0.53

0.0030

0.0027

(Well, I don't know of a built-in method, but as it can be done like this, why should there be an extra built-in method? In the end, the operator forms are mere syntax sugar - sugar that has to be paid by a lot of documentation. (Btw., I usually do not hesitate to write longer code for better performance.))

Answer 4

Unadvised solution:

Unprotect[Select];
Select[crit_, n_Integer] = Select[#, crit, n] &;
Protect[Select];

so that

Select[{1, 2, 3, 4}, GreaterThan[2], 1]
Select[GreaterThan[2], 1]@{1, 2, 3, 4}
% == %% (* True *)

Answer 5

I think what you're looking for is simply Select[#, #a==1&, 1000]& @ dataset. For example,

select1s[n_] := Select[#, #a==1&, n]&
select1s[1000] @ dataset

does what you want.