def test_prep_df_for_csv_output_include_index(self):
schema_data = {
'schema': "bltypes/v1",
'fields': {
"date": {
"type": "date",
"index": 1,
},
"time": {
"type": "time",
"index": 2,
},
"timetz": {
"type": "timetz",
"index": 3,
},
}
}
records_format = DelimitedRecordsFormat(variant='bluelabs')
records_schema = RecordsSchema.from_data(schema_data)
processing_instructions = ProcessingInstructions()
# us_eastern = pytz.timezone('US/Eastern')
data = {
'time': [
pd.Timestamp(year=1970, month=1, day=1,
hour=12, minute=33, second=53, microsecond=1234)
],
# timetz is not well supported in records mover yet. For
# instance, specifying how it's turned into a CSV is not
# currently part of the records spec:
#
# https://github.com/bluelabsio/records-mover/issues/76
#
# In addition, Vertica suffers from a driver limitation:
#
# https://github.com/bluelabsio/records-mover/issues/77
#
# 'timetz': [
# us_eastern.localize(pd.Timestamp(year=1970, month=1, day=1,
# hour=12, minute=33, second=53,
# microsecond=1234)),
# ],
}
df = pd.DataFrame(data,
index=[pd.Timestamp(year=1970, month=1, day=1)],
columns=['time', 'timetz'])
new_df = prep_df_for_csv_output(df=df,
include_index=True,
records_schema=records_schema,
records_format=records_format,
processing_instructions=processing_instructions)
self.assertEqual(new_df.index[0], '1970-01-01')
self.assertEqual(new_df['time'][0], '12:33:53')
# self.assertEqual(new_df['timetz'][0], '12:33:53-05')
self.assertIsNotNone(new_df)
def save_df(df: 'DataFrame', output_filename: str) -> None:
df = prep_df_for_csv_output(
df,
include_index=self.include_index,
records_schema=records_schema,
records_format=delimited_records_format,
processing_instructions=processing_instructions)
df.to_csv(path_or_buf=output_filename,
index=self.include_index,
**options)
logger.info('CSV file written')
def test_timeonlyformat(self):
schema_data = {
'schema': "bltypes/v1",
'fields': {
"time_as_timestamp": {
"type": "time",
"index": 1,
},
"time_as_time": {
"type": "time",
"index": 2,
},
}
}
records_schema = RecordsSchema.from_data(schema_data)
processing_instructions = ProcessingInstructions()
for timeonlyformat in TIMEONLY_CASES:
records_format = DelimitedRecordsFormat(variant='bluelabs',
hints={
'timeonlyformat': timeonlyformat,
})
# us_eastern = pytz.timezone('US/Eastern')
time_as_timestamp = pd.Timestamp(year=SAMPLE_YEAR, month=SAMPLE_MONTH, day=SAMPLE_DAY,
hour=SAMPLE_HOUR, minute=SAMPLE_MINUTE,
second=SAMPLE_SECOND)
time_as_time = datetime.time(hour=SAMPLE_HOUR,
minute=SAMPLE_MINUTE,
second=SAMPLE_SECOND)
data = {
'time_as_timestamp': [
time_as_timestamp
],
'time_as_time': [
time_as_time
],
}
df = pd.DataFrame(data, columns=['time_as_timestamp', 'time_as_time'])
new_df = prep_df_for_csv_output(df=df,
include_index=False,
records_schema=records_schema,
records_format=records_format,
processing_instructions=processing_instructions)
self.assertEqual(new_df['time_as_timestamp'][0],
create_sample(timeonlyformat),
timeonlyformat)
self.assertEqual(new_df['time_as_time'][0],
create_sample(timeonlyformat),
timeonlyformat)
# self.assertEqual(new_df['timetz'][0], '12:33:53-05')
self.assertIsNotNone(new_df)
def test_dateformat(self):
schema_data = {
'schema': "bltypes/v1",
'fields': {
"date_as_timestamp": {
"type": "date",
"index": 1,
},
"date_as_date": {
"type": "date",
"index": 1,
},
}
}
records_schema = RecordsSchema.from_data(schema_data)
processing_instructions = ProcessingInstructions()
for dateformat in DATE_CASES:
records_format = DelimitedRecordsFormat(variant='bluelabs',
hints={
'dateformat': dateformat
})
# us_eastern = pytz.timezone('US/Eastern')
data = {
'date_as_timestamp': [
pd.Timestamp(year=SAMPLE_YEAR, month=SAMPLE_MONTH, day=SAMPLE_DAY)
],
'date_as_date': [
datetime.date(year=SAMPLE_YEAR, month=SAMPLE_MONTH, day=SAMPLE_DAY)
],
}
df = pd.DataFrame(data,
columns=['date_as_timestamp', 'date_as_date'])
new_df = prep_df_for_csv_output(df=df,
include_index=False,
records_schema=records_schema,
records_format=records_format,
processing_instructions=processing_instructions)
self.assertEqual(new_df['date_as_timestamp'][0],
create_sample(dateformat))
self.assertEqual(new_df['date_as_date'][0],
create_sample(dateformat))
self.assertIsNotNone(new_df)
def test_datetimeformat(self):
schema_data = {
'schema': "bltypes/v1",
'fields': {
"datetimez": {
"type": "datetime",
"index": 1,
},
}
}
records_schema = RecordsSchema.from_data(schema_data)
processing_instructions = ProcessingInstructions()
for datetimeformat in DATETIME_CASES:
records_format = DelimitedRecordsFormat(variant='bluelabs',
hints={
'datetimeformat': datetimeformat
})
# us_eastern = pytz.timezone('US/Eastern')
timestamp = pd.Timestamp(year=SAMPLE_YEAR, month=SAMPLE_MONTH, day=SAMPLE_DAY,
hour=SAMPLE_HOUR, minute=SAMPLE_MINUTE,
second=SAMPLE_SECOND)
data = {
'datetime': [
timestamp
],
}
df = pd.DataFrame(data, columns=['datetime'])
new_df = prep_df_for_csv_output(df=df,
include_index=False,
records_schema=records_schema,
records_format=records_format,
processing_instructions=processing_instructions)
# No conversion is done of datetime as pandas' CSV
# outputter handles it properly, so we should expect the
# original again
self.assertEqual(new_df['datetime'][0],
timestamp,
create_sample(datetimeformat))
# self.assertEqual(new_df['timetz'][0], '12:33:53-05')
self.assertIsNotNone(new_df)
def write_dfs(path_or_buf: Union[str, IO[str]]) -> int:
first_row = True
move_count = 0
for df in dfs_source.dfs:
logger.info("Appending from dataframe...")
# Include the header at most once in the file
include_header_row = options['header'] and first_row
first_row = False
options['header'] = include_header_row
df = prep_df_for_csv_output(df,
include_index=dfs_source.include_index,
records_schema=records_schema,
records_format=records_format,
processing_instructions=processing_instructions)
df.to_csv(path_or_buf=path_or_buf,
mode="a",
index=dfs_source.include_index,
**options)
move_count += len(df.index)
return move_count