| # Licensed to the Apache Software Foundation (ASF) under one | |
| # or more contributor license agreements. See the NOTICE file | |
| # distributed with this work for additional information | |
| # regarding copyright ownership. The ASF licenses this file | |
| # to you under the Apache License, Version 2.0 (the | |
| # "License"); you may not use this file except in compliance | |
| # with the License. You may obtain a copy of the License at | |
| # | |
| # http://www.apache.org/licenses/LICENSE-2.0 | |
| # | |
| # Unless required by applicable law or agreed to in writing, | |
| # software distributed under the License is distributed on an | |
| # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY | |
| # KIND, either express or implied. See the License for the | |
| # specific language governing permissions and limitations | |
| # under the License. | |
| import os | |
| import random | |
| import shutil | |
| import tempfile | |
| import unittest | |
| from typing import List | |
| import pyarrow as pa | |
| import pytest | |
| from pypaimon import CatalogFactory, Schema | |
| from pypaimon.manifest.schema.manifest_entry import ManifestEntry | |
| from pypaimon.manifest.schema.simple_stats import SimpleStats | |
| from pypaimon.read.scanner.file_scanner import FileScanner | |
| from pypaimon.table.row.generic_row import GenericRow, GenericRowDeserializer | |
| def _random_format(): | |
| return random.choice(['parquet', 'avro', 'orc']) | |
| class BinaryRowTest(unittest.TestCase): | |
| @classmethod | |
| def setUpClass(cls): | |
| cls.tempdir = tempfile.mkdtemp() | |
| cls.warehouse = os.path.join(cls.tempdir, 'warehouse') | |
| cls.catalog = CatalogFactory.create({'warehouse': cls.warehouse}) | |
| cls.catalog.create_database('default', False) | |
| pa_schema = pa.schema([ | |
| pa.field('f0', pa.int64(), nullable=False), | |
| ('f1', pa.string()), | |
| ('f2', pa.int64()), | |
| ]) | |
| cls.catalog.create_table('default.test_append', Schema.from_pyarrow_schema( | |
| pa_schema, partition_keys=['f0'], options={'file.format': _random_format()}), False) | |
| cls.catalog.create_table('default.test_pk', Schema.from_pyarrow_schema( | |
| pa_schema, partition_keys=['f2'], primary_keys=['f0'], | |
| options={'bucket': '1', 'file.format': _random_format()}), False) | |
| cls.data = pa.Table.from_pydict({ | |
| 'f0': [1, 2, 3, 4, 5], | |
| 'f1': ['abc', 'abbc', 'bc', 'd', None], | |
| 'f2': [6, 7, 8, 9, 10], | |
| }, schema=pa_schema) | |
| append_table = cls.catalog.get_table('default.test_append') | |
| write_builder = append_table.new_batch_write_builder() | |
| write = write_builder.new_write() | |
| commit = write_builder.new_commit() | |
| write.write_arrow(cls.data) | |
| commit.commit(write.prepare_commit()) | |
| write.close() | |
| commit.close() | |
| pk_table = cls.catalog.get_table('default.test_pk') | |
| write_builder = pk_table.new_batch_write_builder() | |
| write = write_builder.new_write() | |
| commit = write_builder.new_commit() | |
| write.write_arrow(cls.data) | |
| commit.commit(write.prepare_commit()) | |
| write.close() | |
| commit.close() | |
| @classmethod | |
| def tearDownClass(cls): | |
| shutil.rmtree(cls.tempdir, ignore_errors=True) | |
| @pytest.mark.python_plan | |
| def test_not_equal_append(self): | |
| table = self.catalog.get_table('default.test_append') | |
| self._overwrite_manifest_entry(table) | |
| read_builder = table.new_read_builder() | |
| predicate_builder = read_builder.new_predicate_builder() | |
| predicate = predicate_builder.not_equal('f2', 6) # test stats filter when filtering ManifestEntry | |
| splits, actual = self._read_result(read_builder.with_filter(predicate)) | |
| expected = self.data.slice(1, 4) | |
| self.assertEqual(expected, actual) | |
| self.assertEqual(len(expected), len(splits)) | |
| @pytest.mark.python_plan | |
| def test_less_than_append(self): | |
| table = self.catalog.get_table('default.test_append') | |
| self._overwrite_manifest_entry(table) | |
| read_builder = table.new_read_builder() | |
| predicate_builder = read_builder.new_predicate_builder() | |
| predicate = predicate_builder.less_than('f2', 8) | |
| splits, actual = self._read_result(read_builder.with_filter(predicate)) | |
| expected = self.data.slice(0, 2) | |
| self.assertEqual(actual, expected) | |
| self.assertEqual(len(expected), len(splits)) # test stats filter when filtering ManifestEntry | |
| def test_is_null_append(self): | |
| table = self.catalog.get_table('default.test_append') | |
| read_builder = table.new_read_builder() | |
| predicate_builder = read_builder.new_predicate_builder() | |
| predicate = predicate_builder.is_null('f1') # value_stats_cols=None | |
| splits, actual = self._read_result(read_builder.with_filter(predicate)) | |
| expected = self.data.slice(4, 1) | |
| self.assertEqual(expected, actual) | |
| self.assertEqual(len(expected), len(splits)) | |
| @pytest.mark.python_plan | |
| def test_is_not_null_append(self): | |
| table = self.catalog.get_table('default.test_append') | |
| file_scanner = FileScanner(table, lambda: ([], None)) | |
| latest_snapshot = file_scanner.snapshot_manager.get_latest_snapshot() | |
| manifest_files = file_scanner.manifest_list_manager.read_all(latest_snapshot) | |
| manifest_entries = file_scanner.manifest_file_manager.read(manifest_files[0].file_name) | |
| self._transform_manifest_entries(manifest_entries, []) | |
| l = ['abc', 'abbc', 'bc', 'd', None] | |
| for i, entry in enumerate(manifest_entries): | |
| entry.file.value_stats_cols = ['f1'] | |
| entry.file.value_stats = SimpleStats( | |
| GenericRow([l[i]], [table.fields[1]]), | |
| GenericRow([l[i]], [table.fields[1]]), | |
| [1 if l[i] is None else 0], | |
| ) | |
| file_scanner.manifest_file_manager.write(manifest_files[0].file_name, manifest_entries) | |
| read_builder = table.new_read_builder() | |
| predicate_builder = read_builder.new_predicate_builder() | |
| predicate = predicate_builder.is_not_null('f1') | |
| splits, actual = self._read_result(read_builder.with_filter(predicate)) | |
| expected = self.data.slice(0, 4) | |
| self.assertEqual(expected, actual) | |
| self.assertEqual(len(expected), len(splits)) | |
| @pytest.mark.python_plan | |
| def test_is_in_append(self): | |
| table = self.catalog.get_table('default.test_append') | |
| self._overwrite_manifest_entry(table) | |
| read_builder = table.new_read_builder() | |
| predicate_builder = read_builder.new_predicate_builder() | |
| predicate = predicate_builder.is_in('f2', [6, 8]) | |
| splits, actual = self._read_result(read_builder.with_filter(predicate)) | |
| expected = self.data.take([0, 2]) | |
| self.assertEqual(expected, actual) | |
| self.assertEqual(len(expected), len(splits)) | |
| def test_equal_pk(self): | |
| table = self.catalog.get_table('default.test_pk') | |
| read_builder = table.new_read_builder() | |
| predicate_builder = read_builder.new_predicate_builder() | |
| predicate = predicate_builder.equal('f2', 6) | |
| splits, actual = self._read_result(read_builder.with_filter(predicate)) | |
| expected = self.data.slice(0, 1) | |
| self.assertEqual(expected, actual) | |
| self.assertEqual(len(splits), len(expected)) # test partition filter when filtering ManifestEntry | |
| @pytest.mark.python_plan | |
| def test_not_equal_pk(self): | |
| table = self.catalog.get_table('default.test_pk') | |
| read_builder = table.new_read_builder() | |
| predicate_builder = read_builder.new_predicate_builder() | |
| predicate = predicate_builder.not_equal('f2', 6) | |
| splits, actual = self._read_result(read_builder.with_filter(predicate)) | |
| expected = self.data.slice(1, 4) | |
| self.assertEqual(actual, expected) | |
| self.assertEqual(len(splits), len(expected)) # test partition filter when filtering ManifestEntry | |
| @pytest.mark.python_plan | |
| def test_less_than_pk(self): | |
| table = self.catalog.get_table('default.test_pk') | |
| read_builder = table.new_read_builder() | |
| predicate_builder = read_builder.new_predicate_builder() | |
| predicate = predicate_builder.less_than('f0', 3) | |
| splits, actual = self._read_result(read_builder.with_filter(predicate)) | |
| expected = self.data.slice(0, 2) | |
| self.assertEqual(expected, actual) | |
| self.assertEqual(len(expected), len(splits)) # test key stats filter when filtering ManifestEntry | |
| def test_is_null_pk(self): | |
| table = self.catalog.get_table('default.test_pk') | |
| read_builder = table.new_read_builder() | |
| predicate_builder = read_builder.new_predicate_builder() | |
| predicate = predicate_builder.is_null('f1') | |
| splits, actual = self._read_result(read_builder.with_filter(predicate)) | |
| expected = self.data.slice(4, 1) | |
| self.assertEqual(actual, expected) | |
| @pytest.mark.python_plan | |
| def test_is_not_null_pk(self): | |
| table = self.catalog.get_table('default.test_pk') | |
| read_builder = table.new_read_builder() | |
| predicate_builder = read_builder.new_predicate_builder() | |
| predicate = predicate_builder.is_not_null('f1') | |
| splits, actual = self._read_result(read_builder.with_filter(predicate)) | |
| expected = self.data.slice(0, 4) | |
| self.assertEqual(actual, expected) | |
| @pytest.mark.python_plan | |
| def test_is_in_pk(self): | |
| table = self.catalog.get_table('default.test_pk') | |
| read_builder = table.new_read_builder() | |
| predicate_builder = read_builder.new_predicate_builder() | |
| predicate = predicate_builder.is_in('f0', [1, 5]) | |
| splits, actual = self._read_result(read_builder.with_filter(predicate)) | |
| # expected rows: indices [0, 3] | |
| expected = self.data.take([0, 4]) | |
| self.assertEqual(actual, expected) | |
| self.assertEqual(len(splits), len(expected)) # test key stats filter when filtering ManifestEntry | |
| @pytest.mark.python_plan | |
| def test_append_multi_cols(self): | |
| # Create a 10-column append table and write 10 rows | |
| pa_schema = pa.schema([ | |
| ('f0', pa.int64()), | |
| ('f1', pa.string()), | |
| ('f2', pa.int64()), | |
| ('f3', pa.string()), | |
| ('f4', pa.int64()), | |
| ('f5', pa.string()), | |
| ('f6', pa.int64()), | |
| ('f7', pa.string()), | |
| ('f8', pa.int64()), | |
| ('f9', pa.string()), | |
| ]) | |
| schema = Schema.from_pyarrow_schema( | |
| pa_schema, | |
| partition_keys=['f0'], | |
| options={'file.format': _random_format()} | |
| ) | |
| self.catalog.create_table('default.test_append_10cols', schema, False) | |
| table = self.catalog.get_table('default.test_append_10cols') | |
| write_builder = table.new_batch_write_builder() | |
| table_write = write_builder.new_write() | |
| table_commit = write_builder.new_commit() | |
| data = { | |
| 'f0': list(range(1, 11)), # 0..9 | |
| 'f1': ['a0', 'bb', 'a2', 'a3', 'a4', 'a5', 'a6', 'a7', 'a8', 'a9'], # contains 'bb' at index 1 | |
| 'f2': [10, 20, 30, 40, 50, 60, 70, 80, 90, 100], | |
| 'f3': ['x0', 'x1', 'x2', 'x3', 'x4', 'x5', 'x6', 'x7', 'x8', 'x9'], | |
| 'f4': [0, 1, 0, 1, 0, 1, 0, 1, 0, 1], | |
| 'f5': ['y0', 'y1', 'y2', 'y3', 'y4', 'y5', 'y6', 'y7', 'y8', 'y9'], | |
| 'f6': [100, 200, 300, 400, 500, 600, 700, 800, 900, 1000], | |
| 'f7': ['z0', 'z1', 'z2', 'z3', 'z4', 'z5', 'z6', 'z7', 'z8', 'z9'], | |
| 'f8': [5, 4, 3, 2, 1, 0, -1, -2, -3, -4], | |
| 'f9': ['w0', 'w1', 'w2', 'w3', 'w4', 'w5', 'w6', 'w7', 'w8', 'w9'], | |
| } | |
| pa_table = pa.Table.from_pydict(data, schema=pa_schema) | |
| table_write.write_arrow(pa_table) | |
| table_commit.commit(table_write.prepare_commit()) | |
| table_write.close() | |
| table_commit.close() | |
| file_scanner = FileScanner(table, lambda: ([], None)) | |
| latest_snapshot = file_scanner.snapshot_manager.get_latest_snapshot() | |
| manifest_files = file_scanner.manifest_list_manager.read_all(latest_snapshot) | |
| manifest_entries = file_scanner.manifest_file_manager.read(manifest_files[0].file_name) | |
| self._transform_manifest_entries(manifest_entries, []) | |
| for entry in manifest_entries: | |
| # Manifest entry order is not part of the file format. Derive the | |
| # synthetic statistics from this file's partition value. | |
| partition_id = entry.partition.values[0] | |
| entry.file.value_stats_cols = ['f2', 'f6', 'f8'] | |
| entry.file.value_stats = SimpleStats( | |
| GenericRow([10 * partition_id, 100 * partition_id, 6 - partition_id], | |
| [table.fields[2], table.fields[6], table.fields[8]]), | |
| GenericRow([10 * partition_id, 100 * partition_id, 6 - partition_id], | |
| [table.fields[2], table.fields[6], table.fields[8]]), | |
| [0, 0, 0], | |
| ) | |
| file_scanner.manifest_file_manager.write(manifest_files[0].file_name, manifest_entries) | |
| # Build multiple predicates and combine them | |
| read_builder = table.new_read_builder() | |
| predicate_builder = read_builder.new_predicate_builder() | |
| p_in = predicate_builder.is_in('f6', [100, 600, 1000]) | |
| p_contains = predicate_builder.less_or_equal('f8', -3) | |
| p_not_null = predicate_builder.is_not_null('f2') | |
| p_ge = predicate_builder.greater_or_equal('f2', 50) | |
| p_or = predicate_builder.or_predicates([p_in, p_contains]) | |
| combined = predicate_builder.and_predicates([p_or, p_not_null, p_ge]) | |
| splits, actual = self._read_result(read_builder.with_filter(combined)) | |
| # Expected rows after filter: indices 5 and 9 | |
| expected_data = {'f0': [6, 9, 10], | |
| 'f1': ['a5', 'a8', 'a9'], | |
| 'f2': [60, 90, 100], | |
| 'f3': ['x5', 'x8', 'x9'], | |
| 'f4': [1, 0, 1], | |
| 'f5': ['y5', 'y8', 'y9'], | |
| 'f6': [600, 900, 1000], | |
| 'f7': ['z5', 'z8', 'z9'], | |
| 'f8': [0, -3, -4], | |
| 'f9': ['w5', 'w8', 'w9'] | |
| } | |
| self.assertEqual(expected_data, actual.sort_by('f0').to_pydict()) | |
| file_scanner = FileScanner(table, lambda: ([], None)) | |
| latest_snapshot = file_scanner.snapshot_manager.get_latest_snapshot() | |
| manifest_files = file_scanner.manifest_list_manager.read_all(latest_snapshot) | |
| manifest_entries = file_scanner.manifest_file_manager.read(manifest_files[0].file_name) | |
| self._transform_manifest_entries(manifest_entries, []) | |
| for entry in manifest_entries: | |
| partition_id = entry.partition.values[0] | |
| entry.file.value_stats_cols = ['f2', 'f6', 'f8'] | |
| entry.file.value_stats = SimpleStats( | |
| GenericRow([0, 100 * partition_id, 6 - partition_id], | |
| [table.fields[2], table.fields[6], table.fields[8]]), | |
| GenericRow([0, 100 * partition_id, 6 - partition_id], | |
| [table.fields[2], table.fields[6], table.fields[8]]), | |
| [0, 0, 0], | |
| ) | |
| file_scanner.manifest_file_manager.write(manifest_files[0].file_name, manifest_entries) | |
| splits, actual = self._read_result(read_builder.with_filter(combined)) | |
| self.assertFalse(actual) | |
| def _read_result(self, read_builder): | |
| scan = read_builder.new_scan() | |
| read = read_builder.new_read() | |
| splits = scan.plan().splits() | |
| actual = read.to_arrow(splits) | |
| return splits, actual | |
| def _transform_manifest_entries(self, manifest_entries: List[ManifestEntry], trimmed_pk_fields): | |
| for entry in manifest_entries: | |
| entry.file.key_stats.min_values = GenericRowDeserializer.from_bytes(entry.file.key_stats.min_values.data, | |
| trimmed_pk_fields) | |
| entry.file.key_stats.max_values = GenericRowDeserializer.from_bytes(entry.file.key_stats.max_values.data, | |
| trimmed_pk_fields) | |
| def _overwrite_manifest_entry(self, table): | |
| file_scanner = FileScanner(table, lambda: ([], None)) | |
| latest_snapshot = file_scanner.snapshot_manager.get_latest_snapshot() | |
| manifest_files = file_scanner.manifest_list_manager.read_all(latest_snapshot) | |
| manifest_entries = file_scanner.manifest_file_manager.read(manifest_files[0].file_name) | |
| self._transform_manifest_entries(manifest_entries, []) | |
| for i, entry in enumerate(manifest_entries): | |
| entry.file.value_stats_cols = ['f2'] | |
| entry.file.value_stats = SimpleStats( | |
| GenericRow([6 + i], [table.fields[2]]), | |
| GenericRow([6 + i], [table.fields[2]]), | |
| [0], | |
| ) | |
| file_scanner.manifest_file_manager.write(manifest_files[0].file_name, manifest_entries) |