emkornfield commented on a change in pull request #10230:
URL: https://github.com/apache/arrow/pull/10230#discussion_r663310270
##########
File path: python/pyarrow/_csv.pyx
##########
@@ -1010,15 +1015,44 @@ def write_csv(data, output_file, write_options=None,
get_writer(output_file, &stream)
c_memory_pool = maybe_unbox_memory_pool(memory_pool)
+ c_write_options.io_context = CIOContext(c_memory_pool)
if isinstance(data, RecordBatch):
batch = pyarrow_unwrap_batch(data).get()
with nogil:
- check_status(WriteCSV(deref(batch), c_write_options, c_memory_pool,
- stream.get()))
+ check_status(WriteCSV(deref(batch), c_write_options, stream.get()))
elif isinstance(data, Table):
table = pyarrow_unwrap_table(data).get()
with nogil:
- check_status(WriteCSV(deref(table), c_write_options, c_memory_pool,
- stream.get()))
+ check_status(WriteCSV(deref(table), c_write_options, stream.get()))
else:
raise TypeError(f"Expected Table or RecordBatch, got '{type(data)}'")
+
+
+cdef class CSVWriter(_CRecordBatchWriter):
+ """Writer to create a CSV file.
+
+ Parameters
+ ----------
+ sink: string, path, pyarrow.OutputStream or file-like object
+ The location where to write the CSV data.
+ schema: pyarrow.Schema
+ The schema of the data to be written.
+ write_options: pyarrow.csv.WriteOptions
+ Options to configure writing the CSV data.
+ memory_pool: MemoryPool, optional
+ Pool for temporary allocations.
+ """
+
+ def __init__(self, sink, Schema schema, *,
+ WriteOptions write_options=None, MemoryPool memory_pool=None):
+ cdef:
+ shared_ptr[COutputStream] c_stream
+ shared_ptr[CSchema] c_schema = pyarrow_unwrap_schema(schema)
+ CCSVWriteOptions c_write_options
+ CMemoryPool* c_memory_pool = maybe_unbox_memory_pool(memory_pool)
+ _get_write_options(write_options, &c_write_options)
+ c_write_options.io_context = CIOContext(c_memory_pool)
Review comment:
IOContext is new to me in general. Should we be making new API's take
that instead of Memory pool?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
