[python] Use Native for Java calendar query authorization masks

paimon-python/pypaimon/common/predicate_json_parser.py

@@ -54,13 +54,8 @@ def _wrap_int32(value: int) -> int:
     "TRAILING": (pc.utf8_rtrim, str.rstrip),
 }
 
-# Calendar-field extractions on a DATE / TIMESTAMP field, mirroring Java's
-# DateExtractTransform subclasses (name -> the Arrow kernel over a timestamp).
-# QUARTER matches (month - 1) / 3 + 1; the day-of-week variants match Java's
-# java.time numbering: WEEKDAY is Monday=0..Sunday=6, ISO_DAY_OF_WEEK is
-# Monday=1..Sunday=7, DAY_OF_WEEK is Sunday=1..Saturday=7. WEEK / YEAR_OF_WEEK
-# are the ISO week-of-week-based-year and the ISO week-based year. The rest map
-# one to one.
+# Java DateExtractTransform names. LTZ is excluded because a calendar field
+# would depend on a session time zone, which authorization rules do not carry.
 _DATE_EXTRACT = {
     "YEAR": pc.year,
     "MONTH": pc.month,
@@ -199,16 +194,93 @@ def _apply_predicate_transform(transform: dict, batch: pa.RecordBatch,
 
 
 def _date_extract(name: str, field_ref: dict, batch: pa.RecordBatch) -> pa.Array:
-    """Extract a calendar field from a DATE / TIMESTAMP column, as INT.
+    """Java calendar extraction on DATE days or Timestamp milliseconds.
 
-    A DATE is read at the start of its day (Java's LocalDate.atStartOfDay),
-    so it is cast to a timestamp first and HOUR/MINUTE/SECOND come out as 0.
-    Nulls propagate. Mirrors Java DateExtractTransform.
+    Timestamp.toLocalDateTime narrows the epoch day to a signed int. Work on
+    that representation directly, including DATE32 values which cannot be
+    cast to an Arrow microsecond timestamp or a Python datetime.
     """
+    from pypaimon.casting.temporal_utils import _signed
+    stored = field_ref.get('type')
+    if not isinstance(stored, str):
+        raise ValueError("Calendar transform field needs a DATE or TIMESTAMP type")
+    # Validate the complete stored SQL spelling before schema-name rebinding.
+    match = re.fullmatch(r"\s*(DATE(?=\s|\Z)|TIMESTAMP(?=\s|\(|\Z)(?:\s*\(\s*([0-9]+)\s*\))?"
+                         r"(?:\s*WITHOUT\s+TIME\s+ZONE(?=\s|\Z))?)"
+                         r"(?:\s*(?:NOT\s+)?NULL)?\s*", stored, flags=re.IGNORECASE)
+    if match is None or (match.group(2) is not None and int(match.group(2)) > 9):
+        raise ValueError("Calendar transforms require DATE or TIMESTAMP without time zone")
     column = _field_column(field_ref, batch)
-    if pa.types.is_date(column.type):
-        column = pc.cast(column, pa.timestamp("us"))
-    return pc.cast(_DATE_EXTRACT[name](column), pa.int32())
+    date = pa.types.is_date32(column.type)
+    timestamp = pa.types.is_timestamp(column.type) and column.type.tz is None
+    if not date and not timestamp:
+        raise ValueError("Calendar transforms require DATE or TIMESTAMP without time zone")
+    raw = column.cast(pa.int32() if date else pa.int64())
+    unit = column.type.unit if timestamp else None
+    # Arrow calendar kernels agree with Java within years 1..9999. Keep their
+    # vectorized execution for ordinary batches; outside this range use Java's
+    # epoch-day arithmetic, avoiding Arrow's bounded calendar representation.
+    bounds = pc.min_max(raw).as_py()
+    ticks_per_day = 1 if date else 86400 * {'s': 1, 'ms': 1000, 'us': 1000000, 'ns': 1000000000}[unit]
+    if (bounds['min'] is None or
+            (-719162 * ticks_per_day <= bounds['min'] and
+             bounds['max'] < 2932897 * ticks_per_day)):
+        source = column.cast(pa.timestamp('us')) if date else column
+        return _DATE_EXTRACT[name](source).cast(pa.int32())
+    values = raw.to_pylist()
+    results = []
+    for value in values:
+        if value is None:
+            results.append(None)
+            continue
+        if date:
+            day, millis = value, 0
+        else:
+            if unit == 's':
+                millis = _signed(value * 1000, 64)
+            elif unit == 'ms':
+                millis = value
+            else:
+                millis = value // (1000 if unit == 'us' else 1000000)
+            day, millis = _signed(millis // 86400000, 32), millis % 86400000
+        results.append(_calendar_field(name, day, millis))
+    return pa.array(results, type=pa.int32())
+
+
+def _calendar_field(name, day, millis):
+    from pypaimon.casting.temporal_utils import _epoch_day_to_date
+    year, month, date = _epoch_day_to_date(day)
+    weekday = (day + 3) % 7
+    if name == 'YEAR':
+        return year
+    if name == 'MONTH':
+        return month
+    if name == 'DAY':
+        return date
+    if name == 'HOUR':
+        return millis // 3600000
+    if name == 'MINUTE':
+        return millis // 60000 % 60
+    if name == 'SECOND':
+        return millis // 1000 % 60
+    if name == 'QUARTER':
+        return (month - 1) // 3 + 1
+    if name == 'DAY_OF_YEAR':
+        return day - _january_first(year) + 1
+    if name == 'WEEKDAY':
+        return weekday
+    if name == 'ISO_DAY_OF_WEEK':
+        return weekday + 1
+    if name == 'DAY_OF_WEEK':
+        return (weekday + 1) % 7 + 1
+    thursday = day + 3 - weekday
+    week_year = _epoch_day_to_date(thursday)[0]
+    return week_year if name == 'YEAR_OF_WEEK' else (thursday - _january_first(week_year)) // 7 + 1
+
+
+def _january_first(year):
+    era, yoe = divmod(year - 1, 400)
+    return era * 146097 + yoe * 365 + yoe // 4 - yoe // 100 + 306 - 719468
 
 
 def _substring(inputs, batch: pa.RecordBatch) -> pa.Array:
@@ -799,9 +871,7 @@ def _collect_all_field_refs_from_transform(transform: dict, fields: set = None)
     if fields is None:
         fields = set()
     name = transform.get("name")
-    if name == "FIELD_REF" and "fieldRef" in transform:
-        fields.add(transform["fieldRef"]["name"])
-    elif name == "CAST" and "fieldRef" in transform:
+    if (name in ("FIELD_REF", "CAST") or name in _DATE_EXTRACT) and "fieldRef" in transform:
         fields.add(transform["fieldRef"]["name"])
     else:
         for inp in transform.get("inputs", []):

paimon-python/pypaimon/tests/native_calendar_mask_test.py

@@ -0,0 +1,154 @@
+# 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.
+
+"""Existing calendar authorization masks through Native and Python readers."""
+
+import json
+from unittest.mock import patch
+
+import pyarrow as pa
+import pytest
+
+from pypaimon import Schema
+from pypaimon.api.api_response import AuthTableQueryResponse
+from pypaimon.common.predicate_builder import PredicateBuilder
+from pypaimon.read.native_plan import native_runtime_available
+from pypaimon.read.split_serializer import deserialize_split_v1, serialize_split_v1
+from pypaimon.schema.schema_change import SchemaChange
+from pypaimon.tests import native_query_auth_test
+
+rest_catalog = native_query_auth_test.rest_catalog
+_read = native_query_auth_test._read
+
+pytestmark = [pytest.mark.native_plan, pytest.mark.skipif(
+    not native_runtime_available(), reason='Rust reader required')]
+_PARTS = ['YEAR', 'MONTH', 'DAY', 'HOUR', 'MINUTE', 'SECOND', 'QUARTER', 'DAY_OF_YEAR',
+          'WEEKDAY', 'ISO_DAY_OF_WEEK', 'DAY_OF_WEEK', 'WEEK', 'YEAR_OF_WEEK']
+
+
+@pytest.fixture(params=[
+    ('append', {}), ('primary-key', {'bucket': '1'}),
+    ('data-evolution', {'row-tracking.enabled': 'true', 'data-evolution.enabled': 'true'}),
+], ids=lambda mode: mode[0])
+def calendar_table(rest_catalog, request):
+    catalog, server = rest_catalog
+    mode, options = request.param
+    schema = pa.schema([('id', pa.int32()), ('n', pa.int32()), ('d', pa.date32()),
+                        ('ts', pa.timestamp('ms')), ('us', pa.timestamp('us')), ('ns', pa.timestamp('ns'))])
+    catalog.create_table('default.calendar', Schema.from_pyarrow_schema(
+        schema, primary_keys=['id'] if mode == 'primary-key' else None, options=options), False)
+    table = catalog.get_table('default.calendar')
+    builder = table.new_batch_write_builder()
+    writer, commit = builder.new_write(), builder.new_commit()
+    try:
+        # Leap day, ISO year boundary, pre-epoch fractional second, NULL.
+        days = [19782, 20087, -1, None]
+        millis = [1709251199123, 1735516800000, -1, None]
+        writer.write_arrow(pa.table({
+            'id': [1, 2, 3, 4], 'n': [0] * 4, 'd': days, 'ts': millis,
+            'us': [value * 1000 + 456 if value is not None else None for value in millis],
+            'ns': [value * 1000000 + 456789 if value is not None else None for value in millis],
+        }, schema=schema))
+        commit.commit(writer.prepare_commit())
+    finally:
+        writer.close()
+        commit.close()
+    catalog.alter_table(table.identifier, [SchemaChange.set_option('query-auth.enabled', 'true')])
+    return catalog.get_table('default.calendar'), server
+
+
+def _rule(part, field, type_):
+    return json.dumps({'name': part, 'fieldRef': {'name': field, 'index': 99, 'type': type_}})
+
+
+@pytest.mark.parametrize('field,type_', [('d', 'DATE'), ('ts', 'TIMESTAMP(3)'),
+                                         ('us', 'TIMESTAMP(6)'), ('ns', 'TIMESTAMP(9)')])
+def test_all_calendar_masks_read_hidden_operands_natively(calendar_table, field, type_):
+    table, server = calendar_table
+    # Expected Java values for DATE midnight and TIMESTAMP's wall clock.
+    expected = [
+        [2024, 2024, 1969, None], [2, 12, 12, None], [29, 30, 31, None],
+        [0, 0, 0, None], [0, 0, 0, None], [0, 0, 0, None], [1, 4, 4, None],
+        [60, 365, 365, None], [3, 0, 2, None], [4, 1, 3, None], [5, 2, 4, None],
+        [9, 1, 1, None], [2024, 2025, 1970, None],
+    ]
+    if field != 'd':
+        expected[3:6] = [[23, 0, 23, None], [59, 0, 59, None], [59, 0, 59, None]]
+    for part, values in zip(_PARTS, expected):
+        server.table_query_auth_store[table.identifier.get_full_name()] = AuthTableQueryResponse(
+            column_masking={'n': _rule(part, field, type_)})
+        for native in [False, True]:
+            result = _read(table, native, ['id', 'n']).sort_by([('id', 'ascending')])
+            assert result.column('n').type == pa.int32()
+            assert result.column('n').to_pylist() == values, (part, field, native)
+
+
+@pytest.mark.parametrize('mode', ['wire', 'setup-failure', 'disable-native'])
+def test_calendar_fallback_keeps_hidden_dependency_and_captured_rule(calendar_table, mode):
+    table, server = calendar_table
+    server.table_query_auth_store[table.identifier.get_full_name()] = AuthTableQueryResponse(
+        column_masking={'n': _rule('YEAR_OF_WEEK', 'd', 'DATE')})
+    table = table.copy({'read.native.enabled': 'true', 'scan.native-plan.enabled': 'true'})
+    builder = (table.new_read_builder().with_projection(['id']).with_limit(1)
+               .with_filter(PredicateBuilder(table.fields).equal('n', 2025)))
+    scan = builder.new_scan()
+    with patch.object(scan.file_scanner, 'scan', side_effect=AssertionError('Native plan fell back')):
+        splits = scan.plan().splits()
+    server.table_query_auth_store[table.identifier.get_full_name()] = AuthTableQueryResponse()
+    if mode == 'wire':
+        splits = [deserialize_split_v1(serialize_split_v1(split), []) for split in splits]
+    if mode == 'disable-native':
+        builder = (table.copy({'read.native.enabled': 'false'}).new_read_builder().with_projection(['id'])
+                   .with_filter(PredicateBuilder(table.fields).equal('n', 2025)).with_limit(1))
+    if mode == 'setup-failure':
+        with patch('pypaimon.read.native_plan.native_read', side_effect=RuntimeError('setup failed')):
+            rows = builder.new_read().to_arrow(splits).to_pylist()
+    else:
+        rows = builder.new_read().to_arrow(splits).to_pylist()
+    assert rows == [{'id': 2}]
+
+
+@pytest.mark.parametrize('native', [False, True])
+def test_calendar_mask_cannot_expose_another_masked_operand(calendar_table, native):
+    table, server = calendar_table
+    server.table_query_auth_store[table.identifier.get_full_name()] = AuthTableQueryResponse(
+        column_masking={'n': _rule('YEAR', 'd', 'DATE'), 'd': json.dumps({'name': 'NULL'})})
+    table = table.copy({'scan.native-plan.enabled': str(native).lower()})
+    with pytest.raises(ValueError, match="masked too"):
+        table.new_read_builder().with_projection(['n']).new_scan().plan()
+
+
+@pytest.mark.parametrize('native', [False, True])
+def test_stale_calendar_operand_fails_closed(calendar_table, native):
+    table, server = calendar_table
+    server.table_query_auth_store[table.identifier.get_full_name()] = AuthTableQueryResponse(
+        column_masking={'n': _rule('YEAR', 'missing', 'DATE')})
+    table = table.copy({'scan.native-plan.enabled': str(native).lower()})
+    with pytest.raises(ValueError, match="missing"):
+        table.new_read_builder().with_projection(['n']).new_scan().plan()
+
+
+@pytest.mark.parametrize('field,type_', [('d', 'TIMESTAMP(6)'), ('ts', 'DATE')])
+def test_calendar_mask_rebinds_stored_temporal_family(calendar_table, field, type_):
+    table, server = calendar_table
+    # Java validates the stored temporal family, then copies the FieldRef from
+    # the current schema. The actual column decides the executed calendar value.
+    server.table_query_auth_store[table.identifier.get_full_name()] = AuthTableQueryResponse(
+        column_masking={'n': _rule('YEAR_OF_WEEK', field, type_)})
+    for native in [False, True]:
+        result = _read(table, native, ['n', 'id']).sort_by([('id', 'ascending')])
+        assert result.column('n').to_pylist() == [2024, 2025, 1970, None]

paimon-python/pypaimon/tests/predicate_json_parser_test.py

@@ -19,6 +19,7 @@
 import datetime
 import json
 import unittest
+from pathlib import Path
 
 import pyarrow as pa
 
@@ -861,6 +862,70 @@ def test_date_field_starts_at_midnight(self):
         self.assertEqual(self._apply("MONTH", batch, "d", "DATE"), [1, None])
         self.assertEqual(self._apply("HOUR", batch, "d", "DATE"), [0, None])
 
+    def test_calendar_parts_match_java_golden_for_every_storage_unit(self):
+        parts = ['YEAR', 'MONTH', 'DAY', 'HOUR', 'MINUTE', 'SECOND', 'QUARTER', 'DAY_OF_YEAR',
+                 'WEEKDAY', 'ISO_DAY_OF_WEEK', 'DAY_OF_WEEK', 'WEEK', 'YEAR_OF_WEEK']
+        golden = Path(__file__).parent / 'resources' / 'calendar_extract_java.jsonl'
+        for line in golden.read_text().splitlines():
+            case = json.loads(line)
+            unit = case['unit']
+            type_ = pa.date32() if unit == 'date' else pa.timestamp(unit)
+            wire = 'DATE' if unit == 'date' else 'TIMESTAMP(9)'
+            batch = pa.RecordBatch.from_pydict({'t': pa.array([case['value'], None], type=type_)})
+            for part, expected in zip(parts, case['parts']):
+                with self.subTest(unit=unit, value=case['value'], part=part):
+                    self.assertEqual(self._apply(part, batch, 't', wire), [expected, None])
+
+    def test_empty_and_all_null_calendar_batches(self):
+        for unit in ['date', 's', 'ms', 'us', 'ns']:
+            type_ = pa.date32() if unit == 'date' else pa.timestamp(unit)
+            wire = 'DATE' if unit == 'date' else 'TIMESTAMP(9)'
+            for values in [[], [None, None]]:
+                batch = pa.RecordBatch.from_pydict({'t': pa.array(values, type=type_)})
+                for part in ['YEAR', 'MONTH', 'DAY', 'HOUR', 'MINUTE', 'SECOND', 'QUARTER', 'DAY_OF_YEAR',
+                             'WEEKDAY', 'ISO_DAY_OF_WEEK', 'DAY_OF_WEEK', 'WEEK', 'YEAR_OF_WEEK']:
+                    with self.subTest(unit=unit, values=values, part=part):
+                        self.assertEqual(self._apply(part, batch, 't', wire), values)
+
+    def test_calendar_operand_is_a_dependency_even_if_hidden(self):
+        for part in ['YEAR', 'MONTH', 'DAY', 'HOUR', 'MINUTE', 'SECOND', 'QUARTER', 'DAY_OF_YEAR',
+                     'WEEKDAY', 'ISO_DAY_OF_WEEK', 'DAY_OF_WEEK', 'WEEK', 'YEAR_OF_WEEK']:
+            rule = {'name': part, 'fieldRef': {'name': 'hidden', 'index': 99, 'type': 'DATE'}}
+            predicate = {'kind': 'LEAF', 'function': 'IS_NULL', 'transform': rule}
+            self.assertEqual(extract_referenced_fields(json.dumps(predicate)), {'hidden'})
+
+    def test_stored_and_rebound_calendar_types_are_validated(self):
+        for stored in ['INT', 'BIGINT', 'TIME(3)', 'TIMESTAMP_LTZ(3)', 'TIMESTAMP(10)', 'DATENULL',
+                       'TIMESTAMP(٣)', 'TIMESTAMP(３)',
+                       'TIMESTAMP(3) WITHOUT TIME ZONENULL',
+                       'TIMESTAMP(3) WITHOUT TIME ZONENOT NULL',
+                       {'type': 'DATE'}, {'type': 'TIMESTAMP(6)'}, None]:
+            with self.subTest(stored=stored), self.assertRaises(ValueError):
+                self._apply('YEAR', self._ts_batch(), 't', stored)
+        for type_, value in [(pa.int32(), 1), (pa.time32('ms'), 1), (pa.timestamp('ms', tz='UTC'), 1)]:
+            batch = pa.RecordBatch.from_pydict({'t': pa.array([value], type=type_)})
+            with self.subTest(type=type_), self.assertRaises(ValueError):
+                self._apply('YEAR', batch, 't', 'DATE')
+        # As in Java copyWithNewInputs, DATE and local TIMESTAMP may rebind to
+        # one another; the actual column decides the executed calendar value.
+        for stored in ['DATE', 'DATE NULL', 'TIMESTAMP(3)NOT NULL', 'TIMESTAMP WITHOUT TIME ZONE',
+                       'TIMESTAMP(3) WITHOUT TIME ZONE',
+                       'TIMESTAMP(3)WITHOUT TIME ZONE']:
+            with self.subTest(stored=stored):
+                self.assertEqual(self._apply('YEAR', self._ts_batch(), 't', stored), [2024, 2024, None])
+
+    def test_full_date32_domain_does_not_require_a_timestamp_cast(self):
+        batch = pa.RecordBatch.from_pydict({"d": pa.array([2147483647, -2147483648, None], type=pa.date32())})
+        self.assertEqual(self._apply("HOUR", batch, "d", "DATE"), [0, 0, None])
+
+    def test_timestamp_epoch_day_narrows_as_java_int_before_calendar_extraction(self):
+        # Java Timestamp narrows its millisecond epoch day to a signed int.
+        ticks = (2147483648 * 86400000)
+        batch = pa.RecordBatch.from_pydict({"t": pa.array([ticks, -1, None], type=pa.timestamp("ms"))})
+        days = pa.RecordBatch.from_pydict({"d": pa.array([-2147483648, -1, None], type=pa.date32())})
+        self.assertEqual(self._apply("YEAR", batch, "t", "TIMESTAMP(3)"),
+                         self._apply("YEAR", days, "d", "DATE"))
+
     def test_wired_through_a_leaf_filter(self):
         pred = json.dumps({
             "kind": "LEAF",