Fokko commented on code in PR #7831:
URL: https://github.com/apache/iceberg/pull/7831#discussion_r1264727823
##########
python/pyiceberg/io/pyarrow.py:
##########
@@ -1013,3 +1025,271 @@ def map_key_partner(self, partner_map:
Optional[pa.Array]) -> Optional[pa.Array]
def map_value_partner(self, partner_map: Optional[pa.Array]) ->
Optional[pa.Array]:
return partner_map.items if isinstance(partner_map, pa.MapArray) else
None
+
+
+class StatsAggregator:
+ def __init__(self, type_string: str, trunc_length: Optional[int] = None)
-> None:
+ self.current_min: Any = None
+ self.current_max: Any = None
+ self.trunc_length = trunc_length
+ self.primitive_type: Optional[PrimitiveType] = None
+
+ if type_string == "BOOLEAN":
+ self.primitive_type = BooleanType()
+ elif type_string == "INT32":
+ self.primitive_type = IntegerType()
+ elif type_string == "INT64":
+ self.primitive_type = LongType()
+ elif type_string == "INT96":
+ raise NotImplementedError("Statistics not implemented for INT96
physical type")
+ elif type_string == "FLOAT":
+ self.primitive_type = FloatType()
+ elif type_string == "DOUBLE":
+ self.primitive_type = DoubleType()
+ elif type_string == "BYTE_ARRAY":
+ self.primitive_type = BinaryType()
+ elif type_string == "FIXED_LEN_BYTE_ARRAY":
+ self.primitive_type = BinaryType()
+ else:
+ raise AssertionError(f"Unknown physical type {type_string}")
+
+ def serialize(self, value: Any) -> bytes:
+ if type(value) == str:
+ value = value.encode()
+ assert self.primitive_type is not None # appease mypy
+ return to_bytes(self.primitive_type, value)
+
+ def add_min(self, val: Any) -> None:
+ if self.current_min is None:
+ self.current_min = val
+ else:
+ self.current_min = min(val, self.current_min)
+
+ def add_max(self, val: Any) -> None:
+ if self.current_max is None:
+ self.current_max = val
+ else:
+ self.current_max = max(self.current_max, val)
+
+ def get_min(self) -> bytes:
+ return self.serialize(self.current_min)[: self.trunc_length]
+
+ def get_max(self) -> bytes:
+ return self.serialize(self.current_max)[: self.trunc_length]
+
+
+DEFAULT_TRUNCATION_LENGHT = 16
+TRUNCATION_EXPR = r"^truncate\((\d+)\)$"
+
+
+class MetricModeTypes(Enum):
+ TRUNCATE = "truncate"
+ NONE = "none"
+ COUNTS = "counts"
+ FULL = "full"
+
+
+DEFAULT_METRICS_MODE_KEY = "write.metadata.metrics.default"
+COLUMN_METRICS_MODE_KEY = "write.metadata.metrics.column"
+
+
+@dataclass(frozen=True)
+class MetricsMode(Singleton):
+ type: MetricModeTypes
+ length: Optional[int] = None
+
+
+def match_metrics_mode(mode: str) -> MetricsMode:
+ m = re.match(TRUNCATION_EXPR, mode, re.IGNORECASE)
+ if m:
+ length = int(m[1])
+ if length < 1:
+ raise AssertionError("Truncation length must be larger than 0")
+ return MetricsMode(MetricModeTypes.TRUNCATE, int(m[1]))
+ elif re.match("^none$", mode, re.IGNORECASE):
+ return MetricsMode(MetricModeTypes.NONE)
+ elif re.match("^counts$", mode, re.IGNORECASE):
+ return MetricsMode(MetricModeTypes.COUNTS)
+ elif re.match("^full$", mode, re.IGNORECASE):
+ return MetricsMode(MetricModeTypes.FULL)
+ else:
+ raise AssertionError(f"Unsupported metrics mode {mode}")
+
+
+@dataclass(frozen=True)
+class StatisticsCollector:
+ field_id: int
+ iceberg_type: PrimitiveType
+ mode: MetricsMode
+ column_name: str
+
+
+class
PyArrowStatisticsCollector(PreOrderSchemaVisitor[List[StatisticsCollector]]):
+ _field_id = 0
+ _schema: Schema
+ _properties: Dict[str, str]
+
+ def __init__(self, schema: Schema, properties: Dict[str, str]):
+ self._schema = schema
+ self._properties = properties
+
+ def schema(self, schema: Schema, struct_result: Callable[[],
List[StatisticsCollector]]) -> List[StatisticsCollector]:
+ return struct_result()
+
+ def struct(
+ self, struct: StructType, field_results: List[Callable[[],
List[StatisticsCollector]]]
+ ) -> List[StatisticsCollector]:
+ return list(chain(*[result() for result in field_results]))
+
+ def field(self, field: NestedField, field_result: Callable[[],
List[StatisticsCollector]]) -> List[StatisticsCollector]:
+ self._field_id = field.field_id
+ result = field_result()
+ return result
+
+ def list(self, list_type: ListType, element_result: Callable[[],
List[StatisticsCollector]]) -> List[StatisticsCollector]:
+ self._field_id = list_type.element_id
+ return element_result()
+
+ def map(
+ self,
+ map_type: MapType,
+ key_result: Callable[[], List[StatisticsCollector]],
+ value_result: Callable[[], List[StatisticsCollector]],
+ ) -> List[StatisticsCollector]:
+ self._field_id = map_type.key_id
+ k = key_result()
+ self._field_id = map_type.value_id
+ v = value_result()
+ return k + v
+
+ def primitive(self, primitive: PrimitiveType) -> List[StatisticsCollector]:
+ column_name = self._schema.find_column_name(self._field_id)
+ assert column_name is not None, f"Column for field {self._field_id}
not found"
+
+ metrics_mode = MetricsMode(MetricModeTypes.TRUNCATE,
DEFAULT_TRUNCATION_LENGHT)
+
+ default_mode = self._properties.get(DEFAULT_METRICS_MODE_KEY)
+ if default_mode:
+ metrics_mode = match_metrics_mode(default_mode)
+
+ col_mode =
self._properties.get(f"{COLUMN_METRICS_MODE_KEY}.{column_name}")
+ if col_mode:
+ metrics_mode = match_metrics_mode(col_mode)
+
+ return [StatisticsCollector(field_id=self._field_id,
iceberg_type=primitive, mode=metrics_mode, column_name=column_name)]
+
+
+def fill_parquet_file_metadata(
+ df: DataFile,
+ parquet_metadata: pq.FileMetaData,
+ file_size: int,
+ table_metadata: TableMetadata,
+) -> None:
+ """
+ Computes and fills the following fields of the DataFile object.
+
+ - file_format
+ - record_count
+ - file_size_in_bytes
+ - column_sizes
+ - value_counts
+ - null_value_counts
+ - nan_value_counts
+ - lower_bounds
+ - upper_bounds
+ - split_offsets
+
+ Args:
+ df (DataFile): A DataFile object representing the Parquet file for
which metadata is to be filled.
+ parquet_metadata (pyarrow.parquet.FileMetaData): A pyarrow metadata
object.
+ file_size (int): The total compressed file size cannot be retrieved
from the metadata and hence has to
+ be passed here. Depending on the kind of file system and pyarrow
library call used, different
+ ways to obtain this value might be appropriate.
+ table_metadata (pyiceberg.table.metadata.TableMetadata): The Iceberg
table metadata. It is required to
+ compute the mapping if column position to iceberg schema type id.
It's also used to set the mode
+ for column metrics collection
+ """
+ schema = next(filter(lambda s: s.schema_id ==
table_metadata.current_schema_id, table_metadata.schemas))
+
+ stats_columns = pre_order_visit(schema, PyArrowStatisticsCollector(schema,
table_metadata.properties))
+ assert parquet_metadata.num_columns == len(
+ stats_columns
+ ), f"Number of columns in metadata ({len(stats_columns)}) is different
from the number of columns in pyarrow table ({parquet_metadata.num_columns})"
+
+ column_sizes: Dict[int, int] = {}
+ value_counts: Dict[int, int] = {}
+ split_offsets: List[int] = []
+
+ null_value_counts: Dict[int, int] = {}
+ nan_value_counts: Dict[int, int] = {}
+
+ col_aggs = {}
+
+ for r in range(parquet_metadata.num_row_groups):
+ # References:
+ #
https://github.com/apache/iceberg/blob/fc381a81a1fdb8f51a0637ca27cd30673bd7aad3/parquet/src/main/java/org/apache/iceberg/parquet/ParquetUtil.java#L232
+ #
https://github.com/apache/parquet-mr/blob/ac29db4611f86a07cc6877b416aa4b183e09b353/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/metadata/ColumnChunkMetaData.java#L184
+
+ row_group = parquet_metadata.row_group(r)
+
+ data_offset = row_group.column(0).data_page_offset
+ dictionary_offset = row_group.column(0).dictionary_page_offset
+
+ if row_group.column(0).has_dictionary_page and dictionary_offset <
data_offset:
+ split_offsets.append(dictionary_offset)
+ else:
+ split_offsets.append(data_offset)
+
+ for c in range(parquet_metadata.num_columns):
+ col_id = stats_columns[c].field_id
+
+ column = row_group.column(c)
+
+ column_sizes[col_id] = column_sizes.get(col_id, 0) +
column.total_compressed_size
+
+ metrics_mode = stats_columns[c].mode
+
+ if metrics_mode == MetricsMode(MetricModeTypes.NONE):
+ continue
+
+ value_counts[col_id] = value_counts.get(col_id, 0) +
column.num_values
+
+ if column.is_stats_set:
+ try:
+ statistics = column.statistics
+
+ null_value_counts[col_id] = null_value_counts.get(col_id,
0) + statistics.null_count
+
+ if metrics_mode == MetricsMode(MetricModeTypes.COUNTS):
+ continue
+
+ if col_id not in col_aggs:
+ col_aggs[col_id] =
StatsAggregator(statistics.physical_type, metrics_mode.length)
+
+ col_aggs[col_id].add_min(statistics.min)
+ col_aggs[col_id].add_max(statistics.max)
Review Comment:
What do you think of the following? This way we have fewer lookups. Also, I
think it is best to stick with the Iceberg naming (field-id over column-id).
```suggestion
for pos, stats_col in enumerate(stats_columns):
field_id = stats_col.field_id
column = row_group.column(pos)
column_sizes[field_id] = column_sizes.get(field_id, 0) +
column.total_compressed_size
if stats_col.mode == MetricsMode(MetricModeTypes.NONE):
continue
value_counts[field_id] = value_counts.get(field_id, 0) +
column.num_values
if column.is_stats_set:
try:
statistics = column.statistics
null_value_counts[field_id] =
null_value_counts.get(field_id, 0) + statistics.null_count
if stats_col.mode == MetricsMode(MetricModeTypes.COUNTS):
continue
if field_id not in col_aggs:
col_aggs[field_id] =
StatsAggregator(statistics.physical_type, stats_col.mode.length)
col_aggs[field_id].add_min(statistics.min)
col_aggs[field_id].add_max(statistics.max)
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
