add use_dtypes_as_filter in python

rapidsai · Feb 7, 2024 · f827ae2 · f827ae2
1 parent 6f8a5bc
commit f827ae2
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 1 deletion.
diff --git a/python/cudf/cudf/_lib/cpp/io/json.pxd b/python/cudf/cudf/_lib/cpp/io/json.pxd
@@ -28,6 +28,7 @@ cdef extern from "cudf/io/json.hpp" \
         size_type get_byte_range_size() except +
         bool is_enabled_lines() except +
         bool is_enabled_mixed_types_as_string() except +
+        bool is_enabled_use_dtypes_as_filter() except +
         bool is_enabled_dayfirst() except +
         bool is_enabled_experimental() except +
 
@@ -41,6 +42,7 @@ cdef extern from "cudf/io/json.hpp" \
         void set_byte_range_size(size_type size) except +
         void enable_lines(bool val) except +
         void enable_mixed_types_as_string(bool val) except +
+        void enable_use_dtypes_as_filter(bool val) except +
         void enable_dayfirst(bool val) except +
         void enable_experimental(bool val) except +
         void enable_keep_quotes(bool val) except +
@@ -79,6 +81,9 @@ cdef extern from "cudf/io/json.hpp" \
         json_reader_options_builder& mixed_types_as_string(
             bool val
         ) except +
+        json_reader_options_builder& use_dtypes_as_filter(
+            bool val
+        ) except +
         json_reader_options_builder& dayfirst(
             bool val
         ) except +

diff --git a/python/cudf/cudf/_lib/json.pyx b/python/cudf/cudf/_lib/json.pyx
@@ -51,7 +51,8 @@ cpdef read_json(object filepaths_or_buffers,
                 object byte_range,
                 bool legacy,
                 bool keep_quotes,
-                bool mixed_types_as_string):
+                bool mixed_types_as_string,
+                bool use_dtypes_as_filter):
     """
     Cython function to call into libcudf API, see `read_json`.
 
@@ -130,6 +131,7 @@ cpdef read_json(object filepaths_or_buffers,
 
     opts.enable_keep_quotes(keep_quotes)
     opts.enable_mixed_types_as_string(mixed_types_as_string)
+    opts.enable_use_dtypes_as_filter(use_dtypes_as_filter)
     # Read JSON
     cdef cudf_io_types.table_with_metadata c_result
 

diff --git a/python/cudf/cudf/io/json.py b/python/cudf/cudf/io/json.py
@@ -26,6 +26,7 @@ def read_json(
     keep_quotes=False,
     storage_options=None,
     mixed_types_as_string=False,
+    use_dtypes_as_filter=False,
     *args,
     **kwargs,
 ):
@@ -120,6 +121,7 @@ def read_json(
             engine == "cudf_legacy",
             keep_quotes,
             mixed_types_as_string,
+            use_dtypes_as_filter,
         )
     else:
         warnings.warn(