You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@arrow.apache.org by "paleolimbot (via GitHub)" <gi...@apache.org> on 2024/02/13 19:42:45 UTC
[PR] GH-40066: [Python] Support `requested_schema` in `__arrow_c__stream__()` [arrow]
paleolimbot opened a new pull request, #40070:
URL: https://github.com/apache/arrow/pull/40070
### Rationale for this change
The `requested_schema` portion of the `__arrow_c_stream__()` protocol methods errored in all cases if passed an unequal schema. There was a note about figuring out how to check the cast before doing it and a comment in https://github.com/apache/arrow/issues/40066 about how it should be done lazily. This PR (hopefully) solves both!
### What changes are included in this PR?
- Added `arrow::py::CastingRecordBatchReader`, which wraps a `arrow::RecordBatchReader`, casting each batch as it is pulled.
### Are these changes tested?
Needs more tests! Happy to add more after there's a vague consensus on the approach.
### Are there any user-facing changes?
Yes: the current approach adds `RecordBatchReader.cast()` as the way to access the casting reader.
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40066: [Python] Support `requested_schema` in `__arrow_c__stream__()` [arrow]
Posted by "paleolimbot (via GitHub)" <gi...@apache.org>.
paleolimbot commented on code in PR #40070:
URL: https://github.com/apache/arrow/pull/40070#discussion_r1504494067
##########
python/pyarrow/tests/test_ipc.py:
##########
@@ -51,16 +51,16 @@ def get_source(self):
def write_batches(self, num_batches=5, as_table=False):
nrows = 5
- schema = pa.schema([('one', pa.float64()), ('two', pa.utf8())])
+ schema = pa.schema([("one", pa.float64()), ("two", pa.utf8())])
Review Comment:
😬 (Fixed now!)
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40066: [Python] Support `requested_schema` in `__arrow_c__stream__()` [arrow]
Posted by "pitrou (via GitHub)" <gi...@apache.org>.
pitrou commented on code in PR #40070:
URL: https://github.com/apache/arrow/pull/40070#discussion_r1490614522
##########
python/pyarrow/src/arrow/python/ipc.cc:
##########
@@ -63,5 +64,59 @@ Result<std::shared_ptr<RecordBatchReader>> PyRecordBatchReader::Make(
return reader;
}
+CastingRecordBatchReader::CastingRecordBatchReader() {}
+
+Status CastingRecordBatchReader::Init(std::shared_ptr<RecordBatchReader> parent,
+ std::shared_ptr<Schema> schema) {
+ std::shared_ptr<Schema> src = parent->schema();
+
+ // Check for conformable number of columns
+ int num_fields = schema->num_fields();
+ if (src->num_fields() != num_fields) {
+ return Status::Invalid("Source has ", src->num_fields(), " but requested schema has ",
+ num_fields);
+ }
+
+ // Try to cast an empty version of all the columns before succceeding
+ compute::CastOptions options;
+ for (int i = 0; i < num_fields; i++) {
+ ARROW_ASSIGN_OR_RAISE(auto empty_array, MakeEmptyArray(src->field(i)->type()));
Review Comment:
Instead, you can probably call `CanCast` on the pairs of types?
##########
python/pyarrow/src/arrow/python/ipc.cc:
##########
@@ -19,6 +19,7 @@
#include <memory>
+#include "arrow/compute/api.h"
Review Comment:
Nit: `arrow/compute/cast.h` is probably sufficient and will pull less headers.
##########
python/pyarrow/src/arrow/python/ipc.cc:
##########
@@ -63,5 +64,59 @@ Result<std::shared_ptr<RecordBatchReader>> PyRecordBatchReader::Make(
return reader;
}
+CastingRecordBatchReader::CastingRecordBatchReader() {}
+
+Status CastingRecordBatchReader::Init(std::shared_ptr<RecordBatchReader> parent,
+ std::shared_ptr<Schema> schema) {
+ std::shared_ptr<Schema> src = parent->schema();
+
+ // Check for conformable number of columns
+ int num_fields = schema->num_fields();
+ if (src->num_fields() != num_fields) {
+ return Status::Invalid("Source has ", src->num_fields(), " but requested schema has ",
+ num_fields);
+ }
+
+ // Try to cast an empty version of all the columns before succceeding
+ compute::CastOptions options;
+ for (int i = 0; i < num_fields; i++) {
+ ARROW_ASSIGN_OR_RAISE(auto empty_array, MakeEmptyArray(src->field(i)->type()));
+ options.to_type = schema->field(i)->type();
+ ARROW_ASSIGN_OR_RAISE(auto emtpy_array_dst, compute::Cast(empty_array, options));
+ }
+
+ parent_ = std::move(parent);
+ schema_ = std::move(schema);
+
+ return Status::OK();
+}
+
+std::shared_ptr<Schema> CastingRecordBatchReader::schema() const { return schema_; }
+
+Status CastingRecordBatchReader::ReadNext(std::shared_ptr<RecordBatch>* batch) {
+ std::shared_ptr<RecordBatch> out;
+ ARROW_RETURN_NOT_OK(parent_->ReadNext(&out));
+ if (!out) {
+ return Status::OK();
+ }
+
+ auto num_columns = out->num_columns();
+ ArrayVector columns(num_columns);
+ for (int i = 0; i < num_columns; i++) {
+ ARROW_ASSIGN_OR_RAISE(columns[i],
+ compute::Cast(*out->column(i), schema_->field(i)->type()));
Review Comment:
Do we want to check for nulls if the destination fields is non-nullable?
##########
python/pyarrow/tests/test_cffi.py:
##########
@@ -591,6 +590,24 @@ def test_roundtrip_reader_capsule(constructor):
assert batch.equals(expected)
+def test_roundtrip_batch_reader_capsule_requested_schema():
+ batch = make_batch()
+ requested_schema = pa.schema([('ints', pa.list_(pa.int64()))])
+ requested_capsule = requested_schema.__arrow_c_schema__()
+ # RecordBatch has no cast() method
Review Comment:
Do we have a GH issue open for this?
##########
python/pyarrow/ipc.pxi:
##########
@@ -772,6 +772,28 @@ cdef class RecordBatchReader(_Weakrefable):
def __exit__(self, exc_type, exc_val, exc_tb):
self.close()
+ def cast(self, schema):
+ """
+ Wraps this reader with one that casts each batch lazily as it is pulled.
Review Comment:
Nit: infinitive
```suggestion
Wrap this reader with one that casts each batch lazily as it is pulled.
```
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40066: [Python] Support `requested_schema` in `__arrow_c__stream__()` [arrow]
Posted by "pitrou (via GitHub)" <gi...@apache.org>.
pitrou commented on code in PR #40070:
URL: https://github.com/apache/arrow/pull/40070#discussion_r1502930081
##########
python/pyarrow/table.pxi:
##########
@@ -2681,6 +2681,28 @@ cdef class RecordBatch(_Tabular):
return pyarrow_wrap_batch(c_batch)
+ def cast(self, Schema target_schema, safe=None, options=None):
+ """
+ Cast batch values to another schema.
+
+ Parameters
+ ----------
+ target_schema : Schema
+ Schema to cast to, the names and order of fields must match.
+ safe : bool, default True
+ Check for overflows or other unsafe conversions.
+ options : CastOptions, default None
+ Additional checks pass by CastOptions
+
+ Returns
+ -------
+ RecordBatch
+ """
+ # Wrap the more general Table cast implementation
+ tbl = Table.from_batches([self])
+ casted_tbl = tbl.cast(target_schema, safe=safe, options=options)
+ return list(casted_tbl.to_batches())[0]
Review Comment:
Hmm, we should ensure that we don't ignore any extraneous batch. Something like?
```suggestion
casted_batch, = casted_tbl.to_batches()
return casted_batch
```
##########
python/pyarrow/src/arrow/python/ipc.cc:
##########
@@ -63,5 +64,70 @@ Result<std::shared_ptr<RecordBatchReader>> PyRecordBatchReader::Make(
return reader;
}
+CastingRecordBatchReader::CastingRecordBatchReader() {}
Review Comment:
Nit
```suggestion
CastingRecordBatchReader::CastingRecordBatchReader() = default;
```
##########
python/pyarrow/src/arrow/python/ipc.cc:
##########
@@ -63,5 +64,70 @@ Result<std::shared_ptr<RecordBatchReader>> PyRecordBatchReader::Make(
return reader;
}
+CastingRecordBatchReader::CastingRecordBatchReader() {}
+
+Status CastingRecordBatchReader::Init(std::shared_ptr<RecordBatchReader> parent,
+ std::shared_ptr<Schema> schema) {
+ std::shared_ptr<Schema> src = parent->schema();
+
+ // The check for names has already been done in Python where it's easier to
+ // generate a nice error message.
+ int num_fields = schema->num_fields();
+ if (src->num_fields() != num_fields) {
+ return Status::Invalid("Number of fields not equal");
+ }
+
+ // Ensure all columns can be cast before succeeding
+ for (int i = 0; i < num_fields; i++) {
+ if (!compute::CanCast(*src->field(i)->type(), *schema->field(i)->type())) {
+ return Status::NotImplemented("Field ", i, " cannot be cast from ",
Review Comment:
`Status::TypeError` sounds better IMHO. `NotImplemented` implies that the corresponding cast should be implemented some day.
##########
python/pyarrow/tests/test_ipc.py:
##########
@@ -51,16 +51,16 @@ def get_source(self):
def write_batches(self, num_batches=5, as_table=False):
nrows = 5
- schema = pa.schema([('one', pa.float64()), ('two', pa.utf8())])
+ schema = pa.schema([("one", pa.float64()), ("two", pa.utf8())])
Review Comment:
Any reason for all these style changes? These don't seem related. Did you apply a formatting tool by mistake?
##########
python/pyarrow/table.pxi:
##########
@@ -2995,7 +3017,7 @@ cdef class RecordBatch(_Tabular):
----------
requested_schema : PyCapsule | None
A PyCapsule containing a C ArrowSchema representation of a requested
- schema. PyArrow will attempt to cast the batch to this schema.
+ schema. PyArrow will attempt to cast each batch to this schema.
Review Comment:
Why this change?
##########
python/pyarrow/ipc.pxi:
##########
@@ -772,6 +772,38 @@ cdef class RecordBatchReader(_Weakrefable):
def __exit__(self, exc_type, exc_val, exc_tb):
self.close()
+ def cast(self, target_schema):
+ """
+ Wrap this reader with one that casts each batch lazily as it is pulled.
+ Currently only a safe cast to target_schema is implemented.
+
+ Parameters
+ ----------
+ target_schema : Schema
+ Schema to cast to, the names and order of fields must match.
+
+ Returns
+ -------
+ RecordBatchReader
+ """
+ cdef:
+ shared_ptr[CSchema] c_schema
+ shared_ptr[CRecordBatchReader] c_reader
+ RecordBatchReader out
+
+ if self.schema.names != target_schema.names:
+ raise ValueError("Target schema's field names are not matching "
Review Comment:
Nit, but you can use f-strings now rather than explicit `format` calls.
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40066: [Python] Support `requested_schema` in `__arrow_c__stream__()` [arrow]
Posted by "paleolimbot (via GitHub)" <gi...@apache.org>.
paleolimbot commented on code in PR #40070:
URL: https://github.com/apache/arrow/pull/40070#discussion_r1501015349
##########
python/pyarrow/tests/test_cffi.py:
##########
@@ -591,6 +590,24 @@ def test_roundtrip_reader_capsule(constructor):
assert batch.equals(expected)
+def test_roundtrip_batch_reader_capsule_requested_schema():
+ batch = make_batch()
+ requested_schema = pa.schema([('ints', pa.list_(pa.int64()))])
+ requested_capsule = requested_schema.__arrow_c_schema__()
+ # RecordBatch has no cast() method
Review Comment:
I wrapped the `Table` implementation (but could also remove that and open an issue). I considered pasting the implementation as well but assembling the options is not trivial and copying that also seemed problematic.
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40066: [Python] Support `requested_schema` in `__arrow_c_stream__()` [arrow]
Posted by "pitrou (via GitHub)" <gi...@apache.org>.
pitrou commented on PR #40070:
URL: https://github.com/apache/arrow/pull/40070#issuecomment-1968790961
@jorisvandenbossche @wjones127 Does one of you want to take a quick look here? Otherwise I'll merge.
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40066: [Python] Support `requested_schema` in `__arrow_c__stream__()` [arrow]
Posted by "paleolimbot (via GitHub)" <gi...@apache.org>.
paleolimbot commented on code in PR #40070:
URL: https://github.com/apache/arrow/pull/40070#discussion_r1504494474
##########
python/pyarrow/table.pxi:
##########
@@ -2995,7 +3017,7 @@ cdef class RecordBatch(_Tabular):
----------
requested_schema : PyCapsule | None
A PyCapsule containing a C ArrowSchema representation of a requested
- schema. PyArrow will attempt to cast the batch to this schema.
+ schema. PyArrow will attempt to cast each batch to this schema.
Review Comment:
I don't remember! (Reverted!)
##########
python/pyarrow/src/arrow/python/ipc.cc:
##########
@@ -63,5 +64,70 @@ Result<std::shared_ptr<RecordBatchReader>> PyRecordBatchReader::Make(
return reader;
}
+CastingRecordBatchReader::CastingRecordBatchReader() {}
+
+Status CastingRecordBatchReader::Init(std::shared_ptr<RecordBatchReader> parent,
+ std::shared_ptr<Schema> schema) {
+ std::shared_ptr<Schema> src = parent->schema();
+
+ // The check for names has already been done in Python where it's easier to
+ // generate a nice error message.
+ int num_fields = schema->num_fields();
+ if (src->num_fields() != num_fields) {
+ return Status::Invalid("Number of fields not equal");
+ }
+
+ // Ensure all columns can be cast before succeeding
+ for (int i = 0; i < num_fields; i++) {
+ if (!compute::CanCast(*src->field(i)->type(), *schema->field(i)->type())) {
+ return Status::NotImplemented("Field ", i, " cannot be cast from ",
Review Comment:
Done!
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40066: [Python] Support `requested_schema` in `__arrow_c__stream__()` [arrow]
Posted by "paleolimbot (via GitHub)" <gi...@apache.org>.
paleolimbot commented on PR #40070:
URL: https://github.com/apache/arrow/pull/40070#issuecomment-1964167802
@pitrou I think I've implemented your suggestions whenever you have time to take another look!
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40066: [Python] Support `requested_schema` in `__arrow_c__stream__()` [arrow]
Posted by "paleolimbot (via GitHub)" <gi...@apache.org>.
paleolimbot commented on code in PR #40070:
URL: https://github.com/apache/arrow/pull/40070#discussion_r1501013503
##########
python/pyarrow/src/arrow/python/ipc.cc:
##########
@@ -63,5 +64,59 @@ Result<std::shared_ptr<RecordBatchReader>> PyRecordBatchReader::Make(
return reader;
}
+CastingRecordBatchReader::CastingRecordBatchReader() {}
+
+Status CastingRecordBatchReader::Init(std::shared_ptr<RecordBatchReader> parent,
+ std::shared_ptr<Schema> schema) {
+ std::shared_ptr<Schema> src = parent->schema();
+
+ // Check for conformable number of columns
+ int num_fields = schema->num_fields();
+ if (src->num_fields() != num_fields) {
+ return Status::Invalid("Source has ", src->num_fields(), " but requested schema has ",
+ num_fields);
+ }
+
+ // Try to cast an empty version of all the columns before succceeding
+ compute::CastOptions options;
+ for (int i = 0; i < num_fields; i++) {
+ ARROW_ASSIGN_OR_RAISE(auto empty_array, MakeEmptyArray(src->field(i)->type()));
+ options.to_type = schema->field(i)->type();
+ ARROW_ASSIGN_OR_RAISE(auto emtpy_array_dst, compute::Cast(empty_array, options));
+ }
+
+ parent_ = std::move(parent);
+ schema_ = std::move(schema);
+
+ return Status::OK();
+}
+
+std::shared_ptr<Schema> CastingRecordBatchReader::schema() const { return schema_; }
+
+Status CastingRecordBatchReader::ReadNext(std::shared_ptr<RecordBatch>* batch) {
+ std::shared_ptr<RecordBatch> out;
+ ARROW_RETURN_NOT_OK(parent_->ReadNext(&out));
+ if (!out) {
+ return Status::OK();
+ }
+
+ auto num_columns = out->num_columns();
+ ArrayVector columns(num_columns);
+ for (int i = 0; i < num_columns; i++) {
+ ARROW_ASSIGN_OR_RAISE(columns[i],
+ compute::Cast(*out->column(i), schema_->field(i)->type()));
Review Comment:
Done!
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40066: [Python] Support `requested_schema` in `__arrow_c__stream__()` [arrow]
Posted by "paleolimbot (via GitHub)" <gi...@apache.org>.
paleolimbot commented on code in PR #40070:
URL: https://github.com/apache/arrow/pull/40070#discussion_r1504495011
##########
python/pyarrow/ipc.pxi:
##########
@@ -772,6 +772,38 @@ cdef class RecordBatchReader(_Weakrefable):
def __exit__(self, exc_type, exc_val, exc_tb):
self.close()
+ def cast(self, target_schema):
+ """
+ Wrap this reader with one that casts each batch lazily as it is pulled.
+ Currently only a safe cast to target_schema is implemented.
+
+ Parameters
+ ----------
+ target_schema : Schema
+ Schema to cast to, the names and order of fields must match.
+
+ Returns
+ -------
+ RecordBatchReader
+ """
+ cdef:
+ shared_ptr[CSchema] c_schema
+ shared_ptr[CRecordBatchReader] c_reader
+ RecordBatchReader out
+
+ if self.schema.names != target_schema.names:
+ raise ValueError("Target schema's field names are not matching "
Review Comment:
Done!
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40066: [Python] Support `requested_schema` in `__arrow_c_stream__()` [arrow]
Posted by "conbench-apache-arrow[bot] (via GitHub)" <gi...@apache.org>.
conbench-apache-arrow[bot] commented on PR #40070:
URL: https://github.com/apache/arrow/pull/40070#issuecomment-1970344543
After merging your PR, Conbench analyzed the 7 benchmarking runs that have been run so far on merge-commit d6b9051fa05b230ce288eb0b7ab6b596200c45b1.
There were no benchmark performance regressions. 🎉
The [full Conbench report](https://github.com/apache/arrow/runs/22108642836) has more details. It also includes information about 2 possible false positives for unstable benchmarks that are known to sometimes produce them.
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40066: [Python] Support `requested_schema` in `__arrow_c__stream__()` [arrow]
Posted by "github-actions[bot] (via GitHub)" <gi...@apache.org>.
github-actions[bot] commented on PR #40070:
URL: https://github.com/apache/arrow/pull/40070#issuecomment-1942281104
:warning: GitHub issue #40066 **has been automatically assigned in GitHub** to PR creator.
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40066: [Python] Support `requested_schema` in `__arrow_c__stream__()` [arrow]
Posted by "paleolimbot (via GitHub)" <gi...@apache.org>.
paleolimbot commented on code in PR #40070:
URL: https://github.com/apache/arrow/pull/40070#discussion_r1501013262
##########
python/pyarrow/src/arrow/python/ipc.cc:
##########
@@ -19,6 +19,7 @@
#include <memory>
+#include "arrow/compute/api.h"
Review Comment:
Done!
##########
python/pyarrow/src/arrow/python/ipc.cc:
##########
@@ -63,5 +64,59 @@ Result<std::shared_ptr<RecordBatchReader>> PyRecordBatchReader::Make(
return reader;
}
+CastingRecordBatchReader::CastingRecordBatchReader() {}
+
+Status CastingRecordBatchReader::Init(std::shared_ptr<RecordBatchReader> parent,
+ std::shared_ptr<Schema> schema) {
+ std::shared_ptr<Schema> src = parent->schema();
+
+ // Check for conformable number of columns
+ int num_fields = schema->num_fields();
+ if (src->num_fields() != num_fields) {
+ return Status::Invalid("Source has ", src->num_fields(), " but requested schema has ",
+ num_fields);
+ }
+
+ // Try to cast an empty version of all the columns before succceeding
+ compute::CastOptions options;
+ for (int i = 0; i < num_fields; i++) {
+ ARROW_ASSIGN_OR_RAISE(auto empty_array, MakeEmptyArray(src->field(i)->type()));
Review Comment:
Done!
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40066: [Python] Support `requested_schema` in `__arrow_c_stream__()` [arrow]
Posted by "jorisvandenbossche (via GitHub)" <gi...@apache.org>.
jorisvandenbossche merged PR #40070:
URL: https://github.com/apache/arrow/pull/40070
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org