使用Polars Series作为Scikit-Learn TfidfVectorizer输入报错求助

我们正计划为string_grouper添加Polars支持，第一步需要实现用Scikit-Learn的TfidfVectorizer处理Polars Series。根据Polars官方文档，该功能应该被支持，但实际运行代码时出现错误，想确认是操作有误还是确实不支持。

运行代码 #

from sklearn.feature_extraction.text import TfidfVectorizer

company_names = df.select(pl.col('Company Name'))
vectorizer = TfidfVectorizer(min_df=1)
tf_idf_matrix = vectorizer.fit_transform(company_names)

company_names输出 #

print(company_names)

shape: (663_000, 1)
┌─────────────────────────────────┐
│ Company Name                    │
│ ---                             │
│ str                             │
╞═════════════════════════════════╡
│ !J INC                          │
│ #1 A LIFESAFER HOLDINGS, INC.   │
│ #1 ARIZONA DISCOUNT PROPERTIES… │
│ #1 PAINTBALL CORP               │
│ $ LLC                           │

错误信息 #

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
Cell In[8], line 1
----> 1 tf_idf_matrix = vectorizer.fit_transform(company_names)

File ~/dev/python/sg_test/venv/lib/python3.12/site-packages/sklearn/feature_extraction/text.py:2104, in TfidfVectorizer.fit_transform(self, raw_documents, y)
   2097 self._check_params()
   2098 self._tfidf = TfidfTransformer(
   2099     norm=self.norm,
   2100     use_idf=self.use_idf,
   2101     smooth_idf=self.smooth_idf,
   2102     sublinear_tf=self.sublinear_tf,
   2103 )
-> 2104 X = super().fit_transform(raw_documents)
   2105 self._tfidf.fit(X)
   2106 # X is already a transformed view of raw_documents so
   2107 # we set copy to False

File ~/dev/python/sg_test/venv/lib/python3.12/site-packages/sklearn/base.py:1389, in _fit_context.<locals>.decorator.<locals>.wrapper(estimator, *args, **kwargs)
   1382     estimator._validate_params()
   1384 with config_context(
   1385     skip_parameter_validation=(
   1386         prefer_skip_nested_validation or global_skip_validation
   1387     )
   1388 ):
-> 1389     return fit_method(estimator, *args, **kwargs)

File ~/dev/python/sg_test/venv/lib/python3.12/site-packages/sklearn/feature_extraction/text.py:1376, in CountVectorizer.fit_transform(self, raw_documents, y)
   1368             warnings.warn(
   1369                 "Upper case characters found in"
   1370                 " vocabulary while 'lowercase'"
   1371                 " is True. These entries will not"
   1372                 " be matched with any documents"
   1373             )
   1374             break
-> 1376 vocabulary, X = self._count_vocab(raw_documents, self.fixed_vocabulary_)
   1378 if self.binary:
   1379     X.data.fill(1)

File ~/dev/python/sg_test/venv/lib/python3.12/site-packages/sklearn/feature_extraction/text.py:1263, in CountVectorizer._count_vocab(self, raw_documents, fixed_vocab)
   1261 for doc in raw_documents:
   1262     feature_counter = {}
-> 1263     for feature in analyze(doc):
   1264         try:
   1265             feature_idx = vocabulary[feature]

File ~/dev/python/sg_test/venv/lib/python3.12/site-packages/sklearn/feature_extraction/text.py:104, in _analyze(doc, analyzer, tokenizer, ngrams, preprocessor, decoder, stop_words)
    102 else:
    103     if preprocessor is not None:
-> 104         doc = preprocessor(doc)
    105     if tokenizer is not None:
    106         doc = tokenizer(doc)

File ~/dev/python/sg_test/venv/lib/python3.12/site-packages/sklearn/feature_extraction/text.py:62, in _preprocess(doc, accent_function, lower)
     43 """Chain together an optional series of text preprocessing steps to
     44 apply to a document.
     45 
   (...)
     59     preprocessed string
     60 """
     61 if lower:
-> 62     doc = doc.lower()
     63 if accent_function is not None:
     64     doc = accent_function(doc)

AttributeError: 'Series' object has no attribute 'lower'

解决方案 #

错误根源是df.select(pl.col('Company Name'))返回的是单列Polars DataFrame，而TfidfVectorizer需要接收可迭代的字符串序列（比如Polars Series）。Scikit-Learn支持Polars Series，但不支持直接传入单列DataFrame——此时迭代DataFrame会得到每一行的Series对象，而非字符串，导致调用lower()时出错。

只需修改获取列的方式，提取Polars Series即可：

from sklearn.feature_extraction.text import TfidfVectorizer
import polars as pl

# 提取Series而非单列DataFrame
company_names = df['Company Name']  # 等价于 df.get_column('Company Name')
vectorizer = TfidfVectorizer(min_df=1)
tf_idf_matrix = vectorizer.fit_transform(company_names)

此时传入的Series会被Scikit-Learn正确解析为字符串集合，每个元素对应一个文档，就能正常完成TF-IDF转换。

内容的提问来源于stack exchange，提问作者Chris van den Berg