Skip to content

Commit c4db42c

Browse files
authored
[CHORE]: Don't use record fields (#2306)
I would like to clean up the `record_fields()` method: apache/iceberg-python#580 You should not check the name of the `_partition` since it can be anything (technically). In Iceberg we look up the partition by position based on the partition spec.
1 parent 06b00a6 commit c4db42c

1 file changed

Lines changed: 8 additions & 6 deletions

File tree

daft/iceberg/iceberg_scan.py

Lines changed: 8 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -101,16 +101,18 @@ def display_name(self) -> str:
101101
def partitioning_keys(self) -> list[PartitionField]:
102102
return self._partition_keys
103103

104-
def _iceberg_record_to_partition_spec(self, record: Record) -> daft.table.Table | None:
104+
def _iceberg_record_to_partition_spec(self, spec: IcebergPartitionSpec, record: Record) -> daft.table.Table | None:
105+
partition_fields = iceberg_partition_spec_to_fields(self._table.schema(), spec)
105106
arrays = dict()
106-
assert len(record._position_to_field_name) == len(self._partition_keys)
107-
for name, value, pfield in zip(record._position_to_field_name, record.record_fields(), self._partition_keys):
107+
assert len(record._position_to_field_name) == len(partition_fields)
108+
for idx, pfield in enumerate(partition_fields):
108109
field = Field._from_pyfield(pfield.field)
109110
field_name = field.name
110111
field_dtype = field.dtype
111112
arrow_type = field_dtype.to_arrow_dtype()
112-
assert name == field_name
113-
arrays[name] = daft.Series.from_arrow(pa.array([value], type=arrow_type), name=name).cast(field_dtype)
113+
arrays[field_name] = daft.Series.from_arrow(pa.array([record[idx]], type=arrow_type), name=field_name).cast(
114+
field_dtype
115+
)
114116
if len(arrays) > 0:
115117
return daft.table.Table.from_pydict(arrays)
116118
else:
@@ -162,7 +164,7 @@ def to_scan_tasks(self, pushdowns: Pushdowns) -> Iterator[ScanTask]:
162164
raise NotImplementedError("Iceberg Merge-on-Read currently not supported, please make an issue!")
163165

164166
# TODO: Thread in Statistics to each ScanTask: P2
165-
pspec = self._iceberg_record_to_partition_spec(file.partition)
167+
pspec = self._iceberg_record_to_partition_spec(self._table.specs()[file.spec_id], file.partition)
166168
st = ScanTask.catalog_scan_task(
167169
file=path,
168170
file_format=file_format_config,

0 commit comments

Comments
 (0)