You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@arrow.apache.org by jo...@apache.org on 2024/03/15 15:15:46 UTC
(arrow) branch main updated: GH-30915: [C++][Python] Add missing methods to `RecordBatch` (#39506)
This is an automated email from the ASF dual-hosted git repository.
jorisvandenbossche pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow.git
The following commit(s) were added to refs/heads/main by this push:
new 7f361fd806 GH-30915: [C++][Python] Add missing methods to `RecordBatch` (#39506)
7f361fd806 is described below
commit 7f361fd80651010f3dc91ec6302f661a16892291
Author: Judah Rand <17...@users.noreply.github.com>
AuthorDate: Fri Mar 15 15:15:38 2024 +0000
GH-30915: [C++][Python] Add missing methods to `RecordBatch` (#39506)
### Rationale for this change
These methods are present on `Table` but missing on `RecordBatch`:
* `add_column`
* `append_column`
* `remove_column`
* `set_column`
* `drop_columns`
* `rename_columns`
* `cast`
We also should probably accept a `dict` as input to `pa.record_batch` like we do for `pa.table`.
### What changes are included in this PR?
Add the methods.
### Are these changes tested?
Yes.
* Parent issue: https://github.com/apache/arrow/issues/36399
* Related: #30559
* Closes #30915
* GitHub Issue: #30915
Lead-authored-by: Judah Rand <17...@users.noreply.github.com>
Co-authored-by: Joris Van den Bossche <jo...@gmail.com>
Signed-off-by: Joris Van den Bossche <jo...@gmail.com>
---
cpp/src/arrow/record_batch.cc | 29 +++
cpp/src/arrow/record_batch.h | 7 +
cpp/src/arrow/record_batch_test.cc | 26 ++
python/pyarrow/includes/libarrow.pxd | 7 +
python/pyarrow/table.pxi | 474 +++++++++++++++++++++++++++--------
python/pyarrow/tests/test_table.py | 202 +++++++++------
6 files changed, 561 insertions(+), 184 deletions(-)
diff --git a/cpp/src/arrow/record_batch.cc b/cpp/src/arrow/record_batch.cc
index d23b2b584b..d52ebe053b 100644
--- a/cpp/src/arrow/record_batch.cc
+++ b/cpp/src/arrow/record_batch.cc
@@ -395,6 +395,35 @@ Result<std::shared_ptr<RecordBatch>> RecordBatch::ReplaceSchema(
return RecordBatch::Make(std::move(schema), num_rows(), columns());
}
+std::vector<std::string> RecordBatch::ColumnNames() const {
+ std::vector<std::string> names(num_columns());
+ for (int i = 0; i < num_columns(); ++i) {
+ names[i] = schema()->field(i)->name();
+ }
+ return names;
+}
+
+Result<std::shared_ptr<RecordBatch>> RecordBatch::RenameColumns(
+ const std::vector<std::string>& names) const {
+ int n = num_columns();
+
+ if (static_cast<int>(names.size()) != n) {
+ return Status::Invalid("tried to rename a record batch of ", n, " columns but only ",
+ names.size(), " names were provided");
+ }
+
+ ArrayVector columns(n);
+ FieldVector fields(n);
+
+ for (int i = 0; i < n; ++i) {
+ columns[i] = column(i);
+ fields[i] = schema()->field(i)->WithName(names[i]);
+ }
+
+ return RecordBatch::Make(::arrow::schema(std::move(fields)), num_rows(),
+ std::move(columns));
+}
+
Result<std::shared_ptr<RecordBatch>> RecordBatch::SelectColumns(
const std::vector<int>& indices) const {
int n = static_cast<int>(indices.size());
diff --git a/cpp/src/arrow/record_batch.h b/cpp/src/arrow/record_batch.h
index 8a2c1ba6d7..16d721caad 100644
--- a/cpp/src/arrow/record_batch.h
+++ b/cpp/src/arrow/record_batch.h
@@ -227,6 +227,13 @@ class ARROW_EXPORT RecordBatch {
/// \return PrettyPrint representation suitable for debugging
std::string ToString() const;
+ /// \brief Return names of all columns
+ std::vector<std::string> ColumnNames() const;
+
+ /// \brief Rename columns with provided names
+ Result<std::shared_ptr<RecordBatch>> RenameColumns(
+ const std::vector<std::string>& names) const;
+
/// \brief Return new record batch with specified columns
Result<std::shared_ptr<RecordBatch>> SelectColumns(
const std::vector<int>& indices) const;
diff --git a/cpp/src/arrow/record_batch_test.cc b/cpp/src/arrow/record_batch_test.cc
index db68a9a937..45cf7cae65 100644
--- a/cpp/src/arrow/record_batch_test.cc
+++ b/cpp/src/arrow/record_batch_test.cc
@@ -315,6 +315,32 @@ TEST_F(TestRecordBatch, RemoveColumn) {
AssertBatchesEqual(*new_batch, *batch4);
}
+TEST_F(TestRecordBatch, RenameColumns) {
+ const int length = 10;
+
+ auto field1 = field("f1", int32());
+ auto field2 = field("f2", uint8());
+ auto field3 = field("f3", int16());
+
+ auto schema1 = ::arrow::schema({field1, field2, field3});
+
+ random::RandomArrayGenerator gen(42);
+
+ auto array1 = gen.ArrayOf(int32(), length);
+ auto array2 = gen.ArrayOf(uint8(), length);
+ auto array3 = gen.ArrayOf(int16(), length);
+
+ auto batch = RecordBatch::Make(schema1, length, {array1, array2, array3});
+ EXPECT_THAT(batch->ColumnNames(), testing::ElementsAre("f1", "f2", "f3"));
+
+ ASSERT_OK_AND_ASSIGN(auto renamed, batch->RenameColumns({"zero", "one", "two"}));
+ EXPECT_THAT(renamed->ColumnNames(), testing::ElementsAre("zero", "one", "two"));
+ EXPECT_THAT(renamed->columns(), testing::ElementsAre(array1, array2, array3));
+ ASSERT_OK(renamed->ValidateFull());
+
+ ASSERT_RAISES(Invalid, batch->RenameColumns({"hello", "world"}));
+}
+
TEST_F(TestRecordBatch, SelectColumns) {
const int length = 10;
diff --git a/python/pyarrow/includes/libarrow.pxd b/python/pyarrow/includes/libarrow.pxd
index e44c699fc5..9e5e3d3fa6 100644
--- a/python/pyarrow/includes/libarrow.pxd
+++ b/python/pyarrow/includes/libarrow.pxd
@@ -961,8 +961,15 @@ cdef extern from "arrow/api.h" namespace "arrow" nogil:
shared_ptr[CArray] column(int i)
const c_string& column_name(int i)
+ CResult[shared_ptr[CRecordBatch]] AddColumn(
+ int i, shared_ptr[CField] field, shared_ptr[CArray] column)
+ CResult[shared_ptr[CRecordBatch]] RemoveColumn(int i)
+ CResult[shared_ptr[CRecordBatch]] SetColumn(
+ int i, shared_ptr[CField] field, shared_ptr[CArray] column)
+
const vector[shared_ptr[CArray]]& columns()
+ CResult[shared_ptr[CRecordBatch]] RenameColumns(const vector[c_string]&)
CResult[shared_ptr[CRecordBatch]] SelectColumns(const vector[int]&)
int num_columns()
diff --git a/python/pyarrow/table.pxi b/python/pyarrow/table.pxi
index 9f60150427..2ec19db1d1 100644
--- a/python/pyarrow/table.pxi
+++ b/python/pyarrow/table.pxi
@@ -2150,6 +2150,117 @@ cdef class _Tabular(_PandasConvertible):
pieces.append('...')
return '\n'.join(pieces)
+ def remove_column(self, int i):
+ # implemented in RecordBatch/Table subclasses
+ raise NotImplementedError
+
+ def drop_columns(self, columns):
+ """
+ Drop one or more columns and return a new Table or RecordBatch.
+
+ Parameters
+ ----------
+ columns : str or list[str]
+ Field name(s) referencing existing column(s).
+
+ Raises
+ ------
+ KeyError
+ If any of the passed column names do not exist.
+
+ Returns
+ -------
+ Table or RecordBatch
+ A tabular object without the column(s).
+
+ Examples
+ --------
+ Table (works similarly for RecordBatch)
+
+ >>> import pyarrow as pa
+ >>> import pandas as pd
+ >>> df = pd.DataFrame({'n_legs': [2, 4, 5, 100],
+ ... 'animals': ["Flamingo", "Horse", "Brittle stars", "Centipede"]})
+ >>> table = pa.Table.from_pandas(df)
+
+ Drop one column:
+
+ >>> table.drop_columns("animals")
+ pyarrow.Table
+ n_legs: int64
+ ----
+ n_legs: [[2,4,5,100]]
+
+ Drop one or more columns:
+
+ >>> table.drop_columns(["n_legs", "animals"])
+ pyarrow.Table
+ ...
+ ----
+ """
+ if isinstance(columns, str):
+ columns = [columns]
+
+ indices = []
+ for col in columns:
+ idx = self.schema.get_field_index(col)
+ if idx == -1:
+ raise KeyError("Column {!r} not found".format(col))
+ indices.append(idx)
+
+ indices.sort()
+ indices.reverse()
+
+ res = self
+ for idx in indices:
+ res = res.remove_column(idx)
+
+ return res
+
+ def add_column(self, int i, field_, column):
+ # implemented in RecordBatch/Table subclasses
+ raise NotImplementedError
+
+ def append_column(self, field_, column):
+ """
+ Append column at end of columns.
+
+ Parameters
+ ----------
+ field_ : str or Field
+ If a string is passed then the type is deduced from the column
+ data.
+ column : Array or value coercible to array
+ Column data.
+
+ Returns
+ -------
+ Table or RecordBatch
+ New table or record batch with the passed column added.
+
+ Examples
+ --------
+ >>> import pyarrow as pa
+ >>> import pandas as pd
+ >>> df = pd.DataFrame({'n_legs': [2, 4, 5, 100],
+ ... 'animals': ["Flamingo", "Horse", "Brittle stars", "Centipede"]})
+ >>> table = pa.Table.from_pandas(df)
+
+ Append column at the end:
+
+ >>> year = [2021, 2022, 2019, 2021]
+ >>> table.append_column('year', [year])
+ pyarrow.Table
+ n_legs: int64
+ animals: string
+ year: int64
+ ----
+ n_legs: [[2,4,5,100]]
+ animals: [["Flamingo","Horse","Brittle stars","Centipede"]]
+ year: [[2021,2022,2019,2021]]
+ """
+ return self.add_column(self.num_columns, field_, column)
+
cdef class RecordBatch(_Tabular):
"""
@@ -2483,6 +2594,214 @@ cdef class RecordBatch(_Tabular):
def __sizeof__(self):
return super(RecordBatch, self).__sizeof__() + self.nbytes
+ def add_column(self, int i, field_, column):
+ """
+ Add column to RecordBatch at position i.
+
+ A new record batch is returned with the column added, the original record batch
+ object is left unchanged.
+
+ Parameters
+ ----------
+ i : int
+ Index to place the column at.
+ field_ : str or Field
+ If a string is passed then the type is deduced from the column
+ data.
+ column : Array or value coercible to array
+ Column data.
+
+ Returns
+ -------
+ RecordBatch
+ New record batch with the passed column added.
+
+ Examples
+ --------
+ >>> import pyarrow as pa
+ >>> import pandas as pd
+ >>> df = pd.DataFrame({'n_legs': [2, 4, 5, 100],
+ ... 'animals': ["Flamingo", "Horse", "Brittle stars", "Centipede"]})
+ >>> batch = pa.RecordBatch.from_pandas(df)
+
+ Add column:
+
+ >>> year = [2021, 2022, 2019, 2021]
+ >>> batch.add_column(0,"year", year)
+ pyarrow.RecordBatch
+ year: int64
+ n_legs: int64
+ animals: string
+ ----
+ year: [2021,2022,2019,2021]
+ n_legs: [2,4,5,100]
+ animals: ["Flamingo","Horse","Brittle stars","Centipede"]
+
+ Original record batch is left unchanged:
+
+ >>> batch
+ pyarrow.RecordBatch
+ n_legs: int64
+ animals: string
+ ----
+ n_legs: [2,4,5,100]
+ animals: ["Flamingo","Horse","Brittle stars","Centipede"]
+ """
+ cdef:
+ shared_ptr[CRecordBatch] c_batch
+ Field c_field
+ Array c_arr
+
+ if isinstance(column, Array):
+ c_arr = column
+ else:
+ c_arr = array(column)
+
+ if isinstance(field_, Field):
+ c_field = field_
+ else:
+ c_field = field(field_, c_arr.type)
+
+ with nogil:
+ c_batch = GetResultValue(self.batch.AddColumn(
+ i, c_field.sp_field, c_arr.sp_array))
+
+ return pyarrow_wrap_batch(c_batch)
+
+ def remove_column(self, int i):
+ """
+ Create new RecordBatch with the indicated column removed.
+
+ Parameters
+ ----------
+ i : int
+ Index of column to remove.
+
+ Returns
+ -------
+ Table
+ New record batch without the column.
+
+ Examples
+ --------
+ >>> import pyarrow as pa
+ >>> import pandas as pd
+ >>> df = pd.DataFrame({'n_legs': [2, 4, 5, 100],
+ ... 'animals': ["Flamingo", "Horse", "Brittle stars", "Centipede"]})
+ >>> batch = pa.RecordBatch.from_pandas(df)
+ >>> batch.remove_column(1)
+ pyarrow.RecordBatch
+ n_legs: int64
+ ----
+ n_legs: [2,4,5,100]
+ """
+ cdef shared_ptr[CRecordBatch] c_batch
+
+ with nogil:
+ c_batch = GetResultValue(self.batch.RemoveColumn(i))
+
+ return pyarrow_wrap_batch(c_batch)
+
+ def set_column(self, int i, field_, column):
+ """
+ Replace column in RecordBatch at position.
+
+ Parameters
+ ----------
+ i : int
+ Index to place the column at.
+ field_ : str or Field
+ If a string is passed then the type is deduced from the column
+ data.
+ column : Array or value coercible to array
+ Column data.
+
+ Returns
+ -------
+ RecordBatch
+ New record batch with the passed column set.
+
+ Examples
+ --------
+ >>> import pyarrow as pa
+ >>> import pandas as pd
+ >>> df = pd.DataFrame({'n_legs': [2, 4, 5, 100],
+ ... 'animals': ["Flamingo", "Horse", "Brittle stars", "Centipede"]})
+ >>> batch = pa.RecordBatch.from_pandas(df)
+
+ Replace a column:
+
+ >>> year = [2021, 2022, 2019, 2021]
+ >>> batch.set_column(1,'year', year)
+ pyarrow.RecordBatch
+ n_legs: int64
+ year: int64
+ ----
+ n_legs: [2,4,5,100]
+ year: [2021,2022,2019,2021]
+ """
+ cdef:
+ shared_ptr[CRecordBatch] c_batch
+ Field c_field
+ Array c_arr
+
+ if isinstance(column, Array):
+ c_arr = column
+ else:
+ c_arr = array(column)
+
+ if isinstance(field_, Field):
+ c_field = field_
+ else:
+ c_field = field(field_, c_arr.type)
+
+ with nogil:
+ c_batch = GetResultValue(self.batch.SetColumn(
+ i, c_field.sp_field, c_arr.sp_array))
+
+ return pyarrow_wrap_batch(c_batch)
+
+ def rename_columns(self, names):
+ """
+ Create new record batch with columns renamed to provided names.
+
+ Parameters
+ ----------
+ names : list of str
+ List of new column names.
+
+ Returns
+ -------
+ RecordBatch
+
+ Examples
+ --------
+ >>> import pyarrow as pa
+ >>> import pandas as pd
+ >>> df = pd.DataFrame({'n_legs': [2, 4, 5, 100],
+ ... 'animals': ["Flamingo", "Horse", "Brittle stars", "Centipede"]})
+ >>> batch = pa.RecordBatch.from_pandas(df)
+ >>> new_names = ["n", "name"]
+ >>> batch.rename_columns(new_names)
+ pyarrow.RecordBatch
+ n: int64
+ name: string
+ ----
+ n: [2,4,5,100]
+ name: ["Flamingo","Horse","Brittle stars","Centipede"]
+ """
+ cdef:
+ shared_ptr[CRecordBatch] c_batch
+ vector[c_string] c_names
+
+ for name in names:
+ c_names.push_back(tobytes(name))
+
+ with nogil:
+ c_batch = GetResultValue(self.batch.RenameColumns(move(c_names)))
+
+ return pyarrow_wrap_batch(c_batch)
+
def serialize(self, memory_pool=None):
"""
Write RecordBatch to Buffer as encapsulated IPC message, which does not
@@ -2744,7 +3063,7 @@ cdef class RecordBatch(_Tabular):
def cast(self, Schema target_schema, safe=None, options=None):
"""
- Cast batch values to another schema.
+ Cast record batch values to another schema.
Parameters
----------
@@ -2758,12 +3077,52 @@ cdef class RecordBatch(_Tabular):
Returns
-------
RecordBatch
+
+ Examples
+ --------
+ >>> import pyarrow as pa
+ >>> import pandas as pd
+ >>> df = pd.DataFrame({'n_legs': [2, 4, 5, 100],
+ ... 'animals': ["Flamingo", "Horse", "Brittle stars", "Centipede"]})
+ >>> batch = pa.RecordBatch.from_pandas(df)
+ >>> batch.schema
+ n_legs: int64
+ animals: string
+ -- schema metadata --
+ pandas: '{"index_columns": [{"kind": "range", "name": null, "start": 0, ...
+
+ Define new schema and cast batch values:
+
+ >>> my_schema = pa.schema([
+ ... pa.field('n_legs', pa.duration('s')),
+ ... pa.field('animals', pa.string())]
+ ... )
+ >>> batch.cast(target_schema=my_schema)
+ pyarrow.RecordBatch
+ n_legs: duration[s]
+ animals: string
+ ----
+ n_legs: [2,4,5,100]
+ animals: ["Flamingo","Horse","Brittle stars","Centipede"]
"""
- # Wrap the more general Table cast implementation
- tbl = Table.from_batches([self])
- casted_tbl = tbl.cast(target_schema, safe=safe, options=options)
- casted_batch, = casted_tbl.to_batches()
- return casted_batch
+ cdef:
+ Array column, casted
+ Field field
+ list newcols = []
+
+ if self.schema.names != target_schema.names:
+ raise ValueError("Target schema's field names are not matching "
+ "the record batch's field names: {!r}, {!r}"
+ .format(self.schema.names, target_schema.names))
+
+ for column, field in zip(self.itercolumns(), target_schema):
+ if not field.nullable and column.null_count > 0:
+ raise ValueError("Casting field {!r} with null values to non-nullable"
+ .format(field.name))
+ casted = column.cast(field.type, safe=safe, options=options)
+ newcols.append(casted)
+
+ return RecordBatch.from_arrays(newcols, schema=target_schema)
def _to_pandas(self, options, **kwargs):
return Table.from_batches([self])._to_pandas(options, **kwargs)
@@ -4664,46 +5023,6 @@ cdef class Table(_Tabular):
return pyarrow_wrap_table(c_table)
- def append_column(self, field_, column):
- """
- Append column at end of columns.
-
- Parameters
- ----------
- field_ : str or Field
- If a string is passed then the type is deduced from the column
- data.
- column : Array, list of Array, or values coercible to arrays
- Column data.
-
- Returns
- -------
- Table
- New table with the passed column added.
-
- Examples
- --------
- >>> import pyarrow as pa
- >>> import pandas as pd
- >>> df = pd.DataFrame({'n_legs': [2, 4, 5, 100],
- ... 'animals': ["Flamingo", "Horse", "Brittle stars", "Centipede"]})
- >>> table = pa.Table.from_pandas(df)
-
- Append column at the end:
-
- >>> year = [2021, 2022, 2019, 2021]
- >>> table.append_column('year', [year])
- pyarrow.Table
- n_legs: int64
- animals: string
- year: int64
- ----
- n_legs: [[2,4,5,100]]
- animals: [["Flamingo","Horse","Brittle stars","Centipede"]]
- year: [[2021,2022,2019,2021]]
- """
- return self.add_column(self.num_columns, field_, column)
-
def remove_column(self, int i):
"""
Create new Table with the indicated column removed.
@@ -4838,67 +5157,6 @@ cdef class Table(_Tabular):
return pyarrow_wrap_table(c_table)
- def drop_columns(self, columns):
- """
- Drop one or more columns and return a new table.
-
- Parameters
- ----------
- columns : str or list[str]
- Field name(s) referencing existing column(s).
-
- Raises
- ------
- KeyError
- If any of the passed column names do not exist.
-
- Returns
- -------
- Table
- New table without the column(s).
-
- Examples
- --------
- >>> import pyarrow as pa
- >>> import pandas as pd
- >>> df = pd.DataFrame({'n_legs': [2, 4, 5, 100],
- ... 'animals': ["Flamingo", "Horse", "Brittle stars", "Centipede"]})
- >>> table = pa.Table.from_pandas(df)
-
- Drop one column:
-
- >>> table.drop_columns("animals")
- pyarrow.Table
- n_legs: int64
- ----
- n_legs: [[2,4,5,100]]
-
- Drop one or more columns:
-
- >>> table.drop_columns(["n_legs", "animals"])
- pyarrow.Table
- ...
- ----
- """
- if isinstance(columns, str):
- columns = [columns]
-
- indices = []
- for col in columns:
- idx = self.schema.get_field_index(col)
- if idx == -1:
- raise KeyError("Column {!r} not found".format(col))
- indices.append(idx)
-
- indices.sort()
- indices.reverse()
-
- table = self
- for idx in indices:
- table = table.remove_column(idx)
-
- return table
-
def drop(self, columns):
"""
Drop one or more columns and return a new table.
@@ -5258,8 +5516,10 @@ def record_batch(data, names=None, schema=None, metadata=None):
# need to cast it if the producer wasn't able to cast to exact schema.
batch = batch.cast(schema)
return batch
+
elif _pandas_api.is_data_frame(data):
return RecordBatch.from_pandas(data, schema=schema)
+
else:
raise TypeError("Expected pandas DataFrame or list of arrays")
diff --git a/python/pyarrow/tests/test_table.py b/python/pyarrow/tests/test_table.py
index 87b17c3501..374a29e956 100644
--- a/python/pyarrow/tests/test_table.py
+++ b/python/pyarrow/tests/test_table.py
@@ -505,12 +505,14 @@ def test_recordbatch_basics():
batch.nbytes == (5 * 2) + (5 * 4 + 1)
assert sys.getsizeof(batch) >= object.__sizeof__(
batch) + batch.get_total_buffer_size()
+
pydict = batch.to_pydict()
assert pydict == OrderedDict([
('c0', [0, 1, 2, 3, 4]),
('c1', [-10, -5, 0, None, 10])
])
assert isinstance(pydict, dict)
+ assert batch == pa.record_batch(pydict, schema=batch.schema)
with pytest.raises(IndexError):
# bounds checking
@@ -1259,12 +1261,14 @@ def test_table_basics():
assert table.nbytes == 2 * (5 * 8)
assert sys.getsizeof(table) >= object.__sizeof__(
table) + table.get_total_buffer_size()
+
pydict = table.to_pydict()
assert pydict == OrderedDict([
('a', [0, 1, 2, 3, 4]),
('b', [-10, -5, 0, 5, 10])
])
assert isinstance(pydict, dict)
+ assert table == pa.table(pydict, schema=table.schema)
columns = []
for col in table.itercolumns():
@@ -1405,60 +1409,81 @@ def test_table_column_with_duplicates():
table.column('a')
-def test_table_add_column():
+@pytest.mark.parametrize(
+ ('cls'),
+ [
+ (pa.Table),
+ (pa.RecordBatch)
+ ]
+)
+def test_table_add_column(cls):
data = [
pa.array(range(5)),
pa.array([-10, -5, 0, 5, 10]),
pa.array(range(5, 10))
]
- table = pa.Table.from_arrays(data, names=('a', 'b', 'c'))
+ table = cls.from_arrays(data, names=('a', 'b', 'c'))
new_field = pa.field('d', data[1].type)
t2 = table.add_column(3, new_field, data[1])
t3 = table.append_column(new_field, data[1])
- expected = pa.Table.from_arrays(data + [data[1]],
- names=('a', 'b', 'c', 'd'))
+ expected = cls.from_arrays(data + [data[1]],
+ names=('a', 'b', 'c', 'd'))
assert t2.equals(expected)
assert t3.equals(expected)
t4 = table.add_column(0, new_field, data[1])
- expected = pa.Table.from_arrays([data[1]] + data,
- names=('d', 'a', 'b', 'c'))
+ expected = cls.from_arrays([data[1]] + data,
+ names=('d', 'a', 'b', 'c'))
assert t4.equals(expected)
-def test_table_set_column():
+@pytest.mark.parametrize(
+ ('cls'),
+ [
+ (pa.Table),
+ (pa.RecordBatch)
+ ]
+)
+def test_table_set_column(cls):
data = [
pa.array(range(5)),
pa.array([-10, -5, 0, 5, 10]),
pa.array(range(5, 10))
]
- table = pa.Table.from_arrays(data, names=('a', 'b', 'c'))
+ table = cls.from_arrays(data, names=('a', 'b', 'c'))
new_field = pa.field('d', data[1].type)
t2 = table.set_column(0, new_field, data[1])
expected_data = list(data)
expected_data[0] = data[1]
- expected = pa.Table.from_arrays(expected_data,
- names=('d', 'b', 'c'))
+ expected = cls.from_arrays(expected_data,
+ names=('d', 'b', 'c'))
assert t2.equals(expected)
-def test_table_drop_columns():
+@pytest.mark.parametrize(
+ ('cls'),
+ [
+ (pa.Table),
+ (pa.RecordBatch)
+ ]
+)
+def test_table_drop_columns(cls):
""" drop one or more columns given labels"""
a = pa.array(range(5))
b = pa.array([-10, -5, 0, 5, 10])
c = pa.array(range(5, 10))
- table = pa.Table.from_arrays([a, b, c], names=('a', 'b', 'c'))
+ table = cls.from_arrays([a, b, c], names=('a', 'b', 'c'))
t2 = table.drop_columns(['a', 'b'])
t3 = table.drop_columns('a')
- exp_t2 = pa.Table.from_arrays([c], names=('c',))
+ exp_t2 = cls.from_arrays([c], names=('c',))
assert exp_t2.equals(t2)
- exp_t3 = pa.Table.from_arrays([b, c], names=('b', 'c',))
+ exp_t3 = cls.from_arrays([b, c], names=('b', 'c',))
assert exp_t3.equals(t3)
# -- raise KeyError if column not in Table
@@ -1486,26 +1511,40 @@ def test_table_drop():
table.drop(['d'])
-def test_table_remove_column():
+@pytest.mark.parametrize(
+ ('cls'),
+ [
+ (pa.Table),
+ (pa.RecordBatch)
+ ]
+)
+def test_table_remove_column(cls):
data = [
pa.array(range(5)),
pa.array([-10, -5, 0, 5, 10]),
pa.array(range(5, 10))
]
- table = pa.Table.from_arrays(data, names=('a', 'b', 'c'))
+ table = cls.from_arrays(data, names=('a', 'b', 'c'))
t2 = table.remove_column(0)
t2.validate()
- expected = pa.Table.from_arrays(data[1:], names=('b', 'c'))
+ expected = cls.from_arrays(data[1:], names=('b', 'c'))
assert t2.equals(expected)
-def test_table_remove_column_empty():
+@pytest.mark.parametrize(
+ ('cls'),
+ [
+ (pa.Table),
+ (pa.RecordBatch)
+ ]
+)
+def test_table_remove_column_empty(cls):
# ARROW-1865
data = [
pa.array(range(5)),
]
- table = pa.Table.from_arrays(data, names=['a'])
+ table = cls.from_arrays(data, names=['a'])
t2 = table.remove_column(0)
t2.validate()
@@ -1533,20 +1572,27 @@ def test_empty_table():
assert table.equals(pa.Table.from_arrays([], []))
-def test_table_rename_columns():
+@pytest.mark.parametrize(
+ ('cls'),
+ [
+ (pa.Table),
+ (pa.RecordBatch)
+ ]
+)
+def test_table_rename_columns(cls):
data = [
pa.array(range(5)),
pa.array([-10, -5, 0, 5, 10]),
pa.array(range(5, 10))
]
- table = pa.Table.from_arrays(data, names=['a', 'b', 'c'])
+ table = cls.from_arrays(data, names=['a', 'b', 'c'])
assert table.column_names == ['a', 'b', 'c']
t2 = table.rename_columns(['eh', 'bee', 'sea'])
t2.validate()
assert t2.column_names == ['eh', 'bee', 'sea']
- expected = pa.Table.from_arrays(data, names=['eh', 'bee', 'sea'])
+ expected = cls.from_arrays(data, names=['eh', 'bee', 'sea'])
assert t2.equals(expected)
@@ -1769,12 +1815,19 @@ def test_table_negative_indexing():
table[4]
-def test_table_cast_to_incompatible_schema():
+@pytest.mark.parametrize(
+ ('cls'),
+ [
+ (pa.Table),
+ (pa.RecordBatch)
+ ]
+)
+def test_table_cast_to_incompatible_schema(cls):
data = [
pa.array(range(5)),
pa.array([-10, -5, 0, 5, 10]),
]
- table = pa.Table.from_arrays(data, names=tuple('ab'))
+ table = cls.from_arrays(data, names=tuple('ab'))
target_schema1 = pa.schema([
pa.field('A', pa.int32()),
@@ -1783,22 +1836,35 @@ def test_table_cast_to_incompatible_schema():
target_schema2 = pa.schema([
pa.field('a', pa.int32()),
])
- message = ("Target schema's field names are not matching the table's "
- "field names:.*")
+
+ if cls is pa.Table:
+ cls_name = 'table'
+ else:
+ cls_name = 'record batch'
+ message = ("Target schema's field names are not matching the "
+ f"{cls_name}'s field names:.*")
+
with pytest.raises(ValueError, match=message):
table.cast(target_schema1)
with pytest.raises(ValueError, match=message):
table.cast(target_schema2)
-def test_table_safe_casting():
+@pytest.mark.parametrize(
+ ('cls'),
+ [
+ (pa.Table),
+ (pa.RecordBatch)
+ ]
+)
+def test_table_safe_casting(cls):
data = [
pa.array(range(5), type=pa.int64()),
pa.array([-10, -5, 0, 5, 10], type=pa.int32()),
pa.array([1.0, 2.0, 3.0, 4.0, 5.0], type=pa.float64()),
pa.array(['ab', 'bc', 'cd', 'de', 'ef'], type=pa.string())
]
- table = pa.Table.from_arrays(data, names=tuple('abcd'))
+ table = cls.from_arrays(data, names=tuple('abcd'))
expected_data = [
pa.array(range(5), type=pa.int32()),
@@ -1806,7 +1872,7 @@ def test_table_safe_casting():
pa.array([1, 2, 3, 4, 5], type=pa.int64()),
pa.array(['ab', 'bc', 'cd', 'de', 'ef'], type=pa.string())
]
- expected_table = pa.Table.from_arrays(expected_data, names=tuple('abcd'))
+ expected_table = cls.from_arrays(expected_data, names=tuple('abcd'))
target_schema = pa.schema([
pa.field('a', pa.int32()),
@@ -1819,14 +1885,21 @@ def test_table_safe_casting():
assert casted_table.equals(expected_table)
-def test_table_unsafe_casting():
+@pytest.mark.parametrize(
+ ('cls'),
+ [
+ (pa.Table),
+ (pa.RecordBatch)
+ ]
+)
+def test_table_unsafe_casting(cls):
data = [
pa.array(range(5), type=pa.int64()),
pa.array([-10, -5, 0, 5, 10], type=pa.int32()),
pa.array([1.1, 2.2, 3.3, 4.4, 5.5], type=pa.float64()),
pa.array(['ab', 'bc', 'cd', 'de', 'ef'], type=pa.string())
]
- table = pa.Table.from_arrays(data, names=tuple('abcd'))
+ table = cls.from_arrays(data, names=tuple('abcd'))
expected_data = [
pa.array(range(5), type=pa.int32()),
@@ -1834,7 +1907,7 @@ def test_table_unsafe_casting():
pa.array([1, 2, 3, 4, 5], type=pa.int64()),
pa.array(['ab', 'bc', 'cd', 'de', 'ef'], type=pa.string())
]
- expected_table = pa.Table.from_arrays(expected_data, names=tuple('abcd'))
+ expected_table = cls.from_arrays(expected_data, names=tuple('abcd'))
target_schema = pa.schema([
pa.field('a', pa.int32()),
@@ -2707,20 +2780,34 @@ def test_table_join_many_columns():
})
-def test_table_cast_invalid():
+@pytest.mark.parametrize(
+ ('cls'),
+ [
+ (pa.Table),
+ (pa.RecordBatch)
+ ]
+)
+def test_table_cast_invalid(cls):
# Casting a nullable field to non-nullable should be invalid!
- table = pa.table({'a': [None, 1], 'b': [None, True]})
+ table = cls.from_pydict({'a': [None, 1], 'b': [None, True]})
new_schema = pa.schema([pa.field("a", "int64", nullable=True),
pa.field("b", "bool", nullable=False)])
with pytest.raises(ValueError):
table.cast(new_schema)
- table = pa.table({'a': [None, 1], 'b': [False, True]})
+ table = cls.from_pydict({'a': [None, 1], 'b': [False, True]})
assert table.cast(new_schema).schema == new_schema
-def test_table_sort_by():
- table = pa.table([
+@pytest.mark.parametrize(
+ ('cls'),
+ [
+ (pa.Table),
+ (pa.RecordBatch)
+ ]
+)
+def test_table_sort_by(cls):
+ table = cls.from_arrays([
pa.array([3, 1, 4, 2, 5]),
pa.array(["b", "a", "b", "a", "c"]),
], names=["values", "keys"])
@@ -2735,7 +2822,7 @@ def test_table_sort_by():
"values": [5, 4, 3, 2, 1]
}
- tab = pa.Table.from_arrays([
+ tab = cls.from_arrays([
pa.array([5, 7, 7, 35], type=pa.int64()),
pa.array(["foo", "car", "bar", "foobar"])
], names=["a", "b"])
@@ -2751,45 +2838,6 @@ def test_table_sort_by():
assert sorted_tab_dict["b"] == ["foo", "car", "bar", "foobar"]
-def test_record_batch_sort():
- rb = pa.RecordBatch.from_arrays([
- pa.array([7, 35, 7, 5], type=pa.int64()),
- pa.array([4, 1, 3, 2], type=pa.int64()),
- pa.array(["foo", "car", "bar", "foobar"])
- ], names=["a", "b", "c"])
-
- sorted_rb = rb.sort_by([("a", "descending"), ("b", "descending")])
- sorted_rb_dict = sorted_rb.to_pydict()
- assert sorted_rb_dict["a"] == [35, 7, 7, 5]
- assert sorted_rb_dict["b"] == [1, 4, 3, 2]
- assert sorted_rb_dict["c"] == ["car", "foo", "bar", "foobar"]
-
- sorted_rb = rb.sort_by([("a", "ascending"), ("b", "ascending")])
- sorted_rb_dict = sorted_rb.to_pydict()
- assert sorted_rb_dict["a"] == [5, 7, 7, 35]
- assert sorted_rb_dict["b"] == [2, 3, 4, 1]
- assert sorted_rb_dict["c"] == ["foobar", "bar", "foo", "car"]
-
-
-def test_record_batch_cast():
- rb = pa.RecordBatch.from_arrays([
- pa.array([None, 1]),
- pa.array([False, True])
- ], names=["a", "b"])
- new_schema = pa.schema([pa.field("a", "int64", nullable=True),
- pa.field("b", "bool", nullable=False)])
-
- assert rb.cast(new_schema).schema == new_schema
-
- # Casting a nullable field to non-nullable is invalid
- rb = pa.RecordBatch.from_arrays([
- pa.array([None, 1]),
- pa.array([None, True])
- ], names=["a", "b"])
- with pytest.raises(ValueError):
- rb.cast(new_schema)
-
-
@pytest.mark.parametrize("constructor", [pa.table, pa.record_batch])
def test_numpy_asarray(constructor):
table = constructor([[1, 2, 3], [4.0, 5.0, 6.0]], names=["a", "b"])