This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/paimon.git
The following commit(s) were added to refs/heads/master by this push:
new acd7471eb7 [python] Speed up selected-key projection for ordinary MAP
columns (#10147)
acd7471eb7 is described below
commit acd7471eb796473d1a7543ec380edf41d532291c
Author: zhigang <[email protected]>
AuthorDate: Thu Sep 24 10:49:02 2026 +0800
[python] Speed up selected-key projection for ordinary MAP columns (#10147)
---
.../pypaimon/data/map_shared_shredding.py | 22 ++++
.../tests/normal_map_selected_keys_test.py | 111 +++++++++++++++++++++
2 files changed, 133 insertions(+)
diff --git a/paimon-python/pypaimon/data/map_shared_shredding.py
b/paimon-python/pypaimon/data/map_shared_shredding.py
index 1cfa783161..1fe0a61627 100644
--- a/paimon-python/pypaimon/data/map_shared_shredding.py
+++ b/paimon-python/pypaimon/data/map_shared_shredding.py
@@ -495,6 +495,28 @@ def assemble_normal_map_selected_keys(
"""Materialize selected values when an older file stores a normal MAP."""
if not pa.types.is_map(column.type):
raise TypeError("Selected-key MAP must be stored as a map or shared
struct")
+ # map_lookup preserves first-match semantics, including a null first value.
+ # Older supported Arrow releases do not provide this kernel.
+ map_lookup = getattr(pc, "map_lookup", None)
+ if map_lookup is not None and (
+ pa.types.is_string(column.type.key_type) or
pa.types.is_large_string(column.type.key_type)
+ ):
+ try:
+ children = [
+ _restore_orc_temporal_values(
+ map_lookup(column, pa.scalar(key,
type=column.type.key_type), "first"),
+ value_type,
+ )
+ for key in selected_keys
+ ]
+ except pa.ArrowNotImplementedError:
+ # Some value types support take but not the map_lookup builder.
+ pass
+ else:
+ fields = [pa.field(key, value_type) for key in selected_keys]
+ mask = column.is_null() if column.null_count else None
+ return pa.StructArray.from_arrays(children, fields=fields,
mask=mask)
+
offsets, start, end = _normalized_offsets(column)
keys = column.keys.slice(start, end - start).to_pylist()
values = _restore_orc_temporal_values(
diff --git a/paimon-python/pypaimon/tests/normal_map_selected_keys_test.py
b/paimon-python/pypaimon/tests/normal_map_selected_keys_test.py
new file mode 100644
index 0000000000..3d69010543
--- /dev/null
+++ b/paimon-python/pypaimon/tests/normal_map_selected_keys_test.py
@@ -0,0 +1,111 @@
+# 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.
+
+from datetime import date, datetime, time, timezone
+from decimal import Decimal
+import unittest
+
+import pyarrow as pa
+
+from pypaimon.data.map_shared_shredding import
assemble_normal_map_selected_keys
+
+
+class NormalMapSelectedKeysTest(unittest.TestCase):
+ def test_first_match_nulls_and_slices(self):
+ cases = [
+ (pa.bool_(), True),
+ (pa.int8(), -128),
+ (pa.int16(), 32767),
+ (pa.int32(), -2147483648),
+ (pa.int64(), 9007199254740993),
+ (pa.float32(), float("inf")),
+ (pa.float64(), float("-inf")),
+ (pa.decimal128(28, 6), Decimal("12345678901234567890.123456")),
+ (pa.date32(), date(1960, 1, 1)),
+ (pa.time32("ms"), time(12, 34, 56, 123000)),
+ (pa.timestamp("us"), datetime(1960, 1, 1, 0, 0, 0, 123456)),
+ (pa.timestamp("us", "UTC"), datetime(1960, 1, 1,
tzinfo=timezone.utc)),
+ (pa.binary(), b"\x00\xff"),
+ (pa.string(), "中文"),
+ (pa.list_(pa.int64()), [None, 9007199254740993]),
+ (pa.struct([("x", pa.int64())]), {"x": None}),
+ (pa.map_(pa.string(), pa.int64()), [("x", 1), ("x", None)]),
+ ]
+ for key_type in [pa.string(), pa.large_string()]:
+ for value_type, value in cases:
+ with self.subTest(key_type=key_type, value_type=value_type):
+ rows = [
+ [("a", value)],
+ [("a", None), ("a", value)],
+ [("b", value), ("a", value), ("a", None)],
+ [],
+ None,
+ ]
+ column = pa.array(rows, type=pa.map_(key_type, value_type))
+ expected = [
+ {"a": value, "missing": None},
+ {"a": None, "missing": None},
+ {"a": value, "missing": None},
+ {"a": None, "missing": None},
+ None,
+ ]
+ for offset, size in [(0, 5), (1, 4), (2, 0), (2, 1)]:
+ actual = assemble_normal_map_selected_keys(
+ column.slice(offset, size), ["a", "missing"],
value_type
+ )
+ end = offset + size
+ self.assertEqual(expected[offset:end],
actual.to_pylist())
+ self.assertEqual(value_type, actual.type[0].type)
+
+ def test_null_map_hides_physical_entries(self):
+ column = pa.MapArray.from_arrays(
+ pa.array([0, None, 2, 2], type=pa.int32()), pa.array(["a", "a"]),
pa.array([10, 20])
+ )
+ result = assemble_normal_map_selected_keys(column, ["a"], pa.int64())
+ self.assertEqual([{"a": 10}, None, {"a": None}], result.to_pylist())
+ self.assertEqual([10, None, None], result.field(0).to_pylist())
+
+ def test_restores_orc_time_after_selection(self):
+ column = pa.array([[("a", 1234)], [], None], type=pa.map_(pa.string(),
pa.int32()))
+ actual = assemble_normal_map_selected_keys(column, ["a"],
pa.time32("ms"))
+ self.assertEqual([{"a": time(0, 0, 1, 234000)}, {"a": None}, None],
actual.to_pylist())
+
+ def test_non_string_keys_keep_comparison_semantics(self):
+ for key_type, key in [(pa.binary(), b"a"), (pa.int32(), 1)]:
+ column = pa.array([[(key, 10)]], type=pa.map_(key_type,
pa.int64()))
+ result = assemble_normal_map_selected_keys(column, ["a"],
pa.int64())
+ self.assertEqual([{"a": None}], result.to_pylist())
+
+ def test_extension_values_use_take_fallback(self):
+ class IntExtension(pa.ExtensionType):
+ def __init__(self):
+ super().__init__(pa.int64(), "paimon.test.map-int")
+
+ def __arrow_ext_serialize__(self):
+ return b""
+
+ @classmethod
+ def __arrow_ext_deserialize__(cls, storage_type, serialized):
+ return cls()
+
+ value_type = IntExtension()
+ values = pa.ExtensionArray.from_storage(value_type, pa.array([10,
None, 20]))
+ column = pa.MapArray.from_arrays(
+ pa.array([0, 1, 3], type=pa.int32()), pa.array(["a", "a", "a"]),
values
+ )
+ result = assemble_normal_map_selected_keys(column, ["a"], value_type)
+ self.assertEqual(value_type, result.field(0).type)
+ self.assertEqual([10, None], result.field(0).storage.to_pylist())