Adds support for json lines format to the nested JSON reader

cpp/src/io/json/experimental/read_json.cpp

@@ -50,14 +50,13 @@ table_with_metadata read_json(host_span<std::unique_ptr<datasource>> sources,
   auto const dtypes_empty =
     std::visit([](const auto& dtypes) { return dtypes.empty(); }, reader_opts.get_dtypes());
   CUDF_EXPECTS(dtypes_empty, "user specified dtypes are not yet supported");
-  CUDF_EXPECTS(not reader_opts.is_enabled_lines(), "JSON Lines format is not yet supported");
   CUDF_EXPECTS(reader_opts.get_byte_range_offset() == 0 and reader_opts.get_byte_range_size() == 0,
                "specifying a byte range is not yet supported");
 
   auto const buffer = ingest_raw_input(sources, reader_opts.get_compression());
   auto data = host_span<char const>(reinterpret_cast<char const*>(buffer.data()), buffer.size());
 
-  return cudf::io::json::detail::parse_nested_json(data, stream, mr);
+  return cudf::io::json::detail::parse_nested_json(data, reader_opts, stream, mr);
 }
 
 }  // namespace cudf::io::detail::json::experimental

cpp/src/io/json/nested_json.hpp

@@ -16,6 +16,7 @@
 
 #pragma once
 
+#include <cudf/io/json.hpp>
 #include <cudf/io/types.hpp>
 #include <cudf/types.hpp>
 #include <cudf/utilities/bit.hpp>
@@ -261,50 +262,54 @@ enum token_t : PdaTokenT {
 };
 
 namespace detail {
+
+// TODO: return device_uvector instead of passing pre-allocated memory
 /**
  * @brief Identifies the stack context for each character from a JSON input. Specifically, we
  * identify brackets and braces outside of quoted fields (e.g., field names, strings).
  * At this stage, we do not perform bracket matching, i.e., we do not verify whether a closing
  * bracket would actually pop a the corresponding opening brace.
  *
- * @param[in] d_json_in The string of input characters
+ * @param[in] json_in The string of input characters
  * @param[out] d_top_of_stack Will be populated with what-is-on-top-of-the-stack for any given input
  * character of \p d_json_in, where a '{' represents that the corresponding input character is
  * within the context of a struct, a '[' represents that it is within the context of an array, and a
  * '_' symbol that it is at the root of the JSON.
  * @param[in] stream The cuda stream to dispatch GPU kernels to
  */
-void get_stack_context(device_span<SymbolT const> d_json_in,
+void get_stack_context(device_span<SymbolT const> json_in,
                        SymbolT* d_top_of_stack,
                        rmm::cuda_stream_view stream);
 
 /**
  * @brief Parses the given JSON string and emits a sequence of tokens that demarcate relevant
  * sections from the input.
  *
- * @param[in] d_json_in The JSON input
- * @param[out] d_tokens Device memory to which the parsed tokens are written
- * @param[out] d_tokens_indices Device memory to which the indices are written, where each index
- * represents the offset within \p d_json_in that cause the input being written
- * @param[out] d_num_written_tokens The total number of tokens that were parsed
- * @param[in] stream The CUDA stream to which kernels are dispatched
+ * @param json_in The JSON input
+ * @param options Parsing options specifying the parsing behaviour
+ * @param stream The CUDA stream to which kernels are dispatched
+ * @param mr Optional, resource with which to allocate
+ * @return Pair of device vectors, where the first vector represents the token types and the second
+ * vector represents the index within the input corresponding to each token
  */
-void get_token_stream(device_span<SymbolT const> d_json_in,
-                      PdaTokenT* d_tokens,
-                      SymbolOffsetT* d_tokens_indices,
-                      SymbolOffsetT* d_num_written_tokens,
-                      rmm::cuda_stream_view stream);
+std::pair<rmm::device_uvector<PdaTokenT>, rmm::device_uvector<SymbolOffsetT>> get_token_stream(
+  device_span<SymbolT const> json_in,
+  cudf::io::json_reader_options const& options,
+  rmm::cuda_stream_view stream,
+  rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource());
 
 /**
  * @brief Parses the given JSON string and generates table from the given input.
  *
  * @param input The JSON input
+ * @param options Parsing options specifying the parsing behaviour
  * @param stream The CUDA stream to which kernels are dispatched
- * @param mr Optional, resource with which to allocate.
+ * @param mr Optional, resource with which to allocate
  * @return The data parsed from the given JSON input
  */
 table_with_metadata parse_nested_json(
   host_span<SymbolT const> input,
+  cudf::io::json_reader_options const& options,
   rmm::cuda_stream_view stream        = cudf::default_stream_value,
   rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource());