def tzone_compute(con, guid, tz):
schema = ibis.schema([('ts', dt.Timestamp(tz)), ('b', 'double'),
('c', 'string')])
con.create_table(guid, schema=schema)
t = con.table(guid)
n = 10
df = pd.DataFrame({
'ts': pd.date_range('2017-04-01', periods=n, tz=tz).values,
'b': np.arange(n).astype('float64'),
'c': list(string.ascii_lowercase[:n]),
})
df.to_sql(
guid,
con.con,
index=False,
if_exists='append',
dtype={
'ts': sa.TIMESTAMP(timezone=True),
'b': sa.FLOAT,
'c': sa.TEXT
},
)
try:
yield t
finally:
con.drop_table(guid)
assert guid not in con.list_tables()
def sqlalchemy_type_to_ibis_type(column_type, nullable=True):
type_class = type(column_type)
if isinstance(column_type, sa.types.NUMERIC):
return dt.Decimal(column_type.precision,
column_type.scale,
nullable=nullable)
else:
if type_class in _sqla_type_to_ibis:
ibis_class = _sqla_type_to_ibis[type_class]
elif isinstance(column_type, sa.DateTime):
ibis_class = dt.Timestamp()
elif isinstance(column_type, sa.ARRAY):
dimensions = column_type.dimensions
if dimensions is not None and dimensions != 1:
raise NotImplementedError(
'Nested array types not yet supported')
value_type = sqlalchemy_type_to_ibis_type(column_type.item_type)
def make_array_type(nullable, value_type=value_type):
return dt.Array(value_type, nullable=nullable)
ibis_class = make_array_type
else:
try:
ibis_class = next(v for k, v in _sqla_type_mapping.items()
if isinstance(column_type, k))
except StopIteration:
raise NotImplementedError(
'Unable to convert SQLAlchemy type {} to ibis type'.format(
column_type))
return ibis_class(nullable)
def schema_from_table(table):
# Convert SQLA table to Ibis schema
names = table.columns.keys()
types = []
for c in table.columns.values():
type_class = type(c.type)
if isinstance(c.type, sa.types.NUMERIC):
t = dt.Decimal(c.type.precision, c.type.scale, nullable=c.nullable)
else:
if c.type in _sqla_type_to_ibis:
ibis_class = _sqla_type_to_ibis[c.type]
elif type_class in _sqla_type_to_ibis:
ibis_class = _sqla_type_to_ibis[type_class]
elif isinstance(c.type, sa.DateTime):
ibis_class = dt.Timestamp()
else:
for k, v in _sqla_type_to_ibis.items():
if isinstance(c.type, type(k)):
ibis_class = v
break
else:
raise NotImplementedError(c.type)
t = ibis_class(c.nullable)
types.append(t)
return dt.Schema(names, types)
def pandas_dtypes_to_ibis_schema(df, schema):
dtypes = df.dtypes
pairs = []
for column_name, dtype in dtypes.iteritems():
if not isinstance(column_name, six.string_types):
raise TypeError(
'Column names must be strings to use the pandas backend')
if column_name in schema:
ibis_type = dt.validate_type(schema[column_name])
elif dtype == np.object_:
inferred_dtype = infer_dtype(df[column_name].dropna())
if inferred_dtype == 'mixed':
raise TypeError(
'Unable to infer type of column {0!r}. Try instantiating '
'your table from the client with client.table('
"'my_table', schema={{{0!r}: <explicit type>}})".format(
column_name))
ibis_type = _INFERRED_DTYPE_TO_IBIS_TYPE[inferred_dtype]
elif hasattr(dtype, 'tz'):
ibis_type = dt.Timestamp(str(dtype.tz))
else:
dtype_string = str(dtype)
ibis_type = _DTYPE_TO_IBIS_TYPE.get(dtype_string, dtype_string)
pairs.append((column_name, ibis_type))
return ibis.schema(pairs)
def test_kudu_schema_convert(self):
spec = [
# name, type, is_nullable, is_primary_key
('a', dt.Int8(False), 'int8', False, True),
('b', dt.Int16(False), 'int16', False, True),
('c', dt.Int32(False), 'int32', False, False),
('d', dt.Int64(True), 'int64', True, False),
('e', dt.String(True), 'string', True, False),
('f', dt.Boolean(False), 'bool', False, False),
('g', dt.Float(False), 'float', False, False),
('h', dt.Double(True), 'double', True, False),
# TODO
# ('i', 'binary', False, False),
('j', dt.Timestamp(True), 'timestamp', True, False),
]
builder = kudu.schema_builder()
primary_keys = []
ibis_types = []
for name, itype, type_, is_nullable, is_primary_key in spec:
builder.add_column(name, type_, nullable=is_nullable)
if is_primary_key:
primary_keys.append(name)
ibis_types.append((name, itype))
builder.set_primary_keys(primary_keys)
kschema = builder.build()
ischema = ksupport.schema_kudu_to_ibis(kschema)
expected = ibis.schema(ibis_types)
assert_equal(ischema, expected)
def test_timestamp_with_timezone():
df = pd.DataFrame(
{'A': pd.date_range('20130101', periods=3, tz='US/Eastern')})
schema = sch.infer(df)
expected = ibis.schema([('A', "timestamp('US/Eastern')")])
assert schema.equals(expected)
assert schema.types[0].equals(dt.Timestamp('US/Eastern'))
def test_timestamp_type_accepts_all_timezones(con):
assert all(
dt.Timestamp(row.name).timezone == row.name
for row in con.con.execute(
'SELECT name FROM pg_timezone_names'
)
)
def _timestamp_from_str(value: str,
timezone: str | None = None) -> ir.TimestampScalar:
try:
value = pd.Timestamp(value, tz=timezone)
except pd.errors.OutOfBoundsDatetime:
value = dateutil.parser.parse(value)
dtype = dt.Timestamp(
timezone=timezone if timezone is not None else value.tzname())
return literal(value, type=dtype)
def test_timestamp_cast_noop(alltypes, translate):
target = dt.Timestamp(nullable=False)
result1 = alltypes.timestamp_col.cast(target)
result2 = alltypes.int_col.cast(target)
assert isinstance(result1, ir.TimestampColumn)
assert isinstance(result2, ir.TimestampColumn)
assert translate(result1) == '`timestamp_col`'
assert translate(result2) == 'CAST(`int_col` AS DateTime)'
def test_times_ops_with_tz(t, df, tz, rconstruct, column):
expected = rconstruct(len(df), dtype=bool)
expr = t[column].time().between('01:00', '02:00', timezone=tz)
result = expr.execute()
tm.assert_numpy_array_equal(result, expected)
# Test that casting behavior is the same as using the timezone kwarg
ts = t[column].cast(dt.Timestamp(timezone=tz))
expr = ts.time().between('01:00', '02:00')
result = expr.execute()
tm.assert_numpy_array_equal(result, expected)
def pandas_col_to_ibis_type(col):
import numpy as np
dty = col.dtype
# datetime types
if pdcom.is_datetime64tz_dtype(dty):
return dt.Timestamp(str(dty.tz))
if pdcom.is_datetime64_dtype(dty):
if pdcom.is_datetime64_ns_dtype(dty):
return dt.timestamp
else:
raise com.IbisTypeError("Column {0} has dtype {1}, which is "
"datetime64-like but does "
"not use nanosecond units".format(
col.name, dty))
if pdcom.is_timedelta64_dtype(dty):
print("Warning: encoding a timedelta64 as an int64")
return dt.int64
if pdcom.is_categorical_dtype(dty):
return dt.Category(len(col.cat.categories))
if pdcom.is_bool_dtype(dty):
return dt.boolean
# simple numerical types
if issubclass(dty.type, np.int8):
return dt.int8
if issubclass(dty.type, np.int16):
return dt.int16
if issubclass(dty.type, np.int32):
return dt.int32
if issubclass(dty.type, np.int64):
return dt.int64
if issubclass(dty.type, np.float32):
return dt.float
if issubclass(dty.type, np.float64):
return dt.double
if issubclass(dty.type, np.uint8):
return dt.int16
if issubclass(dty.type, np.uint16):
return dt.int32
if issubclass(dty.type, np.uint32):
return dt.int64
if issubclass(dty.type, np.uint64):
raise com.IbisTypeError("Column {} is an unsigned int64".format(
col.name))
if pdcom.is_object_dtype(dty):
return _infer_object_dtype(col)
raise com.IbisTypeError("Column {0} is dtype {1}".format(col.name, dty))
def test_times_ops_with_tz(t, df, tz, rconstruct, column):
expected = dd.from_array(rconstruct(len(df), dtype=bool), )
time = t[column].time()
expr = time.between('01:00', '02:00', timezone=tz)
result = expr.compile()
tm.assert_series_equal(result.compute(), expected.compute())
# Test that casting behavior is the same as using the timezone kwarg
ts = t[column].cast(dt.Timestamp(timezone=tz))
expr = ts.time().between('01:00', '02:00')
result = expr.compile()
tm.assert_series_equal(result.compute(), expected.compute())
def between(
self,
lower: str | datetime.time | TimeValue,
upper: str | datetime.time | TimeValue,
timezone: str | None = None,
) -> ir.BooleanValue:
"""Check if the expr falls between `lower` and `upper`, inclusive.
Adjusts according to `timezone` if provided.
Parameters
----------
lower
Lower bound
upper
Upper bound
timezone
Time zone
Returns
-------
BooleanValue
Whether `self` is between `lower` and `upper`, adjusting `timezone`
as needed.
"""
import ibis.expr.datatypes as dt
import ibis.expr.operations as ops
op = self.op()
if isinstance(op, ops.Time):
# Here we pull out the first argument to the underlying Time
# operation which is by definition (in _timestamp_value_methods) a
# TimestampValue. We do this so that we can potentially specialize
# the "between time" operation for
# timestamp_value_expr.time().between(). A similar mechanism is
# triggered when creating expressions like
# t.column.distinct().count(), which is turned into
# t.column.nunique().
arg = op.arg
if timezone is not None:
arg = arg.cast(dt.Timestamp(timezone=timezone))
op_cls = ops.BetweenTime
else:
arg = self
op_cls = ops.Between
return op_cls(arg, lower, upper).to_expr()
def test_timestamp_with_timezone_repr():
ts = dt.Timestamp('UTC')
assert repr(ts) == "Timestamp(timezone='UTC', nullable=True)"
(np.int16(-1), dt.int16),
(np.int32(2), dt.int32),
(np.int64(-5), dt.int64),
(np.uint8(5), dt.uint8),
(np.uint16(50), dt.uint16),
(np.uint32(500), dt.uint32),
(np.uint64(5000), dt.uint64),
(np.float32(5.5), dt.float32),
(np.float64(5.55), dt.float64),
(np.bool_(True), dt.boolean),
(np.bool_(False), dt.boolean),
(np.arange(5, dtype='int32'), dt.Array(dt.int32)),
# pandas types
(
pd.Timestamp('2015-01-01 12:00:00', tz='US/Eastern'),
dt.Timestamp('US/Eastern'),
),
# TODO - add in Period/Interval/Int/Categorical
],
)
def test_infer_dtype(value, expected_dtype):
assert dt.infer(value) == expected_dtype
@pytest.mark.parametrize(
('numpy_dtype', 'ibis_dtype'),
[
(np.bool_, dt.boolean),
(np.int8, dt.int8),
(np.int16, dt.int16),
(np.int32, dt.int32),
def spark_timestamp_dtype_to_ibis_dtype(spark_type_obj, nullable=True):
return dt.Timestamp(nullable=nullable)
def test_timestamp_with_timezone_is_inferred_correctly(t, df):
assert t.plain_datetimes_naive.type().equals(dt.timestamp)
assert t.plain_datetimes_ny.type().equals(dt.Timestamp('America/New_York'))
assert t.plain_datetimes_utc.type().equals(dt.Timestamp('UTC'))
@pytest.mark.parametrize(('to', 'expected'), [('double', 'float64'),
('string', 'object')])
def test_cast_string(t, df, from_, to, expected):
c = t[from_].cast(to)
result = c.execute()
assert str(result.dtype) == expected
@pytest.mark.parametrize(
('to', 'expected'),
[
('string', 'object'),
('int64', 'int64'),
param('double', 'float64', marks=pytest.mark.xfail(raises=TypeError)),
(
dt.Timestamp('America/Los_Angeles'),
'datetime64[ns, America/Los_Angeles]',
),
(
"timestamp('America/Los_Angeles')",
'datetime64[ns, America/Los_Angeles]',
),
],
)
@pytest.mark.parametrize(
'column',
['plain_datetimes_naive', 'plain_datetimes_ny', 'plain_datetimes_utc'],
)
def test_cast_timestamp_column(t, df, column, to, expected):
c = t[column].cast(to)
result = c.execute()
def _(value: str, timezone: str | None = None) -> ir.TimestampScalar:
try:
value = pd.Timestamp(value, tz=timezone)
except pd.errors.OutOfBoundsDatetime:
value = dateutil.parser.parse(value)
return literal(value, type=dt.Timestamp(timezone=timezone))
def pa_timestamp_type(arrow_type, nullable=True):
return dt.Timestamp(arrow_type.tz, nullable=nullable)
def test_ts_timezone_is_preserved(tzone_compute, tz):
assert dt.Timestamp(tz).equals(tzone_compute.ts.type())
def from_pandas_tzdtype(value):
return dt.Timestamp(timezone=str(value.tz))
def infer_pandas_timestamp(value):
if value.tz is not None:
return dt.Timestamp(timezone=str(value.tz))
else:
return dt.timestamp
def test_timestamp_with_timezone_str():
ts = dt.Timestamp('UTC')
assert str(ts) == "timestamp('UTC')"
def type(self):
return dt.Timestamp(timezone=self._timezone)
def my_string_length(series, **kwargs):
return series.str.len() * 2
@elementwise(input_type=[dt.double, dt.double], output_type=dt.double)
def my_add(series1, series2, **kwargs):
return series1 + series2
@reduction(['double'], 'double')
def my_mean(series):
return series.mean()
@reduction(
input_type=[dt.Timestamp(timezone="UTC")],
output_type=dt.Timestamp(timezone="UTC"),
)
def my_tz_min(series):
return series.min()
@reduction(input_type=[dt.string], output_type=dt.int64)
def my_string_length_sum(series, **kwargs):
return (series.str.len() * 2).sum()
@reduction(input_type=[dt.double, dt.double], output_type=dt.double)
def my_corr(lhs, rhs, **kwargs):
return lhs.corr(rhs)
def sa_datetime(_, satype, nullable=True, default_timezone='UTC'):
timezone = default_timezone if satype.timezone else None
return dt.Timestamp(timezone=timezone, nullable=nullable)
(dt.int8, 'INT64'),
(dt.int16, 'INT64'),
(dt.int32, 'INT64'),
(dt.int64, 'INT64'),
(dt.string, 'STRING'),
(dt.Array(dt.int64), 'ARRAY<INT64>'),
(dt.Array(dt.string), 'ARRAY<STRING>'),
(
dt.Struct.from_tuples([('a', dt.int64), ('b', dt.string),
('c', dt.Array(dt.string))]),
'STRUCT<a INT64, b STRING, c ARRAY<STRING>>',
),
(dt.date, 'DATE'),
(dt.timestamp, 'TIMESTAMP'),
param(
dt.Timestamp(timezone='US/Eastern'),
'TIMESTAMP',
marks=pytest.mark.xfail(raises=TypeError,
reason='Not supported in BigQuery'),
),
('array<struct<a: string>>', 'ARRAY<STRUCT<a STRING>>'),
param(
dt.Decimal(38, 9),
'NUMERIC',
marks=pytest.mark.xfail(raises=TypeError,
reason='Not supported in BigQuery'),
),
],
)
def test_simple(datatype, expected):
context = TypeTranslationContext()
def test_timestamp_timezone_type(tz):
expr = ibis.timestamp('2017-01-01', timezone=tz)
expected = dt.Timestamp(timezone=tz)
assert expected == expr.op().dtype
def _(value, timezone: str | None = None) -> ir.TimestampScalar:
return literal(value, type=dt.Timestamp(timezone=timezone))
