new writer for json

Author: chitralverma

core/src/main/scala/org/polars/scala/polars/api/io/Writeable.scala

@@ -10,7 +10,7 @@ import org.polars.scala.polars.internal.jni.io.write._
 class Writeable private[polars] (ptr: Long) {
   import org.polars.scala.polars.jsonMapper
 
-  private val _options: MutableMap[String, String] = MutableMap("writeMode" -> "errorifexists")
+  private val _options: MutableMap[String, String] = MutableMap("write_mode" -> "errorifexists")
 
   def option(key: String, value: String): Writeable = synchronized {
     if (Option(key).exists(_.trim.isEmpty) || Option(value).exists(_.trim.isEmpty)) {
@@ -51,4 +51,22 @@ class Writeable private[polars] (ptr: Long) {
       filePath = filePath,
       options = jsonMapper.writeValueAsString(_options)
     )
+
+  def json(filePath: String): Unit = {
+    option("write_json_format", "json")
+    writeJson(
+      ptr = ptr,
+      filePath = filePath,
+      options = jsonMapper.writeValueAsString(_options)
+    )
+  }
+
+  def json_lines(filePath: String): Unit = {
+    option("write_json_format", "json_lines")
+    writeJson(
+      ptr = ptr,
+      filePath = filePath,
+      options = jsonMapper.writeValueAsString(_options)
+    )
+  }
 }

core/src/main/scala/org/polars/scala/polars/internal/jni/io/write.scala

@@ -22,4 +22,10 @@ object write extends Natively {
       options: String
   ): Unit
 
+  @native def writeJson(
+      ptr: Long,
+      filePath: String,
+      options: String
+  ): Unit
+
 }

native/src/internal_jni/io/write/json.rs

@@ -0,0 +1,41 @@
+#![allow(non_snake_case)]
+
+use jni::objects::{JObject, JString};
+use jni::sys::jlong;
+use jni::JNIEnv;
+use jni_fn::jni_fn;
+use polars::prelude::*;
+
+use crate::internal_jni::io::write::{get_df_and_writer, parse_json_to_options, DynWriter};
+
+#[jni_fn("org.polars.scala.polars.internal.jni.io.write$")]
+pub fn writeJson(
+    mut env: JNIEnv,
+    _object: JObject,
+    df_ptr: jlong,
+    filePath: JString,
+    options: JString,
+) {
+    let mut options = parse_json_to_options(&mut env, options).unwrap();
+
+    let json_format = options
+        .remove("write_json_format")
+        .and_then(|s| match s.to_lowercase().as_str() {
+            "json" => Some(JsonFormat::Json),
+            "json_lines" => Some(JsonFormat::JsonLines),
+            _ => None,
+        })
+        .unwrap_or(JsonFormat::Json);
+
+    let overwrite_mode = options
+        .remove("write_mode")
+        .map(|s| matches!(s.to_lowercase().as_str(), "overwrite"))
+        .unwrap_or(false);
+
+    let (mut dataframe, writer): (DataFrame, DynWriter) =
+        get_df_and_writer(&mut env, df_ptr, filePath, overwrite_mode, options).unwrap();
+
+    let mut json_writer = JsonWriter::new(writer).with_json_format(json_format);
+
+    json_writer.finish(&mut dataframe).unwrap();
+}

native/src/internal_jni/io/write/mod.rs

@@ -1,5 +1,6 @@
 pub mod avro;
 pub mod ipc;
+pub mod json;
 pub mod parquet;
 
 use std::fs::File;