# Buggy code
def _chunkwise_moments(df):
    df2 = cudf.DataFrame()
    for col in df.columns:
        df2[col] = df[col].astype("float64").pow(2)
    vals = {
        "df-count": df.count().to_frame().transpose(),
        "df-sum": df.sum().to_frame().transpose(),
        "df2-sum": df2.sum().to_frame().transpose(),
    }
    # NOTE: Perhaps we should convert to pandas here
    # (since we know the results should be small)?
    del df2
    return vals


# Traceback (ValueError)
/opt/conda/envs/rapids/lib/python3.7/site-packages/pandas/core/series.py:726: RuntimeWarning: invalid value encountered in sqrt
result = getattr(ufunc, method)(*inputs, **kwargs)
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<timed eval> in <module>

/nvtabular0.3/NVTabular/nvtabular/workflow.py in apply(self, dataset, apply_offline, record_stats, shuffle, output_path, output_format, out_files_per_proc, num_io_threads, dtypes)
869                 out_files_per_proc=out_files_per_proc,
870                 num_io_threads=num_io_threads,
--> 871                 dtypes=dtypes,
872             )
873         else:

/nvtabular0.3/NVTabular/nvtabular/workflow.py in build_and_process_graph(self, dataset, end_phase, output_path, record_stats, shuffle, output_format, out_files_per_proc, apply_ops, num_io_threads, dtypes)
968             self._base_phase = 0  # Set _base_phase
969             for idx, _ in enumerate(self.phases[:end]):
--> 970                 self.exec_phase(idx, record_stats=record_stats, update_ddf=(idx == (end - 1)))
971             self._base_phase = 0  # Re-Set _base_phase
972

/nvtabular0.3/NVTabular/nvtabular/workflow.py in exec_phase(self, phase_index, record_stats, update_ddf)
755         _ddf = self.get_ddf()
756         if transforms:
--> 757             _ddf = self._aggregated_dask_transform(_ddf, transforms)
758
759         stats = []

/nvtabular0.3/NVTabular/nvtabular/workflow.py in _aggregated_dask_transform(self, ddf, transforms)
724         for transform in transforms:
725             columns_ctx, cols_grp, target_cols, logic, stats_context = transform
--> 726             meta = logic(meta, columns_ctx, cols_grp, target_cols, stats_context)
727         return ddf.map_partitions(self.__class__._aggregated_op, transforms, meta=meta)
728

/nvtabular0.3/NVTabular/nvtabular/ops/transform_operator.py in apply_op(self, gdf, columns_ctx, input_cols, target_cols, stats_context)
89         new_gdf = self.op_logic(gdf, target_columns, stats_context=stats_context)
90         self.update_columns_ctx(columns_ctx, input_cols, new_gdf.columns, target_columns)
---> 91         return self.assemble_new_df(gdf, new_gdf, target_columns)
92
93     def assemble_new_df(self, origin_gdf, new_gdf, target_columns):

/nvtabular0.3/NVTabular/nvtabular/ops/transform_operator.py in assemble_new_df(self, origin_gdf, new_gdf, target_columns)
96                 return new_gdf
97             else:
---> 98                 origin_gdf[target_columns] = new_gdf
99                 return origin_gdf
100         return cudf.concat([origin_gdf, new_gdf], axis=1)

/opt/conda/envs/rapids/lib/python3.7/contextlib.py in inner(*args, **kwds)
72         def inner(*args, **kwds):
73             with self._recreate_cm():
---> 74                 return func(*args, **kwds)
75         return inner
76

/opt/conda/envs/rapids/lib/python3.7/site-packages/cudf/core/dataframe.py in __setitem__(self, arg, value)
777                         replace_df=value,
778                         input_cols=arg,
--> 779                         mask=None,
780                     )
781                 else:

/opt/conda/envs/rapids/lib/python3.7/site-packages/cudf/core/dataframe.py in _setitem_with_dataframe(input_df, replace_df, input_cols, mask)
7266     if len(input_cols) != len(replace_df.columns):
7267         raise ValueError(
-> 7268             "Number of Input Columns must be same replacement Dataframe"
7269         )
7270

ValueError: Number of Input Columns must be same replacement Dataframe

# Fixed code
def _chunkwise_moments(df):
    df2 = cudf.DataFrame()
    for col in df.columns:
        df2[col] = df[col].astype("float64").pow(2)
    vals = {
        "df-count": df.count().to_frame().transpose(),
        "df-sum": df.sum().astype("float64").to_frame().transpose(),
        "df2-sum": df2.sum().to_frame().transpose(),
    }
    # NOTE: Perhaps we should convert to pandas here
    # (since we know the results should be small)?
    del df2
    return vals