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
open Ast

(*
--# Number of rows
--#
--# Returns the number of rows in a DataFrame or the length of a vector.
--#
--# @name nrow
--# @param x :: DataFrame | Vector The input data.
--# @return :: Int The number of rows/elements.
--# @example
--#   nrow(mtcars)
--# @family dataframe
--# @seealso ncol, length
--# @export
*)
let register env =
  Env.add "nrow"
    (make_builtin ~name:"nrow" 1 (fun args _env ->
      match args with
      | [VDataFrame { arrow_table; _ }] -> VInt (Arrow_table.num_rows arrow_table)
      | [VVector v] -> VInt (Array.length v)
      | [VNA _] -> Error.type_error "Function `nrow` expects a DataFrame or vector, got NA."
      | [_] -> Error.type_error "Function `nrow` expects a DataFrame or vector."
      | _ -> Error.arity_error_named "nrow" 1 (List.length args)
    ))
    env