update README.md

.gitignore

@@ -0,0 +1,3 @@
+test.py
+test_eval.py
+*.pyc

LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

README-zh.md

@@ -0,0 +1,141 @@
+# 引言
+[Rain's SQLCoder](https://huggingface.co/SuanChang/rain-SQLCoder) 是自然语言生成 SparkSQL 的 SOTA 大型语言模型（LLM），拥有 32B 参数，基于 [Qwen2.5-Coder-32B-Instruct](https://huggingface.co/Qwen/Qwen2.5-Coder-32B-Instruct) 微调。 Rain's SQLCoder 针对自然语言到 SparkSQL 转换任务进行了优化，能够有效处理最长达 32k 个 token 的上下文，尤其适用于复杂且大规模的 SQL 查询生成任务。
+
+<p align="center">
+          🤗 <a href="https://huggingface.co/SuanChang/rain-SQLCoder">Hugging Face</a> | 🖥️ <a href="https://www.suan-chang.com/">演示</a> | 💬 <a href="./figures/wechat.png">微信</a> | <a href="https://github.com/suan-chang/rain-SQLCoder">GitHub</a>
+</p>
+
+[English](./README.md) | [中文](./README-zh.md)
+
+# 提示词
+Rain's SQLCoder 采用了 [Alpaca](https://github.com/tatsu-lab/stanford_alpaca) 模板，使用的提示词如下。
+````
+Below is an instruction that describes a task. 
+Write a response that appropriately completes the request.
+
+### Instruction:
+[BEGIN OF TASK INSTRUCTION]
+You are an expert in composing Spark SQL queries. You are given a user query and a set of table schemas.
+Based on the user query, you need to generate one Spark SQL query to achieve the purpose.
+{task description for date hint and related question and sqls}
+[END OF TASK INSTRUCTION]
+
+[BEGIN OF TABLE SCHEMAS]
+{schemas}
+[END OF TABLE SCHEMAS]
+
+[BEGIN OF GENERATION HINT]
+{date hint}
+[END OF GENERATION HINT]
+
+[BEGIN OF RELATED QUERIES]
+{related question and sqls}
+[END OF RELATED QUERIES]
+
+[BEGIN OF FORMAT INSTRUCTION]
+The output MUST strictly adhere to the following format, and NO other text MUST be included.
+```sql
+your output Spark SQL query
+``` 
+[END OF FORMAT INSTRUCTION]
+
+[BEGIN OF QUERY]
+User Query: {user question}
+[END OF QUERY]
+
+### Response:
+````
+
+# 评估
+我们沿用了 [SQL-Eval](https://github.com/defog-ai/sql-eval) 中评估预测结果与标准结果的逻辑：
+1. 如果预测的数据块和标准数据块完全一致，则预测结果正确；
+2. 标准SQL中不包含排序逻辑，且预测数据块和标准数据块在排序之后完全一致，则预测结果正确；
+3. 如果标准数据块的列是预测数据块的子集，则预测结果正确；
+4. 其余情况均认为预测结果错误。
+
+# 实验结果
+我们在两个测试集上对比了Rain's SQLCoder与国内外先进自然语言大模型的生成准确率。其中，基准测试集（Benchmark Dataset）包含基础样本，而增强测试集（Enhanced Dataset）则是在基准测试集的基础上，通过分层抽样方法选取20%的样本，并补充了相关的用户查询及对应的SparkSQL语句，以评估模型在增强上下文信息下的性能表现。实验结果表明，Rain's SQLCoder在查询意图理解、SQL语法准确性和复杂查询处理等方面均展现出显著优势。
+
+## 基准测试集
+<img src="./figures/benchmark_dataset_result.png" alt="benchmark" width=800>
+
+## 增强测试集
+<img src="./figures/enhanced_dataset_result.png" alt="enhanced" width=800>
+
+# 快速开始
+我们在此处提供示例，帮助您快速掌握如何加载并使用我们的模型。
+>注意: Rain's SQLCoder 只被训练用于生成 `SELECT` 语句，当表结构无法支持回答用户问题时，模型会拒绝回答。
+
+````python
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from utils.prompt import SQLGeneratePrompt
+
+model_name = "SuanChang/rain-SQLCoder"
+
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    torch_dtype=torch.bfloat16,
+    device_map="auto",
+)
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+question = "What is the name of the department that offers a course that has a description including the word 'Statistics'?"
+schemas = [
+'''CREATE TABLE `course` (
+    `crs_code` STRING,
+    `dept_code` STRING,
+    `crs_description` STRING,
+    `crs_credit` DOUBLE
+);''',
+'''CREATE TABLE `department` (
+    `dept_code` STRING,
+    `dept_name` STRING,
+    `school_code` STRING,
+    `emp_num` INT,
+    `dept_address` STRING,
+    `dept_extension` INT
+);''',
+'''CREATE TABLE `student` (
+    `stu_num` INT,
+    `stu_lname` STRING,
+    `stu_fname` STRING,
+    `stu_init` STRING,
+    `stu_dob` STRING,
+    `stu_hrs` INT,
+    `stu_class` STRING,
+    `stu_gpa` DOUBLE,
+    `stu_transfer` INT,
+    `dept_code` STRING,
+    `stu_phone` INT,
+    `prof_num` INT
+);'''
+]
+hint = "- Today is 2025-02-01."
+data = dict(
+    question=question,
+    schema="\n\n".join(schemas),
+    hint=hint,
+    related_question_sqls=None,
+)
+text, _, _ = SQLGeneratePrompt.prompt(data)
+
+model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
+
+generated_ids = model.generate(
+    **model_inputs,
+    max_new_tokens=32768
+)
+generated_ids = [
+    output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
+]
+response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
+
+print(response)
+
+'''
+```sql
+SELECT d.dept_name FROM department d JOIN course c ON d.dept_code = c.dept_code WHERE c.crs_description LIKE '%Statistics%';
+```
+'''
+````

README.md

@@ -1,3 +1,141 @@
----
-license: apache-2.0
----
+# Introduction
+[Rain's SQLCoder](https://huggingface.co/SuanChang/rain-SQLCoder) is a state-of-the-art large language model (LLM) designed for natural language-to-SparkSQL generation. Rain's SQLCoder, with 32B parameters, is fine-tuned from the [Qwen2.5-Coder-32B-Instruct](https://huggingface.co/Qwen/Qwen2.5-Coder-32B-Instruct). Optimized for natural language-to-SparkSQL conversion tasks, Rain's SQLCoder effectively handles contexts of up to 32k tokens, making it particularly suitable for generating complex and large-scale SQL queries.
+
+<p align="center">
+          🤗 <a href="https://huggingface.co/SuanChang/rain-SQLCoder">Hugging Face</a> | 🖥️ <a href="https://www.suan-chang.com/">Demo</a> | 💬 <a href="./figures/wechat.png">WeChat</a> | <a href="https://github.com/suan-chang/rain-SQLCoder">GitHub</a> 
+</p>
+
+[English](./README.md) | [中文](./README-zh.md)
+
+# Prompt
+Rain's SQLCoder adopted the [Alpaca](https://github.com/tatsu-lab/stanford_alpaca) template, with the following prompt format.
+````
+Below is an instruction that describes a task. 
+Write a response that appropriately completes the request.
+
+### Instruction:
+[BEGIN OF TASK INSTRUCTION]
+You are an expert in composing Spark SQL queries. You are given a user query and a set of table schemas.
+Based on the user query, you need to generate one Spark SQL query to achieve the purpose.
+{task description for date hint and related question and sqls}
+[END OF TASK INSTRUCTION]
+
+[BEGIN OF TABLE SCHEMAS]
+{schemas}
+[END OF TABLE SCHEMAS]
+
+[BEGIN OF GENERATION HINT]
+{date hint}
+[END OF GENERATION HINT]
+
+[BEGIN OF RELATED QUERIES]
+{related question and sqls}
+[END OF RELATED QUERIES]
+
+[BEGIN OF FORMAT INSTRUCTION]
+The output MUST strictly adhere to the following format, and NO other text MUST be included.
+```sql
+your output Spark SQL query
+``` 
+[END OF FORMAT INSTRUCTION]
+
+[BEGIN OF QUERY]
+User Query: {user question}
+[END OF QUERY]
+
+### Response:
+````
+
+# Evaluation
+We followed the evaluation logic from [SQL-Eval](https://github.com/defog-ai/sql-eval) to compare predicted results with ground truth:
+1. If the predicted data frame exactly matches the ground truth data frame, the prediction is considered correct.
+2. If the ground truth SQL does not contain sorting logic, and the predicted data frame matches the ground truth data frame after sorting, the prediction is considered correct.
+3. If the columns in the ground truth data frame are a subset of the predicted data frame, the prediction is considered correct.
+4. In all other cases, the prediction is considered incorrect.
+
+# Experimental Results
+We compared the generation accuracy of Rain's SQLCoder with state-of-the-art natural language large models, both domestic and international, on two test datasets. The Benchmark Dataset contains baseline samples, while the Enhanced Dataset is constructed by applying stratified sampling to 20% of the Benchmark Dataset and supplementing it with relevant user questions and corresponding SparkSQL statements to evaluate the model's performance under enhanced contextual information. The experimental results demonstrate that Rain's SQLCoder exhibits significant advantages in query intent understanding, SQL syntax accuracy, and complex query processing.
+
+## Benchmark Dataset
+<img src="./figures/benchmark_dataset_result.png" alt="benchmark" width=800>
+
+## Enhanced Dataset
+<img src="./figures/enhanced_dataset_result.png" alt="enhanced" width=800>
+
+# Quick Start
+We provide examples here to help you quickly learn how to load and use our model.
+>Tips: Rain's SQLCoder is trained solely for generating `SELECT` statements, and when the table schemas cannot support answering the user's question, the model will refuse to respond.
+
+````python
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from utils.prompt import SQLGeneratePrompt
+
+model_name = "SuanChang/rain-SQLCoder"
+
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    torch_dtype=torch.bfloat16,
+    device_map="auto",
+)
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+question = "What is the name of the department that offers a course that has a description including the word 'Statistics'?"
+schemas = [
+'''CREATE TABLE `course` (
+    `crs_code` STRING,
+    `dept_code` STRING,
+    `crs_description` STRING,
+    `crs_credit` DOUBLE
+);''',
+'''CREATE TABLE `department` (
+    `dept_code` STRING,
+    `dept_name` STRING,
+    `school_code` STRING,
+    `emp_num` INT,
+    `dept_address` STRING,
+    `dept_extension` INT
+);''',
+'''CREATE TABLE `student` (
+    `stu_num` INT,
+    `stu_lname` STRING,
+    `stu_fname` STRING,
+    `stu_init` STRING,
+    `stu_dob` STRING,
+    `stu_hrs` INT,
+    `stu_class` STRING,
+    `stu_gpa` DOUBLE,
+    `stu_transfer` INT,
+    `dept_code` STRING,
+    `stu_phone` INT,
+    `prof_num` INT
+);'''
+]
+hint = "- Today is 2025-02-01."
+data = dict(
+    question=question,
+    schema="\n\n".join(schemas),
+    hint=hint,
+    related_question_sqls=None,
+)
+text, _, _ = SQLGeneratePrompt.prompt(data)
+
+model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
+
+generated_ids = model.generate(
+    **model_inputs,
+    max_new_tokens=32768
+)
+generated_ids = [
+    output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
+]
+response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
+
+print(response)
+
+'''
+```sql
+SELECT d.dept_name FROM department d JOIN course c ON d.dept_code = c.dept_code WHERE c.crs_description LIKE '%Statistics%';
+```
+'''
+````

requirements.txt

@@ -0,0 +1,4 @@
+func_timeout==4.3.5
+pandas==2.2.3
+torch==2.3.1
+transformers==4.43.4

utils/eval.py

@@ -0,0 +1,261 @@
+import re
+import pandas as pd
+from func_timeout import func_timeout, FunctionTimedOut
+import pandas as pd
+from pandas.testing import assert_frame_equal, assert_series_equal
+import collections
+
+
+LIKE_PATTERN = r"LIKE[\s\S]*'"
+
+
+def deduplicate_columns(df: pd.DataFrame) -> pd.DataFrame:
+    cols = df.columns.tolist()
+    if len(cols) != len(set(cols)):
+        duplicates = [
+            item for item, count in collections.Counter(cols).items() if count > 1
+        ]
+        for dup in duplicates:
+            indices = [i for i, x in enumerate(cols) if x == dup]
+            for i in indices:
+                cols[i] = f"{dup}_{i}"
+        df.columns = cols
+    return df
+
+
+def serializate_columns(df: pd.DataFrame):
+    for col in df.columns:
+        if df[col].apply(lambda x: isinstance(x, (list, pd.Series))).any():
+            df[col] = df[col].apply(
+                lambda x: str(sorted(x)) if isinstance(x, (list, pd.Series)) else x
+            )
+    return df
+
+
+def normalize_table(
+    df: pd.DataFrame, query_category: str, if_order: bool, sql: str = None
+) -> pd.DataFrame:
+    """
+    Normalizes a dataframe by:
+    1. removing all duplicate rows
+    2. sorting columns in alphabetical order
+    3. sorting rows using values from first column to last (if query_category is not 'order_by' and question does not ask for ordering)
+    4. resetting index
+    """
+    df = serializate_columns(df)
+
+    # remove duplicate rows, if any
+    df = df.drop_duplicates()
+
+    # sort columns in alphabetical order of column names
+    df = deduplicate_columns(df)
+    sorted_df = df.reindex(sorted(df.columns), axis=1)
+
+    # check if query_category is 'order_by' and if question asks for ordering
+    has_order_by = False
+
+    if query_category == "order_by" or if_order:
+        has_order_by = True
+
+        if sql:
+            # determine which columns are in the ORDER BY clause of the sql generated, using regex
+            pattern = re.compile(r"ORDER BY[\s\S]*", re.IGNORECASE)
+            order_by_clause = re.search(pattern, sql)
+            if order_by_clause:
+                order_by_clause = order_by_clause.group(0)
+                # get all columns in the ORDER BY clause, by looking at the text between ORDER BY and the next semicolon, comma, or parantheses
+                pattern = re.compile(r"(?<=ORDER BY)(.*?)(?=;|,|\)|$)", re.IGNORECASE)
+                order_by_columns = re.findall(pattern, order_by_clause)
+                order_by_columns = (
+                    order_by_columns[0].split() if order_by_columns else []
+                )
+                order_by_columns = [
+                    col.strip().rsplit(".", 1)[-1] for col in order_by_columns
+                ]
+
+                ascending = False
+                # if there is a DESC or ASC in the ORDER BY clause, set the ascending to that
+                if "DESC" in [i.upper() for i in order_by_columns]:
+                    ascending = False
+                elif "ASC" in [i.upper() for i in order_by_columns]:
+                    ascending = True
+
+                # remove whitespace, commas, and parantheses
+                order_by_columns = [col.strip() for col in order_by_columns]
+                order_by_columns = [
+                    col.replace(",", "").replace("(", "") for col in order_by_columns
+                ]
+                order_by_columns = [
+                    i
+                    for i in order_by_columns
+                    if i.lower()
+                    not in ["desc", "asc", "nulls", "last", "first", "limit"]
+                ]
+
+                # get all columns in sorted_df that are not in order_by_columns
+                other_columns = [
+                    i for i in sorted_df.columns.tolist() if i not in order_by_columns
+                ]
+
+                # only choose order_by_columns that are in sorted_df
+                order_by_columns = [
+                    i for i in order_by_columns if i in sorted_df.columns.tolist()
+                ]
+                sorted_df = sorted_df.sort_values(
+                    by=order_by_columns + other_columns, ascending=ascending
+                )
+
+                sorted_df = sorted_df[other_columns + order_by_columns]
+
+    if not has_order_by:
+        # sort rows using values from first column to last
+        sorted_df = sorted_df.sort_values(by=list(sorted_df.columns))
+
+    # reset index
+    sorted_df = deduplicate_columns(sorted_df)
+    sorted_df = sorted_df.reset_index(drop=True)
+    return sorted_df
+
+
+def compare_df(
+    df_gold: pd.DataFrame,
+    df_gen: pd.DataFrame,
+    query_category: str,
+    question: str,
+    query_gold: str = None,
+    query_gen: str = None,
+) -> bool:
+    """
+    Compares two dataframes and returns True if they are the same, else False.
+    query_gold and query_gen are the original queries that generated the respective dataframes.
+    """
+    # drop duplicates to ensure equivalence
+    if df_gen.empty or df_gold.empty:
+        return False
+    try:
+        is_equal = df_gold.values == df_gen.values
+        if is_equal.all():
+            return True
+    except:
+        try:
+            is_equal = df_gold.values == df_gen.values
+            if is_equal:
+                return True
+        except:
+            pass
+
+    pattern = re.compile(r"ORDER BY[\s\S]*", re.IGNORECASE)
+    is_order = re.search(pattern, query_gold)
+
+    df_gold = normalize_table(df_gold, query_category, is_order, query_gold)
+    df_gen = normalize_table(df_gen, query_category, is_order, query_gen)
+
+    # perform same checks again for normalized tables
+    if df_gold.shape != df_gen.shape:
+        return False
+    # fill NaNs with -99999 to handle NaNs in the dataframes for comparison
+    df_gen.fillna(-99999, inplace=True)
+    df_gold.fillna(-99999, inplace=True)
+    is_equal = df_gold.values == df_gen.values
+
+    try:
+        return is_equal.all()
+    except:
+        return is_equal
+
+
+def subset_df(
+    df_sub: pd.DataFrame,
+    df_super: pd.DataFrame,
+    query_category: str,
+    question: str,
+    query_super: str = None,
+    query_sub: str = None,
+    verbose: bool = False,
+) -> bool:
+    """
+    Checks if df_sub is a subset of df_super.
+    """
+    if df_sub.empty and df_super.empty:
+        return True  # handle cases for empty dataframes
+
+    if df_sub.empty:
+        return False
+
+    is_order = False
+    if query_sub:
+        pattern = re.compile(r"ORDER BY[\s\S]*", re.IGNORECASE)
+        is_order = re.search(pattern, query_sub)
+
+    # make a copy of df_super so we don't modify the original while keeping track of matches
+    df_super_temp = df_super.copy(deep=True)
+    matched_columns = []
+    df_sub = deduplicate_columns(df_sub)
+    df_super_temp = deduplicate_columns(df_super_temp)
+    for col_sub_name in df_sub.columns:
+        col_match = False
+        for col_super_name in df_super_temp.columns:
+            col_sub = df_sub[col_sub_name].sort_values().reset_index(drop=True)
+            col_super = (
+                df_super_temp[col_super_name].sort_values().reset_index(drop=True)
+            )
+
+            try:
+                assert_series_equal(
+                    col_sub, col_super, check_dtype=False, check_names=False
+                )
+                col_match = True
+                matched_columns.append(col_super_name)
+                # remove col_super_name to prevent us from matching it again
+                df_super_temp = df_super_temp.drop(columns=[col_super_name])
+                break
+            except AssertionError:
+                continue
+
+        if not col_match:
+            if verbose:
+                print(f"no match for {col_sub_name}")
+            return False
+
+    df_sub_normalized = normalize_table(df_sub, query_category, is_order, query_sub)
+
+    # get matched columns from df_super, and rename them with columns from df_sub, then normalize
+    df_super_matched = df_super[matched_columns].rename(
+        columns=dict(zip(matched_columns, df_sub.columns))
+    )
+    df_super_matched = normalize_table(
+        df_super_matched, query_category, is_order, query_super
+    )
+
+    try:
+        assert_frame_equal(df_sub_normalized, df_super_matched, check_dtype=False)
+        return True
+    except AssertionError:
+        return False
+
+
+def _check_df(
+    gt_df: pd.DataFrame, pre_df: pd.DataFrame, gt_sql: str, pre_sql: str
+) -> bool:
+    try:
+        if gt_df.empty or pre_df.empty:
+            return False
+        result = compare_df(gt_df, pre_df, "", "", gt_sql, pre_sql)
+        if result:
+            return True
+        result = subset_df(gt_df, pre_df, "", "", query_sub=gt_sql, query_super=pre_sql)
+        return result
+    except Exception as e:
+        return False
+
+
+def check_df(
+    gt_df: pd.DataFrame, pre_df: pd.DataFrame, gt_sql: str, pre_sql: str
+) -> bool:
+    try:
+        res = func_timeout(10, _check_df, args=(gt_df, pre_df, gt_sql, pre_sql))
+        return res
+    except FunctionTimedOut as e:
+        return False
+    except Exception as e:
+        return False

utils/prompt.py

@@ -0,0 +1,244 @@
+from abc import ABC, abstractmethod
+from typing import Dict, Any, List, Tuple, Optional, Union
+
+
+class AlpacaTemplate:
+    DEFAULT_SYSTEM = (
+        "Below is an instruction that describes a task. "
+        "Write a response that appropriately completes the request."
+    )
+
+    @classmethod
+    def template(
+        cls,
+        user_content: str,
+        system_content: Union[str, None] = None,
+        response: Union[str, None] = None,
+    ) -> str:
+        template: str = ""
+
+        if system_content:
+            template += f"{system_content}\n\n"
+        else:
+            template += f"{cls.DEFAULT_SYSTEM}\n\n"
+
+        template += f"### Instruction:\n{user_content}\n\n### Response:\n"
+        if response:
+            template += response
+
+        return template
+
+
+class SftPrompt(ABC):
+    @classmethod
+    @abstractmethod
+    def prompt_user_content(cls, data: Dict[str, Any]) -> str:
+        pass
+
+    @classmethod
+    @abstractmethod
+    def prompt_system_content(cls, data: Dict[str, Any]) -> Optional[str]:
+        pass
+
+    @classmethod
+    @abstractmethod
+    def prompt_target(cls, sql: str) -> str:
+        pass
+
+    @classmethod
+    def prompt(cls, data: Dict[str, Any]) -> Tuple[str, Optional[str], str]:
+        user_content = cls.prompt_user_content(data)
+        system_content = cls.prompt_system_content(data)
+        if "spark_sql" in data:
+            target = cls.prompt_target(data["spark_sql"])
+        else:
+            target = None
+        return user_content, system_content, target
+
+
+class SQLGeneratePrompt(SftPrompt):
+    @classmethod
+    def split_res(
+        cls, res: Union[str, List[str], None]
+    ) -> Union[str, List[Union[str, None]], None]:
+        if res == None:
+            return res
+        if isinstance(res, list):
+            for i in range(len(res)):
+                res[i] = cls.split_res(res[i])
+            return res
+        else:
+            res = res.strip()
+            if not (res.startswith("```sql") and res.endswith("```")):
+                return None
+            return res[6:-3].strip()
+
+    @classmethod
+    def compose_extra_task_desc(
+        cls,
+        hint: Union[str, None],
+        related_question_sqls: Union[List[Dict[str, str]], None],
+    ) -> str:
+        if not hint and not related_question_sqls:
+            return ""
+        elif hint and not related_question_sqls:
+            return (
+                "I provide generation hint, you can refer to it to help you generate.\n"
+            )
+        elif not hint and related_question_sqls:
+            return "I provide other user queries related to the user query with the corresponding Spark SQL queries, you can refer to them to help you generate.\n"
+        else:
+            return "I provide generation hint and other user queries related to the user query with the corresponding Spark SQL queries, you can refer to them to help you generate.\n"
+
+    @classmethod
+    def compose_hint_content(cls, hint: Union[str, None]) -> str:
+        if hint == None:
+            return ""
+        if len(hint) == 0:
+            return ""
+        hint = hint.strip()
+        content = (
+            "[BEGIN OF GENERATION HINT]\n" f"{hint}\n" "[END OF GENERATION HINT]\n" "\n"
+        )
+        return content
+
+    @classmethod
+    def compose_related_question_sqls_content(
+        cls, related_question_sqls: Union[List[Dict[str, str]], None]
+    ) -> str:
+        if related_question_sqls == None:
+            return ""
+        if len(related_question_sqls) == 0:
+            return ""
+        content = "[BEGIN OF RELATED QUERIES]\n"
+        for i, question_sql in enumerate(related_question_sqls):
+            question, sql = question_sql["question"], question_sql["spark_sql"]
+
+            question = question.strip()
+            question = question.replace("\n", " ")
+            sql = sql.strip()
+
+            sub_content = (
+                f"# Related Query {i+1}\n"
+                "## User Query\n"
+                f"`{question}`\n"
+                "## Spark SQL Query\n"
+                "```sql\n"
+                f"{sql}\n"
+                "```\n"
+                "\n"
+            )
+            content = content + sub_content
+        content = content.strip() + "\n"
+        content += "[END OF RELATED QUERIES]\n\n"
+        return content
+
+    @classmethod
+    def extract_table_schema(cls, user_content: str) -> str:
+        start_idx = user_content.find("[BEGIN OF TABLE SCHEMAS]") + len(
+            "[BEGIN OF TABLE SCHEMAS]"
+        )
+        end_idx = user_content.find("[END OF TABLE SCHEMAS]")
+        return user_content[start_idx:end_idx].strip()
+
+    @classmethod
+    def extract_user_query(cls, user_content: str) -> str:
+        start_idx = user_content.find("[BEGIN OF QUERY]\nUser Query: ") + len(
+            "[BEGIN OF QUERY]\nUser Query: "
+        )
+        end_idx = user_content.find("[END OF QUERY]")
+        return user_content[start_idx:end_idx].strip()
+
+    @classmethod
+    def extract_hint(cls, user_content: str) -> Union[str, None]:
+        start_idx = user_content.find("[BEGIN OF GENERATION HINT]") + len(
+            "[BEGIN OF GENERATION HINT]"
+        )
+        end_idx = user_content.find("[END OF GENERATION HINT]")
+        if end_idx == -1:
+            return None
+        return user_content[start_idx:end_idx].strip()
+
+    @classmethod
+    def extract_related_question_sqls(
+        cls, user_content: str
+    ) -> Union[List[Dict[str, str]], None]:
+        start_idx = user_content.find("[BEGIN OF RELATED QUERIES]") + len(
+            "[BEGIN OF RELATED QUERIES]"
+        )
+        end_idx = user_content.find("[END OF RELATED QUERIES]")
+        if end_idx == -1:
+            return None
+        related_question_sqls = user_content[start_idx:end_idx].strip().split("\n\n")
+        res = []
+        for question_sql in related_question_sqls:
+            question_start_idx = question_sql.find("`") + 1
+            question_end_idx = question_sql.find("`\n##")
+            question = question_sql[question_start_idx:question_end_idx]
+
+            sql_start_idx = question_sql.find("```sql\n") + 7
+            sql_end_idx = -3
+            sql = question_sql[sql_start_idx:sql_end_idx]
+            res.append({"question": question, "spark_sql": sql})
+        return res
+
+    @classmethod
+    def prompt_user_content(cls, data: Dict[str, Any]) -> str:
+        question, schema = data["question"].strip(), data["schema"].strip()
+        hint, related_question_sqls = (
+            data["hint"],
+            data["related_question_sqls"],
+        )
+        extra_task_desc = cls.compose_extra_task_desc(hint, related_question_sqls)
+        hint_content = cls.compose_hint_content(hint)
+        related_question_sqls_content = cls.compose_related_question_sqls_content(
+            related_question_sqls
+        )
+
+        user_content = (
+            "[BEGIN OF TASK INSTRUCTION]\n"
+            "You are an expert in composing Spark SQL queries. You are given a user query and a set of table schemas.\n"
+            "Based on the user query, you need to generate one Spark SQL query to achieve the purpose.\n"
+            f"{extra_task_desc}"
+            "[END OF TASK INSTRUCTION]\n"
+            "\n"
+            "[BEGIN OF TABLE SCHEMAS]\n"
+            f"{schema}\n"
+            "[END OF TABLE SCHEMAS]\n"
+            "\n"
+            f"{hint_content}"
+            f"{related_question_sqls_content}"
+            "[BEGIN OF FORMAT INSTRUCTION]\n"
+            "The output MUST strictly adhere to the following format, and NO other text MUST be included.\n"
+            "```sql\n"
+            "your output Spark SQL query\n"
+            "```\n"
+            "[END OF FORMAT INSTRUCTION]\n"
+            "\n"
+            "[BEGIN OF QUERY]\n"
+            f"User Query: {question}\n"
+            "[END OF QUERY]\n"
+        )
+
+        user_content = AlpacaTemplate.template(user_content)
+        return user_content
+
+    @classmethod
+    def prompt_system_content(cls, data: Dict[str, Any]) -> Union[str, None]:
+        return None
+
+    @classmethod
+    def prompt_target(cls, sql: Union[str, List[str]]) -> Union[str, List[str]]:
+        def map_func(tmp_sql: str) -> str:
+            tmp_sql = tmp_sql.strip()
+            tmp_sql = tmp_sql.strip(";")
+            tmp_sql = tmp_sql.strip()
+            tmp_sql += ";"
+            return "```sql\n" f"{tmp_sql}\n" "```"
+
+        if isinstance(sql, str):
+            target = map_func(sql)
+        else:
+            target = [map_func(_sql) for _sql in sql]
+
+        return target