stream: allow stream to stay open after take

doc/api/stream.md

@@ -2437,6 +2437,9 @@ added:
 * `options` {Object}
   * `signal` {AbortSignal} allows destroying the stream if the signal is
     aborted.
+  * `destroyStream` {boolean} When set to `false`, the stream will not be
+    closed after take is finished unless the stream had an error.
+    **Default:** `true`.
 * Returns: {Readable} a stream with `limit` chunks taken.
 
 This method returns a new stream with the first `limit` chunks.
@@ -2447,6 +2450,26 @@ import { Readable } from 'node:stream';
 await Readable.from([1, 2, 3, 4]).take(2).toArray(); // [1, 2]
 ```
 
+Using the `destroyStream: false` option prevents `take` from closing the stream,
+so that the remaining stream data can be consumed later on
+
+```mjs
+import fs from 'node:fs';
+import { Readable } from 'node:stream';
+
+const csvParsedStream = fs
+  .createReadStream('file.csv')
+  .compose(myAwesomeParseCSV());
+
+const [columns] = await csvParsedStream
+  .take(1)
+  .toArray();
+
+const parsed = await csvParsedStream
+  .map((row) => parseRowByColumns(row, columns))
+  .toArray();
+```
+
 ##### `readable.asIndexedPairs([options])`
 
 <!-- YAML

lib/internal/streams/operators.js

@@ -15,6 +15,7 @@ const {
   validateAbortSignal,
   validateInteger,
   validateObject,
+  validateBoolean,
 } = require('internal/validators');
 const { kWeakHandler } = require('internal/event_target');
 const { finished } = require('internal/streams/end-of-stream');
@@ -397,21 +398,37 @@ function take(number, options = undefined) {
   if (options?.signal != null) {
     validateAbortSignal(options.signal, 'options.signal');
   }
+  if (options?.destroyStream != null) {
+    validateBoolean(options.destroyStream, 'options.destroyStream');
+  }
 
   number = toIntegerOrInfinity(number);
   return async function* take() {
     if (options?.signal?.aborted) {
       throw new AbortError();
     }
-    for await (const val of this) {
-      if (options?.signal?.aborted) {
-        throw new AbortError();
+    try {
+      for await (const val of this.iterator({ destroyOnReturn: options?.destroyStream ?? true })) {
+        if (options?.signal?.aborted) {
+          throw new AbortError();
+        }
+        if (number-- > 0) {
+          yield val;
+
+          // Don't get another item from iterator in case we reached the end
+          if (number === 0) {
+            return;
+          }
+        } else {
+          return;
+        }
       }
-      if (number-- > 0) {
-        yield val;
-      } else {
-        return;
+    } catch (e) {
+      if (!this.destroyed) {
+        this.destroy(e);
       }
+
+      throw e;
     }
   }.call(this);
 }

test/parallel/test-stream-drop-take.js

@@ -4,7 +4,7 @@ const common = require('../common');
 const {
   Readable,
 } = require('stream');
-const { deepStrictEqual, rejects, throws } = require('assert');
+const { deepStrictEqual, rejects, throws, strictEqual } = require('assert');
 
 const { from } = Readable;
 
@@ -100,3 +100,52 @@ const naturals = () => from(async function*() {
   throws(() => Readable.from([1]).take(1, 1), /ERR_INVALID_ARG_TYPE/);
   throws(() => Readable.from([1]).take(1, { signal: true }), /ERR_INVALID_ARG_TYPE/);
 }
+
+{
+  (async () => {
+    const streamShouldCloseWithoutOption = from([1, 2, 3, 4, 5]);
+
+    // Close stream by default
+    await streamShouldCloseWithoutOption.take(2).toArray();
+    strictEqual(streamShouldCloseWithoutOption.destroyed, true);
+  })().then(common.mustCall());
+}
+
+{
+  (async () => {
+    const streamShouldCloseWithOption = from([1, 2, 3, 4, 5]);
+
+    await streamShouldCloseWithOption.take(2, { destroyStream: true }).toArray();
+    strictEqual(streamShouldCloseWithOption.destroyed, true);
+  })().then(common.mustCall());
+}
+
+{
+  (async () => {
+    const streamShouldNotClose = from([1, 2, 3, 4, 5]);
+
+    // Do not close stream
+    await streamShouldNotClose.take(2, { destroyStream: false }).toArray();
+    strictEqual(streamShouldNotClose.destroyed, false);
+
+    deepStrictEqual(await streamShouldNotClose.toArray(), [3, 4, 5]);
+    strictEqual(streamShouldNotClose.destroyed, true);
+  })().then(common.mustCall());
+}
+
+{
+  const errorToThrow = new Error('should close');
+
+  const streamShouldNotClose = from((function *() {
+    yield 1;
+    throw errorToThrow;
+  })());
+
+  streamShouldNotClose.take(3, { destroyStream: false })
+    .toArray()
+    .then(common.mustNotCall())
+    .catch(common.mustCall((error) => {
+      strictEqual(streamShouldNotClose.destroyed, true);
+      strictEqual(error, errorToThrow);
+    }));
+}