paimon-python/pypaimon/read/scanner/data_evolution_stats.py - paimon - Git at Google

 # 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 typing import Callable, Dict, List, Tuple

 from pypaimon.common.predicate import Predicate
 from pypaimon.manifest.schema.data_file_meta import DataFileMeta
 from pypaimon.manifest.schema.simple_stats import SimpleStats
 from pypaimon.read.push_down_utils import rewrite_predicate_indices
 from pypaimon.schema.data_types import DataField
 from pypaimon.table.special_fields import SpecialFields
 from pypaimon.table.row.generic_row import GenericRow


 _KNOWN = 0
 _MISSING = 1
 _UNKNOWN = 2


 class _FileLayout:

     def __init__(self, data_fields, stats_offsets):
         self.data_fields = {field.id: field for field in data_fields}
         self.stats_offsets = stats_offsets


 class _StatsProvider:

     def __init__(self, file, layout):
         self.file = file
         self.layout = layout
         self.tied = False


 class DataEvolutionGroupStatsFilter:
     """Conservatively filters logical row-id groups by merged column stats."""

     def __init__(
         self,
         predicate: Predicate,
         table_fields: List[DataField],
         schema_fields: Callable[[int], List[DataField]],
     ):
         self.predicate = rewrite_predicate_indices(predicate, table_fields)
         self.table_fields = table_fields
         self.schema_fields = schema_fields
         self._layout_cache: Dict[Tuple, _FileLayout] = {}

     def may_match(self, files: List[DataFileMeta]) -> bool:
         if not files:
             return True
         try:
             stats, states, row_count = self._group_stats(files)
             return self._predicate_may_match(
                 self.predicate, stats, states, row_count)
         except Exception:
             # Stats pruning is optional. Unknown schemas, corrupt stats, and
             # incompatible types must retain the complete logical group.
             return True

     def _group_stats(self, files):
         group_start = min(file.non_null_row_id_range().from_ for file in files)
         group_end = max(file.non_null_row_id_range().to for file in files)
         row_count = group_end - group_start + 1

         normal_files = []
         special_field_ids = set()
         for file in files:
             layout = self._layout(file)
             if DataFileMeta.is_blob_file(file.file_name) \
                     or DataFileMeta.is_vector_file(file.file_name):
                 special_field_ids.update(layout.data_fields)
             else:
                 normal_files.append((file, layout))
         providers = {}
         for file, layout in normal_files:
             for field_id in layout.data_fields:
                 current = providers.get(field_id)
                 if (current is None
                         or file.max_sequence_number
                         > current.file.max_sequence_number):
                     providers[field_id] = _StatsProvider(file, layout)
                 elif (file.max_sequence_number
                       == current.file.max_sequence_number):
                     current.tied = True

         min_values = []
         max_values = []
         null_counts = []
         states = []
         for field_index, field in enumerate(self.table_fields):
             provider = providers.get(field.id)
             if provider is None:
                 if field.id in special_field_ids:
                     self._append_unknown(
                         min_values, max_values, null_counts, states)
                 else:
                     min_values.append(None)
                     max_values.append(None)
                     null_counts.append(row_count)
                     states.append(_MISSING)
                 continue

             if provider.tied:
                 self._append_unknown(
                     min_values, max_values, null_counts, states)
                 continue

             file = provider.file
             layout = provider.layout
             file_range = file.non_null_row_id_range()
             source_field = layout.data_fields[field.id]
             # Partial-file stats do not describe the complete logical group.
             if (file_range.from_ != group_start
                     or file_range.to != group_end
                     or source_field.type != field.type
                     or field.id not in layout.stats_offsets):
                 self._append_unknown(
                     min_values, max_values, null_counts, states)
                 continue

             stats = file.value_stats
             stats_offset = layout.stats_offsets[field.id]
             min_value = stats.min_values.get_field(stats_offset)
             max_value = stats.max_values.get_field(stats_offset)
             null_count = (
                 stats.null_counts[stats_offset]
                 if (stats.null_counts is not None
                     and stats_offset < len(stats.null_counts))
                 else None
             )
             self._validate_stats(
                 min_value, max_value, null_count, row_count)
             min_values.append(min_value)
             max_values.append(max_value)
             null_counts.append(null_count)
             states.append(_KNOWN)

         return (
             SimpleStats(
                 GenericRow(min_values, self.table_fields),
                 GenericRow(max_values, self.table_fields),
                 null_counts,
             ),
             states,
             row_count,
         )

     @staticmethod
     def _append_unknown(min_values, max_values, null_counts, states):
         min_values.append(None)
         max_values.append(None)
         null_counts.append(None)
         states.append(_UNKNOWN)

     def _layout(self, file):
         key = (
             file.schema_id,
             tuple(file.write_cols) if file.write_cols is not None else None,
             (tuple(file.value_stats_cols)
              if file.value_stats_cols is not None else None),
         )
         layout = self._layout_cache.get(key)
         if layout is not None:
             return layout

         schema = self.schema_fields(file.schema_id)
         schema_fields = schema.fields if hasattr(schema, 'fields') else schema
         fields_by_name = {field.name: field for field in schema_fields}
         data_fields = (
             schema.data_file_fields(None)
             if file.write_cols is None
             and hasattr(schema, 'data_file_fields')
             else self._project_fields(
                 schema_fields, fields_by_name, file.write_cols)
         )
         stats_fields = self._project_fields(
             data_fields,
             {field.name: field for field in data_fields},
             file.value_stats_cols,
         )
         layout = _FileLayout(
             data_fields,
             {field.id: index for index, field in enumerate(stats_fields)},
         )
         self._layout_cache[key] = layout
         return layout

     @staticmethod
     def _project_fields(default_fields, fields_by_name, names):
         if names is None:
             return default_fields
         if len(names) != len(set(names)):
             raise ValueError("Duplicate fields in file stats metadata.")
         unknown = [
             name for name in names
             if name not in fields_by_name
             and not SpecialFields.is_system_field(name)
         ]
         if unknown:
             raise ValueError("Unknown fields in file stats metadata: %s" % unknown)
         return [fields_by_name[name] for name in names if name in fields_by_name]

     @staticmethod
     def _validate_stats(min_value, max_value, null_count, row_count):
         if (null_count is not None
                 and (isinstance(null_count, bool)
                      or not isinstance(null_count, int)
                      or null_count < 0
                      or null_count > row_count)):
             raise ValueError("Invalid null count in file stats.")
         if (min_value is None) != (max_value is None):
             raise ValueError("Incomplete min/max values in file stats.")
         if min_value is not None:
             try:
                 ordered = min_value <= max_value
             except TypeError as exc:
                 raise ValueError("Incomparable min/max values in file stats.") from exc
             if not ordered:
                 raise ValueError("Invalid min/max order in file stats.")
             if null_count == row_count:
                 raise ValueError("All-null stats contain non-null bounds.")

     def _predicate_may_match(self, predicate, stats, states, row_count):
         if predicate.method == 'and':
             return all(
                 self._predicate_may_match(child, stats, states, row_count)
                 for child in predicate.literals
             )
         if predicate.method == 'or':
             return any(
                 self._predicate_may_match(child, stats, states, row_count)
                 for child in predicate.literals
             )

         index = predicate.index
         if index is None or index < 0 or index >= len(states):
             return True
         if states[index] == _UNKNOWN:
             return True
         if states[index] == _MISSING:
             tester = Predicate.testers.get(predicate.method)
             return True if tester is None else tester.test_by_value(
                 None, predicate.literals)
         field_type = getattr(self.table_fields[index].type, 'type', None)
         if (field_type in ('FLOAT', 'DOUBLE')
                 and predicate.method in ('notEqual', 'notIn')):
             # PyArrow min/max can omit NaN, which still matches negative
             # predicates.
             return True
         return predicate.test_by_simple_stats(stats, row_count)
	# 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 typing import Callable, Dict, List, Tuple

	from pypaimon.common.predicate import Predicate
	from pypaimon.manifest.schema.data_file_meta import DataFileMeta
	from pypaimon.manifest.schema.simple_stats import SimpleStats
	from pypaimon.read.push_down_utils import rewrite_predicate_indices
	from pypaimon.schema.data_types import DataField
	from pypaimon.table.special_fields import SpecialFields
	from pypaimon.table.row.generic_row import GenericRow


	_KNOWN = 0
	_MISSING = 1
	_UNKNOWN = 2


	class _FileLayout:

	def __init__(self, data_fields, stats_offsets):
	self.data_fields = {field.id: field for field in data_fields}
	self.stats_offsets = stats_offsets


	class _StatsProvider:

	def __init__(self, file, layout):
	self.file = file
	self.layout = layout
	self.tied = False


	class DataEvolutionGroupStatsFilter:
	"""Conservatively filters logical row-id groups by merged column stats."""

	def __init__(
	self,
	predicate: Predicate,
	table_fields: List[DataField],
	schema_fields: Callable[[int], List[DataField]],
	):
	self.predicate = rewrite_predicate_indices(predicate, table_fields)
	self.table_fields = table_fields
	self.schema_fields = schema_fields
	self._layout_cache: Dict[Tuple, _FileLayout] = {}

	def may_match(self, files: List[DataFileMeta]) -> bool:
	if not files:
	return True
	try:
	stats, states, row_count = self._group_stats(files)
	return self._predicate_may_match(
	self.predicate, stats, states, row_count)
	except Exception:
	# Stats pruning is optional. Unknown schemas, corrupt stats, and
	# incompatible types must retain the complete logical group.
	return True

	def _group_stats(self, files):
	group_start = min(file.non_null_row_id_range().from_ for file in files)
	group_end = max(file.non_null_row_id_range().to for file in files)
	row_count = group_end - group_start + 1

	normal_files = []
	special_field_ids = set()
	for file in files:
	layout = self._layout(file)
	if DataFileMeta.is_blob_file(file.file_name) \
	or DataFileMeta.is_vector_file(file.file_name):
	special_field_ids.update(layout.data_fields)
	else:
	normal_files.append((file, layout))
	providers = {}
	for file, layout in normal_files:
	for field_id in layout.data_fields:
	current = providers.get(field_id)
	if (current is None
	or file.max_sequence_number
	> current.file.max_sequence_number):
	providers[field_id] = _StatsProvider(file, layout)
	elif (file.max_sequence_number
	== current.file.max_sequence_number):
	current.tied = True

	min_values = []
	max_values = []
	null_counts = []
	states = []
	for field_index, field in enumerate(self.table_fields):
	provider = providers.get(field.id)
	if provider is None:
	if field.id in special_field_ids:
	self._append_unknown(
	min_values, max_values, null_counts, states)
	else:
	min_values.append(None)
	max_values.append(None)
	null_counts.append(row_count)
	states.append(_MISSING)
	continue

	if provider.tied:
	self._append_unknown(
	min_values, max_values, null_counts, states)
	continue

	file = provider.file
	layout = provider.layout
	file_range = file.non_null_row_id_range()
	source_field = layout.data_fields[field.id]
	# Partial-file stats do not describe the complete logical group.
	if (file_range.from_ != group_start
	or file_range.to != group_end
	or source_field.type != field.type
	or field.id not in layout.stats_offsets):
	self._append_unknown(
	min_values, max_values, null_counts, states)
	continue

	stats = file.value_stats
	stats_offset = layout.stats_offsets[field.id]
	min_value = stats.min_values.get_field(stats_offset)
	max_value = stats.max_values.get_field(stats_offset)
	null_count = (
	stats.null_counts[stats_offset]
	if (stats.null_counts is not None
	and stats_offset < len(stats.null_counts))
	else None
	)
	self._validate_stats(
	min_value, max_value, null_count, row_count)
	min_values.append(min_value)
	max_values.append(max_value)
	null_counts.append(null_count)
	states.append(_KNOWN)

	return (
	SimpleStats(
	GenericRow(min_values, self.table_fields),
	GenericRow(max_values, self.table_fields),
	null_counts,
	),
	states,
	row_count,
	)

	@staticmethod
	def _append_unknown(min_values, max_values, null_counts, states):
	min_values.append(None)
	max_values.append(None)
	null_counts.append(None)
	states.append(_UNKNOWN)

	def _layout(self, file):
	key = (
	file.schema_id,
	tuple(file.write_cols) if file.write_cols is not None else None,
	(tuple(file.value_stats_cols)
	if file.value_stats_cols is not None else None),
	)
	layout = self._layout_cache.get(key)
	if layout is not None:
	return layout

	schema = self.schema_fields(file.schema_id)
	schema_fields = schema.fields if hasattr(schema, 'fields') else schema
	fields_by_name = {field.name: field for field in schema_fields}
	data_fields = (
	schema.data_file_fields(None)
	if file.write_cols is None
	and hasattr(schema, 'data_file_fields')
	else self._project_fields(
	schema_fields, fields_by_name, file.write_cols)
	)
	stats_fields = self._project_fields(
	data_fields,
	{field.name: field for field in data_fields},
	file.value_stats_cols,
	)
	layout = _FileLayout(
	data_fields,
	{field.id: index for index, field in enumerate(stats_fields)},
	)
	self._layout_cache[key] = layout
	return layout

	@staticmethod
	def _project_fields(default_fields, fields_by_name, names):
	if names is None:
	return default_fields
	if len(names) != len(set(names)):
	raise ValueError("Duplicate fields in file stats metadata.")
	unknown = [
	name for name in names
	if name not in fields_by_name
	and not SpecialFields.is_system_field(name)
	]
	if unknown:
	raise ValueError("Unknown fields in file stats metadata: %s" % unknown)
	return [fields_by_name[name] for name in names if name in fields_by_name]

	@staticmethod
	def _validate_stats(min_value, max_value, null_count, row_count):
	if (null_count is not None
	and (isinstance(null_count, bool)
	or not isinstance(null_count, int)
	or null_count < 0
	or null_count > row_count)):
	raise ValueError("Invalid null count in file stats.")
	if (min_value is None) != (max_value is None):
	raise ValueError("Incomplete min/max values in file stats.")
	if min_value is not None:
	try:
	ordered = min_value <= max_value
	except TypeError as exc:
	raise ValueError("Incomparable min/max values in file stats.") from exc
	if not ordered:
	raise ValueError("Invalid min/max order in file stats.")
	if null_count == row_count:
	raise ValueError("All-null stats contain non-null bounds.")

	def _predicate_may_match(self, predicate, stats, states, row_count):
	if predicate.method == 'and':
	return all(
	self._predicate_may_match(child, stats, states, row_count)
	for child in predicate.literals
	)
	if predicate.method == 'or':
	return any(
	self._predicate_may_match(child, stats, states, row_count)
	for child in predicate.literals
	)

	index = predicate.index
	if index is None or index < 0 or index >= len(states):
	return True
	if states[index] == _UNKNOWN:
	return True
	if states[index] == _MISSING:
	tester = Predicate.testers.get(predicate.method)
	return True if tester is None else tester.test_by_value(
	None, predicate.literals)
	field_type = getattr(self.table_fields[index].type, 'type', None)
	if (field_type in ('FLOAT', 'DOUBLE')
	and predicate.method in ('notEqual', 'notIn')):
	# PyArrow min/max can omit NaN, which still matches negative
	# predicates.
	return True
	return predicate.test_by_simple_stats(stats, row_count)