This is an automated email from the ASF dual-hosted git repository.
lzljs3620320 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/paimon-python.git
The following commit(s) were added to refs/heads/main by this push:
new 7703bcd Add tests for new reader and writer api (#17)
7703bcd is described below
commit 7703bcdc8ecf0769fd478259665159cc6666695e
Author: yuzelin <[email protected]>
AuthorDate: Wed Oct 9 18:04:16 2024 +0800
Add tests for new reader and writer api (#17)
---
paimon_python_java/tests/test_write_and_read.py | 78 +++++++++++++++++++++++++
1 file changed, 78 insertions(+)
diff --git a/paimon_python_java/tests/test_write_and_read.py
b/paimon_python_java/tests/test_write_and_read.py
index 06cadd9..c4e78c9 100644
--- a/paimon_python_java/tests/test_write_and_read.py
+++ b/paimon_python_java/tests/test_write_and_read.py
@@ -219,3 +219,81 @@ class TableWriteReadTest(unittest.TestCase):
# check data (ignore index)
pd.testing.assert_frame_equal(
result.reset_index(drop=True), expected_df.reset_index(drop=True))
+
+ def testAllWriteAndReadApi(self):
+ schema = Schema(self.simple_pa_schema)
+ self.catalog.create_table('default.test_all_api', schema, False)
+ table = self.catalog.get_table('default.test_all_api')
+ write_builder = table.new_batch_write_builder()
+
+ # write_arrow
+ table_write = write_builder.new_write()
+ table_commit = write_builder.new_commit()
+ data1 = {
+ 'f0': [1, 2, 3],
+ 'f1': ['a', 'b', 'c'],
+ }
+ pa_table = pa.Table.from_pydict(data1, schema=self.simple_pa_schema)
+ table_write.write_arrow(pa_table)
+ table_commit.commit(table_write.prepare_commit())
+ table_write.close()
+ table_commit.close()
+
+ # write_arrow_batch
+ table_write = write_builder.new_write()
+ table_commit = write_builder.new_commit()
+ data2 = {
+ 'f0': [4, 5, 6],
+ 'f1': ['d', 'e', 'f'],
+ }
+ df = pd.DataFrame(data2)
+ record_batch = pa.RecordBatch.from_pandas(df,
schema=self.simple_pa_schema)
+ table_write.write_arrow_batch(record_batch)
+ table_commit.commit(table_write.prepare_commit())
+ table_write.close()
+ table_commit.close()
+
+ # write_pandas
+ table_write = write_builder.new_write()
+ table_commit = write_builder.new_commit()
+ data3 = {
+ 'f0': [7, 8, 9],
+ 'f1': ['g', 'h', 'i'],
+ }
+ df = pd.DataFrame(data3)
+ table_write.write_pandas(df)
+ table_commit.commit(table_write.prepare_commit())
+ table_write.close()
+ table_commit.close()
+
+ read_builder = table.new_read_builder()
+ table_scan = read_builder.new_scan()
+ table_read = read_builder.new_read()
+ splits = table_scan.plan().splits()
+
+ # to_arrow
+ actual = table_read.to_arrow(splits)
+ expected = pa.Table.from_pydict({
+ 'f0': [1, 2, 3, 4, 5, 6, 7, 8, 9],
+ 'f1': ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i'],
+ }, schema=self.simple_pa_schema)
+ self.assertEqual(actual, expected)
+
+ # to_arrow_batch_reader
+ data_frames = [
+ batch.to_pandas()
+ for batch in table_read.to_arrow_batch_reader(splits)
+ ]
+ actual = pd.concat(data_frames)
+ expected = pd.DataFrame({
+ 'f0': [1, 2, 3, 4, 5, 6, 7, 8, 9],
+ 'f1': ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i'],
+ })
+ expected['f0'] = expected['f0'].astype('int32')
+ pd.testing.assert_frame_equal(
+ actual.reset_index(drop=True), expected.reset_index(drop=True))
+
+ # to_pandas
+ actual = table_read.to_pandas(splits)
+ pd.testing.assert_frame_equal(
+ actual.reset_index(drop=True), expected.reset_index(drop=True))
