def cast(self, data_type: Any) -> ColumnExpr:
other = self._copy()
other._as_name = self.as_name
other._distinct = self.is_distinct
other._as_type = None if data_type is None else to_pa_datatype(
data_type)
return other
def cast(self, data_type: Any) -> "ColumnExpr":
if self.wildcard and data_type is not None:
raise ValueError("'*' can't cast")
other = _NamedColumnExpr(self.name)
other._as_name = self.as_name
other._as_type = None if data_type is None else to_pa_datatype(
data_type)
return other
def __setitem__( # type: ignore
self, name: str, value: Any, *args: List[Any],
**kwds: Dict[str, Any]) -> None:
assert_arg_not_none(value, "value")
if not validate_column_name(name):
raise SchemaError(f"Invalid column name {name}")
if name in self: # update existing value is not allowed
raise SchemaError(f"{name} already exists in {self}")
if isinstance(value, pa.Field):
assert_or_throw(name == value.name,
SchemaError(f"{name} doesn't match {value}"))
elif isinstance(value, pa.DataType):
value = pa.field(name, value)
else:
value = pa.field(name, to_pa_datatype(value))
assert_or_throw(is_supported(value.type),
SchemaError(f"{value} is not supported"))
super().__setitem__(name, value, *args, **kwds) # type: ignore
def test_to_pa_datatype():
assert pa.int32() == to_pa_datatype(pa.int32())
assert pa.int32() == to_pa_datatype("int")
assert pa.int64() == to_pa_datatype(int)
assert pa.float64() == to_pa_datatype(float)
assert pa.float64() == to_pa_datatype(np.float64)
assert TRIAD_DEFAULT_TIMESTAMP == to_pa_datatype(datetime)
assert pa.date32() == to_pa_datatype(date)
assert pa.date32() == to_pa_datatype("date")
assert pa.binary() == to_pa_datatype("bytes")
assert pa.binary() == to_pa_datatype("binary")
raises(TypeError, lambda: to_pa_datatype(123))
raises(TypeError, lambda: to_pa_datatype(None))
def visitFugueSchemaSimpleType(
self, ctx: fp.FugueSchemaSimpleTypeContext
) -> pa.DataType:
return to_pa_datatype(self.ctxToStr(ctx))
def infer_type(self, schema: Schema) -> Optional[pa.DataType]:
if self.value is None:
return self.as_type
return self.as_type or to_pa_datatype(type(self.value))
def cast(self, data_type: Any) -> ColumnExpr:
other = _LiteralColumnExpr(self.value)
other._as_name = self.as_name
other._as_type = None if data_type is None else to_pa_datatype(
data_type)
return other