Migrate column factories to pylibcudf (#15257)

This PR implements `column_factories.hpp` using `pylibcudf` and migrates the cuDF cython to use them cc @vyasr Authors: - https://github.com/brandon-b-miller - Lawrence Mitchell (https://github.com/wence-) Approvers: - Vyas Ramasubramani (https://github.com/vyasr) - Lawrence Mitchell (https://github.com/wence-) URL: #15257
rapidsai · Jun 4, 2024 · eb46016 · eb46016
1 parent 382de32
commit eb46016
Show file tree

Hide file tree

Showing 17 changed files with 767 additions and 29 deletions.
diff --git a/cpp/src/column/column_factories.cpp b/cpp/src/column/column_factories.cpp
@@ -65,7 +65,8 @@ std::size_t size_of(data_type element_type)
 std::unique_ptr<column> make_empty_column(data_type type)
 {
   CUDF_EXPECTS(type.id() == type_id::EMPTY || !cudf::is_nested(type),
-               "make_empty_column is invalid to call on nested types");
+               "make_empty_column is invalid to call on nested types",
+               cudf::data_type_error);
   return std::make_unique<column>(type, 0, rmm::device_buffer{}, rmm::device_buffer{}, 0);
 }
 
@@ -80,7 +81,9 @@ std::unique_ptr<column> make_numeric_column(data_type type,
                                             rmm::device_async_resource_ref mr)
 {
   CUDF_FUNC_RANGE();
-  CUDF_EXPECTS(is_numeric(type), "Invalid, non-numeric type.");
+  CUDF_EXPECTS(type.id() != type_id::EMPTY && is_numeric(type),
+               "Invalid, non-numeric type.",
+               cudf::data_type_error);
   CUDF_EXPECTS(size >= 0, "Column size cannot be negative.");
 
   return std::make_unique<column>(
@@ -100,7 +103,7 @@ std::unique_ptr<column> make_fixed_point_column(data_type type,
                                                 rmm::device_async_resource_ref mr)
 {
   CUDF_FUNC_RANGE();
-  CUDF_EXPECTS(is_fixed_point(type), "Invalid, non-fixed_point type.");
+  CUDF_EXPECTS(is_fixed_point(type), "Invalid, non-fixed_point type.", cudf::data_type_error);
   CUDF_EXPECTS(size >= 0, "Column size cannot be negative.");
 
   return std::make_unique<column>(
@@ -120,7 +123,7 @@ std::unique_ptr<column> make_timestamp_column(data_type type,
                                               rmm::device_async_resource_ref mr)
 {
   CUDF_FUNC_RANGE();
-  CUDF_EXPECTS(is_timestamp(type), "Invalid, non-timestamp type.");
+  CUDF_EXPECTS(is_timestamp(type), "Invalid, non-timestamp type.", cudf::data_type_error);
   CUDF_EXPECTS(size >= 0, "Column size cannot be negative.");
 
   return std::make_unique<column>(
@@ -140,7 +143,7 @@ std::unique_ptr<column> make_duration_column(data_type type,
                                              rmm::device_async_resource_ref mr)
 {
   CUDF_FUNC_RANGE();
-  CUDF_EXPECTS(is_duration(type), "Invalid, non-duration type.");
+  CUDF_EXPECTS(is_duration(type), "Invalid, non-duration type.", cudf::data_type_error);
   CUDF_EXPECTS(size >= 0, "Column size cannot be negative.");
 
   return std::make_unique<column>(
@@ -160,7 +163,9 @@ std::unique_ptr<column> make_fixed_width_column(data_type type,
                                                 rmm::device_async_resource_ref mr)
 {
   CUDF_FUNC_RANGE();
-  CUDF_EXPECTS(is_fixed_width(type), "Invalid, non-fixed-width type.");
+  CUDF_EXPECTS(type.id() != type_id::EMPTY && is_fixed_width(type),
+               "Invalid, non-fixed-width type.",
+               cudf::data_type_error);
 
   // clang-format off
   if      (is_timestamp  (type)) return make_timestamp_column  (type, size, state, stream, mr);

diff --git a/cpp/tests/column/factories_test.cpp b/cpp/tests/column/factories_test.cpp
@@ -164,7 +164,7 @@ TEST_P(NonNumericFactoryTest, NonNumericThrow)
     auto column = cudf::make_numeric_column(
       cudf::data_type{GetParam()}, this->size(), cudf::mask_state::UNALLOCATED);
   };
-  EXPECT_THROW(construct(), cudf::logic_error);
+  EXPECT_THROW(construct(), cudf::data_type_error);
 }
 
 INSTANTIATE_TEST_CASE_P(NonNumeric,
@@ -307,7 +307,7 @@ TEST_P(NonFixedWidthFactoryTest, NonFixedWidthThrow)
     auto column = cudf::make_fixed_width_column(
       cudf::data_type{GetParam()}, this->size(), cudf::mask_state::UNALLOCATED);
   };
-  EXPECT_THROW(construct(), cudf::logic_error);
+  EXPECT_THROW(construct(), cudf::data_type_error);
 }
 
 INSTANTIATE_TEST_CASE_P(NonFixedWidth,

diff --git a/cpp/tests/fixed_point/fixed_point_tests.cpp b/cpp/tests/fixed_point/fixed_point_tests.cpp
@@ -498,7 +498,7 @@ TYPED_TEST(FixedPointTestAllReps, FixedPointColumnWrapper)
 TYPED_TEST(FixedPointTestAllReps, NoScaleOrWrongTypeID)
 {
   EXPECT_THROW(cudf::make_fixed_point_column(cudf::data_type{cudf::type_id::INT32}, 0),
-               cudf::logic_error);
+               cudf::data_type_error);
 }
 
 TYPED_TEST(FixedPointTestAllReps, SimpleFixedPointColumnWrapper)

diff --git a/docs/cudf/source/user_guide/api_docs/pylibcudf/column_factories.rst b/docs/cudf/source/user_guide/api_docs/pylibcudf/column_factories.rst
@@ -0,0 +1,6 @@
+================
+column_factories
+================
+
+.. automodule:: cudf._lib.pylibcudf.column_factories
+   :members:
diff --git a/docs/cudf/source/user_guide/api_docs/pylibcudf/index.rst b/docs/cudf/source/user_guide/api_docs/pylibcudf/index.rst
@@ -11,6 +11,7 @@ This page provides API documentation for pylibcudf.
     aggregation
     binaryop
     column
+    column_factories
     concatenate
     copying
     filling

diff --git a/python/cudf/cudf/_lib/column.pyx b/python/cudf/cudf/_lib/column.pyx
@@ -39,14 +39,10 @@ from cudf._lib.types cimport (
 from cudf._lib.null_mask import bitmask_allocation_size_bytes
 from cudf._lib.types import dtype_from_pylibcudf_column
 
-# TODO: We currently need this for "casting" empty pylibcudf columns in
-# from_pylibcudf by instead creating an empty numeric column. We will be able
-# to remove this once column factories are exposed to pylibcudf.
 
 cimport cudf._lib.pylibcudf.libcudf.copying as cpp_copying
 cimport cudf._lib.pylibcudf.libcudf.types as libcudf_types
 cimport cudf._lib.pylibcudf.libcudf.unary as libcudf_unary
-from cudf._lib.pylibcudf cimport Column as plc_Column
 from cudf._lib.pylibcudf.libcudf.column.column cimport column, column_contents
 from cudf._lib.pylibcudf.libcudf.column.column_factories cimport (
     make_column_from_scalar as cpp_make_column_from_scalar,
@@ -623,22 +619,17 @@ cdef class Column:
         pylibcudf.Column
             A new pylibcudf.Column referencing the same data.
         """
-        cdef libcudf_types.data_type new_dtype
         if col.type().id() == pylibcudf.TypeId.TIMESTAMP_DAYS:
             col = pylibcudf.unary.cast(
                 col, pylibcudf.DataType(pylibcudf.TypeId.TIMESTAMP_SECONDS)
             )
         elif col.type().id() == pylibcudf.TypeId.EMPTY:
-            new_dtype = libcudf_types.data_type(libcudf_types.type_id.INT8)
-            # TODO: This function call is what requires cimporting pylibcudf.
-            # We can remove the cimport once we can directly do
-            # pylibcudf.column_factories.make_numeric_column or equivalent.
-            col = plc_Column.from_libcudf(
-                move(
-                    make_numeric_column(
-                        new_dtype, col.size(), libcudf_types.mask_state.ALL_NULL
-                        )
-                    )
+            new_dtype = pylibcudf.DataType(pylibcudf.TypeId.INT8)
+
+            col = pylibcudf.column_factories.make_numeric_column(
+                new_dtype,
+                col.size(),
+                pylibcudf.column_factories.MaskState.ALL_NULL
             )
 
         dtype = dtype_from_pylibcudf_column(col)

diff --git a/python/cudf/cudf/_lib/pylibcudf/CMakeLists.txt b/python/cudf/cudf/_lib/pylibcudf/CMakeLists.txt
@@ -16,6 +16,7 @@ set(cython_sources
     aggregation.pyx
     binaryop.pyx
     column.pyx
+    column_factories.pyx
     concatenate.pyx
     copying.pyx
     filling.pyx

diff --git a/python/cudf/cudf/_lib/pylibcudf/__init__.pxd b/python/cudf/cudf/_lib/pylibcudf/__init__.pxd
@@ -4,6 +4,7 @@
 from . cimport (
     aggregation,
     binaryop,
+    column_factories,
     concatenate,
     copying,
     filling,
@@ -40,6 +41,7 @@ __all__ = [
     "binaryop",
     "concatenate",
     "copying",
+    "column_factories",
     "filling",
     "gpumemoryview",
     "groupby",

diff --git a/python/cudf/cudf/_lib/pylibcudf/__init__.py b/python/cudf/cudf/_lib/pylibcudf/__init__.py
@@ -3,6 +3,7 @@
 from . import (
     aggregation,
     binaryop,
+    column_factories,
     concatenate,
     copying,
     filling,
@@ -27,7 +28,7 @@
 from .gpumemoryview import gpumemoryview
 from .scalar import Scalar
 from .table import Table
-from .types import DataType, TypeId
+from .types import DataType, MaskState, TypeId
 
 __all__ = [
     "Column",
@@ -39,6 +40,7 @@
     "binaryop",
     "concatenate",
     "copying",
+    "column_factories",
     "filling",
     "gpumemoryview",
     "groupby",

diff --git a/python/cudf/cudf/_lib/pylibcudf/column_factories.pxd b/python/cudf/cudf/_lib/pylibcudf/column_factories.pxd
@@ -0,0 +1,52 @@
+# Copyright (c) 2024, NVIDIA CORPORATION.
+from libcpp.memory cimport unique_ptr
+from libcpp.utility cimport move
+
+from cudf._lib.pylibcudf.libcudf.types cimport mask_state, size_type
+
+from .column cimport Column
+from .types cimport DataType, size_type, type_id
+
+ctypedef fused MakeEmptyColumnOperand:
+    DataType
+    type_id
+    object
+
+ctypedef fused MaskArg:
+    mask_state
+    object
+
+
+cpdef Column make_empty_column(
+    MakeEmptyColumnOperand type_or_id
+)
+
+cpdef Column make_numeric_column(
+    DataType type_,
+    size_type size,
+    MaskArg mask,
+)
+
+cpdef Column make_fixed_point_column(
+    DataType type_,
+    size_type size,
+    MaskArg mask,
+)
+
+cpdef Column make_timestamp_column(
+    DataType type_,
+    size_type size,
+    MaskArg mask,
+)
+
+cpdef Column make_duration_column(
+    DataType type_,
+    size_type size,
+    MaskArg mask,
+)
+
+cpdef Column make_fixed_width_column(
+    DataType type_,
+    size_type size,
+    MaskArg mask,
+)