diff --git a/CHANGELOG.md b/CHANGELOG.md index 924fc3c2a5a..95117190ab1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -28,6 +28,7 @@ - Added support for `DatetimeIndex.mean` and `DatetimeIndex.std` methods. - Added support for `Resampler.asfreq`. - Added support for `resample` frequency `W`, `ME`, `YE` with `closed = "left"`. +- Added support for file writes. This feature is currently in private preview. #### Bug Fixes diff --git a/src/snowflake/snowpark/files.py b/src/snowflake/snowpark/files.py index bb8b700c933..cdbfbaf2f89 100644 --- a/src/snowflake/snowpark/files.py +++ b/src/snowflake/snowpark/files.py @@ -15,6 +15,8 @@ import sys from io import RawIOBase +from snowflake.snowpark._internal.utils import private_preview + # Python 3.8 needs to use typing.Iterable because collections.abc.Iterable is not subscriptable # Python 3.9 can use both # Python 3.10 needs to use collections.abc.Iterable because typing.Iterable is removed @@ -44,6 +46,7 @@ def __init__( is_owner_file: bool = False, *, require_scoped_url: bool = True, + from_result_api: bool = False, ) -> None: super().__init__() # The URL/URI of the file to be opened by the SnowflakeFile object @@ -122,6 +125,20 @@ def isatty(self) -> None: """ raise NotImplementedError(_DEFER_IMPLEMENTATION_ERR_MSG) + @classmethod + @private_preview(version="1.22.1") + def open_new_result(cls, mode: str = "w") -> SnowflakeFile: + """ + Returns a :class:`~snowflake.snowpark.file.SnowflakeFile`. + In UDF and Stored Procedures, the object works like a Python IOBase object and as a wrapper for an IO stream of remote files. The IO Stream is to support the file operations defined in this class. + + This stream will open a writable result file that will be materialized when it's returned by a UDF or Stored Procedure. When the file is materialized then file_uri will be set to a scoped URL that temporarily references this file. + + Args: + mode: A string used to mark the type of an IO stream. + """ + return cls("new results file", mode, require_scoped_url=0, from_result_api=True) + def read(self, size: int = -1) -> None: """ See https://docs.python.org/3/library/io.html#io.RawIOBase.read diff --git a/tests/unit/test_files.py b/tests/unit/test_files.py index 2ea3aa80268..f7b33adadcc 100644 --- a/tests/unit/test_files.py +++ b/tests/unit/test_files.py @@ -6,6 +6,7 @@ import pytest +from snowflake.snowpark._internal.utils import private_preview from snowflake.snowpark.files import _DEFER_IMPLEMENTATION_ERR_MSG, SnowflakeFile @@ -15,6 +16,13 @@ def test_create_snowflakefile(): assert snowflake_file._mode == "r" +@private_preview(version="1.22.1") +def test_write_snowflakefile(): + with SnowflakeFile.open_new_result("w") as snowflake_file: + assert snowflake_file._file_location == "new results file" + assert snowflake_file._mode == "w" + + def test_snowflake_file_attribute(): with SnowflakeFile.open("test_file_location") as snowflake_file: assert snowflake_file.buffer is None