Update pandas to polars for feature processing (Fix #86)

CHANGELOG.md

@@ -4,6 +4,10 @@
 [Doing] Add support for saving pb file, exporting embeddings
 [Doing] Add support of NVTabular data
 
+**FuxiCTR v2.2.2, 2024-04-18**
++ [Feature] Update to use polars instead of pandas for faster feature processing
++ [Fix] When num_workers > 1, NpzBlockDataLoader cannot keep the reading order of samples ([#86](https://github.com/xue-pai/FuxiCTR/issues/86))
+
 **FuxiCTR v2.2.1, 2024-04-16**
 + [Fix] Fix issue of evaluation not performed at epoch end when streaming=True ([#85](https://github.com/xue-pai/FuxiCTR/issues/85))
 + [Fix] Fix issue when loading pretrain_emb in npz format ([#84](https://github.com/xue-pai/FuxiCTR/issues/84))

fuxictr/preprocess/build_dataset.py

@@ -21,6 +21,7 @@
 import numpy as np
 import gc
 import multiprocessing as mp
+import polars as pl
 
 
 def split_train_test(train_ddf=None, valid_ddf=None, test_ddf=None, valid_size=0, 
@@ -53,28 +54,28 @@ def save_npz(darray_dict, data_path):
  np.savez(data_path, **darray_dict)
 
 
-def transform_block(feature_encoder, df_block, filename, preprocess=False):
- if preprocess:
- df_block = feature_encoder.preprocess(df_block)
+def transform_block(feature_encoder, df_block, filename):
  darray_dict = feature_encoder.transform(df_block)
  save_npz(darray_dict, os.path.join(feature_encoder.data_dir, filename))
 
 
-def transform(feature_encoder, ddf, filename, preprocess=False, block_size=0):
+def transform(feature_encoder, ddf, filename, block_size=0):
  if block_size > 0:
  pool = mp.Pool(mp.cpu_count() // 2)
  block_id = 0
  for idx in range(0, len(ddf), block_size):
- df_block = ddf[idx: (idx + block_size)]
- pool.apply_async(transform_block, args=(feature_encoder,
- df_block,
- '{}/part_{:05d}.npz'.format(filename, block_id),
- preprocess))
+ df_block = ddf.iloc[idx:(idx + block_size)]
+ pool.apply_async(
+ transform_block, 
+ args=(feature_encoder,
+ df_block,
+ '{}/part_{:05d}.npz'.format(filename, block_id))
+ )
  block_id += 1
  pool.close()
  pool.join()
  else:
- transform_block(feature_encoder, ddf, filename, preprocess)
+ transform_block(feature_encoder, ddf, filename)
 
 
 def build_dataset(feature_encoder, train_data=None, valid_data=None, test_data=None, valid_size=0, 
@@ -96,28 +97,30 @@ def build_dataset(feature_encoder, train_data=None, valid_data=None, test_data=N
  valid_ddf = feature_encoder.read_csv(valid_data, **kwargs)
  test_ddf = feature_encoder.read_csv(test_data, **kwargs)
  train_ddf, valid_ddf, test_ddf = split_train_test(train_ddf, valid_ddf, test_ddf, 
- valid_size, test_size, split_type)
+  valid_size, test_size, split_type)
 
  # fit and transform train_ddf
  train_ddf = feature_encoder.preprocess(train_ddf)
  feature_encoder.fit(train_ddf, **kwargs)
- transform(feature_encoder, train_ddf, 'train', preprocess=False, block_size=data_block_size)
+ transform(feature_encoder, train_ddf, 'train', block_size=data_block_size)
  del train_ddf
  gc.collect()
 
  # Transfrom valid_ddf
  if valid_ddf is None and (valid_data is not None):
  valid_ddf = feature_encoder.read_csv(valid_data, **kwargs)
  if valid_ddf is not None:
- transform(feature_encoder, valid_ddf, 'valid', preprocess=True, block_size=data_block_size)
+ valid_ddf = feature_encoder.preprocess(valid_ddf)
+ transform(feature_encoder, valid_ddf, 'valid', block_size=data_block_size)
  del valid_ddf
  gc.collect()
 
  # Transfrom test_ddf
  if test_ddf is None and (test_data is not None):
  test_ddf = feature_encoder.read_csv(test_data, **kwargs)
  if test_ddf is not None:
- transform(feature_encoder, test_ddf, 'test', preprocess=True, block_size=data_block_size)
+ test_ddf = feature_encoder.preprocess(test_ddf)
+ transform(feature_encoder, test_ddf, 'test', block_size=data_block_size)
  del test_ddf
  gc.collect()
  logging.info("Transform csv data to npz done.")

fuxictr/preprocess/feature_processor.py

@@ -18,12 +18,14 @@
 import numpy as np
 from collections import Counter, OrderedDict
 import pandas as pd
+import polars as pl
 import pickle
 import os
 import logging
 import json
 import re
 import shutil
+import glob
 from pathlib import Path
 import sklearn.preprocessing as sklearn_preprocess
 from fuxictr.features import FeatureMap
@@ -65,11 +67,14 @@ def _complete_feature_cols(self, feature_cols):
  full_feature_cols.append(col)
  return full_feature_cols
 
- def read_csv(self, data_path, sep=",", nrows=None, **kwargs):
+ def read_csv(self, data_path, sep=",", n_rows=None, **kwargs):
  logging.info("Reading file: " + data_path)
- usecols_fn = lambda x: x in self.dtype_dict
- ddf = pd.read_csv(data_path, sep=sep, usecols=usecols_fn, 
- dtype=object, memory_map=True, nrows=nrows)
+ file_names = sorted(glob.glob(data_path))
+ assert len(file_names) > 0, f"Invalid data path: {data_path}"
+ # Require python >= 3.8 for use polars to scan multiple csv files
+ file_names = file_names[0]
+ ddf = pl.scan_csv(source=file_names, separator=sep, dtypes=self.dtype_dict,
+ low_memory=False, n_rows=n_rows)
  return ddf
 
  def preprocess(self, ddf):
@@ -78,50 +83,38 @@ def preprocess(self, ddf):
  for col in all_cols:
  name = col["name"]
  if name in ddf.columns:
- if ddf[name].isnull().values.any():
-  ddf[name] = self._fill_na_(col, ddf[name])
- ddf[name] = ddf[name].astype(self.dtype_dict[name])
+ fill_na = "" if col["dtype"] in ["str", str] else 0
+ fill_na = col.get("fill_na", fill_na)
+ ddf = ddf.with_columns(pl.col(name).fill_null(fill_na))
  if col.get("preprocess"):
- preprocess_splits = re.split(r"\(|\)", col["preprocess"])
- preprocess_fn = getattr(self, preprocess_splits[0])
- if len(preprocess_splits) > 1:
- ddf[name] = preprocess_fn(ddf, preprocess_splits[1])
- else:
- ddf[name] = preprocess_fn(ddf, name)
- ddf[name] = ddf[name].astype(self.dtype_dict[name])
+ preprocess_args = re.split(r"\(|\)", col["preprocess"])
+ preprocess_fn = getattr(self, preprocess_args[0])
+ ddf = preprocess_fn(ddf, name, *preprocess_args[1:-1])
+ ddf = ddf.with_columns(pl.col(name).cast(self.dtype_dict[name]))
  active_cols = [col["name"] for col in all_cols if col.get("active") != False]
- ddf = ddf.loc[:, active_cols]
+ ddf = ddf.select(active_cols)
  return ddf
 
- def _fill_na_(self, col, series):
- na_value = col.get("fill_na")
- if na_value is not None:
- return series.fillna(na_value)
- elif col["dtype"] in ["str", str]:
- return series.fillna("")
- else:
- raise RuntimeError("Feature column={} requires to assign fill_na value!".format(col["name"]))
-
  def fit(self, train_ddf, min_categr_count=1, num_buckets=10, **kwargs): 
  logging.info("Fit feature processor...")
  for col in self.feature_cols:
  name = col["name"]
  if col["active"]:
  logging.info("Processing column: {}".format(col))
+ col_series = train_ddf.select(name).collect().to_series().to_pandas()
  if col["type"] == "meta": # e.g. group_id
- self.fit_meta_col(col, train_ddf[name].values)
+ self.fit_meta_col(col)
  elif col["type"] == "numeric":
- self.fit_numeric_col(col, train_ddf[name].values)
+ self.fit_numeric_col(col, col_series)
  elif col["type"] == "categorical":
-
- self.fit_categorical_col(col, train_ddf[name].values, 
+ self.fit_categorical_col(col, col_series,
  min_categr_count=min_categr_count,
  num_buckets=num_buckets)
  elif col["type"] == "sequence":
- self.fit_sequence_col(col, train_ddf[name].values, 
+ self.fit_sequence_col(col, col_series, 
  min_categr_count=min_categr_count)
  else:
- raise NotImplementedError("feature_col={}".format(feature_col))
+ raise NotImplementedError("feature type={}".format(col["type"]))
 
  # Expand vocab from pretrained_emb
  os.makedirs(self.data_dir, exist_ok=True)
@@ -166,17 +159,16 @@ def fit(self, train_ddf, min_categr_count=1, num_buckets=10, **kwargs):
  self.feature_map.save(self.json_file)
  logging.info("Set feature processor done.")
 
- def fit_meta_col(self, col, col_values):
+ def fit_meta_col(self, col):
  name = col["name"]
  feature_type = col["type"]
  self.feature_map.features[name] = {"type": feature_type}
- # assert col.get("remap") == False, "Meta feature currently only supports `remap=False`, \
- # since it needs to map train and validation sets together."
  if col.get("remap", True):
+ # No need to fit, update vocab in encode_meta()
  tokenizer = Tokenizer(min_freq=1, remap=True)
  self.processor_dict[name + "::tokenizer"] = tokenizer
 
- def fit_numeric_col(self, col, col_values):
+ def fit_numeric_col(self, col, col_series):
  name = col["name"]
  feature_type = col["type"]
  feature_source = col.get("source", "")
@@ -186,10 +178,10 @@ def fit_numeric_col(self, col, col_values):
  self.feature_map.features[name]["feature_encoder"] = col["feature_encoder"]
  if "normalizer" in col:
  normalizer = Normalizer(col["normalizer"])
- normalizer.fit(col_values)
+ normalizer.fit(col_series.dropna().values)
  self.processor_dict[name + "::normalizer"] = normalizer
 
- def fit_categorical_col(self, col, col_values, min_categr_count=1, num_buckets=10):
+ def fit_categorical_col(self, col, col_series, min_categr_count=1, num_buckets=10):
  name = col["name"]
  feature_type = col["type"]
  feature_source = col.get("source", "")
@@ -206,7 +198,7 @@ def fit_categorical_col(self, col, col_values, min_categr_count=1, num_buckets=1
  tokenizer = Tokenizer(min_freq=min_categr_count, 
  na_value=col.get("fill_na", ""), 
  remap=col.get("remap", True))
- tokenizer.fit_on_texts(col_values)
+ tokenizer.fit_on_texts(col_series)
  if "share_embedding" in col:
  self.feature_map.features[name]["share_embedding"] = col["share_embedding"]
  tknzr_name = col["share_embedding"] + "::tokenizer"
@@ -225,20 +217,18 @@ def fit_categorical_col(self, col, col_values, min_categr_count=1, num_buckets=1
  if category_processor == "quantile_bucket": # transform numeric value to bucket
  num_buckets = col.get("num_buckets", num_buckets)
  qtf = sklearn_preprocess.QuantileTransformer(n_quantiles=num_buckets + 1)
- qtf.fit(col_values)
+ qtf.fit(col_series.values)
  boundaries = qtf.quantiles_[1:-1]
  self.feature_map.features[name]["vocab_size"] = num_buckets
  self.processor_dict[name + "::boundaries"] = boundaries
  elif category_processor == "hash_bucket":
  num_buckets = col.get("num_buckets", num_buckets)
- uniques = Counter(col_values)
- num_buckets = min(num_buckets, len(uniques))
  self.feature_map.features[name]["vocab_size"] = num_buckets
  self.processor_dict[name + "::num_buckets"] = num_buckets
  else:
  raise NotImplementedError("category_processor={} not supported.".format(category_processor))
 
- def fit_sequence_col(self, col, col_values, min_categr_count=1):
+ def fit_sequence_col(self, col, col_series, min_categr_count=1):
  name = col["name"]
  feature_type = col["type"]
  feature_source = col.get("source", "")
@@ -259,7 +249,7 @@ def fit_sequence_col(self, col, col_values, min_categr_count=1):
  tokenizer = Tokenizer(min_freq=min_categr_count, splitter=splitter, 
  na_value=na_value, max_len=max_len, padding=padding,
  remap=col.get("remap", True))
- tokenizer.fit_on_texts(col_values)
+ tokenizer.fit_on_texts(col_series)
  if "share_embedding" in col:
  self.feature_map.features[name]["share_embedding"] = col["share_embedding"]
  tknzr_name = col["share_embedding"] + "::tokenizer"
@@ -275,39 +265,39 @@ def fit_sequence_col(self, col, col_values, min_categr_count=1):
  "vocab_size": tokenizer.vocab_size()})
 
  def transform(self, ddf):
- logging.info("Transform feature columns...")
+ logging.info("Transform feature columns with ID mapping...")
  data_dict = dict()
  for feature, feature_spec in self.feature_map.features.items():
  if feature in ddf.columns:
  feature_type = feature_spec["type"]
- col_values = ddf.loc[:, feature].values
+ col_series = ddf[feature]
  if feature_type == "meta":
  if feature + "::tokenizer" in self.processor_dict:
  tokenizer = self.processor_dict[feature + "::tokenizer"]
- data_dict[feature] = tokenizer.encode_meta(col_values)
+ data_dict[feature] = tokenizer.encode_meta(col_series)
  # Update vocab in tokenizer
  self.processor_dict[feature + "::tokenizer"] = tokenizer
  else:
- data_dict[feature] = col_values.astype(self.dtype_dict[feature])
+ data_dict[feature] = col_series.values
  elif feature_type == "numeric":
- col_values = col_values.astype(float)
+ col_values = col_series.values
  normalizer = self.processor_dict.get(feature + "::normalizer")
  if normalizer:
  col_values = normalizer.transform(col_values)
  data_dict[feature] = col_values
  elif feature_type == "categorical":
  category_processor = feature_spec.get("category_processor")
  if category_processor is None:
- data_dict[feature] = self.processor_dict.get(feature + "::tokenizer").encode_category(col_values)
+ data_dict[feature] = self.processor_dict.get(feature + "::tokenizer").encode_category(col_series)
  elif category_processor == "numeric_bucket":
  raise NotImplementedError
  elif category_processor == "hash_bucket":
  raise NotImplementedError
  elif feature_type == "sequence":
- data_dict[feature] = self.processor_dict.get(feature + "::tokenizer").encode_sequence(col_values)
+ data_dict[feature] = self.processor_dict.get(feature + "::tokenizer").encode_sequence(col_series)
  for label in self.feature_map.labels:
  if label in ddf.columns:
- data_dict[label] = ddf.loc[:, label].values
+ data_dict[label] = ddf[label].values
  return data_dict
 
  def load_pickle(self, pickle_file=None):
@@ -335,6 +325,6 @@ def save_vocab(self, vocab_file):
  with open(vocab_file, "w") as fd:
  fd.write(json.dumps(vocab, indent=4))
 
- def copy_from(self, ddf, src_name):
- return ddf[src_name]
-
+ def copy_from(self, ddf, name, src_name):
+ ddf = ddf.with_columns(pl.col(src_name).alias(name))
+ return ddf

fuxictr/preprocess/normalizer.py

@@ -33,8 +33,7 @@ def __init__(self, normalizer):
 
  def fit(self, X):
  if not self.callable:
- null_index = np.isnan(X)
- self.normalizer.fit(X[~null_index].reshape(-1, 1))
+ self.normalizer.fit(X.reshape(-1, 1))
 
  def transform(self, X):
  if self.callable: