Struct SessionContext

pub struct SessionContext { /* private fields */ }

Main interface for executing queries with DataFusion. Maintains the state of the connection between a user and an instance of the DataFusion engine.

See examples below for how to use the SessionContext to execute queries and how to configure the session.

Overview

SessionContext provides the following functionality:

• Create a DataFrame from a CSV or Parquet data source.
• Register a CSV or Parquet data source as a table that can be referenced from a SQL query.
• Register a custom data source that can be referenced from a SQL query.
• Execution a SQL query

Example: DataFrame API

The following example demonstrates how to use the context to execute a query against a CSV data source using the DataFrame API:

use datafusion::prelude::*;
let ctx = SessionContext::new();
let df = ctx.read_csv("tests/data/example.csv", CsvReadOptions::new()).await?;
let df = df.filter(col("a").lt_eq(col("b")))?
           .aggregate(vec![col("a")], vec![min(col("b"))])?
           .limit(0, Some(100))?;
let results = df
  .collect()
  .await?;
assert_batches_eq!(
 &[
   "+---+----------------+",
   "| a | min(?table?.b) |",
   "+---+----------------+",
   "| 1 | 2              |",
   "+---+----------------+",
 ],
 &results
);

Example: SQL API

The following example demonstrates how to execute the same query using SQL:

use datafusion::prelude::*;
let ctx = SessionContext::new();
ctx.register_csv("example", "tests/data/example.csv", CsvReadOptions::new()).await?;
let results = ctx
  .sql("SELECT a, min(b) FROM example GROUP BY a LIMIT 100")
  .await?
  .collect()
  .await?;
assert_batches_eq!(
 &[
   "+---+----------------+",
   "| a | min(example.b) |",
   "+---+----------------+",
   "| 1 | 2              |",
   "+---+----------------+",
 ],
 &results
);

Example: Configuring SessionContext

The SessionContext can be configured by creating a SessionState using SessionStateBuilder:

// Configure a 4k batch size
let config = SessionConfig::new() .with_batch_size(4 * 1024);

// configure a memory limit of 1GB with 20%  slop
 let runtime_env = RuntimeEnvBuilder::new()
    .with_memory_limit(1024 * 1024 * 1024, 0.80)
    .build_arc()
    .unwrap();

// Create a SessionState using the config and runtime_env
let state = SessionStateBuilder::new()
  .with_config(config)
  .with_runtime_env(runtime_env)
  // include support for built in functions and configurations
  .with_default_features()
  .build();

// Create a SessionContext
let ctx = SessionContext::from(state);

Relationship between SessionContext, SessionState, and TaskContext

The state required to optimize, and evaluate queries is broken into three levels to allow tailoring

The objects are:

1. SessionContext: Most users should use a SessionContext. It contains all information required to execute queries including high level APIs such as SessionContext::sql. All queries run with the same SessionContext share the same configuration and resources (e.g. memory limits).

2. SessionState: contains information required to plan and execute an individual query (e.g. creating a LogicalPlan or ExecutionPlan). Each query is planned and executed using its own SessionState, which can be created with SessionContext::state. SessionState allows finer grained control over query execution, for example disallowing DDL operations such as CREATE TABLE.

3. TaskContext contains the state required for query execution (e.g. ExecutionPlan::execute). It contains a subset of information in SessionState. TaskContext allows executing ExecutionPlans PhysicalExprs without requiring a full SessionState.

Implementations

impl SessionContext

pub async fn read_avro<P: DataFilePaths>( &self, table_paths: P, options: AvroReadOptions<'_>, ) -> Result<DataFrame>

Creates a DataFrame for reading an Avro data source.

For more control such as reading multiple files, you can use read_table with a super::ListingTable.

For an example, see read_csv

pub async fn register_avro( &self, table_ref: impl Into<TableReference>, table_path: impl AsRef<str>, options: AvroReadOptions<'_>, ) -> Result<()>

Registers an Avro file as a table that can be referenced from SQL statements executed against this context.

impl SessionContext

pub async fn read_csv<P: DataFilePaths>( &self, table_paths: P, options: CsvReadOptions<'_>, ) -> Result<DataFrame>

Creates a DataFrame for reading a CSV data source.

For more control such as reading multiple files, you can use read_table with a super::ListingTable.

Example usage is given below:

use datafusion::prelude::*;
let ctx = SessionContext::new();
// You can read a single file using `read_csv`
let df = ctx.read_csv("tests/data/example.csv", CsvReadOptions::new()).await?;
// you can also read multiple files:
let df = ctx.read_csv(vec!["tests/data/example.csv", "tests/data/example.csv"], CsvReadOptions::new()).await?;

pub async fn register_csv( &self, table_ref: impl Into<TableReference>, table_path: impl AsRef<str>, options: CsvReadOptions<'_>, ) -> Result<()>

Registers a CSV file as a table which can referenced from SQL statements executed against this context.

pub async fn write_csv( &self, plan: Arc<dyn ExecutionPlan>, path: impl AsRef<str>, ) -> Result<()>

Executes a query and writes the results to a partitioned CSV file.

impl SessionContext

pub async fn read_json<P: DataFilePaths>( &self, table_paths: P, options: NdJsonReadOptions<'_>, ) -> Result<DataFrame>

Creates a DataFrame for reading an JSON data source.

For more control such as reading multiple files, you can use read_table with a super::ListingTable.

For an example, see read_csv

pub async fn register_json( &self, table_ref: impl Into<TableReference>, table_path: impl AsRef<str>, options: NdJsonReadOptions<'_>, ) -> Result<()>

Registers a JSON file as a table that it can be referenced from SQL statements executed against this context.

pub async fn write_json( &self, plan: Arc<dyn ExecutionPlan>, path: impl AsRef<str>, ) -> Result<()>

Executes a query and writes the results to a partitioned JSON file.

impl SessionContext

pub async fn read_parquet<P: DataFilePaths>( &self, table_paths: P, options: ParquetReadOptions<'_>, ) -> Result<DataFrame>

Available on crate feature parquet only.

Creates a DataFrame for reading a Parquet data source.

For more control such as reading multiple files, you can use read_table with a super::ListingTable.

For an example, see read_csv

pub async fn register_parquet( &self, table_ref: impl Into<TableReference>, table_path: impl AsRef<str>, options: ParquetReadOptions<'_>, ) -> Result<()>

Available on crate feature parquet only.

Registers a Parquet file as a table that can be referenced from SQL statements executed against this context.

pub async fn write_parquet( &self, plan: Arc<dyn ExecutionPlan>, path: impl AsRef<str>, writer_properties: Option<WriterProperties>, ) -> Result<()>

Available on crate feature parquet only.

Executes a query and writes the results to a partitioned Parquet file.

impl SessionContext

pub fn new() -> Self

Creates a new SessionContext using the default SessionConfig.

pub async fn refresh_catalogs(&self) -> Result<()>

Finds any ListingSchemaProviders and instructs them to reload tables from “disk”

pub fn new_with_config(config: SessionConfig) -> Self

Creates a new SessionContext using the provided SessionConfig and a new RuntimeEnv.

See Self::new_with_config_rt for more details on resource limits.

pub fn new_with_config_rt( config: SessionConfig, runtime: Arc<RuntimeEnv>, ) -> Self

Creates a new SessionContext using the provided SessionConfig and a RuntimeEnv.

Resource Limits

By default, each new SessionContext creates a new RuntimeEnv, and therefore will not enforce memory or disk limits for queries run on different SessionContexts.

To enforce resource limits (e.g. to limit the total amount of memory used) across all DataFusion queries in a process, all SessionContext’s should be configured with the same RuntimeEnv.

pub fn new_with_state(state: SessionState) -> Self

Creates a new SessionContext using the provided SessionState

pub fn enable_url_table(self) -> Self

Enable querying local files as tables.

This feature is security sensitive and should only be enabled for systems that wish to permit direct access to the file system from SQL.

When enabled, this feature permits direct access to arbitrary files via SQL like

SELECT * from 'my_file.parquet'

See DynamicFileCatalog for more details

let ctx = SessionContext::new()
  .enable_url_table(); // permit local file access
let results = ctx
  .sql("SELECT a, MIN(b) FROM 'tests/data/example.csv' as example GROUP BY a LIMIT 100")
  .await?
  .collect()
  .await?;
assert_batches_eq!(
 &[
   "+---+----------------+",
   "| a | min(example.b) |",
   "+---+----------------+",
   "| 1 | 2              |",
   "+---+----------------+",
 ],
 &results
);

pub fn into_state_builder(self) -> SessionStateBuilder

Convert the current SessionContext into a SessionStateBuilder

This is useful to switch back to SessionState with custom settings such as Self::enable_url_table.

Avoids cloning the SessionState if possible.

Example

let my_rule = PushDownFilter{}; // pretend it is a new rule
// Create a new builder with a custom optimizer rule
let context: SessionContext = SessionStateBuilder::new()
  .with_optimizer_rule(Arc::new(my_rule))
  .build()
  .into();
// Enable local file access and convert context back to a builder
let builder = context
  .enable_url_table()
  .into_state_builder();

pub fn session_start_time(&self) -> DateTime<Utc>

Returns the time this SessionContext was created

pub fn with_function_factory( self, function_factory: Arc<dyn FunctionFactory>, ) -> Self

Registers a FunctionFactory to handle CREATE FUNCTION statements

pub fn add_optimizer_rule( &self, optimizer_rule: Arc<dyn OptimizerRule + Send + Sync>, )

Adds an optimizer rule to the end of the existing rules.

See SessionState for more control of when the rule is applied.

pub fn add_analyzer_rule( &self, analyzer_rule: Arc<dyn AnalyzerRule + Send + Sync>, )

Adds an analyzer rule to the end of the existing rules.

See SessionState for more control of when the rule is applied.

pub fn register_object_store( &self, url: &Url, object_store: Arc<dyn ObjectStore>, ) -> Option<Arc<dyn ObjectStore>>

Registers an ObjectStore to be used with a specific URL prefix.

See RuntimeEnv::register_object_store for more details.

Example: register a local object store for the “file://” URL prefix

let object_store_url = ObjectStoreUrl::parse("file://").unwrap();
let object_store = object_store::local::LocalFileSystem::new();
let ctx = SessionContext::new();
// All files with the file:// url prefix will be read from the local file system
ctx.register_object_store(object_store_url.as_ref(), Arc::new(object_store));

pub fn register_batch( &self, table_name: &str, batch: RecordBatch, ) -> Result<Option<Arc<dyn TableProvider>>>

Registers the RecordBatch as the specified table name

pub fn runtime_env(&self) -> Arc<RuntimeEnv>

Return the RuntimeEnv used to run queries with this SessionContext

pub fn session_id(&self) -> String

Returns an id that uniquely identifies this SessionContext.

pub fn table_factory( &self, file_type: &str, ) -> Option<Arc<dyn TableProviderFactory>>

Return the TableProviderFactory that is registered for the specified file type, if any.

pub fn enable_ident_normalization(&self) -> bool

Return the enable_ident_normalization of this Session

pub fn copied_config(&self) -> SessionConfig

Return a copied version of config for this Session

pub fn copied_table_options(&self) -> TableOptions

Return a copied version of table options for this Session

pub async fn sql(&self, sql: &str) -> Result<DataFrame>

Creates a DataFrame from SQL query text.

Note: This API implements DDL statements such as CREATE TABLE and CREATE VIEW and DML statements such as INSERT INTO with in-memory default implementations. See Self::sql_with_options.

Example: Running SQL queries

See the example on Self

Example: Creating a Table with SQL

use datafusion::prelude::*;
let ctx = SessionContext::new();
ctx
  .sql("CREATE TABLE foo (x INTEGER)")
  .await?
  .collect()
  .await?;
assert!(ctx.table_exist("foo").unwrap());

pub async fn sql_with_options( &self, sql: &str, options: SQLOptions, ) -> Result<DataFrame>

Creates a DataFrame from SQL query text, first validating that the queries are allowed by options

Example: Preventing Creating a Table with SQL

If you want to avoid creating tables, or modifying data or the session, set SQLOptions appropriately:

use datafusion::prelude::*;
let ctx = SessionContext::new();
let options = SQLOptions::new()
  .with_allow_ddl(false);
let err = ctx.sql_with_options("CREATE TABLE foo (x INTEGER)", options)
  .await
  .unwrap_err();
assert!(
  err.to_string().starts_with("Error during planning: DDL not supported: CreateMemoryTable")
);

pub fn parse_sql_expr(&self, sql: &str, df_schema: &DFSchema) -> Result<Expr>

Creates logical expressions from SQL query text.

Example: Parsing SQL queries

// datafusion will parse number as i64 first.
let sql = "a > 10";
let expected = col("a").gt(lit(10 as i64));
// provide type information that `a` is an Int32
let schema = Schema::new(vec![Field::new("a", DataType::Int32, true)]);
let df_schema = DFSchema::try_from(schema).unwrap();
let expr = SessionContext::new()
 .parse_sql_expr(sql, &df_schema)?;
assert_eq!(expected, expr);

pub async fn execute_logical_plan(&self, plan: LogicalPlan) -> Result<DataFrame>

Execute the LogicalPlan, return a DataFrame. This API is not featured limited (so all SQL such as CREATE TABLE and COPY will be run).

If you wish to limit the type of plan that can be run from SQL, see Self::sql_with_options and SQLOptions::verify_plan.

pub fn create_physical_expr( &self, expr: Expr, df_schema: &DFSchema, ) -> Result<Arc<dyn PhysicalExpr>>

Create a PhysicalExpr from an Expr after applying type coercion and function rewrites.

Note: The expression is not simplified or otherwise optimized: a = 1 + 2 will not be simplified to a = 3 as this is a more involved process. See the expr_api example for how to simplify expressions.

Example

// a = 1 (i64)
let expr = col("a").eq(lit(1i64));
// provide type information that `a` is an Int32
let schema = Schema::new(vec![Field::new("a", DataType::Int32, true)]);
let df_schema = DFSchema::try_from(schema).unwrap();
// Create a PhysicalExpr. Note DataFusion automatically coerces (casts) `1i64` to `1i32`
let physical_expr = SessionContext::new()
  .create_physical_expr(expr, &df_schema).unwrap();

See Also

• SessionState::create_physical_expr for a lower level API

pub fn register_variable( &self, variable_type: VarType, provider: Arc<dyn VarProvider + Send + Sync>, )

Registers a variable provider within this context.

pub fn register_udtf(&self, name: &str, fun: Arc<dyn TableFunctionImpl>)

Register a table UDF with this context

pub fn register_udf(&self, f: ScalarUDF)

Registers a scalar UDF within this context.

Note in SQL queries, function names are looked up using lowercase unless the query uses quotes. For example,

• SELECT MY_FUNC(x)... will look for a function named "my_func"
• SELECT "my_FUNC"(x) will look for a function named "my_FUNC"

Any functions registered with the udf name or its aliases will be overwritten with this new function

pub fn register_udaf(&self, f: AggregateUDF)

Registers an aggregate UDF within this context.

Note in SQL queries, aggregate names are looked up using lowercase unless the query uses quotes. For example,

• SELECT MY_UDAF(x)... will look for an aggregate named "my_udaf"
• SELECT "my_UDAF"(x) will look for an aggregate named "my_UDAF"

pub fn register_udwf(&self, f: WindowUDF)

Registers a window UDF within this context.

Note in SQL queries, window function names are looked up using lowercase unless the query uses quotes. For example,

• SELECT MY_UDWF(x)... will look for a window function named "my_udwf"
• SELECT "my_UDWF"(x) will look for a window function named "my_UDWF"

pub fn deregister_udf(&self, name: &str)

Deregisters a UDF within this context.

pub fn deregister_udaf(&self, name: &str)

Deregisters a UDAF within this context.

pub fn deregister_udwf(&self, name: &str)

Deregisters a UDWF within this context.

pub fn deregister_udtf(&self, name: &str)

Deregisters a UDTF within this context.

pub async fn read_arrow<P: DataFilePaths>( &self, table_paths: P, options: ArrowReadOptions<'_>, ) -> Result<DataFrame>

Creates a DataFrame for reading an Arrow data source.

For more control such as reading multiple files, you can use read_table with a ListingTable.

For an example, see read_csv

pub fn read_empty(&self) -> Result<DataFrame>

Creates an empty DataFrame.

pub fn read_table(&self, provider: Arc<dyn TableProvider>) -> Result<DataFrame>

Creates a DataFrame for a TableProvider such as a ListingTable or a custom user defined provider.

pub fn read_batch(&self, batch: RecordBatch) -> Result<DataFrame>

Creates a DataFrame for reading a RecordBatch

pub fn read_batches( &self, batches: impl IntoIterator<Item = RecordBatch>, ) -> Result<DataFrame>

Create a DataFrame for reading a [Vec[RecordBatch]]

pub async fn register_listing_table( &self, table_ref: impl Into<TableReference>, table_path: impl AsRef<str>, options: ListingOptions, provided_schema: Option<SchemaRef>, sql_definition: Option<String>, ) -> Result<()>

Registers a ListingTable that can assemble multiple files from locations in an ObjectStore instance into a single table.

This method is async because it might need to resolve the schema.

pub async fn register_arrow( &self, name: &str, table_path: &str, options: ArrowReadOptions<'_>, ) -> Result<()>

Registers an Arrow file as a table that can be referenced from SQL statements executed against this context.

pub fn register_catalog( &self, name: impl Into<String>, catalog: Arc<dyn CatalogProvider>, ) -> Option<Arc<dyn CatalogProvider>>

Registers a named catalog using a custom CatalogProvider so that it can be referenced from SQL statements executed against this context.

Returns the CatalogProvider previously registered for this name, if any

pub fn catalog_names(&self) -> Vec<String>

Retrieves the list of available catalog names.

pub fn catalog(&self, name: &str) -> Option<Arc<dyn CatalogProvider>>

Retrieves a CatalogProvider instance by name

pub fn register_table( &self, table_ref: impl Into<TableReference>, provider: Arc<dyn TableProvider>, ) -> Result<Option<Arc<dyn TableProvider>>>

Registers a TableProvider as a table that can be referenced from SQL statements executed against this context.

If a table of the same name was already registered, returns “Table already exists” error.

pub fn deregister_table( &self, table_ref: impl Into<TableReference>, ) -> Result<Option<Arc<dyn TableProvider>>>

Deregisters the given table.

Returns the registered provider, if any

pub fn table_exist(&self, table_ref: impl Into<TableReference>) -> Result<bool>

Return true if the specified table exists in the schema provider.

pub async fn table( &self, table_ref: impl Into<TableReference>, ) -> Result<DataFrame>

Retrieves a DataFrame representing a table previously registered by calling the register_table function.

Returns an error if no table has been registered with the provided reference.

pub fn table_function(&self, name: &str) -> Result<Arc<TableFunction>>

Retrieves a TableFunction reference by name.

Returns an error if no table function has been registered with the provided name.

pub async fn table_provider( &self, table_ref: impl Into<TableReference>, ) -> Result<Arc<dyn TableProvider>>

Return a TableProvider for the specified table.

pub fn task_ctx(&self) -> Arc<TaskContext>

Get a new TaskContext to run in this session

pub fn state(&self) -> SessionState

Return a new SessionState suitable for executing a single query.

Notes:

1. query_execution_start_time is set to the current time for the returned state.

2. The returned state is not shared with the current session state and this changes to the returned SessionState such as changing ConfigOptions will not be reflected in this SessionContext.

pub fn state_ref(&self) -> Arc<RwLock<SessionState>>

Get reference to SessionState

pub fn state_weak_ref(&self) -> Weak<RwLock<SessionState>>

Get weak reference to SessionState

pub fn register_catalog_list(&self, catalog_list: Arc<dyn CatalogProviderList>)

Register CatalogProviderList in SessionState

pub fn register_table_options_extension<T: ConfigExtension>(&self, extension: T)

Registers a ConfigExtension as a table option extension that can be referenced from SQL statements executed against this context.

Trait Implementations

impl Clone for SessionContext

fn clone(&self) -> SessionContext

Returns a copy of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

impl Default for SessionContext

fn default() -> Self

Returns the “default value” for a type.

impl From<&SessionContext> for TaskContext

Create a new task context instance from SessionContext

fn from(session: &SessionContext) -> Self

Converts to this type from the input type.

impl From<SessionContext> for SessionStateBuilder

fn from(session: SessionContext) -> Self

Converts to this type from the input type.

impl From<SessionState> for SessionContext

fn from(state: SessionState) -> Self

Converts to this type from the input type.

impl FunctionRegistry for SessionContext

fn udfs(&self) -> HashSet<String>

Set of all available udfs.

fn udf(&self, name: &str) -> Result<Arc<ScalarUDF>>

Returns a reference to the user defined scalar function (udf) named name.

fn udaf(&self, name: &str) -> Result<Arc<AggregateUDF>>

Returns a reference to the user defined aggregate function (udaf) named name.

fn udwf(&self, name: &str) -> Result<Arc<WindowUDF>>

Returns a reference to the user defined window function (udwf) named name.

fn register_udf( &mut self, udf: Arc<ScalarUDF>, ) -> Result<Option<Arc<ScalarUDF>>>

Registers a new ScalarUDF, returning any previously registered implementation.

fn register_udaf( &mut self, udaf: Arc<AggregateUDF>, ) -> Result<Option<Arc<AggregateUDF>>>

Registers a new AggregateUDF, returning any previously registered implementation.

fn register_udwf( &mut self, udwf: Arc<WindowUDF>, ) -> Result<Option<Arc<WindowUDF>>>

Registers a new WindowUDF, returning any previously registered implementation.

fn register_function_rewrite( &mut self, rewrite: Arc<dyn FunctionRewrite + Send + Sync>, ) -> Result<()>

Registers a new FunctionRewrite with the registry.

fn expr_planners(&self) -> Vec<Arc<dyn ExprPlanner>>

Set of all registered ExprPlanners

fn register_expr_planner( &mut self, expr_planner: Arc<dyn ExprPlanner>, ) -> Result<()>

Registers a new ExprPlanner with the registry.

fn deregister_udf( &mut self, _name: &str, ) -> Result<Option<Arc<ScalarUDF>>, DataFusionError>

Deregisters a ScalarUDF, returning the implementation that was deregistered.

fn deregister_udaf( &mut self, _name: &str, ) -> Result<Option<Arc<AggregateUDF>>, DataFusionError>

Deregisters a AggregateUDF, returning the implementation that was deregistered.

fn deregister_udwf( &mut self, _name: &str, ) -> Result<Option<Arc<WindowUDF>>, DataFusionError>

Deregisters a WindowUDF, returning the implementation that was deregistered.