Merge pull request #129 from bryevdv/bryanv/sequence

add static Series.sequence method

modules/cudf/src/column.cpp

@@ -171,6 +171,8 @@ Napi::Object Column::Init(Napi::Env env, Napi::Object exports) {
                   // column/stream_compaction.cpp
                   InstanceMethod<&Column::drop_nulls>("drop_nulls"),
                   InstanceMethod<&Column::drop_nans>("drop_nans"),
+                  // column/filling.cpp
+                  StaticMethod<&Column::sequence>("sequence"),
                   // column/transform.cpp
                   InstanceMethod<&Column::nans_to_nulls>("nans_to_nulls"),
                   // column/reduction.cpp

modules/cudf/src/column.ts

@@ -71,6 +71,33 @@ export type ColumnProps<T extends DataType = any> = {
 interface ColumnConstructor {
   readonly prototype: Column;
   new<T extends DataType = any>(props: ColumnProps<T>): Column<T>;
+
+  /**
+   * Fills a column with a sequence of values specified by an initial value and a step of 1.
+   *
+   * @param size Size of the output column
+   * @param init First value in the sequence
+   * @param memoryResource The optional MemoryResource used to allocate the result column's device
+   *   memory
+   * @returns column with the sequence
+   */
+  sequence<U extends DataType>(size: number, init: Scalar<U>, memoryResource?: MemoryResource):
+    Column<U>;
+
+  /**
+   * Fills a column with a sequence of values specified by an initial value and a step.
+   *
+   * @param size Size of the output column
+   * @param init First value in the sequence
+   * @param step Increment value
+   * @param memoryResource The optional MemoryResource used to allocate the result column's device
+   *   memory
+   * @returns column with the sequence
+   */
+  sequence<U extends DataType>(size: number,
+                               init: Scalar<U>,
+                               step: Scalar<U>,
+                               memoryResource?: MemoryResource): Column<U>;
 }
 
 /**

modules/cudf/src/column/filling.cpp

@@ -0,0 +1,77 @@
+// Copyright (c) 2021, NVIDIA CORPORATION.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <node_cudf/column.hpp>
+#include <node_rmm/device_buffer.hpp>
+#include <nv_node/utilities/wrap.hpp>
+
+#include <napi.h>
+#include <cudf/column/column.hpp>
+#include <cudf/filling.hpp>
+#include <cudf/types.hpp>
+#include <rmm/device_buffer.hpp>
+
+namespace nv {
+
+ObjectUnwrap<Column> Column::sequence(Napi::Env const& env,
+                                      cudf::size_type size,
+                                      cudf::scalar const& init,
+                                      rmm::mr::device_memory_resource* mr) {
+  try {
+    return Column::New(cudf::sequence(size, init, mr));
+  } catch (cudf::logic_error const& err) { NAPI_THROW(Napi::Error::New(env, err.what())); }
+}
+
+ObjectUnwrap<Column> Column::sequence(Napi::Env const& env,
+                                      cudf::size_type size,
+                                      cudf::scalar const& init,
+                                      cudf::scalar const& step,
+                                      rmm::mr::device_memory_resource* mr) {
+  try {
+    return Column::New(cudf::sequence(size, init, step, mr));
+  } catch (cudf::logic_error const& err) { NAPI_THROW(Napi::Error::New(env, err.what())); }
+}
+
+Napi::Value Column::sequence(Napi::CallbackInfo const& info) {
+  CallbackArgs args{info};
+
+  if (args.Length() != 3 and args.Length() != 4) {
+    NAPI_THROW(Napi::Error::New(
+      info.Env(), "sequence expects a size, init, and optionally a step and/or memory resource"));
+  }
+
+  if (!args[0].IsNumber()) {
+    throw Napi::Error::New(info.Env(), "sequence size argument expects a number");
+  }
+  cudf::size_type size = args[0];
+
+  if (!Scalar::is_instance(args[1])) {
+    throw Napi::Error::New(info.Env(), "sequence init argument expects a scalar");
+  }
+  auto& init = *Scalar::Unwrap(args[1]);
+
+  if (args.Length() == 3) {
+    rmm::mr::device_memory_resource* mr = args[2];
+    return Column::sequence(info.Env(), size, init, mr);
+  } else {
+    if (!Scalar::is_instance(args[2])) {
+      throw Napi::Error::New(info.Env(), "sequence step argument expects a scalar");
+    }
+    auto& step                          = *Scalar::Unwrap(args[2]);
+    rmm::mr::device_memory_resource* mr = args[3];
+    return Column::sequence(info.Env(), size, init, step, mr);
+  }
+
+}  // namespace nv
+}  // namespace nv

modules/cudf/src/node_cudf/column.hpp

@@ -596,6 +596,20 @@ class Column : public Napi::ObjectWrap<Column> {
   ObjectUnwrap<Column> drop_nans(
     rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource()) const;
 
+  // column/filling.cpp
+  static ObjectUnwrap<Column> sequence(
+    Napi::Env const& env,
+    cudf::size_type size,
+    cudf::scalar const& init,
+    rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource());
+
+  static ObjectUnwrap<Column> sequence(
+    Napi::Env const& env,
+    cudf::size_type size,
+    cudf::scalar const& init,
+    cudf::scalar const& step,
+    rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource());
+
   // column/transform.cpp
   std::pair<std::unique_ptr<rmm::device_buffer>, cudf::size_type> nans_to_nulls(
     rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource()) const;
@@ -713,6 +727,9 @@ class Column : public Napi::ObjectWrap<Column> {
   Napi::Value drop_nulls(Napi::CallbackInfo const& info);
   Napi::Value drop_nans(Napi::CallbackInfo const& info);
 
+  // column/filling.cpp
+  static Napi::Value sequence(Napi::CallbackInfo const& info);
+
   // column/transform.cpp
   Napi::Value nans_to_nulls(Napi::CallbackInfo const& info);
 

modules/cudf/src/series.ts

@@ -33,16 +33,15 @@ import {
   Int64,
   Int8,
   List,
+  Numeric,
   Struct,
   Uint16,
   Uint32,
   Uint64,
   Uint8,
   Utf8String,
 } from './types/dtypes';
-import {
-  NullOrder,
-} from './types/enums';
+import {NullOrder} from './types/enums';
 import {ArrowToCUDFType, arrowToCUDFType} from './types/mappings';
 
 export type SeriesProps<T extends DataType = any> = {
@@ -82,6 +81,14 @@ export type SeriesProps<T extends DataType = any> = {
   children?: ReadonlyArray<Series>|null;
 };
 
+export type SequenceOptions<U extends Numeric = any> = {
+  type: U,
+  size: number,
+  init: number,
+  step?: number,
+  memoryResource?: MemoryResource
+};
+
 export type Series<T extends arrow.DataType = any> = {
   [arrow.Type.NONE]: never,  // TODO
   [arrow.Type.Null]: never,  // TODO
@@ -291,6 +298,23 @@ export class AbstractSeries<T extends DataType = any> {
     return new DataFrame({0: this}).toArrow().getChildAt<T>(0)!.chunks[0] as VectorType<T>;
   }
 
+  /**
+   * Fills a Series with a sequence of values.
+   *
+   * If step is omitted, it takes a value of 1.
+   *
+   * @param opts Options for creating the sequence
+   * @returns Series with the sequence
+   */
+  public static sequence<U extends Numeric>(opts: SequenceOptions<U>): Series<U> {
+    const init = new Scalar({type: opts.type, value: opts.init});
+    if (opts.step === undefined || opts.step == 1) {
+      return Series.new(Column.sequence<U>(opts.size, init, opts.memoryResource));
+    }
+    const step = new Scalar({type: opts.type, value: opts.step});
+    return Series.new(Column.sequence<U>(opts.size, init, step, opts.memoryResource));
+  }
+
   /**
    * Generate an ordering that sorts the Series in a specified way.
    *

modules/cudf/test/cudf-series-test.ts

@@ -13,7 +13,17 @@
 // limitations under the License.
 
 import {Float32Buffer, Int32Buffer, setDefaultAllocator, Uint8Buffer} from '@nvidia/cuda';
-import {Bool8, Column, Float32, Int32, NullOrder, Series, Uint8, Utf8String} from '@rapidsai/cudf';
+import {
+  Bool8,
+  Column,
+  Float32,
+  Float64,
+  Int32,
+  NullOrder,
+  Series,
+  Uint8,
+  Utf8String
+} from '@rapidsai/cudf';
 import {CudaMemoryResource, DeviceBuffer} from '@rapidsai/rmm';
 import {Uint8Vector, Utf8Vector} from 'apache-arrow';
 import {BoolVector} from 'apache-arrow';
@@ -314,3 +324,18 @@ test('FloatSeries.nansToNulls', () => {
   expect(result.nullCount).toEqual(1);
   expect(col.nullCount).toEqual(0);
 });
+
+describe.each([new Int32, new Float32, new Float64])('Series.sequence({type=%p,, ...})', (typ) => {
+  test('no step', () => {
+    const col = Series.sequence({type: typ, size: 10, init: 0});
+    expect([...col.toArrow()]).toEqual([0, 1, 2, 3, 4, 5, 6, 7, 8, 9]);
+  });
+  test('step=1', () => {
+    const col = Series.sequence({type: typ, size: 10, step: 1, init: 0});
+    expect([...col.toArrow()]).toEqual([0, 1, 2, 3, 4, 5, 6, 7, 8, 9]);
+  });
+  test('step=2', () => {
+    const col = Series.sequence({type: typ, size: 10, step: 2, init: 0});
+    expect([...col.toArrow()]).toEqual([0, 2, 4, 6, 8, 10, 12, 14, 16, 18]);
+  });
+});