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())