1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
open Ast
(*
--# Get the last n rows/items
--#
--# Returns the last n items from a List, Vector, or DataFrame.
--# For DataFrames, it returns the bottom n rows.
--#
--# @name tail
--# @param data :: DataFrame | List | Vector The collection to slice.
--# @param n :: Int = 5 Number of items to return.
--# @return :: DataFrame | List | Vector A subset of the input containing the last n items.
--# @example
--# tail([1, 2, 3, 4, 5, 6], n = 3)
--# -- Returns = [4, 5, 6]
--#
--# df |> tail(n = 10)
--# @family core
--# @export
*)
let register env =
Env.add "tail"
(make_builtin_named ~name:"tail" ~variadic:true 1 (fun named_args _env ->
(* Extract named arguments *)
let n_named = List.fold_left (fun acc (name, v) ->
match name, v with
| Some "n", VInt n when n >= 0 -> Some n
| _ -> acc
) None named_args in
(* Extract positional arguments *)
let args = List.filter (fun (name, _) ->
name <> Some "n"
) named_args |> List.map snd in
let take_tail_df arrow_table group_keys n =
let nrows = Arrow_table.num_rows arrow_table in
let take_n = min n nrows in
let start = nrows - take_n in
let indices = List.init take_n (fun i -> start + i) in
let new_table = Arrow_table.take_rows arrow_table indices in
VDataFrame { arrow_table = new_table; group_keys }
in
match args with
| [VDataFrame { arrow_table; group_keys }] ->
let n = match n_named with Some n -> n | None -> 5 in
take_tail_df arrow_table group_keys n
| [VDataFrame { arrow_table; group_keys }; VInt n] when n >= 0 ->
take_tail_df arrow_table group_keys n
| [VList items] ->
(match n_named with
| Some n ->
let len = List.length items in
let take_n = min n len in
VList (List.filteri (fun i _ -> i >= len - take_n) items)
| None -> (match items with _ :: rest -> VList rest | [] -> VList []))
| [VList items; VInt n] when n >= 0 ->
let len = List.length items in
let take_n = min n len in
VList (List.filteri (fun i _ -> i >= len - take_n) items)
| [VVector arr] ->
(match n_named with
| Some n ->
let len = Array.length arr in
let take_n = min n len in
VVector (Array.sub arr (len - take_n) take_n)
| None ->
if Array.length arr > 0 then
VVector (Array.sub arr 1 (Array.length arr - 1))
else VVector [||])
| [VVector arr; VInt n] when n >= 0 ->
let len = Array.length arr in
let take_n = min n len in
VVector (Array.sub arr (len - take_n) take_n)
| [VNA _] -> Error.type_error "Function `tail` cannot be called on NA."
| [_] -> Error.type_error "Function `tail` expects a DataFrame, List, or Vector."
| _ -> Error.arity_error_named "tail" 2 (List.length args)
))
env