def CanHandle(arrow_schema: pa.Schema,
tensor_representation: schema_pb2.TensorRepresentation) -> bool:
depth, value_type = _GetNestDepthAndValueType(
arrow_schema.field_by_name(
tensor_representation.varlen_sparse_tensor.column_name))
# Currently can only handle 1-nested lists, but can easily support
# arbitrarily nested ListArrays.
return depth == 1 and _IsSupportedArrowValueType(value_type)
def BaseCanHandle(
arrow_schema: pa.Schema,
tensor_representation: schema_pb2.TensorRepresentation) -> bool:
depth, value_type = _GetNestDepthAndValueType(
arrow_schema.field_by_name(
tensor_representation.dense_tensor.column_name))
# Can only handle 1-nested lists.
return depth == 1 and _IsSupportedArrowValueType(value_type)
def CanHandle(
arrow_schema: pa.Schema,
tensor_representation: schema_pb2.TensorRepresentation) -> bool:
"""Returns whether `tensor_representation` can be handled."""
sparse_representation = tensor_representation.sparse_tensor
if (len(sparse_representation.dense_shape.dim) !=
len(sparse_representation.index_column_names)):
return False
if any([d.size <= 0 for d in sparse_representation.dense_shape.dim]):
return False
# All the index columns must be of integral types.
for index_column in sparse_representation.index_column_names:
depth, value_type = _GetNestDepthAndValueType(
arrow_schema.field_by_name(index_column))
if depth != 1 or not pa.types.is_integer(value_type):
return False
depth, value_type = _GetNestDepthAndValueType(
arrow_schema.field_by_name(sparse_representation.value_column_name))
return depth == 1 and _IsSupportedArrowValueType(value_type)
