enh: Adding sinkNdJson and sinkIpc (#354)

Adding sinkNdJson and sinkIpc to lazyDf to close #353

Co-authored-by: Bidek56 <[email protected]>

Author: Bidek56

__tests__/lazyframe.test.ts

@@ -1351,6 +1351,52 @@ describe("lazyframe", () => {
     expect(newDF.sort("foo")).toFrameEqual(actualDf);
     fs.rmSync("./test.parquet");
   });
+  test("sinkNdJson:path", async () => {
+    const ldf = pl
+      .DataFrame([
+        pl.Series("foo", [1, 2, 3], pl.Int64),
+        pl.Series("bar", ["a", "b", "c"]),
+      ])
+      .lazy();
+    await ldf.sinkNdJson("./test.ndjson").collect();
+    let df = pl.scanJson("./test.ndjson").collectSync();
+    expect(df.shape).toEqual({ height: 3, width: 2 });
+
+    await ldf
+      .sinkNdJson("./test.ndjson", {
+        retries: 1,
+        syncOnClose: "all",
+        maintainOrder: false,
+      })
+      .collect();
+    df = pl.scanJson("./test.ndjson").collectSync();
+    expect(df.shape).toEqual({ height: 3, width: 2 });
+
+    fs.rmSync("./test.ndjson");
+  });
+  test("sinkIpc:path", async () => {
+    const ldf = pl
+      .DataFrame([
+        pl.Series("foo", [1, 2, 3], pl.Int64),
+        pl.Series("bar", ["a", "b", "c"]),
+      ])
+      .lazy();
+    await ldf.sinkIpc("./test.ipc").collect();
+    let df = pl.scanIPC("./test.ipc").collectSync();
+    expect(df.shape).toEqual({ height: 3, width: 2 });
+
+    await ldf
+      .sinkIpc("./test.ipc", {
+        retries: 1,
+        syncOnClose: "all",
+        maintainOrder: false,
+        compression: "lz4",
+      })
+      .collect();
+    df = pl.scanIPC("./test.ipc").collectSync();
+    expect(df.shape).toEqual({ height: 3, width: 2 });
+    fs.rmSync("./test.ipc");
+  });
   test("unpivot renamed", () => {
     const ldf = pl
       .DataFrame({

polars/lazy/dataframe.ts

@@ -14,6 +14,8 @@ import type {
   LazyJoinOptions,
   LazyOptions,
   LazySameNameColumnJoinOptions,
+  SinkIpcOptions,
+  SinkJsonOptions,
   SinkParquetOptions,
 } from "../types";
 import {
@@ -616,10 +618,6 @@ export interface LazyDataFrame<S extends Schema = any>
    *
    * Evaluate the query in streaming mode and write to a Parquet file.
 
-    .. warning::
-        Streaming mode is considered **unstable**. It may be changed
-        at any point without it being considered a breaking change.
-
     This allows streaming results that are larger than RAM to be written to disk.
 
     Parameters
@@ -667,6 +665,82 @@ export interface LazyDataFrame<S extends Schema = any>
     >>> lf.sinkParquet("out.parquet").collect()  # doctest: +SKIP
    */
   sinkParquet(path: string, options?: SinkParquetOptions): LazyDataFrame;
+
+  /**
+   * 
+   * Evaluate the query in streaming mode and write to an NDJSON file.
+   * This allows streaming results that are larger than RAM to be written to disk.
+   * 
+   * Parameters
+    @param path - File path to which the file should be written.
+    @param options.maintainOrder - Maintain the order in which data is processed. Default -> true
+        Setting this to `False` will  be slightly faster.
+    @param options.mkdir - Recursively create all the directories in the path. Default -> false
+    @param options.retries - Number of retries if accessing a cloud instance fails. Default = 2
+    @param options.syncOnClose - { None, 'data', 'all' } Default -> 'all'
+            Sync to disk when before closing a file.
+
+            * `None` does not sync.
+            * `data` syncs the file contents.
+            * `all` syncs the file contents and metadata.
+    @param options.cloudOptions - Options that indicate how to connect to a cloud provider.
+        If the cloud provider is not supported by Polars, the storage options are passed to `fsspec.open()`.
+
+        The cloud providers currently supported are AWS, GCP, and Azure.
+        See supported keys here:
+
+        * `aws <https://docs.rs/object_store/latest/object_store/aws/enum.AmazonS3ConfigKey.html>`_
+        * `gcp <https://docs.rs/object_store/latest/object_store/gcp/enum.GoogleConfigKey.html>`_
+        * `azure <https://docs.rs/object_store/latest/object_store/azure/enum.AzureConfigKey.html>`_
+
+        If `cloudOptions` is not provided, Polars will try to infer the information from environment variables.
+    @return DataFrame
+    Examples
+    --------
+    >>> const lf = pl.scanCsv("/path/to/my_larger_than_ram_file.csv")  # doctest: +SKIP
+    >>> lf.sinkNdJson("out.ndjson").collect()
+   */
+  sinkNdJson(path: string, options?: SinkJsonOptions): LazyDataFrame;
+  /**
+   * 
+   * Evaluate the query in streaming mode and write to an IPC file.
+   * This allows streaming results that are larger than RAM to be written to disk.
+   * 
+   * Parameters
+    @param path - File path to which the file should be written.
+    @param options.compression : {'uncompressed', 'lz4', 'zstd'}
+            Choose "zstd" for good compression performance.
+            Choose "lz4" for fast compression/decompression.
+    @param options.compatLevel : { 'newest', 'oldest' } Default -> newest
+            Use a specific compatibility level when exporting Polars' internal data structures.
+    @param options.maintainOrder - Maintain the order in which data is processed. Default -> true
+        Setting this to `False` will  be slightly faster.
+    @param options.mkdir - Recursively create all the directories in the path. Default -> false
+    @param options.retries - Number of retries if accessing a cloud instance fails. Default = 2
+    @param options.syncOnClose - { None, 'data', 'all' } Default -> 'all'
+            Sync to disk when before closing a file.
+
+            * `None` does not sync.
+            * `data` syncs the file contents.
+            * `all` syncs the file contents and metadata.
+    @param options.cloudOptions - Options that indicate how to connect to a cloud provider.
+        If the cloud provider is not supported by Polars, the storage options are passed to `fsspec.open()`.
+
+        The cloud providers currently supported are AWS, GCP, and Azure.
+        See supported keys here:
+
+        * `aws <https://docs.rs/object_store/latest/object_store/aws/enum.AmazonS3ConfigKey.html>`_
+        * `gcp <https://docs.rs/object_store/latest/object_store/gcp/enum.GoogleConfigKey.html>`_
+        * `azure <https://docs.rs/object_store/latest/object_store/azure/enum.AzureConfigKey.html>`_
+
+        If `cloudOptions` is not provided, Polars will try to infer the information from environment variables.
+    @return DataFrame
+    Examples
+    --------
+    >>> const lf = pl.scanCsv("/path/to/my_larger_than_ram_file.csv")  # doctest: +SKIP
+    >>> lf.sinkIpc("out.arrow").collect()
+   */
+  sinkIpc(path: string, options?: SinkIpcOptions): LazyDataFrame;
 }
 
 const prepareGroupbyInputs = (by) => {
@@ -1179,6 +1253,22 @@ export const _LazyDataFrame = (_ldf: any): LazyDataFrame => {
       };
       return _ldf.sinkParquet(path, options);
     },
+    sinkNdJson(path: string, options: SinkJsonOptions = {}) {
+      options.retries = options.retries ?? 2;
+      options.syncOnClose = options.syncOnClose ?? "all";
+      options.maintainOrder = options.maintainOrder ?? true;
+      options.mkdir = options.mkdir ?? true;
+      return _ldf.sinkJson(path, options);
+    },
+    sinkIpc(path: string, options: SinkIpcOptions = {}) {
+      options.compatLevel = options.compatLevel ?? "newest";
+      options.compression = options.compression ?? "uncompressed";
+      options.retries = options.retries ?? 2;
+      options.syncOnClose = options.syncOnClose ?? "all";
+      options.maintainOrder = options.maintainOrder ?? true;
+      options.mkdir = options.mkdir ?? true;
+      return _ldf.sinkIpc(path, options);
+    },
   };
 };
 

polars/types.ts

@@ -69,7 +69,7 @@ export interface SinkOptions {
 }
 
 /**
- * Options for @see {@link LazyDataFrame.sinkParquet}
+ * Options for @see {@link LazyDataFrame.sinkParquet }
  * @category Options
  */
 export interface SinkParquetOptions {
@@ -89,6 +89,30 @@ export interface SinkParquetOptions {
   retries?: number;
   sinkOptions?: SinkOptions;
 }
+/**
+ * Options for @see {@link LazyDataFrame.sinkNdJson}
+ * @category Options
+ */
+export interface SinkJsonOptions {
+  cloudOptions?: Map<string, string>;
+  retries?: number;
+  syncOnClose?: string; // Call sync when closing the file.
+  maintainOrder?: boolean; // The output file needs to maintain order of the data that comes in.
+  mkdir?: boolean; // Recursively create all the directories in the path.
+}
+/**
+ * Options for @see {@link LazyDataFrame.sinkIpc}
+ * @category Options
+ */
+export interface SinkIpcOptions {
+  compression?: string;
+  compatLevel?: string;
+  cloudOptions?: Map<string, string>;
+  retries?: number;
+  syncOnClose?: string; // Call sync when closing the file.
+  maintainOrder?: boolean; // The output file needs to maintain order of the data that comes in.
+  mkdir?: boolean; // Recursively create all the directories in the path.
+}
 /**
  * Options for {@link DataFrame.writeJSON}
  * @category Options

src/conversion.rs

@@ -414,6 +414,16 @@ impl FromNapiValue for Wrap<Option<IpcCompression>> {
     }
 }
 
+impl ToNapiValue for Wrap<Option<IpcCompression>> {
+    unsafe fn to_napi_value(env: sys::napi_env, val: Self) -> Result<sys::napi_value> {
+        let s = match val.0.unwrap() {
+            IpcCompression::LZ4 => "lz4",
+            IpcCompression::ZSTD => "zstd",
+        };
+        String::to_napi_value(env, s.to_owned())
+    }
+}
+
 impl FromNapiValue for Wrap<UniqueKeepStrategy> {
     unsafe fn from_napi_value(env: sys::napi_env, napi_val: sys::napi_value) -> JsResult<Self> {
         let method = String::from_napi_value(env, napi_val)?;
@@ -648,6 +658,26 @@ pub struct SinkParquetOptions {
     pub sink_options: JsSinkOptions,
 }
 
+#[napi(object)]
+pub struct SinkJsonOptions {
+    pub maintain_order: Option<bool>,
+    pub cloud_options: Option<HashMap<String, String>>,
+    pub retries: Option<u32>,
+    pub sync_on_close: Wrap<SyncOnCloseType>,
+    pub mkdir: Option<bool>,
+}
+
+#[napi(object)]
+pub struct SinkIpcOptions {
+    pub compat_level: Option<String>,
+    pub compression: Wrap<Option<IpcCompression>>,
+    pub maintain_order: Option<bool>,
+    pub cloud_options: Option<HashMap<String, String>>,
+    pub retries: Option<u32>,
+    pub sync_on_close: Wrap<SyncOnCloseType>,
+    pub mkdir: Option<bool>,
+}
+
 #[napi(object)]
 pub struct ScanParquetOptions {
     pub n_rows: Option<i64>,

src/lazy/dataframe.rs

@@ -609,6 +609,62 @@ impl JsLazyFrame {
             .map_err(JsPolarsErr::from)?;
         Ok(rldf.into())
     }
+
+    #[napi(catch_unwind)]
+    pub fn sink_json(&self, path: String, options: SinkJsonOptions) -> napi::Result<JsLazyFrame> {
+        let cloud_options = parse_cloud_options(&path, options.cloud_options, options.retries);
+        let sink_options: SinkOptions = SinkOptions {
+            maintain_order: options.maintain_order.unwrap_or(true),
+            sync_on_close: options.sync_on_close.0,
+            mkdir: options.mkdir.unwrap_or(true),
+        };
+        let sink_target = SinkTarget::Path(PlPath::new(&path));
+        let rldf = self
+            .ldf
+            .clone()
+            .sink_json(
+                sink_target,
+                JsonWriterOptions {},
+                cloud_options,
+                sink_options,
+            )
+            .map_err(JsPolarsErr::from)?;
+        Ok(rldf.into())
+    }
+
+    #[napi(catch_unwind)]
+    pub fn sink_ipc(&self, path: String, options: SinkIpcOptions) -> napi::Result<JsLazyFrame> {
+        let cloud_options = parse_cloud_options(&path, options.cloud_options, options.retries);
+        let sink_options: SinkOptions = SinkOptions {
+            maintain_order: options.maintain_order.unwrap_or(true),
+            sync_on_close: options.sync_on_close.0,
+            mkdir: options.mkdir.unwrap_or(true),
+        };
+
+        let compat_level: CompatLevel = match options.compat_level.unwrap().as_str() {
+            "newest" => CompatLevel::newest(),
+            "oldest" => CompatLevel::oldest(),
+            _ => {
+                return Err(napi::Error::from_reason(
+                    "use one of {'newest', 'oldest'}".to_owned(),
+                ))
+            }
+        };
+
+        let ipc_options = IpcWriterOptions {
+            compression: options.compression.0,
+            compat_level: compat_level,
+            ..Default::default()
+        };
+
+        let sink_target = SinkTarget::Path(PlPath::new(&path));
+        let rldf = self
+            .ldf
+            .clone()
+            .sink_ipc(sink_target, ipc_options, cloud_options, sink_options)
+            .map_err(JsPolarsErr::from)?;
+        Ok(rldf.into())
+    }
 }
 
 #[napi(object)]