| ################################################################################ |
| # 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 tempfile |
| import unittest |
| |
| import pandas as pd |
| import pyarrow as pa |
| |
| from pypaimon import CatalogFactory, Schema |
| from pypaimon.table.row.generic_row import GenericRowDeserializer |
| |
| |
| def _check_filtered_result(read_builder, expected_df): |
| scan = read_builder.new_scan() |
| read = read_builder.new_read() |
| actual_df = read.to_pandas(scan.plan().splits()) |
| pd.testing.assert_frame_equal( |
| actual_df.reset_index(drop=True), expected_df.reset_index(drop=True)) |
| |
| |
| def _random_format(): |
| return random.choice(['parquet', 'avro', 'orc']) |
| |
| |
| class PredicateTest(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([ |
| ('f0', pa.int64()), |
| ('f1', pa.string()), |
| ]) |
| cls.catalog.create_table('default.test_append', Schema.from_pyarrow_schema( |
| pa_schema, options={'file.format': _random_format()}), False) |
| cls.catalog.create_table('default.test_pk', Schema.from_pyarrow_schema( |
| pa_schema, primary_keys=['f0'], options={'bucket': '1', 'file.format': _random_format()}), False) |
| |
| df = pd.DataFrame({ |
| 'f0': [1, 2, 3, 4, 5], |
| 'f1': ['abc', 'abbc', 'bc', 'd', None], |
| }) |
| |
| 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_pandas(df) |
| 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_pandas(df) |
| commit.commit(write.prepare_commit()) |
| write.close() |
| commit.close() |
| |
| cls.df = df |
| |
| def test_wrong_field_name(self): |
| table = self.catalog.get_table('default.test_append') |
| predicate_builder = table.new_read_builder().new_predicate_builder() |
| with self.assertRaises(ValueError) as e: |
| predicate_builder.equal('f2', 'a') |
| self.assertEqual(str(e.exception), "The field f2 is not in field list ['f0', 'f1'].") |
| |
| def test_append_with_duplicate(self): |
| pa_schema = pa.schema([ |
| ('f0', pa.int64()), |
| ('f1', pa.string()), |
| ]) |
| self.catalog.create_table('default.test_append_with_duplicate', Schema.from_pyarrow_schema(pa_schema), False) |
| |
| df = pd.DataFrame({ |
| 'f0': [1, 1, 2, 2], |
| 'f1': ['a', 'b', 'c', 'd'], |
| }) |
| |
| table = self.catalog.get_table('default.test_append_with_duplicate') |
| write_builder = table.new_batch_write_builder() |
| write = write_builder.new_write() |
| commit = write_builder.new_commit() |
| write.write_pandas(df) |
| commit.commit(write.prepare_commit()) |
| write.close() |
| commit.close() |
| |
| predicate_builder = table.new_read_builder().new_predicate_builder() |
| |
| predicate = predicate_builder.equal('f0', 1) |
| _check_filtered_result(table.new_read_builder().with_filter(predicate), df.loc[0:1]) |
| |
| predicate = predicate_builder.equal('f0', 0) |
| read_builder = table.new_read_builder().with_filter(predicate) |
| scan = read_builder.new_scan() |
| read = read_builder.new_read() |
| actual_df = read.to_pandas(scan.plan().splits()) |
| self.assertEqual(len(actual_df), 0) |
| |
| def test_all_field_types_with_equal(self): |
| pa_schema = pa.schema([ |
| # int |
| ('_tinyint', pa.int8()), |
| ('_smallint', pa.int16()), |
| ('_int', pa.int32()), |
| ('_bigint', pa.int64()), |
| # float |
| ('_float16', pa.float32()), # NOTE: cannot write pa.float16() data into Paimon |
| ('_float32', pa.float32()), |
| ('_double', pa.float64()), |
| # string |
| ('_string', pa.string()), |
| # bool |
| ('_boolean', pa.bool_()) |
| ]) |
| self.catalog.create_table('default.test_all_field_types', |
| Schema.from_pyarrow_schema(pa_schema, options={'file.format': _random_format()}), |
| False) |
| table = self.catalog.get_table('default.test_all_field_types') |
| write_builder = table.new_batch_write_builder() |
| write = write_builder.new_write() |
| commit = write_builder.new_commit() |
| |
| df = pd.DataFrame({ |
| '_tinyint': pd.Series([1, 2], dtype='int8'), |
| '_smallint': pd.Series([10, 20], dtype='int16'), |
| '_int': pd.Series([100, 200], dtype='int32'), |
| '_bigint': pd.Series([1000, 2000], dtype='int64'), |
| '_float16': pd.Series([1.0, 2.0], dtype='float16'), |
| '_float32': pd.Series([1.00, 2.00], dtype='float32'), |
| '_double': pd.Series([1.000, 2.000], dtype='double'), |
| '_string': pd.Series(['A', 'B'], dtype='object'), |
| '_boolean': [True, False] |
| }) |
| record_batch = pa.RecordBatch.from_pandas(df, schema=pa_schema) |
| # prepare for assertion |
| df['_float16'] = df['_float16'].astype('float32') |
| |
| write.write_arrow_batch(record_batch) |
| commit.commit(write.prepare_commit()) |
| write.close() |
| commit.close() |
| |
| predicate_builder = table.new_read_builder().new_predicate_builder() |
| |
| predicate = predicate_builder.equal('_tinyint', 1) |
| _check_filtered_result(table.new_read_builder().with_filter(predicate), df.loc[[0]]) |
| |
| predicate = predicate_builder.equal('_smallint', 20) |
| _check_filtered_result(table.new_read_builder().with_filter(predicate), df.loc[[1]]) |
| |
| predicate = predicate_builder.equal('_int', 100) |
| _check_filtered_result(table.new_read_builder().with_filter(predicate), df.loc[[0]]) |
| |
| predicate = predicate_builder.equal('_bigint', 2000) |
| _check_filtered_result(table.new_read_builder().with_filter(predicate), df.loc[[1]]) |
| |
| predicate = predicate_builder.equal('_float16', 1.0) |
| _check_filtered_result(table.new_read_builder().with_filter(predicate), df.loc[[0]]) |
| |
| predicate = predicate_builder.equal('_float32', 2.00) |
| _check_filtered_result(table.new_read_builder().with_filter(predicate), df.loc[[1]]) |
| |
| predicate = predicate_builder.equal('_double', 1.000) |
| _check_filtered_result(table.new_read_builder().with_filter(predicate), df.loc[[0]]) |
| |
| predicate = predicate_builder.equal('_string', 'B') |
| _check_filtered_result(table.new_read_builder().with_filter(predicate), df.loc[[1]]) |
| |
| predicate = predicate_builder.equal('_boolean', True) |
| _check_filtered_result(table.new_read_builder().with_filter(predicate), df.loc[[0]]) |
| |
| def test_equal_pk(self): |
| table = self.catalog.get_table('default.test_pk') |
| predicate_builder = table.new_read_builder().new_predicate_builder() |
| predicate = predicate_builder.equal('f0', 1) |
| _check_filtered_result(table.new_read_builder().with_filter(predicate), self.df.loc[[0]]) |
| |
| def test_not_equal_append(self): |
| table = self.catalog.get_table('default.test_append') |
| predicate_builder = table.new_read_builder().new_predicate_builder() |
| predicate = predicate_builder.not_equal('f0', 1) |
| _check_filtered_result(table.new_read_builder().with_filter(predicate), self.df.loc[1:4]) |
| |
| def test_not_equal_pk(self): |
| table = self.catalog.get_table('default.test_pk') |
| predicate_builder = table.new_read_builder().new_predicate_builder() |
| predicate = predicate_builder.not_equal('f0', 1) |
| _check_filtered_result(table.new_read_builder().with_filter(predicate), self.df.loc[1:4]) |
| |
| def test_less_than_append(self): |
| table = self.catalog.get_table('default.test_append') |
| predicate_builder = table.new_read_builder().new_predicate_builder() |
| predicate = predicate_builder.less_than('f0', 3) |
| _check_filtered_result(table.new_read_builder().with_filter(predicate), self.df.loc[0:1]) |
| |
| def test_less_than_pk(self): |
| table = self.catalog.get_table('default.test_pk') |
| predicate_builder = table.new_read_builder().new_predicate_builder() |
| predicate = predicate_builder.less_than('f0', 3) |
| _check_filtered_result(table.new_read_builder().with_filter(predicate), self.df.loc[0:1]) |
| |
| def test_less_or_equal_append(self): |
| table = self.catalog.get_table('default.test_append') |
| predicate_builder = table.new_read_builder().new_predicate_builder() |
| predicate = predicate_builder.less_or_equal('f0', 3) |
| _check_filtered_result(table.new_read_builder().with_filter(predicate), self.df.loc[0:2]) |
| |
| def test_less_or_equal_pk(self): |
| table = self.catalog.get_table('default.test_pk') |
| predicate_builder = table.new_read_builder().new_predicate_builder() |
| predicate = predicate_builder.less_or_equal('f0', 3) |
| _check_filtered_result(table.new_read_builder().with_filter(predicate), self.df.loc[0:2]) |
| |
| def test_greater_than_append(self): |
| table = self.catalog.get_table('default.test_append') |
| predicate_builder = table.new_read_builder().new_predicate_builder() |
| predicate = predicate_builder.greater_than('f0', 3) |
| _check_filtered_result(table.new_read_builder().with_filter(predicate), self.df.loc[3:4]) |
| |
| def test_greater_than_pk(self): |
| table = self.catalog.get_table('default.test_pk') |
| predicate_builder = table.new_read_builder().new_predicate_builder() |
| predicate = predicate_builder.greater_than('f0', 3) |
| _check_filtered_result(table.new_read_builder().with_filter(predicate), self.df.loc[3:4]) |
| |
| def test_greater_or_equal_append(self): |
| table = self.catalog.get_table('default.test_append') |
| predicate_builder = table.new_read_builder().new_predicate_builder() |
| predicate = predicate_builder.greater_or_equal('f0', 3) |
| _check_filtered_result(table.new_read_builder().with_filter(predicate), self.df.loc[2:4]) |
| |
| def test_greater_or_equal_pk(self): |
| table = self.catalog.get_table('default.test_pk') |
| predicate_builder = table.new_read_builder().new_predicate_builder() |
| predicate = predicate_builder.greater_or_equal('f0', 3) |
| _check_filtered_result(table.new_read_builder().with_filter(predicate), self.df.loc[2:4]) |
| |
| def test_is_null_append(self): |
| table = self.catalog.get_table('default.test_append') |
| predicate_builder = table.new_read_builder().new_predicate_builder() |
| predicate = predicate_builder.is_null('f1') |
| _check_filtered_result(table.new_read_builder().with_filter(predicate), self.df.loc[[4]]) |
| |
| def test_is_null_pk(self): |
| table = self.catalog.get_table('default.test_pk') |
| predicate_builder = table.new_read_builder().new_predicate_builder() |
| predicate = predicate_builder.is_null('f1') |
| _check_filtered_result(table.new_read_builder().with_filter(predicate), self.df.loc[[4]]) |
| |
| def test_is_not_null_append(self): |
| table = self.catalog.get_table('default.test_append') |
| predicate_builder = table.new_read_builder().new_predicate_builder() |
| predicate = predicate_builder.is_not_null('f1') |
| _check_filtered_result(table.new_read_builder().with_filter(predicate), self.df.loc[0:3]) |
| |
| def test_is_not_null_pk(self): |
| table = self.catalog.get_table('default.test_pk') |
| predicate_builder = table.new_read_builder().new_predicate_builder() |
| predicate = predicate_builder.is_not_null('f1') |
| _check_filtered_result(table.new_read_builder().with_filter(predicate), self.df.loc[0:3]) |
| |
| def test_startswith_append(self): |
| table = self.catalog.get_table('default.test_append') |
| predicate_builder = table.new_read_builder().new_predicate_builder() |
| predicate = predicate_builder.startswith('f1', 'ab') |
| _check_filtered_result(table.new_read_builder().with_filter(predicate), self.df.loc[0:1]) |
| |
| def test_startswith_pk(self): |
| table = self.catalog.get_table('default.test_pk') |
| predicate_builder = table.new_read_builder().new_predicate_builder() |
| predicate = predicate_builder.startswith('f1', 'ab') |
| _check_filtered_result(table.new_read_builder().with_filter(predicate), self.df.loc[0:1]) |
| |
| def test_endswith_append(self): |
| table = self.catalog.get_table('default.test_append') |
| predicate_builder = table.new_read_builder().new_predicate_builder() |
| predicate = predicate_builder.endswith('f1', 'bc') |
| _check_filtered_result(table.new_read_builder().with_filter(predicate), self.df.loc[0:2]) |
| |
| def test_endswith_pk(self): |
| table = self.catalog.get_table('default.test_pk') |
| predicate_builder = table.new_read_builder().new_predicate_builder() |
| predicate = predicate_builder.endswith('f1', 'bc') |
| _check_filtered_result(table.new_read_builder().with_filter(predicate), self.df.loc[0:2]) |
| |
| def test_contains_append(self): |
| table = self.catalog.get_table('default.test_append') |
| predicate_builder = table.new_read_builder().new_predicate_builder() |
| predicate = predicate_builder.contains('f1', 'bb') |
| _check_filtered_result(table.new_read_builder().with_filter(predicate), self.df.loc[[1]]) |
| |
| def test_contains_pk(self): |
| table = self.catalog.get_table('default.test_pk') |
| predicate_builder = table.new_read_builder().new_predicate_builder() |
| predicate = predicate_builder.contains('f1', 'bb') |
| _check_filtered_result(table.new_read_builder().with_filter(predicate), self.df.loc[[1]]) |
| |
| def test_is_in_append(self): |
| table = self.catalog.get_table('default.test_append') |
| predicate_builder = table.new_read_builder().new_predicate_builder() |
| predicate = predicate_builder.is_in('f0', [1, 2]) |
| _check_filtered_result(table.new_read_builder().with_filter(predicate), self.df.loc[0:1]) |
| |
| def test_is_in_pk(self): |
| table = self.catalog.get_table('default.test_pk') |
| predicate_builder = table.new_read_builder().new_predicate_builder() |
| predicate = predicate_builder.is_in('f1', ['abc', 'd']) |
| _check_filtered_result(table.new_read_builder().with_filter(predicate), self.df.loc[[0, 3]]) |
| |
| def test_is_not_in_append(self): |
| table = self.catalog.get_table('default.test_append') |
| predicate_builder = table.new_read_builder().new_predicate_builder() |
| predicate = predicate_builder.is_not_in('f0', [1, 2]) |
| _check_filtered_result(table.new_read_builder().with_filter(predicate), self.df.loc[2:4]) |
| |
| def test_is_not_in_pk(self): |
| table = self.catalog.get_table('default.test_pk') |
| predicate_builder = table.new_read_builder().new_predicate_builder() |
| predicate = predicate_builder.is_not_in('f1', ['abc', 'abbc']) |
| _check_filtered_result(table.new_read_builder().with_filter(predicate), self.df.loc[2:4]) |
| |
| def test_between_append(self): |
| table = self.catalog.get_table('default.test_append') |
| predicate_builder = table.new_read_builder().new_predicate_builder() |
| predicate = predicate_builder.between('f0', 1, 3) |
| _check_filtered_result(table.new_read_builder().with_filter(predicate), self.df.loc[0:2]) |
| |
| def test_between_pk(self): |
| table = self.catalog.get_table('default.test_pk') |
| predicate_builder = table.new_read_builder().new_predicate_builder() |
| predicate = predicate_builder.between('f0', 1, 3) |
| _check_filtered_result(table.new_read_builder().with_filter(predicate), self.df.loc[0:2]) |
| |
| def test_and_predicates(self): |
| table = self.catalog.get_table('default.test_append') |
| predicate_builder = table.new_read_builder().new_predicate_builder() |
| predicate1 = predicate_builder.greater_than('f0', 1) |
| predicate2 = predicate_builder.startswith('f1', 'ab') |
| predicate = predicate_builder.and_predicates([predicate1, predicate2]) |
| _check_filtered_result(table.new_read_builder().with_filter(predicate), self.df.loc[[1]]) |
| |
| def test_or_predicates(self): |
| table = self.catalog.get_table('default.test_append') |
| predicate_builder = table.new_read_builder().new_predicate_builder() |
| predicate1 = predicate_builder.greater_than('f0', 3) |
| predicate2 = predicate_builder.less_than('f0', 2) |
| predicate = predicate_builder.or_predicates([predicate1, predicate2]) |
| _check_filtered_result(table.new_read_builder().with_filter(predicate), |
| self.df.loc[[0, 3, 4]]) |
| |
| def test_pk_reader_with_filter(self): |
| pa_schema = pa.schema([ |
| pa.field('key1', pa.int32(), nullable=False), |
| pa.field('key2', pa.string(), nullable=False), |
| ('behavior', pa.string()), |
| pa.field('dt1', pa.string(), nullable=False), |
| pa.field('dt2', pa.int32(), nullable=False) |
| ]) |
| schema = Schema.from_pyarrow_schema(pa_schema, |
| partition_keys=['dt1', 'dt2'], |
| primary_keys=['key1', 'key2'], |
| options={'bucket': '1'}) |
| self.catalog.create_table('default.test_pk_filter', schema, False) |
| table = self.catalog.get_table('default.test_pk_filter') |
| |
| write_builder = table.new_batch_write_builder() |
| table_write = write_builder.new_write() |
| table_commit = write_builder.new_commit() |
| data1 = { |
| 'key1': [1, 2, 3, 4], |
| 'key2': ['h', 'g', 'f', 'e'], |
| 'behavior': ['a', 'b', 'c', None], |
| 'dt1': ['p1', 'p1', 'p2', 'p1'], |
| 'dt2': [2, 2, 1, 2], |
| } |
| pa_table = pa.Table.from_pydict(data1, schema=pa_schema) |
| table_write.write_arrow(pa_table) |
| table_commit.commit(table_write.prepare_commit()) |
| table_write.close() |
| table_commit.close() |
| |
| table_write = write_builder.new_write() |
| table_commit = write_builder.new_commit() |
| data1 = { |
| 'key1': [5, 2, 7, 8], |
| 'key2': ['d', 'g', 'b', 'a'], |
| 'behavior': ['e', 'b-new', 'g', 'h'], |
| 'dt1': ['p2', 'p1', 'p1', 'p2'], |
| 'dt2': [2, 2, 1, 2] |
| } |
| pa_table = pa.Table.from_pydict(data1, schema=pa_schema) |
| table_write.write_arrow(pa_table) |
| table_commit.commit(table_write.prepare_commit()) |
| table_write.close() |
| table_commit.close() |
| |
| # test filter by partition |
| predicate_builder = table.new_read_builder().new_predicate_builder() |
| p1 = predicate_builder.startswith('dt1', "p1") |
| p2 = predicate_builder.is_in('dt1', ["p2"]) |
| p3 = predicate_builder.or_predicates([p1, p2]) |
| p4 = predicate_builder.equal('dt2', 2) |
| g1 = predicate_builder.and_predicates([p3, p4]) |
| # (dt1 startswith 'p1' or dt1 is_in ["p2"]) and dt2 == 2 |
| read_builder = table.new_read_builder().with_filter(g1) |
| splits = read_builder.new_scan().plan().splits() |
| self.assertEqual(len(splits), 2) |
| self.assertEqual(splits[0].partition.to_dict()["dt2"], 2) |
| self.assertEqual(splits[1].partition.to_dict()["dt2"], 2) |
| |
| # test filter by stats |
| predicate_builder = table.new_read_builder().new_predicate_builder() |
| p1 = predicate_builder.equal('key1', 7) |
| p2 = predicate_builder.is_in('key2', ["e", "f"]) |
| p3 = predicate_builder.or_predicates([p1, p2]) |
| p4 = predicate_builder.greater_than('key1', 3) |
| g1 = predicate_builder.and_predicates([p3, p4]) |
| # (key1 == 7 or key2 is_in ["e", "f"]) and key1 > 3 |
| read_builder = table.new_read_builder().with_filter(g1) |
| splits = read_builder.new_scan().plan().splits() |
| # initial splits meta: |
| # p1, 2 -> 2g, 2g; 1e, 4h |
| # p2, 1 -> 3f, 3f |
| # p2, 2 -> 5a, 8d |
| # p1, 1 -> 7b, 7b |
| self.assertEqual(len(splits), 3) |
| # expect to filter out `p1, 2 -> 2g, 2g` and `p2, 1 -> 3f, 3f` |
| count = 0 |
| for split in splits: |
| if split.partition.values == ["p1", 2]: |
| count += 1 |
| self.assertEqual(len(split.files), 1) |
| min_values = GenericRowDeserializer.from_bytes(split.files[0].key_stats.min_values.data, |
| table.primary_keys_fields).to_dict() |
| max_values = GenericRowDeserializer.from_bytes(split.files[0].key_stats.max_values.data, |
| table.primary_keys_fields).to_dict() |
| self.assertTrue(min_values["key1"] == 1 and min_values["key2"] == "e" |
| and max_values["key1"] == 4 and max_values["key2"] == "h") |
| elif split.partition.values == ["p2", 2]: |
| count += 1 |
| min_values = GenericRowDeserializer.from_bytes(split.files[0].key_stats.min_values.data, |
| table.primary_keys_fields).to_dict() |
| max_values = GenericRowDeserializer.from_bytes(split.files[0].key_stats.max_values.data, |
| table.primary_keys_fields).to_dict() |
| self.assertTrue(min_values["key1"] == 5 and min_values["key2"] == "a" |
| and max_values["key1"] == 8 and max_values["key2"] == "d") |
| elif split.partition.values == ["p1", 1]: |
| count += 1 |
| min_values = GenericRowDeserializer.from_bytes(split.files[0].key_stats.min_values.data, |
| table.primary_keys_fields).to_dict() |
| max_values = GenericRowDeserializer.from_bytes(split.files[0].key_stats.max_values.data, |
| table.primary_keys_fields).to_dict() |
| self.assertTrue(min_values["key1"] == max_values["key1"] == 7 |
| and max_values["key2"] == max_values["key2"] == "b") |
| self.assertEqual(count, 3) |
