def assign_fresh_ids(type_var, next_id): from ..schema import Schema if isinstance(type_var, Type): return visit(type_var, AssignFreshIds(next_id)) elif isinstance(type_var, Schema): schema = type_var return Schema(list(visit(schema.as_struct(), AssignFreshIds(next_id)) .as_nested_type().fields))
def select(schema, field_ids): import iceberg.api.schema if schema is None: raise RuntimeError("Schema cannot be None") if field_ids is None: raise RuntimeError("Field ids cannot be None") result = visit(schema, PruneColumns(field_ids)) if schema.as_struct() == result: return schema elif result is not None: if schema.get_aliases() is not None: return iceberg.api.schema.Schema(result.as_nested_type().fields, schema.get_aliases()) else: return iceberg.api.schema.Schema(result.as_nested_type().fields) return iceberg.api.schema.Schema(list(), schema.get_aliases())
def visit_custom_order(arg, visitor): from ..schema import Schema if isinstance(arg, Schema): schema = arg return visitor.schema(arg, VisitFuture(schema.as_struct(), visitor)) elif isinstance(arg, Type): type_var = arg if type_var.type_id == TypeID.STRUCT: struct = type_var.as_nested_type().as_struct_type() results = list() fields = struct.fields for field in fields: results.append(VisitFieldFuture(field, visitor)) struct = visitor.struct(struct, [x.get() for x in results]) return struct return visitor.primitive(type_var.as_primitive_type())